pax_global_header00006660000000000000000000000064152106726250014520gustar00rootroot0000000000000052 comment=3050b0a7bd48e04f853027c5fa1f5ab7bc20b856 tvm-ffi-0.1.12/000077500000000000000000000000001521067262500131515ustar00rootroot00000000000000tvm-ffi-0.1.12/.asf.yaml000066400000000000000000000042541521067262500146710ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. github: description: "Open ABI and FFI for Machine Learning Systems" homepage: https://tvm.apache.org/ffi labels: - ffi - gpu - machine-learning features: # Enable issue management issues: true # Enable projects for project management boards projects: true protected_branches: # NOTE: when updating the config # make sure to first update the branch protection name # in a separate branch(e.g. dev), confirm things works # before changing it to main main: # note: do not set status check for now # committer will check ci status before merging required_pull_request_reviews: required_approving_review_count: 1 enabled_merge_buttons: # enable squash button: squash: true # default commit message when merging with a squash commit # can either be: DEFAULT | PR_TITLE | PR_TITLE_AND_COMMIT_DETAILS | PR_TITLE_AND_DESC squash_commit_message: PR_TITLE_AND_DESC # enable merge button: merge: false # default commit message when merging with a merge commit # can either be: DEFAULT | PR_TITLE | PR_TITLE_AND_DESC merge_commit_message: DEFAULT # disable rebase button: rebase: false notifications: commits: commits@tvm.apache.org issues: discuss-archive@tvm.apache.org pullrequests: discuss-archive@tvm.apache.org jobs: discuss-archive@tvm.apache.org discussions: discuss-archive@tvm.apache.org tvm-ffi-0.1.12/.clang-format000066400000000000000000000004401521067262500155220ustar00rootroot00000000000000# Run the following command to reformat a file: # clang-format -i -style=Google # Or use clang-format-diff to only reformat the changed lines: # https://clang.llvm.org/docs/ClangFormat.html BasedOnStyle: Google DerivePointerAlignment: false ColumnLimit: 100 PointerAlignment: Left tvm-ffi-0.1.12/.clang-tidy000066400000000000000000000014751521067262500152140ustar00rootroot00000000000000# -clang-analyzer-unix.Malloc is disabled due to false positives Checks: > clang-diagnostic-*, clang-analyzer-*, modernize-*, bugprone-*, performance-*, portability-*, google-*, -modernize-use-trailing-return-type, -modernize-concat-nested-namespaces, -modernize-use-auto, -modernize-use-nodiscard, -modernize-return-braced-init-list, -modernize-avoid-c-arrays, -bugprone-easily-swappable-parameters, -google-build-using-namespace, -performance-enum-size, -clang-analyzer-unix.Malloc, -performance-avoid-endl WarningsAsErrors: '*' HeaderFilterRegex: '.*(include/tvm/ffi|src/ffi)/.*\.h$' HeaderFileExtensions: ['h'] ImplementationFileExtensions: ['cc'] CheckOptions: - key: misc-include-cleaner.UnusedIncludes value: 'true' - key: misc-include-cleaner.MissingIncludes value: 'true' tvm-ffi-0.1.12/.claude/000077500000000000000000000000001521067262500144645ustar00rootroot00000000000000tvm-ffi-0.1.12/.claude/launch.json000066400000000000000000000006651521067262500166400ustar00rootroot00000000000000{ "version": "0.0.1", "configurations": [ { "name": "docs", "runtimeExecutable": "uv", "runtimeArgs": [ "run", "--group", "docs", "sphinx-autobuild", "docs", "docs/_build/html", "--ignore", "docs/reference/cpp/generated" ], "port": 8000, "env": { "BUILD_CPP_DOCS": "1", "BUILD_RUST_DOCS": "1" } } ] } tvm-ffi-0.1.12/.claude/skills/000077500000000000000000000000001521067262500157655ustar00rootroot00000000000000tvm-ffi-0.1.12/.claude/skills/devtools/000077500000000000000000000000001521067262500176245ustar00rootroot00000000000000tvm-ffi-0.1.12/.claude/skills/devtools/SKILL.md000066400000000000000000000216651521067262500210360ustar00rootroot00000000000000 --- name: devtools description: Developer reference for Apache TVM-FFI. argument-hint: "[lint | cpp | python | docs]" --- # TVM-FFI Developer Guide Condensed reference from `docs/dev/`. Use this when working on the TVM-FFI codebase. ## Prerequisites - **Python**: 3.9+ (managed via `uv`; default virtualenv at `.venv`) - **Compiler**: C++17-capable toolchain (GCC/Clang on Linux, Apple Clang on macOS, MSVC on Windows) - **Build tools**: CMake 3.18+, Ninja - **Source**: Always clone with `--recursive`, or run `git submodule update --init --recursive` All Python-related commands below use [`uv`](https://docs.astral.sh/uv/). The default virtual environment is `.venv` in the repo root. --- ## 1. Run Linters and clang-tidy ### Pre-commit hooks (primary linting workflow) Install and register git hooks: ```bash uv tool install pre-commit pre-commit install # register git hooks (runs automatically before each commit) ``` Run hooks manually: ```bash pre-commit run --all-files # all hooks on every file pre-commit run # all hooks on staged files only pre-commit run --all-files # single hook, e.g. ruff-check, clang-format ``` **Linters by language:** | Language | Formatter | Linter / Type Checker | |----------|---------------------------------|----------------------------------| | Python | `ruff` (format) | `ruff` (lint), `ty` (type check) | | C/C++ | `clang-format` | `clang-tidy` (see below) | | Cython | `double-quote-cython-strings` | `cython-lint` | | CMake | `cmake-format` | `cmake-lint` | | Shell | `shfmt` | `shellcheck` | | YAML | `yamllint` | | | TOML | `taplo-format` | | | Markdown | `markdownlint-cli2` | | | RST | `rstcheck` | | **Troubleshooting pre-commit:** - **Version problems**: Ensure pre-commit 2.18.0+ (`pre-commit --version`). - **Stale cache**: Run `pre-commit clean` to clear the hook cache. - **Auto-fixed files**: Most formatting hooks fix issues in place. Review changes, stage with `git add -u`, and commit again. ### clang-tidy (separate from pre-commit) `clang-tidy` is **not** a pre-commit hook. It runs as a separate CI job and only checks changed C++ files. To reproduce locally: ```bash # Run on specific files uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=$(sysctl -n hw.ncpu) \ include/tvm/ffi/c_api.h src/some_file.cc # Run on all C++ sources uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=$(sysctl -n hw.ncpu) \ ./src/ ./include ./tests ``` > On macOS, `clang-tidy` is resolved via `xcrun`; the wrapper script handles this automatically. > On Linux, replace `$(sysctl -n hw.ncpu)` with `$(nproc)`. --- ## 2. Build and Test the C++ Package ### Build the C++ library (no Python) ```bash cmake . -B build_cpp -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build_cpp --parallel --config RelWithDebInfo --target tvm_ffi_shared cmake --install build_cpp --config RelWithDebInfo --prefix ./dist ``` After installation: - Headers: `dist/include/` - Libraries: `dist/lib/` ### Build and run C++ tests Set parallel build level first: ```bash export CMAKE_BUILD_PARALLEL_LEVEL=$(sysctl -n hw.ncpu) # macOS # export CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) # Linux ``` Configure, build, and run: ```bash cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests ctest -V -C Debug --test-dir build_test --output-on-failure ``` ### Key CMake options | Option | Default | Description | |----------------------------------|---------|----------------------------------------------| | `TVM_FFI_BUILD_TESTS` | `OFF` | Enable C++ test targets | | `TVM_FFI_ATTACH_DEBUG_SYMBOLS` | `OFF` | Attach debug symbols in release mode | | `TVM_FFI_USE_LIBBACKTRACE` | `ON` | Enable libbacktrace | | `TVM_FFI_USE_EXTRA_CXX_API` | `ON` | Enable extra C++ API in shared lib | | `TVM_FFI_BACKTRACE_ON_SEGFAULT` | `ON` | Print backtrace on segfault | | `CMAKE_EXPORT_COMPILE_COMMANDS` | `OFF` | Generate `compile_commands.json` for clangd | > On Windows, run from a **Developer Command Prompt for VS** or ensure the MSVC toolchain is on your `PATH`. --- ## 3. Build and Test the Python Package ### Build (editable install) The Python build uses `scikit-build-core` which drives CMake to compile the C++ core and Cython extension: ```bash uv pip install --force-reinstall --verbose -e . ``` - `--force-reinstall` forces a full rebuild. - `-e` (editable) means pure Python changes are reflected immediately without rebuilding. - **C++/Cython changes always require re-running this command.** Pass CMake flags via `--config-settings`: ```bash uv pip install --force-reinstall --verbose -e . \ --config-settings cmake.define.TVM_FFI_ATTACH_DEBUG_SYMBOLS=ON ``` Verify the install: ```bash uv run python -c "import tvm_ffi; print(tvm_ffi.__version__)" ``` ### Update Python stubs After building (or after C++/Cython reflection changes), regenerate inline type stubs: ```bash uv run tvm-ffi-stubgen python ``` This updates inline stub blocks (between `tvm-ffi-stubgen(begin)` / `tvm-ffi-stubgen(end)` markers) inside `.py` files with type annotations derived from the C++ reflection registry (field types, method signatures, global function schemas). ### Run Python tests Install with test dependencies, then run pytest: ```bash uv pip install --force-reinstall --verbose --group test -e . uv run pytest -vvs tests/python ``` ### Run Rust tests Rust tests require the Python package to be installed first (Rust FFI bindings link against the built shared library): ```bash cd rust && cargo test ``` ### Troubleshooting - **Rebuilding after C++/Cython changes**: Re-run `uv pip install --force-reinstall -e .`. Editable installs only auto-reflect pure Python changes. - **Submodules missing**: Run `git submodule update --init --recursive` from the repo root. - **Library not found at import time**: Ensure the dynamic loader can find the shared library. Add the `lib` directory to `LD_LIBRARY_PATH` (Linux), `DYLD_LIBRARY_PATH` (macOS), or `PATH` (Windows). --- ## 4. Build Documentation Website Building documentation requires the Python package to be installed first (see section 3). ### Interactive build (auto-reload) Serves locally at `http://127.0.0.1:8000` with live reload: ```bash uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/cpp/generated ``` ### One-off build ```bash uv run --group docs sphinx-build -M html docs docs/_build ``` ### With C++ API reference (requires Doxygen) ```bash # Install Doxygen brew install doxygen # macOS # sudo apt install doxygen # Linux # Interactive BUILD_CPP_DOCS=1 uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/cpp/generated --watch include # One-off BUILD_CPP_DOCS=1 uv run --group docs sphinx-build -M html docs docs/_build ``` ### With Rust API reference (requires cargo) ```bash # Interactive BUILD_RUST_DOCS=1 uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/rust/generated --watch rust # One-off BUILD_RUST_DOCS=1 uv run --group docs sphinx-build -M html docs docs/_build ``` ### Build all documentation ```bash BUILD_CPP_DOCS=1 BUILD_RUST_DOCS=1 uv run --group docs sphinx-build \ -M html docs docs/_build ``` ### Cleanup ```bash rm -rf docs/_build/ docs/reference/python/generated docs/reference/cpp/generated docs/reference/rust/generated ``` tvm-ffi-0.1.12/.cmake-format.json000066400000000000000000000034511521067262500164730ustar00rootroot00000000000000{ "format": { "disable": false, "line_width": 100, "tab_size": 2, "use_tabchars": false, "fractional_tab_policy": "use-space", "max_subgroups_hwrap": 2, "max_pargs_hwrap": 6, "max_rows_cmdline": 2, "separate_ctrl_name_with_space": true, "separate_fn_name_with_space": false, "dangle_parens": true, "dangle_align": "prefix", "min_prefix_chars": 4, "max_prefix_chars": 10, "max_lines_hwrap": 2, "line_ending": "unix", "command_case": "canonical", "keyword_case": "unchanged", "always_wrap": [], "enable_sort": true, "autosort": false, "require_valid_layout": false, "layout_passes": {} }, "markup": { "bullet_char": "*", "enum_char": ".", "first_comment_is_literal": true, "literal_comment_pattern": null, "fence_pattern": "^\\s*([`~]{3}[`~]*)(.*)$", "ruler_pattern": "^\\s*[^\\w\\s]{3}.*[^\\w\\s]{3}$", "explicit_trailing_pattern": "#<", "hashruler_min_length": 10, "canonicalize_hashrulers": true, "enable_markup": true }, "lint": { "disabled_codes": [], "function_pattern": "[0-9a-z_]+", "macro_pattern": "[0-9A-Z_]+", "global_var_pattern": "[A-Z][0-9A-Z_]+", "internal_var_pattern": "_[A-Z][0-9A-Z_]+", "local_var_pattern": "[a-z][a-z0-9_]+", "private_var_pattern": "_[0-9a-z_]+", "public_var_pattern": "[A-Z][0-9A-Z_]+", "argument_var_pattern": "[a-z][a-z0-9_]+", "keyword_pattern": "[A-Z][0-9A-Z_]+", "max_conditionals_custom_parser": 2, "min_statement_spacing": 1, "max_statement_spacing": 2, "max_returns": 6, "max_branches": 12, "max_arguments": 5, "max_localvars": 15, "max_statements": 50 }, "encode": { "emit_byteorder_mark": false, "input_encoding": "utf-8", "output_encoding": "utf-8" } } tvm-ffi-0.1.12/.github/000077500000000000000000000000001521067262500145115ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/000077500000000000000000000000001521067262500161515ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/build-orcjit-wheel/000077500000000000000000000000001521067262500216425ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/build-orcjit-wheel/action.yml000066400000000000000000000106411521067262500236440ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Build OrcJIT Wheel description: > Build and test OrcJIT wheels for a given OS/architecture combination. Handles LLVM caching, conda installation, and cibuildwheel execution. inputs: arch: description: "Target architecture (e.g., x86_64, aarch64, arm64, AMD64)" required: true build: description: "cibuildwheel build selector (e.g., cp312-manylinux_x86_64)" required: true checkout_ref: description: "Branch, tag, or SHA to check out before building" required: true runs: using: "composite" steps: - name: Check out source uses: actions/checkout@v5 with: ref: ${{ inputs.checkout_ref }} submodules: recursive - uses: ./.github/actions/detect-env-vars id: env_vars # ---- Cache LLVM prefix ---- - name: Cache LLVM uses: actions/cache@v4 id: llvm-cache with: path: ${{ runner.os == 'Windows' && 'C:/opt/llvm' || '/opt/llvm' }} key: llvm-22.1.0-${{ runner.os }}-${{ inputs.arch }}-v3 # ---- Install LLVM via conda (cache miss only) ---- - name: Setup conda if: steps.llvm-cache.outputs.cache-hit != 'true' uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0 continue-on-error: true id: conda1 with: miniforge-version: latest - name: Setup conda (retry with tar.bz2) if: steps.llvm-cache.outputs.cache-hit != 'true' && steps.conda1.outcome == 'failure' uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0 with: miniforge-version: latest use-only-tar-bz2: true - name: Create /opt/llvm (macOS) if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'macOS' shell: bash run: sudo mkdir -p /opt/llvm && sudo chown -R $(whoami) /opt/llvm - name: Install LLVM (Unix) if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os != 'Windows' shell: bash -l {0} run: | conda create -q -p /opt/llvm -c conda-forge \ llvmdev=22.1.0 clangdev=22.1.0 compiler-rt=22.1.0 zlib zstd-static \ -y - name: Install LLVM (Windows) if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'Windows' shell: cmd /C call {0} run: | conda create -q -p C:\opt\llvm -c conda-forge llvmdev=22.1.0 zlib zstd-static -y # ---- Build and test wheels ---- - name: Build and test wheels uses: pypa/cibuildwheel@298ed2fb2c105540f5ed055e8a6ad78d82dd3a7e # v3.3.1 with: package-dir: addons/tvm_ffi_orcjit output-dir: wheelhouse env: CIBW_BUILD: ${{ inputs.build }} CIBW_ARCHS_LINUX: ${{ inputs.arch }} CIBW_ARCHS_MACOS: ${{ inputs.arch }} CIBW_ARCHS_WINDOWS: ${{ inputs.arch }} CIBW_BUILD_VERBOSITY: 1 CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} CIBW_ENVIRONMENT: LLVM_PREFIX=/opt/llvm CIBW_ENVIRONMENT_WINDOWS: LLVM_PREFIX="C:/opt/llvm" CIBW_CONTAINER_ENGINE: "docker; create_args: --volume /opt/llvm:/opt/llvm" CIBW_TEST_REQUIRES: pytest ninja CIBW_TEST_COMMAND: >- pip install --force-reinstall {project} && pytest {project}/addons/tvm_ffi_orcjit/tests -v && python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang cpp && python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang c CIBW_TEST_COMMAND_WINDOWS: >- pip install --force-reinstall {project} && pytest {project}/addons/tvm_ffi_orcjit/tests -v && python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang c tvm-ffi-0.1.12/.github/actions/build-wheel-for-publish/000077500000000000000000000000001521067262500226025ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/build-wheel-for-publish/action.yml000066400000000000000000000101531521067262500246020ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Build Wheel For Publish description: > Build and test wheels (and optionally sdists) for a given OS/architecture/manylinux combination, suitable for both validation and publish workflows. inputs: os: description: "Runner operating system (e.g., ubuntu-latest, macos-14, windows-latest)" required: true arch: description: "Target architecture for cibuildwheel (e.g., x86_64, aarch64, arm64, AMD64)" required: true linux_image: description: "Manylinux image tag to use on Linux runners (empty string for non-Linux runners)" required: true checkout_ref: description: "Branch, tag, or SHA to check out before building" required: true build_wheels: description: "Set to false to skip wheel builds (useful for sdist-only runs)" default: "true" required: false build_sdist: description: "Set to true to build a source distribution and run metadata checks" default: "false" required: false runs: using: "composite" steps: # Special handling for macOS arm64 + python 3.8. # Install Python 3.8 from `actions/setup-python`, which is an arm64 build. - name: Install Python 3.8 for macOS arm64 if: runner.os == 'macOS' && inputs.arch == 'arm64' uses: actions/setup-python@v5 with: python-version: 3.8 - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 - name: Check out source uses: actions/checkout@v5 with: ref: ${{ inputs.checkout_ref }} submodules: recursive fetch-depth: 0 fetch-tags: true - uses: ./.github/actions/detect-env-vars id: env_vars - name: Print current commit shell: bash run: git log -1 --oneline - name: Run cpp tests if: runner.os != 'Windows' env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} shell: bash run: | cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests ctest -V -C Debug --test-dir build_test --output-on-failure - name: Run cpp tests[windows] if: ${{ runner.os == 'Windows' }} shell: cmd env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: > cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug && cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests && ctest -V -C Debug --test-dir build_test --output-on-failure - name: Build wheels if: ${{ inputs.build_wheels == 'true' }} uses: pypa/cibuildwheel@298ed2fb2c105540f5ed055e8a6ad78d82dd3a7e # v3.3.1 env: CIBW_ARCHS_MACOS: ${{ inputs.arch }} CIBW_ARCHS_LINUX: ${{ inputs.arch }} CIBW_ARCHS_WINDOWS: ${{ inputs.arch }} CIBW_MANYLINUX_X86_64_IMAGE: ${{ inputs.linux_image }} CIBW_MANYLINUX_AARCH64_IMAGE: ${{ inputs.linux_image }} CIBW_BUILD_VERBOSITY: 1 CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} with: package-dir: . output-dir: wheelhouse - name: Build sdist if: ${{ inputs.build_sdist == 'true' }} shell: bash run: | uv tool run --from build pyproject-build --sdist --outdir dist . uv tool run twine check dist/* tvm-ffi-0.1.12/.github/actions/detect-env-vars/000077500000000000000000000000001521067262500211605ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/detect-env-vars/action.yml000066400000000000000000000027241521067262500231650ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Detect Environment Variables description: Detects environment variables such as CPU count and sets them as outputs. runs: using: "composite" steps: - name: Run Python to detect environment variables shell: python id: detect run: | import multiprocessing, os output_file = open(os.environ.get('GITHUB_OUTPUT'), 'a') def write_env_var(name, value): output_file.write(f"{name}={value}\n") print(f"Detected environment variable: {name}={value}") write_env_var('cpu_count', multiprocessing.cpu_count()) outputs: cpu_count: description: "The number of CPU cores" value: '${{ steps.detect.outputs.cpu_count }}' tvm-ffi-0.1.12/.github/actions/detect-skip-ci/000077500000000000000000000000001521067262500207565ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/actions/detect-skip-ci/action.yaml000066400000000000000000000112521521067262500231200ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Detect Skip CI description: > Detects if commit message contains [skip ci] to skip CI. Also detects if changes are docs-only and can be skipped. inputs: github_event_name: description: 'GitHub event name (push or pull_request)' required: true pr_base_ref: description: 'Base branch ref of the pull request (only required for pull_request events)' required: false pr_head_sha: description: 'Head SHA of the pull request (only required for pull_request events)' required: false runs: using: "composite" steps: - name: Checkout code uses: actions/checkout@v5 with: fetch-depth: 0 - name: Detect [skip ci] and docs-only changes shell: bash id: detect run: | set -e echo "Event: ${{ inputs.github_event_name }}" if [[ "${{ inputs.github_event_name }}" == "workflow_dispatch" ]]; then echo "should_skip_ci_commit=false" >> $GITHUB_OUTPUT echo "should_skip_ci_docs_only=false" >> $GITHUB_OUTPUT exit 0 fi if [[ "${{ inputs.github_event_name }}" == "pull_request" ]]; then echo "Fetching base branch '${{ inputs.pr_base_ref }}' for diff..." # Get commit message from PR head commit first if ! commit_message=$(git log -1 --pretty=%B ${{ inputs.pr_head_sha }} 2>/dev/null); then echo "Warning: Failed to get commit message, assuming no skip" commit_message="" fi echo "Commit message: $commit_message" if ! git fetch origin "${{ inputs.pr_base_ref }}" 2>/dev/null; then echo "Warning: Failed to fetch base branch, assuming all files changed" changed_files=$(git ls-files) else # Get list of changed files between base and PR head changed_files=$(git diff --name-only origin/${{ inputs.pr_base_ref }} ${{ inputs.pr_head_sha }}) fi else # For push event, get head commit message and diff against previous commit if ! commit_message=$(git log -1 --pretty=%B 2>/dev/null); then echo "Warning: Failed to get commit message, assuming no skip" commit_message="" fi echo "Commit message: $commit_message" # Handle case where there's no previous commit (initial commit) if ! changed_files=$(git diff --name-only HEAD^ HEAD 2>/dev/null); then echo "Warning: No previous commit found, assuming all files changed" changed_files=$(git ls-files) fi fi echo "Changed files:" echo "$changed_files" # Check for [bypass ci] in commit message if [[ "$commit_message" == "[bypass ci]"* ]]; then echo "[INFO] [bypass ci] detected in commit message, skipping CI ..." echo "should_skip_ci_commit=true" >> $GITHUB_OUTPUT else echo "should_skip_ci_commit=false" >> $GITHUB_OUTPUT fi # check if changes are docs-only # Check if all changed files are docs-only: # Customize the docs paths/extensions below as needed are_changes_docs_only=true for file in $changed_files; do if [[ ! "$file" =~ ^docs/ ]] && [[ ! "$file" =~ \.md$ ]]; then are_changes_docs_only=false break fi done if [[ "$are_changes_docs_only" == "true" ]]; then echo "[INFO] Changes are docs-only, skipping tests ..." echo "should_skip_ci_docs_only=true" >> $GITHUB_OUTPUT else echo "should_skip_ci_docs_only=false" >> $GITHUB_OUTPUT fi outputs: should_skip_ci_commit: description: 'Set to true if [bypass ci] detected in commit message' value: '${{ steps.detect.outputs.should_skip_ci_commit }}' should_skip_ci_docs_only: description: 'Set to true if changes are docs-only and CI can be skipped' value: '${{ steps.detect.outputs.should_skip_ci_docs_only }}' tvm-ffi-0.1.12/.github/workflows/000077500000000000000000000000001521067262500165465ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/workflows/ci_mainline_only.yml000066400000000000000000000143571521067262500226130ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: mainline-only on: workflow_dispatch: push: branches: - main jobs: prepare: name: Prepare runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: fetch-depth: 0 fetch-tags: true - name: Detect skip ci and docs changes id: detect uses: ./.github/actions/detect-skip-ci with: github_event_name: ${{ github.event_name }} pr_base_ref: ${{ github.event.pull_request.base.ref || '' }} pr_head_sha: ${{ github.event.pull_request.head.sha || '' }} clang-tidy: needs: [prepare] runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: submodules: recursive fetch-depth: 0 fetch-tags: true - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 - uses: ./.github/actions/detect-env-vars id: env_vars - name: Run clang-tidy run: | uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=${{ steps.env_vars.outputs.cpu_count }} \ ./src/ ./include ./tests build-wheels: name: Build wheels needs: [prepare] runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: include: - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux2014, build_sdist: "true"} - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux_2_28, build_sdist: "false"} - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux2014, build_sdist: "false"} - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux_2_28, build_sdist: "false"} - {os: windows-latest, arch: AMD64, linux_image: "", build_sdist: "false"} - {os: macos-14, arch: arm64, linux_image: "", build_sdist: "false"} steps: - name: Checkout repository (for local composite action) uses: actions/checkout@v5 with: fetch-depth: 1 - name: Build wheel uses: ./.github/actions/build-wheel-for-publish with: os: ${{ matrix.os }} arch: ${{ matrix.arch }} linux_image: ${{ matrix.linux_image }} checkout_ref: ${{ github.sha }} build_sdist: ${{ matrix.build_sdist }} examples: name: Run examples needs: [prepare] runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: include: - {os: ubuntu-latest, python_version: "3.14"} - {os: macos-14, python_version: "3.13"} - {os: windows-latest, python_version: "3.12"} steps: - uses: actions/checkout@v5 with: submodules: recursive fetch-depth: 0 fetch-tags: true - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 with: python-version: ${{ matrix.python_version }} activate-environment: true - uses: ./.github/actions/detect-env-vars id: env_vars - name: Install dependencies for examples env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | uv pip install --reinstall --verbose --group torch -e . uv pip install --reinstall --verbose numpy scikit-build-core torch-c-dlpack-ext - name: Run example/quickstart (CPU) [posix] if: ${{ runner.os != 'Windows' }} env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | pushd examples/quickstart rm -rf build bash run_all_cpu.sh popd - name: Run example/quickstart (CPU) [windows] if: ${{ runner.os == 'Windows' }} shell: cmd env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | cd examples\quickstart call run_all_cpu.bat - name: Run example/stable_c_abi [posix] if: ${{ runner.os != 'Windows' }} env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | pushd examples/stable_c_abi rm -rf build bash run_all.sh popd - name: Run example/stable_c_abi [windows] if: ${{ runner.os == 'Windows' }} shell: cmd env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | cd examples\stable_c_abi call run_all.bat - name: Run example/python_packaging [posix] if: ${{ runner.os != 'Windows' }} env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | pushd examples/python_packaging # This directory will be auto-generated in `CMakeLists.txt` by setting `STUB_INIT ON` rm -rf python/my_ffi_extension uv pip install --verbose . --no-build-isolation python run_example.py popd - name: Run example/python_packaging [windows] if: ${{ runner.os == 'Windows' }} shell: pwsh env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | Set-Location examples/python_packaging Remove-Item -Recurse -Force python/my_ffi_extension uv pip install --verbose . --no-build-isolation python run_example.py tvm-ffi-0.1.12/.github/workflows/ci_test.yml000066400000000000000000000213211521067262500207220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: CI on: workflow_dispatch: pull_request: push: branches: - main # dev branch is used for testing purposes - dev jobs: prepare: name: Prepare runs-on: ubuntu-latest outputs: should_skip_ci_commit: ${{ steps.detect.outputs.should_skip_ci_commit }} should_skip_ci_docs_only: ${{ steps.detect.outputs.should_skip_ci_docs_only }} cpp_changed: ${{ steps.cpp_files.outputs.changed }} cpp_files: ${{ steps.cpp_files.outputs.files }} orcjit_changed: ${{ steps.orcjit_files.outputs.changed }} steps: - uses: actions/checkout@v5 with: fetch-depth: 0 fetch-tags: true - name: Detect skip ci and docs changes id: detect uses: ./.github/actions/detect-skip-ci with: github_event_name: ${{ github.event_name }} pr_base_ref: ${{ github.event.pull_request.base.ref || '' }} pr_head_sha: ${{ github.event.pull_request.head.sha || '' }} - name: Get changed C++ files id: cpp_files run: | FILES=$(git diff --name-only --diff-filter=ACMR origin/${{ github.base_ref }}...HEAD -- \ src/ tests/ | grep -E '\.(c|cc|cpp|cxx)$' | tr '\n' ' ') echo "files=$FILES" >> $GITHUB_OUTPUT [ -n "$FILES" ] && echo "changed=true" >> $GITHUB_OUTPUT || echo "changed=false" >> $GITHUB_OUTPUT - name: Get changed OrcJIT files id: orcjit_files run: | FILES=$(git diff --name-only --diff-filter=ACMR origin/${{ github.base_ref }}...HEAD -- \ addons/tvm_ffi_orcjit/ .github/actions/build-orcjit-wheel/ | tr '\n' ' ') echo "files=$FILES" >> $GITHUB_OUTPUT [ -n "$FILES" ] && echo "changed=true" >> $GITHUB_OUTPUT || echo "changed=false" >> $GITHUB_OUTPUT lint: needs: [prepare] runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: fetch-depth: 0 fetch-tags: true - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 - name: Set up Python environment run: uv sync --group dev --no-install-project - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 clang-tidy: needs: [prepare] if: needs.prepare.outputs.cpp_changed == 'true' runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: submodules: recursive fetch-depth: 0 fetch-tags: true - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 with: python-version: 3.13 - uses: ./.github/actions/detect-env-vars id: env_vars - name: Run clang-tidy run: | uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=${{ steps.env_vars.outputs.cpu_count }} \ ${{ needs.prepare.outputs.cpp_files }} doc: needs: [lint, prepare] runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: submodules: recursive fetch-depth: 0 fetch-tags: true - name: Set up uv uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 with: python-version: 3.13 - name: Generate docs run: | sudo apt install doxygen BUILD_CPP_DOCS=1 uv run --group docs sphinx-build -W --keep-going -b html docs docs/_build/html test: needs: [lint, prepare] if: > needs.prepare.outputs.should_skip_ci_commit != 'true' && needs.prepare.outputs.should_skip_ci_docs_only != 'true' name: ${{ matrix.os }} runs-on: ${{ matrix.os }} strategy: fail-fast: true matrix: include: - {os: ubuntu-latest, arch: x86_64, python_version: '3.14t'} - {os: ubuntu-24.04-arm, arch: aarch64, python_version: '3.8'} - {os: windows-latest, arch: AMD64, python_version: '3.9'} - {os: macos-14, arch: arm64, python_version: '3.13'} steps: - uses: actions/checkout@v5 with: submodules: recursive fetch-depth: 0 fetch-tags: true - name: Print current commit run: git log -1 --oneline # Detect CPU count - uses: ./.github/actions/detect-env-vars id: env_vars - name: Print CPU count run: | echo "CPU count: ${{ steps.env_vars.outputs.cpu_count }}" # Run C++ tests - name: Run cpp tests if: ${{ matrix.os != 'windows-latest' }} env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests ctest -V -C Debug --test-dir build_test --output-on-failure - name: Run cpp tests[windows] if: ${{ matrix.os == 'windows-latest' }} shell: cmd env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: > cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug && cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests && ctest -V -C Debug --test-dir build_test --output-on-failure - name: Locate and Set VsDevCmd Path [windows] if: ${{ matrix.os == 'windows-latest' }} shell: pwsh run: | # Captures the output (path) from your Python script $vsPath = python .github/workflows/utils/locate_vsdevcmd_bat.py # Sets an environment variable for all subsequent steps in this job "VS_DEV_CMD_PATH=$vsPath" | Out-File -FilePath $env:GITHUB_ENV -Append # Run Python tests - name: Setup Python ${{ matrix.python_version }} uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 with: python-version: ${{ matrix.python_version }} activate-environment: true - name: Build and install python package env: CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} run: | uv pip install --reinstall --verbose --group test -e . - name: Run python tests if: ${{ matrix.os != 'windows-latest' }} run: | pytest -vvs tests/python - name: Run python tests [windows] if: ${{ matrix.os == 'windows-latest' }} shell: cmd run: | call "%VS_DEV_CMD_PATH%" pytest -vvs tests/python # Run Rust tests, must happen after installing the pip package. - name: Run rust tests working-directory: rust run: | cargo test orcjit: needs: [lint, prepare] if: > needs.prepare.outputs.orcjit_changed == 'true' && needs.prepare.outputs.should_skip_ci_commit != 'true' && needs.prepare.outputs.should_skip_ci_docs_only != 'true' name: orcjit ${{ matrix.os }} (${{ matrix.arch }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: include: - {os: ubuntu-latest, arch: x86_64, build: "cp312-manylinux_x86_64"} - {os: ubuntu-24.04-arm, arch: aarch64, build: "cp312-manylinux_aarch64"} - {os: macos-14, arch: arm64, build: "cp312-macosx_arm64"} - {os: windows-latest, arch: AMD64, build: "cp312-win_amd64"} steps: - uses: actions/checkout@v5 with: fetch-depth: 1 - name: Build OrcJIT wheel uses: ./.github/actions/build-orcjit-wheel with: arch: ${{ matrix.arch }} build: ${{ matrix.build }} checkout_ref: ${{ github.sha }} - uses: actions/upload-artifact@v4 with: name: cibw-orcjit-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }} path: wheelhouse/*.whl tvm-ffi-0.1.12/.github/workflows/publish_orcjit_wheel.yml000066400000000000000000000052241521067262500235000ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Publish OrcJIT wheel on: workflow_dispatch: inputs: tag: description: "Tag to publish (manual run)" required: true jobs: build_wheels: name: ${{ matrix.os }} (${{ matrix.arch }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: include: - {os: ubuntu-latest, arch: x86_64, build: "cp312-manylinux_x86_64"} - {os: ubuntu-24.04-arm, arch: aarch64, build: "cp312-manylinux_aarch64"} - {os: macos-14, arch: arm64, build: "cp312-macosx_arm64"} - {os: windows-latest, arch: AMD64, build: "cp312-win_amd64"} steps: - name: Checkout repository (for local composite action) uses: actions/checkout@v5 with: fetch-depth: 1 - name: Build OrcJIT wheel uses: ./.github/actions/build-orcjit-wheel with: arch: ${{ matrix.arch }} build: ${{ matrix.build }} checkout_ref: ${{ inputs.tag }} - name: Upload wheels uses: actions/upload-artifact@v4 with: name: cibw-orcjit-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }} path: ./wheelhouse/*.whl upload_pypi: needs: [build_wheels] runs-on: ubuntu-latest environment: pypi permissions: id-token: write attestations: write if: github.event_name == 'workflow_dispatch' steps: - uses: actions/download-artifact@v4 with: pattern: cibw-orcjit-* path: dist merge-multiple: true - name: Generate artifact attestation for wheels uses: actions/attest-build-provenance@v1 with: subject-path: dist/* - name: Publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 with: attestations: true verbose: true tvm-ffi-0.1.12/.github/workflows/publish_wheel.yml000066400000000000000000000064351521067262500221330ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: Publish wheel on: workflow_dispatch: inputs: tag: description: "Tag to publish (manual run)" required: true jobs: build_wheels: runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: include: - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux2014, build_sdist: "true"} - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux_2_28, build_sdist: "false"} - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux2014, build_sdist: "false"} - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux_2_28, build_sdist: "false"} - {os: windows-latest, arch: AMD64, linux_image: "", build_sdist: "false"} - {os: macos-14, arch: arm64, linux_image: "", build_sdist: "false"} steps: - name: Checkout repository (for local composite action) uses: actions/checkout@v5 with: fetch-depth: 1 - name: Build wheel uses: ./.github/actions/build-wheel-for-publish with: os: ${{ matrix.os }} arch: ${{ matrix.arch }} linux_image: ${{ matrix.linux_image }} checkout_ref: ${{ inputs.tag }} build_sdist: ${{ matrix.build_sdist }} - name: Upload wheels uses: actions/upload-artifact@v4 with: name: cibw-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }} path: ./wheelhouse/*.whl - name: Upload sdist if: ${{ matrix.build_sdist == 'true' }} uses: actions/upload-artifact@v4 with: name: cibw-sdist path: dist/*.tar.gz upload_pypi: needs: [build_wheels] runs-on: ubuntu-latest environment: pypi permissions: id-token: write attestations: write if: github.event_name == 'workflow_dispatch' # <-- publish only on manual trigger steps: - uses: actions/download-artifact@v4 with: # unpacks all CIBW artifacts into dist/ pattern: cibw-* path: dist merge-multiple: true - name: Generate artifact attestation for sdist and wheels uses: actions/attest-build-provenance@v1 with: subject-path: dist/* - name: Publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 with: attestations: true verbose: true # testing publish url # repository-url: https://test.pypi.org/legacy/ tvm-ffi-0.1.12/.github/workflows/torch_c_dlpack.yml000066400000000000000000000146561521067262500222440ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. name: torch c dlpack on: workflow_dispatch: inputs: branch: description: "Branch or tag to publish (manual run)" required: true jobs: build_wheels_linux: strategy: fail-fast: false matrix: arch: ["x86_64", "aarch64"] python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"] runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }} steps: - uses: jlumbroso/free-disk-space@main if: ${{ matrix.arch == 'x86_64'}} - uses: actions/checkout@v5 with: repository: apache/tvm-ffi ref: main fetch-depth: 0 submodules: recursive path: tvm-ffi - name: prepare docker environment run: | docker pull quay.io/pypa/manylinux_2_28_${{ matrix.arch }}:latest docker run --name build_lib -d quay.io/pypa/manylinux_2_28_${{ matrix.arch }}:latest tail -f /dev/null docker cp tvm-ffi build_lib:/tvm-ffi wget https://developer.download.nvidia.com/compute/cuda/13.0.2/local_installers/cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm > /dev/null 2>&1 docker cp cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm build_lib:/ rm cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm docker exec build_lib bash -c "rpm -i /cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm \ && dnf clean all \ && dnf -y install cuda-toolkit-13-0 \ && rm /cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch }}.rpm \ && dnf clean all" - name: build torch libs and wheels run: | docker exec -w /tvm-ffi build_lib bash ./addons/torch_c_dlpack_ext/build_aot_wheels.sh ${{ matrix.arch }} ${{ matrix.python-version }} - name: collect built wheels run: | mkdir wheelhouse docker cp build_lib:/tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/ . - uses: actions/upload-artifact@v4 with: name: pypi-wheels-linux-${{ matrix.arch }}-${{ matrix.python-version }} path: ./wheelhouse/*.whl build_wheels_windows: strategy: fail-fast: false matrix: arch: ["x86_64"] python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"] runs-on: windows-latest steps: - uses: actions/checkout@v5 with: repository: apache/tvm-ffi ref: main fetch-depth: 0 submodules: recursive path: tvm-ffi - uses: Jimver/cuda-toolkit@6008063726ffe3309d1b22e413d9e88fed91a2f2 id: cuda-toolkit - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 - name: build wheels env: CUDA_HOME: ${{steps.cuda-toolkit.outputs.CUDA_PATH}} shell: cmd working-directory: ./tvm-ffi run: ./addons/torch_c_dlpack_ext/build_aot_wheels.bat ${{ matrix.arch }} ${{ matrix.python-version }} - uses: actions/upload-artifact@v4 with: name: pypi-wheels-windows-${{ matrix.arch }}-${{ matrix.python-version }} path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl build_wheels_macos: strategy: fail-fast: false matrix: arch: ["arm64"] python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"] runs-on: macos-14 steps: - uses: actions/checkout@v5 with: repository: apache/tvm-ffi ref: main fetch-depth: 0 submodules: recursive path: tvm-ffi - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 - name: build torch libs and wheels working-directory: ./tvm-ffi run: | bash ./addons/torch_c_dlpack_ext/build_aot_wheels.sh ${{ matrix.arch }} ${{ matrix.python-version }} - uses: actions/upload-artifact@v4 with: name: pypi-wheels-macos-${{ matrix.arch }}-${{ matrix.python-version }} path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl build_source: runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 with: repository: apache/tvm-ffi ref: main fetch-depth: 0 submodules: recursive path: tvm-ffi - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # v7.3.1 with: python-version: "3.12" activate-environment: true - name: build source run: | uv pip install build cd ./tvm-ffi/addons/torch_c_dlpack_ext python -m build -s - uses: actions/upload-artifact@v4 with: name: pypi-source path: ./tvm-ffi/addons/torch_c_dlpack_ext/dist/*tar.gz upload_pypi: needs: [build_wheels_linux, build_wheels_windows, build_wheels_macos, build_source] runs-on: ubuntu-latest environment: pypi permissions: id-token: write attestations: write if: github.event_name == 'workflow_dispatch' # <-- publish only on manual trigger steps: - uses: actions/download-artifact@v4 with: pattern: pypi-* path: dist merge-multiple: true - name: Generate artifact attestation for sdist and wheels uses: actions/attest-build-provenance@v1 with: subject-path: dist/* - name: Publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 with: attestations: true verbose: true # testing publish url # repository-url: https://test.pypi.org/legacy/ tvm-ffi-0.1.12/.github/workflows/utils/000077500000000000000000000000001521067262500177065ustar00rootroot00000000000000tvm-ffi-0.1.12/.github/workflows/utils/locate_vsdevcmd_bat.py000066400000000000000000000040501521067262500242470ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Locate the VsDevCmd.bat file for the current Visual Studio installation.""" import os import subprocess from pathlib import Path def main() -> None: """Locate the VsDevCmd.bat file for the current Visual Studio installation. Raise exception if not found. If found, print the path to stdout. """ # Path to vswhere.exe vswhere_path = str( Path(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)")) / "Microsoft Visual Studio" / "Installer" / "vswhere.exe" ) if not Path(vswhere_path).exists(): raise FileNotFoundError("vswhere.exe not found.") # Find the Visual Studio installation path vs_install_path = subprocess.run( [ vswhere_path, "-latest", "-prerelease", "-products", "*", "-property", "installationPath", ], capture_output=True, text=True, check=True, ).stdout.strip() if not vs_install_path: raise FileNotFoundError("No Visual Studio installation found.") # Construct the path to the VsDevCmd.bat file vsdevcmd_path = str(Path(vs_install_path) / "Common7" / "Tools" / "VsDevCmd.bat") print(vsdevcmd_path) if __name__ == "__main__": main() tvm-ffi-0.1.12/.gitignore000066400000000000000000000071541521067262500151500ustar00rootroot00000000000000# version file is auto-generated python/tvm_ffi/_version.py /tmp/ *.bak # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class .DS_Store *.S # C extensions *.so build/ *.ll .npm # Distribution / packaging .Python env/ build/ build-*/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ wheelhouse/ pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST .conda/ # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Generated by python/gen_requirements.py python/requirements/*.txt # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ /Testing/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ docs/_staging/ # PyBuilder target/ /target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule # Allow composite actions under .github/actions despite build-* ignore pattern !.github/actions/build-wheel-for-publish/ !.github/actions/build-wheel-for-publish/action.yml !.github/actions/build-orcjit-wheel/ !.github/actions/build-orcjit-wheel/action.yml celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject *~ *.pyc *~ config.mk /config.cmake Win32 *.dir perf *.wasm .emscripten ## IOS DerivedData/ ## Java *.class *.worksheet *.idea *.iml *.classpath *.project *.settings */node_modules/ ## Various settings *.pbxuser !default.pbxuser *.mode1v3 !default.mode1v3 *.mode2v3 !default.mode2v3 *.perspectivev3 !default.perspectivev3 xcuserdata/ .pkl_memoize_* .emscripten* .m2 # Compiled Dynamic libraries *.so *.dylib *.dll # Compiled Object files *.slo *.lo *.o *.obj # Precompiled Headers *.gch *.pch # Compiled Static libraries *.lai *.la *.a *.lib # Executables *.exe *.out *.app ## Other *.moved-aside *.xccheckout *.xcscmblueprint .DS_Store tags cscope* *.lock # vim temporary files *.swp *.swo .bash_history # *.json *.params *.ro *.onnx *.h5 # Mac OS X .DS_Store # Jetbrain .idea .ipython .jupyter .nv .pylint.d .python_history .pytest_cache .local cmake-build-debug # Visual Studio .vs # Visual Studio Code .vscode # clangd language server .clangd # tmp file .nfs* # keys *.pem *.p12 *.pfx *.cer *.crt *.der # patch sentinel patched.txt # Python type checking .mypy_cache/ .pyre/ # pipenv files Pipfile Pipfile.lock # conda package artifacts conda/Dockerfile.cuda* conda/pkg .node_repl_history # nix files .envrc *.nix # Docker files .sudo_as_admin_successful # Local docs build _docs/ .config/configstore/ .ci-py-scripts/ # Used in CI to communicate between Python and Jenkins .docker-image-names/ # GDB history file .gdb_history build/ *.cubin *.fatbin build_cpp/ tvm-ffi-0.1.12/.gitmodules000066400000000000000000000003241521067262500153250ustar00rootroot00000000000000[submodule "3rdparty/dlpack"] path = 3rdparty/dlpack url = https://github.com/dmlc/dlpack [submodule "3rdparty/libbacktrace"] path = 3rdparty/libbacktrace url = https://github.com/ianlancetaylor/libbacktrace tvm-ffi-0.1.12/.markdownlint-cli2.yaml000066400000000000000000000015041521067262500174530ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. ignores: - ".claude/**" config: MD013: false tvm-ffi-0.1.12/.pre-commit-config.yaml000066400000000000000000000074361521067262500174440ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. default_install_hook_types: - pre-commit repos: - repo: local hooks: - id: check-asf-header name: check ASF Header entry: python tests/lint/check_asf_header.py --check language: python language_version: python3 pass_filenames: false verbose: false - repo: local hooks: - id: check-file-type name: check file types entry: python tests/lint/check_file_type.py language: python language_version: python3 pass_filenames: false verbose: false - repo: local hooks: - id: check-version-consistency name: check version consistency entry: python tests/lint/check_version.py --cpp # TODO: add `--rust` once Rust binding matures language: python language_version: python3 additional_dependencies: - setuptools-scm - packaging - tomli pass_filenames: false verbose: false - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 hooks: - id: check-added-large-files - id: check-case-conflict - id: check-merge-conflict - id: check-symlinks - id: end-of-file-fixer - id: mixed-line-ending - id: requirements-txt-fixer - id: trailing-whitespace - id: check-yaml - id: check-toml - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.9 hooks: - id: ruff-check types_or: [python, pyi, jupyter] args: [--fix] - id: ruff-format types_or: [python, pyi, jupyter] - repo: local hooks: - id: ty name: ty check entry: uvx ty@0.0.15 check --error-on-warning language: system pass_filenames: false types: [python] - repo: https://github.com/pre-commit/mirrors-clang-format rev: "v21.1.7" hooks: - id: clang-format - repo: https://github.com/adrienverge/yamllint rev: v1.37.1 hooks: - id: yamllint args: - --config-file - .yamllint.yaml - repo: https://github.com/ComPWA/taplo-pre-commit rev: v0.9.3 hooks: - id: taplo-format - repo: https://github.com/MarcoGorelli/cython-lint rev: v0.18.1 hooks: - id: cython-lint args: [--max-line-length=120] - id: double-quote-cython-strings - repo: https://github.com/scop/pre-commit-shfmt rev: v3.12.0-2 hooks: - id: shfmt args: [--indent=2] - repo: https://github.com/shellcheck-py/shellcheck-py rev: v0.11.0.1 hooks: - id: shellcheck - repo: https://github.com/DavidAnson/markdownlint-cli2 rev: v0.20.0 hooks: - id: markdownlint-cli2 - repo: https://github.com/rstcheck/rstcheck rev: v6.2.5 hooks: - id: rstcheck additional_dependencies: - rstcheck[sphinx] args: - --config - docs/.rstcheck.cfg - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 hooks: - id: cmake-format - id: cmake-lint tvm-ffi-0.1.12/.yamllint.yaml000066400000000000000000000017471521067262500157550ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. extends: default rules: document-start: disable line-length: max: 120 level: warning truthy: allowed-values: - "on" - "off" - "yes" - "no" - "true" - "false" tvm-ffi-0.1.12/3rdparty/000077500000000000000000000000001521067262500147215ustar00rootroot00000000000000tvm-ffi-0.1.12/3rdparty/dlpack/000077500000000000000000000000001521067262500161575ustar00rootroot00000000000000tvm-ffi-0.1.12/3rdparty/libbacktrace/000077500000000000000000000000001521067262500173275ustar00rootroot00000000000000tvm-ffi-0.1.12/CLAUDE.md000066400000000000000000000150451521067262500144350ustar00rootroot00000000000000 # CLAUDE.md — Apache TVM FFI ## What is this project? TVM FFI is an open ABI and FFI (Foreign Function Interface) for machine learning systems. It provides a stable C ABI, C++17 API, Python bindings (via Cython), and Rust bindings. The core abstractions are type-erased values (`Any`/`AnyView`), reference-counted objects (`Object`/`ObjectRef`), and packed functions (`Function`). ## Repository layout ```text include/tvm/ffi/ C++ public headers (core API) src/ffi/ C++ implementation python/tvm_ffi/ Python package (Cython bindings in cython/) rust/ Rust crate workspace (tvm-ffi, tvm-ffi-sys, tvm-ffi-macros) tests/cpp/ GoogleTest C++ tests tests/python/ pytest Python tests tests/lint/ Lint scripts (ASF header, file type, version check) docs/ Sphinx documentation (RST + Markdown) examples/ Runnable examples cmake/Utils/ CMake utility modules 3rdparty/ Vendored deps (dlpack, libbacktrace) addons/ Optional addons (torch_c_dlpack_ext) ``` ## Building All Python commands use `uv`. The default virtualenv is `.venv` in the repo root. ### Python editable install (primary workflow) ```bash uv pip install --force-reinstall --verbose -e . ``` C++/Cython changes always require re-running this command. Pure Python changes are reflected immediately. ### Update Python stubs After building (or after C++/Cython reflection changes), regenerate inline type stubs: ```bash uv run tvm-ffi-stubgen python ``` This updates inline stub blocks (between `tvm-ffi-stubgen(begin)` / `tvm-ffi-stubgen(end)` markers) inside `.py` files with type annotations derived from the C++ reflection registry (field types, method signatures, global function schemas). ### C++-only build ```bash cmake . -B build_cpp -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build_cpp --parallel --config RelWithDebInfo --target tvm_ffi_shared ``` ### Prerequisites - Python 3.9+, C++17 compiler, CMake 3.18+, Ninja - Submodules: `git submodule update --init --recursive` ## Testing ### C++ tests ```bash cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests ctest -V -C Debug --test-dir build_test --output-on-failure ``` ### Python tests ```bash uv pip install --force-reinstall --verbose --group test -e . uv run pytest -vvs tests/python ``` ### Rust tests ```bash cd rust && cargo test # requires Python package installed first ``` ## Linting ### Pre-commit (primary lint workflow) ```bash pre-commit run --all-files # all hooks pre-commit run --all-files # single hook (e.g. ruff-check, clang-format) ``` Key linters: `ruff` (Python), `clang-format` (C++, Google style, 100-col), `cython-lint`, `cmake-format`, `shfmt`/`shellcheck`, `markdownlint-cli2`. ### clang-tidy (separate from pre-commit) ```bash uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=$(sysctl -n hw.ncpu) \ ./src/ ./include ./tests ``` ## Code conventions ### C++ - Source files: `.cc` (not `.cpp`). Headers: `.h`. - Style: Google (via clang-format), 100-col limit, pointer-left (`int* ptr`). - Namespaces: `namespace tvm { namespace ffi { ... } }` (no C++17 nested form). - Header guards: `#ifndef TVM_FFI__H_`. - Object pattern: `FooObj` (data) + `Foo` (ref wrapper extending `ObjectRef`). - `TVM_FFI_DECLARE_OBJECT_INFO("key", FooObj, ParentObj)` in the Obj class. - `TVM_FFI_DEFINE_OBJECT_REF_METHODS(Foo, ParentRef, FooObj)` in the Ref class. - Errors: `TVM_FFI_THROW(ErrorType) << "message"`. - Doc comments: Doxygen (`/*! \brief ... */`). - Every file needs an Apache 2.0 license header. ### Python - `from __future__ import annotations` at the top of every file. - Style: `ruff` (100-col, double quotes, Google docstrings). - Register objects: `@register_object("type.Key")`. - Register functions: `register_global_func("name", fn)`. ### Commit messages Tag-based style: `[FEAT]`, `[FIX]`, `[ERROR]`, `[TEST]`, `[CORE]`, `[EXTRA]`, etc. Conventional style (`feat:`, `fix:`, `doc:`) is also used. Include PR number. ## Key architecture concepts - **Any/AnyView**: Type-erased value containers. `AnyView` is non-owning, `Any` owns. - **Object system**: Ref-counted heap objects. `ObjectObj` holds data, `Object` is the ref wrapper. Created via `make_object(args...)`. - **Function**: Type-erased callable with packed calling convention `(const AnyView* args, int32_t num_args, Any* rv)`. - **Global registry**: Functions registered by string name, accessible cross-language via `register_global_func`/`get_global_func`. - **Containers**: `Array` (immutable), `List` (mutable), `Map` (immutable), `Dict` (mutable), `String`, `Tensor`, `Shape`, `Tuple`, `Variant`. - **Reflection**: `ObjectDef` builder with `def_field`/`def_method`. Exposed to Python as `tvm_ffi.dataclasses.c_class`. - **Module system**: `load_module("path.so")` wraps shared libraries, exposing functions via `__tvm_ffi_` symbol prefix. ## CI Runs on: Linux x86_64 + aarch64, macOS arm64, Windows AMD64. Jobs: lint -> clang-tidy (if C++ changed) -> doc build -> test (C++, Python, Rust). ## Further reading The `docs/` directory contains the full Sphinx documentation site, including: - `docs/concepts/` — design docs (ABI overview, Any, Object/Class, Tensor, Function, etc.) - `docs/guides/` — usage guides (exporting functions/classes, kernel libraries, C++/Python/Rust) - `docs/get_started/` — quickstart and stable C ABI - `docs/dev/` — developer instructions (source build, CI/CD, release process, doc build) - `docs/packaging/` — Python packaging guide tvm-ffi-0.1.12/CMakeLists.txt000066400000000000000000000337511521067262500157220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.18) project(tvm_ffi LANGUAGES CXX C) option(TVM_FFI_USE_LIBBACKTRACE "Enable libbacktrace" ON) option(TVM_FFI_USE_EXTRA_CXX_API "Enable extra CXX API in shared lib" ON) option(TVM_FFI_USE_THREADS "Link against threads in shared lib" ON) option(TVM_FFI_USE_DL_LIBS "Link against dl libs in shared lib" ON) option(TVM_FFI_BACKTRACE_ON_SEGFAULT "Set signal handler to print backtrace on segfault" ON) include(${CMAKE_CURRENT_LIST_DIR}/cmake/Utils/DetectTargetTriple.cmake) if (TVM_FFI_USE_LIBBACKTRACE) include(${CMAKE_CURRENT_LIST_DIR}/cmake/Utils/AddLibbacktrace.cmake) endif () include(${CMAKE_CURRENT_LIST_DIR}/cmake/Utils/Library.cmake) # ######### Target: `tvm_ffi_header` ########## # they can be used in cases where user do not want to link into the library in cases like deferred # linking add_library(tvm_ffi_header INTERFACE) add_library(tvm_ffi::header ALIAS tvm_ffi_header) target_compile_features(tvm_ffi_header INTERFACE cxx_std_17) if (CMAKE_CXX_BYTE_ORDER STREQUAL "BIG_ENDIAN") target_compile_definitions(tvm_ffi_header INTERFACE TVM_FFI_CMAKE_LITTLE_ENDIAN=0) elseif (CMAKE_CXX_BYTE_ORDER STREQUAL "LITTLE_ENDIAN") target_compile_definitions(tvm_ffi_header INTERFACE TVM_FFI_CMAKE_LITTLE_ENDIAN=1) else () message(STATUS "Endianness could not be determined, skip setting") endif () target_include_directories( tvm_ffi_header INTERFACE $ $ ) target_include_directories( tvm_ffi_header INTERFACE $ $ ) # ######### Target: `tvm_ffi_objs` ########## set(_tvm_ffi_objs_sources "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/backtrace.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/backtrace_win.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/object.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/error.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/function.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/tensor.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/dtype.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/container.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/init_once.cc" ) set(_tvm_ffi_extra_objs_sources "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_equal.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_hash.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/visit_error_context.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_parser.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/dataclass.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_system_lib.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_dynamic_lib.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/env_context.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/env_c_api.cc" ) if (TVM_FFI_USE_EXTRA_CXX_API) list(APPEND _tvm_ffi_objs_sources ${_tvm_ffi_extra_objs_sources}) endif () add_library(tvm_ffi_objs OBJECT ${_tvm_ffi_objs_sources}) target_compile_features(tvm_ffi_objs PRIVATE cxx_std_17) set_target_properties( tvm_ffi_objs PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_EXTENSIONS OFF CXX_STANDARD_REQUIRED ON CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON PREFIX "lib" ) # add the include path as public so they are visible to downstreams target_link_libraries(tvm_ffi_objs PUBLIC tvm_ffi_header) if (TVM_FFI_USE_LIBBACKTRACE) message(STATUS "Setting C++ macro TVM_FFI_USE_LIBBACKTRACE - 1") target_compile_definitions(tvm_ffi_objs PRIVATE TVM_FFI_USE_LIBBACKTRACE=1) else () message(STATUS "Setting C++ macro TVM_FFI_USE_LIBBACKTRACE - 0") target_compile_definitions(tvm_ffi_objs PRIVATE TVM_FFI_USE_LIBBACKTRACE=0) endif () if (TVM_FFI_BACKTRACE_ON_SEGFAULT) message(STATUS "Setting C++ macro TVM_FFI_BACKTRACE_ON_SEGFAULT - 1") target_compile_definitions(tvm_ffi_objs PRIVATE TVM_FFI_BACKTRACE_ON_SEGFAULT=1) else () message(STATUS "Setting C++ macro TVM_FFI_BACKTRACE_ON_SEGFAULT - 0") target_compile_definitions(tvm_ffi_objs PRIVATE TVM_FFI_BACKTRACE_ON_SEGFAULT=0) endif () tvm_ffi_add_msvc_flags(tvm_ffi_objs) tvm_ffi_add_target_from_obj(tvm_ffi tvm_ffi_objs) if (TVM_FFI_USE_THREADS) find_package(Threads REQUIRED) target_link_libraries(tvm_ffi_shared PRIVATE Threads::Threads) target_link_libraries(tvm_ffi_static INTERFACE Threads::Threads) endif () if (TVM_FFI_USE_EXTRA_CXX_API AND CMAKE_DL_LIBS AND TVM_FFI_USE_DL_LIBS ) target_link_libraries(tvm_ffi_shared PRIVATE ${CMAKE_DL_LIBS}) target_link_libraries(tvm_ffi_static INTERFACE ${CMAKE_DL_LIBS}) endif () if (TARGET libbacktrace) target_link_libraries(tvm_ffi_objs PRIVATE libbacktrace) target_link_libraries(tvm_ffi_shared PRIVATE libbacktrace) target_link_libraries(tvm_ffi_static PRIVATE libbacktrace) endif () if (MSVC) target_link_libraries(tvm_ffi_objs PRIVATE DbgHelp.lib) target_link_libraries(tvm_ffi_shared PRIVATE DbgHelp.lib) target_link_libraries(tvm_ffi_static PRIVATE DbgHelp.lib) # produce pdb file target_link_options(tvm_ffi_shared PRIVATE /DEBUG) endif () # expose the headers as public dependencies target_link_libraries(tvm_ffi_objs PUBLIC tvm_ffi_header) target_link_libraries(tvm_ffi_shared PUBLIC tvm_ffi_header) target_link_libraries(tvm_ffi_static PUBLIC tvm_ffi_header) # ######### Target: `tvm_ffi_testing` ########## # Build testing utilities as a separate shared library that can be loaded on demand # `tvm_ffi_testing` won't be inlcuded in `libtvm_ffi` and contains functions that are registered # only for testing purposes target_sources(tvm_ffi_testing PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/testing/testing.cc") target_compile_definitions(tvm_ffi_testing PRIVATE TVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1) target_compile_features(tvm_ffi_testing PRIVATE cxx_std_17) set_target_properties( tvm_ffi_testing PROPERTIES CXX_EXTENSIONS OFF CXX_STANDARD_REQUIRED ON LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) target_link_libraries(tvm_ffi_testing PRIVATE tvm_ffi_shared) target_link_libraries(tvm_ffi_testing PUBLIC tvm_ffi_header) tvm_ffi_add_msvc_flags(tvm_ffi_testing) tvm_ffi_add_apple_dsymutil(tvm_ffi_testing) # Set the install RPATH for tvm_ffi_testing so it can find tvm_ffi.so relatively if (APPLE) # macOS uses @loader_path set_target_properties(tvm_ffi_testing PROPERTIES INSTALL_RPATH "@loader_path") elseif (UNIX AND NOT APPLE) # Linux uses $ORIGIN set_target_properties(tvm_ffi_testing PROPERTIES INSTALL_RPATH "\$ORIGIN") endif () # ---------------------------------------------------------------------------- # The following code section only is triggered when the project is the root and will be skipped when # the project is a subproject. # ---------------------------------------------------------------------------- if (NOT ${PROJECT_NAME} STREQUAL ${CMAKE_PROJECT_NAME}) return() endif () option(TVM_FFI_ATTACH_DEBUG_SYMBOLS "Attach debug symbols even in release mode" OFF) option(TVM_FFI_BUILD_TESTS "Adding test targets." OFF) if (TVM_FFI_ATTACH_DEBUG_SYMBOLS) if (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") target_compile_options(tvm_ffi_objs PRIVATE -g1) endif () endif () include(cmake/Utils/CxxWarning.cmake) include(cmake/Utils/Sanitizer.cmake) # remap the file name to the source directory so we can see the exact file name in backtrace # relative to the project source root tvm_ffi_add_prefix_map(tvm_ffi_objs ${CMAKE_SOURCE_DIR}) # ######### Adding cpp tests ########## # logics below are only executed when the project is the root project. but not when the project is a # subproject. if (TVM_FFI_BUILD_TESTS) enable_testing() message(STATUS "Enable Testing") include(cmake/Utils/AddGoogleTest.cmake) add_subdirectory(tests/cpp/) tvm_ffi_add_cxx_warning(tvm_ffi_objs) endif () # ######### Adding python module ########## option(TVM_FFI_BUILD_PYTHON_MODULE "Adding python module." OFF) if (TVM_FFI_BUILD_PYTHON_MODULE) # Helper function to build the cython module message(STATUS "Building cython module..") # prefer virtualenv when searching for python set(Python_FIND_VIRTUALENV FIRST) # cmake-lint: disable=C0103 find_package( Python COMPONENTS Interpreter Development.Module Development.SABIModule REQUIRED ) set(_core_cpp ${CMAKE_CURRENT_BINARY_DIR}/core.cpp) set(_core_pyx ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/core.pyx) set(_cython_sources ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/core.pyx ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/base.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/device.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/dtype.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/error.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/function.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/tensor.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/object.pxi ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/string.pxi ) # Run a Python script to check for free-threaded build execute_process( COMMAND ${Python_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_config_var('Py_GIL_DISABLED') == 1)" OUTPUT_VARIABLE PYTHON_IS_FREE_THREADED OUTPUT_STRIP_TRAILING_WHITESPACE ) if (PYTHON_IS_FREE_THREADED) message(STATUS "Free-threaded Python detected.") endif () add_custom_command( OUTPUT ${_core_cpp} COMMAND ${Python_EXECUTABLE} -m cython --cplus ${_core_pyx} -o ${_core_cpp} --module-name "tvm_ffi.core" WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Transpiling ${_core_pyx} to ${_core_cpp}" DEPENDS ${_cython_sources} VERBATIM ) if (Python_VERSION VERSION_GREATER_EQUAL "3.12" AND NOT PYTHON_IS_FREE_THREADED) # >= Python3.12, use Use_SABI version python_add_library(tvm_ffi_cython MODULE "${_core_cpp}" USE_SABI 3.12) target_link_libraries(tvm_ffi_cython PRIVATE Python::SABIModule) set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core") if (NOT WIN32) target_link_libraries(tvm_ffi_cython PRIVATE Python::Module) set_target_properties(tvm_ffi_cython PROPERTIES SUFFIX ".abi3.so") endif () else () # before Python3.12, use WITH_SOABI version python_add_library(tvm_ffi_cython MODULE "${_core_cpp}" WITH_SOABI) set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core") endif () target_include_directories( tvm_ffi_cython PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython ) target_compile_features(tvm_ffi_cython PRIVATE cxx_std_17) target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_header) target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_shared) # link against testing to ensure right unloading order (cython first then testing) target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_testing) # Set RPATH for tvm_ffi_cython to find tvm_ffi_shared.so relatively if (APPLE) # macOS uses @loader_path set_target_properties(tvm_ffi_cython PROPERTIES INSTALL_RPATH "@loader_path/lib") elseif (UNIX AND NOT APPLE) # Linux uses $ORIGIN set_target_properties(tvm_ffi_cython PROPERTIES INSTALL_RPATH "\$ORIGIN/lib") endif () install(TARGETS tvm_ffi_cython DESTINATION .) # ######### Installing the source ########## install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dlpack/include/ DESTINATION 3rdparty/dlpack/include/ ) install( DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/libbacktrace/ DESTINATION 3rdparty/libbacktrace/ PATTERN ".git" EXCLUDE PATTERN ".git*" EXCLUDE PATTERN "*.tmp" EXCLUDE ) install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/ DESTINATION src/ffi/) install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Utils/ DESTINATION share/cmake/tvm_ffi/Utils) install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt DESTINATION .) install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm_ffi-config.cmake DESTINATION share/cmake/tvm_ffi ) install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/tvm_ffi_python_helpers.h DESTINATION include/ ) endif () # ######### Install the related for normal cmake library ########## install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/tvm/ffi/ DESTINATION include/tvm/ffi/) install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dlpack/include/ DESTINATION include/) install(TARGETS tvm_ffi_shared DESTINATION lib) # if tvm_ffi_testing is built, we also install it if (TARGET tvm_ffi_testing) install( TARGETS tvm_ffi_testing DESTINATION lib OPTIONAL ) endif () # ship additional dSYM files for debugging symbols on if available if (APPLE) install( DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/lib/ DESTINATION lib FILES_MATCHING PATTERN "*.dSYM" ) endif () if (NOT TVM_FFI_BUILD_PYTHON_MODULE) # when building wheel, we do not ship static as we already ships source and dll install( TARGETS tvm_ffi_static DESTINATION lib OPTIONAL ) endif () tvm-ffi-0.1.12/CONTRIBUTING.md000066400000000000000000000043301521067262500154020ustar00rootroot00000000000000 # Contributing to TVM FFI We welcome contributions of all kinds, including bug fixes, documentation improvements, enhancements, and more. - Fork the repository and create a new branch for your work. - Push your changes to your fork and open a pull request to the main repository. - Please provide a clear description of your changes and link to the relevant issue if one exists. - Create necessary test cases and documentation. - Work with the community by incorporating feedback from reviewers until the change is ready to be merged. For significant changes, it's often a good idea to open a GitHub issue first (with `[RFC] title`) to discuss your proposal. It is optional, but can be very helpful as it allows the maintainers and the community to provide feedback and helps ensure your work aligns with the project's goals. The full developer manual is hosted at . Key pages: - [Build from Source](https://tvm.apache.org/ffi/dev/source_build.html) -- editable install, CMake flags, C++-only build - [Build This Doc Site](https://tvm.apache.org/ffi/dev/doc_build.html) -- building the documentation locally - [Reproduce CI/CD](https://tvm.apache.org/ffi/dev/ci_cd.html) -- linters, pre-commit, unit tests, wheel building - [Release Process](https://tvm.apache.org/ffi/dev/release_process.html) -- versioning and release workflow tvm-ffi-0.1.12/KEYS000066400000000000000000000450611521067262500136550ustar00rootroot00000000000000This file contains the PGP keys of various developers. Please don't use them for email unless you have to. Their main purpose is code signing. Examples of importing this file in your keystore: gpg --import KEYS Examples of adding your key to this file: (gpg --list-sigs && gpg --armor --export ) >> this file. ----------------------------------------------------------------------------------- pub rsa4096 2019-11-15 [SC] EF52D68AD5276994249816836754EA97C55E3DEB uid [ultimate] Tianqi Chen (CODE SIGNING KEY) sig 3 6754EA97C55E3DEB 2019-11-15 Tianqi Chen (CODE SIGNING KEY) sub rsa4096 2019-11-15 [E] sig 6754EA97C55E3DEB 2019-11-15 Tianqi Chen (CODE SIGNING KEY) -----BEGIN PGP PUBLIC KEY BLOCK----- mQINBF3OK24BEADD4hxjrsgb4jIDIACHS15X+5YP/YaUF5UDDQs/bNn/xGJGVl4/ 4sJ6qKZcvMDrWTmnNItYBuaHi1qhGvlcASBekm/9PU2U8lZmAF1lZkKIIYZkX+If s8PEYurE8cDr65orrdsFF8Zwb+u6x+gMsHNivsU2Kn3xbQjGmeW44UA+aaXzcJp6 sVk3aX5DypoYJNBmbASyOjZVWkcrJ+NKEfJ1dKtka5/siqOjuvCd8NT5dJVhZbm3 Sf8iclEMqog1LhdI/FhE2fB3C5hJkzcinq2v55qDaGqsL+qgT7agf9b4t0EgjbVh cs6jlCglad+Oz27BQIjt06HE1OB5T/Gxa080FK4JZMpxZJ5tDA2/7DQM2MyN84z/ s62JuBJnsrzr4w8D/QcAyzAmyzAqvxLR/aqLgJTIcQiw6AenHovKkNbEQOBYE2T5 ms7uVO2E2Tv42J4Te4OKhpId9mK+7elCLvOb2DfAJDdYxDN9c8dJTls+G6xmv0h9 bb2+QRjkpDiFeu1hKNEe0/ST/YXDfRYpKl+1t/QZ+JccLgEdEwuo/IQ1e4POH2h0 Zqvy7TR5obeTf0TvmLzW+i3s1oUkmSAnQEncSGnGnlugYk0BLuMMi9Fhx6qcC5pC cA3nsRqFKebtnpop+m+psFkmd//xKSXJt9IYVEbQVNiUKm9uYq6RxZEAmQARAQAB tDJUaWFucWkgQ2hlbiAoQ09ERSBTSUdOSU5HIEtFWSkgPHRxY2hlbkBhcGFjaGUu b3JnPokCTgQTAQgAOBYhBO9S1orVJ2mUJJgWg2dU6pfFXj3rBQJdzituAhsDBQsJ CAcCBhUKCQgLAgQWAgMBAh4BAheAAAoJEGdU6pfFXj3rVJIQALBArXEaFDdTw8wl 65nPLU6+QPc6eMn7mz6BDp1V7xL6Lq1GbArLpmQHIFhfQ/5Qmg80wuFBU1CNSRHd tdZq3v8tB9Txvhy6bLQ+IijWH/TxSEPqnrkNsWBQLqAygDC5O3Ook/T6B5kuc176 Kz+w+YhzPS5hoPfJK6xGoKDNlkhmI/EnUjAq459VNpXeoeemiydzvApiCHH0VfOj XnmgAJsAJA21EfT5Wuh/WODsf0HkaXB0xoWZfE/ugIQBLhZi9nUTYgwU2r4a+v4A 4C2T1OyJ3mDU+Oi/z6d0WJvsIrLCFcF4Q7b/6+MGkgLDGlsEKK2LZMrulGzQ1QY/ O4ck3dVDseqT2urplrTamDIh1IQmOt1FqMFwugdjfQwJ5HQeX6IeUGZei2Av/IZR 8Vw5Wxtm1Aksz3Js6iP3QmAh7txDUKO+eT5zLSXBoPmkleLnvCdtlvwaSNCAudHw 12h10IV286OetJvyyjmh/q/30sKNGiuucLMzPMwtLNW/j3cts3fqRHIHxepT6m94 FoYIlwVu4afiGgSi/7cN4p9GgfwnFGeETd25pgNG0KdXbVWniO1dTEKzOtvtuPYK Y88ZAfdOgj4dyeI9ZnJV8RaZvpImDPVHGQm69/071jBxyWZnVi/YtOm+DjHfw0Vi uiUdzoIb54oWW8tbiNg/nfiLUaJBuQINBF3OK24BEAC9W8Cwubu4Dpr4m0IIrLF5 zRRqQm9QIcEC0QHf6w1c2NWQTJP+MQY/jZLjtKw5yCQDghT+qsil2p8xCM0EqRd6 6NqxsAoweTCoV0MwolQv5T3KuP54SlNWjO+6gT73LkKuOHoIyy5cS9pIITlExHy+ XHtfQi1keDpWUEyvSRG9slu1DcxAeo6nFEpCuoQ+xx/lrCMxDlyZJCDhj2fXs2hK 8oKLV5NbIuifbXbCiOvZUdBHk0yLCEc6wNsVR30yLijSiPCKsAPcsG0PjQnz3eTb 0czq+6g50zUVOTioUghIlZ1DhCsxQGnlxoLY71pnmc7qVszdXPV2Mp7/KSIhDJFQ LN0enDVz9aRXfpEK3SifxaPVNd61O/BGziza+XCK5qpEQL95UM2NdQCWixYmIOJE k95tpnagtNupMkrY6WEa0CjVBzF1kdr5WpeUd6w85rA/opcqpQ8yLmvpyJ4tXZhN 7oAWZSUzyB904FMswUEhaS7pEJIlACeFcPwm31Jv/637gw1CopZpDxDUaW5/boG5 9Gp9D/GV2gyMrHAcwA1gZSbmolv5ZYcnUmwTPijVNZ+o70HBbvbNZqziPgy9G+L/ oGBkY/fpg7qfaGtAbOUbx1ck04CbafSUQIxpCG8in6zwrIRnn4uj6q4wIZ8SnvQ0 h3Ug0DmdsxvB/xdfillH/QARAQABiQI2BBgBCAAgFiEE71LWitUnaZQkmBaDZ1Tq l8VePesFAl3OK24CGwwACgkQZ1Tql8VePeuZ1Q//csRsGDKNrW5e0EitEcfPZ0PC teEw7A16dniXiCQF39KxxLzjCjUq7U8iWNm7bn1zdXcSVYZow+i5hFWXgZLKTKep tQoocJmQ7kPV5oiTBewFy9T4BICUekj/EhXhSz1wxb3GSc+uHL2IUlFkixTY4k4B 9zq49gkNkTM02Or3quu1ZWAgeol1BSyV0tcI1h3M0OXtrN6idLyzQJFRyMYtzfwp Pd2+hdaKAl8mKANs/GMJni3QvyVXzuJxMP6SNOFx4mWj0UVFVZvosv1lLXDesvwY sNZmz5IkfuU4DHz1ZzZc3sThkpBdBiadvyKtNsenNh5nEXtwVhpiFf3IdZAvG7Ks 7i3Fx1/ObbvxMCWeFoB6oP/swHr9i6dqntiJoB6Gl5y1ye3qte8PiNuwRVhz+YOK 58Ga3wWMvODpi2AgSFv7cd1OFXXsoonORfmpcfAp+h6dIr/ttQMP2929/NoX3Cs4 /pXoG9L5EOpMfj0Q24sAGW8VzuCAHL3e7QSijFuSHZxz9oe4C28/mAY+KP0dif0Q O3rq4kpqlhseyzcRyE1LWBvzuCeSTui2OPmyivFY57TOPnMHm5sXVby1VUiwm0B0 RgBtZDRLv765lAFGtp43sccZ7zfRaKhkVmzh3bAZ62nJyQNGw0TWg96Pf7Kjb0Bv ha8fS9ysWDy/Ye65MP4= =MSiP -----END PGP PUBLIC KEY BLOCK----- pub rsa4096 2023-05-05 [SC] 664EF29634C05669C3DCF83106D051CA84EF3749 uid [ultimate] Siyuan Feng (CODE SIGNING KEY) sig 3 06D051CA84EF3749 2023-05-05 [self-signature] sub rsa4096 2023-05-05 [E] sig 06D051CA84EF3749 2023-05-05 [self-signature] -----BEGIN PGP PUBLIC KEY BLOCK----- mQINBGRUb34BEACyedBJjD2tYY86mIr1OR42eR+w0dvh/cECgp4UIm5QG6z9YXrU KjHtv426uKCloYiWU5+b8ASPMbtP5q1rvrRKQuapbDBN2qlS7E/PScFHDK50ydOA v8melfp8pWLW48kE5EeSdvhF8U2QzEqT6EYmNIExLYjSV1+Jck40DbVL+ak/4clB Qv/l0DW2Fw1u6GAKaNgnWgZDhc3os176rS0ERzflZeF+rqskYSb8Uy37vBB9By64 edIsaxxOEqvpJfa2Ar/6vDnBplYSfMHq9kqs0vuAVu0r7Dadl199oVMmUwBPD4uq IVaWoetCbR+DpVIZZ74i4kOH47xfh3kZ5zvWo5E+PL4XH/6us5tp+nAnCjhthxY2 tpBBZo7M4qgwNOxk/3zsysZPYPhQaZUx7/LBaApOhyWrmQYrtZpvNjHsoEToHtDj wCSwDxf2c6mrncnmKg6UfsWSEiIiPcsYqXkp2Bh4xKzdI54qCt8LTPGKjTPj4FO/ EccT5Ad8+lOKSOIugbNmECEbpVUhrlDa/yYtkosHbnQZFgBBaa+RCsRdbGpuuc+E hqEpvwbPq0J8zPJpKmyeu/S5gWww29ix39J+F5ZxjQZBSUUPRsnCwHDdkgFBlDdK ZeQqlKjr7Mfp+LlI+7HIIxO9HOPN3WIjFyWPm10d6cWHN7MDxMySP8l1+wARAQAB tDJTaXl1YW4gRmVuZyAoQ09ERSBTSUdOSU5HIEtFWSkgPHN5ZmVuZ0BhcGFjaGUu b3JnPokCTgQTAQoAOBYhBGZO8pY0wFZpw9z4MQbQUcqE7zdJBQJkVG9+AhsDBQsJ CAcDBRUKCQgLBRYCAwEAAh4FAheAAAoJEAbQUcqE7zdJsZQP/0EX1XKFDF37c8cI jQEeQ44Z3F5C6cxM7X9efqLQzDgvVRR525qjpM9uk/usUdKupwwX0GkJQu+JbRus jFCi9RaOO88+w7ihS4qOFcXV0CXHxxSKkKfKU7DZhEprJtyQ1QE38gWlHKB/E6Y+ oJy9EKMZMXP28gj/tUT56IABI6X+b1BSTT5PV8QURkkTPDVQDpWo/AmtPdei2bvd KnTZGAjxv98rMYvjJUGMPx8oA3cqDRlIltgSyXlwht8Ig/wbUzW/oRmzLk3TdeoE 42/QHTLjfDlSm72V2B67QToYo/URIxJimUvg4/+VT8ByLxPzkL3jYkDh6K87JaB3 9XO8HjNCH/xfZtMSQjmApInpk7VjDVhDC4Dlf5t41pUK5KvGsU7eLAE0jL/R/aA8 z5pvf3afLK3Bpj2zKvy2rFkRmCKIS7mycogBUdOk4GT8ZoLDuaTmUcbfx9H4/9Zw UYWCw6cJzD3qICqCcszlfW+99b92JddCU5ITMfwuWuY/OX/LfpwibAjzor2TFWya CVI7kkQj9C0vcbpxgCMd4HRMV9p2CQUkvPEKfaPg+kzJC1Yz87DC6aSrLIzVvcIj LZ2yOzR4QIeTS6hRsMmQRGPZO0KFres4760BiUCH0gid6LWDNq2YTXqdNu1ffQXe PV8Risr23rrxOTJqlYX3GF+Xd7C1uQINBGRUb34BEAC2q4MdKGYgsl9BpvOA7TnN kBtc8Gmg+DOdjBhG5BCo6h6U15RxfIvSikRi0Sz3F3YZymGKIeJp8ug6brY4KWjA 7dtwqlvnthyWa0mPrgHZvkIM86URO5wSvRMXx1x/qWJ8BrOoCDji+fmC3uI9IbY5 RvkzHACYz4duZM54ZlhM6lOL3TtgF2OyXod2MFwuC5WAAPuqAG5MF+gNdf5JA+p2 RfDIGeZNOWQVWi9CrHWt8fC80WG/7r6Ta84yV6KTfqhsXToFZICVXt2BEg7K/UzI Ip7Zf5rKsDT0iDxtiJIbwFBbTS0hE3ICrPWVHPNRVsqp3wHDkjjnt0aPL+G938i/ dzOwHZIf9nPjIrX94DvPpOXGrHsW6JyMHZ/3diROpWy7DzplW2i/wVftzkhT4GSN xARgLJM/iZriOMYvHafpOm5OxkfFzeJpnjZrRTJKNCAQFbdxI9pzR/2ingvZekUs FYcauQVuL1MbamEhf7pRXHOS0bOPAONpeXl7aNAH4jf8x/3iRHRQYcpGRT1ESddO /swv+Cj2qj80vzF/oT/QMWrWDGHIiriRDSf1WpFHVqzki1jY6znrw6EvGsowbzha kGC9dRzLm8P1aCwjjsU1J82yEKA4oHwAX6rlkRkWVUkTMk3G1coandbgcbnI6ngG 7V932RapIOZB06rh/6mGgQARAQABiQI2BBgBCgAgFiEEZk7yljTAVmnD3PgxBtBR yoTvN0kFAmRUb34CGwwACgkQBtBRyoTvN0kLfg/+JOmX8SZLksdEo4H5KmeHVOQk EpWlFk/SmSoxV1k+kz68B5gxBPWQwRj61cGoBFKdvP2s3BSnKy7+iow0uwh6KIy0 zMooEOqCr/kEeFLxq79kFzxwwDSkzUO1UwGWCzVGj4V3UCq72xt2r3mJxRLNijTr JJFe6+pLFXRbgrZ5ulWGxkiZRK007fPqtkretLiTyUXcJzU6HBUi0/pnyA2B0mWL E0HO0TdOPIDTH/t2vtLZNhWl2T0lbjtdL7IcPAKQoNd07GyK1pPGpLQqP5gRrpFR zrbslqRvKtpVjL5iPQCtv8Tc9ovOIQVRXzuXm9W1OtgFVYH2GQe+vbpYhLL7dOfF Cwo6YLQnc2DRLAffy6G5weLJYYE52gMzb8z++Ys0A+XBkuVvHX1EBeyWy+OSuFSl Aujm/jMMgK7dtNIHXgtVGEAKKJb4amc1wsZ9dmUyf1UyFxWlDUEaLq8+5Ut8a+tW pTwwQLAXPVElpi32gLDP2rvHzIw1Hs0MpoxwOOjH/QCeRQ/V3acAUVv1JX96On0t NqlR/Q5S24ktyC1uy9oLdIZmgllKUb8i6s6+XSkWRata3HTsfySDXMdntZV1Zrjx WTgrESErlqNLN5ZTTW/1jBELJCfJKxgHUip+Yo6qNZoWwNLP1BaIcoA3miSG3DXf wS/UuN04NxDy7V6mPXE= =MTba -----END PGP PUBLIC KEY BLOCK----- pub rsa4096 2025-11-27 [SC] [expires: 2026-05-26] C06AC8A0A0ADC6E5DE82DD82B87BBA6381A37E54 uid [ultimate] Junru Shao (CODE SIGNING KEY) sig 3 B87BBA6381A37E54 2025-11-27 [self-signature] sub rsa4096 2025-11-27 [E] [expires: 2026-05-26] sig B87BBA6381A37E54 2025-11-27 [self-signature] -----BEGIN PGP PUBLIC KEY BLOCK----- mQINBGkn0pEBEADN/3Nc/5Mousseqg20t+BCf/TNY0nM2LQBndsFgAGUE8cGiU4e 1LR502TO5MMUqpMuIBcE+EIYuIZ03Zz9XwCLGlM2g7kOx2yMHuV8qZmxs8YSNE33 VcIuSpxfQjGYYk3HtGacvijWb79CVqHxGOAC5dMel4gV2UDerZfXSI/9NCnr0P86 7gG2zYYHYySfT7B7AieByQ7JCINTqQcVp12jGW7yhOFWlyEqQagBbKSbX6mzbBxf 03CJXdVZ9KKe0s+6qmnyMwcrg+Bn8nZbiCzZpbgbueawl3NxIGqHygIgl/KTr40I BsxCyZE3YJzpnDu2HAasAqb0L9G7y97SDDrlLvEsa/GNvBsiyXdqB/J7+O08UiM7 42nZNo+s+MlE6i9wsGhm0Y0lfucUe91E57d/meU2TrpN+3TS1p7uZARRnSPBq3OT BfWEjSeJ/PMqgjL0NvwofYpZORBfRIeWAfPQ2MwDceIy7saRj5ijy0r+0zv9yM2C I0Xuo2EnCkdB/rHNVWCiwn96lAN9sfz+WZjL7PnB6255SPJWaAEmLzN33eILvFBB r2lMF/v9n2llKZmiEXC2u8qinymKUmWUu3568vDCXj5k/51yqsWxlZzWhG6bUmCs +QuaQ2zF40QwMMKnqxPO8ch1AjD+AeK5ootCx9f0A7mvcr4gj8sLSpAPiwARAQAB iQI2BCABCAAgFiEEEYhoPIYNLw/gRZ70iQut4l0rCEUFAmkn1MECHQMACgkQiQut 4l0rCEUojw/+M709KwqglaXri1J+TQaiCNIC+7LNMEKvVcP5EBSI6/v2Ey5raZeQ d65iJxgH7JB+NpvZCW5frJY5C2MhICkeKmbr/jS7RcMCE0pO1vbxUQH5h+C5tsYO 4UiNU0/N0sef61BaUMexQfa0sMbgC8tCD5COYmNJStF44Eba+iBUbFBf9Hjirzzd Ecx3j7OUYMKsxCpytyxXrxKnWI2AEDU9YWAW2oZUn5wmnex8JZEULle4xi3C5MnA TCQT5sQDyK+bHFIhhGAD08uZyReCK1qtvCm6PUsaw739b7sDznjYJ9r01MSN/vRy Cvs4safB64NrrVenKVF+sp9T40nE8o7NBARUWQybj/qpn66vbHlr1FrglrBgBZwD 07kMYUVuJ9UY+f+66y+Up/KhzwQFFl4vcrXWB8lQFjZRRlyYpNuyqEL9zBx/z45B CQTMzTiR1NKmF6ETuowvKCfp3TbN+b6nEh8ZFytR4M0cACWSw0o5UPLOcXwmaDib 59BlopyszjL7rGsS85DVIWkT9+361Ru3xmQKhTY4Y3g5Y/pbn3y1kwrnJqcGDtXy HMxXTEs3F1+3uZ4QDKpv5fi97fR7vwVa/DSFH1Uy4vsOIgEjQhx1h2xXb06XQ816 J6ufe7MmrN52PqUp++gb9GoVWNNEWhfVAf4LBr4H6rXyBzl7xdB3aD20NEp1bnJ1 IFNoYW8gKENPREUgU0lHTklORyBLRVkpIDxqdW5ydXNoYW9AYXBhY2hlLm9yZz6J AlcEEwEIAEEWIQQRiGg8hg0vD+BFnvSJC63iXSsIRQUCaSfSkQIbAwUJA8JnAAUL CQgHAgIiAgYVCgkICwIEFgIDAQIeBwIXgAAKCRCJC63iXSsIRUjMEADFPHG+unmL bcQJkZhnnhdQUKmCM/wC3NVPZjBkqYVOHMfqxnVdVa4hd4AsiJbJjOEuMlyChP8/ zP16sCS/xmW1AbgruuMvz5aG1Cv7zSlWigGK+jERjnTculMWkLjb2BwkLsus7Jta pXEpxJ/ssU/gdRBCIhP0kgNxHCq8tw8dNYa53XsQPaCJSMQw7zvxB29tQ+ZO6FOt r5KATeloHun+MKrP0qmC28dN2D5Tdsxgga9mGTbbXYhrEQ8QbObQdMjYPXlnwUWS IdxQH+Nxh9t2sdKOh+I9EPIrU7l/19Rs1+V6fnTvQHR+f5j+hIaPPz+R61UEgkjM klAqyHcLvNNw72FJ0IGLn6FVe7QocXFTX3nKdXwQDRNkUZSPdPAMSyGcQXsonN1t 1u9IX8suVn5MJCSbrXei+jTwLRaj9Gj/5z1dqX5lzDtv+855b6MiAqx50tIw9omI upcqzr6xPgr87EBVGvtxd6VMZy9IvemrcL2p9I4RYGLLiySKNDUuwc+MhREs91x1 1YjgfsGgz6DHK9Oe0GJKKFL9eQBR2ZiEnJo2iQ9Gjb3I/Ewq7ChEto2SHzFIAUy+ J7nePjfbyJqQgF+RYfr4mLDVK54erBTMHoxXIvqWDI6hw4nhWPQsyYD/nc/MFwmn 0dDTiPqN3WQWE4OHf9EgvSi6V8TaccHBfbkCDQRpJ9KRARAAx0Ecke6VCqkL/Ycm ZAOHyXfz6m8Y9TJrEiiAyrUMXqpzeIjsW/BZ3NJol0GvBqrxuCsGHzhs9vhBpyqV 9XbjI1ckPfhPH2W7VIKMMCBOEoGCcckpjLPBko0wo3ddqcBuTrYoszrDWgxERAGF nzBsRU1ptjl+0q/prdXdrxHXfLjvMkTgTeOeGyuvBGOJmxkf4Mk/rYSVgITEK/nL wyIwVmHezwJMqu90eFu8OvYVZ9iG2LuvxA2R8ds+nEvagq0TMuZzQrMBViABgrN6 OoWstooRFaStqmFfh184B/I93UKuWbrvu6tTPJRTweg/CHEqZPz+EfSsbaAAORTg uJisYOJiGuR27I/E7hI3meco7L6YvMhQHyzPikY143JgaDP0TDMDWJVJgwy0bnjQ theIP/05/uRE6/ZBVgp+qNZGpyE5WBztRR6SA3igYxjVu64jsGaERdYRuxmebWPO PKGxzE8uPU//xG6g4qBL7Jzq5zsB/AQSoLd2CcLkcJiw0INHvf5VE5lolDQTlflX frdd5xmQkl2Ohv7Qry3XjIufoh4kcsJreehtwbu/75L4BkxXUQQJnqiCfVBywYWN JQycKgiQCsQCF19VsSv+jZXw1+zOWdpZcDHZfNN08hXLZhV9MSV9BreUR9LqeG1o HO3CtOQkIRBRHQKUD9voNWrkVhcAEQEAAYkCPAQYAQgAJhYhBBGIaDyGDS8P4EWe 9IkLreJdKwhFBQJpJ9KRAhsMBQkDwmcAAAoJEIkLreJdKwhF7dMP/AzXEQlYSZCu 9GXn92/Jl36VrYJlC9ymaDcMZdNrpyXvnwGJZQ8IBMwUh6Yf5Ldx1tIGr7Ty4VLR RBcMFxi81Wo5rF92hfnA/NAWTP+4nCRsPj0kVNOfViHjfbPIQj7I3Bw8JHC39mV6 XwwPiA24tj8SSEyS20NwRcE2ubrE1t7ktBv52rUrJoCguRQYoMA4a6QAyg1wH2yk zS88zh9vBRbTOwlJj5lymU18y7V3FNoX1wOh45tvxf8xe85q/gyoKEbD75eesXmC agmWMjqWChWhXetZbK73qpyaPqrh4fixq9ghxoAIsMzJPDKwtfr+H9A7dXOZpP33 qqrYWIdWsbnXvZCMNL7myKlH3BVB0cQywlOaEBE52QLueDoFhL8mD50ppVNCBoJV sOzkDSFrsnjlNtN5Kq3DoWZTPfEe5pEEaVF7vydt1zQceVryOSsfK9IQbmzGG7JW lEYwV2fj0oWpXYTK67chflSToMDdmGQbqvIFKV1uOsjB9xbz1S1SeSPlj65J0pV3 BKFIXtPQpuaUlj1gNZEZwK0O1079KCCtyCHqAwDAo69vgtzoJ2UbtkhRU5oZmKr7 /CBGdcQX/HG5NGzfmXB/s6jN/3zuhdsS8DJA1bR8RZS49RlIlDScdTcRiya4mHgy mq9cXog2JCfAidnDDTAds7p8S6Twym/YmDMEaSeNIRYJKwYBBAHaRw8BAQdABfJd 75WhhdCxMJkGkMa606BYfycDmOYhfJcmRR+Hp5CIeAQgFgoAIBYhBJo/a5WQgYFo uQUJeOUCB0kamipTBQJpJ9TjAh0DAAoJEOUCB0kamipT3oQA/iS3Q0wndksyN7v8 AjYZkb9kjMXFz32Sb9OILK/32ShyAP9VcGdbdRh65N7nX4KCdeuf0vTtQCNgDDAd M+rmRE1vC7QqSnVucnUgU2hhbyAoQXBhY2hlKSA8anVucnVzaGFvQGFwYWNoZS5v cmc+iJMEExYKADsWIQSaP2uVkIGBaLkFCXjlAgdJGpoqUwUCaSeNIQIbAwULCQgH AgIiAgYVCgkICwIEFgIDAQIeBwIXgAAKCRDlAgdJGpoqU0zzAP9njtZ8xsXP372V 8cwGg7Xusf/5dhExaT27dfwYxzfvuQEAlyQnE3+w5la3EGAjOpzC6TtR4JS8m9E4 OLWDFXN9TAm4OARpJ40hEgorBgEEAZdVAQUBAQdAJD7TrMeqBDrZBQ/YmGBbe0ES csL9yGXcK4vTlWXegEsDAQgHiHgEGBYKACAWIQSaP2uVkIGBaLkFCXjlAgdJGpoq UwUCaSeNIQIbDAAKCRDlAgdJGpoqUzJhAQD+XrYmQ2moYfHIPmIZoxPH5ZxYoOSj CLhiC1mlwQzL0QEAyKWgj/1Db4r/KlnpgrMVzwiBOceb5Kvgt5fSd+97OwCZAg0E aSfXjgEQALZnUwLAvBtLyEl+eRFMDHXmbE1TmIxOlhV931x/EBaMDNW5lVfVt/0G WMDxx46PMLYTtOH+rF1vz/fnCgfGnKdF2UFMJdCYVDRmwukvSgQHEe1TvX2hW+NM 1+F6ZsboNnKBlRTK7agD2NdOSIJKMnDeAuDt9Kt2Khw7VBPMiUx8CSaARwOoJR8Y Vvo7xQ7OvlH05aLGUj0yG54LAL3Zfu8zOPx5qlMXhGSLXY1RDF1b1RHuuMILUiyb wcxPiLak1/08UguoYGrmvh1eDFbuU9vGc17t/hjYWhiE8CIYG8R3LHCqGjzNboeT HksBcLcH5EzZGjuvM0754Vjtd2b6auzv/Wp9wre/GYxhmRRs46oPOxmGThCjPTTg MPpKcXPVOE6asBg28A0uL3hHdz5IVNXacaq8Wsz6djpuX+EWMZ2ad6Nn5LehygBq mzD98OtthPIg/zHr8ncMjSYbFkd9e/7exPiUNE+gOtUa2jQeVK2wahLzpww6CcQS PlCnL5P74S7QW1dmXI5lGUzU2/MPv6EN8hPyOUKUooJdc+S42EE+FjlvTKUcsVy2 XIxM9o2I+zA/gvvHV7Y20toimMc27uaU+iBqW9EhxPm58xO5p+9YPEt+WvIhhjdB yqS+1e34UQEMNb+GJladrQI6o8E9s3Hy8x++GuAw/cH8qbDBEmZVABEBAAG0NEp1 bnJ1IFNoYW8gKENPREUgU0lHTklORyBLRVkpIDxqdW5ydXNoYW9AYXBhY2hlLm9y Zz6JAlcEEwEIAEEWIQTAasigoK3G5d6C3YK4e7pjgaN+VAUCaSfXjgIbAwUJAO1O AAULCQgHAgIiAgYVCgkICwIEFgIDAQIeBwIXgAAKCRC4e7pjgaN+VNRwD/4jALf9 vLQNKpIw09uZbjvtq2CXC+Gg3VsX1dJWz5x7GBbUQ/SfvLdrxrgHm43jVci3+UYi XoTFkVlW+MqK8mewjkgSa2/IPfcMx/ms5hzAmERUxJPJe57dATnSNwQg2D0/m9oK aNG550wGBwS0DxhvC+wc1EQRoGHRfwQ6wIRqcNaRZ28zuzw8hxReEISvqaCjt2gf sj0ZmARiQhQLCKX0GH+70ArTnX3XHyNTVMao6nnqRnEYV3jARvNIF4mOywIfvgDO o4KdCFgg8b5vhM6sRB0p8Y4DT6+S3pCHmBFnjKBKSKvYjr2CKSqGHtvSVxeg+UuV 7YMPb9PJ5Sz0VgljbVfTR1Gm3Tyc8u+Hk20P/bP9Sm2TsMt6vDDJJmTcK0jGKStU nxm+n8J4H1V1EjMEbNdUw6z3x1tEpP3AtPmSvJbKTbZE7QFxa5y1VuJ6J2ybibnx 91SPexxg11mSWr4+zvsEM8hQjhlUOz+3GfKV7U59Zr65VsokgRz1ZXV1UxDPQlIJ plEB4TbarVVk4gUURXFulTFKGJl4s+R2F3QaQK0OXBW2OcYhTYFanshZGuYTXLn9 Auob0ui4Hw8uAa4PE3QpDz38+XGvMJhyPFeAh3LWjTozv6eLD3hAIGaYANaGmAKP 7JhmfxrTIBginXTgFVPxD3c3rn/VLUJSdjRSybkCDQRpJ9eOARAAxDskpHStEu0j AJTDmIQ3IzfVZ8F+UQFRzNMKOFxjTO2sYULYRKU90WsUGfABdbTlAiebu2tmrf4K qfdk1leoZ8ioLGnk6LaD3XgpdsmsCVSPRKouQ1z3BF3gljD0m5yavK020UpKtKRd 98mqIEASzIua0mfn3gRZ1ju+6G6K6JGavcTFS8ii4HwsHFk4FFajaAPFVkwlrfvy wzkWmq1e6mIScNKTJslsvnBScdtoQLRnZiZO3BjAPe1VPzpvQbD/ywgzHk9gCD8K hcFx6iUdsitNaE/6f/vSzJsweqiud15EvNwd8AOmSavnzrnEP/e4XcE7jOVdiw20 l5Iqw7AzfehgPtUNZuoYGtQxAG5wCt5CgNIqCVPjjDsdO0i51iWjnNinbdRvq5oA pz3tXvw+mM7d+sKoaPAy3CCYvEjj/R6r1aFqBbo/zNQdLRinDu26WkBWmwFWl+zZ n4AGtgb/cwmFPS68kx9UBza+d9pSP/mlgxznRw/d8q2PCRJCGWUa/S+xG3aCx8De tYPDoKB1iGsHB+mAPCQ/79YawUXUa/UM/Mcllx+2M6fBXYOHiKZxw0iHzFUwRpeD 7jI0b/aBd2Kq7ArYJ+wxXtb5xCQMWZ4hfllKHiD2n8HMEIc/5fsXgQB8nh2h4OOd ATHxcQZemZ6aDWr0/Frg8GI3sCi1mIMAEQEAAYkCPAQYAQgAJhYhBMBqyKCgrcbl 3oLdgrh7umOBo35UBQJpJ9eOAhsMBQkA7U4AAAoJELh7umOBo35USBwQAJK0R8hn 5v5CCRxuAjMmaG3FYVZ+mm4rvvyqtWahygK8kK3usrHN/5OxjDtUTdOWOVi7W5fU bHIk0x99MAP45eo0s4l+RIVLO6YTJRNfr4Azzr/9EbU+wjVqWBVSuHWZ/eTCNsnF 3zb9vQN4enP+ydjAZ1i8vk6HvzQ8dP2cKkpYLV+lRvcgfInG8Ix/ZPbBvX7qjd2d 7J5l/SwuA+qNaqjv+aVVp9CutHy4YBVhZP0HLMQpvsx9Mm5UMMYqN8SkxQELvrqf yLVVxqgkamNz8s0t+CGoFLntBbWRWsNNFlzB9SvLn+0gjQLYXWE4GTUE09ccVfoM 7mNCtEK/Xh7Z0M0nHS9tH4xW2LJJhrwtypOD8XO7fJVQsSzq42e3xzz+VdlNjJ9N h0xf3amAqVyfjkN0NYHhXZKgQOqebPXdnXCKxInW2pyey7gmSH/kDvqbmHJ+ZJh0 hwCA0hO/6HyiMg63AbThR7Sjt1/13qDzXzbRDgVjNlIIWIwcBzc2k9qn+sC8uOjw 1G40wCaPTroHGCi9gqfMQUYvUSI47JM8kLsraTSJC3ftbKqAvzn58Fh6ZuD7oS+j AcHrLGWwmFGWL/PYxWbPFZbRkTMfwyliIchwLixWSANELrflWhqAgvHhNMWHWMaA SS3FTwtyLjNqu+j6TzL8Hju43zyKgI4M61Z5 =fzwj -----END PGP PUBLIC KEY BLOCK----- pub rsa4096 2026-02-24 [SC] B516D13DF4D5D5A2EA117DE6638CD75EDAE74883 uid [ultimate] Yong Wu (CODE SIGNING KEY) sig 3 638CD75EDAE74883 2026-02-24 [self-signature] sub rsa4096 2026-02-24 [E] sig 638CD75EDAE74883 2026-02-24 [self-signature] -----BEGIN PGP PUBLIC KEY BLOCK----- mQINBGmd3akBEAC7h8xzjpuiWTPqoGIT2Hk/c2xBGLlVs+cmBuhPv/VbJe2IkBvU J7FmvuaiAnjW+bR+Bj8ZsI5vkV0Kxhw+ZIsJihMe2SzhsjnRkwN8tVZMiFJ7VAt8 lYhI+egc4LxnnAAFI4Xg1MoCOXDWN/fCwSnoZBEQOGO0yOtii8Rydyw+nw+622jC otXpxQINM1Gf/u0PlmG/lqF8fd7/cDE15cssx52qYyAd/QA37xGV+hho2CU8zcwS tpxvoHBoselND8kXmTa9AiLFm7qzORInn57p3jTGk9dnZ8Q8hfscDn2L44vyL4J6 ye+X/fzq2GsxelHDKjkywR1hDzAZ08Gy/uaRgPAiVg+R6zvXxd1EYXytkFxYNKOi jQFkUYnfTL/OvWL2+Hm/67atqoXEbrt6InLMyVte5+paXm4zkt0LC52u+7aJVW0u gWJj7ADmhqUsPOdc4YTZOwBYeUNwsZ91znwlbQqzftHXzswn83tt4gvIlNbSNlV6 LcqcZOj4TbXUS0mTK2uGvREFQSLPuIxznLlkIy+Cu8Wr5ddOWe29vPYmSnoRQs/L RHfgTtDVDwGx2tXXBVq05pvtYaaQ7s3kDQO5tZig3fct26a92kC7yoLSnD3PdPrg NSG8XBLgBDJzNiAK0ZvUQIQyLoIuBgISRCr4yrERDoLB150x0IYqBzwGfwARAQAB tC9Zb25nIFd1IChDT0RFIFNJR05JTkcgS0VZKSA8eW9uZ3d3d0BhcGFjaGUub3Jn PokCUQQTAQgAOxYhBLUW0T301dWi6hF95mOM117a50iDBQJpnd2pAhsDBQsJCAcC AiICBhUKCQgLAgQWAgMBAh4HAheAAAoJEGOM117a50iDRcwP/0Y/407ehWPrdHY7 ckXkv61QeOXBM5S9J4Xw71g8Pv98Z5diqkNjdXMw45abDuGGa7HyWEJr6kj7YS9E 8F1lVo4zQQOAAniGYfWByY0z+0RH5yUu8f35nC71c3jk2o1qAYA09Jpw+sx8WmU6 Ztc7Iv2TztPKOUOD0wTET/vj5C1FBsyYbZJfiQfgVg3csPU6cwln0WUUKMK5Ew5t arYRIgZTtI9osTUa322bukgOOnB2mxbsStMwwn1uH2AB6nQwFJkjQLHCl1s8flaV taHaFfaE/k9SPzHDhPoKhnyELBkuDoihBpq2oroX6LE2456eFeJmeunrAk+0Bj1d pmIlcC79rC1aMvyMGYEEn7eYV+Wfxtb/rdRKw0a2hJiWSC7btdl+5c+DsB+3SnM6 /U5udjj1bi/2ikYTBJ0XdJCJKgdQ1EBzec/Zb5amIn6SMwNtnbunLOrtkBdhVfaJ j+s2/1FJCaFaHfVMcM56FOylTc1xy5C+4jY8iz/oGQvKBTkf6glZ0/Pt7FpZpdFS rTeQkxz3LQRK3yzB4nsvuJjJOYWqRW8n29lnsLcV4IGMLWOojE2ypvw7JLpukmTJ 5YcJ7ytMjK3U1V59aj2Ei7fFE/8dTm/NhLqDujWQ/MLp9QSYpd+NnWRHyJF+QlxC zV6/660vVAc73XZiFpU+GsxAsrqluQINBGmd3akBEACvzHtPl5sdroybt/HrtiPL 187gPqdYwxU84T/EJBSjtWRTNJ3+9w/WWW0TMYQ2sGBtOdem04828zzYhSY+H3I9 hyalCN04uF8UwhjBEOtI4sJxFCvaamASqg9E7qVXOEFGWPbmXHN27mIoLgfyz2Zo 9RbS9qcV9m/SngkhZ795Zl6fdMPn/FY/9O832SRZibEROVBcjUqReN73IKAa+YKh dH2ui3OvTcy0x7Yzc/t5lvMBDNeBLvFcnaj18SKvdHQ2WJZopXjkfHdA9EeazwIp C5d5280OqlFwKGyGWYWfgVwoODydYmN+xVNPxCa3dJ5S0xYzPYmf/Sl2EWs7Lw6N e2yO1AhxJN/TsWzZfHUpaHjrsMzlJTNHfVUPaOwelCbFDBfVybubVKlC4rpT4OAY 06cofUpZcj97shlGARiY+dHYJRmLpmbAcsQz0utht+47kJLl2wt0oHI6tvV6ESI+ w7Dx5P+aj9VJSGPr4vO7laAv2hsRHvdSvYOfpzVJSyGIUxjEE1HLjwtlyxbjfWAY e2JRLTkxHGyitXPMFsAm06Je7CA6tyx7q7shkOZXFkwNYuVYGABFri0miWkeCf8U SFHLyIwD7dwYDB9ZPmoyU1ElM+ZDJjwomSsAmrm5ejvaTAH39IUVOLacsmYaKGBV CzU6NrjAM074SzaQNK129QARAQABiQI2BBgBCAAgFiEEtRbRPfTV1aLqEX3mY4zX XtrnSIMFAmmd3akCGwwACgkQY4zXXtrnSIPPKA//QAk5rOJXxe+Asg3rN4hh+hz7 Wd8ytey7qFTN3mNUUYFMgB3gw7N3zRx4M6AQwRND0nSTZgA0F0mk3iCgDs5AtDgm cL+/JatVYW6onVsG72wB68sOrMVni5mgvP42tuUOkOpBT14Eba3OHR/J0KrrTj3/ S0qqUyRBcdO/mJM7CS6UpROfSNAVvkSExkC/Ly8f0zI66WBABKz/IUX9gOSdGUjh DRju/P1W526sPEU0kRmDgr7/Up4cBx8D07T0DbZfCVIuInvuk/c1BFpgUuXvplhi Lf2PeyEnYNuaVFQrzr8Jk7P/dKbENkQb+oqc2ozs84Fr23INLI79GxaNREnfsHxQ LW2fKlDk0EXwdec6i66cT31io3jI0k3EUiC3g4j9PMp3+pL95SLXwebLbVNlqPiz Ins3ux8Y0Jpsw6tf99m/ZnPtyIFuQGlBl+GbnbtblQctFdyUT41Nk5bDygFmGJs6 an4K55q/sV8FmA+eE92JzOyu5JFnBnwXJ7/8vP5d25hoXCI1fsfm9M4QgA4Z3tDo XuHg/PpBgqQNIvDuDGv1wbdRZZTCyz/dvE2Zo4LMjHcrll0xyAfkjLtsnPV28k2z jhgdmBlJMCRLBkWRTHcNwKUbd51MLfD5j47FUGiedzCoDZ/gw7imYq5KZu0Pqzx4 v45Ng1UNbmbJVnCLxKU= =cZH7 -----END PGP PUBLIC KEY BLOCK----- tvm-ffi-0.1.12/LICENSE000066400000000000000000000270711521067262500141650ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "{}" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright {yyyy} {name of copyright owner} Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ------------------------------------------------------------------------------------ This product bundles various third-party components under other open source licenses. This section summarizes those components and their licenses. See licenses/ for text of these licenses. Apache Software Foundation License 2.0 -------------------------------------- 3rdparty/dlpack BSD 3-Clause "New" or "Revised" License --------------------------------------- 3rdparty/libbacktrace tvm-ffi-0.1.12/NOTICE000066400000000000000000000002531521067262500140550ustar00rootroot00000000000000Apache TVM FFI Copyright 2024-present The Apache Software Foundation This product includes software developed at The Apache Software Foundation (http://www.apache.org/). tvm-ffi-0.1.12/README.md000066400000000000000000000125711521067262500144360ustar00rootroot00000000000000 # TVM FFI: Open ABI and FFI for Machine Learning Systems 📚 [Documentation](https://tvm.apache.org/ffi/) | 🚀 [Quickstart](https://tvm.apache.org/ffi/get_started/quickstart.html) Apache TVM FFI is an open ABI and FFI for machine learning systems. It is a minimal, framework-agnostic, yet flexible open convention with the following systems in mind: - **Kernel libraries** - ship one wheel to support multiple frameworks, Python versions, and different languages. [[FlashInfer](https://docs.flashinfer.ai/)] - **Kernel DSLs** - reusable open ABI for JIT and AOT kernel exposure frameworks and runtimes. [[TileLang](https://tilelang.com/)][[cuteDSL](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.html)] - **Frameworks and runtimes** - a uniform extension point for ABI-compliant libraries and DSLs. [[PyTorch](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-pytorch)][[JAX](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-jax)][[PaddlePaddle](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-paddle)][[NumPy/CuPy](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-numpy)] - **ML infrastructure** - out-of-box bindings and interop across languages. [[Python](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-python)][[C++](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-cpp)][[Rust](https://tvm.apache.org/ffi/get_started/quickstart.html#ship-to-rust)][[XGrammar](https://github.com/mlc-ai/xgrammar)] - **Coding agents** - a unified mechanism for shipping generated code in production. ## Features - **Stable, minimal C ABI** designed for kernels, DSLs, and runtime extensibility. - **Zero-copy interop** across PyTorch, JAX, and CuPy using [DLPack protocol](https://data-apis.org/array-api/2024.12/design_topics/data_interchange.html). - **Compact value and call convention** covering common data types for ultra low-overhead ML applications. - **Multi-language support** out of the box: Python, C++, and Rust (with a path towards more languages). These enable broad **interoperability** across frameworks, libraries, DSLs, and agents; the ability to **ship one wheel** for multiple frameworks and Python versions (including free-threaded Python); and consistent infrastructure across environments. ## Getting Started Install TVM-FFI with pip, uv or from source: ```bash pip install apache-tvm-ffi pip install torch-c-dlpack-ext # compatibility package for torch <= 2.9 ``` ## Status and Release Versioning **C ABI stability** is our top priority. **Status: RFC** Main features are complete and ABI stable. We recognize potential needs for evolution to ensure it works best for the machine learning systems community, and would like to work together with the community for such evolution. We plan to stay in the RFC stage for three months from the v0.1.0 release. Releases during the RFC stage will be `0.X.Y`, where bumps in `X` indicate C ABI-breaking changes and `Y` indicates other changes. We anticipate the RFC stage will last for three months, then we will start following [Semantic Versioning](https://packaging.python.org/en/latest/discussions/versioning/) (`major.minor.patch`) going forward. ## Documentation Our [documentation site](https://tvm.apache.org/ffi/) includes: ### Get Started - [Quick Start](https://tvm.apache.org/ffi/get_started/quickstart.html) - [Stable C ABI](https://tvm.apache.org/ffi/get_started/stable_c_abi.html) ### Guides - [Export Functions & Classes](https://tvm.apache.org/ffi/guides/export_func_cls.html) - [Kernel Library Guide](https://tvm.apache.org/ffi/guides/kernel_library_guide.html) ### Concepts - [ABI Overview](https://tvm.apache.org/ffi/concepts/abi_overview.html) - [Any](https://tvm.apache.org/ffi/concepts/any.html) - [Object & Class](https://tvm.apache.org/ffi/concepts/object_and_class.html) - [Tensor](https://tvm.apache.org/ffi/concepts/tensor.html) - [Function & Module](https://tvm.apache.org/ffi/concepts/func_module.html) - [Exception Handling](https://tvm.apache.org/ffi/concepts/exception_handling.html) ### Packaging - [Python Packaging](https://tvm.apache.org/ffi/packaging/python_packaging.html) - [Stub Generation](https://tvm.apache.org/ffi/packaging/stubgen.html) - [C++ Tooling](https://tvm.apache.org/ffi/packaging/cpp_tooling.html) ### Developer Manual - [Build from Source](https://tvm.apache.org/ffi/dev/source_build.html) - [Reproduce CI/CD](https://tvm.apache.org/ffi/dev/ci_cd.html) - [Release Process](https://tvm.apache.org/ffi/dev/release_process.html) tvm-ffi-0.1.12/addons/000077500000000000000000000000001521067262500144215ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/000077500000000000000000000000001521067262500202405ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/LICENSE000066400000000000000000000261351521067262500212540ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "{}" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright {yyyy} {name of copyright owner} Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/NOTICE000066400000000000000000000002651521067262500211470ustar00rootroot00000000000000Torch C DLPack Extension Copyright 2024-present The Apache Software Foundation This product includes software developed at The Apache Software Foundation (http://www.apache.org/). tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/README.md000066400000000000000000000027051521067262500215230ustar00rootroot00000000000000 # Torch C DLPack Extension This folder contains the source for the `torch-c-dlpack-ext` package, which provides an Ahead-Of-Time (AOT) compiled module to support faster DLPack conversion in DLPack v1.2. By default, `tvm-ffi` will JIT-compile a version of this functionality during loading, and use a safe-fallback if JIT-compilation fails. Installing this wheel allows users to avoid this JIT compilation overhead and also avoid the cases where the user environment does not necessarily have a compiler toolchain to run JIT-compilation. ```bash pip install torch-c-dlpack-ext ``` tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/build_aot_wheels.bat000066400000000000000000000072641521067262500242520ustar00rootroot00000000000000@REM Licensed to the Apache Software Foundation (ASF) under one @REM or more contributor license agreements. See the NOTICE file @REM distributed with this work for additional information @REM regarding copyright ownership. The ASF licenses this file @REM to you under the Apache License, Version 2.0 (the @REM "License"); you may not use this file except in compliance @REM with the License. You may obtain a copy of the License at @REM @REM http://www.apache.org/licenses/LICENSE-2.0 @REM @REM Unless required by applicable law or agreed to in writing, @REM software distributed under the License is distributed on an @REM "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY @REM KIND, either express or implied. See the License for the @REM specific language governing permissions and limitations @REM under the License. @echo off setlocal enabledelayedexpansion set arch=%~1 set python_version=%~2 set tvm_ffi=%cd% set torch_c_dlpack_ext=%tvm_ffi%\addons\torch_c_dlpack_ext if not exist "%tvm_ffi%\.venv" mkdir "%tvm_ffi%\.venv" if not exist "%tvm_ffi%\lib" mkdir "%tvm_ffi%\lib" for %%P in (2.4 2.5 2.6 2.7 2.8 2.9) do ( call :build_libs %%P ) copy %tvm_ffi%\lib\*.dll %torch_c_dlpack_ext%\torch_c_dlpack_ext uv venv %tvm_ffi%\.venv\build --python %python_version% call %tvm_ffi%\.venv\build\Scripts\activate uv pip install build wheel cd %torch_c_dlpack_ext% python -m build -w dir dist for %%f in (dist\*.whl) do python -m wheel tags "%%f" --python-tag=%python_version% --abi-tag=%python_version% --platform-tag=win_amd64 dir dist mkdir wheelhouse copy dist\*-win_amd64.whl wheelhouse dir wheelhouse endlocal exit /b :build_libs set torch_version=%1 call :check_availability if %errorlevel%==0 ( call :get_torch_url uv venv %tvm_ffi%\.venv\torch%torch_version% --python %python_version% call %tvm_ffi%\.venv\torch%torch_version%\Scripts\activate uv pip install setuptools ninja uv pip install torch==%torch_version% --index-url !torch_url! uv pip install -v . python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir %tvm_ffi%\lib python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir %tvm_ffi%\lib --build-with-cuda call deactivate rmdir -s -q %tvm_ffi%\.venv\torch%torch_version% ) else ( echo Skipping build for torch %torch_version% on %arch% with python %python_version% as it is not available. ) exit /b 0 :check_availability if %torch_version%==2.4 ( if %python_version%==cp313 exit /b 1 if %python_version%==cp314 exit /b 1 exit /b 0 ) if %torch_version%==2.5 ( if %python_version%==cp314 exit /b 1 exit /b 0 ) if %torch_version%==2.6 ( if %python_version%==cp314 exit /b 1 exit /b 0 ) if %torch_version%==2.7 ( if %python_version%==cp314 exit /b 1 exit /b 0 ) if %torch_version%==2.8 ( if %python_version%==cp314 exit /b 1 exit /b 0 ) if %torch_version%==2.9 ( if %python_version%==cp39 exit /b 1 exit /b 0 ) echo Unknown or unsupported torch version: %torch_version% >&2 exit /b 1 :get_torch_url set cuda_version= if %torch_version%==2.4 set cuda_version=cu124 if %torch_version%==2.5 set cuda_version=cu124 if %torch_version%==2.6 set cuda_version=cu126 if %torch_version%==2.7 set cuda_version=cu128 if %torch_version%==2.8 set cuda_version=cu129 if %torch_version%==2.9 set cuda_version=cu129 if defined cuda_version ( set torch_url=https://download.pytorch.org/whl/%cuda_version% exit /b 0 ) echo Unknown or unsupported torch version: %torch_version% >&2 exit /b 1 tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/build_aot_wheels.sh000077500000000000000000000110371521067262500241120ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # shellcheck disable=SC1090,1091 set -eux arch=$1 python_version=$2 os=$(uname -s) case "$os" in "Linux" | "Darwin") ;; *) echo "Unknown OS: $os" return 1 ;; esac tvm_ffi="$PWD" torch_c_dlpack_ext="$tvm_ffi"/addons/torch_c_dlpack_ext function get_torch_url() { local version="$1" case "$version" in "2.4" | "2.5") echo "https://download.pytorch.org/whl/cu124" ;; "2.6") echo "https://download.pytorch.org/whl/cu126" ;; "2.7") echo "https://download.pytorch.org/whl/cu128" ;; "2.8" | "2.9") echo "https://download.pytorch.org/whl/cu129" ;; *) echo "Unknown or unsupported torch version: $version" >&2 return 1 ;; esac } function check_availability() { local torch_version="$1" case "$torch_version" in "2.4") ! [[ "$arch" == "aarch64" || "$python_version" == "cp313" || "$python_version" == "cp314" ]] ;; "2.5") ! [[ ("$os" == "Linux" && ("$arch" == "aarch64" || "$python_version" == "cp314")) || ("$os" == "Darwin" && ("$python_version" == "cp313" || "$python_version" == "cp314")) ]] ;; "2.6") ! [[ "$arch" == "aarch64" || "$python_version" == "cp314" ]] ;; "2.7" | "2.8") ! [[ "$python_version" == "cp314" ]] ;; "2.9") ! [[ "$python_version" == "cp39" ]] ;; *) echo "Unknown or unsupported torch version: $torch_version" >&2 return 1 ;; esac } function build_libs() { local torch_version=$1 if check_availability "$torch_version"; then uv venv "$tvm_ffi"/.venv/torch"$torch_version" --python "$python_version" source "$tvm_ffi"/.venv/torch"$torch_version"/bin/activate uv pip install setuptools ninja if [[ "$os" == "Linux" ]]; then uv pip install torch=="$torch_version" --index-url "$(get_torch_url "$torch_version")" else uv pip install torch=="$torch_version" fi uv pip install -v . python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir "$tvm_ffi"/lib if [[ "$os" == "Linux" ]]; then python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir "$tvm_ffi"/lib --build-with-cuda fi ls "$tvm_ffi"/lib deactivate rm -rf "$tvm_ffi"/.venv/torch"$torch_version" else echo "Skipping build for torch $torch_version on $arch with python $python_version as it is not available." fi } mkdir -p "$tvm_ffi"/.venv mkdir -p "$tvm_ffi"/lib torch_versions=("2.4" "2.5" "2.6" "2.7" "2.8" "2.9") for version in "${torch_versions[@]}"; do build_libs "$version" done cp "$tvm_ffi"/lib/*.so "$torch_c_dlpack_ext"/torch_c_dlpack_ext uv venv "$tvm_ffi"/.venv/build --python "$python_version" source "$tvm_ffi"/.venv/build/bin/activate uv pip install build wheel cd "$torch_c_dlpack_ext" python -m build -w ls dist if [[ "$os" == "Linux" ]]; then python -m wheel tags dist/*.whl --python-tag="$python_version" --abi-tag="$python_version" --remove uv pip install auditwheel auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libc10.so --exclude libtorch_python.so --exclude libtorch_cuda.so --exclude libc10_cuda.so dist/*.whl -w wheelhouse else python -m wheel tags dist/*.whl --python-tag="$python_version" --abi-tag="$python_version" --platform-tag=macosx_11_0_arm64 --remove uv pip install delocate delocate-wheel -v --ignore-missing-dependencies --exclude libtorch.dylib,libtorch_cpu.dylib,libc10.dylib,libtorch_python.dylib dist/*.whl -w wheelhouse fi ls wheelhouse tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/build_backend.py000066400000000000000000000075111521067262500233640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """build backend for torch c dlpack ext.""" from __future__ import annotations import os import subprocess import sys from contextlib import suppress from pathlib import Path from setuptools import build_meta as orig _root = Path(__file__).parent.resolve() _package_path = _root / "torch_c_dlpack_ext" get_requires_for_build_sdist = orig.get_requires_for_build_sdist get_requires_for_build_editable = orig.get_requires_for_build_editable prepare_metadata_for_build_wheel = orig.prepare_metadata_for_build_wheel prepare_metadata_for_build_editable = orig.prepare_metadata_for_build_editable build_sdist = orig.build_sdist build_editable = orig.build_editable def _is_lib_prebuilt() -> bool: if sys.platform.startswith("win32"): extension = "dll" else: extension = "so" return next(_package_path.rglob(f"*.{extension}"), None) is not None def get_requires_for_build_wheel( config_settings: orig._ConfigSettings = None, ) -> list[str]: """Get build requirements for wheel, conditionally including torch and apache-tvm-ffi.""" requires = orig.get_requires_for_build_wheel(config_settings) if not _is_lib_prebuilt(): # build wheel from sdist package, requires apache-tvm-ffi>=0.1.1 to build lib requires.append("apache-tvm-ffi>=0.1.1") return requires def build_wheel( wheel_directory: orig.StrPath, config_settings: orig._ConfigSettings = None, metadata_directory: orig.StrPath | None = None, ) -> str: """Build wheel.""" torch = None with suppress(ModuleNotFoundError): import torch # noqa: PLC0415 if torch is not None and not _is_lib_prebuilt(): # build wheel from sdist package, compile the torch c dlpack ext library locally. if hasattr(torch.Tensor, "__dlpack_c_exchange_api__") or hasattr( torch.Tensor, "__c_dlpack_exchange_api__" ): print( "torch.Tensor already has attribute __dlpack_c_exchange_api__. " "No need to build any torch c dlpackc libs." ) else: extra_args = [] # First use "torch.cuda.is_available()" to check whether GPU environment # is available. Then determine the GPU type. if torch.cuda.is_available(): if torch.version.cuda is not None: extra_args.append("--build-with-cuda") elif torch.version.hip is not None: extra_args.append("--build-with-rocm") else: raise ValueError("Cannot determine whether to build with CUDA or ROCm.") subprocess.run( [ sys.executable, "-m", "tvm_ffi.utils._build_optional_torch_c_dlpack", "--output-dir", str(_package_path), *extra_args, ], check=True, env={**os.environ, "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1"}, ) return orig.build_wheel(wheel_directory, config_settings, metadata_directory) tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/pyproject.toml000066400000000000000000000030451521067262500231560ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [project] name = "torch_c_dlpack_ext" version = "0.1.5" requires-python = ">=3.9" description = "torch c dlpack ext" dependencies = ["torch"] authors = [{ name = "TVM FFI team" }] readme = "README.md" license = { file = "LICENSE" } classifiers = [ "License :: OSI Approved :: Apache Software License", "Development Status :: 4 - Beta", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", ] [build-system] requires = ["setuptools>=61.0"] build-backend = "build_backend" backend-path = ["."] [tool.setuptools] include-package-data = true py-modules = ["build_backend"] [tool.setuptools.packages.find] include = ["torch_c_dlpack_ext*"] [tool.setuptools.package-data] torch_c_dlpack_ext = ["*.so", "*.dll", "*.dylib"] tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/000077500000000000000000000000001521067262500240575ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/__init__.py000066400000000000000000000015071521067262500261730ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """torch c dlpack ext pakcage.""" from . import core tvm-ffi-0.1.12/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py000066400000000000000000000064561521067262500253740ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """torch c dlpack ext core methods.""" import ctypes import sys from pathlib import Path from typing import Any import torch from packaging.version import Version def _create_dlpack_exchange_api_capsule(ptr_as_int: int) -> Any: """Create a PyCapsule wrapping the DLPack exchange API pointer.""" capsule_name = b"dlpack_exchange_api" pythonapi = ctypes.pythonapi pythonapi.PyCapsule_New.restype = ctypes.py_object pythonapi.PyCapsule_New.argtypes = [ ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ] capsule = pythonapi.PyCapsule_New(ctypes.c_void_p(ptr_as_int), capsule_name, None) return capsule def _torch_extension_device(torch_module: Any) -> str: """Return the torch backend name used in the optional extension library name.""" if torch_module.cuda.is_available(): if getattr(torch_module.version, "cuda", None) is not None: return "cuda" if getattr(torch_module.version, "hip", None) is not None: return "rocm" return "cuda" return "cpu" def load_torch_c_dlpack_extension() -> None: """Load the torch c dlpack extension based on torch version.""" if hasattr(torch.Tensor, "__dlpack_c_exchange_api__") or hasattr( torch.Tensor, "__c_dlpack_exchange_api__" ): return None version = Version(torch.__version__) if sys.platform.startswith("win32"): extension = "dll" elif sys.platform.startswith("darwin"): extension = "dylib" else: extension = "so" device = _torch_extension_device(torch) lib_path = ( Path(__file__).parent / f"libtorch_c_dlpack_addon_torch{version.major}{version.minor}-{device}.{extension}" ) if not lib_path.exists() or not lib_path.is_file(): raise ImportError("No matching prebuilt torch c dlpack extension") lib = ctypes.CDLL(str(lib_path)) func = lib.TorchDLPackExchangeAPIPtr func.restype = ctypes.c_uint64 func.argtypes = [] # note: we need to keep this behavior for a while # to ensure backward compatibility with older versions dependencies # that relies on the value being int. # We will do eager upgrade to PyCapsule in the tvm-ffi side instead. dlpack_exchange_api_ptr_as_int = func() setattr(torch.Tensor, "__c_dlpack_exchange_api__", dlpack_exchange_api_ptr_as_int) setattr( torch.Tensor, "__dlpack_c_exchange_api__", _create_dlpack_exchange_api_capsule(dlpack_exchange_api_ptr_as_int), ) return lib _lib = load_torch_c_dlpack_extension() tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/000077500000000000000000000000001521067262500174255ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/CMakeLists.txt000066400000000000000000000130431521067262500221660ustar00rootroot00000000000000cmake_minimum_required(VERSION 3.20) project( tvm_ffi_orcjit VERSION 0.1.0 LANGUAGES C CXX ) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) # ---- LLVM_PREFIX handling ---- # scikit-build-core overrides CMAKE_PREFIX_PATH via its init cache, so we read LLVM_PREFIX from the # environment instead. if (DEFINED ENV{LLVM_PREFIX}) list(APPEND CMAKE_PREFIX_PATH "$ENV{LLVM_PREFIX}") if (EXISTS "$ENV{LLVM_PREFIX}/Library") list(APPEND CMAKE_PREFIX_PATH "$ENV{LLVM_PREFIX}/Library") endif () endif () # ---- Find packages ---- find_package(LLVM REQUIRED CONFIG) message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION} in ${LLVM_DIR}") find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # ---- Build shared library ---- add_library( tvm_ffi_orcjit SHARED src/ffi/orcjit_session.cc src/ffi/orcjit_dylib.cc src/ffi/orcjit_memory_manager.cc src/ffi/orcjit_slab.cc src/ffi/llvm_patches/gotpcrelx_fix.cc src/ffi/llvm_patches/init_fini_plugin.cc src/ffi/llvm_patches/macho_cxa_atexit_shim.cc src/ffi/llvm_patches/win_coff_pdata_strip.cc src/ffi/llvm_patches/win_dll_import_generator.cc ) set_target_properties( tvm_ffi_orcjit PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON ) target_include_directories( tvm_ffi_orcjit PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${LLVM_INCLUDE_DIRS} ) separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS}) target_compile_definitions(tvm_ffi_orcjit PRIVATE ${LLVM_DEFINITIONS_LIST}) # ---- Static LLVM linking via llvm-config ---- set(_llvm_config_hints ${LLVM_TOOLS_BINARY_DIR} "${LLVM_DIR}/../../../bin") if (DEFINED ENV{LLVM_PREFIX}) list(APPEND _llvm_config_hints "$ENV{LLVM_PREFIX}/bin" "$ENV{LLVM_PREFIX}/Library/bin") endif () find_program( LLVM_CONFIG_EXE NAMES llvm-config-${LLVM_VERSION_MAJOR} llvm-config HINTS ${_llvm_config_hints} ) if (NOT LLVM_CONFIG_EXE) message(FATAL_ERROR "llvm-config not found") endif () execute_process( COMMAND ${LLVM_CONFIG_EXE} --link-static --libs Core OrcJIT Support native OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE _llvm_libs ) execute_process( COMMAND ${LLVM_CONFIG_EXE} --link-static --ldflags OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE _llvm_ldflags ) separate_arguments(_llvm_libs_list NATIVE_COMMAND "${_llvm_libs}") separate_arguments(_llvm_ldflags_list NATIVE_COMMAND "${_llvm_ldflags}") # ---- zlib (static, from LLVM prefix) ---- cmake_path(GET LLVM_DIR PARENT_PATH _llvm_prefix) cmake_path(GET _llvm_prefix PARENT_PATH _llvm_prefix) cmake_path(GET _llvm_prefix PARENT_PATH _llvm_prefix) set(_lib_hints "${_llvm_prefix}/lib" "${_llvm_prefix}/lib64" "${_llvm_prefix}/Library/lib") if (WIN32) find_library( ZLIB_STATIC NAMES zlibstatic z HINTS ${_lib_hints} REQUIRED ) else () find_library( ZLIB_STATIC NAMES libz.a HINTS ${_lib_hints} REQUIRED ) endif () # ---- zstd (static, from LLVM prefix) ---- if (WIN32) find_library( ZSTD_STATIC NAMES zstd_static zstd HINTS ${_lib_hints} REQUIRED ) else () find_library( ZSTD_STATIC NAMES libzstd.a HINTS ${_lib_hints} REQUIRED ) endif () # ---- Link everything ---- target_link_libraries( tvm_ffi_orcjit PRIVATE tvm_ffi::header tvm_ffi::shared ${_llvm_ldflags_list} ${_llvm_libs_list} ${ZLIB_STATIC} ${ZSTD_STATIC} ) # LLVM system libs (ntdll/psapi on Windows, rt/dl on Linux), minus zlib/zstd (linked above). execute_process( COMMAND ${LLVM_CONFIG_EXE} --system-libs OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE _sys_libs RESULT_VARIABLE _sys_libs_rc ) if (_sys_libs_rc EQUAL 0 AND _sys_libs) separate_arguments(_sys_libs_list NATIVE_COMMAND "${_sys_libs}") list(FILTER _sys_libs_list EXCLUDE REGEX "(zstd|^-lz$|^z\\.lib$|xml2)") target_link_libraries(tvm_ffi_orcjit PRIVATE ${_sys_libs_list}) endif () # ---- Platform-specific fixups ---- if (APPLE) add_custom_command( TARGET tvm_ffi_orcjit POST_BUILD COMMAND install_name_tool -change @rpath/libc++.1.dylib /usr/lib/libc++.1.dylib $ COMMENT "Fixing libc++ rpath to use system library" ) elseif (UNIX AND NOT WIN32) target_link_options(tvm_ffi_orcjit PRIVATE -static-libstdc++ -static-libgcc) endif () # ---- Find and bundle liborc_rt ---- # Platform notes: # # * Linux: liborc_rt is used for ExecutorNativePlatform → ELFNixPlatform. # * macOS: MachOPlatform is skipped to sidestep the compact-unwind delta bug (see # src/ffi/orcjit_session.cc); liborc_rt is unused. # * Windows: COFFPlatform is not hooked up; liborc_rt is unused. set(ORC_RT_PATH "") if (NOT APPLE AND NOT WIN32) if (DEFINED ENV{ORC_RT_PATH}) set(ORC_RT_PATH "$ENV{ORC_RT_PATH}") else () file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/lib/clang/*/lib/liborc_rt*.a") if (_orc_rt_candidates) list(GET _orc_rt_candidates 0 ORC_RT_PATH) endif () endif () if (NOT ORC_RT_PATH) message(WARNING "Could not find liborc_rt. ORC runtime features will be disabled.") endif () endif () if (ORC_RT_PATH) message(STATUS "Found liborc_rt: ${ORC_RT_PATH}") file(COPY "${ORC_RT_PATH}" DESTINATION "${CMAKE_BINARY_DIR}") install(FILES "${ORC_RT_PATH}" DESTINATION lib) endif () # ---- Install ---- install( TARGETS tvm_ffi_orcjit LIBRARY DESTINATION lib ARCHIVE DESTINATION lib RUNTIME DESTINATION lib ) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md000066400000000000000000000552031521067262500220240ustar00rootroot00000000000000 # ORC JIT v2 Primer — Background for the `tvm_ffi_orcjit` Addon This document explains the background knowledge needed to understand the `tvm_ffi_orcjit` addon: object file formats, the classical linker model, and LLVM's ORC JIT v2 architecture. It then maps those concepts onto the addon's implementation. --- ## 1. Object File Formats A **compiled object file** (`.o` / `.obj`) is not an executable. It is an intermediate container that holds machine code, data, and metadata about unresolved references. The three dominant formats are: | Format | Platform | File extensions | | --- | --- | --- | | **ELF** (Executable and Linkable Format) | Linux, most Unix | `.o`, `.so`, `.elf` | | **Mach-O** (Mach Object) | macOS, iOS | `.o`, `.dylib`, `.macho` | | **COFF/PE** (Common Object File Format / Portable Executable) | Windows | `.obj`, `.dll`, `.exe` | Despite surface differences, all three share the same conceptual structure: ```text ┌─────────────────────────────────┐ │ File Header │ magic, target arch, section count ├─────────────────────────────────┤ │ Section Headers │ name, type, file offset, size, flags ├─────────────────────────────────┤ │ .text (code) │ machine instructions │ .rodata (read-only data) │ string literals, constants │ .data (writable data) │ initialized globals │ .bss (zero-init data) │ uninitialized globals (no file bytes) │ .init_array (constructors) │ array of function pointers — C++ ctors │ .fini_array (destructors) │ array of function pointers — C++ dtors │ .eh_frame (unwind info) │ exception/stack-unwind tables │ … other sections … │ ├─────────────────────────────────┤ │ Symbol Table │ name → (section, offset, binding, type) ├─────────────────────────────────┤ │ Relocation Tables │ (section, offset, symbol, type, addend) └─────────────────────────────────┘ ``` ### Symbols A **symbol** is a named location in the file — a function entry point, a global variable, or a section boundary. Symbols have: - **binding**: `LOCAL` (file-private), `GLOBAL` (externally visible), `WEAK` (override-able default) - **definition status**: *defined* (has an address in this file) vs. *undefined* (imported — must be resolved at link time) ### Relocations A **relocation** is a "fixup recipe" stored alongside the machine code. It says: > "At byte offset X in section S, patch in the address of symbol FOO, using formula R." The formula R is a **relocation type** that encodes what arithmetic to apply: | Type | Meaning | | --- | --- | | `R_X86_64_64` | Absolute 64-bit address of the symbol | | `R_X86_64_PC32` | Symbol address minus the patch location (PC-relative 32-bit) | | `R_X86_64_PLT32` | PC-relative call through the PLT (procedure linkage table) | | `IMAGE_REL_AMD64_ADDR32NB` (COFF) | Address relative to image base (Pointer32NB) | The linker (or JIT) processes these tables to produce the final binary. ### Platform-Specific Initialization Sections C++ global constructors and destructors must run before/after `main`. Compilers encode them as arrays of function pointers in special sections: | Platform | Constructors | Destructors | Notes | | --- | --- | --- | --- | | ELF (Linux) | `.init_array` | `.fini_array` | Priority suffix `.init_array.NNN` (lower = earlier) | | ELF (legacy) | `.ctors` | `.dtors` | Older GCC style | | Mach-O | `__DATA,__mod_init_func` | `__DATA,__mod_term_func` | Processed by dyld | | COFF (Windows) | `.CRT$XCU` (default) | `.CRT$XTZ` | Suffix encodes priority via ASCII ordering | The linker or OS loader is responsible for iterating these arrays and calling each pointer in priority order before handing control to user code. A JIT that loads `.o` files must replicate this behavior itself — this is a key responsibility of `tvm_ffi_orcjit`. --- ## 2. The Classical Linker Model A **static linker** (`ld`, `link.exe`) takes multiple object files and libraries and produces a single loadable image (executable or shared library). The process is: ```text Object files + Archives │ ▼ ┌───────────────────┐ │ Symbol Resolution │ Match every undefined symbol to a definition. │ │ Archive (.a) members pulled in on demand. └────────┬──────────┘ │ ▼ ┌───────────────────┐ │ Section Merging │ All .text sections → one .text; same for .data, etc. │ │ Assign virtual addresses (VMA) to each section. └────────┬──────────┘ │ ▼ ┌───────────────────┐ │ Relocation │ Apply each relocation fixup now that all VMAs are known. │ │ Patch raw bytes in the output image. └────────┬──────────┘ │ ▼ ┌───────────────────┐ │ Output Emission │ Write ELF/PE/Mach-O header + sections + dynamic table. └───────────────────┘ ``` ### Dynamic Linking At **load time** (when `dlopen` / `LoadLibrary` runs), the OS dynamic linker (`ld.so`, `dyld`, `ntdll`) performs a *reduced* version of the above: 1. Map shared library sections into the process address space. 2. Resolve cross-library symbol references (PLT/GOT stubs). 3. Apply load-time relocations for position-dependent code. 4. Run `.init_array` / `__mod_init_func` constructors. ### Why JIT Is Different A JIT allocates memory at arbitrary addresses *at runtime*. It cannot use precomputed link-time addresses. Every relocation must be re-evaluated against the JIT-allocated addresses. Furthermore, if the JIT loads object files incrementally (one at a time), symbol resolution must be deferred until all relevant objects are present. --- ## 3. LLVM ORC JIT v2 **ORC** stands for *On Request Compilation*. LLVM ORC JIT v2 (introduced in LLVM 9, stabilized in LLVM 13+) is a complete redesign of LLVM's JIT infrastructure. It is designed to be composable, asynchronous, and correct for production use (unlike the older `MCJIT` which had several fundamental limitations around multi-module linking). ### 3.1 Core Concepts #### ExecutionSession `llvm::orc::ExecutionSession` is the root object of any ORC JIT instance. It owns: - A set of `JITDylib`s (the symbol namespaces). - The dispatch mechanism for asynchronous compilation tasks. - The global symbol interning table (maps string → `SymbolStringPtr`). Think of it as the "JIT process" — one per logical JIT environment. #### JITDylib (JIT Dynamic Library) A `JITDylib` is a **symbol namespace** that loosely mirrors a shared library. It: - Holds a **symbol table**: `name → {flags, address}`. - Has a **link order** (a list of other `JITDylib`s to search for unresolved symbols). - Can be populated via *materialization units* (object files, LLVM IR, inline asm, etc.). Multiple `JITDylib`s can coexist in one `ExecutionSession`, enabling isolation: e.g., one library per compiled kernel, sharing a common runtime library. ```text ExecutionSession ├── JITDylib "main" ← default, process symbols ├── JITDylib "libA" ← user kernel A │ link_order: [libA, libB, main] └── JITDylib "libB" ← user kernel B (shared by A) ``` #### LLJIT `llvm::orc::LLJIT` is the high-level, batteries-included wrapper around `ExecutionSession`. It: - Sets up a target machine and data layout. - Creates a default `JITDylib` ("main") with process symbol support. - Configures the linking pipeline (see below). - Exposes `addObjectFile()`, `addIRModule()`, and `lookup()`. `LLJIT` is what `tvm_ffi_orcjit` wraps in `ORCJITExecutionSessionObj`. #### MaterializationUnit and MaterializationResponsibility A **MaterializationUnit** is a lazy producer of symbols. When a symbol is first looked up and not yet defined, ORC triggers its materialization unit to produce the definition asynchronously. This is the "on request" in ORC. An object file becomes a `StaticLibraryDefinitionGenerator` or `ObjectLayer`-level unit: the object is parsed, linked, and its symbols resolved only when someone asks for them. #### Layers ORC processes objects through a pipeline of **layers**, each transforming the input: ```text addObjectFile(buffer) │ ▼ ┌─────────────────────┐ │ ObjectTransformLayer│ (optional) transform raw object bytes before linking │ e.g. strip .pdata │ └────────┬────────────┘ │ ▼ ┌─────────────────────┐ │ ObjectLinkingLayer │ runs JITLink to: │ (uses JITLink) │ • parse object format │ │ • resolve symbols across JITDylibs │ │ • apply relocations │ │ • allocate + write JIT memory └────────┬────────────┘ │ ▼ JITDylib symbol table updated; code is live ``` #### JITLink `JITLink` is LLVM's low-level, graph-based linker used inside `ObjectLinkingLayer`. It: 1. **Parses** the object file into an in-memory `LinkGraph` (nodes = sections/atoms, edges = relocations). 2. **Runs pass pipelines** (pre-prune, post-allocation, post-fixup) — plugins can inspect and modify the graph at each stage. 3. **Allocates** JIT memory (code + data) via a `JITLinkMemoryManager`. 4. **Resolves** relocations using the current `ExecutionSession` symbol lookup. 5. **Finalizes** by writing machine code into the allocated pages and marking them executable. The pass pipeline is where `tvm_ffi_orcjit`'s `InitFiniPlugin` does its work. ### 3.2 Symbol Lookup and Resolution When `session.lookup(search_order, symbol_name)` is called: ```text 1. Check each JITDylib in search_order, in order. 2. If symbol is defined → return its ExecutorAddr. 3. If symbol is in a MaterializationUnit not yet realized → trigger materialization. 4. Materialization runs JITLink for the relevant object → symbol becomes defined. 5. Return the address. ``` Symbol names are **mangled** (C++ name mangling) or explicitly prefixed. `tvm_ffi_orcjit` uses the `__tvm_ffi_` prefix to namespace exported functions. ### 3.3 DefinitionGenerators A `DefinitionGenerator` is a fallback attached to a `JITDylib`. When a symbol is not found in the dylib's own table, the generator is invoked to dynamically create a definition — typically by looking up a symbol in the host process or another library. `LLJIT` attaches a `DynamicLibrarySearchGenerator` ("ProcessSymbols") to the main dylib, which resolves symbols like `malloc`, `printf`, or any C runtime function by looking them up in the running process. On Windows, `tvm_ffi_orcjit` adds a custom `DLLImportDefinitionGenerator` that handles `__imp_XXX` import stubs that MSVC-compiled objects expect. ### 3.4 Platform Support #### What an ORC Platform Is When the OS dynamic linker loads a shared library, it does more than just map bytes into memory — it runs constructors, registers `atexit` handlers, and sets up thread-local storage. An **ORC platform** is the JIT-side object that replicates this OS loader behavior for JIT-linked code. Concretely, an ORC platform: 1. **Intercepts `__cxa_atexit`**: C++ objects with destructors register their cleanup via `__cxa_atexit`. The platform installs an interposer so these registrations are scoped to a `JITDylib` and can be drained when that dylib is torn down — instead of running at process exit like normal. 2. **Drives initialization**: After an object file is linked into a `JITDylib`, calling `jit_->initialize(dylib)` asks the platform to run constructors (e.g., iterate `__mod_init_func` on macOS, `.init_array` on ELF). 3. **Drives deinitialization**: `jit_->deinitialize(dylib)` drains the `atexit` handlers registered during step 2 and runs destructors. This requires a small **ORC runtime library** (part of LLVM's `compiler-rt`) to be compiled for the target and loaded into the JIT — it provides the actual `__cxa_atexit` interposer and initialization trampolines that the platform object coordinates with. The three platform objects in LLVM are: | Platform | OS | Init section driven | | --- | --- | --- | | `MachOPlatform` | macOS / iOS | `__DATA,__mod_init_func` | | `ELFNativePlatform` | Linux / ELF | `.init_array`, TLS | | `COFFPlatform` | Windows | `.CRT$XC*` init, `__cxa_atexit` interop | `ExecutorNativePlatform` is a convenience builder that auto-selects the right platform for the host OS and loads the ORC runtime from a given path. #### How `tvm_ffi_orcjit` Uses (or Avoids) ORC Platforms The addon takes a different approach on each platform: - **macOS**: ORC platform support is *optional*. When the caller passes an ORC runtime path to `ExecutionSession`, `ExecutorNativePlatform` activates `MachOPlatform`. `jit_->initialize(dylib)` and `jit_->deinitialize(dylib)` then drive `__mod_init_func` and `__cxa_atexit` teardown natively. Without the path, the addon falls back to its own `InitFiniPlugin`. - **Windows**: `COFFPlatform` is skipped entirely because it requires MSVC CRT symbols (`_CxxThrowException`, RTTI vtables, iostream objects) that are not resolvable in the JIT context. Instead, `InitFiniPlugin` manually handles `.CRT$XC*` / `.CRT$XT*` init/fini sections. - **Linux**: `ELFNativePlatform` is not used. `InitFiniPlugin` handles `.init_array` / `.fini_array` / `.ctors` / `.dtors` directly, without the ORC runtime. --- ## 4. How `tvm_ffi_orcjit` Uses ORC JIT v2 With the background above, here is how the addon maps onto the concepts. ### 4.1 Object Model ```text Python / C++ API LLVM ORC v2 concept ───────────────────────────────────────────────────────────────── ExecutionSession LLJIT + ExecutionSession DynamicLibrary JITDylib dylib.add("foo.o") ObjectLinkingLayer.add(buffer) dylib.get_function("add") ExecutionSession.lookup("__tvm_ffi_add") dylib.set_link_order([a, b]) JITDylib.setLinkOrder([a, b, main, ...]) ``` ### 4.2 Object File Loading Pipeline ```text dylib.add("foo.o") │ ▼ ORCJITDynamicLibraryObj::AddObjectFile() │ jit_->addObjectFile(*dylib_, MemoryBuffer) │ ▼ ObjectTransformLayer (macOS: skipped; Linux/Win: may strip .pdata) │ Windows: strip .pdata/.xdata avoids JITLink COMDAT limitation │ fix __ImageBase fixes Pointer32NB relocations │ ▼ ObjectLinkingLayer (JITLink) │ Parse object → LinkGraph │ Run InitFiniPlugin passes: │ PrePrunePasses: mark .init_array / .ctors blocks as live │ PostAllocationPasses: (Windows) set __ImageBase, strip SEH sections │ PostFixupPasses: extract resolved init/fini function pointers │ → session.AddPendingInitializer(dylib, entry) │ Resolve relocations via ExecutionSession.lookup() │ Allocate JIT memory, write code, mark executable │ ▼ JITDylib symbol table updated (symbols are defined but constructors not yet run) ``` ### 4.3 Symbol Lookup and Initialization ```text dylib.get_function("add") │ ▼ ORCJITDynamicLibraryObj::GetFunction("add") │ → GetSymbol("__tvm_ffi_add") │ build JITDylibSearchOrder: [this, linked dylibs, LLJIT default] │ jit_->getExecutionSession().lookup(order, "__tvm_ffi_add") │ ← ExecutorAddr │ ▼ Linux/Windows: RunPendingInitializers() │ Sort entries by priority │ Call each function pointer (C++ ctors, .init_array entries) │ ▼ macOS: jit_->initialize(*dylib_) │ ORC MachOPlatform calls __mod_init_func pointers │ ▼ Wrap raw function pointer as tvm_ffi::Function via Function::FromPacked lambda (marshals AnyView args ↔ TVMFFIAny) ``` ### 4.4 InitFiniPlugin Detail `InitFiniPlugin` is an `ObjectLinkingLayer::Plugin` — it receives callbacks during JITLink's pass pipeline for every object being linked. ```text JITLink pass pipeline for each object file: ───────────────────────────────────────────────────────────── PrePrunePasses → InitFiniPlugin::modifyPassConfig (PrePrunePasses) For each section named .init_array* / .ctors* / .fini_array* / .dtors* (ELF) or __DATA,__mod_init_func (Mach-O) or .CRT$XC* (COFF): Mark all blocks in that section as "keep" (live) → Prevents dead-code elimination from removing constructors PostAllocationPasses (Windows only) → InitFiniPlugin::modifyPassConfig (PostAllocationPasses) Set __ImageBase = lowest allocated block address Strip .pdata / .xdata exception handler sections PostFixupPasses → InitFiniPlugin::modifyPassConfig (PostFixupPasses) Iterate all blocks in init/fini sections Read each 8-byte slot as an ExecutorAddr Parse section name to determine: priority, is_init vs is_fini Call session->AddPendingInitializer(dylib, InitFiniEntry{addr, section, priority}) (or AddPendingDeinitializer for dtors/fini_array) ``` ### 4.5 Cross-Library Symbol Resolution ```text libA depends on a symbol defined in libB: libA.set_link_order([libB]) → JITDylibSearchOrder for libA: [libA, libB, main(ProcessSymbols)] When libA's object file has an unresolved symbol "foo": JITLink asks ExecutionSession.lookup([libA, libB, main], "foo") → found in libB → returns libB's ExecutorAddr for "foo" → relocation in libA patched with that address ``` ### 4.6 Windows DLL Import Stubs Windows MSVC objects reference DLL functions through `__imp_XXX` pointer stubs (the Import Address Table pattern). At static link time the linker creates these stubs. In JIT mode there is no linker, so `DLLImportDefinitionGenerator` creates them on demand: ```text JITLink encounters undefined symbol "__imp_malloc" │ ▼ DLLImportDefinitionGenerator::tryToGenerate() │ Search ucrtbase.dll, msvcrt.dll, then all process modules │ for the real address of "malloc" │ ▼ Allocate two JIT-memory stubs: │ __imp_malloc → 8-byte slot containing &malloc (host address) │ malloc → x86_64 jmp [__imp_malloc] trampoline │ ▼ Define both symbols in the JITDylib → JITLink can now apply PCRel32 reloc to __imp_malloc (stub is close in JIT memory) ``` The stubs must live in JIT-allocated memory (not at the host process address) because x86_64 `PCRel32` relocations can only reach ±2 GB. The host's `malloc` may be farther than 2 GB from the JIT allocation. --- ## 5. End-to-End Example ```python import tvm_ffi_orcjit as oj # 1. Create an ExecutionSession (wraps LLJIT) sess = oj.ExecutionSession() # 2. Create a JITDylib lib = sess.create_dynamic_library() # 3. Load a compiled object file # → object parsed, JITLink links it, InitFiniPlugin collects ctors lib.add("add.o") # 4. Look up a function # → LLVM resolves "__tvm_ffi_add", RunPendingInitializers() fires ctors add = lib.get_function("add") # returns tvm_ffi.Function # 5. Call it result = add(3, 4) # → 7 ``` The corresponding C++ side of `add.o`: ```cpp // add.cc — compiled to add.o with clang++ -c -O2 add.cc #include static tvm::ffi::Function add_impl = [](int a, int b) { return a + b; }; // Exports symbol "__tvm_ffi_add" using TVMFFISafeCallType ABI TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); ``` --- ## 6. Key Concepts Summary | Concept | What it is | Where in the addon | | --- | --- | --- | | Object file | Container of machine code, data, symbols, relocations | Input to `dylib.add()` | | Relocation | Recipe to patch a code address at link/JIT time | Applied by JITLink | | `.init_array` / `.ctors` | Array of C++ constructor pointers in ELF objects | Collected by `InitFiniPlugin` | | `ExecutionSession` | Root of the ORC JIT environment | `ORCJITExecutionSessionObj` | | `LLJIT` | High-level ORC JIT wrapper | Stored in `ORCJITExecutionSessionObj::jit_` | | `JITDylib` | Symbol namespace / virtual shared library | `ORCJITDynamicLibraryObj::dylib_` | | `JITLink` | LLVM's JIT-aware linker | Used inside `ObjectLinkingLayer` | | JITLink pass pipeline | Pre-prune → post-alloc → post-fixup hooks | Where `InitFiniPlugin` runs | | `DefinitionGenerator` | Fallback symbol provider | `DLLImportDefinitionGenerator` (Win) | | Link order | Search path across JITDylibs for symbol resolution | `SetLinkOrder()` | | `__tvm_ffi_` prefix | Namespace for TVM-FFI exported functions | Used in `GetFunction()` | --- ## 7. Further Reading - LLVM ORC JIT documentation: - JITLink design: - ELF specification: - PE/COFF specification: - Mach-O reference: - Ian Lance Taylor's linker series (20-part blog): foundational reading on linkers tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/README.md000066400000000000000000000204521521067262500207070ustar00rootroot00000000000000 # TVM-FFI OrcJIT A Python package that enables dynamic loading of compiled object files (`.o`) using LLVM ORC JIT v2, providing a flexible JIT execution environment for TVM-FFI exported functions. ## Features - **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2 - **Multiple Libraries**: Create separate dynamic libraries with independent symbol namespaces - **Incremental Loading**: Add multiple object files to the same library incrementally - **Symbol Isolation**: Different libraries can define the same symbol without conflicts - **Init/Fini Support**: Handles static constructors/destructors across ELF (`.init_array`/`.ctors`), Mach-O (`__mod_init_func`), and COFF (`.CRT$XC*`/`.CRT$XT*`) - **Cross-Platform**: Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64) - **Multi-Compiler**: Tested with LLVM Clang, GCC, Apple Clang, MSVC, and clang-cl - **TVM-FFI Integration**: Seamlessly works with TVM-FFI's stable C ABI - **Python API**: Simple Pythonic interface for JIT compilation and execution ## Supported Platforms and Compilers Object files compiled with any of the following compiler/platform combinations can be loaded and executed by the ORC JIT: | Platform | Compilers | C | C++ | | -------- | --------- | :-: | :-: | | Linux (x86_64, aarch64) | LLVM Clang, GCC | yes | yes | | macOS (arm64) | LLVM Clang, Apple Clang | yes | yes | | Windows (AMD64) | LLVM Clang, MSVC, clang-cl | yes | no | Windows is C-only across all compilers. C++ objects compiled with `TVM_FFI_DLL_EXPORT_TYPED_FUNC` use `try`/`catch` (via `TVM_FFI_SAFE_CALL_BEGIN/END`), which requires Itanium exception ABI symbols (`__cxa_begin_catch`, `__gxx_personality_v0`, etc.) that the MSVC-built host process cannot provide. Pure C objects using the `TVMFFISafeCallType` ABI work on all platforms. ## Installation ### Install from PyPI ```bash pip install apache-tvm-ffi apache-tvm_ffi_orcjit ``` ### Build from Source #### Prerequisites - Python 3.10+, CMake 3.20+, C++17 compiler - LLVM 22+ development libraries (`llvmdev`, `llvm-config`) - Static `zlib` and `zstd` libraries (in the same prefix as LLVM) #### Install LLVM via conda-forge The easiest way to get all dependencies is via conda-forge: ```bash conda create -p /opt/llvm -c conda-forge \ llvmdev=22.1.0 clangdev=22.1.0 compiler-rt=22.1.0 zlib zstd-static -y export LLVM_PREFIX=/opt/llvm ``` On Windows: ```cmd conda create -p C:\opt\llvm -c conda-forge llvmdev=22.1.0 zlib zstd-static -y set LLVM_PREFIX=C:\opt\llvm ``` #### Build and install ```bash git clone --recursive https://github.com/apache/tvm-ffi.git cd tvm-ffi # Install tvm-ffi first pip install -e . # Build and install the orcjit addon cd addons/tvm_ffi_orcjit pip install -e . ``` The `LLVM_PREFIX` environment variable tells CMake where to find LLVM. If LLVM is installed in a conda env or a standard system path, CMake can auto-discover it and `LLVM_PREFIX` is not needed. ## Usage ### Basic Example ```python from tvm_ffi_orcjit import ExecutionSession # Create an execution session session = ExecutionSession() # Create a dynamic library lib = session.create_library() # Load an object file lib.add("example.o") # Get and call a function add_func = lib.get_function("add") result = add_func(1, 2) print(f"Result: {result}") # Output: Result: 3 ``` ### Multiple Libraries with Symbol Isolation ```python session = ExecutionSession() lib1 = session.create_library("lib1") lib2 = session.create_library("lib2") lib1.add("implementation_v1.o") lib2.add("implementation_v2.o") add_v1 = lib1.get_function("add") add_v2 = lib2.get_function("add") print(add_v1(5, 3)) # Uses implementation from lib1 print(add_v2(5, 3)) # Uses implementation from lib2 ``` ### Cross-Library Linking ```python session = ExecutionSession() base_lib = session.create_library("base") base_lib.add("math_ops.o") caller_lib = session.create_library("caller") caller_lib.set_link_order(base_lib) # Can resolve symbols from base_lib caller_lib.add("caller.o") result = caller_lib.get_function("call_math")(10, 20) ``` ## Writing Functions for OrcJIT ### C++ (Linux/macOS) ```cpp #include TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, [](int a, int b) { return a + b; }); ``` Compile: `clang++ -std=c++17 -fPIC -O2 -c -o example.o example.cc` ### Pure C (all platforms including Windows) ```c #include TVM_FFI_DLL_EXPORT int __tvm_ffi_add( void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->v_int64 = args[0].v_int64 + args[1].v_int64; return 0; } ``` Compile: `clang -O2 -c -o example.o example.c` ## How It Works - **LLJIT**: Built on LLVM's ORC JIT v2 with `ObjectLinkingLayer` (JITLink) for all platforms. - **InitFiniPlugin**: Custom `ObjectLinkingLayer::Plugin` that collects function pointers from init/fini sections (ELF `.init_array`/`.ctors`/`.fini_array`/`.dtors`, Mach-O `__mod_init_func`/`__mod_term_func`, COFF `.CRT$XC*`/`.CRT$XT*`) and runs them in priority order at symbol lookup / library teardown. - **DLL Import Stubs** (Windows): Custom `DefinitionGenerator` that resolves host process symbols from all loaded DLLs and creates `__imp_*` pointer stubs in JIT memory, keeping all fixups within PCRel32 range. - **SEH Stripping** (Windows): `ObjectTransformLayer` strips `.pdata`/`.xdata` relocations from COFF objects before JITLink graph building, working around a JITLink limitation with COMDAT section symbols. Please refers to [ORCJIT_PRIMER.md](./ORCJIT_PRIMER.md) to learn more about object file, linking, llvm orcjit v2, and how the addon works. ## Project Structure ```text tvm_ffi_orcjit/ ├── CMakeLists.txt # Build configuration ├── pyproject.toml # Python package metadata ├── src/ffi/ │ ├── orcjit_session.cc # ExecutionSession (LLJIT setup, plugins) │ ├── orcjit_session.h │ ├── orcjit_dylib.cc # DynamicLibrary (object loading, symbol lookup) │ ├── orcjit_dylib.h │ └── orcjit_utils.h # LLVM error handling utilities ├── python/tvm_ffi_orcjit/ │ ├── __init__.py # Module exports and library loading │ ├── session.py # Python ExecutionSession wrapper │ └── dylib.py # Python DynamicLibrary wrapper ├── tests/ # See tests/README.md └── examples/quick-start/ # Complete example with CMake ``` ## CI Runs on Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64) via `cibuildwheel`. Each platform builds test objects with multiple compilers and runs the full test suite. See `.github/workflows/tvm_ffi_orcjit.yml`. ## Troubleshooting ### "Cannot find global function" error The shared library wasn't loaded. Reinstall: `pip install --force-reinstall apache-tvm_ffi_orcjit` ### "Duplicate definition of symbol" error Use separate libraries for different implementations of the same symbol. ### "Symbol not found" error Ensure functions are exported with TVM-FFI macros (`TVM_FFI_DLL_EXPORT_TYPED_FUNC` for C++, or `__tvm_ffi_` prefix for C). ### Relocation errors on Windows MSVC/clang-cl objects must be compiled with `/GS-` to disable buffer security checks (`__security_cookie`) which are CRT symbols the JIT cannot resolve. ### LLVM version mismatch The package requires LLVM 22+. Set `LLVM_PREFIX` to the LLVM install prefix: ```bash export LLVM_PREFIX=/path/to/llvm ``` ## License Apache License 2.0 tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/000077500000000000000000000000001521067262500212435ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/000077500000000000000000000000001521067262500235125ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/CMakeLists.txt000066400000000000000000000043441521067262500262570ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.18) project(tvm_ffi_orcjit_example) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Helper: build a source file into an object and copy as .o function (add_example_object name source) get_filename_component(_ext "${source}" EXT) add_library(${name}_obj OBJECT "${source}") target_link_libraries(${name}_obj PRIVATE tvm_ffi::header) if (_ext STREQUAL ".c" AND MSVC) target_compile_options(${name}_obj PRIVATE /O2) else () target_compile_options(${name}_obj PRIVATE -O2) if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64") target_compile_options(${name}_obj PRIVATE -mno-outline-atomics) endif () endif () add_custom_target( copy_${name} ALL COMMAND ${CMAKE_COMMAND} -E copy $ ${CMAKE_CURRENT_SOURCE_DIR}/${name}.o DEPENDS ${name}_obj COMMENT "Copying ${name}.o to example directory" ) endfunction () # C++ object — skip on Windows (C-only strategy for ORC JIT on Windows) if (NOT WIN32) add_example_object(add add.cc) endif () # Pure C object — built on all platforms enable_language(C) add_example_object(add_c add_c.c) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/README.md000066400000000000000000000061571521067262500250020ustar00rootroot00000000000000 # Quick Start Example Demonstrates basic usage of tvm_ffi_orcjit: compile functions to object files, load them into the ORC JIT at runtime, and call them from Python. ## Files | File | Description | | ------ | ------------- | | `add.cc` | C++ source using `TVM_FFI_DLL_EXPORT_TYPED_FUNC` (automatic type marshaling, supports `std::string`) | | `add_c.c` | Pure C source using `TVMFFISafeCallType` ABI (`__tvm_ffi_` prefix). No C++ runtime needed. | | `run.py` | Python script that compiles, loads, and calls the functions | | `CMakeLists.txt` | Alternative CMake build configuration | ## Prerequisites - Python 3.10+, C/C++ compiler - `apache-tvm-ffi` and `apache-tvm_ffi_orcjit` packages installed ## Run `run.py` compiles the source files to object files using `tvm_ffi.cpp.build`, then loads them into the ORC JIT: ```bash # C++ variant (Linux/macOS) python run.py # Pure C variant (all platforms including Windows) python run.py --lang c ``` ## How It Works **C++ variant** (`add.cc`): Functions are exported with `TVM_FFI_DLL_EXPORT_TYPED_FUNC`, which wraps a typed C++ lambda/function into TVM-FFI's packed calling convention. Supports C++ types like `std::string`. **C variant** (`add_c.c`): Functions follow the `TVMFFISafeCallType` ABI directly — each function is named `__tvm_ffi_` and manually packs arguments/results via `TVMFFIAny`. Zero C++ dependencies, works on all platforms including Windows with MSVC or clang-cl. **Python side** (`run.py`): ```python from tvm_ffi_orcjit import ExecutionSession session = ExecutionSession() # Create ORC JIT session lib = session.create_library() # Create a JITDylib lib.add("add_c.o") # Load object file into JIT add = lib.get_function("add") # Look up symbol print(add(10, 20)) # Call like a normal function → 30 ``` ## Platform Notes | Platform | C++ (`add.o`) | C (`add_c.o`) | | ---------- | :-: | :-: | | Linux (Clang/GCC) | yes | yes | | macOS (Clang/Apple Clang) | yes | yes | | Windows (all compilers) | no | yes | C++ is not supported for ORC JIT on Windows. The `TVM_FFI_DLL_EXPORT_TYPED_FUNC` macro uses `try`/`catch` which requires Itanium exception ABI symbols that the MSVC-built host process cannot provide. Use the pure C variant on Windows. tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/add.cc000066400000000000000000000031301521067262500245460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Quick Start Example - Simple Math Functions * * This file demonstrates how to export C++ functions using TVM-FFI * so they can be loaded dynamically at runtime with tvm_ffi_orcjit. */ #include // Simple addition function int add_impl(int a, int b) { return a + b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); // Multiplication function int multiply_impl(int a, int b) { return a * b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(multiply, multiply_impl); // Fibonacci function (recursive) int fib_impl(int n) { if (n <= 1) return n; return fib_impl(n - 1) + fib_impl(n - 2); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(fibonacci, fib_impl); // String concatenation example std::string concat_impl(std::string a, std::string b) { return a + b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(concat, concat_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/add_c.c000066400000000000000000000041431521067262500247120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Quick Start Example - Simple Math Functions (Pure C) * * This file demonstrates how to export C functions using TVM-FFI's * C ABI (TVMFFISafeCallType) so they can be loaded dynamically at * runtime with tvm_ffi_orcjit. No C++ runtime dependencies. */ #include /* add: add two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64; return 0; } /* multiply: multiply two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_multiply(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 * args[1].v_int64; return 0; } /* fibonacci: recursive fibonacci */ static int64_t fib(int64_t n) { if (n <= 1) return n; return fib(n - 1) + fib(n - 2); } TVM_FFI_DLL_EXPORT int __tvm_ffi_fibonacci(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = fib(args[0].v_int64); return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/examples/quick-start/run.py000077500000000000000000000104351521067262500246760ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Quick Start Example - Compile, load, and call functions via ORC JIT. This script demonstrates how to: 1. Compile C/C++ source files to object files using tvm_ffi.cpp.build 2. Load them into an ORC JIT ExecutionSession 3. Get functions by name 4. Call them like regular Python functions Usage: python run.py # Load C++ object file (add.o) python run.py --lang c # Load pure C object file (add_c.o) """ from __future__ import annotations import argparse import platform import sys from pathlib import Path import tvm_ffi.cpp from tvm_ffi_orcjit import ExecutionSession SCRIPT_DIR = Path(__file__).resolve().parent def _extra_cflags() -> list[str]: """Return extra compiler flags for the current platform.""" machine = platform.machine() if machine in ("aarch64", "arm64"): return ["-mno-outline-atomics"] return [] def _build_object(lang: str) -> str: """Compile the source file to a relocatable object file.""" if lang == "c": return tvm_ffi.cpp.build( name="add_c", sources=[str(SCRIPT_DIR / "add_c.c")], output="add_c.o", extra_cflags=_extra_cflags(), ) else: return tvm_ffi.cpp.build( name="add", sources=[str(SCRIPT_DIR / "add.cc")], output="add.o", extra_cflags=_extra_cflags(), ) def _run_tests(obj_file: str, lang: str) -> None: """Load object file and run all test assertions. All JIT references (functions, lib, session) are released automatically when this function returns. """ print(f"Loading object file: {obj_file} (lang={lang})") # Create execution session and dynamic library session = ExecutionSession() lib = session.create_library() lib.add(obj_file) print("Object file loaded successfully\n") # Get and call the 'add' function print("=== Testing add function ===") add = lib.get_function("add") result = add(10, 20) print(f"add(10, 20) = {result}") assert result == 30, f"Expected 30, got {result}" # Get and call the 'multiply' function print("\n=== Testing multiply function ===") multiply = lib.get_function("multiply") result = multiply(7, 6) print(f"multiply(7, 6) = {result}") assert result == 42, f"Expected 42, got {result}" # Get and call the 'fibonacci' function print("\n=== Testing fibonacci function ===") fibonacci = lib.get_function("fibonacci") result = fibonacci(10) print(f"fibonacci(10) = {result}") assert result == 55, f"Expected 55, got {result}" if lang == "cpp": # String concatenation only available in C++ variant (uses std::string) print("\n=== Testing concat function ===") concat = lib.get_function("concat") result = concat("Hello, ", "World!") print(f"concat('Hello, ', 'World!') = '{result}'") assert result == "Hello, World!", f"Expected 'Hello, World!', got '{result}'" print("\n" + "=" * 50) print("All tests passed successfully!") print("=" * 50) def main() -> int: """Run the quick start example.""" parser = argparse.ArgumentParser(description="Quick Start Example") parser.add_argument( "--lang", choices=["cpp", "c"], default="cpp", help="Language variant to load: 'cpp' for add.o (default), 'c' for add_c.o", ) args = parser.parse_args() obj_file = _build_object(args.lang) _run_tests(obj_file, args.lang) return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/pyproject.toml000066400000000000000000000055161521067262500223500ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [build-system] requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"] build-backend = "scikit_build_core.build" [project] name = "apache-tvm_ffi_orcjit" version = "0.1.0" description = "Load TVM-FFI exported object files using LLVM ORC JIT v2" readme = "README.md" requires-python = ">=3.10" license = { text = "Apache-2.0" } authors = [{ name = "TVM-FFI OrcJIT Contributors" }] keywords = ["tvm-ffi", "llvm", "jit", "orcjit"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: C++", ] dependencies = ["apache-tvm-ffi>=0.1.0"] [project.urls] Homepage = "https://github.com/apache/tvm-ffi" Repository = "https://github.com/apache/tvm-ffi" [tool.scikit-build] cmake.build-type = "Release" wheel.py-api = "py3" build-dir = "build" build.verbose = true editable.rebuild = false editable.verbose = true wheel.packages = ["python/tvm_ffi_orcjit"] wheel.install-dir = "tvm_ffi_orcjit" sdist.include = [ "/README.md", "/LICENSE", "/pyproject.toml", "/CMakeLists.txt", "/src/**/*.h", "/src/**/*.cc", "/src/**/*.cpp", "/python/**/*.py", ] [tool.scikit-build.cmake.define] CMAKE_EXPORT_COMPILE_COMMANDS = "ON" [tool.cibuildwheel] # One build per platform is sufficient: the shared library does not link against # Python C API, so wheel.py-api = "py3" produces a version-agnostic wheel. build = "cp312-*" skip = "*-win32 *-manylinux_i686 *-musllinux*" [tool.cibuildwheel.linux] # LLVM_PREFIX is unnecessary when LLVM is installed in a conda env # (cmake auto-discovers via CONDA_PREFIX) or in a standard system path. repair-wheel-command = "auditwheel repair --exclude libtvm_ffi.so -w {dest_dir} {wheel}" [tool.cibuildwheel.macos] repair-wheel-command = "delocate-wheel --ignore-missing-dependencies --exclude libtvm_ffi.dylib -w {dest_dir} -v {wheel}" tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/000077500000000000000000000000001521067262500207465ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/000077500000000000000000000000001521067262500237525ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py000066400000000000000000000057761521067262500261020ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """TVM-FFI OrcJIT. This module provides functionality to load object files (.o) compiled with TVM-FFI exports using LLVM ORC JIT v2. Example: >>> from tvm_ffi_orcjit import ExecutionSession >>> session = ExecutionSession() >>> lib = session.create_library() >>> lib.add("example.o") >>> func = lib.get_function("my_function") >>> result = func(arg1, arg2) """ import ctypes import os import platform import sys from pathlib import Path from tvm_ffi import load_module # Determine the library name based on platform if platform.system() == "Windows": _LIB_NAME = "tvm_ffi_orcjit.dll" elif platform.system() == "Darwin": _LIB_NAME = "libtvm_ffi_orcjit.dylib" else: _LIB_NAME = "libtvm_ffi_orcjit.so" # Load the orcjit extension library # - lib/: normal install (wheel) # - ../../build/: editable install (cmake build output relative to python/tvm_ffi_orcjit/) _LIB_PATH = [ Path(__file__).parent / "lib" / _LIB_NAME, Path(__file__).parent.parent.parent / "build" / _LIB_NAME, ] _lib_dir = None for path in _LIB_PATH: if path.exists(): _ = load_module(str(path)) _lib_dir = path.parent if _lib_dir is None: raise RuntimeError( f"Could not find {_LIB_NAME}. " f"Searched in {_LIB_PATH} and site-packages. " f"Please ensure the package is installed correctly." ) # Explicitly initialize the library to register functions # This is needed because static initializers may not run when loaded via dlopen try: # The dll search path need to be added explicitly in windows if sys.platform.startswith("win32"): os.add_dll_directory(str(_lib_dir)) # Load the library with ctypes and call the initialization function c_lib = ctypes.CDLL(str(_lib_dir / _LIB_NAME), mode=ctypes.RTLD_GLOBAL) init_func = c_lib.TVMFFIOrcJITInitialize init_func.restype = None init_func() except Exception as e: import warnings warnings.warn(f"Failed to explicitly initialize orcjit library: {e}") from .dylib import DynamicLibrary from .session import ExecutionSession __all__ = ["DynamicLibrary", "ExecutionSession"] try: from importlib.metadata import version __version__ = version("apache-tvm_ffi_orcjit") except Exception: __version__ = "0.0.0.dev0" tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/_ffi_api.py000066400000000000000000000015461521067262500260660ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """FFI APIs for orcjit.""" import tvm_ffi tvm_ffi.init_ffi_api("orcjit", __name__) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/dylib.py000066400000000000000000000054541521067262500254370ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """ORC JIT Dynamic Library.""" from __future__ import annotations from pathlib import Path from tvm_ffi import Module class DynamicLibrary(Module): """ORC JIT Dynamic Library (JITDylib). Represents a collection of symbols that can be loaded from object files and linked against other dynamic libraries. Supports JIT compilation and symbol resolution. Examples -------- >>> session = ExecutionSession() >>> lib = session.create_library() >>> lib.add("add.o") >>> lib.add("multiply.o") >>> add_func = lib.get_function("add") >>> result = add_func(1, 2) """ def add(self, object_file: str | Path) -> None: """Add an object file to this dynamic library. Parameters ---------- object_file : str or Path Path to the object file to load. Examples -------- >>> lib.add("add.o") >>> lib.add(Path("multiply.o")) """ if isinstance(object_file, Path): object_file = str(object_file) self.get_function("orcjit.add_object_file")(object_file) def set_link_order(self, *libraries: DynamicLibrary) -> None: """Set the link order for symbol resolution. When resolving symbols, this library will search in the specified libraries in the order provided. This replaces any previous link order. Parameters ---------- *libraries : DynamicLibrary One or more dynamic libraries to search for symbols (in order). Examples -------- >>> session = ExecutionSession() >>> lib_utils = session.create_library() >>> lib_utils.add("utils.o") >>> lib_core = session.create_library() >>> lib_core.add("core.o") >>> lib_main = session.create_library() >>> lib_main.add("main.o") >>> # main can call symbols from utils and core (utils searched first) >>> lib_main.set_link_order(lib_utils, lib_core) """ self.get_function("orcjit.set_link_order")(list(libraries)) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py000066400000000000000000000131701521067262500260110ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """ORC JIT Execution Session.""" from __future__ import annotations import sys from tvm_ffi import Object, register_object from . import _ffi_api, _lib_dir from .dylib import DynamicLibrary def _find_orc_rt_library() -> str | None: """Find the bundled liborc_rt library in the same directory as the .so/.dll.""" # Windows: skip ORC runtime entirely. LLVM's COFFPlatform (loaded via # ExecutorNativePlatform with liborc_rt) depends on MSVC C++ runtime symbols # that are not available in the JIT environment. On Windows, ORC JIT uses a # C-only strategy: JIT objects are compiled as pure C (TVMFFISafeCallType ABI), # avoiding all C++ runtime dependencies (magic statics, RTTI, sized delete, # SEH, COMDAT). Our custom InitFiniPlugin handles .CRT$XC*/.CRT$XT* init/fini # sections, and DLLImportDefinitionGenerator resolves __imp_ DLL import stubs. # # macOS: skip ORC runtime too. ExecutorNativePlatform would install # MachOPlatform, which triggers a compact-unwind 32-bit-delta bug in # JITLink when a user graph mmaps below the per-JITDylib Mach-O header # (see repo-root fix-machoplatform-libunwind-dso-base.patch). Our # InitFiniPlugin handles __mod_init_func / __mod_term_func instead. # Tradeoff: no C++ exception unwinding across JIT frames on macOS. if sys.platform in ("win32", "darwin"): return None patterns = ["liborc_rt*.a"] for pattern in patterns: for lib_path in _lib_dir.glob(pattern): return str(lib_path) return None @register_object("orcjit.ExecutionSession") class ExecutionSession(Object): """ORC JIT Execution Session. Manages the LLVM ORC JIT execution environment and creates dynamic libraries (JITDylibs). This is the top-level context for JIT compilation and symbol management. Examples -------- >>> session = ExecutionSession() >>> lib = session.create_library(name="main") >>> lib.add("add.o") >>> add_func = lib.get_function("add") """ def __init__(self, orc_rt_path: str | None = None, slab_size: int = 0) -> None: """Initialize ExecutionSession. Args: orc_rt_path: Optional path to the liborc_rt library. If not provided, it will be automatically discovered using clang. slab_size: Per-slab capacity in bytes for the JIT memory manager. Linux only — ignored on macOS and Windows, where the slab allocator is compiled out. 0 = arch default (64 MB; initial slab halves on mmap failure down to 8 MB under RLIMIT_AS / container limits), >0 = custom size, <0 = disable slab allocator (LLJIT uses its default scattered-mmap allocator). The session holds a growable pool of slabs: a fresh slab is mmap'd on demand when no existing one can fit a graph. Graphs that don't fit a normal slab trigger a power-of-2 larger slab (slab_size, 2*slab_size, ...) sized to fit. Drained slabs stay mapped until the session is destroyed or ``clear_free_slabs()`` is called. """ if orc_rt_path is None: orc_rt_path = _find_orc_rt_library() if orc_rt_path is None: orc_rt_path = "" self.__init_handle_by_constructor__(_ffi_api.ExecutionSession, orc_rt_path, slab_size) # type: ignore def create_library(self, name: str = "") -> DynamicLibrary: """Create a new dynamic library associated with this execution session. Args: name: Optional name for the library. If empty, a unique name will be generated. Returns: A new DynamicLibrary instance. """ handle = _ffi_api.ExecutionSessionCreateDynamicLibrary(self, name) # type: ignore lib = DynamicLibrary.__new__(DynamicLibrary) lib.__move_handle_from__(handle) return lib def clear_free_slabs(self) -> int: """Release drained slabs (no live JIT allocations) back to the OS. Call this after dropping a batch of libraries to reclaim RSS. Fresh slabs that have never been allocated on are preserved, so the session remains ready to accept new work. Safety: call when no JIT work is in flight on another thread. From single-threaded Python this is always safe; once ``del lib`` has returned, the C++ destructor has finished and the slab's live count reflects the drop. Returns: Number of slabs actually munmap'd. Returns 0 on macOS/Windows (slab pool compiled out) or when the pool is disabled via ``slab_size=-1``. """ return int(_ffi_api.ExecutionSessionClearFreeSlabs(self)) # type: ignore tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/000077500000000000000000000000001521067262500202145ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/000077500000000000000000000000001521067262500207605ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/000077500000000000000000000000001521067262500234415ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md000066400000000000000000000055331521067262500247260ustar00rootroot00000000000000 # LLVM patches This directory holds workarounds for upstream LLVM bugs and missing features. **A file belongs here if and only if its entire reason to exist is an LLVM defect, and we would delete the whole file once upstream catches up.** Features that happen to coexist with a bug (e.g. the arena memory manager, which we keep for THP and contiguous layout regardless of LLVM state) live at the top level of `src/ffi/`, not here. Each patch file opens with a fixed-shape header describing: - which LLVM issue it addresses (link or "not yet filed"), - affected version range, - exact trigger conditions, - symptom without the patch, and - a `## Removal` section listing the `#include` and plugin-registration line(s) to delete when the upstream fix lands and the project's minimum LLVM version bumps past it. ## Index - **GOTPCRELX relaxation** (`gotpcrelx_fix.{h,cc}`) LLVM issue: TBD — issue not yet filed. Upstream status: open. Remove when: LLVM floor bumps past the release that contains the fix. - **ELF init/fini** (Linux branch of `init_fini_plugin.{h,cc}`) LLVM issue: [llvm/llvm-project#175981](https://github.com/llvm/llvm-project/issues/175981). Upstream status: open, patch submitted. Remove when: LLVM floor bumps past the release that contains the fix. - **COFF ctor/dtor** (Windows branch of `init_fini_plugin.{h,cc}`) LLVM issue: COFFPlatform stalled. Upstream status: stalled 2+ years. Remove when: COFFPlatform becomes usable end-to-end with clang-cl / MSVC objects. macOS already has working `MachOPlatform`, so no patch file is needed for that platform. ## Removal checklist When deleting a patch file: 1. Delete the `.h` and `.cc` pair. 2. Remove the matching `#include "llvm_patches/.h"` in `orcjit_session.cc`. 3. Remove the plugin-registration line(s) in `orcjit_session.cc` identified in the file's `## Removal` header block. 4. Remove the corresponding sources from `addons/tvm_ffi_orcjit/CMakeLists.txt`. 5. Update the index above. tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.cc000066400000000000000000000134511521067262500266310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file gotpcrelx_fix.cc * \brief LLVM JITLink GOTPCRELX relaxation bug workaround (x86_64). * * See gotpcrelx_fix.h for the trigger, symptom, and removal procedure. */ #include "gotpcrelx_fix.h" #if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64)) #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { namespace { /*! \brief Correct broken GOTPCRELX relaxations produced by * optimizeGOTAndStubAccesses(). * * Strategy: * 1. Build target-symbol → GOT-entry-symbol map (O(B+S) up front). * 2. For every Pointer32 edge whose preceding bytes are 67 e8 * (relaxed call) or e9 (relaxed jmp): * - If the target is reachable via a signed 32-bit PC-relative * displacement, change the edge to BranchPCRel32. * - Otherwise revert the relaxation: restore the original * indirect-call/jmp opcode bytes (ff 15 / ff 25), retarget * the edge to the GOT entry, and use PCRel32 with addend 0 * (JITLink normalises GOTPCRELX addends to 0). */ llvm::Error fixBrokenGOTPCRELXRelaxation(llvm::jitlink::LinkGraph& G) { using namespace llvm::jitlink; // Build block → first symbol at offset 0 (for GOT entry symbol lookup). llvm::DenseMap BlockToSym; for (auto* Sym : G.defined_symbols()) { if (Sym->getOffset() == 0 && !BlockToSym.count(&Sym->getBlock())) { BlockToSym[&Sym->getBlock()] = Sym; } } // Build target symbol → GOT entry symbol map. // GOT entries are pointer-sized blocks with exactly one Pointer64 edge. llvm::DenseMap SymToGOTSym; for (auto* B : G.blocks()) { if (B->getSize() != G.getPointerSize()) continue; if (B->edges_size() != 1) continue; auto& E = *B->edges().begin(); if (E.getKind() == x86_64::Pointer64) { auto It = BlockToSym.find(B); if (It != BlockToSym.end()) { SymToGOTSym[&E.getTarget()] = It->second; } } } for (auto* B : G.blocks()) { for (auto& E : B->edges()) { if (E.getKind() != x86_64::Pointer32) continue; if (E.getOffset() < 2) continue; auto MutableContent = B->getMutableContent(G); auto* FixupData = reinterpret_cast(MutableContent.data()) + E.getOffset(); uint8_t Prev2 = FixupData[-2]; uint8_t Prev1 = FixupData[-1]; bool isRelaxedCall = (Prev2 == 0x67 && Prev1 == 0xe8); bool isRelaxedJmp = (Prev1 == 0xe9); if (!isRelaxedCall && !isRelaxedJmp) continue; // Check if PC-relative displacement would fit. auto TargetAddr = E.getTarget().getAddress(); auto FixupAddr = B->getFixupAddress(E); int64_t Displacement = TargetAddr.getValue() - (FixupAddr.getValue() + 4) + E.getAddend(); if (llvm::isInt<32>(Displacement)) { E.setKind(x86_64::BranchPCRel32); continue; } // Distance doesn't fit — revert to indirect call/jmp through GOT. auto It = SymToGOTSym.find(&E.getTarget()); if (It == SymToGOTSym.end()) { return llvm::make_error( "Cannot revert GOTPCRELX relaxation: no GOT entry for " + (E.getTarget().hasName() ? std::string(*E.getTarget().getName()) : std::string("")), llvm::inconvertibleErrorCode()); } Symbol* GOTSym = It->second; if (isRelaxedCall) { // Restore: 67 e8 → ff 15 (call *[rip+disp32]) FixupData[-2] = 0xff; FixupData[-1] = 0x15; } else { // Restore: e9 XX XX XX XX 90 → ff 25 XX XX XX XX FixupData[-1] = 0xff; FixupData[0] = 0x25; // For jmp, the optimization shifted offset by -1; shift back. E.setOffset(E.getOffset() + 1); } E.setKind(x86_64::PCRel32); E.setTarget(*GOTSym); E.setAddend(0); } } return llvm::Error::success(); } } // namespace void GOTPCRELXFixPlugin::modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G, llvm::jitlink::PassConfiguration& Config) { Config.PreFixupPasses.emplace_back(fixBrokenGOTPCRELXRelaxation); } llvm::Error GOTPCRELXFixPlugin::notifyFailed(llvm::orc::MaterializationResponsibility& MR) { return llvm::Error::success(); } llvm::Error GOTPCRELXFixPlugin::notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) { return llvm::Error::success(); } void GOTPCRELXFixPlugin::notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey, llvm::orc::ResourceKey SrcKey) {} } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __linux__ && __x86_64__ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.h000066400000000000000000000076531521067262500265020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file gotpcrelx_fix.h * \brief LLVM JITLink GOTPCRELX relaxation bug workaround (x86_64). * * LLVM issue: not yet filed, internal TODO. * Affected versions: observed on LLVM 20.x and 21.x; likely all versions * since `optimizeGOTAndStubAccesses` landed in * JITLink/x86_64.cpp. * Trigger: x86_64 JITLink with external symbols whose resolved addresses * fit in uint32 (e.g. libc PLT entries in a non-PIE executable, * or any low-VA process symbol) while JIT code is at high * addresses (as produced by our arena memory manager). * Symptom: `optimizeGOTAndStubAccesses` relaxes * `call *foo@GOTPCREL(%rip)` (ff 15) * into * `addr32 call foo` (67 e8) * and sets the edge kind to `Pointer32` (absolute 32-bit). But * `call rel32` is always PC-relative, so the absolute fixup * produces a garbage displacement. Result: SIGSEGV at JIT * execution or during ORC-runtime teardown. * * The `GOTPCRELXFixPlugin` registers a PreFixupPass that runs *after* * `optimizeGOTAndStubAccesses`, detects `Pointer32` edges preceded by * `67 e8` / `e9` bytes, and either converts them to `BranchPCRel32` (if * the PC-relative displacement fits in int32) or reverts the relaxation * to an indirect call/jmp through the GOT (`ff 15` / `ff 25`, edge kind * `PCRel32`, addend 0). * * ## Removal * * When the upstream fix lands and the project's minimum LLVM version * bumps past the first release containing it, delete this file and * remove: * - the `#include "llvm_patches/gotpcrelx_fix.h"` in orcjit_session.cc * - the `OLL->addPlugin(std::make_unique())` call * inside the `setObjectLinkingLayerCreator` lambda in orcjit_session.cc. * - the `llvm_patches/gotpcrelx_fix.cc` entry in * addons/tvm_ffi_orcjit/CMakeLists.txt. */ #ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_ #define TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_ #if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64)) #include namespace tvm { namespace ffi { namespace orcjit { /*! \brief PreFixupPass plugin that corrects broken GOTPCRELX * relaxations produced by JITLink's `optimizeGOTAndStubAccesses`. * * See the file-level docstring above for the trigger, symptom, and * removal procedure. */ class GOTPCRELXFixPlugin : public llvm::orc::ObjectLinkingLayer::Plugin { public: void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G, llvm::jitlink::PassConfiguration& Config) override; llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override; llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override; void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey, llvm::orc::ResourceKey SrcKey) override; }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __linux__ && __x86_64__ #endif // TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.cc000066400000000000000000000240731521067262500273040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file init_fini_plugin.cc * \brief Init/fini section handling for ELF, MachO, and COFF JIT objects. * * See init_fini_plugin.h for the trigger, symptom, and removal * procedure for each platform. */ #include "init_fini_plugin.h" #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { void InitFiniPlugin::modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G, llvm::jitlink::PassConfiguration& Config) { auto& jit_dylib = MR.getTargetJITDylib(); // Mark all init/fini section blocks and their edge targets as live // so they survive dead-stripping. Config.PrePrunePasses.emplace_back([](llvm::jitlink::LinkGraph& G) { for (auto& Section : G.sections()) { auto section_name = Section.getName(); // ELF: .init_array*, .fini_array*, .ctors*, .dtors* // Mach-O: __DATA,__mod_init_func, __DATA,__mod_term_func // COFF: .CRT$XC* (ctors), .CRT$XT* (dtors) if (section_name.starts_with(".init_array") || section_name.starts_with(".fini_array") || section_name.starts_with(".ctors") || section_name.starts_with(".dtors") || section_name == "__DATA,__mod_init_func" || section_name == "__DATA,__mod_term_func" || section_name.starts_with(".CRT$XC") || section_name.starts_with(".CRT$XT")) { for (auto* Block : Section.blocks()) { bool has_live_sym = false; for (auto* Sym : G.defined_symbols()) { if (&Sym->getBlock() == Block) { Sym->setLive(true); has_live_sym = true; } } // MSVC may emit .CRT$XC* blocks with data but no symbol table // entries (static variables in __declspec(allocate) sections). // Add an anonymous symbol so the block survives dead-stripping. if (!has_live_sym) { G.addAnonymousSymbol(*Block, 0, Block->getSize(), false, true).setLive(true); } for (auto& Edge : Block->edges()) { Edge.getTarget().setLive(true); } } } } return llvm::Error::success(); }); #ifdef _WIN32 // Without COFFPlatform, __ImageBase (used by IMAGE_REL_AMD64_ADDR32NB / // Pointer32NB relocations) defaults to 0. This causes all Pointer32 fixups // to overflow since JIT addresses don't fit in 32 bits. // // Fix: set __ImageBase to the lowest block address in the graph after // allocation. This makes all intra-graph Pointer32NB offsets small. // // Also strip .pdata/.xdata edges: SEH unwind data references external // handlers (e.g., __CxxFrameHandler3) in DLLs that may be >4GB from // __ImageBase. Since we don't call RtlAddFunctionTable, SEH data is // unused anyway. Config.PostAllocationPasses.emplace_back([](llvm::jitlink::LinkGraph& G) { // Set __ImageBase to the lowest allocated block address. auto ImageBaseName = G.intern("__ImageBase"); llvm::jitlink::Symbol* ImageBase = nullptr; for (auto* Sym : G.external_symbols()) { if (Sym->getName() == ImageBaseName) { ImageBase = Sym; break; } } if (ImageBase) { llvm::orc::ExecutorAddr BaseAddr; for (auto* B : G.blocks()) { if (!BaseAddr || B->getAddress() < BaseAddr) { BaseAddr = B->getAddress(); } } ImageBase->getAddressable().setAddress(BaseAddr); } // Strip .pdata/.xdata edges: external handlers may be >4GB from // __ImageBase, and we don't register SEH data anyway. for (auto& Sec : G.sections()) { if (Sec.getName() == ".pdata" || Sec.getName().starts_with(".xdata")) { for (auto* B : Sec.blocks()) { while (!B->edges_empty()) { B->removeEdge(B->edges().begin()); } } } } return llvm::Error::success(); }); #endif // After fixups, read resolved function pointers from all init/fini data sections. // Handles ELF (.init_array, .ctors, .fini_array, .dtors), // Mach-O (__DATA,__mod_init_func, __DATA,__mod_term_func), // and COFF (.CRT$XC*, .CRT$XT*) section conventions. Config.PostFixupPasses.emplace_back([this, &jit_dylib](llvm::jitlink::LinkGraph& G) { using Entry = ORCJITExecutionSessionObj::InitFiniEntry; for (auto& Sec : G.sections()) { auto section_name = Sec.getName(); // --- ELF sections --- bool is_init_array = section_name.starts_with(".init_array"); bool is_ctors = section_name.starts_with(".ctors"); bool is_fini_array = section_name.starts_with(".fini_array"); bool is_dtors = section_name.starts_with(".dtors"); // --- Mach-O sections --- bool is_mod_init = (section_name == "__DATA,__mod_init_func"); bool is_mod_term = (section_name == "__DATA,__mod_term_func"); // --- COFF sections --- bool is_crt_xc = section_name.starts_with(".CRT$XC"); bool is_crt_xt = section_name.starts_with(".CRT$XT"); if (!is_init_array && !is_ctors && !is_fini_array && !is_dtors && !is_mod_init && !is_mod_term && !is_crt_xc && !is_crt_xt) continue; int priority = 0; Entry::Section sec; bool is_init; // ELF default priority for sections without a numeric suffix is 65535. // Lower priority numbers run first for .init_array; .fini_array and .ctors // negate so that higher-numbered entries run first (reverse order). if (is_init_array) { if (section_name.consume_front(".init_array.")) { section_name.getAsInteger(10, priority); } else { priority = 65535; } sec = Entry::Section::kInitArray; is_init = true; } else if (is_ctors) { if (section_name.consume_front(".ctors.") && !section_name.getAsInteger(10, priority)) { priority = -priority; } sec = Entry::Section::kCtors; is_init = true; } else if (is_fini_array) { if (section_name.consume_front(".fini_array.") && !section_name.getAsInteger(10, priority)) { priority = -priority; } else { priority = -65535; } sec = Entry::Section::kFiniArray; is_init = false; } else if (is_dtors) { if (section_name.consume_front(".dtors.")) { section_name.getAsInteger(10, priority); } sec = Entry::Section::kDtors; is_init = false; } else if (is_mod_init) { // Mach-O __mod_init_func: no priority system, treated as init_array sec = Entry::Section::kInitArray; is_init = true; } else if (is_mod_term) { // Mach-O __mod_term_func: no priority system, treated as fini_array sec = Entry::Section::kFiniArray; is_init = false; } else if (is_crt_xc) { // COFF .CRT$XC[suffix]: C++ constructors, sorted alphabetically by suffix. // Convert suffix to integer priority that preserves alphabetical ordering. // E.g., .CRT$XCA → 'A'*100000=6500000, .CRT$XCU → 'U'*100000=8500000, // .CRT$XCT00200 → 'T'*100000+200=8400200 sec = Entry::Section::kInitArray; is_init = true; auto suffix = section_name.substr(7); // after ".CRT$XC" if (!suffix.empty()) { priority = static_cast(suffix[0]) * 100000; if (suffix.size() > 1) { int num = 0; suffix.substr(1).getAsInteger(10, num); priority += num; } } } else { // COFF .CRT$XT[suffix]: C++ destructors, same suffix-to-priority scheme. sec = Entry::Section::kFiniArray; is_init = false; auto suffix = section_name.substr(7); // after ".CRT$XT" if (!suffix.empty()) { priority = static_cast(suffix[0]) * 100000; if (suffix.size() > 1) { int num = 0; suffix.substr(1).getAsInteger(10, num); priority += num; } } } for (auto* Block : Sec.blocks()) { auto Content = Block->getContent(); size_t PtrSize = G.getPointerSize(); for (size_t Offset = 0; Offset + PtrSize <= Content.size(); Offset += PtrSize) { uint64_t FnAddr = 0; memcpy(&FnAddr, Content.data() + Offset, PtrSize); if (FnAddr != 0) { Entry entry{llvm::orc::ExecutorAddr(FnAddr), sec, priority}; if (is_init) { session_->AddPendingInitializer(&jit_dylib, entry); } else { session_->AddPendingDeinitializer(&jit_dylib, entry); } } } } } return llvm::Error::success(); }); } llvm::Error InitFiniPlugin::notifyFailed(llvm::orc::MaterializationResponsibility& MR) { return llvm::Error::success(); } llvm::Error InitFiniPlugin::notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) { return llvm::Error::success(); } void InitFiniPlugin::notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey, llvm::orc::ResourceKey SrcKey) {} } // namespace orcjit } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h000066400000000000000000000130661521067262500271460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file init_fini_plugin.h * \brief Init/fini section handling for ELF, MachO, and COFF JIT objects. * * Emulates the missing/broken/bypassed LLVM ORC platform support for * init/fini sections on all three host platforms. Collects function * pointers from `.init_array` / `.fini_array` / `.ctors` / `.dtors` * (ELF), `__DATA,__mod_init_func` / `__DATA,__mod_term_func` (MachO), * and `.CRT$XC*` / `.CRT$XT*` (COFF); ties them to the containing * `JITDylib`; and runs them in priority order through * `ORCJITExecutionSessionObj::Run{Pending{Init,De}initializers}`. * * On Windows the plugin additionally patches `__ImageBase` (set to the * lowest block address so `IMAGE_REL_AMD64_ADDR32NB` fixups don't * overflow) and strips `.pdata` / `.xdata` edges (SEH handlers live in * DLLs > 4 GB from `__ImageBase` and are never registered with * `RtlAddFunctionTable` anyway). Those pieces also disappear once * `COFFPlatform` becomes usable. * * Trigger: any JIT module on any platform containing constructors, * destructors, or `__attribute__((constructor))` / * MSVC `#pragma init_seg` equivalents. * Symptom without the patch: * - ELF: constructors/destructors never run (`ELFNixPlatform` * enumerates but does not invoke them before * llvm/llvm-project#175981). * - MachO: we skip `MachOPlatform` entirely to sidestep the * compact-unwind 32-bit delta bug (see `orcjit_session.cc`), so no * platform runs init/fini; this plugin is the only mechanism. * - COFF: relocation overflow / unresolved-SEH crashes * (`COFFPlatform` is not hooked up because its MSVC CRT symbol * requirements cannot be satisfied). * * ## Removal — Linux * * LLVM issue: https://github.com/llvm/llvm-project/issues/175981 * When the upstream fix lands and the project's minimum LLVM version * bumps past the first release containing it, replace this plugin's * Linux usage with `ELFNixPlatform` and delete the ELF handling path * from this file. Concretely: * - Remove the ELF-section branches (`.init_array`, `.ctors`, * `.fini_array`, `.dtors`) from `InitFiniPlugin::modifyPassConfig`. * - If no platform still needs this plugin, delete this file outright * and follow the checklist in `llvm_patches/README.md`. * * ## Removal — macOS * * Tied to re-enabling `MachOPlatform`. That requires the compact-unwind * per-graph `dso_base` fix (see `fix-machoplatform-libunwind-dso-base.patch` * in the repo root) to land in our LLVM. Until then we skip MachOPlatform * and this plugin handles `__mod_init_func` / `__mod_term_func`. When * MachOPlatform is re-enabled, delete the MachO-section branches from * `InitFiniPlugin::modifyPassConfig` and drop the macOS side of the * `addPlugin` call in `orcjit_session.cc`. * * ## Removal — Windows * * LLVM status: `COFFPlatform` has been stalled for 2+ years because it * requires MSVC CRT symbols (`_CxxThrowException`, RTTI vtables, ...) * that LLVM's COFF ORC runtime cannot provide. * When `COFFPlatform` becomes usable end-to-end with clang-cl / MSVC * objects, replace this plugin's Windows usage with it and delete the * COFF handling path (including the `__ImageBase` fixup and the * `.pdata` / `.xdata` edge stripping). */ #ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_ #define TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_ #include #include "../orcjit_session.h" namespace tvm { namespace ffi { namespace orcjit { /*! \brief Init/fini section collector and runner for ELF, MachO, and COFF. * * See the file-level docstring above for the three-platform strategy * and the removal procedure for each platform. */ class InitFiniPlugin : public llvm::orc::ObjectLinkingLayer::Plugin { // Store a raw pointer to avoid a reference cycle: // Session → LLJIT → ObjectLinkingLayer → Plugin → Session // The plugin's lifetime is bounded by the ObjectLinkingLayer which is // owned by LLJIT which is owned by the session, so the pointer is always valid. ORCJITExecutionSessionObj* session_; public: explicit InitFiniPlugin(ORCJITExecutionSessionObj* session) : session_(session) {} void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G, llvm::jitlink::PassConfiguration& Config) override; llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override; llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override; void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey, llvm::orc::ResourceKey SrcKey) override; }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc000066400000000000000000000064131521067262500302740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file macho_cxa_atexit_shim.cc * \brief Per-JITDylib `__cxa_atexit` interposer implementation. * * See macho_cxa_atexit_shim.h for the trigger, symptom, and removal * procedure. */ #include "macho_cxa_atexit_shim.h" #ifdef __APPLE__ #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { namespace { // TLS pointer to the currently-active dylib's records vector. // Each CxaAtexitRecordsScope saves the previous pointer and restores it // on exit, so nested scopes compose correctly across re-entrant init/fini. // // We cannot override ___dso_handle (LLJIT's Platform has already defined // it in every user JITDylib), so we don't rely on the `dso_handle` arg // passed to the shim. Instead, each ORCJITDynamicLibraryObj publishes its // own records vector via this TLS slot, scoped around any JIT entry point // that may run ctors / dtors. The shim pushes (fn, arg) into the // TLS-pointed vector, or silently drops if no scope is active (which // would be a stray call from outside any JIT execution — acceptable // degradation). thread_local CxaAtexitRecords* g_active_cxa_records = nullptr; extern "C" int tvm_ffi_cxa_atexit_shim(void (*fn)(void*), void* arg, void* /*dso_handle*/) noexcept { if (fn == nullptr || g_active_cxa_records == nullptr) return 0; g_active_cxa_records->emplace_back(fn, arg); return 0; } } // namespace CxaAtexitRecordsScope::CxaAtexitRecordsScope(CxaAtexitRecords* records) : prev_(g_active_cxa_records) { g_active_cxa_records = records; } CxaAtexitRecordsScope::~CxaAtexitRecordsScope() { g_active_cxa_records = prev_; } void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd) { llvm::orc::SymbolMap shim_syms; shim_syms[ES.intern("___cxa_atexit")] = { llvm::orc::ExecutorAddr::fromPtr(reinterpret_cast(&tvm_ffi_cxa_atexit_shim)), llvm::JITSymbolFlags::Exported | llvm::JITSymbolFlags::Callable}; llvm::cantFail(jd.define(llvm::orc::absoluteSymbols(std::move(shim_syms)))); } void DrainCxaAtexit(CxaAtexitRecords& records) { CxaAtexitRecordsScope scope(&records); while (!records.empty()) { auto [fn, arg] = records.back(); records.pop_back(); fn(arg); } } } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __APPLE__ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h000066400000000000000000000112741521067262500301370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file macho_cxa_atexit_shim.h * \brief Per-JITDylib `__cxa_atexit` interposer for macOS JIT. * * We skip `MachOPlatform` entirely on macOS to sidestep the * compact-unwind 32-bit-delta bug in JITLink's `CompactUnwindSupport` * (see the analysis in orcjit_session.cc and * fix-machoplatform-libunwind-dso-base.patch at the repo root). With * no Platform in the picture, clang-lowered * `__attribute__((destructor))` and C++ global dtors — which register * through `__cxa_atexit(fn, arg, &__dso_handle)` during init — would * fall through to libSystem's `___cxa_atexit`, orphaning those * callbacks from our drop-time drain. * * This shim: * 1. Installs an absolute symbol for `___cxa_atexit` on each user * JITDylib that points at our own capture function * (`InstallCxaAtexitShim`). * 2. Publishes the owning dylib's `CxaAtexitRecords` vector in TLS * (`CxaAtexitRecordsScope`) so the capture function knows where * to push `(fn, arg)` pairs. * 3. Drains the captured records LIFO at dylib destruction time * (`DrainCxaAtexit`). * * Trigger: any macOS JIT object containing static destructors, * `__attribute__((destructor))` functions, or C++ global * objects with non-trivial dtors. * Symptom without the shim: destructors never run; any resource * held by a JIT global leaks for the lifetime of the host * process. * * ## Removal * * Tied to re-enabling `MachOPlatform`. See the macOS removal notes * in init_fini_plugin.h. When `MachOPlatform` is restored, delete * this file, drop the `#include` in orcjit_session.cc / orcjit_dylib.h, * and remove the `cxa_atexit_records_` field from `ORCJITDynamicLibraryObj`. */ #ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_ #define TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_ #ifdef __APPLE__ #include #include #include namespace tvm { namespace ffi { namespace orcjit { /*! \brief `(dtor, arg)` pairs captured from per-dylib `__cxa_atexit` calls. */ using CxaAtexitRecords = std::vector>; /*! \brief RAII scope that publishes a per-dylib `CxaAtexitRecords` vector * to the TLS slot consulted by the `___cxa_atexit` shim. * * The shim reads the TLS pointer and pushes `(fn, arg)` into the pointed-to * vector; outside any active scope the shim silently drops registrations. * Wrap any JIT entry point that can run ctors / dtors (e.g. `GetSymbol`'s * initialize call, the dtor's drain loop) with one of these, constructed * with `&cxa_atexit_records_` from the owning `ORCJITDynamicLibraryObj`. */ class CxaAtexitRecordsScope { public: explicit CxaAtexitRecordsScope(CxaAtexitRecords* records); ~CxaAtexitRecordsScope(); CxaAtexitRecordsScope(const CxaAtexitRecordsScope&) = delete; CxaAtexitRecordsScope& operator=(const CxaAtexitRecordsScope&) = delete; private: CxaAtexitRecords* prev_; }; /*! \brief Install the `___cxa_atexit` absolute-symbol shim on \p jd. * * Must be called once per user JITDylib, before any JIT code on that dylib * materializes. Placing the definition on the dylib itself (rather than * injecting into the link order) ensures it wins over ``'s * libSystem fallback — JITDylib::define-time symbols are searched before * the link order. */ void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd); /*! \brief Drain captured `(fn, arg)` records LIFO, running each dtor. * * Pop-and-call order handles re-entrant registrations from within a dtor — * the internal `CxaAtexitRecordsScope` keeps the TLS pointer live so any * `___cxa_atexit` call from inside a dtor also lands in \p records. */ void DrainCxaAtexit(CxaAtexitRecords& records); } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __APPLE__ #endif // TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.cc000066400000000000000000000104631521067262500301400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file win_coff_pdata_strip.cc * \brief COFF `.pdata` / `.xdata` relocation stripper implementation. * * See win_coff_pdata_strip.h for the trigger, symptom, and removal * procedure. */ #include "win_coff_pdata_strip.h" #ifdef _WIN32 #include #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { llvm::Expected> StripCoffPdataXdata( std::unique_ptr Buf) { const char* Data = Buf->getBufferStart(); size_t Size = Buf->getBufferSize(); if (Size < 20) return std::move(Buf); // Parse COFF header (regular or bigobj format) uint16_t w0, w1; std::memcpy(&w0, Data, 2); std::memcpy(&w1, Data + 2, 2); bool bigobj = (w0 == 0 && w1 == 0xFFFF); uint16_t machine; uint32_t num_sections, ptr_to_symtab, num_symbols; size_t sec_hdr_start, sym_entry_size; if (bigobj) { if (Size < 56) return std::move(Buf); std::memcpy(&machine, Data + 6, 2); std::memcpy(&num_sections, Data + 44, 4); std::memcpy(&ptr_to_symtab, Data + 48, 4); std::memcpy(&num_symbols, Data + 52, 4); sec_hdr_start = 56; sym_entry_size = 20; } else { machine = w0; uint16_t ns, opt_hdr_size; std::memcpy(&ns, Data + 2, 2); std::memcpy(&opt_hdr_size, Data + 16, 2); std::memcpy(&ptr_to_symtab, Data + 8, 4); std::memcpy(&num_symbols, Data + 12, 4); num_sections = ns; sec_hdr_start = 20 + opt_hdr_size; sym_entry_size = 18; } if (machine != 0x8664) return std::move(Buf); // String table follows the symbol table size_t strtab_start = ptr_to_symtab + static_cast(num_symbols) * sym_entry_size; // Resolve a section name (inline 8-byte or "/offset" string table ref) constexpr size_t kSecHdrSize = 40; auto resolve_name = [&](size_t hdr_off) -> llvm::StringRef { const char* raw = Data + hdr_off; if (raw[0] == '/' && raw[1] >= '0' && raw[1] <= '9') { uint32_t offset = 0; for (int j = 1; j < 8 && raw[j] >= '0' && raw[j] <= '9'; ++j) offset = offset * 10 + (raw[j] - '0'); size_t pos = strtab_start + offset; if (pos < Size) { size_t len = 0; while (pos + len < Size && Data[pos + len]) ++len; return {Data + pos, len}; } } size_t len = 0; while (len < 8 && raw[len]) ++len; return {raw, len}; }; // Collect section header offsets needing relocation stripping llvm::SmallVector strip_offsets; for (uint32_t i = 0; i < num_sections; ++i) { size_t off = sec_hdr_start + i * kSecHdrSize; if (off + kSecHdrSize > Size) break; auto name = resolve_name(off); if (name.starts_with(".pdata") || name.starts_with(".xdata")) { uint16_t num_relocs; std::memcpy(&num_relocs, Data + off + 32, 2); if (num_relocs > 0) strip_offsets.push_back(off); } } if (strip_offsets.empty()) return std::move(Buf); // Create mutable copy, zero out PointerToRelocations and NumberOfRelocations llvm::SmallVector MutableBuf(Data, Data + Size); for (auto off : strip_offsets) { std::memset(&MutableBuf[off + 24], 0, 4); // PointerToRelocations std::memset(&MutableBuf[off + 32], 0, 2); // NumberOfRelocations } return llvm::MemoryBuffer::getMemBufferCopy(llvm::StringRef(MutableBuf.data(), MutableBuf.size()), Buf->getBufferIdentifier()); } } // namespace orcjit } // namespace ffi } // namespace tvm #endif // _WIN32 tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.h000066400000000000000000000061201521067262500277750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file win_coff_pdata_strip.h * \brief COFF `.pdata` / `.xdata` relocation stripper for Windows JIT. * * Installed as an `ObjTransformLayer` transform on Windows. clang-cl * places static functions in COMDAT sections, and the `.pdata` SEH * unwind data has relocations targeting COMDAT leader symbols. * JITLink's `COFFLinkGraphBuilder` doesn't register COMDAT leaders in * its symbol table when the second COMDAT symbol is `CLASS_STATIC` * (not `CLASS_EXTERNAL`), causing "Could not find symbol" errors at * graph-build time. * * `init_fini_plugin.cc` also strips `.pdata` / `.xdata` edges in a * `PostAllocationPass`; this transform moves the stripping earlier * (pre-graph-build) to prevent the graph builder error. Both pieces * disappear once `COFFPlatform` becomes usable. * * Trigger: Windows x86_64 JIT of clang-cl objects that contain * static-linkage functions (i.e. almost any C++ input). * Symptom without the transform: JITLink graph-build failure with * "Could not find symbol" referencing a COMDAT leader. * * We avoid `llvm/Object/COFF.h` because `windows.h` (included * transitively by `LLJIT.h`) defines `IMAGE_*` macros that conflict * with LLVM's COFF enums; we parse the COFF header with raw * `memcpy`s instead. * * ## Removal * * Tied to enabling `COFFPlatform`. See the Windows removal notes in * init_fini_plugin.h. */ #ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_ #define TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_ #ifdef _WIN32 #include #include #include namespace tvm { namespace ffi { namespace orcjit { /*! \brief Zero out `PointerToRelocations` / `NumberOfRelocations` on * `.pdata` and `.xdata` section headers before the COFF object * reaches JITLink's graph builder. * * Returns the (possibly copy-modified) buffer. Safe to install as the * `ObjTransformLayer` transform — invokes no LLVM object parsers and * short-circuits for any non-x86_64 or too-small input. */ llvm::Expected> StripCoffPdataXdata( std::unique_ptr Buf); } // namespace orcjit } // namespace ffi } // namespace tvm #endif // _WIN32 #endif // TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc000066400000000000000000000136501521067262500310450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file win_dll_import_generator.cc * \brief Windows DLL-import symbol generator implementation. * * See win_dll_import_generator.h for the trigger, symptom, and removal * procedure. */ #include "win_dll_import_generator.h" #ifdef _WIN32 #ifndef NOMINMAX #define NOMINMAX #endif #ifndef WIN32_LEAN_AND_MEAN #define WIN32_LEAN_AND_MEAN #endif #include #include #include #include #include #include #include // windows.h must precede psapi.h — psapi.h uses SIZE_T / DWORD typedefs // defined in windows.h. Left to its own devices clang-format sorts these // alphabetically, landing psapi.h first and breaking the MSVC build. // clang-format off #include #include // clang-format on #include namespace tvm { namespace ffi { namespace orcjit { void* DLLImportDefinitionGenerator::FindInProcessModules(const std::string& Name) { // Try specific runtime DLLs first, then tvm_ffi.dll (loaded by Python), // then all process modules, then LLVM's search. static const char* kRuntimeDLLs[] = { "vcruntime140.dll", "vcruntime140_1.dll", "ucrtbase.dll", "msvcp140.dll", }; // NOTE: We intentionally do not call FreeLibrary() here. These runtime DLLs // (vcruntime140, ucrtbase, etc.) are already loaded by the process and will // remain loaded for its lifetime. LoadLibraryA merely increments the refcount; // the extra refcount is harmless and avoids the overhead of balancing // Get/FreeLibrary for every symbol lookup. for (const char* dll : kRuntimeDLLs) { if (HMODULE hMod = LoadLibraryA(dll)) { if (auto addr = GetProcAddress(hMod, Name.c_str())) { return reinterpret_cast(addr); } } } // Also check tvm_ffi.dll (host process symbol provider) if (HMODULE hTvmFfi = GetModuleHandleA("tvm_ffi.dll")) { if (auto addr = GetProcAddress(hTvmFfi, Name.c_str())) { return reinterpret_cast(addr); } } HMODULE hMods[1024]; DWORD cbNeeded; if (EnumProcessModules(GetCurrentProcess(), hMods, sizeof(hMods), &cbNeeded)) { DWORD count = cbNeeded / sizeof(HMODULE); if (count > 1024) count = 1024; for (DWORD i = 0; i < count; ++i) { if (auto addr = GetProcAddress(hMods[i], Name.c_str())) { return reinterpret_cast(addr); } } } if (void* addr = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(Name)) { return addr; } return nullptr; } llvm::Error DLLImportDefinitionGenerator::tryToGenerate( llvm::orc::LookupState& LS, llvm::orc::LookupKind K, llvm::orc::JITDylib& JD, llvm::orc::JITDylibLookupFlags JDLookupFlags, const llvm::orc::SymbolLookupSet& LookupSet) { // Step 1: Collect unique base names (strip __imp_ prefix) and resolve addresses. llvm::DenseMap Resolved; for (auto& [Name, Flags] : LookupSet) { llvm::StringRef NameStr = *Name; std::string BaseName = NameStr.starts_with("__imp_") ? NameStr.drop_front(6).str() : NameStr.str(); if (BaseName == "__ImageBase") continue; auto InternedBase = ES_.intern(BaseName); if (Resolved.count(InternedBase)) continue; void* Addr = FindInProcessModules(BaseName); if (Addr) { Resolved[InternedBase] = llvm::orc::ExecutorAddr::fromPtr(Addr); } } if (Resolved.empty()) return llvm::Error::success(); // Step 2: Build a LinkGraph with __imp_ pointers and PLT jump stubs. auto G = std::make_unique( "", ES_.getSymbolStringPool(), ES_.getTargetTriple(), llvm::SubtargetFeatures(), llvm::jitlink::getGenericEdgeKindName); auto Prot = static_cast(static_cast(llvm::orc::MemProt::Read) | static_cast(llvm::orc::MemProt::Exec)); auto& Sec = G->createSection("__dll_stubs", Prot); for (auto& [InternedName, Addr] : Resolved) { // Absolute symbol at the real address (local to this graph) auto& Target = G->addAbsoluteSymbol(G->intern(("__real_" + *InternedName).str()), Addr, G->getPointerSize(), llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Local, false); // __imp_XXX pointer (GOT-like entry) auto& Ptr = llvm::jitlink::x86_64::createAnonymousPointer(*G, Sec, &Target); Ptr.setName(G->intern(("__imp_" + *InternedName).str())); Ptr.setLinkage(llvm::jitlink::Linkage::Strong); Ptr.setScope(llvm::jitlink::Scope::Default); // XXX jump stub (PLT-like entry) for direct calls auto& StubBlock = llvm::jitlink::x86_64::createPointerJumpStubBlock(*G, Sec, Ptr); G->addDefinedSymbol(StubBlock, 0, *InternedName, StubBlock.getSize(), llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Default, true, false); } return L_.add(JD, std::move(G)); } } // namespace orcjit } // namespace ffi } // namespace tvm #endif // _WIN32 tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h000066400000000000000000000071621521067262500307100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file win_dll_import_generator.h * \brief Windows DLL-import symbol generator for JIT code (COFF/x86_64). * * On Windows with the MSVC ABI, COFF objects reference DLL-imported * symbols via `__imp_XXX` pointer stubs and direct calls. Without * `COFFPlatform` (which we skip because its MSVC CRT dependencies * cannot be satisfied — see init_fini_plugin.h for the Windows * removal notes), JITLink has no way to resolve those references. * * `DLLImportDefinitionGenerator` is a `DefinitionGenerator` that, for * each requested symbol: * 1. Looks up the real address by walking the MSVC runtime DLLs, * `tvm_ffi.dll`, all loaded process modules, and finally LLVM's * `SearchForAddressOfSymbol`. * 2. Emits a fresh `LinkGraph` containing a GOT-like `__imp_XXX` * pointer holding the real address, plus a PLT-like jump stub * (`jmp [__imp_XXX]`) so direct calls stay in range. All stubs * land in JIT memory, keeping every PCRel32 fixup within ±2 GB * of the JIT arena. * * Trigger: any Windows x86_64 JIT graph that references a DLL-imported * symbol (which, for MSVC-compiled objects, is effectively * any C runtime or libtvm_ffi entry point). * Symptom without the generator: unresolved `__imp_XXX` externals * at JITLink time, or PCRel32 overflow for direct DLL calls. * * ## Removal * * Tied to enabling `COFFPlatform`. When `COFFPlatform` becomes usable * end-to-end with clang-cl / MSVC objects (blocked upstream on MSVC * CRT symbol requirements), delete this file and the corresponding * `addGenerator` call in `orcjit_session.cc`. */ #ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_ #define TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_ #ifdef _WIN32 #include #include namespace tvm { namespace ffi { namespace orcjit { /*! \brief JIT-allocated `__imp_XXX` pointer stubs + PLT jumps for DLL imports. * * See the file-level docstring for the trigger, symptom, and removal * procedure. */ class DLLImportDefinitionGenerator : public llvm::orc::DefinitionGenerator { public: DLLImportDefinitionGenerator(llvm::orc::ExecutionSession& ES, llvm::orc::ObjectLinkingLayer& L) : ES_(ES), L_(L) {} llvm::Error tryToGenerate(llvm::orc::LookupState& LS, llvm::orc::LookupKind K, llvm::orc::JITDylib& JD, llvm::orc::JITDylibLookupFlags JDLookupFlags, const llvm::orc::SymbolLookupSet& LookupSet) override; private: static void* FindInProcessModules(const std::string& Name); llvm::orc::ExecutionSession& ES_; llvm::orc::ObjectLinkingLayer& L_; }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // _WIN32 #endif // TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc000066400000000000000000000223521521067262500237500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_dylib.cc * \brief LLVM ORC JIT DynamicLibrary implementation */ #include "orcjit_dylib.h" #include #include #include #include #include #include #include #include #include #include #include #include #include #include "orcjit_session.h" #include "orcjit_utils.h" namespace tvm { namespace ffi { namespace orcjit { namespace { // When JIT thunks start reading their ctx/handle arg, grow this wrapper with // a DylibFnContext prefix whose first field is the thunk pointer, and flip // `safe_call` at a slab-emitted redirect that dispatches via ctx[0]. Example: // // struct DylibFnContext { TVMFFISafeCallType fn; /* + closure fields */ }; // struct DylibFnContextWithModule { // DylibFnContext ctx; // first — pointer-interconvertible with wrapper // Module module_ref; // }; struct DylibFnContextWithModule { Module module_ref; // keeps the owning dylib (and its slab) alive }; void DeleteDylibFnContextWithModule(void* p) { delete static_cast(p); } } // namespace ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session, llvm::orc::JITDylib* dylib, llvm::orc::LLJIT* jit, String name) : session_(std::move(session)), dylib_(dylib), jit_(jit), name_(std::move(name)) { TVM_FFI_CHECK(dylib_ != nullptr, ValueError) << "JITDylib cannot be null"; TVM_FFI_CHECK(jit_ != nullptr, ValueError) << "LLJIT cannot be null"; if (void** ctx_addr = reinterpret_cast(GetSymbol(ffi::symbol::tvm_ffi_library_ctx))) { *ctx_addr = this; } Module::VisitContextSymbols([this](const ffi::String& name, void* symbol) { if (void** ctx_addr = reinterpret_cast(GetSymbol(name))) { *ctx_addr = symbol; } }); } ORCJITDynamicLibraryObj::~ORCJITDynamicLibraryObj() { // Step 1: run static destructors for code in this JITDylib via our // InitFiniPlugin (uniform across Linux / macOS / Windows — see the // plugin docstring in llvm_patches/init_fini_plugin.h). session_->RunPendingDeinitializers(GetJITDylib()); #ifdef __APPLE__ // Step 1b (macOS only): drain per-dylib __cxa_atexit registrations in // LIFO. Clang on Darwin lowers __attribute__((destructor)) and C++ // global dtors as __cxa_atexit(fn, arg, &__dso_handle) registrations // during init; the shim in llvm_patches/macho_cxa_atexit_shim.cc // captured them into cxa_atexit_records_ via the TLS pointer // published by the GetSymbol-time scope. DrainCxaAtexit(cxa_atexit_records_); #endif // Step 2: remove the JITDylib from the ExecutionSession. Triggers // JITDylib::clear(), which releases all tracked linker resources — in // particular, invokes the ObjectLinkingLayer's ResourceManager which calls // our memory manager's deallocate() for every FinalizedAlloc belonging to // this dylib. Without this call JIT code pages accumulate in the arena // until the whole session is destroyed. No Platform is registered in any // of our LLJIT configurations (see orcjit_session.cc), so there is no // Platform::teardownJITDylib callback to worry about. session_->RemoveDylib(dylib_); dylib_ = nullptr; } void ORCJITDynamicLibraryObj::AddObjectFile(const String& path) { // Read object file auto buffer_or_err = llvm::MemoryBuffer::getFile(path.c_str()); if (!buffer_or_err) { TVM_FFI_THROW(IOError) << "Failed to read object file: " << path; } // Add object file to this JITDylib TVM_FFI_ORCJIT_LLVM_CALL(jit_->addObjectFile(*dylib_, std::move(*buffer_or_err))); } void ORCJITDynamicLibraryObj::SetLinkOrder(const std::vector& dylibs) { // Rebuild the link order: user-specified libraries first, then the LLJIT // default link order (Main → Platform → ProcessSymbols). Preserving the // default link order is essential — without ProcessSymbols, C++ objects // that need host-process symbols (runtime, libtvm_ffi) would fail to link. link_order_.clear(); for (auto* lib : dylibs) { link_order_.emplace_back(lib, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); } for (auto& kv : jit_->defaultLinkOrder()) { link_order_.emplace_back(kv.first, kv.second); } // Set the link order in the LLVM JITDylib dylib_->setLinkOrder(link_order_, false); } void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { // Build search order: this dylib first, then all linked dylibs llvm::orc::JITDylibSearchOrder search_order; search_order.emplace_back(dylib_, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); // Append linked libraries search_order.insert(search_order.end(), link_order_.begin(), link_order_.end()); // Look up symbol using the full search order auto symbol_or_err = jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str())); // Run pending initializers via InitFiniPlugin. RunPendingInitializers // drains and erases the map entry; subsequent calls are no-ops until new // object files add fresh entries — supports incremental loading. // // On macOS the scope publishes this dylib's __cxa_atexit records vector // so the ___cxa_atexit shim (see orcjit_session.cc) can route dtor // registrations here for our destructor to drain. Static init on C // and C++ code typically happens here (first GetSymbol resolves the // `main` entry point, which materializes and fires __mod_init_func). #ifdef __APPLE__ CxaAtexitRecordsScope scope(&cxa_atexit_records_); #endif session_->RunPendingInitializers(GetJITDylib()); // Convert ExecutorAddr to pointer return symbol_or_err ? symbol_or_err->getAddress().toPtr() : nullptr; } llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { TVM_FFI_CHECK(dylib_ != nullptr, InternalError) << "JITDylib is null"; return *dylib_; } Optional ORCJITDynamicLibraryObj::GetFunction(const String& name) { if (name == "orcjit.add_object_file") { return Function::FromTyped([this](const String& path) { AddObjectFile(path); }); } if (name == "orcjit.set_link_order") { return Function::FromTyped([this](const Array& libraries) { std::vector libs; libs.reserve(libraries.size()); for (const ORCJITDynamicLibrary& lib : libraries) { libs.push_back(&lib->GetJITDylib()); } SetLinkOrder(libs); }); } // TVM-FFI exports have __tvm_ffi_ prefix std::string symbol_name = symbol::tvm_ffi_symbol_prefix + std::string(name); // Try to get the symbol - return NullOpt if not found if (void* symbol = GetSymbol(symbol_name)) { TVMFFISafeCallType c_func = reinterpret_cast(symbol); auto* wrapper = new DylibFnContextWithModule{GetRef(this)}; return Function::FromExternC(wrapper, c_func, DeleteDylibFnContextWithModule); } return std::nullopt; } //------------------------------------- // Registration //------------------------------------- static void RegisterOrcJITFunctions() { static bool registered = false; if (registered) return; registered = true; namespace refl = tvm::ffi::reflection; refl::ObjectDef(); refl::GlobalDef() .def("orcjit.ExecutionSession", [](const std::string& orc_rt_path, int64_t slab_size_bytes) { return ORCJITExecutionSession(orc_rt_path, slab_size_bytes); }) .def("orcjit.ExecutionSessionCreateDynamicLibrary", [](const ORCJITExecutionSession& session, const String& name) -> Module { return session->CreateDynamicLibrary(name); }) .def("orcjit.ExecutionSessionClearFreeSlabs", [](const ORCJITExecutionSession& session) -> int64_t { return session->ClearFreeSlabs(); }); } TVM_FFI_STATIC_INIT_BLOCK() { // This block may not execute when loaded via dlopen on some platforms. // Call TVMFFIOrcJITInitialize() explicitly if functions are not registered. RegisterOrcJITFunctions(); } } // namespace orcjit } // namespace ffi } // namespace tvm // C API for explicit initialization extern "C" { TVM_FFI_DLL_EXPORT void TVMFFIOrcJITInitialize() { tvm::ffi::orcjit::RegisterOrcJITFunctions(); } } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h000066400000000000000000000105321521067262500236070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_dylib.h * \brief LLVM ORC JIT DynamicLibrary (JITDylib) wrapper */ #ifndef TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ #define TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ #include #include #include #include #include #include "llvm_patches/macho_cxa_atexit_shim.h" #include "orcjit_session.h" namespace tvm { namespace ffi { namespace orcjit { class ORCJITExecutionSession; class ORCJITDynamicLibraryObj : public ModuleObj { public: /*! * \brief Constructor * \param session The parent execution session * \param dylib The LLVM JITDylib * \param jit The LLJIT instance * \param name The library name */ ORCJITDynamicLibraryObj(ORCJITExecutionSession session, llvm::orc::JITDylib* dylib, llvm::orc::LLJIT* jit, String name); ~ORCJITDynamicLibraryObj(); const char* kind() const final { return "orcjit"; } Optional GetFunction(const String& name) override; private: /*! * \brief Add an object file to this library * \param path Path to the object file to load */ void AddObjectFile(const String& path); /*! * \brief Set the link order for symbol resolution * \param dylibs Vector of libraries to search for symbols (in order) * * When resolving symbols, this library will search in the specified libraries * in the order provided. This replaces any previous link order. */ void SetLinkOrder(const std::vector& dylibs); /*! * \brief Look up a symbol in this library * \param name The symbol name to look up * \return Pointer to the symbol, or nullptr if not found */ void* GetSymbol(const String& name); /*! * \brief Get the underlying LLVM JITDylib * \return Reference to the LLVM JITDylib */ llvm::orc::JITDylib& GetJITDylib(); /*! * \brief Get the name of this library * \return The library name */ String GetName() const { return name_; } /*! \brief Parent execution session (for lifetime management) */ ORCJITExecutionSession session_; /*! \brief The LLVM JITDylib */ llvm::orc::JITDylib* dylib_; /*! \brief The LLJIT instance (for addObjectFile API) */ llvm::orc::LLJIT* jit_; /*! \brief Library name */ String name_; /*! \brief Link order tracking (to support incremental linking) */ llvm::orc::JITDylibSearchOrder link_order_; #ifdef __APPLE__ /*! \brief Per-dylib __cxa_atexit registry. * * Without MachOPlatform, clang-lowered \c __attribute__((destructor)) * and C++ global dtors register through \c __cxa_atexit during init. * We interpose \c ___cxa_atexit per-JITDylib (see * \c orcjit_session.cc); the shim pushes \c (fn, arg) into the vector * published via \c CxaAtexitRecordsScope around each JIT entry. The * destructor drains LIFO before \c RemoveDylib. */ CxaAtexitRecords cxa_atexit_records_; #endif // __APPLE__ }; /*! * \brief DynamicLibrary wrapper for LLVM ORC JIT v2 JITDylib * * This class wraps an LLVM JITDylib and provides functionality to: * - Load object files * - Link against other dynamic libraries * - Look up symbols */ class ORCJITDynamicLibrary : public Module { public: explicit ORCJITDynamicLibrary(const ObjectPtr& ptr) : Module(ptr) {}; TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ORCJITDynamicLibrary, Module, ORCJITDynamicLibraryObj); }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc000066400000000000000000000151651521067262500256530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_memory_manager.cc * \brief Growable per-session pool of `Slab`s. */ #include "orcjit_memory_manager.h" #ifdef __linux__ #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { using llvm::Error; using llvm::Expected; SlabPoolMemoryManager::SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size) : page_size_(page_size), slab_size_(slab_size) { // Reserve the initial slab. Halving retry only applies here: if the // very first mmap fails (RLIMIT_AS, container limits), we halve the // requested size down to kMinSlabSize before giving up. Subsequent // slabs added during allocate() use exactly slab_size_ and propagate // errors on mmap failure. std::size_t floor = std::min(slab_size_, kMinSlabSize); std::size_t cap = slab_size_; while (cap >= floor) { auto slab = std::make_unique(page_size_, cap); if (slab->isValid()) { // Pin the actual initial-slab size to whatever we succeeded with. // If RLIMIT_AS forced us to 8 MB, we keep 8 MB as the working slab // size; growing later at 64 MB would just fail again. slab_size_ = cap; slabs_.push_back(std::move(slab)); return; } cap /= 2; } llvm::report_fatal_error("SlabPoolMemoryManager: failed to reserve at least " + llvm::Twine(floor / (1024 * 1024)) + " MB of virtual address space"); } std::unique_ptr SlabPoolMemoryManager::createSlab(std::size_t capacity) { auto slab = std::make_unique(page_size_, capacity); if (!slab->isValid()) return nullptr; return slab; } void SlabPoolMemoryManager::allocate(const llvm::jitlink::JITLinkDylib* /*JD*/, llvm::jitlink::LinkGraph& G, OnAllocatedFunction OnAllocated) { using AllocResult = Expected>; // Step 1: first-fit over existing slabs. `pool_mu_` only protects // the slabs_ vector — never held across a Slab::allocate call or a // user callback, since the LLJIT linker issues nested lookups (and // thus re-entrant allocate() calls via materialization) from inside // OnAllocated and a coarse lock would deadlock. Snapshot raw pointers // under the lock; slabs are guaranteed to outlive this call because // clearFreeSlabs() is only safe when the session is quiescent. // // Slab::allocate is synchronous (invokes its callback inline on every // code path), so a captured std::optional observes the result before // the call returns. std::vector snapshot; { std::lock_guard lock(pool_mu_); snapshot.reserve(slabs_.size()); for (auto& s : slabs_) snapshot.push_back(s.get()); } for (Slab* slab : snapshot) { std::optional observed; slab->allocate(G, [&](AllocResult R) { observed.emplace(std::move(R)); }); AllocResult result = std::move(*observed); if (result) { OnAllocated(std::move(result)); return; } Error E = result.takeError(); if (E.isA()) { // Retriable: this graph didn't fit this slab's per-pool budget. llvm::consumeError(std::move(E)); continue; } // Terminal (mmap, mprotect, JITLink, BasicLayout). OnAllocated(std::move(E)); return; } // Step 2: grow. Size the fresh slab to fit this graph's per-pool // footprint — normal graphs fall through at slab_size_, skewed or // oversize graphs double up until both pools can host them (see // Slab::capacityForFootprint). A single growth branch replaces the // pre-gate "normal vs oversize" split; we trade one `N × failed // Slab::allocate` scan (step 1) for the extra pre-filter — negligible // since each failed call is a BasicLayout build with no mmap. auto fp = Slab::computeGraphFootprint(G, page_size_); std::size_t cap = Slab::capacityForFootprint(fp, slab_size_); auto slab = createSlab(cap); if (!slab) { OnAllocated( llvm::make_error("SlabPoolMemoryManager: mmap failed for new slab of " + llvm::formatv("{0:x}", cap).str() + " bytes", llvm::inconvertibleErrorCode())); return; } Slab* raw = slab.get(); { std::lock_guard lock(pool_mu_); slabs_.push_back(std::move(slab)); } raw->allocate(G, std::move(OnAllocated)); } void SlabPoolMemoryManager::deallocate(std::vector Allocs, OnDeallocatedFunction OnDeallocated) { Error DeallocErr = Error::success(); for (auto& Alloc : Allocs) { auto* FA = Alloc.release().toPtr(); FA->owner->deallocateOne(FA, DeallocErr); delete FA; } OnDeallocated(std::move(DeallocErr)); } std::size_t SlabPoolMemoryManager::clearFreeSlabs() { // Partition under the lock, move discards to a local vector, drop the // lock, then let ~Slab (which calls munmap) run outside the lock. // Keeping munmap outside pool_mu_ matches the rest of allocate/deallocate, // which also never hold the lock across syscalls that might block. std::vector> discard; { std::lock_guard lock(pool_mu_); auto keep_end = std::partition(slabs_.begin(), slabs_.end(), [](const auto& s) { return !s->isReclaimable(); }); discard.reserve(static_cast(slabs_.end() - keep_end)); for (auto it = keep_end; it != slabs_.end(); ++it) { discard.push_back(std::move(*it)); } slabs_.erase(keep_end, slabs_.end()); } std::size_t reclaimed = discard.size(); // discard goes out of scope — Slab destructors munmap each reservation. return reclaimed; } } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __linux__ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h000066400000000000000000000123461521067262500255130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_memory_manager.h * \brief Per-session growable slab pool. * * `SlabPoolMemoryManager` implements `JITLinkMemoryManager` on top of a * per-session `std::vector>`. On each `allocate` * it picks the first `Slab` that can fit the graph; if none do, it * `mmap`s a fresh slab sized to fit (`Slab::capacityForFootprint`) and * appends it. Normal-size graphs land on a `slab_size`-sized slab; * skewed or oversize graphs land on a power-of-2 larger slab whose * per-pool budgets cover the graph. * * ## Lifecycle * * Once a slab is added to the pool it stays mapped until it is * reclaimed or until the pool (and its enclosing session) is * destroyed. Individual graphs are deallocated via * `FA->owner->deallocateOne(...)`, returning bytes to the slab's free * list. Drained slabs can be returned to the OS via * `clearFreeSlabs()`. * * ## GOTPCRELX relaxation workaround * * See `llvm_patches/gotpcrelx_fix.cc`. The plugin is added per-session * to the `ObjectLinkingLayer` alongside this memory manager and is * orthogonal to pool growth. */ #ifndef TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_ #define TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_ #include #include #include #include #include #include "orcjit_slab.h" namespace tvm { namespace ffi { namespace orcjit { /*! * \brief `JITLinkMemoryManager` backed by a growable pool of `Slab`s. * * The constructor reserves one initial slab (halving its capacity down * to `kMinSlabSize` if `mmap` fails under RLIMIT_AS). Subsequent * slabs are added on demand by `allocate()` at a capacity chosen by * `Slab::capacityForFootprint` — `slab_size_` for normal graphs, the * next power of two up for skewed / oversize graphs. No retry, no * halving on growth; mmap errors propagate. */ class SlabPoolMemoryManager : public llvm::jitlink::JITLinkMemoryManager { public: // Default per-slab capacity. 64 MB is above the p99 size of typical // ML JIT graphs (single-kernel bindings, fused kernels), below the // PC-relative relocation limit, and a multiple of the 2 MB THP // granule. Small enough that a pinned slab only wastes 64 MB of RSS. static constexpr std::size_t kDefaultSlabSize = std::size_t{64} << 20; // 64 MB // Lower bound on initial-slab reservation. If the first `mmap` // fails and halving drops below this, the constructor aborts. // 8 MB is enough for a minimal JITDylib setup under very tight // RLIMIT_AS. static constexpr std::size_t kMinSlabSize = std::size_t{8} << 20; // 8 MB explicit SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size); ~SlabPoolMemoryManager() override = default; SlabPoolMemoryManager(const SlabPoolMemoryManager&) = delete; SlabPoolMemoryManager& operator=(const SlabPoolMemoryManager&) = delete; SlabPoolMemoryManager(SlabPoolMemoryManager&&) = delete; SlabPoolMemoryManager& operator=(SlabPoolMemoryManager&&) = delete; void allocate(const llvm::jitlink::JITLinkDylib* JD, llvm::jitlink::LinkGraph& G, OnAllocatedFunction OnAllocated) override; void deallocate(std::vector Allocs, OnDeallocatedFunction OnDeallocated) override; /*! \brief Number of slabs currently held (test introspection). */ std::size_t numSlabs() const { std::lock_guard lock(pool_mu_); return slabs_.size(); } /*! * \brief Release drained slabs — zero live allocations and at least one * prior allocation — back to the OS via `munmap`. * * Returns the number of slabs reclaimed. Safe to call any time the * session is quiescent (no concurrent JIT work in flight). A typical * pattern is to call this after dropping a batch of libraries: * * for lib in libs: del lib * session.clear_free_slabs() # Python API * * Fresh slabs that have never been allocated on are preserved — the * session remains ready to accept new JIT work. */ std::size_t clearFreeSlabs(); private: /*! \brief Reserve a fresh slab at exactly \p capacity bytes. Returns * nullptr on mmap failure (caller reports the error). */ std::unique_ptr createSlab(std::size_t capacity); std::size_t page_size_; std::size_t slab_size_; mutable std::mutex pool_mu_; std::vector> slabs_; }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc000066400000000000000000000317151521067262500243330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_session.cc * \brief LLVM ORC JIT ExecutionSession implementation */ #include "orcjit_session.h" #include #include #include #include #include #include #include #include #include #include #include "orcjit_dylib.h" #include "orcjit_memory_manager.h" #include "orcjit_utils.h" #if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64)) #include "llvm_patches/gotpcrelx_fix.h" #endif #include "llvm_patches/init_fini_plugin.h" #ifdef __APPLE__ #include "llvm_patches/macho_cxa_atexit_shim.h" #endif #ifdef _WIN32 #include #include "llvm_patches/win_coff_pdata_strip.h" #include "llvm_patches/win_dll_import_generator.h" #endif namespace tvm { namespace ffi { namespace orcjit { // Initialize LLVM native target (only once) struct LLVMInitializer { LLVMInitializer() { llvm::InitializeNativeTarget(); llvm::InitializeNativeTargetAsmPrinter(); llvm::InitializeNativeTargetAsmParser(); } }; static LLVMInitializer llvm_initializer; ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_path, int64_t slab_size_bytes) : jit_(nullptr) { // Create slab-backed memory manager — pre-reserves a contiguous VA region // so all JIT allocations stay within PC-relative relocation range (±2 GB // x86_64, ±4 GB AArch64). Eliminates scattered-mmap relocation overflow // (LLVM #173269). // // slab_size_bytes: 0 = arch default (1 GB x86_64 / AArch64, with fallback), // >0 = custom size, <0 = disable arena (LLJIT uses its // default allocator — scattered mmap, no PC-rel guarantee). // The parameter is Linux-only; on macOS/Windows the arena is compiled out // entirely (see #ifdef below) and the value is ignored. // // `slab_size_bytes` is the per-slab capacity for the growable pool. // Session memory grows in slab-sized increments; graphs that don't // fit a normal slab trigger a power-of-2 larger slab sized to fit // (see `Slab::capacityForFootprint`). // // The default (64 MB) is above typical ML JIT graph sizes while well // under the PC-relative relocation limit. The initial-slab constructor // halves its capacity on mmap failure (RLIMIT_AS, containers) down to // 8 MB; subsequent slabs are reserved at the size returned by // `capacityForFootprint` (>= slab_size) and mmap errors propagate. // // LLJIT auto-configures ObjectLinkingLayer (JITLink) on x86_64 and aarch64 // Linux (see LLJITBuilderState::prepareForConstruction). We override // the layer creator to pass our memory manager. macOS/Windows are gated // off pending testing. (The historical "MachOPlatform teardown crashes // with the arena" concern is moot now that we skip MachOPlatform below, // but enabling the slab on macOS still needs a validation pass.) #ifdef __linux__ if (slab_size_bytes >= 0) { auto page_size = llvm::sys::Process::getPageSizeEstimate(); size_t slab_size; if (slab_size_bytes > 0) { slab_size = static_cast(slab_size_bytes); } else { slab_size = SlabPoolMemoryManager::kDefaultSlabSize; } memory_manager_ = std::make_unique(page_size, slab_size); } #endif auto setup_builder = [this](llvm::orc::LLJITBuilder& builder) { #ifdef __linux__ if (memory_manager_) { builder.setObjectLinkingLayerCreator( [this](llvm::orc::ExecutionSession& ES) -> llvm::Expected> { auto OLL = std::make_unique(ES, *memory_manager_); #if defined(__x86_64__) || defined(_M_X64) OLL->addPlugin(std::make_unique()); #endif return OLL; }); } // if (memory_manager_) #elif defined(__APPLE__) || defined(_WIN32) // Force ObjectLinkingLayer (JITLink) so we can attach InitFiniPlugin. // macOS: LLJIT already defaults to JITLink for Darwin, but the explicit // creator keeps the static_cast in the addPlugin site below type-safe. // Windows: LLJIT defaults to RTDyld; we need JITLink for InitFiniPlugin // and DLLImportDefinitionGenerator. builder.setObjectLinkingLayerCreator( [](llvm::orc::ExecutionSession& ES) -> llvm::Expected> { return std::make_unique(ES); }); #endif #ifdef _WIN32 // Override ProcessSymbols setup to NOT add the default // EPCDynamicLibrarySearchGenerator. That generator resolves symbols to // absolute host-process addresses, which causes PCRel32 overflow when // JIT code calls into DLLs >2GB away. Our DLLImportDefinitionGenerator // (added after construction) wraps every resolved address in a // JIT-allocated PLT stub, keeping all fixups in range. builder.setProcessSymbolsJITDylibSetup( [](llvm::orc::LLJIT& J) -> llvm::Expected { return &J.getExecutionSession().createBareJITDylib(""); }); #endif (void)builder; }; auto builder = llvm::orc::LLJITBuilder(); #ifndef __APPLE__ // macOS: always skip ExecutorNativePlatform / MachOPlatform to sidestep // the compact-unwind 32-bit-delta bug in JITLink's CompactUnwindSupport // (personality delta against a per-JITDylib header base wraps `uint64_t` // and fails `isUInt<32>` when a later user graph mmaps below the header; // see the repo-root fix-machoplatform-libunwind-dso-base.patch for the // full analysis). InitFiniPlugin below handles __mod_init_func / // __mod_term_func instead. Tradeoff: no C++ exception unwinding across // JIT frames on macOS. if (!orc_rt_path.empty()) { builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(orc_rt_path)); } #else (void)orc_rt_path; #endif setup_builder(builder); jit_ = TVM_FFI_ORCJIT_LLVM_CALL(builder.create()); #ifdef _WIN32 // Strip .pdata/.xdata relocations from COFF objects before JITLink graph // building. See llvm_patches/win_coff_pdata_strip.h for the rationale. jit_->getObjTransformLayer().setTransform(&StripCoffPdataXdata); #endif // Use our custom InitFiniPlugin on every platform for init/fini section // collection and priority-ordered execution (ELF .init_array/.fini_array, // MachO __mod_init_func/__mod_term_func, COFF .CRT$XC*/.CRT$XT*). See // llvm_patches/init_fini_plugin.h for per-platform removal criteria. auto& objlayer = jit_->getObjLinkingLayer(); static_cast(objlayer).addPlugin( std::make_unique(this)); #ifdef _WIN32 // On Windows, the default process-symbol generator only searches the main // exe module via GetProcAddress(GetModuleHandle(NULL), ...). Add a // comprehensive generator that searches all loaded DLLs (vcruntime140, // ucrtbase, tvm_ffi, etc.) and creates __imp_* pointer stubs. if (auto PSG = jit_->getProcessSymbolsJITDylib()) { auto& ObjLayer = static_cast(jit_->getObjLinkingLayer()); PSG->addGenerator( std::make_unique(jit_->getExecutionSession(), ObjLayer)); } #endif } ORCJITExecutionSession::ORCJITExecutionSession(const std::string& orc_rt_path, int64_t slab_size_bytes) { ObjectPtr obj = make_object(orc_rt_path, slab_size_bytes); data_ = std::move(obj); } ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary(const String& name) { TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; // Generate name if not provided String lib_name = name; if (lib_name.empty()) { std::ostringstream oss; oss << "dylib_" << dylib_counter_++; lib_name = oss.str(); } llvm::orc::JITDylib& jit_dylib = TVM_FFI_ORCJIT_LLVM_CALL(jit_->getExecutionSession().createJITDylib(lib_name.c_str())); // Use the LLJIT's default link order (Main → Platform → ProcessSymbols). // This provides host process symbols via the ProcessSymbols JITDylib's generator, // while ensuring the platform's __cxa_atexit interposer (in PlatformJD) takes // precedence — so __cxa_atexit handlers are managed by the platform and can be // drained per-JITDylib via __lljit_run_atexits at teardown. for (auto& kv : jit_->defaultLinkOrder()) { jit_dylib.addToLinkOrder(*kv.first, kv.second); } auto dylib_obj = make_object(GetRef(this), &jit_dylib, jit_.get(), lib_name); #ifdef __APPLE__ // Inject ___cxa_atexit on the user JITDylib so it wins over 's // fallback (which resolves to libSystem's and would orphan dtors from // our drop-time drain). See llvm_patches/macho_cxa_atexit_shim.h. InstallCxaAtexitShim(jit_->getExecutionSession(), jit_dylib); #endif return ORCJITDynamicLibrary(std::move(dylib_obj)); } llvm::orc::ExecutionSession& ORCJITExecutionSessionObj::GetLLVMExecutionSession() { TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; return jit_->getExecutionSession(); } llvm::orc::LLJIT& ORCJITExecutionSessionObj::GetLLJIT() { TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; return *jit_; } using CtorDtor = void (*)(); void ORCJITExecutionSessionObj::RunPendingInitializers(llvm::orc::JITDylib& jit_dylib) { auto it = pending_initializers_.find(&jit_dylib); if (it != pending_initializers_.end()) { llvm::sort(it->second, [](const InitFiniEntry& a, const InitFiniEntry& b) { if (a.section != b.section) return static_cast(a.section) < static_cast(b.section); return a.priority < b.priority; }); for (const auto& entry : it->second) { entry.address.toPtr()(); } pending_initializers_.erase(it); } } void ORCJITExecutionSessionObj::RunPendingDeinitializers(llvm::orc::JITDylib& jit_dylib) { auto it = pending_deinitializers_.find(&jit_dylib); if (it != pending_deinitializers_.end()) { llvm::sort(it->second, [](const InitFiniEntry& a, const InitFiniEntry& b) { if (a.section != b.section) return static_cast(a.section) < static_cast(b.section); return a.priority < b.priority; }); for (const auto& entry : it->second) { entry.address.toPtr()(); } pending_deinitializers_.erase(it); } } void ORCJITExecutionSessionObj::AddPendingInitializer(llvm::orc::JITDylib* jit_dylib, const InitFiniEntry& entry) { pending_initializers_[jit_dylib].push_back(entry); } void ORCJITExecutionSessionObj::AddPendingDeinitializer(llvm::orc::JITDylib* jit_dylib, const InitFiniEntry& entry) { pending_deinitializers_[jit_dylib].push_back(entry); } int64_t ORCJITExecutionSessionObj::ClearFreeSlabs() { #ifdef __linux__ if (memory_manager_) { return static_cast(memory_manager_->clearFreeSlabs()); } #endif return 0; } void ORCJITExecutionSessionObj::RemoveDylib(llvm::orc::JITDylib* jit_dylib) { if (jit_dylib == nullptr) return; // Drop any pending init/fini records keyed by this JITDylib*. After removal // the address may be recycled for a freshly-created JITDylib; leftover // entries would then be attributed to the wrong dylib. pending_initializers_.erase(jit_dylib); pending_deinitializers_.erase(jit_dylib); if (jit_ == nullptr) return; // removeJITDylib is best-effort at destruction time: the session may already // be tearing down, the platform may report an error during clear(), etc. // Swallow errors rather than throwing from a destructor; the session // destructor will munmap everything when it runs. if (auto err = jit_->getExecutionSession().removeJITDylib(*jit_dylib)) { llvm::consumeError(std::move(err)); } } } // namespace orcjit } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h000066400000000000000000000125011521067262500241650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_session.h * \brief LLVM ORC JIT ExecutionSession wrapper */ #ifndef TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ #define TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ #include #include #include #include #include #include #include #include #include #include "orcjit_memory_manager.h" namespace tvm { namespace ffi { namespace orcjit { // Forward declaration class ORCJITDynamicLibrary; /*! * \brief ExecutionSession object for LLVM ORC JIT v2 * * This class manages the lifetime of an LLVM ExecutionSession and provides * functionality to create and manage multiple JITDylibs (DynamicLibraries). */ class ORCJITExecutionSessionObj : public Object { public: /*! * \brief Default constructor (for make_object) */ explicit ORCJITExecutionSessionObj(const std::string& orc_rt_path = "", int64_t slab_size_bytes = 0); /*! * \brief Create a new DynamicLibrary (JITDylib) in this session * \param name Optional name for the library (for debugging) * \return The created dynamic library instance */ ORCJITDynamicLibrary CreateDynamicLibrary(const String& name); /*! * \brief Get the underlying LLVM ExecutionSession * \return Reference to the LLVM ExecutionSession */ llvm::orc::ExecutionSession& GetLLVMExecutionSession(); /*! * \brief Get the underlying LLJIT instance * \return Reference to the LLJIT instance */ llvm::orc::LLJIT& GetLLJIT(); static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("orcjit.ExecutionSession", ORCJITExecutionSessionObj, Object); struct InitFiniEntry { enum class Section { kInitArray = 0, kCtors = 1, kDtors = 2, kFiniArray = 3, }; llvm::orc::ExecutorAddr address; Section section; int priority; }; void RunPendingInitializers(llvm::orc::JITDylib& jit_dylib); void RunPendingDeinitializers(llvm::orc::JITDylib& jit_dylib); void AddPendingInitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry); void AddPendingDeinitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry); /*! * \brief Remove a JITDylib from the ExecutionSession, releasing its JIT * memory and dropping it from the session's dylib list. * * Invoked by \c ORCJITDynamicLibraryObj's destructor after any required * static-destructor sequence (\c RunPendingDeinitializers on Linux/Windows, * \c LLJIT::deinitialize on macOS) has completed. The caller must ensure no * further use of the \c JITDylib* after this call — it becomes "Closed" and * its address may be reused by a subsequent \c createJITDylib. * * Also erases any pending init/fini map entries keyed by \p jd so that a * subsequent \c JITDylib allocated at the same address starts with a clean * slate. */ void RemoveDylib(llvm::orc::JITDylib* jd); /*! * \brief Release drained slabs (no live JIT allocations) back to the OS. * * Returns the number of slabs reclaimed. No-op on macOS/Windows * where the slab pool is compiled out, or when the pool has been * disabled via `slab_size < 0`. */ int64_t ClearFreeSlabs(); private: /*! \brief Slab-pool memory manager — must be declared before jit_ for destruction order */ std::unique_ptr memory_manager_; /*! \brief The LLVM ORC JIT instance */ std::unique_ptr jit_; /*! \brief Counter for auto-generating library names */ std::atomic dylib_counter_{0}; std::unordered_map> pending_initializers_; std::unordered_map> pending_deinitializers_; }; /*! * \brief Reference wrapper for ORCJITExecutionSessionObj * * A reference wrapper serves as a reference-counted pointer to the session object. */ class ORCJITExecutionSession : public ObjectRef { public: /*! * \brief Create a new ExecutionSession * \return The created execution session instance */ explicit ORCJITExecutionSession(const std::string& orc_rt_path = "", int64_t slab_size_bytes = 0); // Required: define object reference methods TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ORCJITExecutionSession, ObjectRef, ORCJITExecutionSessionObj); }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc000066400000000000000000000717721521067262500236000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_slab.cc * \brief Slab implementation — bump allocator + commit bitmap + free list. */ #include "orcjit_slab.h" #ifdef __linux__ #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { using namespace llvm; using namespace llvm::jitlink; using namespace llvm::orc; // ── SlabPoolExhaustedError ────────────────────────────────────────── char SlabPoolExhaustedError::ID = 0; SlabPoolExhaustedError::SlabPoolExhaustedError(const char* pool, std::size_t used, std::size_t requested, std::size_t limit) : pool_(pool), used_(used), requested_(requested), limit_(limit) {} void SlabPoolExhaustedError::log(raw_ostream& os) const { os << "Slab: " << pool_ << " pool exhausted (used " << formatv("{0:x}", used_) << " + requested " << formatv("{0:x}", requested_) << " > limit " << formatv("{0:x}", limit_) << ")"; } std::error_code SlabPoolExhaustedError::convertToErrorCode() const { return inconvertibleErrorCode(); } // ── Overflow section edge classification ─────────────────────────── // // Conservative whitelist: only known absolute relocation kinds return true. // Unknown or future edge kinds default to PC-relative → sections stay in // the slab (safe: never breaks relocations, just forgoes the overflow // optimization for unknown kinds). namespace { bool isAbsoluteEdge(const Triple& TT, Edge::Kind K) { if (K < Edge::FirstRelocation) return true; // KeepAlive, Invalid — not a relocation constraint if (TT.isAArch64()) { using namespace llvm::jitlink::aarch64; switch (K) { case Pointer64: case Pointer32: case Pointer64Authenticated: case MoveWide16: return true; default: return false; } } if (TT.isX86()) { using namespace llvm::jitlink::x86_64; switch (K) { case Pointer64: case Pointer32: case Pointer32Signed: case Pointer16: case Pointer8: case Size64: case Size32: return true; default: return false; } } return false; // Unknown arch — treat as PC-relative (safe) } /*! \brief Identify sections eligible for the overflow (separate-mmap) * path. * * Name-based candidate selection followed by edge validation: any * PC-relative cross-section edge targeting a candidate disqualifies * it (the section must live inside the slab so the fixup can reach). * Returns the surviving candidate set. */ DenseSet classifyOverflowSections(LinkGraph& G) { DenseSet candidates; for (auto& Sec : G.sections()) { if (Sec.getMemLifetime() == MemLifetime::NoAlloc) continue; StringRef Name = Sec.getName(); if (Name.starts_with(".nv_fatbin")) { candidates.insert(&Sec); } } if (candidates.empty()) return candidates; const auto& TT = G.getTargetTriple(); for (auto& Sec : G.sections()) { for (auto* B : Sec.blocks()) { for (auto& E : B->edges()) { if (!E.isRelocation()) continue; if (isAbsoluteEdge(TT, E.getKind())) continue; if (!E.getTarget().isDefined()) continue; auto* TargetSec = &E.getTarget().getBlock().getSection(); candidates.erase(TargetSec); } } if (candidates.empty()) break; } return candidates; } } // namespace // ── Platform abstraction ──────────────────────────────────────────── void* Slab::reserveVA(std::size_t size) { void* p = ::mmap(nullptr, size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0); if (p == MAP_FAILED) return nullptr; return p; } void Slab::releaseVA(void* addr, std::size_t size) { int rc = ::munmap(addr, size); assert(rc == 0 && "munmap failed in Slab destructor"); (void)rc; } Error Slab::commitPages(void* addr, std::size_t size) { if (size == 0) return Error::success(); // Commit at commit-chunk (2 MB) granularity for THP promotion. std::size_t offset = static_cast(addr) - arena_base_; std::size_t first_chunk = offset / kCommitGranularity; std::size_t last_chunk = (offset + size - 1) / kCommitGranularity; for (std::size_t i = first_chunk; i <= last_chunk; ++i) { if (committed_[i].load(std::memory_order_acquire) != 0) continue; std::size_t chunk_offset = i * kCommitGranularity; std::size_t chunk_len = std::min(kCommitGranularity, arena_capacity_ - chunk_offset); // mprotect is idempotent, so a concurrent racer calling it on the same chunk // is harmless. Only flip the flag after success — otherwise a failed commit // followed by freeRegion() would leave committed_[i] == 1, causing a // later allocation to skip mprotect and write into PROT_NONE memory. if (::mprotect(arena_base_ + chunk_offset, chunk_len, PROT_READ | PROT_WRITE) != 0) { return make_error("Slab: mprotect(RW) failed for chunk at offset " + formatv("{0:x}", chunk_offset) + ": " + std::strerror(errno), inconvertibleErrorCode()); } committed_[i].store(1, std::memory_order_release); } return Error::success(); } void Slab::decommitPages(void* addr, std::size_t size) { // Intentionally a no-op for slab pages. The ORC runtime may still reference // deallocated JIT memory during session teardown (e.g., ELFNixPlatform // deinitializers run after some allocations are freed). Decommitting // (MADV_DONTNEED or mprotect PROT_NONE) would cause segfaults or illegal // instructions during shutdown. // // Physical pages stay committed but are returned to the free list for reuse. // The slab destructor releases all VA and physical memory via munmap. (void)addr; (void)size; } Error Slab::protectPages(void* addr, std::size_t size, MemProt Prot) { int prot = PROT_NONE; if ((Prot & MemProt::Read) != MemProt::None) prot |= PROT_READ; if ((Prot & MemProt::Write) != MemProt::None) prot |= PROT_WRITE; if ((Prot & MemProt::Exec) != MemProt::None) prot |= PROT_EXEC; if (::mprotect(addr, size, prot) != 0) { return make_error("Slab: mprotect failed at " + formatv("{0:x}", addr) + " size " + formatv("{0:x}", size) + ": " + std::strerror(errno), inconvertibleErrorCode()); } if ((Prot & MemProt::Exec) != MemProt::None) { sys::Memory::InvalidateInstructionCache(addr, size); } return Error::success(); } // ── InFlightAlloc ─────────────────────────────────────────────────── class Slab::InFlightAlloc : public JITLinkMemoryManager::InFlightAlloc { public: // A contiguous region within one pool: [offset, offset + standard_size + finalize_size). // Standard-lifetime bytes come first; Finalize-lifetime bytes follow and are freed // at the end of finalize(). Any field may be 0 to indicate no allocation from // that pool on this call. struct PoolRegion { std::size_t offset; std::size_t standard_size; std::size_t finalize_size; }; InFlightAlloc(Slab& S, LinkGraph& G, BasicLayout BL, PoolRegion non_exec, PoolRegion exec, std::vector overflow_blocks) : S(S), G(&G), BL(std::move(BL)), non_exec_(non_exec), exec_(exec), overflow_blocks_(std::move(overflow_blocks)) {} ~InFlightAlloc() override { assert(!G && "Slab::InFlightAlloc destroyed without finalize or abandon"); } void finalize(OnFinalizedFunction OnFinalized) override { // Apply target protections for each slab segment. if (auto Err = applyProtections()) { OnFinalized(std::move(Err)); return; } // Apply target protections for overflow blocks. for (auto& ob : overflow_blocks_) { if (auto Err = S.protectPages(ob.addr, ob.size, ob.prot)) { OnFinalized(std::move(Err)); return; } } // Run finalization actions (e.g., register EH frames). auto DeallocActions = shared::runFinalizeActions(BL.graphAllocActions()); if (!DeallocActions) { OnFinalized(DeallocActions.takeError()); return; } // Decommit finalize-lifetime pages in each pool — they're no longer needed. for (auto& R : {non_exec_, exec_}) { if (R.finalize_size > 0) { S.decommitPages(S.arena_base_ + R.offset + R.standard_size, R.finalize_size); S.freeRegion(R.offset + R.standard_size, R.finalize_size); } } #ifndef NDEBUG G = nullptr; #endif // Create finalized allocation handle. LLVM's FinalizedAlloc stores an // opaque ExecutorAddr (integer), so we must use raw new here. Ownership // transfers to deallocate(), which LLVM guarantees is called for every // finalized allocation. auto* FA = new FinalizedAllocInfo{&S, non_exec_.offset, non_exec_.standard_size, exec_.offset, exec_.standard_size, std::move(*DeallocActions), std::move(overflow_blocks_)}; // Bump the slab's live-alloc counter before publishing the handle so // a concurrent `clearFreeSlabs` in another thread cannot see this // slab as reclaimable between handle publication and the FA reaching // LLJIT's bookkeeping. S.noteAllocated(); OnFinalized(JITLinkMemoryManager::FinalizedAlloc(ExecutorAddr::fromPtr(FA))); } void abandon(OnAbandonedFunction OnAbandoned) override { // Decommit and return each pool's full region to the appropriate free list. for (auto& R : {non_exec_, exec_}) { std::size_t total = R.standard_size + R.finalize_size; if (total > 0) { S.decommitPages(S.arena_base_ + R.offset, total); S.freeRegion(R.offset, total); } } // Release overflow blocks. for (auto& ob : overflow_blocks_) { ::munmap(ob.addr, ob.size); } #ifndef NDEBUG G = nullptr; #endif OnAbandoned(Error::success()); } private: Error applyProtections() { for (auto& KV : BL.segments()) { const auto& AG = KV.first; auto& Seg = KV.second; auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, S.page_size_); if (auto Err = S.protectPages(Seg.WorkingMem, SegSize, AG.getMemProt())) return Err; } return Error::success(); } Slab& S; LinkGraph* G; BasicLayout BL; PoolRegion non_exec_; PoolRegion exec_; std::vector overflow_blocks_; }; // ── Slab ──────────────────────────────────────────────────────────── Slab::Slab(std::size_t page_size, std::size_t capacity) : arena_base_(nullptr), arena_capacity_(capacity), page_size_(page_size), midpoint_(0), exec_bump_limit_(0), non_exec_bump_(0), exec_bump_(0) { arena_base_ = static_cast(reserveVA(capacity)); if (!arena_base_) return; // Caller inspects isValid() and retries. // Partition the slab into two pools at a 2 MB-aligned midpoint. The // exec pool starts at midpoint_, which is therefore on a 2 MB // boundary — r-x segments pack into a minimum number of 2 MB pages. // // Constraint: cross-pool displacements (e.g. .text → .rodata via // ADRP+ADD on aarch64) must fit in ±kPCRelReach. The farthest pair of // bytes is (end of exec, start of non-exec), separated by at most // `exec_bump_limit_`, so we cap the exec pool's upper bound at // kPCRelReach even when the VA reservation is larger. exec_bump_limit_ = std::min(capacity, kPCRelReach); std::size_t raw_midpoint = static_cast(exec_bump_limit_ * kDefaultNonExecFraction); midpoint_ = (raw_midpoint / kCommitGranularity) * kCommitGranularity; if (midpoint_ == 0) midpoint_ = kCommitGranularity; if (midpoint_ >= exec_bump_limit_) midpoint_ = exec_bump_limit_ - kCommitGranularity; non_exec_bump_ = 0; exec_bump_ = midpoint_; // Initialize commit tracking. make_unique(n) value-initializes // the array to zero in C++17. num_commit_chunks_ = (capacity + kCommitGranularity - 1) / kCommitGranularity; committed_ = std::make_unique[]>(num_commit_chunks_); // Hint THP promotion for the entire slab. Intentionally unchecked — // MADV_HUGEPAGE is advisory and may fail if THP is disabled system-wide. (void)::madvise(arena_base_, capacity, MADV_HUGEPAGE); } Slab::~Slab() { if (arena_base_) { releaseVA(arena_base_, arena_capacity_); } } Expected Slab::bumpAllocate(std::size_t size, bool is_exec) { std::lock_guard Lock(mu_); auto& free_list = is_exec ? free_list_exec_ : free_list_non_exec_; auto& bump = is_exec ? exec_bump_ : non_exec_bump_; std::size_t limit = is_exec ? exec_bump_limit_ : midpoint_; // Try free list first (best-fit). O(n) scan — acceptable for the expected // workload of tens of JIT allocations, not thousands. std::size_t best_idx = free_list.size(); std::size_t best_waste = std::numeric_limits::max(); for (std::size_t i = 0; i < free_list.size(); ++i) { if (free_list[i].size >= size && free_list[i].size - size < best_waste) { best_idx = i; best_waste = free_list[i].size - size; if (best_waste == 0) break; } } if (best_idx < free_list.size()) { std::size_t offset = free_list[best_idx].offset; if (free_list[best_idx].size == size) { free_list.erase(free_list.begin() + best_idx); } else { free_list[best_idx].offset += size; free_list[best_idx].size -= size; } return offset; } // Bump allocate within the pool's limit. if (bump + size > limit) { return make_error(is_exec ? "exec" : "non-exec", bump, size, limit); } std::size_t offset = bump; bump += size; return offset; } void Slab::freeRegion(std::size_t offset, std::size_t size) { if (size == 0) return; std::lock_guard Lock(mu_); // Route to the correct pool's free list based on offset. auto& free_list = (offset >= midpoint_) ? free_list_exec_ : free_list_non_exec_; // Insert into free list in sorted order. auto it = std::lower_bound(free_list.begin(), free_list.end(), offset, [](const FreeBlock& fb, std::size_t off) { return fb.offset < off; }); it = free_list.insert(it, FreeBlock{offset, size}); // Coalesce with next. auto next = it + 1; if (next != free_list.end() && it->offset + it->size == next->offset) { it->size += next->size; free_list.erase(next); } // Coalesce with previous. if (it != free_list.begin()) { auto prev = it - 1; if (prev->offset + prev->size == it->offset) { prev->size += it->size; free_list.erase(it); } } } Slab::GraphFootprint Slab::computeGraphFootprint(LinkGraph& G, std::size_t page_size) { // Overflow sections live outside any slab (separate mmap at finalize // time) — exclude them from the in-slab footprint. See // classifyOverflowSections() for the candidate-selection rules. DenseSet overflow_candidates = classifyOverflowSections(G); SmallVector, 4> hidden; for (auto* Sec : overflow_candidates) { hidden.push_back({Sec, Sec->getMemLifetime()}); Sec->setMemLifetime(MemLifetime::NoAlloc); } BasicLayout BL(G); for (auto& [Sec, OrigLifetime] : hidden) { Sec->setMemLifetime(OrigLifetime); } std::size_t ne_total = 0, e_total = 0; for (auto& KV : BL.segments()) { auto& AG = KV.first; auto& Seg = KV.second; auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size); bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None; if (is_exec) { e_total += SegSize; } else { ne_total += SegSize; } } return GraphFootprint{ne_total, e_total}; } std::size_t Slab::capacityForFootprint(GraphFootprint fp, std::size_t base_size) { // Mirrors the split formula in Slab::Slab (see the ctor for the // rationale behind each clamp). Kept local rather than sharing a // helper with the ctor so that changes to the split policy require // a deliberate touch in both places. auto budgets = [](std::size_t cap) { std::size_t exec_limit = std::min(cap, kPCRelReach); std::size_t raw_mid = static_cast(exec_limit * kDefaultNonExecFraction); std::size_t mid = (raw_mid / kCommitGranularity) * kCommitGranularity; if (mid == 0) mid = kCommitGranularity; if (mid >= exec_limit) mid = exec_limit - kCommitGranularity; return std::pair{mid, exec_limit - mid}; }; std::size_t cap = base_size; while (true) { auto [ne_budget, e_budget] = budgets(cap); if (fp.non_exec <= ne_budget && fp.exec <= e_budget) return cap; // Budgets plateau once exec_limit saturates; further VA doesn't help. if (cap >= kPCRelReach) return cap; cap *= 2; } } void Slab::allocate(LinkGraph& G, JITLinkMemoryManager::OnAllocatedFunction OnAllocated) { // ── Overflow section classification ── // // Sections matching known overflow names (e.g. .nv_fatbin — large GPU // device blobs referenced only by absolute relocations) are allocated // outside the slab via separate mmap(), keeping the slab compact for // code + small rodata. See classifyOverflowSections() for the // candidate-selection + edge-validation rules. // // Validated candidates are temporarily set to NoAlloc so BasicLayout // skips them, then immediately restored before returning. By the time // JITLink's fixUpBlocks runs, sections are back to Standard — avoiding // the debug assert that prohibits edges from allocated sections to // NoAlloc sections. DenseSet overflow_candidates = classifyOverflowSections(G); // Apply: temporarily hide validated overflow sections from BasicLayout. SmallVector, 4> overflow_sections; for (auto* Sec : overflow_candidates) { overflow_sections.push_back({Sec, Sec->getMemLifetime()}); Sec->setMemLifetime(MemLifetime::NoAlloc); } BasicLayout BL(G); // Restore overflow sections to their original lifetime immediately. // BasicLayout has already captured its segment list; subsequent LLVM // passes (fixUpBlocks) will see the sections as normal Standard sections. for (auto& [Sec, OrigLifetime] : overflow_sections) { Sec->setMemLifetime(OrigLifetime); } // Compute total sizes grouped by lifetime. auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(page_size_); if (!SegsSizes) { OnAllocated(SegsSizes.takeError()); return; } if (SegsSizes->total() > std::numeric_limits::max()) { OnAllocated(make_error( "Total requested size " + formatv("{0:x}", SegsSizes->total()) + " for graph " + G.getName() + " exceeds address space")); return; } auto TotalSize = static_cast(SegsSizes->total()); if (TotalSize == 0 && overflow_sections.empty()) { // Empty graph — return a no-op allocation. OnAllocated(std::make_unique(*this, G, std::move(BL), InFlightAlloc::PoolRegion{0, 0, 0}, InFlightAlloc::PoolRegion{midpoint_, 0, 0}, std::vector{})); return; } // ── Dual-pool split ── // // Partition each segment into one of four buckets based on (Prot, Lifetime): // non-exec × Standard / Finalize → non-exec pool (below midpoint_) // exec × Standard / Finalize → exec pool (at/above midpoint_) // // Within each pool, Standard segments come first and Finalize segments // second, so the Finalize tail of each pool can be freed after finalize(). std::size_t ne_std_size = 0, ne_fin_size = 0; std::size_t e_std_size = 0, e_fin_size = 0; for (auto& KV : BL.segments()) { auto& AG = KV.first; auto& Seg = KV.second; auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size_); bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None; bool is_finalize = AG.getMemLifetime() == MemLifetime::Finalize; if (is_exec) { (is_finalize ? e_fin_size : e_std_size) += SegSize; } else { (is_finalize ? ne_fin_size : ne_std_size) += SegSize; } } std::size_t ne_total = ne_std_size + ne_fin_size; std::size_t e_total = e_std_size + e_fin_size; InFlightAlloc::PoolRegion ne_region{0, 0, 0}; InFlightAlloc::PoolRegion e_region{midpoint_, 0, 0}; auto allocPool = [&](std::size_t req, bool is_exec) -> Expected { if (req == 0) return std::size_t{0}; auto off = bumpAllocate(req, is_exec); if (!off) return off.takeError(); if (auto Err = commitPages(arena_base_ + *off, req)) { freeRegion(*off, req); return std::move(Err); } // Recycled-region protection reset. commitPages() only mprotects a 2 MB // commit-chunk the first time it's touched (guarded by committed_). // When the region was previously handed out, finalized to r-x/r--/rw- // and later returned to the free list by deallocateOne(), those finalize // protections are still in effect — the memset(0) below would fault on // read-only pages. mprotect(RW) has no decommit effect, so the // physical pages stay resident; we're only restoring write access for // the upcoming zero-fill and subsequent JITLink content writes. if (auto Err = protectPages(arena_base_ + *off, req, MemProt::Read | MemProt::Write)) { freeRegion(*off, req); return std::move(Err); } std::memset(arena_base_ + *off, 0, req); return *off; }; if (ne_total > 0) { auto off = allocPool(ne_total, /*is_exec=*/false); if (!off) { OnAllocated(off.takeError()); return; } ne_region = {*off, ne_std_size, ne_fin_size}; } if (e_total > 0) { auto off = allocPool(e_total, /*is_exec=*/true); if (!off) { // Unwind non-exec allocation on failure to keep the pools consistent. if (ne_total > 0) { decommitPages(arena_base_ + ne_region.offset, ne_total); freeRegion(ne_region.offset, ne_total); } OnAllocated(off.takeError()); return; } e_region = {*off, e_std_size, e_fin_size}; } // Assign addresses to segments from four cursors. Standard comes first in // each pool, then Finalize. auto NeStdCursor = ExecutorAddr::fromPtr(arena_base_ + ne_region.offset); auto NeFinCursor = ExecutorAddr::fromPtr(arena_base_ + ne_region.offset + ne_std_size); auto EStdCursor = ExecutorAddr::fromPtr(arena_base_ + e_region.offset); auto EFinCursor = ExecutorAddr::fromPtr(arena_base_ + e_region.offset + e_std_size); for (auto& KV : BL.segments()) { auto& AG = KV.first; auto& Seg = KV.second; bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None; bool is_finalize = AG.getMemLifetime() == MemLifetime::Finalize; auto& Cursor = is_exec ? (is_finalize ? EFinCursor : EStdCursor) : (is_finalize ? NeFinCursor : NeStdCursor); Seg.WorkingMem = Cursor.toPtr(); Seg.Addr = Cursor; auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size_); Cursor += SegSize; } // Apply layout — copies content and assigns block addresses for slab segments. if (auto Err = BL.apply()) { // On error: decommit and free both pool regions. if (ne_total > 0) { decommitPages(arena_base_ + ne_region.offset, ne_total); freeRegion(ne_region.offset, ne_total); } if (e_total > 0) { decommitPages(arena_base_ + e_region.offset, e_total); freeRegion(e_region.offset, e_total); } OnAllocated(std::move(Err)); return; } // ── Allocate overflow sections via mmap() outside the slab ── std::vector overflow_allocs; for (auto& [Sec, _] : overflow_sections) { // Compute total size for this section's blocks. std::size_t total_sec_size = 0; for (auto* B : Sec->blocks()) { total_sec_size = alignTo(total_sec_size, B->getAlignment()); total_sec_size += B->getSize(); } if (total_sec_size == 0) continue; total_sec_size = alignTo(total_sec_size, page_size_); // mmap outside the slab. void* addr = ::mmap(nullptr, total_sec_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); if (addr == MAP_FAILED) { // Clean up prior overflow allocs, free both pool regions, report error. for (auto& ob : overflow_allocs) ::munmap(ob.addr, ob.size); if (ne_total > 0) { decommitPages(arena_base_ + ne_region.offset, ne_total); freeRegion(ne_region.offset, ne_total); } if (e_total > 0) { decommitPages(arena_base_ + e_region.offset, e_total); freeRegion(e_region.offset, e_total); } OnAllocated(make_error( "Slab: overflow mmap failed for section " + Sec->getName() + ": " + std::strerror(errno), inconvertibleErrorCode())); return; } // Layout blocks within the mmap'd region. char* ptr = static_cast(addr); for (auto* B : Sec->blocks()) { uint64_t align = B->getAlignment(); ptr = reinterpret_cast(alignTo(reinterpret_cast(ptr), align)); std::size_t bsize = B->getSize(); // Copy content and redirect block's mutable content pointer. if (!B->isZeroFill()) { auto content = B->getContent(); std::memcpy(ptr, content.data(), content.size()); B->setMutableContent(MutableArrayRef(ptr, bsize)); } // Assign block address (working mem == executor addr for in-process JIT). B->setAddress(ExecutorAddr::fromPtr(ptr)); ptr += bsize; } overflow_allocs.push_back({addr, total_sec_size, Sec->getMemProt()}); } OnAllocated(std::make_unique(*this, G, std::move(BL), ne_region, e_region, std::move(overflow_allocs))); } void Slab::deallocateOne(FinalizedAllocInfo* FA, Error& err_out) { // Run deallocation actions in reverse order. while (!FA->DeallocActions.empty()) { if (auto Err = FA->DeallocActions.back().runWithSPSRetErrorMerged()) { err_out = joinErrors(std::move(err_out), std::move(Err)); } FA->DeallocActions.pop_back(); } // Decommit and free each pool's Standard region. // // We intentionally do *not* reset page protection here. During session // teardown the ORC runtime may still execute (or read) deallocated JIT // pages from other dylibs while their DeallocActions unwind — same // rationale as decommitPages being a no-op. Protection is reset lazily // in allocate() when the region is re-handed out (see the // "recycled-region protection reset" block in allocPool). if (FA->non_exec_standard_size > 0) { decommitPages(arena_base_ + FA->non_exec_offset, FA->non_exec_standard_size); freeRegion(FA->non_exec_offset, FA->non_exec_standard_size); } if (FA->exec_standard_size > 0) { decommitPages(arena_base_ + FA->exec_offset, FA->exec_standard_size); freeRegion(FA->exec_offset, FA->exec_standard_size); } // Decrement the live-alloc counter. After this point the slab may be // observed as reclaimable by `SlabPoolMemoryManager::clearFreeSlabs`. // Safe: all DeallocActions have already run and the region has been // returned to the free list. noteDeallocated(); // Release overflow blocks. for (auto& ob : FA->overflow_blocks) { ::munmap(ob.addr, ob.size); } } } // namespace orcjit } // namespace ffi } // namespace tvm #endif // __linux__ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h000066400000000000000000000310231521067262500234230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_slab.h * \brief Single contiguous-VA region + dual-pool bump allocator. * * A `Slab` owns one `mmap(PROT_NONE)` reservation of fixed capacity and * bump-allocates from it, keeping all JIT allocations within range of * PC-relative relocations (±2 GB on x86_64, ±4 GB on AArch64). * * The `Slab` is the unit-of-VA-reservation for the OrcJIT memory manager. * Today it is used as a single-slab arena owned by * `ArenaJITLinkMemoryManager`. Stage B of the refactor will introduce a * `SlabPoolMemoryManager` that holds multiple Slabs and grows by mmap-ing * new ones on demand. * * ## Page commit + Transparent Huge Page (THP) support * * Pages are committed in 2 MB chunks (`kCommitGranularity`) — the 2 MB * size matches the Linux huge-page granule on both x86_64 and AArch64, * enabling THP promotion via `madvise(MADV_HUGEPAGE)` on the full * reservation. Each 2 MB commit-chunk is `mprotect`-ed to RW exactly once * via an atomic bitmap flag (`committed_`), avoiding lock contention with * the per-pool allocator mutex. * * ## Dual-pool exec / non-exec split * * The slab is partitioned at a 2 MB-aligned `midpoint_` into two bump * pools: * * non-exec pool = [base, base + midpoint_ ) * exec pool = [base + midpoint_, base + exec_bump_limit_ ) * * Both pools grow upward; cross-pool displacements (.text → .rodata etc.) * must fit in ±`kPCRelReach` — we cap `exec_bump_limit_` at that reach * even when the VA reservation is larger. */ #ifndef TVM_FFI_ORCJIT_ORCJIT_SLAB_H_ #define TVM_FFI_ORCJIT_ORCJIT_SLAB_H_ #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { class Slab; // forward-declared for FinalizedAllocInfo. /*! * \brief Retriable "slab is out of room for this graph" error. * * Raised from `Slab::bumpAllocate` when the requested bytes would * overflow the pool bump cursor (and no free-list region fits). * `SlabPoolMemoryManager::allocate` catches this error type * specifically, consumes it, and retries on the next slab — or * creates a new one. * * All other `Slab::allocate` failures (mmap, mprotect, JITLink) stay * as `StringError` / `JITLinkError` and are propagated to the caller * without retry. */ class SlabPoolExhaustedError : public llvm::ErrorInfo { public: static char ID; SlabPoolExhaustedError(const char* pool, std::size_t used, std::size_t requested, std::size_t limit); void log(llvm::raw_ostream& os) const override; std::error_code convertToErrorCode() const override; private: const char* pool_; std::size_t used_; std::size_t requested_; std::size_t limit_; }; /*! * \brief Metadata for a finalized allocation, stored via FinalizedAlloc * handle. * * Each allocate() call may consume a region from either or both pools. * Standard-lifetime pages remain committed after finalize(); * Finalize-lifetime pages are decommitted at the end of finalize(). * Zero-sized sub-regions indicate no allocation from that pool. * * \p owner points to the Slab that handed out this allocation. With one * slab per session today this is redundant, but stamping it now makes * Stage B's pool-manager routing O(1) without address comparison. */ struct FinalizedAllocInfo { Slab* owner; ///< Slab that owns these offsets. std::size_t non_exec_offset; ///< offset of non-exec Standard region (or 0 if unused). std::size_t non_exec_standard_size; std::size_t exec_offset; ///< offset of exec Standard region (or midpoint_ if unused). std::size_t exec_standard_size; std::vector DeallocActions; struct OverflowBlock { void* addr; ///< separately-mmap'd base (outside the slab). std::size_t size; ///< mapping size (page-aligned). llvm::orc::MemProt prot; ///< target protection for finalize. }; std::vector overflow_blocks; }; /*! * \brief One contiguous JIT memory reservation. * * Exposes the per-graph `allocate` entry point (matching * `JITLinkMemoryManager::allocate`'s callback signature) and the * per-FinalizedAlloc `deallocateOne` used by the outer memory manager to * route deallocation. */ class Slab { public: // Commit / THP granularity. Every 2 MB chunk is mprotect'd RW exactly // once via `committed_`; `madvise(MADV_HUGEPAGE)` can then promote a // fully-faulted chunk into a single huge page. static constexpr std::size_t kCommitGranularity = std::size_t{2} << 20; // 2 MB // PC-relative relocation reach (tightest binding fixup). Cross-pool // references must fit in a signed 32-bit displacement. The binding // constraint on both x86_64 and aarch64 is the signed 32-bit Delta32 // used in .eh_frame unwind records (±2 GB), not the wider ADRP+ADD / // RIP-rel reach. `exec_bump_limit_` is capped at this reach so // cross-pool Delta32 fixups always resolve. static constexpr std::size_t kPCRelReach = (std::size_t{1} << 31) - kCommitGranularity; // ~2 GB // Fraction of the slab reserved for non-exec segments (r--, rw-). The // remainder holds exec (r-x). Typical CUDA binding objects: ~2 parts // rodata+data to 1 part text. static constexpr double kDefaultNonExecFraction = 2.0 / 3.0; /*! \brief Construct a Slab and reserve \p capacity bytes of VA. * * On reservation failure, returns with \c base() == nullptr — the * caller is expected to retry at a smaller capacity or * \c report_fatal_error. */ Slab(std::size_t page_size, std::size_t capacity); ~Slab(); Slab(const Slab&) = delete; Slab& operator=(const Slab&) = delete; Slab(Slab&&) = delete; Slab& operator=(Slab&&) = delete; /*! \brief True iff the reservation succeeded. */ bool isValid() const noexcept { return arena_base_ != nullptr; } /*! * \brief Page-aligned per-pool byte totals for a LinkGraph. * * Does not include sections routed to the overflow (separate-mmap) * path. `SlabPoolMemoryManager` uses this to size a fresh slab — * see `capacityForFootprint`. */ struct GraphFootprint { std::size_t non_exec; std::size_t exec; std::size_t total() const noexcept { return non_exec + exec; } }; /*! * \brief Compute pool footprint for \p G without mutating bookkeeping. * * Temporarily reclassifies overflow-section lifetimes while a * `BasicLayout` is built, then restores them — same prefix as * `allocate()`. Safe to call repeatedly; no state is retained. */ static GraphFootprint computeGraphFootprint(llvm::jitlink::LinkGraph& G, std::size_t page_size); /*! * \brief Power-of-2 capacity that fits \p fp in both pools. * * Starts at \p base_size (the pool's nominal slab size) and doubles * until the Slab split formula yields per-pool budgets ≥ \p fp. The * returned value, when fed back to the Slab ctor, produces a slab * whose non-exec and exec pools are both large enough to host the * graph — no separate "oversize path" needed at the pool layer. * * Stops doubling once the capacity reaches \c kPCRelReach : beyond * that point \c exec_bump_limit_ saturates, so the budgets no longer * grow with VA. For truly gigantic graphs the caller will observe a * subsequent \c Slab::allocate failure — rare in practice because * per-pool budgets at the plateau are ~1.3 GB / ~0.7 GB. */ static std::size_t capacityForFootprint(GraphFootprint fp, std::size_t base_size); /*! \brief Single-graph JIT allocation entry point. */ void allocate(llvm::jitlink::LinkGraph& G, llvm::jitlink::JITLinkMemoryManager::OnAllocatedFunction OnAllocated); /*! \brief Per-FA teardown: run DeallocActions, free pool regions, * release overflow blocks. Caller deletes the FA afterwards. * * Errors are joined into \p err_out; never throws. */ void deallocateOne(FinalizedAllocInfo* FA, llvm::Error& err_out); /*! \brief Address-range ownership check. */ bool contains(const void* addr) const noexcept { auto* p = static_cast(addr); return p >= arena_base_ && p < arena_base_ + arena_capacity_; } char* base() const noexcept { return arena_base_; } std::size_t capacity() const noexcept { return arena_capacity_; } std::size_t page_size() const noexcept { return page_size_; } /*! * \brief Record that a new FinalizedAlloc was published from this slab. * Called by `InFlightAlloc::finalize` just before returning the * FinalizedAlloc handle to the caller. */ void noteAllocated() noexcept { live_count_.fetch_add(1, std::memory_order_relaxed); ever_used_.store(true, std::memory_order_release); } /*! * \brief Record that a FinalizedAlloc on this slab has been released. * Called by `deallocateOne` after the region is returned to the * free list and DeallocActions have run. */ void noteDeallocated() noexcept { live_count_.fetch_sub(1, std::memory_order_acq_rel); } /*! * \brief True iff this slab has ever hosted a FinalizedAlloc and * currently has zero live ones. * * Used by `SlabPoolMemoryManager::clearFreeSlabs` to decide which * slabs can be munmap'd. A fresh slab that has never been used is * *not* reclaimable — the caller still expects to allocate on it. */ bool isReclaimable() const noexcept { return ever_used_.load(std::memory_order_acquire) && live_count_.load(std::memory_order_acquire) == 0; } private: class InFlightAlloc; // defined in orcjit_slab.cc /*! \brief Bump-allocate from the selected pool. Returns offset within * the slab's VA reservation. */ llvm::Expected bumpAllocate(std::size_t size, bool is_exec); /*! \brief Return a region to the appropriate free list. Pool is * identified by comparing offset against midpoint_. */ void freeRegion(std::size_t offset, std::size_t size); // ── Platform abstraction (all implemented in orcjit_slab.cc) ── static void* reserveVA(std::size_t size); static void releaseVA(void* addr, std::size_t size); llvm::Error commitPages(void* addr, std::size_t size); static void decommitPages(void* addr, std::size_t size); static llvm::Error protectPages(void* addr, std::size_t size, llvm::orc::MemProt Prot); char* arena_base_; std::size_t arena_capacity_; std::size_t page_size_; // Dual-pool split. See class docstring. std::size_t midpoint_; std::size_t exec_bump_limit_; std::mutex mu_; std::size_t non_exec_bump_; // next free offset in non-exec pool ∈ [0, midpoint_] std::size_t exec_bump_; // next free offset in exec pool ∈ [midpoint_, exec_bump_limit_] struct FreeBlock { std::size_t offset; std::size_t size; }; std::vector free_list_non_exec_; std::vector free_list_exec_; /*! \brief Per-commit-chunk flags (0 = uncommitted, 1 = committed). * Lock-free: each chunk is mprotect'd exactly once via * compare_exchange. */ std::unique_ptr[]> committed_; std::size_t num_commit_chunks_ = 0; /*! \brief Count of live FinalizedAllocs held on this slab. */ std::atomic live_count_{0}; /*! \brief Becomes true on the first `noteAllocated` call; never reset. * A slab that has never been used is not a reclaim candidate * (callers may still plan to allocate on it). */ std::atomic ever_used_{false}; }; } // namespace orcjit } // namespace ffi } // namespace tvm #endif // TVM_FFI_ORCJIT_ORCJIT_SLAB_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/src/ffi/orcjit_utils.h000066400000000000000000000063401521067262500236460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file orcjit_utils.h * \brief LLVM ORC JIT utility macros and helpers. * * Provides TVM_FFI_ORCJIT_LLVM_CALL, a macro that checks an LLVM * Error or Expected result and throws a TVM FFI InternalError * on failure. */ #ifndef TVM_FFI_ORCJIT_ORCJIT_UTILS_H_ #define TVM_FFI_ORCJIT_ORCJIT_UTILS_H_ #include #include #include #include namespace tvm { namespace ffi { namespace orcjit { namespace detail { /*! \brief Check an llvm::Error and throw on failure. */ inline void CallLLVM(llvm::Error err, const char* expr, const char* file, int line) { if (err) { std::string err_msg; llvm::handleAllErrors(std::move(err), [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); TVM_FFI_THROW(InternalError) << file << ":" << line << ": " << expr << " failed: " << err_msg; } } /*! \brief Check an llvm::Expected and return the value or throw on failure. */ template inline T CallLLVM(llvm::Expected value_or_err, const char* expr, const char* file, int line) { if (value_or_err) return std::move(*value_or_err); std::string err_msg; llvm::handleAllErrors(std::move(value_or_err.takeError()), [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); TVM_FFI_THROW(InternalError) << file << ":" << line << ": " << expr << " failed: " << err_msg; } /*! \brief Check an llvm::Expected and return the reference or throw on failure. */ template inline T& CallLLVM(llvm::Expected value_or_err, const char* expr, const char* file, int line) { if (value_or_err) return *value_or_err; std::string err_msg; llvm::handleAllErrors(std::move(value_or_err.takeError()), [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); TVM_FFI_THROW(InternalError) << file << ":" << line << ": " << expr << " failed: " << err_msg; } } // namespace detail } // namespace orcjit } // namespace ffi } // namespace tvm /*! * \brief Check an LLVM Error or Expected and throw on failure. * * Usage: * TVM_FFI_ORCJIT_LLVM_CALL(builder.create()); * auto jit = TVM_FFI_ORCJIT_LLVM_CALL(builder.create()); * TVM_FFI_ORCJIT_LLVM_CALL(jit->addObjectFile(...)); */ #define TVM_FFI_ORCJIT_LLVM_CALL(expr) \ ::tvm::ffi::orcjit::detail::CallLLVM((expr), #expr, __FILE__, __LINE__) #endif // TVM_FFI_ORCJIT_ORCJIT_UTILS_H_ tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/000077500000000000000000000000001521067262500205675ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/CMakeLists.txt000066400000000000000000000073211521067262500233320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.18) project(tvm_ffi_orcjit_tests) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(TEST_OBJ_INSTALL_SUFFIX "" CACHE STRING "Suffix for install subdirectory (e.g., -gcc)" ) if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64") add_compile_options(-mno-outline-atomics) endif () # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Helper: build a source file into an object and install as /.o Derives the target # name and install subdirectory from the source path. function (add_test_object source) get_filename_component(_name "${source}" NAME_WE) get_filename_component(_ext "${source}" EXT) get_filename_component(_dir "${source}" DIRECTORY) get_filename_component(_subdir "${_dir}" NAME) set(target_name "${_subdir}_${_name}") if (_ext STREQUAL ".c") set(lang C) else () set(lang CXX) endif () add_library(${target_name}_obj OBJECT "${source}") target_link_libraries(${target_name}_obj PRIVATE tvm_ffi::header) if (lang STREQUAL "C" AND MSVC) # /GS- disables security cookie checks (__security_cookie) which are CRT symbols the ORC JIT # cannot resolve. target_compile_options(${target_name}_obj PRIVATE /O2 /GS-) else () target_compile_options(${target_name}_obj PRIVATE -O2) endif () install( FILES $ DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/${_subdir}${TEST_OBJ_INSTALL_SUFFIX} RENAME ${_name}.o ) endfunction () # C++ object files — skip on Windows (C-only strategy for ORC JIT on Windows) if (NOT WIN32) add_test_object(sources/cc/test_funcs.cc) add_test_object(sources/cc/test_funcs2.cc) add_test_object(sources/cc/test_funcs_conflict.cc) add_test_object(sources/cc/test_ctor_dtor.cc) add_test_object(sources/cc/test_call_global.cc) add_test_object(sources/cc/test_types.cc) add_test_object(sources/cc/test_link_order_base.cc) add_test_object(sources/cc/test_link_order_caller.cc) add_test_object(sources/cc/test_error.cc) endif () # Pure C object files — built on all platforms (no C++ runtime deps) enablelanguage(C) add_test_object(sources/c/test_funcs.c) add_test_object(sources/c/test_funcs2.c) add_test_object(sources/c/test_funcs_conflict.c) add_test_object(sources/c/test_call_global.c) add_test_object(sources/c/test_types.c) add_test_object(sources/c/test_link_order_base.c) add_test_object(sources/c/test_link_order_caller.c) add_test_object(sources/c/test_error.c) add_test_object(sources/c/test_ctor_dtor.c) # CUDA object files — optional find_package(CUDAToolkit) if (CUDAToolkit_FOUND) enablelanguage(CUDA) message(STATUS "CUDA found: ${CUDAToolkit_VERSION}") add_test_object(sources/cuda/test_funcs.cu) endif () tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/README.md000066400000000000000000000103751521067262500220540ustar00rootroot00000000000000 # TVM-FFI-OrcJIT Tests ## Quick Start Run pytest directly — test objects are auto-built by `utils.build_test_objects()` into a temporary directory using `tvm_ffi.cpp.build`: ```bash pytest tests/ -v ``` Run the quick-start example: ```bash python examples/quick-start/run.py --lang c ``` ### Alternative: cmake-based build A `CMakeLists.txt` is also provided for cmake-based workflows. See `TEST_OBJ_INSTALL_SUFFIX` for multi-compiler builds. ## Tested Configurations Tests are parametrized over compiler variants. Each platform builds objects with multiple compilers; `test_basic.py` auto-discovers which are available. | Platform | Variant | Subdir | Compiler | Languages | | -------- | ------- | ------ | -------- | --------- | | Linux, macOS | default | `c/`, `cc/` | LLVM Clang | C, C++ | | Linux | gcc | `c-gcc/`, `cc-gcc/` | GCC | C, C++ | | macOS | appleclang | `c-appleclang/`, `cc-appleclang/` | Apple Clang (`/usr/bin/clang`) | C, C++ | | Windows | default | `c/` | LLVM Clang | C only | | Windows | msvc | `c-msvc/` | MSVC (`cl`) | C only | | Windows | clang-cl | `c-clang-cl/` | clang-cl | C only | All Windows variants are C-only. C++ objects use `try`/`catch` (via `TVM_FFI_SAFE_CALL_BEGIN/END`) which requires Itanium exception ABI symbols that the MSVC-built host process cannot resolve. ### Platform-specific compiler flags | Flag | Applies to | Reason | | ---- | ---------- | ------ | | `/O2 /GS-` | MSVC, clang-cl | `/GS-` disables buffer security checks (`__security_cookie`) which are CRT symbols the ORC JIT cannot resolve | | `-O2` | All others | Standard optimization | | `-mno-outline-atomics` | aarch64 | Avoids libgcc outline atomics helper calls that the JIT cannot resolve | ## Test Structure ```text tests/ sources/ c/ C source files (.c) cc/ C++ source files (.cc) cuda/ CUDA source files (optional) utils.py Build helpers (compile sources to objects in /tmp) test_basic.py Python test cases CMakeLists.txt Alternative cmake build ``` Built objects are placed in `{tempdir}/tvm_ffi_orcjit_tests/` with per-compiler variant subdirectories (c/, cc/, c-gcc/, etc.). ### Source files | Source | Description | | ------ | ----------- | | `test_funcs` | Basic arithmetic (add, multiply) | | `test_funcs2` | More arithmetic (subtract, divide) | | `test_funcs_conflict` | Symbol conflict testing (duplicate `add`) | | `test_call_global` | Callbacks into Python-registered global functions | | `test_types` | Zero-arg, multi-arg, float, void return types | | `test_link_order_base` / `test_link_order_caller` | Cross-library symbol resolution | | `test_error` | Error propagation from JIT'd code | | `test_ctor_dtor` | Constructor/destructor and init/fini sections | ### Constructor/destructor test coverage `test_ctor_dtor` verifies that the ORC JIT correctly handles platform-specific static initialization and finalization: | Platform | Init mechanism | Fini mechanism | | -------- | -------------- | -------------- | | Linux (ELF) | `.init_array` (with priorities), `.ctors` (reversed) | `.dtors` | | macOS (Mach-O) | `.init_array` (with priorities), `__mod_init_func` | `.fini_array` (with priorities) | | Windows (COFF) | `.CRT$XCA`..`XCU` (alphabetical order) | `.CRT$XTA`..`XTZ` (alphabetical order) | ## CI The CI workflow (`.github/workflows/tvm_ffi_orcjit.yml`) runs `pytest` and the quick-start example via `cibuildwheel`. tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/000077500000000000000000000000001521067262500222525ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/000077500000000000000000000000001521067262500224745ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/fake_fatbin.c000066400000000000000000000052651521067262500251010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Simulates an NVCC-compiled object with a large .nv_fatbin device blob. * The fatbin data is referenced only by absolute relocations (R_*_64 / * R_AARCH64_ABS64), never by PC-relative relocations. This lets us test * overflow-region classification without needing a real CUDA toolchain. * * KEY DETAIL: References go through a pointer in .data (generates * R_AARCH64_ABS64 / R_X86_64_64), not via ADRP/RIP-relative. This * mirrors real NVCC output where __NV_fatbin_* uses absolute relocations. */ #include #include #ifdef __APPLE__ __attribute__((section("__DATA,.nv_fatbin"), used)) #else __attribute__((section(".nv_fatbin"), used)) #endif static const uint8_t fake_fatbin_data[4 * 1024 * 1024] = {0}; /* Indirect reference: .data holds an absolute-relocation pointer to .nv_fatbin. Code accesses .data via PC-relative (ADRP / RIP), and .data→.nv_fatbin is absolute. No PC-relative edge crosses from any section to .nv_fatbin, matching real NVCC objects. */ static const void* const fatbin_ptr = fake_fatbin_data; static const uint64_t fatbin_size = sizeof(fake_fatbin_data); /* get_fatbin_size: return the size of the fake fatbin blob. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_get_fatbin_size(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)fatbin_size; return 0; } /* get_fatbin_addr: return the address of the fake fatbin data. Used by tests to verify overflow sections land outside the arena. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_get_fatbin_addr(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)(uintptr_t)fatbin_ptr; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_addr.c000066400000000000000000000024601521067262500246130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Returns the code address of this function — for arena co-location tests. * Load into multiple libraries to verify they land in the same arena region. */ #include #include TVM_FFI_DLL_EXPORT int __tvm_ffi_code_address(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)(uintptr_t)&__tvm_ffi_code_address; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_call_global.c000066400000000000000000000072311521067262500261350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C test: look up a host-registered global function by name and call it. * Demonstrates JIT code calling back into the host process via the TVM FFI C API. */ #include /* * test_call_global_add: look up "test_host_add" global function and call it * with two integer arguments. Returns the result from the host function. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_call_global_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { int ret_code = 0; TVMFFIByteArray func_name; TVMFFIObjectHandle func_handle = NULL; TVMFFIAny call_args[2]; TVMFFIAny call_result; /* Look up the host-registered global function */ func_name.data = "test_host_add"; func_name.size = 13; ret_code = TVMFFIFunctionGetGlobal(&func_name, &func_handle); if (ret_code != 0) return ret_code; /* Prepare arguments: pass through the two integer args */ call_args[0].type_index = kTVMFFIInt; call_args[0].zero_padding = 0; call_args[0].v_int64 = args[0].v_int64; call_args[1].type_index = kTVMFFIInt; call_args[1].zero_padding = 0; call_args[1].v_int64 = args[1].v_int64; /* Call the global function */ call_result.type_index = kTVMFFINone; call_result.zero_padding = 0; call_result.v_int64 = 0; ret_code = TVMFFIFunctionCall(func_handle, call_args, 2, &call_result); /* Release the function handle */ TVMFFIObjectDecRef((TVMFFIObject*)func_handle); if (ret_code != 0) return ret_code; /* Forward the result */ result->type_index = call_result.type_index; result->zero_padding = call_result.zero_padding; result->v_int64 = call_result.v_int64; return 0; } /* * test_call_global_mul: look up "test_host_multiply" global function and call it. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_call_global_mul(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { int ret_code = 0; TVMFFIByteArray func_name; TVMFFIObjectHandle func_handle = NULL; TVMFFIAny call_args[2]; TVMFFIAny call_result; func_name.data = "test_host_multiply"; func_name.size = 18; ret_code = TVMFFIFunctionGetGlobal(&func_name, &func_handle); if (ret_code != 0) return ret_code; call_args[0].type_index = kTVMFFIInt; call_args[0].zero_padding = 0; call_args[0].v_int64 = args[0].v_int64; call_args[1].type_index = kTVMFFIInt; call_args[1].zero_padding = 0; call_args[1].v_int64 = args[1].v_int64; call_result.type_index = kTVMFFINone; call_result.zero_padding = 0; call_result.v_int64 = 0; ret_code = TVMFFIFunctionCall(func_handle, call_args, 2, &call_result); TVMFFIObjectDecRef((TVMFFIObject*)func_handle); if (ret_code != 0) return ret_code; result->type_index = call_result.type_index; result->zero_padding = call_result.zero_padding; result->v_int64 = call_result.v_int64; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_ctor_dtor.c000066400000000000000000000144341521067262500257040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C version of the constructor/destructor test. * Uses TVMFFIFunctionGetGlobal + TVMFFIFunctionCall to call the host * "append_log" function from constructors and destructors. * * Platform-specific init/deinit mechanisms tested: * ELF → .init_array/.fini_array (via __attribute__) + .ctors/.dtors * Mach-O → __DATA,__mod_init_func (via __attribute__) + explicit section * COFF → .CRT$XC* sections (init) + .CRT$XT* sections (term) */ #include static void puts_log(const char* msg) { TVMFFIByteArray func_name; TVMFFIObjectHandle func_handle = NULL; TVMFFIAny call_args[1]; TVMFFIAny call_result; func_name.data = "append_log"; func_name.size = 10; if (TVMFFIFunctionGetGlobal(&func_name, &func_handle) != 0) return; call_args[0].type_index = kTVMFFIRawStr; call_args[0].zero_padding = 0; call_args[0].v_c_str = msg; call_result.type_index = kTVMFFINone; call_result.zero_padding = 0; call_result.v_int64 = 0; TVMFFIFunctionCall(func_handle, call_args, 1, &call_result); TVMFFIObjectDecRef((TVMFFIObject*)func_handle); } typedef void (*ctor_t)(void); typedef void (*dtor_t)(void); /* ========================================================================= * GCC / Clang (ELF + Mach-O): __attribute__((constructor/destructor)) * ========================================================================= */ #ifndef _MSC_VER __attribute__((constructor)) void init_array(void) { puts_log(""); } __attribute__((constructor(101))) void init_array_101(void) { puts_log(""); } __attribute__((constructor(102))) void init_array_102(void) { puts_log(""); } __attribute__((constructor(103))) void init_array_103(void) { puts_log(""); } __attribute__((destructor)) void fini_array(void) { puts_log(""); } __attribute__((destructor(101))) void fini_array_101(void) { puts_log(""); } __attribute__((destructor(102))) void fini_array_102(void) { puts_log(""); } __attribute__((destructor(103))) void fini_array_103(void) { puts_log(""); } /* ELF-specific: explicit .ctors/.dtors section placements with priorities. * .ctors priorities are reversed: lower number = later execution. */ #ifdef __ELF__ static void ctors(void) { puts_log(""); } __attribute__((section(".ctors"), used)) static ctor_t ctors_ptr = ctors; static void ctors_101(void) { puts_log(""); } __attribute__((section(".ctors.101"), used)) static ctor_t ctors_1_ptr = ctors_101; static void ctors_102(void) { puts_log(""); } __attribute__((section(".ctors.102"), used)) static ctor_t ctors_2_ptr = ctors_102; static void ctors_103(void) { puts_log(""); } __attribute__((section(".ctors.103"), used)) static ctor_t ctors_3_ptr = ctors_103; static void dtors(void) { puts_log(""); } __attribute__((section(".dtors"), used)) static dtor_t dtors_ptr = dtors; static void dtors_101(void) { puts_log(""); } __attribute__((section(".dtors.101"), used)) static dtor_t dtors_1_ptr = dtors_101; static void dtors_102(void) { puts_log(""); } __attribute__((section(".dtors.102"), used)) static dtor_t dtors_2_ptr = dtors_102; static void dtors_103(void) { puts_log(""); } __attribute__((section(".dtors.103"), used)) static dtor_t dtors_3_ptr = dtors_103; #endif /* __ELF__ */ /* Mach-O-specific: explicit __DATA,__mod_init_func section placement. */ #ifdef __APPLE__ static void mod_init_func(void) { puts_log(""); } __attribute__((section("__DATA,__mod_init_func"), used)) static ctor_t mod_init_ptr = mod_init_func; #endif /* __APPLE__ */ #endif /* !_MSC_VER */ /* ========================================================================= * MSVC (COFF): .CRT$XC* sections + atexit() * Subsections run in alphabetical order: XCA < XCB < XCC < XCU. * * Variables in these sections MUST NOT be static: both MSVC and clang-cl * may optimize away unreferenced static variables even with __declspec(allocate). * External linkage + unique names ensures the section data is emitted. * ========================================================================= */ #ifdef _MSC_VER #pragma section(".CRT$XCA", read) #pragma section(".CRT$XCB", read) #pragma section(".CRT$XCC", read) #pragma section(".CRT$XCU", read) #pragma section(".CRT$XTA", read) #pragma section(".CRT$XTZ", read) static void __cdecl crt_init_a(void) { puts_log(""); } __declspec(allocate(".CRT$XCA")) ctor_t __tvm_test_crt_init_a = crt_init_a; static void __cdecl crt_init_b(void) { puts_log(""); } __declspec(allocate(".CRT$XCB")) ctor_t __tvm_test_crt_init_b = crt_init_b; static void __cdecl crt_init_c(void) { puts_log(""); } __declspec(allocate(".CRT$XCC")) ctor_t __tvm_test_crt_init_c = crt_init_c; static void __cdecl crt_init_u(void) { puts_log(""); } __declspec(allocate(".CRT$XCU")) ctor_t __tvm_test_crt_init_u = crt_init_u; static void __cdecl crt_term_a(void) { puts_log(""); } __declspec(allocate(".CRT$XTA")) dtor_t __tvm_test_crt_term_a = crt_term_a; static void __cdecl crt_term_z(void) { puts_log(""); } __declspec(allocate(".CRT$XTZ")) dtor_t __tvm_test_crt_term_z = crt_term_z; #endif /* _MSC_VER */ /* main: callable entry point */ TVM_FFI_DLL_EXPORT int __tvm_ffi_main(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { puts_log("
"); result->type_index = kTVMFFINone; result->zero_padding = 0; result->v_int64 = 0; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_error.c000066400000000000000000000023551521067262500250350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Error propagation test: JIT function signals error via * TVMFFIErrorSetRaisedFromCStr, which should surface as a Python exception. */ #include /* test_error: always raises a ValueError */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_error(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { TVMFFIErrorSetRaisedFromCStr("ValueError", "test error"); return -1; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_funcs.c000066400000000000000000000032741521067262500250230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C test functions using the TVMFFISafeCallType ABI directly. * No C++ features: no exceptions, RTTI, static variables, or std library. * This isolates the JIT infrastructure from MSVC C++ runtime issues. */ #include /* test_add: add two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64; return 0; } /* test_multiply: multiply two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_multiply(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 * args[1].v_int64; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_funcs2.c000066400000000000000000000032121521067262500250750ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C test functions (subtract, divide) using the TVMFFISafeCallType ABI directly. * C version of test_funcs2.cc — no C++ runtime dependencies. */ #include /* test_subtract: subtract two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_subtract(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 - args[1].v_int64; return 0; } /* test_divide: divide two integers */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_divide(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 / args[1].v_int64; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_funcs_conflict.c000066400000000000000000000034031521067262500266760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C conflicting test functions using the TVMFFISafeCallType ABI directly. * C version of test_funcs_conflict.cc — same symbol names as test_funcs.c * but different implementations to test symbol conflict handling. */ #include /* test_add: conflicting add — adds 1000 to the result */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64 + 1000; return 0; } /* test_multiply: conflicting multiply — doubles the result */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_multiply(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 * args[1].v_int64 * 2; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_hidden_caller.c000066400000000000000000000050731521067262500264610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Caller library for ADRP overflow test (LLVM issue #173269). * * Takes the ADDRESS of hidden_helper_add via ADRP+ADD (no GOT, * because of hidden visibility). When this object and * test_hidden_helper.o are in different mmap allocations >4GB * apart, the ADRP immediate overflows — silent truncation on * AArch64 causes a segfault. * * The arena memory manager fixes this by placing all objects * in contiguous VA space (<< 4GB). */ #include #include /* Same hidden declaration — compiler uses ADRP+ADD to take address */ __attribute__((visibility("hidden"))) extern int64_t hidden_helper_add(int64_t a, int64_t b); typedef int64_t (*binop_t)(int64_t, int64_t); /* call_hidden_add: take address of hidden_helper_add, then call via pointer. On AArch64, generates: ADRP x0, hidden_helper_add@PAGE (R_AARCH64_ADR_PREL_PG_HI21, ±4GB) ADD x0, x0, hidden_helper_add@PAGEOFF (R_AARCH64_ADD_ABS_LO12_NC) When hidden_helper_add is in a different allocation >4GB away, ADRP overflows. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_call_hidden_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { volatile binop_t fn = &hidden_helper_add; result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = fn(args[0].v_int64, args[1].v_int64); return 0; } /* Return the address of this function's code — for co-location tests */ TVM_FFI_DLL_EXPORT int __tvm_ffi_caller_code_address(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)(uintptr_t)&__tvm_ffi_caller_code_address; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_hidden_helper.c000066400000000000000000000041411521067262500264710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Helper library for ADRP overflow test. * Defines a hidden-visibility function whose ADDRESS is taken * by test_hidden_caller.c. On AArch64, the caller uses * ADRP+ADD (no GOT) to compute the address — this overflows * when the two objects are in different allocations >4GB apart. * * Reference: LLVM issue #173269 */ #include #include /* Hidden visibility: caller uses ADRP+ADD instead of GOT */ __attribute__((visibility("hidden"))) int64_t hidden_helper_add(int64_t a, int64_t b) { return a + b; } /* Export a TVM FFI function that calls hidden_helper_add directly */ TVM_FFI_DLL_EXPORT int __tvm_ffi_hidden_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = hidden_helper_add(args[0].v_int64, args[1].v_int64); return 0; } /* Return the address of this function's code — for co-location tests */ TVM_FFI_DLL_EXPORT int __tvm_ffi_helper_code_address(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)(uintptr_t)&__tvm_ffi_helper_code_address; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_base.c000066400000000000000000000024671521067262500270320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Base library for cross-library linking test. * Exports helper_add which is called by test_link_order_caller.c. */ #include /* helper_add: add two integers — used as a dependency by the caller library */ TVM_FFI_DLL_EXPORT int __tvm_ffi_helper_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_caller.c000066400000000000000000000027551521067262500273620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Caller library for cross-library linking test. * References __tvm_ffi_helper_add from test_link_order_base.c via extern * declaration, and exports cross_lib_add which forwards to it. */ #include /* Declare external symbol from the base library */ extern int __tvm_ffi_helper_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result); /* cross_lib_add: forwards to helper_add in the base library */ TVM_FFI_DLL_EXPORT int __tvm_ffi_cross_lib_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { return __tvm_ffi_helper_add(self, args, num_args, result); } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/c/test_types.c000066400000000000000000000045371521067262500250540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Pure C test functions exercising type variety: zero-arg, four-arg, * float, and void return types via the TVMFFISafeCallType ABI. */ #include /* test_zero_arg: ignores arguments, returns integer 42 */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_zero_arg(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = 42; return 0; } /* test_four_args: sum of four integer arguments */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_four_args(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64 + args[2].v_int64 + args[3].v_int64; return 0; } /* test_float_multiply: multiply two doubles */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_float_multiply(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIFloat; result->zero_padding = 0; result->v_float64 = args[0].v_float64 * args[1].v_float64; return 0; } /* test_void_function: does nothing, returns None */ TVM_FFI_DLL_EXPORT int __tvm_ffi_test_void_function(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFINone; result->zero_padding = 0; result->v_int64 = 0; return 0; } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/000077500000000000000000000000001521067262500226375ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/fake_fatbin.cc000066400000000000000000000053021521067262500253770ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Simulates an NVCC-compiled object with a large .nv_fatbin device blob. // The fatbin data is referenced only by absolute relocations (R_*_64 / // R_AARCH64_ABS64), never by PC-relative relocations. This lets us test // overflow-region classification without needing a real CUDA toolchain. // // KEY DETAIL: References go through a pointer in .data (generates // R_AARCH64_ABS64 / R_X86_64_64), not via ADRP/RIP-relative. This // mirrors real NVCC output where __NV_fatbin_* uses absolute relocations. #include #include #ifdef __APPLE__ __attribute__((section("__DATA,.nv_fatbin"), used)) #else __attribute__((section(".nv_fatbin"), used)) #endif static const uint8_t fake_fatbin_data[4 * 1024 * 1024] = {0}; // Indirect reference: .data holds an absolute-relocation pointer to // .nv_fatbin. Code accesses .data via PC-relative (ADRP / RIP), and // .data→.nv_fatbin is absolute. No PC-relative edge crosses from any // section to .nv_fatbin, matching real NVCC objects. static const void* const fatbin_ptr = fake_fatbin_data; static const uint64_t fatbin_size = sizeof(fake_fatbin_data); // get_fatbin_size: return the size of the fake fatbin blob. extern "C" { TVM_FFI_DLL_EXPORT int __tvm_ffi_get_fatbin_size(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = static_cast(fatbin_size); return 0; } // get_fatbin_addr: return the address of the fake fatbin data. // Used by tests to verify overflow sections land outside the arena. TVM_FFI_DLL_EXPORT int __tvm_ffi_get_fatbin_addr(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = reinterpret_cast(fatbin_ptr); return 0; } } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_addr.cc000066400000000000000000000022101521067262500251120ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Returns the code address of this function — for arena co-location tests. // Load into multiple libraries to verify they land in the same arena region. #include #include int64_t code_address_impl() { return reinterpret_cast(&code_address_impl); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(code_address, code_address_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_call_global.cc000066400000000000000000000031461521067262500264440ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // C++ test: look up a host-registered global function by name and call it. // Demonstrates JIT code calling back into the host process via the TVM FFI C++ API. #include using namespace tvm::ffi; // Look up "test_host_add" global function and call it with two integer arguments. int test_call_global_add_impl(int a, int b) { Function host_add = Function::GetGlobalRequired("test_host_add"); return host_add(a, b).cast(); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_call_global_add, test_call_global_add_impl); // Look up "test_host_multiply" global function and call it. int test_call_global_mul_impl(int a, int b) { Function host_mul = Function::GetGlobalRequired("test_host_multiply"); return host_mul(a, b).cast(); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_call_global_mul, test_call_global_mul_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_ctor_dtor.cc000066400000000000000000000067111521067262500262110ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include using namespace tvm::ffi; #define PUTS_LOG(msg) \ auto append_log_func = Function::GetGlobalRequired("append_log"); \ append_log_func(msg); using ctor_t = void (*)(); using dtor_t = void (*)(); // __attribute__((constructor/destructor)) works on both ELF and Mach-O: // ELF → .init_array / .fini_array // Mach-O → __DATA,__mod_init_func / __cxa_atexit __attribute__((constructor)) void init_array() { PUTS_LOG(""); } __attribute__((constructor(101))) void init_array_101() { PUTS_LOG(""); } __attribute__((constructor(102))) void init_array_102() { PUTS_LOG(""); } __attribute__((constructor(103))) void init_array_103() { PUTS_LOG(""); } __attribute__((destructor)) void fini_array() { PUTS_LOG(""); } __attribute__((destructor(101))) void fini_array_101() { PUTS_LOG(""); } __attribute__((destructor(102))) void fini_array_102() { PUTS_LOG(""); } __attribute__((destructor(103))) void fini_array_103() { PUTS_LOG(""); } // ELF-specific: explicit .ctors/.dtors section placements with priorities. // These sections don't exist on Mach-O or COFF. #ifdef __ELF__ static void ctors() { PUTS_LOG(""); } __attribute__((section(".ctors"), used)) static ctor_t ctors_ptr = ctors; static void ctors_101() { PUTS_LOG(""); } __attribute__((section(".ctors.101"), used)) static ctor_t ctors_1_ptr = ctors_101; static void ctors_102() { PUTS_LOG(""); } __attribute__((section(".ctors.102"), used)) static ctor_t ctors_2_ptr = ctors_102; static void ctors_103() { PUTS_LOG(""); } __attribute__((section(".ctors.103"), used)) static ctor_t ctors_3_ptr = ctors_103; static void dtors() { PUTS_LOG(""); } __attribute__((section(".dtors"), used)) static dtor_t dtors_ptr = dtors; static void dtors_101() { PUTS_LOG(""); } __attribute__((section(".dtors.101"), used)) static dtor_t dtors_1_ptr = dtors_101; static void dtors_102() { PUTS_LOG(""); } __attribute__((section(".dtors.102"), used)) static dtor_t dtors_2_ptr = dtors_102; static void dtors_103() { PUTS_LOG(""); } __attribute__((section(".dtors.103"), used)) static dtor_t dtors_3_ptr = dtors_103; #endif // __ELF__ // Mach-O-specific: explicit __DATA,__mod_init_func section placement. #ifdef __APPLE__ static void mod_init_func() { PUTS_LOG(""); } __attribute__((section("__DATA,__mod_init_func"), used)) static ctor_t mod_init_ptr = mod_init_func; #endif // __APPLE__ void main_impl() { PUTS_LOG("
"); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(main, main_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_error.cc000066400000000000000000000021421521067262500253350ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Error propagation test (C++ version): throws std::runtime_error // which TVM_FFI_SAFE_CALL_END catches and converts to an InternalError. #include #include int test_error_impl() { throw std::runtime_error("test error"); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_error, test_error_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_funcs.cc000066400000000000000000000021341521067262500253230ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include // Simple addition function int test_add_impl(int a, int b) { return a + b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_impl); // Multiplication function int test_multiply_impl(int a, int b) { return a * b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_funcs2.cc000066400000000000000000000021331521067262500254040ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include // Subtraction function int test_subtract_impl(int a, int b) { return a - b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_subtract, test_subtract_impl); // Division function int test_divide_impl(int a, int b) { return a / b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_divide, test_divide_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_funcs_conflict.cc000066400000000000000000000023211521067262500272020ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include // Conflicting test_add function - different implementation int test_add_conflict_impl(int a, int b) { return a + b + 1000; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_conflict_impl); // Conflicting test_multiply function - different implementation int test_multiply_conflict_impl(int a, int b) { return a * b * 2; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_conflict_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_hidden_caller.cc000066400000000000000000000042431521067262500267650ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Caller library for ADRP overflow test (LLVM issue #173269). // // Takes the ADDRESS of hidden_helper_add via ADRP+ADD (no GOT, // because of hidden visibility). When this object and // test_hidden_helper.o are in different mmap allocations >4GB // apart, the ADRP immediate overflows — silent truncation on // AArch64 causes a segfault. // // The arena memory manager fixes this by placing all objects // in contiguous VA space (<< 4GB). #include #include // Same hidden declaration — compiler uses ADRP+ADD to take address __attribute__((visibility("hidden"))) extern int64_t hidden_helper_add(int64_t a, int64_t b); using binop_t = int64_t (*)(int64_t, int64_t); // call_hidden_add: take address of hidden_helper_add, then call via pointer. // On AArch64, generates: // ADRP x0, hidden_helper_add@PAGE (R_AARCH64_ADR_PREL_PG_HI21, ±4GB) // ADD x0, x0, hidden_helper_add@PAGEOFF (R_AARCH64_ADD_ABS_LO12_NC) // When hidden_helper_add is in a different allocation >4GB away, ADRP overflows. extern "C" { TVM_FFI_DLL_EXPORT int __tvm_ffi_call_hidden_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { volatile binop_t fn = &hidden_helper_add; result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = fn(args[0].v_int64, args[1].v_int64); return 0; } } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_hidden_helper.cc000066400000000000000000000033111521067262500267750ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Helper library for ADRP overflow test. // Defines a hidden-visibility function whose ADDRESS is taken // by test_hidden_caller.cc. On AArch64, the caller uses // ADRP+ADD (no GOT) to compute the address — this overflows // when the two objects are in different allocations >4GB apart. // // Reference: LLVM issue #173269 #include #include // Hidden visibility: caller uses ADRP+ADD instead of GOT __attribute__((visibility("hidden"))) int64_t hidden_helper_add(int64_t a, int64_t b) { return a + b; } // Export a TVM FFI function that calls hidden_helper_add directly extern "C" { TVM_FFI_DLL_EXPORT int __tvm_ffi_hidden_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = hidden_helper_add(args[0].v_int64, args[1].v_int64); return 0; } } tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_base.cc000066400000000000000000000020641521067262500273310ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Base library for cross-library linking test (C++ version). // Exports helper_add which is called by test_link_order_caller.cc. #include int helper_add_impl(int a, int b) { return a + b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(helper_add, helper_add_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_caller.cc000066400000000000000000000033321521067262500276600ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // Caller library for cross-library linking test (C++ version). // References __tvm_ffi_helper_add from test_link_order_base.cc and // exports cross_lib_add which forwards to it. #include #include extern "C" int __tvm_ffi_helper_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result); int cross_lib_add_impl(int a, int b) { TVMFFIAny call_args[2]; TVMFFIAny call_result; call_args[0].type_index = kTVMFFIInt; call_args[0].zero_padding = 0; call_args[0].v_int64 = a; call_args[1].type_index = kTVMFFIInt; call_args[1].zero_padding = 0; call_args[1].v_int64 = b; call_result.type_index = kTVMFFINone; call_result.zero_padding = 0; call_result.v_int64 = 0; __tvm_ffi_helper_add(nullptr, call_args, 2, &call_result); return static_cast(call_result.v_int64); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(cross_lib_add, cross_lib_add_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cc/test_types.cc000066400000000000000000000026471521067262500253620ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // C++ test functions exercising type variety: zero-arg, four-arg, // float, and void return types. #include int test_zero_arg_impl() { return 42; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_zero_arg, test_zero_arg_impl); int test_four_args_impl(int a, int b, int c, int d) { return a + b + c + d; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_four_args, test_four_args_impl); double test_float_multiply_impl(double a, double b) { return a * b; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_float_multiply, test_float_multiply_impl); void test_void_function_impl() {} TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_void_function, test_void_function_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cuda/000077500000000000000000000000001521067262500231665ustar00rootroot00000000000000tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/sources/cuda/test_funcs.cu000066400000000000000000000042701521067262500256770ustar00rootroot00000000000000// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #include #include // Simple addition function __global__ void test_add_kernel(int* a, int* b, int* c) { *c = *a + *b; } int test_add_impl(int a, int b) { int c; int *d_a, *d_b, *d_c; cudaMalloc(&d_a, sizeof(int)); cudaMalloc(&d_b, sizeof(int)); cudaMalloc(&d_c, sizeof(int)); cudaMemcpy(d_a, &a, sizeof(int), cudaMemcpyHostToDevice); cudaMemcpy(d_b, &b, sizeof(int), cudaMemcpyHostToDevice); test_add_kernel<<<1, 1>>>(d_a, d_b, d_c); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) printf("Kernel launch error: %s\n", cudaGetErrorString(err)); cudaMemcpy(&c, d_c, sizeof(int), cudaMemcpyDeviceToHost); cudaFree(d_a); cudaFree(d_b); cudaFree(d_c); return c; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_impl); // Multiplication function __global__ void test_multiply_kernel(int* a, int* b, int* c) { *c = *a * *b; } int test_multiply_impl(int a, int b) { int c; int *d_a, *d_b, *d_c; cudaMalloc(&d_a, sizeof(int)); cudaMalloc(&d_b, sizeof(int)); cudaMalloc(&d_c, sizeof(int)); cudaMemcpy(d_a, &a, sizeof(int), cudaMemcpyHostToDevice); cudaMemcpy(d_b, &b, sizeof(int), cudaMemcpyHostToDevice); test_multiply_kernel<<<1, 1>>>(d_a, d_b, d_c); cudaMemcpy(&c, d_c, sizeof(int), cudaMemcpyDeviceToHost); cudaFree(d_a); cudaFree(d_b); cudaFree(d_c); return c; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_impl); tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/test_basic.py000066400000000000000000001027541521067262500232720ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Basic tests for tvm_ffi_orcjit functionality.""" from __future__ import annotations import gc import sys import tempfile from pathlib import Path import pytest import tvm_ffi from tvm_ffi_orcjit import ExecutionSession from tvm_ffi_orcjit.dylib import DynamicLibrary from utils import build_test_objects # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- OBJ_DIR = build_test_objects() _KNOWN_SUBDIRS = [ "cc", "c", # default (LLVM clang) "cc-gcc", "c-gcc", # GCC (Linux) "cc-appleclang", "c-appleclang", # Apple Clang (macOS) "c-msvc", # MSVC (Windows, C only) "c-clang-cl", # clang-cl (Windows, C only) ] def obj(name: str) -> str: """Return path to a pre-built test object file, or skip if missing.""" path = OBJ_DIR / f"{name}.o" if not path.exists(): pytest.skip(f"{path.name} not found (not built)") return str(path) def make_lib( *obj_names: str, session: ExecutionSession | None = None, name: str = "" ) -> tuple[ExecutionSession, DynamicLibrary]: """Create a library and load one or more object files into it.""" if session is None: session = ExecutionSession() lib = session.create_library(name) for o in obj_names: lib.add(obj(o)) return session, lib # --------------------------------------------------------------------------- # Variants: C and C++ sources live in separate subdirectories (c/, cc/). # Function names are identical; only the object file path differs. # Tests are parametrized over both variants where the logic is identical. # --------------------------------------------------------------------------- class Variant: """Describes a C or C++ test variant (object file path mapping).""" def __init__(self, subdir: str) -> None: self.subdir = subdir # "cc" for C++, "c" for C def funcs_obj(self) -> str: """Return path prefix for test_funcs object.""" return f"{self.subdir}/test_funcs" def funcs2_obj(self) -> str: """Return path prefix for test_funcs2 object.""" return f"{self.subdir}/test_funcs2" def conflict_obj(self) -> str: """Return path prefix for test_funcs_conflict object.""" return f"{self.subdir}/test_funcs_conflict" def call_global_obj(self) -> str: """Return path prefix for test_call_global object.""" return f"{self.subdir}/test_call_global" def types_obj(self) -> str: """Return path prefix for test_types object.""" return f"{self.subdir}/test_types" def link_order_base_obj(self) -> str: """Return path prefix for test_link_order_base object.""" return f"{self.subdir}/test_link_order_base" def link_order_caller_obj(self) -> str: """Return path prefix for test_link_order_caller object.""" return f"{self.subdir}/test_link_order_caller" def error_obj(self) -> str: """Return path prefix for test_error object.""" return f"{self.subdir}/test_error" def ctor_dtor_obj(self) -> str: """Return path prefix for test_ctor_dtor object.""" return f"{self.subdir}/test_ctor_dtor" def fn(self, base_name: str) -> str: """Return the function name for a given base name.""" return base_name def __repr__(self) -> str: parts = self.subdir.split("-", 1) lang = "C++" if parts[0] == "cc" else "C" return f"{lang}-{parts[1]}" if len(parts) > 1 else lang def _discover_variants() -> list[Variant]: return [Variant(s) for s in _KNOWN_SUBDIRS if (OBJ_DIR / s / "test_funcs.o").exists()] _all_variants = _discover_variants() _cpp_only = [v for v in _all_variants if v.subdir.startswith("cc")] def _variant_id(v: Variant) -> str: return repr(v) # --------------------------------------------------------------------------- # Session / library creation # --------------------------------------------------------------------------- def test_create_session() -> None: """Create an ExecutionSession successfully.""" session = ExecutionSession() assert session is not None def test_create_library() -> None: """Create a DynamicLibrary from an ExecutionSession.""" session = ExecutionSession() lib = session.create_library() assert lib is not None def test_multiple_libraries() -> None: """Create multiple named libraries in one session.""" session = ExecutionSession() lib1 = session.create_library("lib1") lib2 = session.create_library("lib2") assert lib1 is not None assert lib2 is not None # --------------------------------------------------------------------------- # Load & execute — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_load_and_execute(v: Variant) -> None: """Load test_funcs and verify add/multiply.""" _, lib = make_lib(v.funcs_obj()) assert lib.get_function(v.fn("test_add"))(10, 20) == 30 assert lib.get_function(v.fn("test_multiply"))(7, 6) == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_load_and_execute_second_set(v: Variant) -> None: """Load test_funcs2 and verify subtract/divide.""" _, lib = make_lib(v.funcs2_obj()) assert lib.get_function(v.fn("test_subtract"))(10, 3) == 7 assert lib.get_function(v.fn("test_divide"))(20, 4) == 5 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_function_not_found(v: Variant) -> None: """Raise AttributeError for a missing function name.""" _, lib = make_lib(v.funcs_obj()) with pytest.raises(AttributeError, match="Module has no function"): lib.get_function("nonexistent_function") # --------------------------------------------------------------------------- # Multi-object / multi-library — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_gradually_add_objects(v: Variant) -> None: """Add multiple objects incrementally and verify all functions work.""" _session, lib = make_lib(v.funcs_obj()) add_func = lib.get_function(v.fn("test_add")) mul_func = lib.get_function(v.fn("test_multiply")) assert add_func(5, 3) == 8 assert mul_func(4, 5) == 20 lib.add(obj(v.funcs2_obj())) assert lib.get_function(v.fn("test_subtract"))(10, 3) == 7 assert lib.get_function(v.fn("test_divide"))(20, 4) == 5 # First object's functions still work assert add_func(10, 20) == 30 assert mul_func(7, 6) == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_two_separate_libraries(v: Variant) -> None: """Separate libraries expose only their own functions.""" session = ExecutionSession() _, lib1 = make_lib(v.funcs_obj(), session=session, name="lib1") _, lib2 = make_lib(v.funcs2_obj(), session=session, name="lib2") assert lib1.get_function(v.fn("test_add"))(5, 3) == 8 assert lib1.get_function(v.fn("test_multiply"))(4, 5) == 20 assert lib2.get_function(v.fn("test_subtract"))(10, 3) == 7 assert lib2.get_function(v.fn("test_divide"))(20, 4) == 5 with pytest.raises(AttributeError, match="Module has no function"): lib1.get_function(v.fn("test_subtract")) with pytest.raises(AttributeError, match="Module has no function"): lib2.get_function(v.fn("test_add")) # --------------------------------------------------------------------------- # Symbol conflicts — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_symbol_conflict_same_library(v: Variant) -> None: """Duplicate symbol in the same library raises an error.""" _, lib = make_lib(v.funcs_obj()) assert lib.get_function(v.fn("test_add"))(10, 20) == 30 with pytest.raises(Exception): lib.add(obj(v.conflict_obj())) @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_symbol_conflict_different_libraries(v: Variant) -> None: """Same symbol in different libraries resolves independently.""" session = ExecutionSession() _, lib1 = make_lib(v.funcs_obj(), session=session, name="lib1") _, lib2 = make_lib(v.conflict_obj(), session=session, name="lib2") assert lib1.get_function(v.fn("test_add"))(10, 20) == 30 assert lib2.get_function(v.fn("test_add"))(10, 20) == 1030 assert lib1.get_function(v.fn("test_multiply"))(5, 6) == 30 assert lib2.get_function(v.fn("test_multiply"))(5, 6) == 60 # --------------------------------------------------------------------------- # Global function callbacks — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.fixture() def _register_host_functions() -> None: """Register host add/multiply functions for JIT code to call.""" @tvm_ffi.register_global_func("test_host_add", override=True) def _host_add(a: int, b: int) -> int: return a + b @tvm_ffi.register_global_func("test_host_multiply", override=True) def _host_mul(a: int, b: int) -> int: return a * b @pytest.mark.usefixtures("_register_host_functions") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_call_global(v: Variant) -> None: """JIT code calls back into Python-registered global functions.""" _, lib = make_lib(v.call_global_obj()) add_func = lib.get_function(v.fn("test_call_global_add")) assert add_func(10, 20) == 30 assert add_func(100, 200) == 300 mul_func = lib.get_function(v.fn("test_call_global_mul")) assert mul_func(7, 6) == 42 assert mul_func(11, 11) == 121 # --------------------------------------------------------------------------- # Error handling — pure Python (Group 1) # --------------------------------------------------------------------------- def test_empty_library() -> None: """get_function on an empty library raises AttributeError.""" session = ExecutionSession() lib = session.create_library("empty") with pytest.raises(AttributeError, match="Module has no function"): lib.get_function("nonexistent") def test_invalid_object_file_path() -> None: """Adding a nonexistent object file raises an error.""" session = ExecutionSession() lib = session.create_library("bad_path") with pytest.raises(Exception): lib.add("/nonexistent_path/does_not_exist.o") def test_invalid_object_file_content() -> None: """Adding a file with garbage content raises an error.""" with tempfile.NamedTemporaryFile(suffix=".o", delete=False) as f: f.write(b"this is not a valid object file") f.flush() session = ExecutionSession() lib = session.create_library("bad_content") with pytest.raises(Exception): lib.add(f.name) def test_multiple_independent_sessions() -> None: """Two independent sessions don't interfere with each other.""" session1 = ExecutionSession() session2 = ExecutionSession() _, lib1 = make_lib(_all_variants[0].funcs_obj(), session=session1) _, lib2 = make_lib(_all_variants[0].funcs2_obj(), session=session2) assert lib1.get_function("test_add")(5, 3) == 8 assert lib2.get_function("test_subtract")(10, 3) == 7 # Each session's library doesn't see the other's functions with pytest.raises(AttributeError, match="Module has no function"): lib1.get_function("test_subtract") with pytest.raises(AttributeError, match="Module has no function"): lib2.get_function("test_add") # --------------------------------------------------------------------------- # Type variety — parametrized over C / C++ (Group 2) # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_zero_arg_function(v: Variant) -> None: """Zero-arg function returns constant 42.""" _, lib = make_lib(v.types_obj()) assert lib.get_function(v.fn("test_zero_arg"))() == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_four_arg_function(v: Variant) -> None: """Four integer arguments summed.""" _, lib = make_lib(v.types_obj()) assert lib.get_function(v.fn("test_four_args"))(1, 2, 3, 4) == 10 assert lib.get_function(v.fn("test_four_args"))(100, 200, 300, 400) == 1000 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_float_function(v: Variant) -> None: """Float multiply returns approximate result.""" _, lib = make_lib(v.types_obj()) result = lib.get_function(v.fn("test_float_multiply"))(3.14, 2.0) assert result == pytest.approx(6.28) @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_void_function(v: Variant) -> None: """Void function returns None.""" _, lib = make_lib(v.types_obj()) result = lib.get_function(v.fn("test_void_function"))() assert result is None # --------------------------------------------------------------------------- # Advanced — cross-library linking and error propagation (Group 3) # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_set_link_order(v: Variant) -> None: """Cross-library symbol resolution via set_link_order.""" session = ExecutionSession() # Base library exports helper_add lib_base = session.create_library("base") lib_base.add(obj(v.link_order_base_obj())) # Caller library references helper_add from base lib_caller = session.create_library("caller") lib_caller.set_link_order(lib_base) lib_caller.add(obj(v.link_order_caller_obj())) cross_add = lib_caller.get_function(v.fn("cross_lib_add")) assert cross_add(10, 20) == 30 assert cross_add(100, 200) == 300 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_error_propagation(v: Variant) -> None: """JIT function that signals an error raises a Python exception.""" _, lib = make_lib(v.error_obj()) with pytest.raises(Exception, match="test error"): lib.get_function(v.fn("test_error"))() # --------------------------------------------------------------------------- # CUDA (optional) # --------------------------------------------------------------------------- def test_load_and_execute_cuda_function() -> None: """Load and execute CUDA-compiled test objects.""" _, lib = make_lib("cuda/test_funcs") assert lib.get_function("test_add")(10, 20) == 30 assert lib.get_function("test_multiply")(7, 6) == 42 # --------------------------------------------------------------------------- # Constructor / destructor (parametrized over C / C++) # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_ctor_dtor(v: Variant) -> None: """Verify constructor/destructor ordering across platforms.""" log = "" @tvm_ffi.register_global_func("append_log", override=True) def _append_ctor_log(x: str) -> None: nonlocal log log += x _, lib = make_lib(v.ctor_dtor_obj()) lib.get_function(v.fn("main"))() del lib main_idx = log.index("
") pre = log[:main_idx] post = log[main_idx:] if sys.platform == "win32": # Windows (all compilers): COFF .CRT$XC* constructors + .CRT$XT* terminators. # All Windows compilers (MSVC, clang-cl, and LLVM Clang targeting MSVC ABI) # define _MSC_VER, so the source uses #pragma section / __declspec(allocate). assert "" in pre, f"CRT initializers not found in log: {log!r}" assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert "" in post, f"CRT terminators not found in log: {log!r}" assert post.index("") < post.index("") elif sys.platform == "linux": # ELF: init_array (priority order) + .ctors (reversed priority) assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert "" in post elif sys.platform == "darwin": # Mach-O: init_array (priority order) + explicit __mod_init_func assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert pre.index("") < pre.index("") assert "" in pre assert post.index("") < post.index("") assert post.index("") < post.index("") assert post.index("") < post.index("") assert "" not in log assert "" not in log # --------------------------------------------------------------------------- # Dylib removal — dropping a DynamicLibrary while its session is still alive. # # Dropping the last Python reference to a DynamicLibrary removes the JITDylib # from the ExecutionSession (releasing its JIT memory) in addition to running # any pending static destructors. These tests pin that contract and guard # against regressions in the slab-pool memory-manager refactor, which assumes # per-dylib deallocation actually happens at drop time. # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_empty_library(v: Variant) -> None: """Drop an empty library, then reuse the session.""" session = ExecutionSession() lib = session.create_library("empty") del lib # Session still works. lib2 = session.create_library("after") lib2.add(obj(v.funcs_obj())) assert lib2.get_function(v.fn("test_add"))(1, 2) == 3 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_loaded_library_then_recreate(v: Variant) -> None: """Drop a library with live JIT code, then create and use a fresh one.""" session = ExecutionSession() lib1 = session.create_library("lib1") lib1.add(obj(v.funcs_obj())) assert lib1.get_function(v.fn("test_add"))(3, 4) == 7 del lib1 lib2 = session.create_library("lib2") lib2.add(obj(v.funcs_obj())) assert lib2.get_function(v.fn("test_add"))(100, 1) == 101 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_repeated_create_drop_many_iterations(v: Variant) -> None: """Long create / load / call / drop cycle must not degrade or crash. Exercises the memory manager's recycled-region path: each iteration frees JIT memory that a later iteration may be handed back by the arena. """ session = ExecutionSession() for i in range(32): lib = session.create_library(f"iter_{i}") lib.add(obj(v.funcs_obj())) assert lib.get_function(v.fn("test_multiply"))(i + 1, 2) == (i + 1) * 2 del lib @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_one_library_does_not_affect_another(v: Variant) -> None: """Dropping one library must leave unrelated sibling libraries working.""" session = ExecutionSession() keep = session.create_library("keep") keep.add(obj(v.funcs_obj())) drop = session.create_library("drop") drop.add(obj(v.funcs2_obj())) assert drop.get_function(v.fn("test_subtract"))(10, 3) == 7 del drop gc.collect() # Untouched sibling still works; room now exists for a fresh library. assert keep.get_function(v.fn("test_add"))(5, 5) == 10 fresh = session.create_library("fresh") fresh.add(obj(v.funcs2_obj())) assert fresh.get_function(v.fn("test_divide"))(20, 4) == 5 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_captured_function_keeps_library_alive(v: Variant) -> None: """A captured Function keeps the library alive after the lib handle is dropped.""" session = ExecutionSession() lib = session.create_library("hold") lib.add(obj(v.funcs_obj())) captured = lib.get_function(v.fn("test_add")) del lib gc.collect() assert captured(11, 31) == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_runs_static_destructors(v: Variant) -> None: """Drop runs static destructors immediately rather than at session teardown.""" log: list[str] = [] @tvm_ffi.register_global_func("append_log", override=True) def _append(x: str) -> None: log.append(x) session = ExecutionSession() lib = session.create_library("ctor_dtor") lib.add(obj(v.ctor_dtor_obj())) lib.get_function(v.fn("main"))() ctor_len = len(log) assert "
" in log, f"main not observed: {log}" del lib gc.collect() post = log[ctor_len:] # Platform-specific fini markers: ELF (.fini_array / .dtors), Mach-O # (__mod_term_func surfaces as fini_array), COFF (.CRT$XT*). if sys.platform == "win32": markers = ("crt.XT",) else: markers = ("dtors", "fini_array") assert any(m in entry for entry in post for m in markers), ( f"no static destructors on drop (platform={sys.platform}): pre={log[:ctor_len]} post={post}" ) @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_caller_leaves_base_usable(v: Variant) -> None: """Dropping a set_link_order caller leaves its base library usable.""" session = ExecutionSession() base = session.create_library("base") base.add(obj(v.link_order_base_obj())) caller = session.create_library("caller") caller.set_link_order(base) caller.add(obj(v.link_order_caller_obj())) assert caller.get_function(v.fn("cross_lib_add"))(1, 2) == 3 del caller gc.collect() caller2 = session.create_library("caller2") caller2.set_link_order(base) caller2.add(obj(v.link_order_caller_obj())) assert caller2.get_function(v.fn("cross_lib_add"))(10, 20) == 30 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_drop_all_libraries_then_session(v: Variant) -> None: """Dropping every library before the session still leaves clean teardown.""" session = ExecutionSession() libs = [] for i in range(4): lib = session.create_library(f"lib_{i}") lib.add(obj(v.funcs_obj())) libs.append(lib) for lib in reversed(libs): del lib libs.clear() gc.collect() # Session destructor runs at end-of-test; a regression in removal would # surface as a crash under pytest. # --------------------------------------------------------------------------- # Slab-pool growth (Stage B). # # A session now holds a growable pool of Slabs, each `slab_size` bytes. # When a JITLink graph won't fit in any existing slab, the pool mmap's # a new one; graphs larger than a single slab go to a dedicated # oversize slab. These tests pin the behavioral contract: small # slab_size sessions must still link many libraries (by growing) and # must reuse drained bytes within a slab (via the free list). # --------------------------------------------------------------------------- # 8 MB is the practical floor — Slab::kCommitGranularity is 2 MB and the # dual-pool midpoint needs at least two commit chunks of headroom above it, # so smaller capacities break the pool layout. See SlabPoolMemoryManager # kMinSlabSize. _SMALL_SLAB = 8 * 1024 * 1024 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_pool_grows_under_small_slab(v: Variant) -> None: """Tight slab_size forces the pool to grow as more libraries load.""" session = ExecutionSession(slab_size=_SMALL_SLAB) libs = [] for i in range(16): lib = session.create_library(f"grow_{i}") lib.add(obj(v.funcs_obj())) assert lib.get_function(v.fn("test_add"))(i, i + 1) == 2 * i + 1 libs.append(lib) # All libraries remain callable after growth. Catches bugs where the # pool routes deallocate to the wrong Slab via FA->owner. for i, lib in enumerate(libs): assert lib.get_function(v.fn("test_multiply"))(i, 2) == i * 2 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_small_slab_recycles_after_drop(v: Variant) -> None: """Drop returns bytes to the slab's free list for subsequent reuse. If the free list weren't working, 32 iterations of load/drop under an 8 MB slab would either exhaust VA or force many new slabs; this test passing on CI containers is the recycling evidence. """ session = ExecutionSession(slab_size=_SMALL_SLAB) for i in range(32): lib = session.create_library(f"recycle_{i}") lib.add(obj(v.funcs_obj())) assert lib.get_function(v.fn("test_add"))(i, 1) == i + 1 del lib gc.collect() @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_pool_survives_mixed_load_drop_create(v: Variant) -> None: """Interleaved load / drop / create exercises growth + free-list together.""" session = ExecutionSession(slab_size=_SMALL_SLAB) # Ramp up to 4 concurrent libraries. live = [session.create_library(f"base_{i}") for i in range(4)] for lib in live: lib.add(obj(v.funcs_obj())) # Drop two, add three, verify remaining two still work, verify new # three work. Mixing drop/create in one session exercises slab # growth alongside free-list reuse from the dropped bytes. live[0] = None # type: ignore[assignment] live[2] = None # type: ignore[assignment] gc.collect() new_libs = [] for i in range(3): lib = session.create_library(f"after_drop_{i}") lib.add(obj(v.funcs2_obj())) new_libs.append(lib) assert live[1].get_function(v.fn("test_add"))(10, 5) == 15 assert live[3].get_function(v.fn("test_multiply"))(7, 6) == 42 for i, lib in enumerate(new_libs): assert lib.get_function(v.fn("test_subtract"))(i + 10, i) == 10 # --------------------------------------------------------------------------- # Manual slab reclamation via session.clear_free_slabs() (Stage C). # # After dropping a batch of libraries, the user can call clear_free_slabs() # to release any drained Slabs (zero live JIT allocations) back to the OS. # --------------------------------------------------------------------------- def _build_big_object(tmp_path: Path, byte_size: int, name: str = "big_object") -> str: """Compile a .c file with a large `.rodata` blob to force the oversize path. Returns the path to the generated .o file. Skips the test if the build toolchain is unavailable. ``name`` lets callers produce multiple distinctly-sized objects within the same ``tmp_path``. """ try: import tvm_ffi.cpp # noqa: PLC0415 except ImportError: pytest.skip("tvm_ffi.cpp not available for on-the-fly object build") src = tmp_path / f"{name}.c" # Global BSS array — volatile + a runtime read in big_probe() prevents # the compiler from folding sizeof() out and eliding the array. The # array becomes a ZeroFill segment of `byte_size` bytes in the # LinkGraph, which counts toward the slab footprint # (see Slab::computeGraphFootprint). src.write_text( f""" #include #include volatile unsigned char big_data[{byte_size}]; TVM_FFI_DLL_EXPORT int __tvm_ffi_big_probe(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) {{ (void)self; (void)args; (void)num_args; /* Reads big_data[0] (always zero for BSS) so the compiler must emit the array; returns the array's size for verification. */ result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = (int64_t)sizeof(big_data) + (int64_t)big_data[0]; return 0; }} """ ) try: return tvm_ffi.cpp.build( name=name, sources=[str(src)], output=f"{name}.o", extra_cflags=["-O2"], build_directory=str(tmp_path / f".build_{name}"), ) except Exception as exc: pytest.skip(f"could not build big object for reclaim test: {exc}") @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_no_drained() -> None: """Calling clear_free_slabs with nothing to reclaim returns 0.""" session = ExecutionSession() # Fresh session — the initial slab was never used, so not reclaimable. assert session.clear_free_slabs() == 0 # Load + keep a library: live allocations pin the slab. lib = session.create_library("alive") lib.add(obj(_all_variants[0].funcs_obj())) lib.get_function("test_add")(1, 2) assert session.clear_free_slabs() == 0 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_reclaims_oversize(tmp_path: Path) -> None: """Dropping an oversize-path library lets clear_free_slabs reclaim its slab. Uses a 3 MB rodata blob that exceeds the 4 MB slab's usable-per-slab threshold (~2 MB after midpoint slack), forcing the oversize path. The dedicated oversize slab then hosts exactly one graph — dropping it returns the slab to fully-drained state, so it is reclaimable. """ big_obj = _build_big_object(tmp_path, 3 * 1024 * 1024) # 4 MB slab → usable = slab_size / 2 = 2 MB → 3 MB blob forces oversize. session = ExecutionSession(slab_size=4 * 1024 * 1024) lib = session.create_library("oversize") lib.add(big_obj) assert lib.get_function("big_probe")() == 3 * 1024 * 1024 del lib gc.collect() # The oversize slab holds only the now-dropped lib's graph, so # clearFreeSlabs reclaims it. The initial (never-used) slab stays. reclaimed = session.clear_free_slabs() assert reclaimed >= 1, f"expected to reclaim the drained oversize slab, got {reclaimed}" @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_idempotent(tmp_path: Path) -> None: """A second call after everything has been reclaimed returns 0.""" big_obj = _build_big_object(tmp_path, 3 * 1024 * 1024) session = ExecutionSession(slab_size=4 * 1024 * 1024) lib = session.create_library("x") lib.add(big_obj) lib.get_function("big_probe")() del lib gc.collect() assert session.clear_free_slabs() >= 1 assert session.clear_free_slabs() == 0 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_preserves_live_pool(tmp_path: Path) -> None: """Reclaim runs only on drained slabs; live libraries keep working. Uses two different oversize blob sizes so the two libs land on separately-sized slabs: the 3 MB blob forces the pool to grow to 8 MB (next power of two covering a 4 MB pool's non-exec budget), while the 5 MB blob forces 16 MB. Dropping the 3 MB lib then leaves its 8 MB slab drained while the 16 MB slab stays live. """ small_obj = _build_big_object(tmp_path, 3 * 1024 * 1024, name="small_obj") large_obj = _build_big_object(tmp_path, 5 * 1024 * 1024, name="large_obj") session = ExecutionSession(slab_size=4 * 1024 * 1024) # Drop lib — lands on the 8 MB grown slab. drop_lib = session.create_library("drop") drop_lib.add(small_obj) drop_lib.get_function("big_probe")() del drop_lib # Keep lib — doesn't fit the 8 MB slab (5 MB > its ~4 MB non-exec # budget), so the pool grows to a 16 MB slab for this one. keep_lib = session.create_library("keep") keep_lib.add(large_obj) assert keep_lib.get_function("big_probe")() == 5 * 1024 * 1024 gc.collect() reclaimed = session.clear_free_slabs() assert reclaimed >= 1, f"expected drained slab to be reclaimed, got {reclaimed}" # The kept lib's slab was untouched; it still executes correctly. assert keep_lib.get_function("big_probe")() == 5 * 1024 * 1024 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_disabled_pool() -> None: """When the slab pool is disabled, clear_free_slabs is a no-op (returns 0).""" session = ExecutionSession(slab_size=-1) assert session.clear_free_slabs() == 0 tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/test_memory_manager.py000066400000000000000000000661341521067262500252140ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for JIT memory manager — verifies co-location and relocation safety. Background ---------- LLVM ORC JIT v2 uses ``InProcessMemoryMapper`` (backed by ``MapperJITLinkMemoryManager``) to allocate JIT memory. Each allocation is a separate ``mmap(MAP_ANONYMOUS)`` call whose address the kernel picks. Under virtual-address (VA) pressure — leaked slabs from failed materializations, long-running pytest sessions holding tracebacks, or simply a fragmented address space — the kernel can place successive allocations far apart. This matters for **PC-relative relocations with limited range**: - **x86_64 R_X86_64_PC32 / Delta32**: ±2 GB range. GCC-compiled C++ objects reference ``__dso_handle`` (used by ``__cxa_atexit`` for DSO identification) via PC32 when the symbol has hidden visibility. LLVM's ``ELFNixPlatform`` defines ``__dso_handle`` per JITDylib in a separate ``DSOHandleMaterializationUnit`` — a tiny ``LinkGraph`` allocated independently of the code that references it. If those two allocations land >2 GB apart, the Delta32 fixup overflows. - **AArch64 ADRP+ADD**: ±4 GB range. Hidden-visibility cross-object calls use ADRP (page-relative) which has the same scatter problem at a wider threshold. Our ``SlabPoolMemoryManager`` solves this by pre-reserving one or more contiguous VA slabs via ``mmap(PROT_NONE)`` and bump-allocating within them, guaranteeing allocations that land on the same slab stay within relocation range regardless of external VA pressure. These tests pin ``slab_size`` large enough that every graph they exercise fits on the initial slab, so the "co-located within one slab" property becomes observable end-to-end. Note on ``-fPIC`` vs ``-fpie`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ With ``-fPIC`` (the default for shared-library code), GCC may use ``R_X86_64_GOTPCRELX`` (GOT-relative, load through the GOT) for hidden-visibility externals like ``__dso_handle``. GOT entries are co-located with code, so there is no ±2 GB range issue. With ``-fpie`` (position-independent executable), GCC prefers the shorter direct ``R_X86_64_PC32``, which *does* have the ±2 GB limit. The Delta32 overflow tests (test 6) therefore build with ``-fpie`` to force the problematic relocation type. Test structure -------------- 1. **Co-location** (test 1): a single slab keeps objects within its size. 2. **Scatter baseline** (test 2): with the slab pool disabled, a VA blocker pushes objects far apart — proves the slab pool is responsible for co-location. 3. **Hidden-symbol calls** (test 3): ADRP/PC32 cross-object calls succeed under VA pressure with a slab. 4. **Large data section** (test 4): 4 MB ``.nv_fatbin`` section loads correctly within the slab. 5. **Overflow section** (test 5): ``.nv_fatbin`` data is allocated outside the slab via separate mmap. 6. **Leaked materialization** (test 6): ``__dso_handle`` resolves after prior sessions leaked mmap slabs from failed materializations. 7. **Delta32 overflow** (test 7): ``-fpie`` GCC objects + 3 GB VA blocker. With the slab pool → PASSES; without → Delta32 overflow. All tests use a 256 MB slab and 256 MB-3 GB VA blockers — safe for CI containers. """ from __future__ import annotations import ctypes import ctypes.util import functools import platform import sys from pathlib import Path import pytest from tvm_ffi_orcjit import ExecutionSession from utils import build_test_objects # --------------------------------------------------------------------------- # Setup # --------------------------------------------------------------------------- OBJ_DIR = build_test_objects() _KNOWN_SUBDIRS = [ "c", "c-gcc", "cc", "cc-gcc", "cc-gcc-pie", "c-appleclang", "cc-appleclang", "c-msvc", "c-clang-cl", ] _PIE_VARIANT_MARKER = "-pie" def obj(name: str) -> str: """Return path to a pre-built test object file, or skip if missing.""" path = OBJ_DIR / f"{name}.o" if not path.exists(): pytest.skip(f"{path.name} not found (not built)") return str(path) def _discover_c_variants() -> list[str]: """Discover available C-only compiler variants.""" return [ s for s in _KNOWN_SUBDIRS if s.startswith("c") and not s.startswith("cc") and _PIE_VARIANT_MARKER not in s and (OBJ_DIR / s / "test_funcs.o").exists() ] def _discover_cpp_variants() -> list[str]: """Discover available C++ compiler variants (for __dso_handle tests).""" return [ s for s in _KNOWN_SUBDIRS if s.startswith("cc") and _PIE_VARIANT_MARKER not in s and (OBJ_DIR / s / "test_funcs.o").exists() ] def _discover_gcc_cpp_variants() -> list[str]: """Discover GCC C++ variants (emit R_X86_64_PC32 for __dso_handle).""" return [v for v in _discover_cpp_variants() if "gcc" in v] def _discover_pie_cpp_variants() -> list[str]: """Discover PIE C++ variants built with -fpie. PIE objects force R_X86_64_PC32 (direct, ±2GB) for __dso_handle instead of R_X86_64_GOTPCRELX (GOT-relative, unlimited range). Used exclusively by the Delta32 overflow tests (test 6). """ return [ s for s in _KNOWN_SUBDIRS if _PIE_VARIANT_MARKER in s and (OBJ_DIR / s / "test_funcs.o").exists() ] _c_variants = _discover_c_variants() _cpp_variants = _discover_cpp_variants() _gcc_cpp_variants = _discover_gcc_cpp_variants() _pie_cpp_variants = _discover_pie_cpp_variants() _all_variants = _c_variants + _cpp_variants _is_linux = sys.platform == "linux" _is_x86_64 = platform.machine() in ("x86_64", "AMD64") # Slab test parameters. # # Under the slab-pool design, `slab_size` is the per-slab capacity — each # allocation that exceeds what an existing slab can hold spawns a new one. # These tests assert single-slab invariants (objects within slab_size, # one contiguous VA region), so `_SLAB_SIZE` must be large enough that # every graph we load fits in the first slab. 256 MB comfortably covers # the test objects (all < 5 MB each) plus their overhead. _SLAB_SIZE = 256 * 1024 * 1024 # 256MB — single-slab headroom for these tests _BLOCK_RADIUS = 256 * 1024 * 1024 # 256MB — safe for CI containers _DSO_BLOCK_RADIUS = 3 * 1024 * 1024 * 1024 # 3GB — needed to overflow PC32 (±2GB) _PROT_NONE = 0 _MAP_PRIVATE_ANON = 0x22 # MAP_PRIVATE | MAP_ANONYMOUS _MAP_FIXED_NOREPLACE = 0x100000 # --------------------------------------------------------------------------- # VA blocker — fills nearby free VA gaps to force distant mmap placement # --------------------------------------------------------------------------- @functools.lru_cache(maxsize=1) def _get_libc() -> ctypes.CDLL: """Get a ctypes handle to libc with correct mmap/munmap signatures.""" libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6", use_errno=True) libc.mmap.restype = ctypes.c_void_p libc.mmap.argtypes = [ ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long, ] libc.munmap.restype = ctypes.c_int libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] return libc def _parse_maps() -> list[tuple[int, int]]: """Parse /proc/self/maps into sorted list of (start, end) tuples.""" regions = [] with Path("/proc/self/maps").open() as f: for line in f: addrs = line.split()[0].split("-") regions.append((int(addrs[0], 16), int(addrs[1], 16))) return sorted(regions) def _find_new_mappings( before: set[tuple[int, int]], after: list[tuple[int, int]] ) -> list[tuple[int, int]]: """Find mappings present in *after* but not in *before*.""" return [(s, e) for s, e in after if (s, e) not in before] def block_nearby_va(center: int, radius: int = _BLOCK_RADIUS) -> list[tuple[int, int]]: """Block all free VA gaps within *radius* of *center*. Uses MAP_FIXED_NOREPLACE to place PROT_NONE mappings in every free gap within [center - radius, center + radius]. This forces subsequent mmap(NULL, ...) calls to land outside the blocked region. Returns list of (addr, size) blockers to be freed later. """ libc = _get_libc() maps = _parse_maps() blockers = [] low = max(center - radius, 0) high = center + radius for i in range(len(maps) - 1): gap_start = maps[i][1] gap_end = maps[i + 1][0] if gap_end <= low or gap_start >= high or gap_end <= gap_start: continue block_start = max(gap_start, low) block_end = min(gap_end, high) block_size = block_end - block_start if block_size <= 0: continue addr = libc.mmap( block_start, block_size, _PROT_NONE, _MAP_PRIVATE_ANON | _MAP_FIXED_NOREPLACE, -1, 0 ) if addr != ctypes.c_void_p(-1).value and addr is not None: blockers.append((addr, block_size)) return blockers def free_blockers(blockers: list[tuple[int, int]]) -> None: """Free all VA blockers.""" libc = _get_libc() for addr, size in blockers: libc.munmap(addr, size) # --------------------------------------------------------------------------- # Test 1: Slab co-location — objects stay within slab range # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.parametrize("variant", _all_variants) def test_slab_colocation(variant: str) -> None: """With slab, objects in separate libraries have close code addresses. Uses a 16MB slab and inserts a 256MB VA blocker between object loads. Without the slab, the blocker would push the second object far away. With the slab, both objects land within the 16MB region. """ maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=_SLAB_SIZE) lib1 = session.create_library("lib1") lib1.add(obj(f"{variant}/test_addr")) addr1 = lib1.get_function("code_address")() # Find where LLVM placed the first allocation and block nearby VA maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else addr1 blockers = block_nearby_va(jit_center) try: lib2 = session.create_library("lib2") lib2.add(obj(f"{variant}/test_addr")) addr2 = lib2.get_function("code_address")() finally: free_blockers(blockers) distance = abs(addr1 - addr2) assert distance < _SLAB_SIZE, ( f"Objects should be within {_SLAB_SIZE} bytes, " f"but distance is {distance} ({distance / (1024**2):.1f} MB)" ) # --------------------------------------------------------------------------- # Test 2: Slab effect — compare with-slab vs without-slab under VA pressure # --------------------------------------------------------------------------- def _measure_distance_under_pressure( variant: str, slab_size: int, radius: int = _BLOCK_RADIUS ) -> tuple[int | None, bool]: """Load two objects under VA pressure and return (distance, overflowed). Returns ``(distance_bytes, False)`` when both objects load successfully, or ``(None, True)`` when the second load fails with a relocation overflow (Page21 on AArch64, Delta32 on x86_64). """ maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=slab_size) lib1 = session.create_library("lib1") lib1.add(obj(f"{variant}/test_addr")) addr1 = lib1.get_function("code_address")() maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else addr1 blockers = block_nearby_va(jit_center, radius=radius) try: lib2 = session.create_library("lib2") try: lib2.add(obj(f"{variant}/test_addr")) addr2 = lib2.get_function("code_address")() except Exception: return None, True finally: free_blockers(blockers) return abs(addr1 - addr2), False @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.parametrize("variant", _all_variants) def test_slab_keeps_objects_close(variant: str) -> None: """Slab co-locates objects that would otherwise scatter or overflow. Runs the same workload twice under identical VA pressure — once with the slab and once without — and compares the outcomes: - **With slab**: both objects must land within the slab size (16 MB). - **Without slab**: the blocker should either cause a relocation overflow (proving scatter beyond relocation range) or produce a measurably larger distance. The test proves the slab is responsible for co-location by showing a strictly better outcome with it enabled. If the VA blocker happens to be ineffective (e.g., LLVM slab reuse on 64k-page kernels), the test still passes as long as the slab keeps objects within range. """ # Phase 1: with slab — must always succeed and be within slab range slab_dist, slab_overflow = _measure_distance_under_pressure(variant, slab_size=_SLAB_SIZE) assert not slab_overflow, "Slab session should not overflow" assert slab_dist is not None assert slab_dist < _SLAB_SIZE, ( f"With slab, objects should be within {_SLAB_SIZE} bytes, " f"but distance is {slab_dist} ({slab_dist / (1024**2):.1f} MB)" ) # Phase 2: without slab — expect scatter or overflow no_slab_dist, no_slab_overflow = _measure_distance_under_pressure(variant, slab_size=-1) if no_slab_overflow: # Relocation overflow without slab proves the blocker forced # scatter beyond relocation range — slab prevented this. return assert no_slab_dist is not None if no_slab_dist > slab_dist: # Without slab produced a larger distance — slab effect shown. return # Blocker was ineffective (both distances are small). The slab # assertion above already passed, which is the key property. We # cannot distinguish slab effect from lucky placement here. pytest.skip( f"VA blocker ineffective: slab={slab_dist / 1024:.0f} KB, " f"no-slab={no_slab_dist / 1024:.0f} KB — " f"cannot demonstrate slab effect on this kernel" ) # --------------------------------------------------------------------------- # Test 3: Hidden-symbol ADRP/PC32 relocation with slab + blocker # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.parametrize("variant", _all_variants) def test_slab_hidden_symbol_with_blocker(variant: str) -> None: """Slab prevents hidden-visibility relocation overflow under VA pressure. Loads two objects with hidden-visibility cross-references (ADRP+ADD on AArch64, PC32 on x86_64) with a VA blocker between them. Without slab, the blocker would push objects apart causing overflow. With the slab, both objects are co-located and the call succeeds. """ maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=_SLAB_SIZE) lib = session.create_library("hidden_test") # Load helper and force materialization lib.add(obj(f"{variant}/test_hidden_helper")) assert lib.get_function("hidden_add")(1, 2) == 3 # Block nearby VA to force scatter maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else 0xFFFF00000000 blockers = block_nearby_va(jit_center) try: lib.add(obj(f"{variant}/test_hidden_caller")) fn = lib.get_function("call_hidden_add") assert fn(10, 20) == 30 finally: free_blockers(blockers) # --------------------------------------------------------------------------- # Test 4: Large data section (simulated .nv_fatbin) # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.parametrize("variant", _all_variants) def test_large_data_section(variant: str) -> None: """Load object with a 4MB .nv_fatbin section — basic correctness. The .nv_fatbin section is referenced only by absolute relocations, so it can live anywhere. This test verifies the object loads and the function works. The 4MB section fits in the slab. """ session = ExecutionSession() lib = session.create_library("fatbin") lib.add(obj(f"{variant}/fake_fatbin")) fn = lib.get_function("get_fatbin_size") assert fn() == 4 * 1024 * 1024 # --------------------------------------------------------------------------- # Test 5: Overflow section — .nv_fatbin lands outside the slab # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.parametrize("variant", _all_variants) def test_overflow_section_outside_slab(variant: str) -> None: """Overflow sections (.nv_fatbin) are allocated outside the slab. The slab memory manager detects sections named .nv_fatbin and allocates them via a separate mmap() outside the slab. This keeps the slab compact for code + small rodata, reducing 2MB THP region count and iTLB pressure. Verification: get the fatbin data address and the slab VA range from /proc/self/maps, then assert the fatbin address is NOT within the slab region. """ session = ExecutionSession(slab_size=_SLAB_SIZE) lib = session.create_library("fatbin_overflow") lib.add(obj(f"{variant}/fake_fatbin")) # Verify the function still works correctly. assert lib.get_function("get_fatbin_size")() == 4 * 1024 * 1024 # Get the actual address of the fatbin data in memory. fatbin_addr = lib.get_function("get_fatbin_addr")() # Find the slab mapping: a single large region matching the slab size. # The slab is reserved as PROT_NONE and then committed in slabs, so # look for the contiguous region that spans _SLAB_SIZE. maps = _parse_maps() slab_regions = [(s, e) for s, e in maps if (e - s) >= _SLAB_SIZE] # The fatbin address must not fall within any slab-sized region. for start, end in slab_regions: assert not (start <= fatbin_addr < end), ( f"Fatbin data at {fatbin_addr:#x} should be OUTSIDE the slab " f"[{start:#x}, {end:#x}) but landed inside" ) # --------------------------------------------------------------------------- # Test 6: __dso_handle Delta32 overflow after leaked materialization # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="ELF/GCC-specific __dso_handle test") @pytest.mark.parametrize("variant", _cpp_variants) def test_dso_handle_relocation_after_failed_materialization(variant: str) -> None: """__dso_handle resolves correctly after leaked JIT memory. Mechanism --------- GCC C++ objects call ``__cxa_atexit(&destructor, &obj, __dso_handle)`` for static-storage-duration objects. The ``__dso_handle`` symbol is emitted as ``GLOBAL HIDDEN UND`` in each object file. LLVM's ``ELFNixPlatform`` defines it per JITDylib via a separate ``DSOHandleMaterializationUnit`` — a self-referential pointer block (``void *__dso_handle = &__dso_handle;``) allocated in its own ``LinkGraph`` through ``ObjectLinkingLayer``. When a prior ``lib.add()`` fails (e.g., duplicate symbol), LLVM's ``InProcessMemoryMapper`` leaks the mmap'd slab for that failed materialization. If the process holds references to the old session (e.g., pytest keeping ``sys.exc_info()`` tracebacks alive), the leaked slabs accumulate and push subsequent ``mmap`` allocations to higher addresses. The slab prevents overflow because all allocations — both ``__dso_handle``'s ``LinkGraph`` and the code ``LinkGraph`` — land within the same contiguous pre-reserved VA region. Without slab: may FAIL on x86_64 with GCC PIE objects after repeated leaked materializations push slabs >2 GB apart. With slab: PASSES (all allocations in same slab). """ # Step 1: Trigger leaked materializations to consume low VA space. leaked_sessions = [] for _ in range(3): s0 = ExecutionSession() lib0 = s0.create_library("warmup") lib0.add(obj(f"{variant}/test_funcs")) lib0.get_function("test_add")(10, 20) try: lib0.add(obj(f"{variant}/test_funcs_conflict")) except Exception: pass leaked_sessions.append((s0, lib0)) # Step 2: Fresh session — cross-library resolution must still work. session = ExecutionSession() lib1 = session.create_library("lib1") lib1.add(obj(f"{variant}/test_funcs")) assert lib1.get_function("test_add")(10, 20) == 30 lib2 = session.create_library("lib2") lib2.add(obj(f"{variant}/test_funcs_conflict")) assert lib2.get_function("test_add")(10, 20) == 1030 # --------------------------------------------------------------------------- # Test 6: __dso_handle Delta32 overflow — slab prevents it (x86_64 PIE) # # GCC -fpie objects use R_X86_64_PC32 (±2GB) for __dso_handle. # ELFNixPlatform's DSOHandleMaterializationUnit allocates __dso_handle # in a separate LinkGraph from the code. Under VA pressure, these two # allocations can land >2GB apart, overflowing the Delta32 fixup. # The slab keeps them co-located within relocation range. # --------------------------------------------------------------------------- @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.skipif(not _is_x86_64, reason="Delta32 overflow requires x86_64") @pytest.mark.skipif(not _pie_cpp_variants, reason="No GCC PIE C++ variants built") @pytest.mark.parametrize("variant", _pie_cpp_variants or ["skip"]) def test_dso_handle_delta32_with_slab(variant: str) -> None: """Slab prevents __dso_handle Delta32 overflow under VA pressure. Root cause ---------- GCC C++ objects built with ``-fpie`` emit ``R_X86_64_PC32`` (Delta32, ±2 GB) relocations for ``__dso_handle`` because the symbol has hidden visibility and ``-fpie`` prefers direct PC-relative over GOT-relative. (With ``-fPIC``, GCC uses ``R_X86_64_GOTPCRELX`` which goes through the GOT — always co-located with code, so no range issue.) ``ELFNixPlatform`` defines ``__dso_handle`` per JITDylib in a separate ``DSOHandleMaterializationUnit``. This creates a tiny ``LinkGraph`` (a self-referential pointer: ``void *__dso_handle = &__dso_handle;``) that is allocated through ``ObjectLinkingLayer`` independently of the code ``LinkGraph`` from ``lib.add()``. Both allocations go through ``InProcessMemoryMapper`` → ``mmap(MAP_ANONYMOUS)``, whose placement the kernel decides. Test strategy ------------- 1. Create a session with slab enabled (16 MB). 2. Load PIE GCC objects into lib1 — this triggers materialization of both ``__dso_handle`` (via ``DSOHandleMaterializationUnit``) and the code (via ``lib.add``), all within the slab. 3. Block 3 GB of VA around the first allocation — without slab this would force the next ``mmap`` to land >2 GB away. 4. Load a second PIE GCC object into lib2 — with slab, this still lands within the 16 MB region. 5. Assert the function call succeeds — proves Delta32 is in range. See ``test_dso_handle_delta32_overflow_without_slab`` for the counterpart proving the overflow occurs without slab. """ maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=_SLAB_SIZE) lib1 = session.create_library("lib1") lib1.add(obj(f"{variant}/test_funcs")) assert lib1.get_function("test_add")(10, 20) == 30 # Block 3GB of VA around the first allocation to force scatter maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else 0xFFFF00000000 blockers = block_nearby_va(jit_center, radius=_DSO_BLOCK_RADIUS) try: lib2 = session.create_library("lib2") lib2.add(obj(f"{variant}/test_funcs_conflict")) assert lib2.get_function("test_add")(10, 20) == 1030 finally: free_blockers(blockers) @pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only") @pytest.mark.skipif(not _is_x86_64, reason="Delta32 overflow requires x86_64") @pytest.mark.skipif(not _pie_cpp_variants, reason="No GCC PIE C++ variants built") @pytest.mark.parametrize("variant", _pie_cpp_variants or ["skip"]) def test_dso_handle_delta32_overflow_without_slab(variant: str) -> None: """Without slab, PIE __dso_handle PC32 overflows under VA pressure. Same setup as ``test_dso_handle_delta32_with_slab`` but with slab disabled (``slab_size=-1``). The 3 GB VA blocker fills all free gaps within ±3 GB of the first session's JIT allocations. When lib2 is loaded, ``InProcessMemoryMapper`` calls ``mmap(MAP_ANONYMOUS)`` for a new slab, but the only free VA is >3 GB away. The code ``LinkGraph`` from ``lib2.add()`` lands in that distant slab, while ``__dso_handle`` was already materialized with lib1's ``DSOHandleMaterializationUnit`` in the original region. The ``R_X86_64_PC32`` fixup from code to ``__dso_handle`` now exceeds ±2 GB → JITLink reports ``Delta32 fixup ... is out of range``. The test accepts both outcomes: - **Exception** (PC32 overflow): proves the slab is needed. - **Success** (GOTPCRELX used): GCC chose GOT-relative despite ``-fpie`` — no overflow possible, but the slab is still beneficial for other relocation types. """ maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=-1) # slab disabled lib1 = session.create_library("lib1") lib1.add(obj(f"{variant}/test_addr")) lib1.get_function("code_address")() maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else 0xFFFF00000000 blockers = block_nearby_va(jit_center, radius=_DSO_BLOCK_RADIUS) try: lib2 = session.create_library("lib2") try: lib2.add(obj(f"{variant}/test_funcs_conflict")) result = lib2.get_function("test_add")(10, 20) # If we get here, GCC used GOTPCRELX — no overflow. assert result == 1030 except Exception: # R_X86_64_PC32 overflow as expected — proves slab is needed. pass finally: free_blockers(blockers) tvm-ffi-0.1.12/addons/tvm_ffi_orcjit/tests/utils.py000066400000000000000000000200651521067262500223040ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Build test object files for all available compiler variants. Uses ``tvm_ffi.cpp.build`` to compile C/C++ test sources to relocatable object files. Detects platform and available compilers: - Linux: LLVM Clang (default) + GCC - macOS: LLVM Clang (default) + Apple Clang - Windows: LLVM Clang (default) + MSVC + clang-cl Objects are placed under a temporary directory and skipped if they already exist. LLVM clang is found via ``$LLVM_PREFIX/bin`` or ``$PATH``; no hardcoded install path is assumed. """ from __future__ import annotations import os import platform import shutil import tempfile from pathlib import Path import tvm_ffi.cpp SOURCES_DIR = Path(__file__).resolve().parent / "sources" SOURCES_C = SOURCES_DIR / "c" SOURCES_CC = SOURCES_DIR / "cc" SOURCES_CUDA = SOURCES_DIR / "cuda" _DEFAULT_OUT_DIR = Path(tempfile.gettempdir()) / "tvm_ffi_orcjit_tests" # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _extra_cflags() -> list[str]: machine = platform.machine() if machine in ("aarch64", "arm64"): return ["-mno-outline-atomics"] return [] def _extra_cuda_cflags() -> list[str]: machine = platform.machine() if machine in ("aarch64", "arm64"): return ["-Xcompiler", "-mno-outline-atomics"] return [] def _build_objects( src_dir: Path, out_dir: Path, *, ext_glob: str, extra_cflags: list[str], extra_cuda_cflags: list[str] | None = None, ) -> None: """Compile all sources in *src_dir* to object files in *out_dir*.""" out_dir.mkdir(parents=True, exist_ok=True) for src in sorted(src_dir.glob(ext_glob)): dest = out_dir / f"{src.stem}.o" if dest.exists(): continue build_dir = out_dir / f".build_{src.stem}" obj_path = tvm_ffi.cpp.build( name=src.stem, sources=[str(src)], output=f"{src.stem}.o", extra_cflags=extra_cflags, extra_cuda_cflags=extra_cuda_cflags or [], build_directory=str(build_dir), ) shutil.copy2(obj_path, dest) def _build_variant( name: str, *, cc: str | None, cxx: str | None, extra_cflags: list[str], c_outdir: Path, cc_outdir: Path | None, ) -> None: """Build all test objects for one compiler variant.""" saved_cc = os.environ.get("CC") saved_cxx = os.environ.get("CXX") try: if cc: os.environ["CC"] = cc if cxx: os.environ["CXX"] = cxx if cc: _build_objects(SOURCES_C, c_outdir, ext_glob="*.c", extra_cflags=extra_cflags) if cxx and cc_outdir: _build_objects(SOURCES_CC, cc_outdir, ext_glob="*.cc", extra_cflags=extra_cflags) finally: if saved_cc is None: os.environ.pop("CC", None) else: os.environ["CC"] = saved_cc if saved_cxx is None: os.environ.pop("CXX", None) else: os.environ["CXX"] = saved_cxx # --------------------------------------------------------------------------- # Compiler detection # --------------------------------------------------------------------------- def _find_llvm_clang() -> tuple[str, str] | None: """Find LLVM clang/clang++ via LLVM_PREFIX or PATH. Returns (clang, clang++) paths, or None if not found. """ llvm_prefix = os.environ.get("LLVM_PREFIX") if llvm_prefix: p = Path(llvm_prefix) bin_dir = p / "Library" / "bin" if (p / "Library" / "bin").exists() else p / "bin" cc = bin_dir / ("clang.exe" if platform.system() == "Windows" else "clang") if cc.exists(): cxx = bin_dir / ("clang++.exe" if platform.system() == "Windows" else "clang++") return str(cc), str(cxx) cc = shutil.which("clang") if cc: cxx = shutil.which("clang++") return cc, cxx or cc return None # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def build_test_objects(out_dir: Path | None = None) -> Path: """Build test objects for all available compiler variants. Objects are placed under *out_dir* (default: ``{tempdir}/tvm_ffi_orcjit_tests``). Already-built objects are skipped. Returns the output directory. """ if out_dir is None: out_dir = _DEFAULT_OUT_DIR out_dir.mkdir(parents=True, exist_ok=True) system = platform.system() extra = _extra_cflags() if system in ("Linux", "Darwin"): llvm = _find_llvm_clang() if llvm: clang, clangxx = llvm _build_variant( "LLVM Clang", cc=clang, cxx=clangxx, extra_cflags=extra, c_outdir=out_dir / "c", cc_outdir=out_dir / "cc", ) if system == "Linux" and shutil.which("gcc"): _build_variant( "GCC", cc="gcc", cxx="g++", extra_cflags=extra, c_outdir=out_dir / "c-gcc", cc_outdir=out_dir / "cc-gcc", ) # PIE variant: -fpie forces R_X86_64_PC32 for hidden-visibility # externals like __dso_handle (instead of GOTPCRELX with -fPIC). # Used to reproduce __dso_handle Delta32 overflow on x86_64. _build_variant( "GCC (PIE)", cc=None, cxx="g++", extra_cflags=[*extra, "-fpie"], c_outdir=out_dir / "c-gcc-pie", cc_outdir=out_dir / "cc-gcc-pie", ) if system == "Darwin" and Path("/usr/bin/clang").exists(): _build_variant( "Apple Clang", cc="/usr/bin/clang", cxx="/usr/bin/clang++", extra_cflags=extra, c_outdir=out_dir / "c-appleclang", cc_outdir=out_dir / "cc-appleclang", ) elif system == "Windows": llvm = _find_llvm_clang() if llvm: clang_cl = Path(llvm[0]).parent / "clang-cl.exe" if clang_cl.exists(): _build_variant( "LLVM clang-cl", cc=str(clang_cl), cxx=None, extra_cflags=["/GS-"], c_outdir=out_dir / "c", cc_outdir=None, ) if shutil.which("cl"): _build_variant( "MSVC", cc="cl", cxx=None, extra_cflags=["/GS-"], c_outdir=out_dir / "c-msvc", cc_outdir=None, ) if shutil.which("clang-cl"): _build_variant( "clang-cl", cc="clang-cl", cxx=None, extra_cflags=["/GS-"], c_outdir=out_dir / "c-clang-cl", cc_outdir=None, ) # CUDA (platform-independent, uses nvcc) if shutil.which("nvcc"): _build_objects( SOURCES_CUDA, out_dir / "cuda", ext_glob="*.cu", extra_cflags=[], extra_cuda_cflags=_extra_cuda_cflags(), ) return out_dir tvm-ffi-0.1.12/cmake/000077500000000000000000000000001521067262500142315ustar00rootroot00000000000000tvm-ffi-0.1.12/cmake/Utils/000077500000000000000000000000001521067262500153315ustar00rootroot00000000000000tvm-ffi-0.1.12/cmake/Utils/AddGoogleTest.cmake000066400000000000000000000054301521067262500210220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. include(FetchContent) set(gtest_force_shared_crt # cmake-lint: disable=C0103 ON CACHE BOOL "Always use msvcrt.dll" FORCE ) set(BUILD_GMOCK ON CACHE BOOL "" FORCE ) set(BUILD_GTEST ON CACHE BOOL "" FORCE ) FetchContent_Declare( googletest GIT_REPOSITORY https://github.com/google/googletest.git GIT_TAG v1.14.0 ) FetchContent_GetProperties(googletest) if (NOT googletest_POPULATED) FetchContent_MakeAvailable(googletest) include(GoogleTest) set_target_properties( gtest PROPERTIES EXPORT_COMPILE_COMMANDS OFF EXCLUDE_FROM_ALL ON FOLDER 3rdparty ) set_target_properties( gtest_main PROPERTIES EXPORT_COMPILE_COMMANDS OFF EXCLUDE_FROM_ALL ON FOLDER 3rdparty ) set_target_properties( gmock PROPERTIES EXPORT_COMPILE_COMMANDS OFF EXCLUDE_FROM_ALL ON FOLDER 3rdparty ) set_target_properties( gmock_main PROPERTIES EXPORT_COMPILE_COMMANDS OFF EXCLUDE_FROM_ALL ON FOLDER 3rdparty ) mark_as_advanced( BUILD_GMOCK BUILD_GTEST BUILD_SHARED_LIBS gmock_build_tests gtest_build_samples gtest_build_tests gtest_disable_pthreads gtest_force_shared_crt gtest_hide_internal_symbols ) endif () # ~~~ # TVM_FFI_ADD_GTEST(target_name) # Register a GoogleTest executable as a CTest, link it against gtest_main, # and configure test discovery and properties. # # Parameters: # target_name: Name of the test executable target # ~~~ macro (TVM_FFI_ADD_GTEST target_name) add_test( NAME ${target_name} COMMAND ${target_name} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) target_link_libraries(${target_name} PRIVATE gtest_main) gtest_discover_tests( ${target_name} WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} TEST_DISCOVERY_TIMEOUT 600 DISCOVERY_MODE PRE_TEST PROPERTIES VS_DEBUGGER_WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}" ) set_target_properties(${target_name} PROPERTIES FOLDER tests) endmacro () tvm-ffi-0.1.12/cmake/Utils/AddLibbacktrace.cmake000066400000000000000000000065341521067262500213220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. include(ExternalProject) # ~~~ # _libbacktrace_compile() # Build and install libbacktrace as an ExternalProject, then add an imported # static target `libbacktrace` exposing headers and archive for consumers. # ~~~ function (_libbacktrace_compile) set(libbacktrace_source ${CMAKE_CURRENT_LIST_DIR}/../../3rdparty/libbacktrace) set(libbacktrace_prefix ${CMAKE_CURRENT_BINARY_DIR}/libbacktrace) if (CMAKE_SYSTEM_NAME MATCHES "Darwin" AND (CMAKE_C_COMPILER MATCHES "^/Library" OR CMAKE_C_COMPILER MATCHES "^/Applications") ) set(cmake_c_compiler "/usr/bin/cc") else () set(cmake_c_compiler "${CMAKE_C_COMPILER}") endif () file(MAKE_DIRECTORY ${libbacktrace_prefix}/include) file(MAKE_DIRECTORY ${libbacktrace_prefix}/lib) detect_target_triple(TVM_FFI_MACHINE_NAME) message(STATUS "Detected target triple: ${TVM_FFI_MACHINE_NAME}") # Add symbol hiding flags for GCC, Clang, and AppleClang set(symbol_hiding_flags "") if (CMAKE_C_COMPILER_ID MATCHES "GNU|Clang|AppleClang") set(symbol_hiding_flags "-fvisibility=hidden -fvisibility-inlines-hidden") endif () ExternalProject_Add( project_libbacktrace PREFIX libbacktrace SOURCE_DIR ${libbacktrace_source} BINARY_DIR ${libbacktrace_prefix} LOG_DIR ${libbacktrace_prefix}/logs CONFIGURE_COMMAND "sh" # "${libbacktrace_source}/configure" # "--prefix=${libbacktrace_prefix}" # "--with-pic" # "CC=${cmake_c_compiler}" # "CPP=${cmake_c_compiler} -E" # "CFLAGS=${CMAKE_C_FLAGS} ${symbol_hiding_flags}" # "LDFLAGS=${CMAKE_EXE_LINKER_FLAGS}" # "NM=${CMAKE_NM}" # "STRIP=${CMAKE_STRIP}" # "--host=${TVM_FFI_MACHINE_NAME}" INSTALL_DIR ${libbacktrace_prefix} INSTALL_COMMAND make install BUILD_BYPRODUCTS "${libbacktrace_prefix}/lib/libbacktrace.a" "${libbacktrace_prefix}/include/backtrace.h" LOG_CONFIGURE ON LOG_INSTALL ON LOG_BUILD ON LOG_MERGED_STDOUTERR ON LOG_OUTPUT_ON_FAILURE ON ) ExternalProject_Add_Step( project_libbacktrace checkout DEPENDERS configure DEPENDEES download ) set_target_properties(project_libbacktrace PROPERTIES EXCLUDE_FROM_ALL TRUE) add_library(libbacktrace STATIC IMPORTED) add_dependencies(libbacktrace project_libbacktrace) set_target_properties( libbacktrace PROPERTIES IMPORTED_LOCATION ${libbacktrace_prefix}/lib/libbacktrace.a INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_CURRENT_LIST_DIR}/../../3rdparty/libbacktrace/ ) endfunction () if (NOT MSVC) _libbacktrace_compile() endif () tvm-ffi-0.1.12/cmake/Utils/CxxWarning.cmake000066400000000000000000000027061521067262500204300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ~~~ # tvm_ffi_add_cxx_warning(target_name) # Apply a consistent set of warning flags (or placeholders) depending on the active compiler family. # Parameters: # target_name: CMake target to modify # ~~~ function (tvm_ffi_add_cxx_warning target_name) # GNU, Clang, or AppleClang if (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|AppleClang") target_compile_options( ${target_name} PRIVATE "-Werror" "-Wall" "-Wextra" "-Wpedantic" "-Wno-unused-parameter" ) return() endif () # MSVC if (MSVC) # target_compile_options(${target_name} PRIVATE "/W4" "/WX") return() endif () message(FATAL_ERROR "Unsupported compiler: ${CMAKE_CXX_COMPILER_ID}") endfunction () tvm-ffi-0.1.12/cmake/Utils/DetectTargetTriple.cmake000066400000000000000000000206041521067262500220740ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ~~~ # detect_target_triple(out_var) # Determine the target machine triple and store it in the variable named by # `out_var`. # # The result is determined by (in order of preference): # - querying CMake's own configuration, # - asking the compiler directly via `-dumpmachine` or `--print-target-triple`, # - using CMake's `CMAKE_LIBRARY_ARCHITECTURE` variable (Debian/Ubuntu multiarch hint), # - synthesizing from system information. # ~~~ # machine_triple.cmake function (detect_target_triple out_var) # cmake-lint: disable=R0911,R0912,R0915 # --- 1) Prefer CMake's own notion (e.g. when --target was used) --- foreach (lang C CXX) if (CMAKE_${lang}_COMPILER_TARGET) set(${out_var} "${CMAKE_${lang}_COMPILER_TARGET}" PARENT_SCOPE ) return() endif () endforeach () # --- 2) Ask the compiler directly (works for Clang/GCC, Android NDK, Emscripten) --- set(cc "${CMAKE_C_COMPILER}") if (NOT cc AND CMAKE_CXX_COMPILER) set(cc "${CMAKE_CXX_COMPILER}") endif () if (cc) execute_process( COMMAND "${cc}" -dumpmachine OUTPUT_VARIABLE ret OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET ) if (NOT ret) execute_process( COMMAND "${cc}" --print-target-triple OUTPUT_VARIABLE ret OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET ) endif () if (ret) set(${out_var} "${ret}" PARENT_SCOPE ) return() endif () endif () # --- 3) Platform-specific construction --- # 3a) Emscripten (toolchains usually set CMAKE_SYSTEM_NAME to Emscripten) if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN) set(${out_var} "wasm32-unknown-emscripten" PARENT_SCOPE ) return() endif () # 3b) Android (derive from ANDROID_ABI / ANDROID_PLATFORM) if (ANDROID OR CMAKE_SYSTEM_NAME STREQUAL "Android") set(arch "") set(abi "${ANDROID_ABI}") if (abi STREQUAL "armeabi-v7a") set(arch "armv7a") set(base "linux-androideabi") elseif (abi STREQUAL "arm64-v8a") set(arch "aarch64") set(base "linux-android") elseif (abi STREQUAL "x86") set(arch "i686") set(base "linux-android") elseif (abi STREQUAL "x86_64") set(arch "x86_64") set(base "linux-android") elseif (abi STREQUAL "riscv64") set(arch "riscv64") set(base "linux-android") else () # Fallback from processor if ABI isn't set set(arch "${CMAKE_SYSTEM_PROCESSOR}") string(TOLOWER "${arch}" arch) if (arch MATCHES "armv7") set(arch "armv7a") set(base "linux-androideabi") elseif (arch MATCHES "aarch64|arm64") set(arch "aarch64") set(base "linux-android") elseif (arch MATCHES "x86_64|amd64") set(arch "x86_64") set(base "linux-android") elseif (arch MATCHES "i[3-6]86|x86") set(arch "i686") set(base "linux-android") elseif (arch MATCHES "riscv64") set(arch "riscv64") set(base "linux-android") endif () endif () # Append API level if we can (e.g. aarch64-linux-android21) set(api "") if (DEFINED ANDROID_PLATFORM AND NOT "${ANDROID_PLATFORM}" STREQUAL "") string(REGEX REPLACE "android-?" "" api "${ANDROID_PLATFORM}") endif () if (arch STREQUAL "armv7a") set(ret "${arch}-${base}") else () set(ret "${arch}-${base}") endif () if (api) set(ret "${ret}${api}") endif () set(${out_var} "${ret}" PARENT_SCOPE ) return() endif () # 3c) Apple iOS (device & simulator). Works for Xcode + toolchains. if (APPLE AND (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_OSX_SYSROOT MATCHES "[iI]phone")) # Choose first arch if multi-arch is set set(archs "${CMAKE_OSX_ARCHITECTURES}") if (NOT archs) set(archs "${CMAKE_SYSTEM_PROCESSOR}") endif () list(GET archs 0 arch) string(TOLOWER "${arch}" _arch_l) if (_arch_l MATCHES "aarch64|arm64|arm64e") # iOS uses 'arm64' in triples set(arch "arm64") elseif (_arch_l MATCHES "x86_64|amd64") set(arch "x86_64") endif () # Simulator? set(is_sim OFF) if (CMAKE_OSX_SYSROOT MATCHES "simulator" OR CMAKE_XCODE_EFFECTIVE_PLATFORMS MATCHES "simulator" ) set(is_sim ON) endif () # Deployment target (best-effort) set(ios_ver "") foreach (maybe_ver CMAKE_OSX_DEPLOYMENT_TARGET CMAKE_IOS_DEPLOYMENT_TARGET IOS_DEPLOYMENT_TARGET ) if (DEFINED ${maybe_ver} AND NOT "${${maybe_ver}}" STREQUAL "") set(ios_ver "${${maybe_ver}}") break() endif () endforeach () if (is_sim) set(ret "${arch}-apple-ios${ios_ver}-simulator") else () set(ret "${arch}-apple-ios${ios_ver}") endif () string(REGEX REPLACE "ios$" "ios" ret "${ret}") # normalize empty version case set(${out_var} "${ret}" PARENT_SCOPE ) return() endif () # 3d) Windows + MSVC (cl.exe / clang-cl in MSVC mode) if (MSVC AND CMAKE_SYSTEM_NAME STREQUAL "Windows") set(plat "${CMAKE_GENERATOR_PLATFORM}") if (NOT plat AND DEFINED CMAKE_VS_PLATFORM_NAME) set(plat "${CMAKE_VS_PLATFORM_NAME}") endif () if (plat STREQUAL "Win32") set(arch "i686") elseif (plat MATCHES "^(x64|X64)$") set(arch "x86_64") elseif (plat MATCHES "ARM64") set(arch "arm64") else () # Fallback from pointer size if (CMAKE_SIZEOF_VOID_P EQUAL 8) set(arch "x86_64") else () set(arch "i686") endif () endif () set(${out_var} "${arch}-pc-windows-msvc" PARENT_SCOPE ) return() endif () # 3e) MinGW (handy if you ever hit it) if (MINGW) if (CMAKE_SIZEOF_VOID_P EQUAL 8) set(${out_var} "x86_64-w64-mingw32" PARENT_SCOPE ) else () set(${out_var} "i686-w64-mingw32" PARENT_SCOPE ) return() endif () endif () # --- 4) Debian/Ubuntu multiarch hint provided by CMake --- if (CMAKE_LIBRARY_ARCHITECTURE) set(${out_var} "${CMAKE_LIBRARY_ARCHITECTURE}" PARENT_SCOPE ) return() endif () # --- 5) Canonical, sensible fallback by OS --- set(arch "${CMAKE_SYSTEM_PROCESSOR}") if (NOT arch) if (CMAKE_SIZEOF_VOID_P EQUAL 8) set(arch "x86_64") else () set(arch "i686") endif () endif () string(TOLOWER "${arch}" _arch_l) # Normalize common arch spellings if (_arch_l MATCHES "aarch64|arm64|arm64e") if (APPLE) set(arch "arm64") # Apple uses arm64 else () set(arch "aarch64") endif () elseif (_arch_l MATCHES "x86_64|amd64") set(arch "x86_64") elseif (_arch_l MATCHES "i[3-6]86|x86") set(arch "i686") elseif (_arch_l MATCHES "armv7") set(arch "armv7") elseif (_arch_l MATCHES "riscv64") set(arch "riscv64") endif () if (APPLE) # macOS (Darwin) fallback set(${out_var} "${arch}-apple-darwin" PARENT_SCOPE ) return() elseif (CMAKE_SYSTEM_NAME STREQUAL "Windows") set(${out_var} "${arch}-pc-windows-msvc" PARENT_SCOPE ) return() elseif (CMAKE_SYSTEM_NAME STREQUAL "Linux") # Default to glibc; adjust to 'musl' in your toolchain if needed set(${out_var} "${arch}-unknown-linux-gnu" PARENT_SCOPE ) return() elseif (CMAKE_SYSTEM_NAME MATCHES "FreeBSD") set(${out_var} "${arch}-unknown-freebsd" PARENT_SCOPE ) return() endif () # Last-ditch (keeps your old behavior minimally sane) set(${out_var} "${arch}-${CMAKE_SYSTEM_NAME}" PARENT_SCOPE ) endfunction () tvm-ffi-0.1.12/cmake/Utils/EmbedCubin.cmake000066400000000000000000000122501521067262500203300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # If CMAKE_CUDA_RUNTIME_LIBRARY is not set, we default it to Shared. This prevents static linking of # cudart which requires exact driver version match. if (NOT DEFINED CMAKE_CUDA_RUNTIME_LIBRARY) set(CMAKE_CUDA_RUNTIME_LIBRARY Shared) message(STATUS "CMAKE_CUDA_RUNTIME_LIBRARY not set, defaulting to Shared. " "If you want to use driver API only, set CMAKE_CUDA_RUNTIME_LIBRARY to None." ) endif () set(OBJECT_COPY_UTIL "${CMAKE_CURRENT_LIST_DIR}/ObjectCopyUtil.cmake") # ~~~ # add_tvm_ffi_cubin( CUDA ) # # Creates an object library that compiles CUDA source to CUBIN format. # This function uses CMake's native CUDA support and respects CMAKE_CUDA_ARCHITECTURES. # This is a compatibility util for cmake < 3.27, user can create # cmake target with `CUDA_CUBIN_COMPILATION` for cmake >= 3.27. # # Parameters: # target_name: Name of the object library target # CUDA: One CUDA source file # # Example: # add_tvm_ffi_cubin(my_kernel_cubin CUDA kernel.cu) # ~~~ function (add_tvm_ffi_cubin target_name) cmake_parse_arguments(ARG "" "CUDA" "" ${ARGN}) if (NOT ARG_CUDA) message(FATAL_ERROR "add_tvm_ffi_cubin: CUDA source is required") endif () add_library(${target_name} OBJECT ${ARG_CUDA}) target_compile_options(${target_name} PRIVATE $<$:--cubin>) add_custom_target( ${target_name}_bin ALL COMMAND ${CMAKE_COMMAND} -DOBJECTS="$" -DOUT_DIR="" -DEXT="cubin" -P "${OBJECT_COPY_UTIL}" DEPENDS ${target_name} COMMENT "Generating .cubin files for ${target_name}" VERBATIM ) endfunction () # ~~~ # add_tvm_ffi_fatbin( CUDA ) # # Creates an object library that compiles CUDA source to FATBIN format. # This function uses CMake's native CUDA support and respects CMAKE_CUDA_ARCHITECTURES. # This is a compatibility util for cmake < 3.27, user can create # cmake target with `CUDA_FATBIN_COMPILATION` for cmake >= 3.27. # # Parameters: # target_name: Name of the object library target # CUDA: One CUDA source file # # Example: # add_tvm_ffi_fatbin(my_kernel_cubin CUDA kernel.cu) # ~~~ function (add_tvm_ffi_fatbin target_name) cmake_parse_arguments(ARG "" "CUDA" "" ${ARGN}) if (NOT ARG_CUDA) message(FATAL_ERROR "add_tvm_ffi_fatbin: CUDA source is required") endif () add_library(${target_name} OBJECT ${ARG_CUDA}) target_compile_options(${target_name} PRIVATE $<$:--fatbin>) add_custom_target( ${target_name}_bin ALL COMMAND ${CMAKE_COMMAND} -DOBJECTS="$" -DOUT_DIR="" -DEXT="fatbin" -P "${OBJECT_COPY_UTIL}" DEPENDS ${target_name} COMMENT "Generating .fatbin files for ${target_name}" VERBATIM ) endfunction () # ~~~ # tvm_ffi_embed_bin_into( # SYMBOL # BIN ) # # Embed one cubin/fatbin into given target with specified library name, # can be loaded with `TVM_FFI_EMBED_CUBIN(symbol_name)`. # Can only have one object in target and one cubin/fatbin. # # The reason of this design is to integrate with cmake's workflow. # # Parameters: # target_name: Name of the object library target # symbol_name: Name of the symbol in TVM_FFI_EMBED_CUBIN macro. # BIN: CUBIN or FATBIN file # # Example: # tvm_ffi_embed_bin_into(lib_embedded SYMBOL env BIN "$") # ~~~ function (tvm_ffi_embed_bin_into target_name) cmake_parse_arguments(ARG "" "SYMBOL;BIN" "" ${ARGN}) if (NOT ARG_BIN) message(FATAL_ERROR "tvm_ffi_embed_bin_into: BIN is required") endif () if (NOT ARG_SYMBOL) message(FATAL_ERROR "tvm_ffi_embed_bin_into: SYMBOL is required") endif () set(intermediate_path "${CMAKE_CURRENT_BINARY_DIR}/${ARG_SYMBOL}_intermediate.o") add_custom_command( TARGET ${target_name} PRE_LINK COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${intermediate_path}" COMMENT "Moving $ -> ${intermediate_path}" ) add_custom_command( TARGET ${target_name} PRE_LINK COMMAND ${Python_EXECUTABLE} -m tvm_ffi.utils.embed_cubin --output-obj "$" --name "${ARG_SYMBOL}" --input-obj "${intermediate_path}" --cubin "${ARG_BIN}" COMMENT "Embedding CUBIN into object file (name: ${ARG_SYMBOL})" VERBATIM ) endfunction () tvm-ffi-0.1.12/cmake/Utils/Library.cmake000066400000000000000000000360641521067262500177500ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ~~~ # tvm_ffi_add_prefix_map(target_name, prefix_path) # Add a compile prefix map so absolute paths under `prefix_path` are remapped to a stable, # relative form for reproducible builds and cleaner diagnostics. # # Parameters: # target_name: CMake target to modify # prefix_path: Absolute path prefix to remap # ~~~ function (tvm_ffi_add_prefix_map target_name prefix_path) # Add prefix map so the path displayed becomes relative to prefix_path if (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") target_compile_options(${target_name} PRIVATE "-ffile-prefix-map=${prefix_path}/=") endif () endfunction () # ~~~ # tvm_ffi_add_apple_dsymutil(target_name) # On Apple platforms, run `dsymutil` post-build to generate debug symbols for better backtraces. # No-ops on non-Apple platforms. # # Parameters: # target_name: CMake target to attach post-build step # ~~~ function (tvm_ffi_add_apple_dsymutil target_name) # running dsymutil on macos to generate debugging symbols for backtraces if (APPLE) find_program(DSYMUTIL dsymutil) mark_as_advanced(DSYMUTIL) add_custom_command( TARGET ${target_name} POST_BUILD COMMAND ${DSYMUTIL} ARGS $ COMMENT "[COMMAND] dsymutil $" VERBATIM ) endif () endfunction () # ~~~ # tvm_ffi_add_msvc_flags(target_name) # Apply MSVC-specific definitions and flags to improve build compatibility and warnings behavior # on Windows. # # Parameters: # target_name: CMake target to modify # ~~~ function (tvm_ffi_add_msvc_flags target_name) # running if we are under msvc if (MSVC) target_compile_definitions(${target_name} PUBLIC -DWIN32_LEAN_AND_MEAN) target_compile_definitions(${target_name} PUBLIC -D_CRT_SECURE_NO_WARNINGS) target_compile_definitions(${target_name} PUBLIC -D_SCL_SECURE_NO_WARNINGS) target_compile_definitions(${target_name} PUBLIC -D_ENABLE_EXTENDED_ALIGNED_STORAGE) target_compile_definitions(${target_name} PUBLIC -DNOMINMAX) target_compile_options(${target_name} PRIVATE "/Zi") # Heavy template instantiations in reflection/creator/object.h can exceed MSVC's default # per-object section limit (C1128). Apply /bigobj to every target that uses these flags so # growth in any TU doesn't break the Windows build. target_compile_options(${target_name} PRIVATE "/bigobj") endif () endfunction () # ~~~ # tvm_ffi_add_target_from_obj(target_name, obj_target_name) # Create static and shared library targets from an object library and set output directories # consistently across platforms. Also runs dsymutil on Apple for the shared target. # # Parameters: # target_name: Base name for created targets # obj_target_name: Object library to link into the outputs # ~~~ function (tvm_ffi_add_target_from_obj target_name obj_target_name) add_library(${target_name}_static STATIC $) add_library(${target_name}::static ALIAS ${target_name}_static) set_target_properties( ${target_name}_static PROPERTIES OUTPUT_NAME "${target_name}_static" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) add_library(${target_name}_shared SHARED $) add_library(${target_name}::shared ALIAS ${target_name}_shared) set_target_properties( ${target_name}_shared PROPERTIES OUTPUT_NAME "${target_name}" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) add_library(${target_name}_testing SHARED) set_target_properties( ${target_name}_testing PROPERTIES OUTPUT_NAME "${target_name}_testing" ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) if (WIN32) target_compile_definitions(${obj_target_name} PRIVATE TVM_FFI_EXPORTS) # set the output directory for each config type so msbuild also get into lib without appending # the config type to the output directory do both Release and RELEASE suffix, since while cmake # docs suggest Release is ok. real runs on MSbuild suggest that we might need RELEASE instead foreach (config_type Release RELEASE) set_target_properties( ${target_name}_shared PROPERTIES RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ) set_target_properties( ${target_name}_static PROPERTIES RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ) set_target_properties( ${target_name}_testing PROPERTIES RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib" ) endforeach () endif () tvm_ffi_add_apple_dsymutil(${target_name}_shared) tvm_ffi_add_apple_dsymutil(${target_name}_testing) endfunction () # cmake-lint: disable=C0301,R0912,R0915 # ~~~ # tvm_ffi_configure_target( # target_name # [LINK_SHARED ON|OFF] [LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF] [MSVC_FLAGS ON|OFF] # [STUB_INIT ON|OFF] [STUB_DIR ] [STUB_PKG ] [STUB_PREFIX ] # ) # Configure a target to integrate with TVM-FFI CMake utilities: # - Link against tvm_ffi::header and/or tvm_ffi::shared # - Always apply tvm_ffi_add_prefix_map(target_name ) # - Enable Apple dSYM generation via tvm_ffi_add_apple_dsymutil(target_name) # - Apply MSVC-specific flags via tvm_ffi_add_msvc_flags(target_name) # - Add post-build step to generate Python stubs via tvm_ffi.stub.cli # # Parameters: # target_name: Existing CMake target to modify (positional, required) # # Keyword parameters: # LINK_SHARED: Whether to link tvm_ffi::shared into the target (default: ON; ON/OFF-style) # LINK_HEADER: Whether to link tvm_ffi::header into the target (default: ON; ON/OFF-style) # DEBUG_SYMBOL: Whether to enable debug symbol post-processing hooks. # On Apple this calls tvm_ffi_add_apple_dsymutil(target_name) (default: ON; ON/OFF-style) # On non-Apple platforms this is currently a no-op unless you extend it. (default: ON) # MSVC_FLAGS: Whether to call tvm_ffi_add_msvc_flags(target_name) to apply MSVC-specific flags (default: ON; ON/OFF-style) # STUB_DIR: Stub generation runs when this is set. Directory to generate Python stubs. Relative paths resolve against CMAKE_CURRENT_SOURCE_DIR. # STUB_INIT: Whether to allow generating new directives. Default: OFF (ON/OFF-style) # STUB_PKG: Package name passed to stub generator (requires STUB_DIR and STUB_INIT=ON; default: ${SKBUILD_PROJECT_NAME} if set, otherwise target name) # STUB_PREFIX: Module prefix passed to stub generator (requires STUB_DIR and STUB_INIT=ON; default: ".") # ~~~ function (tvm_ffi_configure_target target) if (NOT target) message( FATAL_ERROR "tvm_ffi_configure_target: missing target name. " "Usage: tvm_ffi_configure_target( [LINK_SHARED ON|OFF] [LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF] [MSVC_FLAGS ON|OFF] [STUB_INIT ON|OFF] [STUB_DIR ] [STUB_PKG ] [STUB_PREFIX ])" ) endif () if (NOT TARGET "${target}") message(FATAL_ERROR "tvm_ffi_configure_target: '${target}' is not an existing CMake target.") endif () # Parse keyword args after the positional target name. set(tvm_ffi_arg_options) # none; require explicit ON/OFF style values set(tvm_ffi_arg_oneValueArgs LINK_SHARED LINK_HEADER DEBUG_SYMBOL MSVC_FLAGS STUB_INIT STUB_DIR STUB_PKG STUB_PREFIX ) set(tvm_ffi_arg_multiValueArgs) cmake_parse_arguments( tvm_ffi_arg_ "${tvm_ffi_arg_options}" "${tvm_ffi_arg_oneValueArgs}" "${tvm_ffi_arg_multiValueArgs}" ${ARGN} ) # Defaults foreach (arg IN ITEMS LINK_SHARED LINK_HEADER DEBUG_SYMBOL MSVC_FLAGS) if (NOT DEFINED tvm_ffi_arg__${arg}) set(tvm_ffi_arg__${arg} ON) endif () endforeach () if (NOT DEFINED tvm_ffi_arg__STUB_INIT) set(tvm_ffi_arg__STUB_INIT OFF) endif () # Validation if ((NOT DEFINED tvm_ffi_arg__STUB_DIR) OR (NOT tvm_ffi_arg__STUB_DIR)) if (DEFINED tvm_ffi_arg__STUB_PKG OR DEFINED tvm_ffi_arg__STUB_PREFIX) message( FATAL_ERROR "tvm_ffi_configure_target(${target}): STUB_PKG/STUB_PREFIX require STUB_DIR to be set." ) endif () endif () if (NOT tvm_ffi_arg__STUB_INIT) if (DEFINED tvm_ffi_arg__STUB_PKG OR DEFINED tvm_ffi_arg__STUB_PREFIX) message( FATAL_ERROR "tvm_ffi_configure_target(${target}): STUB_PKG/STUB_PREFIX cannot be set when STUB_INIT is OFF." ) endif () else () if (NOT DEFINED tvm_ffi_arg__STUB_DIR OR NOT tvm_ffi_arg__STUB_DIR) message( FATAL_ERROR "tvm_ffi_configure_target(${target}): STUB_INIT=ON requires STUB_DIR to be set." ) endif () endif () # STUB_PKG and STUB_PREFIX defaults if (tvm_ffi_arg__STUB_INIT AND tvm_ffi_arg__STUB_DIR) if (NOT DEFINED tvm_ffi_arg__STUB_PKG) if (DEFINED SKBUILD_PROJECT_NAME AND SKBUILD_PROJECT_NAME) set(tvm_ffi_arg__STUB_PKG "${SKBUILD_PROJECT_NAME}") else () set(tvm_ffi_arg__STUB_PKG "${target}") endif () endif () if (NOT DEFINED tvm_ffi_arg__STUB_PREFIX) set(tvm_ffi_arg__STUB_PREFIX "${tvm_ffi_arg__STUB_PKG}.") endif () endif () # Always-on prefix map if (COMMAND tvm_ffi_add_prefix_map) tvm_ffi_add_prefix_map("${target}" "${CMAKE_CURRENT_SOURCE_DIR}") else () message( FATAL_ERROR "tvm_ffi_configure_target(${target}): required function 'tvm_ffi_add_prefix_map' is not defined/included." ) endif () # LINK_HEADER if (tvm_ffi_arg__LINK_HEADER) if (TARGET tvm_ffi::header) target_link_libraries("${target}" PRIVATE tvm_ffi::header) else () message( FATAL_ERROR "tvm_ffi_configure_target(${target}): LINK_HEADER requested but targets 'tvm_ffi::header' do not exist." ) endif () endif () # LINK_SHARED if (tvm_ffi_arg__LINK_SHARED) if (TARGET tvm_ffi::shared) target_link_libraries("${target}" PRIVATE tvm_ffi::shared) else () message( FATAL_ERROR "tvm_ffi_configure_target(${target}): LINK_SHARED requested but targets 'tvm_ffi::shared' do not exist." ) endif () endif () # DEBUG_SYMBOL (default ON). Apple behavior only (hook only; installation handled by # tvm_ffi_install()). if (tvm_ffi_arg__DEBUG_SYMBOL) if (APPLE) if (COMMAND tvm_ffi_add_apple_dsymutil) tvm_ffi_add_apple_dsymutil("${target}") else () message( FATAL_ERROR "tvm_ffi_configure_target(${target}): DEBUG_SYMBOL=ON but 'tvm_ffi_add_apple_dsymutil' is not defined/included." ) endif () endif () endif () # Optional: MSVC flags if (tvm_ffi_arg__MSVC_FLAGS) if (COMMAND tvm_ffi_add_msvc_flags) tvm_ffi_add_msvc_flags("${target}") else () message( FATAL_ERROR "tvm_ffi_configure_target(${target}): MSVC_FLAGS=ON but 'tvm_ffi_add_msvc_flags' is not defined/included." ) endif () endif () if (DEFINED tvm_ffi_arg__STUB_DIR AND tvm_ffi_arg__STUB_DIR) get_filename_component( tvm_ffi_arg__STUB_DIR_ABS "${tvm_ffi_arg__STUB_DIR}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" ) find_package( Python3 COMPONENTS Interpreter REQUIRED ) set(tvm_ffi_stub_cli_args "${tvm_ffi_arg__STUB_DIR_ABS}" --dlls $) if (tvm_ffi_arg__STUB_INIT) list( APPEND tvm_ffi_stub_cli_args --init-lib ${target} --init-pypkg "${tvm_ffi_arg__STUB_PKG}" --init-prefix "${tvm_ffi_arg__STUB_PREFIX}" ) endif () add_custom_command( TARGET ${target} POST_BUILD COMMAND ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args} COMMENT "[COMMAND] Running: ${Python3_EXECUTABLE} -m tvm_ffi.stub.cli ${tvm_ffi_stub_cli_args}" VERBATIM ) endif () endfunction () # ~~~ # tvm_ffi_install(target_name [DESTINATION ]) # Install TVM-FFI related artifacts for a configured target. # # Parameters: # target_name: Existing CMake target whose artifacts should be installed # # Keyword parameters: # DESTINATION: Install destination directory relative to CMAKE_INSTALL_PREFIX (default: ".") # # Behavior: # - On Apple, installs the target's dSYM bundle if it exists. # This uses generator expressions and OPTIONAL so it does not fail if the dSYM is absent. # - On non-Apple platforms, currently no-op (extend as needed for PDB/DWARF packaging). # # Notes: # - This function does not create dSYMs; it only installs them if present. # Pair it with tvm_ffi_configure_target(... DEBUG_SYMBOL ON) to enable dSYM generation hooks. # ~~~ function (tvm_ffi_install target) if (NOT target) message( FATAL_ERROR "tvm_ffi_install: missing target name. Usage: tvm_ffi_install( [DESTINATION ])" ) endif () if (NOT TARGET "${target}") message(FATAL_ERROR "tvm_ffi_install: '${target}' is not an existing CMake target.") endif () set(tvm_ffi_install_options) # none set(tvm_ffi_install_oneValueArgs DESTINATION) set(tvm_ffi_install_multiValueArgs) cmake_parse_arguments( tvm_ffi_install_ "${tvm_ffi_install_options}" "${tvm_ffi_install_oneValueArgs}" "${tvm_ffi_install_multiValueArgs}" ${ARGN} ) if (NOT DEFINED tvm_ffi_install__DESTINATION) set(tvm_ffi_install__DESTINATION ".") endif () if (APPLE) # Install target dSYM bundle if present. install( DIRECTORY "$.dSYM" DESTINATION "${tvm_ffi_install__DESTINATION}" OPTIONAL ) endif () endfunction () tvm-ffi-0.1.12/cmake/Utils/ObjectCopyUtil.cmake000066400000000000000000000040231521067262500212310ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # We need this to simulate `CUDA_{CUBIN,FATBIN}_COMPILATION` in `add_tvm_ffi_{cubin,fatbin}`, to # copy `a.cu.o` to `a.cubin`/`a.fatbin`. # Usage: cmake -DOBJECTS=;...; -DOUT_DIR= # -DEXT= -P # Parameter: OBJECTS: semicolon-separated list of input object files; OUT_DIR: output directory, # empty for the same directory as the object file EXT: extension to rename to string(REPLACE "\"" "" ext_strip "${EXT}") string(REPLACE "\"" "" out_dir_strip "${OUT_DIR}") foreach (obj_raw ${OBJECTS}) string(REPLACE "\"" "" obj "${obj_raw}") # Extract filename: /path/to/kernel.cu.o -> kernel Note: CMake objects are usually named # source.cu.o, so we strip extensions twice. get_filename_component(fname ${obj} NAME_WE) get_filename_component(fname ${fname} NAME_WE) # If OUT_DIR is provided, use it. Otherwise, use the object's directory. if (NOT out_dir_strip STREQUAL "") set(FINAL_DIR "${out_dir_strip}") else () get_filename_component(FINAL_DIR ${obj} DIRECTORY) endif () message("Copying ${obj} -> ${FINAL_DIR}/${fname}.${ext_strip}") execute_process( COMMAND ${CMAKE_COMMAND} -E copy_if_different "${obj}" "${FINAL_DIR}/${fname}.${ext_strip}" ) endforeach () tvm-ffi-0.1.12/cmake/Utils/Sanitizer.cmake000066400000000000000000000035311521067262500203050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ~~~ # add_sanitizer_address(target_name) # Enable AddressSanitizer (ASan) for the given target when supported by the # current compiler. Safe no-op on unsupported compilers. # # Parameters: # target_name: CMake target to instrument with ASan # ~~~ function (add_sanitizer_address target_name) if (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|AppleClang") include(CheckCXXCompilerFlag) set(saved_crf ${CMAKE_REQUIRED_FLAGS}) set(CMAKE_REQUIRED_FLAGS "-fsanitize=address") # cmake-lint: disable=C0103 check_cxx_source_compiles("int main() { return 0; }" COMPILER_SUPPORTS_ASAN) set(CMAKE_REQUIRED_FLAGS ${saved_crf}) # cmake-lint: disable=C0103 get_target_property(saved_type ${target_name} TYPE) if (${saved_type} STREQUAL "INTERFACE_LIBRARY") set(saved_type INTERFACE) else () set(saved_type PRIVATE) endif () target_link_options(${target_name} ${saved_type} "-fsanitize=address") target_compile_options( ${target_name} ${saved_type} "-fsanitize=address" "-fno-omit-frame-pointer" "-g" ) return() endif () endfunction () tvm-ffi-0.1.12/cmake/tvm_ffi-config.cmake000066400000000000000000000044441521067262500201360ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. find_package( Python COMPONENTS Interpreter REQUIRED ) # call tvm_ffi.config to get the cmake directory and set it to tvm_ffi_ROOT execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --includedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_INCLUDE_DIR ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --dlpack-includedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_DLPACK_INCLUDE_DIR ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --libfiles OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_LIB_FILES ) message(STATUS "Finding libfiles ${tvm_ffi_LIB_FILES}") add_library(tvm_ffi::header INTERFACE IMPORTED) target_compile_features(tvm_ffi::header INTERFACE cxx_std_17) target_include_directories(tvm_ffi::header INTERFACE "${tvm_ffi_INCLUDE_DIR}") target_include_directories(tvm_ffi::header INTERFACE "${tvm_ffi_DLPACK_INCLUDE_DIR}") add_library(tvm_ffi::shared SHARED IMPORTED) target_compile_features(tvm_ffi::shared INTERFACE cxx_std_17) if (WIN32) set_target_properties(tvm_ffi::shared PROPERTIES IMPORTED_IMPLIB "${tvm_ffi_LIB_FILES}") else () set_target_properties(tvm_ffi::shared PROPERTIES IMPORTED_LOCATION "${tvm_ffi_LIB_FILES}") endif () set_target_properties( tvm_ffi::shared PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${tvm_ffi_INCLUDE_DIR};${tvm_ffi_DLPACK_INCLUDE_DIR}" ) include(${CMAKE_CURRENT_LIST_DIR}/Utils/Library.cmake) include(${CMAKE_CURRENT_LIST_DIR}/Utils/EmbedCubin.cmake) tvm-ffi-0.1.12/docs/000077500000000000000000000000001521067262500141015ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/.gitignore000066400000000000000000000000261521067262500160670ustar00rootroot00000000000000_build **/generated/* tvm-ffi-0.1.12/docs/.rstcheck.cfg000066400000000000000000000005011521067262500164420ustar00rootroot00000000000000[rstcheck] report_level = warning ignore_directives = automodule, autosummary, currentmodule, toctree, ifconfig, tab-set, collapse, tabs, dropdown, mermaid ignore_roles = ref, cpp:class, cpp:func, py:func, c:macro, external+data-api:doc, external+scikit_build_core:doc, external+dlpack:doc ignore_languages = cpp, python tvm-ffi-0.1.12/docs/README.md000066400000000000000000000017401521067262500153620ustar00rootroot00000000000000 # TVM FFI Documentation See the full guide at: tvm-ffi-0.1.12/docs/_static/000077500000000000000000000000001521067262500155275ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/_static/custom.css000066400000000000000000000002721521067262500175540ustar00rootroot00000000000000/* Avoid the uncessary scrollbar on the primiary sidebar See: https://github.com/executablebooks/sphinx-book-theme/issues/732 */ #rtd-footer-container { margin: 0px !important; } tvm-ffi-0.1.12/docs/_stubs/000077500000000000000000000000001521067262500154005ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/_stubs/cpp_index.rst000066400000000000000000000021311521067262500201000ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. :orphan: ============== Full API Index ============== This page contains stub C++ API indexs to avoid Sphinx warning in CI testing. If you see this page, it means the C++ API docs are not generated via Doxygen. Follow the instructions in `docs/reference/cpp/README.md` to generate the C++ API docs. tvm-ffi-0.1.12/docs/concepts/000077500000000000000000000000001521067262500157175ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/concepts/abi_overview.rst000066400000000000000000000534301521067262500211370ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ABI Overview ============ .. hint:: Authoritative ABI specifications are defined in - C header `tvm/ffi/c_api.h `_, which contains the core ABI, and - C header `tvm/ffi/extra/c_env_api.h `_, which contains extra support features. The TVM-FFI ABI is designed around the following key principles: - **Minimal and efficient.** Keep things simple and deliver close-to-metal performance. - **Stability guarantee.** The ABI remains stable across compiler versions and is independent of host languages or frameworks. - **Expressive for machine learning.** Native support for tensors, shapes, and data types commonly used in ML workloads. - **Extensible.** The ABI supports user-defined types and features through a dynamic type registration system. This tutorial covers common concepts and usage patterns of the TVM-FFI ABI, with low-level C code examples for precise reference. .. important:: C code is used for clarity, precision and friendliness to compiler builders. And C code can be readily translated into code generators such as LLVM IR builder. Any and AnyView --------------- .. seealso:: :doc:`any` for :cpp:class:`~tvm::ffi::Any` and :cpp:class:`~tvm::ffi::AnyView` usage patterns. At the core of TVM-FFI is :cpp:class:`TVMFFIAny`, a 16-byte tagged union that can hold any value recognized by the FFI system. It enables type-erased value passing across language boundaries. .. dropdown:: C ABI Reference: :cpp:class:`TVMFFIAny` :icon: code .. literalinclude:: ../../include/tvm/ffi/c_api.h :language: c :start-after: [TVMFFIAny.begin] :end-before: [TVMFFIAny.end] :caption: tvm/ffi/c_api.h **Ownership.** :cpp:class:`TVMFFIAny` struct can represent either an owning or a borrowing reference. These two ownership patterns are formalized by the C++ wrapper classes :cpp:class:`~tvm::ffi::Any` and :cpp:class:`~tvm::ffi::AnyView`, which have identical memory layouts but different :ref:`ownership semantics `. See :doc:`any` for high-level C++ usage patterns: - **Owning:** :cpp:class:`tvm::ffi::Any` - reference-counted, manages object lifetime - **Borrowing:** :cpp:class:`tvm::ffi::AnyView` - non-owning view, caller must ensure validity .. note:: To convert a borrowing :cpp:class:`~tvm::ffi::AnyView` to an owning :cpp:class:`~tvm::ffi::Any`, use :cpp:func:`TVMFFIAnyViewToOwnedAny`. **Runtime Type Index.** The ``type_index`` field identifies what kind of value is stored: - :ref:`Atomic POD types ` (``type_index`` < :cpp:enumerator:`kTVMFFIStaticObjectBegin `): Stored inline in the payload union without heap allocation or reference counting. - :ref:`Object types ` (``type_index`` >= :cpp:enumerator:`kTVMFFIStaticObjectBegin `): Stored as pointers to heap-allocated, reference-counted TVM-FFI objects. .. important:: The TVM-FFI type index system does not rely on C++ RTTI. Construct Any ~~~~~~~~~~~~~ **From atomic POD types.** The following C code constructs a :cpp:class:`TVMFFIAny` from an integer: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.FromInt_Float.begin] :end-before: [Any_AnyView.FromInt_Float.end] Set the ``type_index`` from :cpp:enum:`TVMFFITypeIndex` and assign the corresponding payload field. .. important:: Always zero the ``zero_padding`` field and any unused bytes in the value union. This invariant enables direct byte comparison and hashing of :cpp:class:`TVMFFIAny` values. **From object types.** The following C code constructs a :cpp:class:`TVMFFIAny` from a heap-allocated object: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.FromObjectPtr.begin] :end-before: [Any_AnyView.FromObjectPtr.end] When ``IS_OWNING_ANY`` is ``true`` (owning :cpp:class:`~tvm::ffi::Any`), this increments the object's reference count. .. _abi-destruct-any: Destruct Any ~~~~~~~~~~~~ The following C code destroys a :cpp:class:`TVMFFIAny`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.Destroy.begin] :end-before: [Any_AnyView.Destroy.end] When ``IS_OWNING_ANY`` is ``true`` (owning :cpp:class:`~tvm::ffi::Any`), this decrements the object's reference count. Extract from Any ~~~~~~~~~~~~~~~~ **Extract an atomic POD.** The following C code extracts an integer or float from a :cpp:class:`TVMFFIAny`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.GetInt_Float.begin] :end-before: [Any_AnyView.GetInt_Float.end] Implicit type conversion may occur. For example, when extracting a float from a :cpp:class:`TVMFFIAny` that holds an integer, the integer is cast to a float. **Extract a DLTensor.** A :c:struct:`DLTensor` may originate from either a raw pointer or a heap-allocated :cpp:class:`~tvm::ffi::TensorObj`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.GetDLTensor.begin] :end-before: [Any_AnyView.GetDLTensor.end] **Extract a TVM-FFI object.** TVM-FFI objects are always heap-allocated and reference-counted, with ``type_index`` >= :cpp:enumerator:`kTVMFFIStaticObjectBegin `: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Any_AnyView.GetObject.begin] :end-before: [Any_AnyView.GetObject.end] To take ownership of the returned value, increment the reference count via :cpp:func:`TVMFFIObjectIncRef`. Release ownership later via :cpp:func:`TVMFFIObjectDecRef`. .. _abi-object: Object ------ .. seealso:: :doc:`object_and_class` for the object system and reflection. TVM-FFI Object (:cpp:class:`TVMFFIObject`) is the cornerstone of TVM-FFI's stable yet extensible type system. .. dropdown:: C ABI Reference: :cpp:class:`TVMFFIObject` :icon: code .. literalinclude:: ../../include/tvm/ffi/c_api.h :language: c :start-after: [TVMFFIObject.begin] :end-before: [TVMFFIObject.end] :caption: tvm/ffi/c_api.h All TVM-FFI objects share these characteristics: - Heap-allocated and reference-counted - Layout-stable 24-byte header containing reference counts, type index, and deleter callback - Type index >= :cpp:enumerator:`kTVMFFIStaticObjectBegin ` **Dynamic Type System.** Classes can be registered at runtime via :cpp:func:`TVMFFITypeGetOrAllocIndex`, with support for single inheritance. See :doc:`object_and_class` for the full object system and :ref:`type-checking-and-casting` for usage details. A small **static section** between :cpp:enumerator:`kTVMFFIStaticObjectBegin ` and :cpp:enumerator:`kTVMFFIDynObjectBegin ` is reserved for static object types, for example, - Strings (:cpp:enumerator:`kTVMFFIStr `) and Bytes (:cpp:enumerator:`kTVMFFIBytes `): Section :ref:`abi-string-and-byte` - Errors (:cpp:enumerator:`kTVMFFIError `): Section :ref:`abi-exception`. - Functions (:cpp:enumerator:`kTVMFFIFunction `): Section :ref:`abi-function`. - Tensors (:cpp:enumerator:`kTVMFFITensor `): Section :ref:`abi-tensor`. - Miscellaneous: Modules (:cpp:enumerator:`kTVMFFIModule `), Arrays (:cpp:enumerator:`kTVMFFIArray `), Maps (:cpp:enumerator:`kTVMFFIMap `), Shapes (:cpp:enumerator:`kTVMFFIShape `), Opaque Python objects (:cpp:enumerator:`kTVMFFIOpaquePyObject `). .. _abi-object-ownership: Ownership Management ~~~~~~~~~~~~~~~~~~~~ Ownership is managed via reference counting, which includes both strong and weak references. Two C APIs manage strong reference counting: - :cpp:func:`TVMFFIObjectIncRef`: Acquire strong ownership by incrementing the reference count - :cpp:func:`TVMFFIObjectDecRef`: Release strong ownership by decrementing the reference count The ``deleter`` callback (:cpp:member:`TVMFFIObject::deleter`) executes when the strong or weak count reaches zero with different flags. See :ref:`object-reference-counting` for details. **Move ownership from Any/AnyView.** The following C code transfers ownership from an owning :cpp:class:`~tvm::ffi::Any` to an object pointer: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Object.MoveFromAny.begin] :end-before: [Object.MoveFromAny.end] Since :cpp:class:`~tvm::ffi::AnyView` is non-owning (``IS_OWNING_ANY`` is ``false``), acquiring ownership requires explicitly incrementing the reference count. **Release ownership.** The following C code releases ownership of a TVM-FFI object: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :name: ABI.Object.Destroy :start-after: [Object.Destroy.begin] :end-before: [Object.Destroy.end] Inheritance Checking ~~~~~~~~~~~~~~~~~~~~ TVM-FFI models single inheritance as a tree where each node points to its parent. Each type has a unique type index, and the system tracks ancestors, inheritance depth, and other metadata. This information is available via :cpp:func:`TVMFFIGetTypeInfo`. The following C code checks whether a type is a subclass of another: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Object.IsInstance.begin] :end-before: [Object.IsInstance.end] .. _abi-tensor: Tensor ------ .. seealso:: :doc:`tensor` for details about TVM-FFI tensors and DLPack interoperability. TVM-FFI provides :cpp:class:`tvm::ffi::TensorObj`, a DLPack-native tensor class that is also a standard TVM-FFI object. This means tensors can be managed using the same reference counting mechanisms as other objects. .. dropdown:: C ABI Reference: :cpp:class:`tvm::ffi::TensorObj` :icon: code .. code-block:: cpp :caption: tvm/ffi/container/tensor.h class TensorObj : public Object, public DLTensor { // no other members besides those from Object and DLTensor }; Access Tensor Metadata ~~~~~~~~~~~~~~~~~~~~~~ The following C code obtains a :c:struct:`DLTensor` pointer from a :cpp:class:`~tvm::ffi::TensorObj`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Tensor.AccessDLTensor.begin] :end-before: [Tensor.AccessDLTensor.end] The :c:struct:`DLTensor` pointer provides access to shape, dtype, device, data pointer, and other tensor metadata. Construct Tensor ~~~~~~~~~~~~~~~~ **Zero-copy conversion.** The following C code constructs a :cpp:class:`~tvm::ffi::TensorObj` from a :c:struct:`DLManagedTensorVersioned`, which shares the underlying data buffer without allocating new memory. .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Tensor_FromDLPack.begin] :end-before: [Tensor_FromDLPack.end] .. hint:: TVM-FFI's Python API automatically wraps framework tensors (e.g., :py:class:`torch.Tensor`) as :cpp:class:`~tvm::ffi::TensorObj`, so manual conversion is typically unnecessary. **Allocate new memory.** Alternatively, if memory allocation is intended, the following C code constructs a :cpp:class:`~tvm::ffi::TensorObj` from a :c:struct:`DLTensor` pointer: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Tensor_Alloc.begin] :end-before: [Tensor_Alloc.end] The ``prototype`` contains the shape, dtype, device, and other tensor metadata that will be used to allocate the new tensor. And the allocator, by default, is the framework's (e.g., PyTorch) allocator, which is automatically set when importing the framework. To override or explicitly look up the allocator, use :cpp:func:`TVMFFIEnvSetDLPackManagedTensorAllocator` and :cpp:func:`TVMFFIEnvGetDLPackManagedTensorAllocator`. .. warning:: In kernel library usecases, it is usually not recommended to dynamically allocate tensors inside a kernel, and instead always pre-allocate outputs, and pass them as :cpp:class:`~tvm::ffi::TensorView` parameters. This approach - avoids memory fragmentation and performance pitfalls, - prevents CUDA graph incompatibilities on GPU, and - allows the outer framework to control allocation policy (pools, device strategies, etc.). Destruct Tensor ~~~~~~~~~~~~~~~ As a standard TVM-FFI object, :cpp:class:`~tvm::ffi::TensorObj` follows the :ref:`standard destruction pattern `. When the reference count reaches zero, the deleter callback (:cpp:member:`TVMFFIObject::deleter`) executes. Export Tensor to DLPack ~~~~~~~~~~~~~~~~~~~~~~~ To share a :cpp:class:`~tvm::ffi::TensorObj` with other frameworks, export it as a :c:struct:`DLManagedTensorVersioned`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Tensor_ToDLPackVersioned.begin] :end-before: [Tensor_ToDLPackVersioned.end] Note that the caller takes ownership of the returned :c:struct:`DLManagedTensorVersioned* ` and must call its ``deleter`` to release the tensor. .. _abi-function: Function -------- .. seealso:: :ref:`sec:function` for a detailed description of TVM-FFI functions. All functions in TVM-FFI follow a unified C calling convention that enables ABI-stable, type-erased, and cross-language function calls, defined by :cpp:type:`TVMFFISafeCallType`. **Calling convention.** The signature includes: - ``handle`` (``void*``): Optional resource handle passed to the callee; typically ``NULL`` for exported symbols - ``args`` (``TVMFFIAny*``) and ``num_args`` (``int``): Array of non-owning :cpp:class:`~tvm::ffi::AnyView` input arguments - ``result`` (``TVMFFIAny*``): Owning :cpp:class:`~tvm::ffi::Any` output value - Return value: ``0`` for success; ``-1`` for errors (see :doc:`exception_handling` and :ref:`sec-exception`) See :ref:`sec:function-calling-convention` for more details. .. important:: The caller must zero-initialize the output argument ``result`` before the call. **Memory layout.** The :cpp:class:`~tvm::ffi::FunctionObj` stores call pointers after the object header. .. dropdown:: C ABI Reference: :cpp:class:`TVMFFIFunctionCell` :icon: code .. literalinclude:: ../../include/tvm/ffi/c_api.h :language: c :start-after: [TVMFFIFunctionCell.begin] :end-before: [TVMFFIFunctionCell.end] :caption: tvm/ffi/c_api.h Construct and Destroy ~~~~~~~~~~~~~~~~~~~~~ .. important:: Dynamic function creation is useful for passing lambdas or closures across language boundaries. The following C code constructs a :cpp:class:`~tvm::ffi::FunctionObj` from a :cpp:type:`TVMFFISafeCallType` and a ``deleter`` callback. The ``deleter`` cleans up resources owned by the function; for global symbols, it is typically ``NULL``. .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Function.Construct.begin] :end-before: [Function.Construct.end] Release a :cpp:class:`~tvm::ffi::FunctionObj` using the :ref:`standard destruction pattern `. Global Registry ~~~~~~~~~~~~~~~ **Retrieve a global function.** The following C code uses :cpp:func:`TVMFFIFunctionGetGlobal` to retrieve a function by name from the global registry: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Function.GetGlobal.begin] :end-before: [Function.GetGlobal.end] .. note:: :cpp:func:`TVMFFIFunctionGetGlobal` returns an owning handle. The caller must release it by calling :cpp:func:`TVMFFIObjectDecRef` when it's no longer needed. **Register a global function.** The following C code uses :cpp:func:`TVMFFIFunctionSetGlobal` to register a function by name in the global registry: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Function.SetGlobal.begin] :end-before: [Function.SetGlobal.end] Call Function ~~~~~~~~~~~~~ The following C code invokes a :cpp:class:`~tvm::ffi::FunctionObj` with arguments: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Function.Call.begin] :end-before: [Function.Call.end] .. _abi-exception: Exception --------- .. seealso:: :doc:`exception_handling` for detailed exception handling patterns. Exceptions are a central part of TVM-FFI's ABI and calling convention. When errors occur, they are stored as objects with a :cpp:class:`TVMFFIErrorCell` payload. .. dropdown:: C ABI Reference: :cpp:class:`TVMFFIErrorCell` :icon: code .. literalinclude:: ../../include/tvm/ffi/c_api.h :language: c :start-after: [TVMFFIErrorCell.begin] :end-before: [TVMFFIErrorCell.end] :caption: tvm/ffi/c_api.h .. important:: Errors from all languages (e.g. Python, C++) will be properly translated into the TVM-FFI error object. Retrieve Error Object ~~~~~~~~~~~~~~~~~~~~~ When a function returns ``-1``, an error object is stored in thread-local storage (TLS). Retrieve it with :cpp:func:`TVMFFIErrorMoveFromRaised`, which returns a :cpp:class:`tvm::ffi::ErrorObj`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.HandleReturnCode.begin] :end-before: [Error.HandleReturnCode.end] This function transfers ownership to the caller and clears the TLS slot. Call :cpp:func:`TVMFFIObjectDecRef` when done to avoid memory leaks. .. admonition:: Print Error Message :class: hint The error payload is a :cpp:type:`TVMFFIErrorCell` structure containing the error kind, message, and backtrace. Access it by skipping the :cpp:type:`TVMFFIObject` header via pointer arithmetic. .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.Print.begin] :end-before: [Error.Print.end] This prints the error message along with its backtrace. Raise Exception ~~~~~~~~~~~~~~~ The following C code sets the TLS error and returns ``-1`` via :cpp:func:`TVMFFIErrorSetRaisedFromCStr`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.RaiseException.begin] :end-before: [Error.RaiseException.end] For non-null-terminated strings, use :cpp:func:`TVMFFIErrorSetRaisedFromCStrParts`, which accepts explicit string lengths. .. note:: You rarely need to create a :cpp:class:`~tvm::ffi::ErrorObj` directly. The C APIs :cpp:func:`TVMFFIErrorSetRaisedFromCStr` and :cpp:func:`TVMFFIErrorSetRaisedFromCStrParts` handle this internally. .. _abi-string-and-byte: 🚧 String and Bytes ------------------- .. warning:: This section is under construction. The ABI supports strings and bytes as first-class citizens. A string can take multiple forms that are identified by its ``type_index``. - ``kTVMFFIRawStr``: raw C string terminated by ``\0``. - ``kTVMFFISmallStr``: small string, the length is stored in ``small_str_len`` and data is stored in ``v_bytes``. - ``kTVMFFIStr``: on-heap string object for strings that are longer than 7 characters. The following code shows the layout of the on-heap string object. .. code-block:: cpp // span-like data structure to store header and length typedef struct { const char* data; size_t size; } TVMFFIByteArray; // showcase the layout of the on-heap string. class StringObj : public ffi::Object, public TVMFFIByteArray { }; The following code shows how to read a string from :cpp:class:`TVMFFIAny` .. code-block:: cpp TVMFFIByteArray ReadString(const TVMFFIAny *value) { TVMFFIByteArray ret; if (value->type_index == kTVMFFIRawStr) { ret.data = value->v_c_str; ret.size = strlen(ret.data); } else if (value->type_index == kTVMFFISmallStr) { ret.data = value->v_bytes; ret.size = value->small_str_len; } else { assert(value->type_index == kTVMFFIStr); ret = *reinterpret_cast( reinterpret_cast(value->v_obj) + sizeof(TVMFFIObject)); } return ret; } Similarly, we have type indices to represent bytes. The C++ API provides classes :cpp:class:`~tvm::ffi::String` and :cpp:class:`~tvm::ffi::Bytes` to enable the automatic conversion of these values with Any storage format. **Rationales**. Separate string and bytes enable clear mappings from the Python side. Small string allows us to store short names on-stack. To favor 8-byte alignment (v_bytes) and keep things simple, we did not further pack characters into the ``small_len`` field. Further Reading --------------- - :doc:`any`: High-level C++ usage of :cpp:class:`~tvm::ffi::Any` and :cpp:class:`~tvm::ffi::AnyView` - :doc:`object_and_class`: The object system and reflection - :doc:`tensor`: Tensor classes and DLPack interoperability - :doc:`func_module`: Functions and modules - :doc:`exception_handling`: Exception handling across language boundaries - :doc:`../get_started/stable_c_abi`: Quick introduction to the stable C ABI tvm-ffi-0.1.12/docs/concepts/any.rst000066400000000000000000000372651521067262500172550ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Any and AnyView =============== TVM-FFI has :cpp:class:`tvm::ffi::Any` and :cpp:class:`tvm::ffi::AnyView`, type-erased containers that hold any supported value and transport it across C, C++, Python, and Rust boundaries through a stable ABI. Similar to ``std::any``, :cpp:class:`~tvm::ffi::Any` is a tagged union that stores values of a wide variety of types, including primitives, objects, and strings. Unlike ``std::any``, it is designed for zero-copy inter-language exchange without RTTI, featuring a fixed 16-byte layout with built-in reference counting and ownership semantics. This tutorial covers common usage patterns, ownership semantics, and memory layout. Common Usage ------------ Function Signatures ~~~~~~~~~~~~~~~~~~~ Use :cpp:class:`~tvm::ffi::AnyView` for function parameters to avoid reference count overhead and unnecessary copies. Use :cpp:class:`~tvm::ffi::Any` for return values to transfer ownership to the caller. .. code-block:: cpp ffi::Any func_cpp_signature(ffi::AnyView arg0, ffi::AnyView arg1) { ffi::Any result = arg0.cast() + arg1.cast(); return result; } // Variant: variadic function void func_cpp_variadic(PackedArgs args, Any* ret) { int32_t num_args = args.size(); int x0 = args[0].cast(); int x1 = args[1].cast(); int y = x0 + x1; *ret = y; } // Variant: variadic function with C ABI signature int func_c_abi_variadic(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* ret) { TVM_FFI_SAFE_CALL_BEGIN(); int x0 = reinterpret_cast(args)[0].cast(); int x1 = reinterpret_cast(args)[1].cast(); int y = x0 + x1; reinterpret_cast(ret)[0] = y; TVM_FFI_SAFE_CALL_END(); } Container Storage ~~~~~~~~~~~~~~~~~ :cpp:class:`~tvm::ffi::Any` can be stored in containers like :cpp:class:`~tvm::ffi::Array`, :cpp:class:`~tvm::ffi::List`, :cpp:class:`~tvm::ffi::Map`, and :cpp:class:`~tvm::ffi::Dict` (see :doc:`containers` for details): .. code-block:: cpp ffi::Map config; config.Set("learning_rate", 0.001); config.Set("batch_size", 32); config.Set("device", DLDevice{kDLCUDA, 0}); Extracting Values ~~~~~~~~~~~~~~~~~ Three methods extract values from :cpp:class:`~tvm::ffi::Any` and :cpp:class:`~tvm::ffi::AnyView`, each with different levels of strictness: .. list-table:: :header-rows: 1 :widths: 20 45 30 * - Method - Behavior - Use When * - :cpp:func:`cast\() ` - Returns ``T`` or throws :cpp:class:`tvm::ffi::Error` - When you know the expected type and want an exception on mismatch * - :cpp:func:`try_cast\() ` - Returns ``std::optional`` - When you want graceful failure and allow type conversions (e.g., int to double) * - :cpp:func:`as\() ` - Returns ``std::optional`` or ``const T*`` (for TVM-FFI object types) - When you need an exact type match with no conversions .. dropdown:: Example of :cpp:func:`cast\() ` :cpp:func:`cast\() ` is the workhorse. It returns the value or throws a :cpp:class:`tvm::ffi::Error` (see :doc:`exception_handling`): .. code-block:: cpp ffi::Any value = 42; int x = value.cast(); // OK: 42 double y = value.cast(); // OK: 42.0 (int → double) try { ffi::String s = value.cast(); // Throws TypeError } catch (const ffi::Error& e) { // "Cannot convert from type `int` to `ffi.Str`" } .. dropdown:: Example of :cpp:func:`try_cast\() ` :cpp:func:`try_cast\() ` allows type coercion: .. code-block:: cpp ffi::Any value = 42; std::optional opt_float = value.try_cast(); // opt_float.has_value() == true, *opt_float == 42.0 std::optional opt_bool = value.try_cast(); // opt_bool.has_value() == true, *opt_bool == true .. dropdown:: Example of :cpp:func:`as\() ` :cpp:func:`as\() ` is strict - it succeeds only if the stored type matches exactly: .. code-block:: cpp ffi::Any value = 42; std::optional opt_int = value.as(); // opt_int.has_value() == true std::optional opt_float = value.as(); // opt_float.has_value() == false (int stored, not float) ffi::Any str_value = ffi::String("hello, world!"); if (const ffi::Object* obj = str_value.as()) { // Use obj without copying } Nullability Checks ~~~~~~~~~~~~~~~~~~ Compare with ``nullptr`` to check for ``None``: .. code-block:: cpp ffi::Any value = std::nullopt; if (value == nullptr) { // Handle None case } else { // Process value } .. _any-ownership: Ownership --------- The core distinction between :cpp:class:`tvm::ffi::Any` and :cpp:class:`tvm::ffi::AnyView` is **ownership**: .. list-table:: :header-rows: 1 :widths: 25 35 40 * - Aspect - :cpp:class:`~tvm::ffi::AnyView` - :cpp:class:`~tvm::ffi::Any` * - Ownership - Non-owning (like ``std::string_view``) - Owning (like ``std::string``) * - Reference counting - No reference count changes on copy - Increments reference count on copy; decrements on destroy * - Lifetime - Valid only while source lives - Extends object lifetime * - Primary use - Function inputs - Return values, storage Examples ~~~~~~~~ :cpp:class:`~tvm::ffi::AnyView` is a lightweight, non-owning view. Copying it simply copies 16 bytes with no reference count updates, making it ideal for passing arguments without overhead: .. code-block:: cpp void process(ffi::AnyView value) {} :cpp:class:`~tvm::ffi::Any` is an owning container. Copying an :cpp:class:`~tvm::ffi::Any` that holds an object increments the reference count; destroying it decrements the count (see :ref:`object-reference-counting` for details on how reference counting works): .. code-block:: cpp ffi::Any create_value() { ffi::Any result; { ffi::String str = "hello"; // refcount = 1 (str created) result = str; // refcount 1 -> 2 (result owns str) } // refcount 2 -> 1 (str is destroyed) return result; // refcount = 1 (result returns to caller) } ABI Boundary ~~~~~~~~~~~~ TVM-FFI's function calling convention follows two simple rules: - **Inputs are non-owning**: Arguments are passed as :cpp:class:`~tvm::ffi::AnyView`. The caller retains ownership, and the callee borrows them for the duration of the call. - **Outputs are owning**: Return values are passed as :cpp:class:`~tvm::ffi::Any`. Ownership transfers to the caller, who becomes responsible for managing the value's lifetime. .. code-block:: cpp // TVM-FFI C ABI int32_t tvm_ffi_c_abi( void* handle, const AnyView* args, // (Non-owning) args: AnyView[num_args] int32_t num_args, Any* result, // (Owning) result: Any (caller takes ownership) ); Destruction Semantics in C ~~~~~~~~~~~~~~~~~~~~~~~~~~ In C, which lacks RAII, you must manually destroy :cpp:class:`~tvm::ffi::Any` objects by calling :cpp:func:`TVMFFIObjectDecRef` for heap-allocated objects. Destroying an :cpp:class:`~tvm::ffi::AnyView` is effectively a no-op - just clear its contents. See :ref:`abi-destruct-any` for C code examples. Layout ------ Tagged Union ~~~~~~~~~~~~ At C ABI level, every value lives in a :cpp:class:`TVMFFIAny`: .. code-block:: cpp typedef struct TVMFFIAny { int32_t type_index; // Bytes 0-3: identifies the stored type union { uint32_t zero_padding; // Bytes 4-7: must be zero (or small_str_len) uint32_t small_str_len; }; union { // Bytes 8-15: the actual value int64_t v_int64; double v_float64; void* v_ptr; DLDataType v_dtype; DLDevice v_device; TVMFFIObject* v_obj; // ... other union members }; } TVMFFIAny; .. tip:: Think of :cpp:class:`TVMFFIAny` as the "layout format", and :cpp:class:`~tvm::ffi::Any`/:cpp:class:`~tvm::ffi::AnyView` as a thin "application layer" that adds type safety, RAII, and ergonomic APIs, which has no change to the layout. .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tvm-ffi/stable-c-abi-layout-any.svg :alt: Layout of the 128-bit Any tagged union :align: center Figure 1. Layout of the :cpp:class:`TVMFFIAny` tagged union in C ABI. :cpp:class:`~tvm::ffi::Any`/:cpp:class:`~tvm::ffi::AnyView` shares the same layout as :cpp:class:`TVMFFIAny`, but adds extra C++ APIs on top of it for type safety, RAII, and ergonomics. It is effectively a layout-stable 16-byte tagged union. * The first 4 bytes (:cpp:member:`TVMFFIAny::type_index`) serve as a tag identifying the stored type. * The last 8 bytes hold the actual value - either stored inline for atomic types (e.g., ``int64_t``, ``float64``, ``void*``) or as a pointer to a heap-allocated object. .. _any-atomic-types: Atomic Types ~~~~~~~~~~~~ Primitive values - integers, floats, booleans, devices, and raw pointers - are stored directly in the 8-byte payload with no heap allocation and no reference counting. .. list-table:: Figure 2. Common atomic types stored directly in :cpp:class:`TVMFFIAny` :header-rows: 1 :name: atomic-types-table :widths: 40 40 30 * - Type - type_index - Payload Field * - ``None`` / ``nullptr`` - :cpp:enumerator:`kTVMFFINone ` = 0 - :cpp:member:`~TVMFFIAny::v_int64` (must be 0) * - ``int64_t`` - :cpp:enumerator:`kTVMFFIInt ` = 1 - :cpp:member:`~TVMFFIAny::v_int64` * - ``bool`` - :cpp:enumerator:`kTVMFFIBool ` = 2 - :cpp:member:`~TVMFFIAny::v_int64` (0 or 1) * - ``float64_t`` - :cpp:enumerator:`kTVMFFIFloat ` = 3 - :cpp:member:`~TVMFFIAny::v_float64` * - ``void*`` (opaque pointer) - :cpp:enumerator:`kTVMFFIOpaquePtr ` = 4 - :cpp:member:`~TVMFFIAny::v_ptr` * - :c:struct:`DLDataType ` - :cpp:enumerator:`kTVMFFIDataType ` = 5 - :cpp:member:`~TVMFFIAny::v_dtype` * - :c:struct:`DLDevice ` - :cpp:enumerator:`kTVMFFIDevice ` = 6 - :cpp:member:`~TVMFFIAny::v_device` * - :c:struct:`DLTensor* ` - :cpp:enumerator:`kTVMFFIDLTensorPtr ` = 7 - :cpp:member:`~TVMFFIAny::v_ptr` * - ``const char*`` (raw string) - :cpp:enumerator:`kTVMFFIRawStr ` = 8 - :cpp:member:`~TVMFFIAny::v_c_str` * - :cpp:class:`TVMFFIByteArray* ` - :cpp:enumerator:`kTVMFFIByteArrayPtr ` = 9 - :cpp:member:`~TVMFFIAny::v_ptr` :ref:`Figure 2 ` shows common atomic types stored in-place inside the :cpp:class:`TVMFFIAny` payload. .. code-block:: cpp AnyView int_val = 42; // v_int64 = 42 AnyView float_val = 3.14; // v_float64 = 3.14 AnyView bool_val = true; // v_int64 = 1 AnyView device = DLDevice{kDLCUDA, 0}; // v_device DLTensor tensor; AnyView view = &tensor; // v_ptr = &tensor Note that raw pointers like :c:struct:`DLTensor* ` and ``char*`` also fit here. These pointers carry no ownership, so the caller must ensure the pointed-to data outlives the :cpp:class:`~tvm::ffi::AnyView` or :cpp:class:`~tvm::ffi::Any`. .. _any-heap-allocated-objects: Heap-Allocated Objects ~~~~~~~~~~~~~~~~~~~~~~ TVM-FFI objects are heap-allocated, reference-counted containers that inherit from :cpp:class:`tvm::ffi::Object`. See :doc:`object_and_class` for details on the object system, and :ref:`object-reference-counting` for how reference counting works. .. list-table:: Figure 3. Common TVM-FFI object types stored as pointers in :cpp:member:`TVMFFIAny::v_obj`. :header-rows: 1 :widths: 40 40 30 * - Type - type_index - Payload Field * - :cpp:class:`ErrorObj* ` - :cpp:enumerator:`kTVMFFIError ` = 67 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`FunctionObj* ` - :cpp:enumerator:`kTVMFFIFunction ` = 68 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`TensorObj* ` - :cpp:enumerator:`kTVMFFITensor ` = 70 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`ArrayObj* ` - :cpp:enumerator:`kTVMFFIArray ` = 71 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`MapObj* ` - :cpp:enumerator:`kTVMFFIMap ` = 72 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`ModuleObj* ` - :cpp:enumerator:`kTVMFFIModule ` = 73 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`ListObj* ` - :cpp:enumerator:`kTVMFFIList ` = 75 - :cpp:member:`~TVMFFIAny::v_obj` * - :cpp:class:`DictObj* ` - :cpp:enumerator:`kTVMFFIDict ` = 76 - :cpp:member:`~TVMFFIAny::v_obj` Heap-allocated objects - :cpp:class:`~tvm::ffi::String`, :cpp:class:`~tvm::ffi::Function`, :cpp:class:`~tvm::ffi::Tensor`, :cpp:class:`~tvm::ffi::Array`, :cpp:class:`~tvm::ffi::List`, :cpp:class:`~tvm::ffi::Map`, :cpp:class:`~tvm::ffi::Dict`, and custom types - are stored as pointers to reference-counted :cpp:class:`TVMFFIObject` headers: .. code-block:: cpp ffi::String str = "hello world"; ffi::Any any_str = str; // v_obj points to StringObj // Object layout in memory: // [TVMFFIObject header (24 bytes)][object-specific data] Caveats ------- Small String Optimization ~~~~~~~~~~~~~~~~~~~~~~~~~ Strings and byte arrays receive special treatment: values of 7 bytes or fewer are stored inline using **small string optimization**, avoiding heap allocation entirely: .. code-block:: cpp ffi::Any small = "hello"; // kTVMFFISmallStr, in v_bytes ffi::Any large = "this is a longer string"; // kTVMFFIStr, heap allocated Further Reading --------------- - :doc:`object_and_class`: How TVM-FFI objects work, including reference counting and type checking - :doc:`func_module`: Function calling conventions and the global registry - :doc:`tensor`: How tensors flow through :cpp:class:`~tvm::ffi::Any` and :cpp:class:`~tvm::ffi::AnyView` - :doc:`exception_handling`: How :cpp:func:`~tvm::ffi::Any::cast` throws exceptions on type mismatch - :doc:`abi_overview`: Low-level C ABI details for working with :cpp:class:`TVMFFIAny` directly tvm-ffi-0.1.12/docs/concepts/containers.rst000066400000000000000000000160771521067262500206310ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Containers ========== TVM-FFI provides five built-in container types for storing and exchanging collections of values across C++, Python, and Rust. They are all heap-allocated, reference-counted objects that can be stored in :cpp:class:`~tvm::ffi::Any` and passed through the FFI boundary. The containers split into two categories: **immutable** containers that use copy-on-write semantics, and **mutable** containers that use shared-reference semantics. Overview -------- .. list-table:: :header-rows: 1 :widths: 14 20 20 16 30 * - Type - C++ Class - Python Class - Mutability - Semantics * - Array - :cpp:class:`Array\ ` - :py:class:`tvm_ffi.Array` - Immutable - Homogeneous sequence with copy-on-write * - List - :cpp:class:`List\ ` - :py:class:`tvm_ffi.List` - Mutable - Homogeneous sequence with shared-reference * - Tuple - :cpp:class:`Tuple\ ` - (backed by :py:class:`tvm_ffi.Array`) - Immutable - Heterogeneous fixed-size sequence (backed by ArrayObj) * - Map - :cpp:class:`Map\ ` - :py:class:`tvm_ffi.Map` - Immutable - Homogeneous key-value mapping with copy-on-write * - Dict - :cpp:class:`Dict\ ` - :py:class:`tvm_ffi.Dict` - Mutable - Homogeneous key-value mapping with shared-reference Immutable Containers (Copy-on-Write) ------------------------------------- Array ~~~~~ ``Array`` is an immutable homogeneous sequence backed by :cpp:class:`~tvm::ffi::ArrayObj`. It implements **copy-on-write** semantics: when a mutation method is called in C++ (e.g. ``push_back``, ``Set``), the array checks whether the backing storage is uniquely owned. If it is shared with other handles, it copies the data first so that existing handles are unaffected. .. code-block:: cpp ffi::Array a = {1, 2, 3}; ffi::Array b = a; // b shares the same ArrayObj a.push_back(4); // copy-on-write: a gets a new backing storage assert(a.size() == 4); assert(b.size() == 3); // b is unchanged In Python, :py:class:`tvm_ffi.Array` implements ``collections.abc.Sequence`` (read-only). When a Python ``list`` or ``tuple`` is passed to an FFI function, it is automatically converted to ``Array``. Tuple ~~~~~ ``Tuple`` is an immutable heterogeneous fixed-size sequence. It is backed by the same :cpp:class:`~tvm::ffi::ArrayObj` as ``Array``, but provides compile-time type safety for each element position via C++ variadic templates. .. code-block:: cpp ffi::Tuple t(42, "hello", true); int x = t.get<0>(); // 42 ffi::String s = t.get<1>(); // "hello" In Python, ``Tuple`` does not have a separate class -- Python tuples passed through the FFI are converted to ``Array``. Map ~~~ ``Map`` is an immutable homogeneous key-value mapping backed by :cpp:class:`~tvm::ffi::MapObj`. It implements **copy-on-write** semantics (same principle as ``Array``). Insertion order is preserved. .. code-block:: cpp ffi::Map m = {{"Alice", 100}, {"Bob", 95}}; ffi::Map m2 = m; // m2 shares the same MapObj m.Set("Charlie", 88); // copy-on-write assert(m.size() == 3); assert(m2.size() == 2); // m2 is unchanged In Python, :py:class:`tvm_ffi.Map` implements ``collections.abc.Mapping`` (read-only). When a Python ``dict`` is passed to an FFI function, it is automatically converted to ``Map``. Mutable Containers (Shared Reference) -------------------------------------- List ~~~~ ``List`` is a mutable homogeneous sequence backed by :cpp:class:`~tvm::ffi::ListObj`. Unlike ``Array``, it does **not** use copy-on-write. Mutations happen directly on the underlying shared object, and **all handles** sharing the same ``ListObj`` see the mutations immediately. .. code-block:: cpp ffi::List a = {1, 2, 3}; ffi::List b = a; // b shares the same ListObj a.push_back(4); // in-place mutation assert(a.size() == 4); assert(b.size() == 4); // b sees the mutation In Python, :py:class:`tvm_ffi.List` implements ``collections.abc.MutableSequence`` and supports ``append``, ``insert``, ``__setitem__``, ``__delitem__``, ``pop``, ``reverse``, ``clear``, and ``extend``. Dict ~~~~ ``Dict`` is a mutable homogeneous key-value mapping backed by :cpp:class:`~tvm::ffi::DictObj`. Like ``List``, mutations happen directly on the shared object with no copy-on-write. .. code-block:: cpp ffi::Dict d = {{"Alice", 100}}; ffi::Dict d2 = d; // d2 shares the same DictObj d.Set("Bob", 95); // in-place mutation assert(d.size() == 2); assert(d2.size() == 2); // d2 sees the mutation In Python, :py:class:`tvm_ffi.Dict` implements ``collections.abc.MutableMapping`` and supports ``__setitem__``, ``__delitem__``, ``pop``, ``clear``, and ``update``. When to Use Each Type --------------------- .. list-table:: :header-rows: 1 :widths: 50 20 * - Use Case - Container * - Immutable snapshot of a sequence (e.g. function arguments) - Array * - Building up or modifying a sequence in-place - List * - Fixed heterogeneous collection (e.g. a multi-typed return value) - Tuple * - Immutable key-value lookup (e.g. configuration) - Map * - Mutable key-value store (e.g. accumulating results) - Dict Thread Safety ------------- **Immutable containers** (``Array``, ``Tuple``, ``Map``) can be safely shared across threads for **read-only** access. Copy-on-write mutations are thread-safe because they create a new backing object when the storage is shared. **Mutable containers** (``List``, ``Dict``) are **NOT** thread-safe. If multiple threads need to read or write the same ``List`` or ``Dict``, external synchronization (e.g. a mutex) is required. Further Reading --------------- - :doc:`../guides/cpp_lang_guide`: C++ examples for all container types - :doc:`../guides/python_lang_guide`: Python examples and conversion rules - :doc:`any`: How containers are stored in the type-erased :cpp:class:`~tvm::ffi::Any` value - :doc:`object_and_class`: The object system underlying all container types tvm-ffi-0.1.12/docs/concepts/exception_handling.rst000066400000000000000000000161061521067262500223170ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. .. _sec-exception-handling: Exception Handling ================== TVM-FFI gracefully handles exceptions across language boundaries without requiring manual error code management. This document covers throwing, catching, and propagating exceptions in TVM-FFI functions. .. important:: Stack traces from all languages are properly preserved and concatenated in the TVM-FFI Stable C ABI. Cross-language exceptions are **first-class citizens** in TVM-FFI. TVM-FFI provides a stable C ABI for propagating exceptions across language boundaries, and wraps this ABI to provide language-native exception handling without exposing the underlying C machinery. **Error in C**. :cpp:class:`tvm::ffi::Error` is a TVM-FFI :doc:`object ` with a :cpp:class:`TVMFFIErrorCell` payload containing: - ``kind``: Error type name (e.g., ``"ValueError"``, ``"RuntimeError"``) - ``message``: Human-readable error message - ``backtrace``: Stack trace from the point of error **Propagating Errors in C**. For call chains, simply propagate return codes—TLS carries the error details: .. code-block:: cpp int my_function(...) { int rc = some_ffi_call(...); if (rc != 0) return rc; // Propagate error // Continue on success return 0; } The following sections describe how to throw and catch exceptions across ABI and language boundaries. Throwing Exceptions ------------------- Python ~~~~~~ Raise native :py:class:`Exception ` instances or derived classes. TVM-FFI catches these at the ABI boundary and converts them to :cpp:class:`tvm::ffi::Error` objects. When C++ code calls into Python and a Python exception occurs, it propagates back to C++ as a :cpp:class:`tvm::ffi::Error`, which C++ code can handle appropriately. .. code-block:: python def my_function(x: int) -> int: if x < 0: raise ValueError(f"x must be non-negative, got {x}") return x + 1 TVM-FFI automatically maps common Python exception types to their corresponding error kinds: ``RuntimeError``, ``ValueError``, ``TypeError``, ``AttributeError``, ``KeyError``, ``IndexError``, ``AssertionError``, and ``MemoryError``. Custom exception types can be registered using :py:func:`tvm_ffi.register_error`. C++ ~~~ Use :cpp:class:`tvm::ffi::Error` or the :c:macro:`TVM_FFI_THROW` macro: .. code-block:: cpp #include void ThrowError(int x) { if (x < 0) { TVM_FFI_THROW(ValueError) << "x must be non-negative, got " << x; } } The :c:macro:`TVM_FFI_THROW` macro captures the current file name, line number, stack trace, and error message, then constructs a :cpp:class:`tvm::ffi::Error` object. At the ABI boundary, this error is stored in TLS and the function returns ``-1`` per the :cpp:type:`TVMFFISafeCallType` calling convention (see :ref:`sec:function-calling-convention`). Additional check macros are available for common validation patterns: .. code-block:: cpp TVM_FFI_CHECK(condition, ErrorKind) << "message"; // Custom error kind TVM_FFI_ICHECK(condition) << "message"; // InternalError TVM_FFI_ICHECK_EQ(x, y) << "message"; // Check equality TVM_FFI_ICHECK_LT(x, y) << "message"; // Check less than // Also: TVM_FFI_ICHECK_GT, TVM_FFI_ICHECK_LE, TVM_FFI_ICHECK_GE, TVM_FFI_ICHECK_NE .. hint:: A detailed implementation of such graceful handling behavior can be found in :c:macro:`TVM_FFI_SAFE_CALL_BEGIN` / :c:macro:`TVM_FFI_SAFE_CALL_END` macros. ANSI C ~~~~~~ For LLVM code generation and other C-based environments, use :cpp:func:`TVMFFIErrorSetRaisedFromCStr` to set the TLS error and return ``-1``: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.RaiseException.begin] :end-before: [Error.RaiseException.end] For constructing error messages from multiple parts (useful in code generators), use :cpp:func:`TVMFFIErrorSetRaisedFromCStrParts`: .. code-block:: cpp const char* parts[] = {"Expected ", "2", " arguments, got ", "1"}; TVMFFIErrorSetRaisedFromCStrParts("ValueError", parts, 4); return -1; .. _sec-exception: Catching Exceptions in C ------------------------ In C++, Python, and many other languages, TVM-FFI exceptions are **first-class citizens**, meaning they can be caught and handled like native exceptions: .. code-block:: cpp try { // Calls a TVM-FFI function that throws an exception } catch (const tvm::ffi::Error& e) { // Handle the exception std::cout << e.kind() << ": " << e.message() << "\n" << e.backtrace() << "\n"; } .. important:: This section covers the **pure C** low-level details of exception handling. For C++ and other languages, the example above is sufficient. Checking Return Codes ~~~~~~~~~~~~~~~~~~~~~ When a TVM-FFI function returns a non-zero code, an error occurred. The error object is stored in thread-local storage (TLS) and can be retrieved with :cpp:func:`TVMFFIErrorMoveFromRaised`: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.HandleReturnCode.begin] :end-before: [Error.HandleReturnCode.end] .. important:: The caller must release the error object via :cpp:func:`TVMFFIObjectDecRef` to avoid memory leaks. Accessing Error Details ~~~~~~~~~~~~~~~~~~~~~~~ The error payload is a :cpp:type:`TVMFFIErrorCell` structure containing the error kind, message, and backtrace. Access it by skipping the :cpp:type:`TVMFFIObject` header: .. literalinclude:: ../../examples/abi_overview/example_code.c :language: c :start-after: [Error.Print.begin] :end-before: [Error.Print.end] Return Code Reference ~~~~~~~~~~~~~~~~~~~~~ - **Error code 0:** Success - **Error code -1:** Error occurred, retrieve via :cpp:func:`TVMFFIErrorMoveFromRaised` Further Reading --------------- - :doc:`func_module`: Functions and modules that use this exception handling mechanism - :doc:`object_and_class`: The object system that backs :cpp:class:`~tvm::ffi::Error` - :doc:`any`: How errors are stored and transported in :cpp:class:`~tvm::ffi::Any` containers - :doc:`abi_overview`: Low-level C ABI details for exceptions - `tvm/ffi/error.h `_: C++ error handling API - `tvm/ffi/c_api.h `_: C ABI error functions tvm-ffi-0.1.12/docs/concepts/func_module.rst000066400000000000000000000402071521067262500207540ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Function and Module =================== TVM-FFI provides a unified and ABI-stable calling convention that enables cross-language function calls between C++, Python, Rust, and other languages. Functions are first-class :doc:`TVM-FFI objects `. This tutorial covers defining, registering, and calling TVM-FFI functions, exception handling, and working with modules. Glossary -------- TVM-FFI ABI, or "Packed Function". :cpp:type:`TVMFFISafeCallType` A stable C calling convention where every function is represented by a single signature, which enables type-erased, cross-language function calls. This calling convention is used across all TVM-FFI function calls at the ABI boundary. See :ref:`Stable C ABI ` for a quick introduction. TVM-FFI Function. :py:class:`tvm_ffi.Function`, :cpp:class:`tvm::ffi::FunctionObj`, :cpp:class:`tvm::ffi::Function` A reference-counted :doc:`function object ` and its managed reference, which wraps any callable, including language-agnostic functions and lambdas (C++, Python, Rust, etc.), member functions, external C symbols, and other callable objects, all sharing the same calling convention. TVM-FFI Module. :py:class:`tvm_ffi.Module`, :cpp:class:`tvm::ffi::ModuleObj`, :cpp:class:`tvm::ffi::Module` A namespace for a collection of functions, loaded from a shared library via ``dlopen`` (Linux, macOS) or ``LoadLibraryW`` (Windows), or statically linked to the current executable. Global Functions and Registry. :py:func:`tvm_ffi.get_global_func` and :py:func:`tvm_ffi.register_global_func` A registry is a table that maps string names to :cpp:class:`~tvm::ffi::Function` objects and their metadata (name, docs, signatures, etc.) for cross-language access. Functions in the registry are called **global functions**. Common Usage ------------ TVM-FFI C Symbols ~~~~~~~~~~~~~~~~~ **Shared library**. Use :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` to export a function as a C symbol that follows the TVM-FFI ABI: .. code-block:: cpp static int AddTwo(int x) { return x + 2; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(/*ExportName=*/add_two, /*Function=*/AddTwo) This creates a C symbol ``__tvm_ffi_`` in the shared library, which can then be loaded and called via :py:func:`tvm_ffi.load_module`: .. code-block:: python import tvm_ffi mod = tvm_ffi.load_module("path/to/library.so") result = mod.add_two(40) # -> 42 **System library**. For symbols bundled in the same executable, use :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol` to register each symbol during static initialization within a :c:macro:`TVM_FFI_STATIC_INIT_BLOCK`. See :py:func:`tvm_ffi.system_lib` for a complete workflow. Global Functions ~~~~~~~~~~~~~~~~ **Register a global function**. In C++, use :cpp:class:`tvm::ffi::reflection::GlobalDef` to register a function: .. code-block:: cpp #include static int AddOne(int x) { return x + 1; } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("my_ext.add_one", AddOne, "Add one to the input"); } The :c:macro:`TVM_FFI_STATIC_INIT_BLOCK` macro ensures that registration occurs during library initialization. The registered function is then accessible from Python by the name ``my_ext.add_one``. In Python, use the decorator :py:func:`tvm_ffi.register_global_func` to register a global function: .. code-block:: python import tvm_ffi @tvm_ffi.register_global_func("my_ext.add_one") def add_one(x: int) -> int: return x + 1 **Retrieve a global function**. After registration, functions are accessible by name. In Python, use :py:func:`tvm_ffi.get_global_func` to retrieve a global function: .. code-block:: python import tvm_ffi # Get a function from the global registry add_one = tvm_ffi.get_global_func("my_ext.add_one") result = add_one(41) # -> 42 In C++, use :cpp:func:`tvm::ffi::Function::GetGlobal` or :cpp:func:`tvm::ffi::Function::GetGlobalRequired` to retrieve a global function: .. code-block:: cpp ffi::Function func = ffi::Function::GetGlobalRequired("my_ext.add_one"); int result = func(41); // -> 42 Create Functions ~~~~~~~~~~~~~~~~ **From C++**. An :cpp:class:`tvm::ffi::Function` can be created via :cpp:func:`tvm::ffi::Function::FromTyped` or :cpp:class:`tvm::ffi::TypedFunction`'s constructor. .. code-block:: cpp // Create type-erased function: add_type_erased ffi::Function add_type_erased = ffi::Function::FromTyped([](int x, int y) { return x + y; }); // Create a typed function: add_typed ffi::TypedFunction add_typed = [](int x, int y) { return x + y; }; // Convert a typed function to a type-erased function ffi::Function generic = add_typed; **From Python**. Any Python :py:class:`Callable ` is automatically converted to a :py:class:`tvm_ffi.Function` at the ABI boundary. The example below demonstrates that in ``my_ext.bind``: - The input ``func`` is automatically converted to a :py:class:`tvm_ffi.Function`. - The returned lambda is also automatically converted to a :py:class:`tvm_ffi.Function`. .. code-block:: python import tvm_ffi @tvm_ffi.register_global_func("my_ext.bind") def bind(func, x): assert isinstance(func, tvm_ffi.Function) return lambda *args: func(x, *args) # converted to `tvm_ffi.Function` def add_x_y(x, y): return x + y func_bind = tvm_ffi.get_global_func("my_ext.bind") add_y = func_bind(add_x_y, 1) # bind x = 1 assert isinstance(add_y, tvm_ffi.Function) print(add_y(2)) # -> 3 :py:func:`tvm_ffi.convert` explicitly converts a Python callable to :py:class:`tvm_ffi.Function`: .. code-block:: python import tvm_ffi def add(x, y): return x + y func_add = tvm_ffi.convert(add) print(func_add(1, 2)) .. _sec:function: Function -------- .. _sec:function-calling-convention: Calling Convention ~~~~~~~~~~~~~~~~~~ All TVM-FFI functions ultimately conform to the :cpp:type:`TVMFFISafeCallType` signature, which provides a stable C ABI for cross-language calls. The C calling convention is defined as: .. code-block:: cpp int tvm_ffi_c_abi( void* handle, // Resource handle const TVMFFIAny* args, // Input arguments (non-owning) int32_t num_args, // Number of input arguments TVMFFIAny* result // Output argument (owning, zero-initialized) ); **Input arguments**. The input arguments are passed as an array of :cpp:class:`tvm::ffi::AnyView` values (see :ref:`any-ownership` for ownership semantics), specified by ``args`` and ``num_args``. **Output argument**. The output argument ``result`` is an owning :cpp:type:`tvm::ffi::Any` that the caller must zero-initialize before the call. .. important:: The caller must zero-initialize the output argument ``result`` before the call. **Return value**. The ABI returns an **error code** that indicates: - **Error code 0**: Success - **Error code -1**: Error occurred, retrievable with :cpp:func:`TVMFFIErrorMoveFromRaised` .. hint:: See :doc:`Any ` for more details on the semantics of :cpp:type:`tvm::ffi::AnyView` and :cpp:type:`tvm::ffi::Any`. This design is called a **packed function**, because it "packs" all arguments into a single array of type-erased :cpp:type:`tvm::ffi::AnyView`, and further unifies calling convention across all languages without resorting to JIT compilation. More specifically, this mechanism enables the following scenarios: - **Dynamic languages**. Well-optimized bindings are provided for, e.g. Python, to translate arguments into packed function format, and translate return value back to the host language. - **Static languages**. Metaprogramming techniques, such as C++ templates, are usually available to directly instantiate packed format on stack, saving the need for dynamic examination. - **Cross-language callbacks**. Language-agnostic :cpp:class:`tvm::ffi::Function` makes it easy to call between languages without depending on language-specific features such as GIL. **Performance Implications**. This approach is in practice highly efficient in machine learning workloads. - In Python/C++ calls, we can get to microsecond level overhead, which is generally similar to overhead for eager mode; - When both sides of calls are static languages, the overhead will go down to tens of nanoseconds. .. note:: Although we found it less necessary in practice, further link time optimization (LTO) is still theoretically possible in scenarios where both sides are static languages with a known symbol and linked into a single binary. In this case, the callee can be inlined into caller side and the stack argument memory can be passed into register passing. .. _sec:function-layout: Layout and ABI ~~~~~~~~~~~~~~ :cpp:class:`tvm::ffi::FunctionObj` stores two call pointers in :cpp:class:`TVMFFIFunctionCell`: - ``safe_call``: Used for cross-ABI function calls; intercepts exceptions and stores them in TLS. - ``cpp_call``: Used within the same DSO; exceptions are thrown directly for better performance. See :ref:`abi-function` for the C struct definition. .. important:: :cpp:func:`TVMFFIFunctionCall` is the idiomatic way to call a :cpp:class:`tvm::ffi::FunctionObj` in C, while ``safe_call`` or ``cpp_call`` remain low-level ABIs for fast access. **Conversion with Any**. Since :py:class:`tvm_ffi.Function` is a TVM-FFI object, it follows the same conversion rules as any other TVM-FFI object. See :ref:`Object Conversion with Any ` for details. Exception Handling ~~~~~~~~~~~~~~~~~~ See :doc:`exception_handling` for details on throwing, catching, and propagating exceptions across language boundaries. Compiler developers commonly need to look up global functions in generated code. Use :cpp:func:`TVMFFIFunctionGetGlobal` to retrieve a function by name, then call it with :cpp:func:`TVMFFIFunctionCall`. See :ref:`abi-function` for C code examples. .. _sec:module: Modules ------- A :py:class:`tvm_ffi.Module` is a namespace for a collection of functions that can be loaded from a shared library or bundled with the current executable. Modules provide namespace isolation and dynamic loading capabilities for TVM-FFI functions. Shared Library ~~~~~~~~~~~~~~ Shared library modules are loaded dynamically at runtime via ``dlopen`` (Linux, macOS) or ``LoadLibraryW`` (Windows). This is the most common way to distribute and load compiled functions. **Export functions from C++**. Use :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` to export a function as a C symbol that follows the TVM-FFI ABI: .. code-block:: cpp #include static int AddTwo(int x) { return x + 2; } // Exports as symbol `__tvm_ffi_add_two` TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_two, AddTwo); **Load and call from Python**. Use :py:func:`tvm_ffi.load_module` to load the shared library: .. code-block:: python import tvm_ffi # Load the shared library mod = tvm_ffi.load_module("path/to/library.so") # Access functions by name result = mod.add_two(40) # -> 42 # Alternative: explicit function retrieval func = mod.get_function("add_two") result = func(40) # -> 42 **Build and load from source**. For rapid prototyping, :py:func:`tvm_ffi.cpp.load` compiles C++/CUDA source files and loads them as a module in one step: .. code-block:: python import tvm_ffi.cpp # Compile and load in one step mod = tvm_ffi.cpp.load( name="my_ops", cpp_files="my_ops.cpp", ) result = mod.add_two(40) Essentially, :py:func:`tvm_ffi.cpp.load` is a convenience function that JIT-compiles the source files and loads the resulting library as a :py:class:`tvm_ffi.Module`. System Library ~~~~~~~~~~~~~~ System library modules contain symbols that are statically linked to the current executable. This technique is useful when you want to simulate dynamic module loading behavior but cannot or prefer not to use ``dlopen`` or ``LoadLibraryW`` (e.g., on iOS). Functions are statically linked to the executable as a system library module. Symbols can be registered via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol` and looked up via :py:func:`tvm_ffi.system_lib`. **Register symbols in C/C++**. Use :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol` to register a symbol during static initialization: .. code-block:: cpp #include #include // A function following the TVM-FFI ABI static int add_one_impl(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { TVM_FFI_SAFE_CALL_BEGIN(); int64_t x = reinterpret_cast(args)[0].cast(); reinterpret_cast(result)[0] = x + 1; TVM_FFI_SAFE_CALL_END(); } // Register during static initialization // The symbol name follows the convention `__tvm_ffi_.` TVM_FFI_STATIC_INIT_BLOCK() { TVMFFIEnvModRegisterSystemLibSymbol( "__tvm_ffi_my_prefix.add_one", reinterpret_cast(add_one_impl) ); } **Access from Python**. Use :py:func:`tvm_ffi.system_lib` to get the system library module: .. code-block:: python import tvm_ffi # Get system library with symbol prefix "my_prefix." # This looks up symbols prefixed with `__tvm_ffi_my_prefix.` mod = tvm_ffi.system_lib("my_prefix.") # Call the registered function func = mod.add_one # looks up `__tvm_ffi_my_prefix.add_one` result = func(10) # -> 11 .. note:: The system library is intended for statically linked symbols that exist for the entire program lifetime. For dynamic loading with the ability to unload, use shared library modules instead. .. _sec:custom-modules: Custom Modules ~~~~~~~~~~~~~~ While the standard shared library and system library modules cover most use cases, some scenarios require a **custom module** that wraps a platform-specific driver API — for example, using ``cuModuleLoad`` to load generated PTX code and expose each kernel as a :cpp:class:`tvm::ffi::Function`. To create a custom module, subclass :cpp:class:`tvm::ffi::ModuleObj` and implement the following: - :cpp:func:`~tvm::ffi::ModuleObj::kind` — return a unique string identifying the module type (e.g., ``"cuda"``). - :cpp:func:`~tvm::ffi::ModuleObj::GetPropertyMask` — return a bitmask indicating the module's capabilities: - ``ffi::Module::kRunnable`` if the module can execute functions. - ``ffi::Module::kBinarySerializable`` if the module supports serialization to and from bytes. - :cpp:func:`~tvm::ffi::ModuleObj::GetFunction` — look up a function by name within the module. - If the module is serializable, override :cpp:func:`~tvm::ffi::ModuleObj::SaveToBytes` and register a global function ``ffi.Module.load_from_bytes.`` so the module can be reconstructed from its serialized form. .. seealso:: :ref:`Embedded Binary Data ` in the export guide describes the ``__tvm_ffi__library_bin`` binary layout used to serialize composite modules that contain custom sub-modules. Further Reading --------------- - :doc:`exception_handling`: Throwing, catching, and propagating exceptions across language boundaries - :doc:`any`: How functions are stored in :cpp:class:`~tvm::ffi::Any` containers - :doc:`object_and_class`: The object system that backs :cpp:class:`~tvm::ffi::FunctionObj` - :doc:`abi_overview`: Low-level C ABI details for functions and exceptions - :doc:`../packaging/python_packaging`: Packaging functions for Python wheels tvm-ffi-0.1.12/docs/concepts/object_and_class.rst000066400000000000000000000400141521067262500217250ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Object and Class ================ TVM-FFI provides a unified object system that enables cross-language interoperability between C++, Python, and Rust. The object system is built around :cpp:class:`tvm::ffi::Object` and :cpp:class:`tvm::ffi::ObjectRef`, which together form the foundation for: - **Type-safe runtime type identification** without relying on C++ RTTI - **Intrusive reference counting** for smart memory management - **Reflection-based class exposure** across programming languages - **Serialization and deserialization** via reflection metadata This tutorial covers defining, using, and extending TVM-FFI objects across languages. Glossary -------- :cpp:class:`tvm::ffi::Object` A heap-allocated, reference-counted container. All TVM-FFI objects inherit from this base class and share a common 24-byte header that stores reference counts, type index, and a deleter callback. :cpp:class:`tvm::ffi::ObjectRef` An intrusive pointer that manages an :cpp:class:`~tvm::ffi::Object`'s lifetime through reference counting. Its subclasses provide type-safe access to specific object types. In its low-level implementation, it is equivalent to a normal C++ pointer to a heap-allocated :cpp:class:`~tvm::ffi::Object`. Type index and type key Type index is an integer that uniquely identifies each object type. Built-in types have statically assigned indices defined in :cpp:enum:`TVMFFITypeIndex` (see :ref:`any-atomic-types` and :ref:`any-heap-allocated-objects` for the complete list), while user-defined types receive indices at startup when first accessed. Type key is a unique string identifier (e.g., ``"my_ext.MyClass"``) that names an object type. It is used for registration, serialization, and cross-language mapping. Common Usage ------------ Define a Class in C++ ~~~~~~~~~~~~~~~~~~~~~ To define a custom object class in normal C++, inherit it from :cpp:class:`tvm::ffi::Object` or its subclasses, and then add one of the following macros that declares its metadata: :c:macro:`TVM_FFI_DECLARE_OBJECT_INFO(TypeKey, TypeName, ParentType) ` Declare an object type that can be subclassed. Type index is assigned dynamically. :c:macro:`TVM_FFI_DECLARE_OBJECT_INFO_FINAL(TypeKey, TypeName, ParentType) ` Declare a final object type (no subclasses). Enables faster type checking. **Example**. The code below shows a minimal example of defining a TVM-FFI object class. It declares a class ``MyObjectObj`` that inherits from :cpp:class:`~tvm::ffi::Object`. .. code-block:: cpp #include namespace ffi = tvm::ffi; class MyObjectObj : public ffi::Object { public: // Normal C++ code: Declare fields, methods, constructor, destructor, etc. int64_t value; ffi::String name; MyObjectObj(int64_t value, ffi::String name) : value(value), name(std::move(name)) {} int64_t GetValue() const { return value; } void AddToValue(int64_t other) { value += other; } // Declare object type info TVM_FFI_DECLARE_OBJECT_INFO( /*type_key=*/"my_ext.MyObject", /*type_name=*/MyObjectObj, /*parent_type=*/ffi::Object); }; **Managed reference**. Optionally, a managed reference class can be defined by inheriting from :cpp:class:`~tvm::ffi::ObjectRef` and using one of the following macros to define the methods. Define its constructor by wrapping the :cpp:func:`tvm::ffi::make_object` function. :c:macro:`TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypeName, ParentType, ObjectName) ` Define a nullable reference class. :c:macro:`TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypeName, ParentType, ObjectName) ` Define a non-nullable reference class. For example, a non-nullable reference class ``MyObject`` can be defined as follows: .. code-block:: cpp class MyObject : public ffi::ObjectRef { public: MyObject(int64_t value, ffi::String name) : ObjectRef(ffi::make_object(value, std::move(name))) {} TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE( /*type_name=*/MyObject, /*parent_type=*/ffi::ObjectRef, /*object_name=*/MyObjectObj); }; // Create a managed object MyObject obj = MyObject(42, "hello"); // Access fields via operator-> std::cout << obj->value << std::endl; // -> 42 Expose a Class in Python ~~~~~~~~~~~~~~~~~~~~~~~~ **Reflection**. The object's metadata is used for reflection. Use :cpp:class:`tvm::ffi::reflection::ObjectDef` to register an object's constructor, fields, and methods. .. code-block:: cpp TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() // Register constructor with signature .def(refl::init()) // Register read-write fields .def_rw("value", &MyObjectObj::value, "The integer value") .def_rw("name", &MyObjectObj::name, "The name string") // Register methods .def("get_value", &MyObjectObj::GetValue, "Returns the value"); } **Python binding**. After registration, the object is automatically available in Python. Use :py:func:`tvm_ffi.register_object` to bind a Python class to a registered C++ type: .. code-block:: python import tvm_ffi from typing import TYPE_CHECKING @tvm_ffi.register_object("my_ext.MyObject") class MyObject(tvm_ffi.Object): # tvm-ffi-stubgen(begin): object/my_ext.MyObject value: int name: str if TYPE_CHECKING: def __init__(self, value: int, name: str) -> None: ... def get_value(self) -> int: ... # tvm-ffi-stubgen(end) # Create and use objects obj = MyObject(42, "hello") print(obj.value) # -> 42 print(obj.get_value()) # -> 42 obj.value = 100 # Mutable field access The decorator looks up the type key ``"my_ext.MyObject"`` in the C++ type registry and binds the Python class to it. Fields and methods registered via :cpp:class:`~tvm::ffi::reflection::ObjectDef` are automatically available on the Python class. The tool ``tvm-ffi-stubgen`` automatically generates the Python type stubs (the code between the markers) from reflection metadata. See :ref:`Stub Generation Tool ` for details. .. _type-checking-and-casting: Type Checking and Casting ~~~~~~~~~~~~~~~~~~~~~~~~~ **Type checking**. Use :cpp:func:`Object::IsInstance\() ` for runtime type checking: .. code-block:: cpp bool CheckType(const ffi::ObjectRef& obj) { if (obj->IsInstance()) { // obj is a MyObjectObj or subclass return true; } return false; } **Type casting**. Use :cpp:func:`ObjectRef::as\() ` for safe downcasting. For strict conversion that throws on mismatch, use :cpp:func:`~tvm::ffi::Any::cast` (see :doc:`exception_handling` for error handling details): .. code-block:: cpp ffi::ObjectRef obj = ...; // as() returns a pointer (nullptr if type doesn't match) if (const MyObjectObj* ptr = obj.as()) { std::cout << ptr->value << std::endl; } // as() returns std::optional if (auto opt = obj.as()) { std::cout << opt->get()->value << std::endl; } **Type info**. Type index is available via :cpp:func:`ObjectRef::type_index() ` and type key is available via :cpp:func:`ObjectRef::GetTypeKey() `. These methods can be used to safely identify object types without relying on C++ RTTI. .. note:: C++ RTTI (e.g. ``typeid``, ``dynamic_cast``) is strictly not useful in TVM-FFI-based approaches. Miscellaneous APIs ~~~~~~~~~~~~~~~~~~ **C++ Serialization**. Use :cpp:func:`tvm::ffi::ToJSONGraph` to serialize an object to a JSON value, and :cpp:func:`tvm::ffi::FromJSONGraph` to deserialize a JSON value to an object. .. code-block:: cpp #include // Serialize to JSON ffi::Any obj = ...; ffi::json::Value json = ffi::ToJSONGraph(obj); // Deserialize from JSON ffi::Any restored = ffi::FromJSONGraph(json); **Python Serialization**. Pickle is overloaded in Python to support TVM-FFI object serialization. Or explicitly use the :py:func:`tvm_ffi.serialization.to_json_graph_str` and :py:func:`tvm_ffi.serialization.from_json_graph_str` to serialize and deserialize an object to a JSON string. .. code-block:: python import pickle obj = MyObject(42, "test") data = pickle.dumps(obj) restored = pickle.loads(data) **Convert between raw and managed references**. Use :cpp:func:`tvm::ffi::GetRef` to convert a raw object pointer to a managed reference, and :cpp:func:`tvm::ffi::ObjectRef::get` to convert a managed reference to a raw object pointer. ABI and Layout -------------- **Stable C Layout**. All subclasses of :cpp:class:`tvm::ffi::Object` share a common 24-byte header (:cpp:class:`TVMFFIObject`) containing reference counts, type index, and a deleter callback. See :ref:`abi-object` for the C struct definition. :cpp:class:`tvm::ffi::ObjectRef` and :cpp:class:`tvm::ffi::ObjectPtr` are smart pointers equivalent to a single ``void*`` pointer. .. _object-reference-counting: Reference Counting ~~~~~~~~~~~~~~~~~~ .. seealso:: :ref:`abi-object-ownership` for C code examples. **Intrusive reference counting**. The reference count is stored directly in the object header, not in a separate control block. This design reduces memory overhead and improves cache locality. The :cpp:member:`TVMFFIObject::combined_ref_count` field packs both strong (lower 32 bits) and weak (upper 32 bits) reference counts in a single 64-bit integer. C APIs are provided to manipulate the reference count: - :cpp:func:`TVMFFIObjectIncRef` to increase the strong reference count - :cpp:func:`TVMFFIObjectDecRef` to decrease the strong reference count **Deleter**. When an object is managed by :cpp:class:`~tvm::ffi::ObjectRef`, the ``deleter`` callback is invoked: - When strong reference count reaches zero: the object's destructor is called. - When weak reference count reaches zero: the memory is freed. The flags in :cpp:enum:`TVMFFIObjectDeleterFlagBitMask` indicate which action to perform. .. _object-conversion-with-any: **Conversion with Any**. At the stable C ABI boundary, TVM-FFI passes values using :cpp:class:`Any ` (owning) or :cpp:class:`AnyView ` (non-owning). Object handles are stored in the :cpp:member:`TVMFFIAny::v_obj` field with a type index >= :cpp:enumerator:`kTVMFFIStaticObjectBegin `. See :ref:`abi-object-ownership` for C code examples demonstrating: - Extracting an object handle from :cpp:class:`TVMFFIAny` - Storing an object handle into :cpp:class:`~tvm::ffi::Any` or :cpp:class:`~tvm::ffi::AnyView` - Managing ownership via reference counting Object Type Registry -------------------- TVM-FFI maintains a global type registry that keeps track of all registered object types, their inheritance relationships, and their reflection metadata. Inheritance and Type Casting ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. warning:: Only single inheritance is supported in TVM-FFI Object system. TVM-FFI implements its own runtime type system that enables type-safe operations without relying on C++ RTTI. Every object carries a runtime type index in its header. **Example**. Code below shows a minimal example of defining a base class and a derived class. .. code-block:: cpp class MyBaseObj : public ffi::Object { public: TVM_FFI_DECLARE_OBJECT_INFO("my_ext.MyBase", MyBaseObj, ffi::Object); }; class MyDerivedObj : public MyBaseObj { public: // Final class: no subclasses allowed TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my_ext.MyDerived", MyDerivedObj, MyBaseObj); }; Registration happens automatically on first access. The :c:macro:`TVM_FFI_DECLARE_OBJECT_INFO` and :c:macro:`TVM_FFI_DECLARE_OBJECT_INFO_FINAL` macros use :cpp:func:`TVMFFITypeGetOrAllocIndex` internally to allocate a type index. See :ref:`type-checking-and-casting` for how to use the type system. Reflect Fields and Methods ~~~~~~~~~~~~~~~~~~~~~~~~~~ The reflection system enables cross-language exposure of C++ classes, their fields, and methods. Use :cpp:class:`ObjectDef\ ` to register reflection metadata for object type ``T``: .. list-table:: :header-rows: 1 :widths: 30 70 * - Method - Description * - ``.def(init())`` - Register a constructor with the given argument types * - ``.def_ro("name", &T::field)`` - Register a read-only field * - ``.def_rw("name", &T::field)`` - Register a read-write field * - ``.def("name", &T::method)`` - Register a member method * - ``.def_static("name", &func)`` - Register a static method **Example**. Code below shows a minimal example of registering a class with reflection metadata. .. code-block:: cpp class IntPairObj : public ffi::Object { public: int64_t a; int64_t b; IntPairObj(int64_t a, int64_t b) : a(a), b(b) {} int64_t Sum() const { return a + b; } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my_ext.IntPair", IntPairObj, ffi::Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def(refl::init()) .def_rw("a", &IntPairObj::a, "the first field") .def_rw("b", &IntPairObj::b, "the second field") .def("sum", &IntPairObj::Sum, "compute a + b"); } **Metadata and Documentation**. Add documentation strings and custom metadata to fields and methods: .. code-block:: cpp // The following example uses MyObjectObj defined earlier to show // how to add documentation and metadata. refl::ObjectDef() .def_rw("value", &MyObjectObj::value, "The numeric value", // docstring refl::default_(0), // default value refl::Metadata{{"min", 0}, {"max", 100}}) // custom metadata .def("add_to_value", &MyObjectObj::AddToValue, "Add a value to the object's value field"); Python Interoperability ~~~~~~~~~~~~~~~~~~~~~~~ **Cross-language lifetime**. Each Python :py:class:`tvm_ffi.Object` instance holds a C handle (``void*``) that references the underlying C++ object. The Python wrapper increments the reference count when constructed and decrements when garbage collected. .. code-block:: python obj = MyObject(42, "test") # C++ object created, C++ refcount = 1 obj2 = obj # Python alias created, C++ refcount unchanged del obj # Python alias removed, C++ refcount unchanged del obj2 # Last Python reference gone, C++ refcount -> 0, object destroyed Further Reading --------------- - :doc:`any`: How objects are stored in :cpp:class:`~tvm::ffi::Any` containers - :doc:`func_module`: Function objects and the global registry - :doc:`tensor`: Tensor objects and DLPack interoperability - :doc:`exception_handling`: Error objects and cross-language exception propagation - :doc:`abi_overview`: Low-level C ABI details for the object system - :doc:`../packaging/python_packaging`: Packaging C++ objects for Python wheels tvm-ffi-0.1.12/docs/concepts/structural_eq_hash.rst000066400000000000000000000742711521067262500223640ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Structural Equality and Hashing =============================== TVM FFI provides ``structural_equal`` and ``structural_hash`` for the object graph. These compare objects by **content** — recursively walking fields — rather than by pointer identity. The behavior is controlled by two layers of annotation on :func:`~tvm_ffi.dataclasses.py_class`: 1. **Type-level** ``structural_eq=`` — what *role* does this type play in the IR graph? 2. **Field-level** ``structural_eq=`` on :func:`~tvm_ffi.dataclasses.field` — should this field be skipped, or does it introduce new variable bindings? This document explains what each annotation means, when to use it, and how they compose. .. note:: Structural equality and hashing **never call Python-level** ``__eq__`` or ``__hash__``. ``structural_equal`` / ``structural_hash`` dispatch entirely through a C++ walker driven by the kind metadata registered via ``structural_eq=``; the Python ``a == b`` / ``hash(a)`` dunders are independent (they default to pointer identity and handle address, inherited from ``Object``). To customize how a specific type participates in *structural* comparison, register the :ref:`sequal-shash` hooks described below — do **not** override ``__eq__`` or ``__hash__``. Type-Level Annotation --------------------- The ``structural_eq`` parameter on ``@py_class`` declares how instances of the type participate in structural equality and hashing: .. code-block:: python @py_class(structural_eq="tree") class Expr(Object): ... Quick reference ~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 18 37 45 * - ``structural_eq=`` - Meaning - Use when... * - ``"tree"`` - A regular IR node - Default for most IR nodes * - ``"const-tree"`` - An immutable value node (with pointer shortcut) - The type has no transitive ``"var"`` children * - ``"dag"`` - A node in a dataflow graph - Pointer sharing is semantically meaningful * - ``"var"`` - A bound variable - The type represents a variable binding * - ``"singleton"`` - A singleton - Exactly one instance per logical identity (e.g. registry entries) * - ``None`` - Not comparable - The type should never be compared structurally ``"tree"`` — The Default ------------------------- .. code-block:: python @py_class(structural_eq="tree") class Add(Object): lhs: Expr rhs: Expr **Meaning**: "This node is defined by its fields. Two nodes are equal if and only if all their fields are recursively equal." This is the right choice for the vast majority of IR nodes: expressions, statements, types, attributes, buffers, etc. **Example.** .. code-block:: text 1 + 2 vs 1 + 2 → Equal 1 + 2 vs 1 + 3 → Not equal (rhs differs) Sharing is invisible ~~~~~~~~~~~~~~~~~~~~ ``"tree"`` treats every reference independently. If the same object is referenced multiple times, each reference is compared by content separately. Sharing is **not** part of the structural identity: .. code-block:: text let s = x + 1 (s, s) ← same object referenced twice (x + 1, x + 1) ← two independent copies with same content These are EQUAL under "tree" — sharing is not detected. The following diagram illustrates this. Under ``"tree"``, the **DAG** on the left and the **tree** on the right are considered structurally equal because every node has the same content: .. mermaid:: graph TD subgraph "DAG — shared node" T1["(_, _)"] S1["s = x + 1"] T1 -->|".0"| S1 T1 -->|".1"| S1 end subgraph "Tree — independent copies" T2["(_, _)"] A1["x + 1"] A2["x + 1"] T2 -->|".0"| A1 T2 -->|".1"| A2 end style S1 fill:#d4edda style A1 fill:#d4edda style A2 fill:#d4edda If sharing needs to matter, use ``"dag"`` instead. ``"const-tree"`` — Tree with a Fast Path ----------------------------------------- .. code-block:: python @py_class(structural_eq="const-tree") class DeviceMesh(Object): shape: list[int] device_ids: list[int] **Meaning**: "Same as ``"tree"``, but if two references point to the same object, they are guaranteed equal — skip the field comparison." This is purely a **performance optimization**. The only behavioral difference from ``"tree"`` is that pointer identity short-circuits to ``True``. When is this safe (and worth it)? ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Three conditions decide whether ``"const-tree"`` is the right choice: 1. **Immutable** — content doesn't change after construction, so same-pointer always implies same-content. 2. **No transitive** ``"var"`` **children** — skipping field traversal won't cause variable mappings to be missed (see :ref:`var-kind` for why this matters). 3. **Sharing is common** — instances are interned or canonicalized, so the same pointer actually appears on both sides of real comparisons. Without interning, the shortcut never fires and ``"const-tree"`` behaves like ``"tree"`` with a dead branch. Conditions 1 and 2 are correctness requirements: violating them is a bug, not a performance regression. Condition 3 is the payoff — ``"const-tree"`` is worth reaching for only when it will actually save work. A useful rule of thumb: *does the system go out of its way to make two equal instances of this type share a pointer?* Canonical types, interned constants, cached shapes, and op metadata usually do. General expression and statement nodes usually don't — and also fail condition 2. Prefer ``"const-tree"`` for the type / attribute / metadata layer of the IR, not the expression / statement layer. Note also that condition 2 is a *whole-subgraph* property: once a field holds an ``Expr`` (which may one day contain a ``Var``), the annotation silently commits the type to that invariant — a later refactor embedding a ``Var`` becomes a correctness break rather than a local change. Why not use it everywhere? ~~~~~~~~~~~~~~~~~~~~~~~~~~ Most IR nodes are immutable, but many transitively contain variables (e.g., ``x + 1`` contains the ``"var"`` node ``x``). The pointer shortcut fires only when both sides of a comparison reference the **same object** — but when that sharing exists, skipping traversal also skips the variable occurrences inside, and mappings that should have been recorded are silently missed. Suppose the ``+`` node were incorrectly annotated as ``"const-tree"``, and consider comparing two tuples that share the ``+`` subtree via pointer identity: .. code-block:: text shared = x + 1 # pointer P, contains var x lhs = (shared, x) # .0 = P, .1 = var x rhs = (shared, y) # .0 = P, .1 = var y (different Var) structural_equal(lhs, rhs, map_free_vars=True) With the ``+`` annotated as plain ``"tree"`` (correct): - ``.0``: traverse into ``shared`` on both sides, visit ``x`` at ``.lhs``, record the mapping ``x ↔ x``. - ``.1``: look up ``x`` → maps to ``x``, but rhs is ``y``. **NOT EQUAL** ✓ With the ``+`` annotated as ``"const-tree"`` (the bug): - ``.0``: pointer shortcut fires on ``shared`` (both sides reference P). Fields are skipped, ``x`` inside is never visited, no mapping is recorded. - ``.1``: compare ``x`` vs ``y``. No existing mapping, and ``map_free_vars=True`` lets a new one be recorded as ``x ↔ y``. **EQUAL** ✗ (wrong) The following diagram illustrates the shared structure. The ``+`` node (``shared``) has two incoming ``.0`` edges — one from each side — which is exactly the situation in which the pointer shortcut fires: .. mermaid:: graph TD LT["lhs: (_, _)"] RT["rhs: (_, _)"] ADD["shared = x + 1
const-tree
same pointer on both sides"] X["x : var"] ONE["1"] Y["y : var"] LT -->|".0"| ADD RT -->|".0"| ADD LT -->|".1"| X RT -->|".1"| Y ADD -->|".lhs"| X ADD -->|".rhs"| ONE style ADD fill:#fff3cd style X fill:#f8d7da style Y fill:#f8d7da The same failure mode arises whenever a shared subtree containing a ``"var"`` is compared inside any definition region (e.g., the body of a ``Lambda`` whose params field is ``structural_eq="def"``), not only under ``map_free_vars=True``. ``"dag"`` — Sharing-Aware Comparison ------------------------------------- .. code-block:: python @py_class(structural_eq="dag") class Binding(Object): var: Var value: Expr **Meaning**: "This node lives in a graph where pointer sharing is semantically meaningful. Two graphs are equal only if they have the same content **and** the same sharing structure." Why it exists ~~~~~~~~~~~~~ In dataflow IR, sharing matters. Consider: .. code-block:: text # Program A: shared — compute once, use twice let s = x + 1 in (s, s) # Program B: independent — compute twice (x + 1, x + 1) Program A computes ``x + 1`` once and references it twice; Program B computes it independently twice. Under ``"tree"`` these are equal; under ``"dag"`` they are **not**: .. mermaid:: graph TD subgraph "Program A — DAG" TA["(_, _)"] SA["s = x + 1"] TA -->|".0"| SA TA -->|".1"| SA end subgraph "Program B — Tree" TB["(_, _)"] A1["x + 1"] A2["x + 1"] TB -->|".0"| A1 TB -->|".1"| A2 end SA -. "NOT EQUAL under dag
(sharing structure differs)" .-> A1 style SA fill:#d4edda style A1 fill:#d4edda style A2 fill:#f8d7da How ``"dag"`` detects sharing ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ``"dag"`` maintains a bijective (one-to-one) mapping between objects that have been successfully compared. When the same object appears again, it checks whether the *pairing* is consistent: .. code-block:: text Comparing Program A vs Program B: .0: s ↔ (x+1)₁ → content equal, record pairing: s ↔ (x+1)₁ .1: s ↔ (x+1)₂ → s already paired with (x+1)₁, not (x+1)₂ → NOT EQUAL The mapping is **bijective**: if ``a`` is paired with ``b``, no other object can pair with either ``a`` or ``b``. This prevents false positives in both directions. **Example of the reverse direction.** .. code-block:: text lhs: (a, b) rhs: (a, a) where a ≅ b (same content) .0: a₁ ↔ a₂ → equal, record a₁ ↔ a₂ .1: b₁ ↔ a₂ → b₁ is new, but a₂ already paired with a₁ → NOT EQUAL Without the reverse check, the second comparison would proceed to content comparison, find ``b₁ ≅ a₂``, and incorrectly succeed. Full comparison: ``"tree"`` vs ``"dag"`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 48 13 13 * - Scenario - ``"tree"`` - ``"dag"`` * - both trees with same content - Equal - Equal * - both DAGs, same sharing shape - Equal - Equal * - ``let s = e in (s, s)`` vs ``(e, e')`` where ``e ≅ e'`` - Equal - **Not equal** * - ``(a, b)`` vs ``(a, a)`` where ``a ≅ b`` - Equal - **Not equal** .. _var-kind: ``"var"`` — Bound Variables ---------------------------- .. code-block:: python @py_class(structural_eq="var") class Var(Object): name: str = field(structural_eq="ignore") # alpha-equivalent vars differ in name type: Type # participates in equality **Meaning**: "This is a variable. Two variables are equal if they are **bound in corresponding positions**, not if they have the same name." The ``name`` field is almost always marked ``structural_eq="ignore"`` because alpha-equivalent variables have different names. Other fields such as ``type`` *are* compared — but only at the binding site (see :ref:`var-fields`). The problem ~~~~~~~~~~~ .. code-block:: text fun x → x + 1 should equal fun y → y + 1 Variables are not defined by their content, such as their name. They are defined by **where they are introduced** and **how they are used**. ``x`` and ``y`` above are interchangeable because they occupy the same binding position and are used in the same way. How it works: definition regions ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ``"var"`` works together with ``field(structural_eq="def")`` (see :ref:`field-annotations`). A field marked ``structural_eq="def"`` is a **definition region** — it's where new variable bindings are introduced. - **Inside a definition region**: encountering two different variables establishes a correspondence ("treat ``x`` as equivalent to ``y``"). - **Outside a definition region**: variables are only equal if a prior correspondence already exists, or they are the same pointer. The following diagram traces the comparison of two alpha-equivalent functions: .. mermaid:: sequenceDiagram participant C as Comparator participant L as lhs: fun x → x + 1 participant R as rhs: fun y → y + 1 Note over C: Field "params" has structural_eq="def" C->>L: get params → [x] C->>R: get params → [y] Note over C: Enter definition region C->>C: Compare x ↔ y: both are Vars Note over C: Record mapping: x ↔ y Note over C: Exit definition region Note over C: Field "body" — normal region C->>L: get body → x + 1 C->>R: get body → y + 1 C->>C: Compare + fields... C->>C: x ↔ y: lookup finds x→y ✓ C->>C: 1 ↔ 1: equal ✓ Note over C: Result: EQUAL ✓ **Without** a definition region, the same variables would **not** be equal: .. code-block:: text # Bare expressions, no enclosing function: x + 1 vs y + 1 → NOT EQUAL (no definition region, different pointers) .. _var-fields: Fields and the sticky mapping ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ A ``"var"`` type still has fields, and non-ignored fields *are* compared — but only on the **first** encounter of a var pair. Once a mapping is recorded, subsequent occurrences look up the mapping and skip field comparison entirely. Take the ``Var`` declaration from the top of this section: ``name`` is ignored, but ``type`` is not. The first time a pair of vars is seen in a definition region, their ``type`` fields are compared and the mapping is only established if they match. After that, the mapping is **sticky** — later occurrences trust the correspondence regardless of those fields: .. list-table:: :header-rows: 1 :widths: 55 45 * - Scenario - Result * - ``Var("x", int)`` vs ``Var("y", int)`` on first encounter - Fields match → mapping ``x ↔ y`` recorded → **Equal** * - ``Var("x", int)`` vs ``Var("y", float)`` on first encounter - Fields differ → **Not equal** * - ``Var("x", int)`` vs ``Var("y", float)`` when ``x ↔ y`` already mapped - Lookup succeeds → **Equal** (types are *not* rechecked) For IRs where type consistency is part of well-formedness, this is usually sufficient: a well-formed program uses each var with a consistent type at every occurrence, so the first-encounter check at the binding site covers the rest. If you truly want types re-verified at every use, they don't belong on the ``"var"`` node — lift them into the surrounding expression/statement node where they participate in normal ``"tree"`` comparison. Full comparison: with and without definition regions ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 55 22 22 * - Scenario - With ``"def"`` - Without * - ``fun x → x + 1`` vs ``fun y → y + 1`` - Equal - n/a * - ``fun x → x + 1`` vs ``fun y → x + 1`` - **Not equal** (body uses ``x`` but mapping says ``y``) - n/a * - ``fun (x, y) → x + y`` vs ``fun (a, b) → a + b`` - Equal (x↔a, y↔b) - n/a * - ``fun (x, y) → x + y`` vs ``fun (a, b) → b + a`` - **Not equal** (x↔a but body uses ``x`` where ``b`` appears) - n/a * - ``x + 1`` vs ``y + 1`` (bare) - n/a - **Not equal** * - ``x + 1`` vs ``x + 1`` (same pointer) - n/a - Equal Inconsistent variable usage ~~~~~~~~~~~~~~~~~~~~~~~~~~~ The bijective mapping catches inconsistencies. Consider: .. code-block:: text fun (x, y) → x + x vs fun (a, b) → a + b .. mermaid:: sequenceDiagram participant C as Comparator participant L as lhs: fun (x, y) → x + x participant R as rhs: fun (a, b) → a + b Note over C: Definition region (params) C->>C: x ↔ a → record x↔a ✓ C->>C: y ↔ b → record y↔b ✓ Note over C: Body: x + x vs a + b C->>C: x ↔ a → lookup x→a, matches ✓ C->>C: x ↔ b → lookup x→a, but rhs is b ≠ a → FAIL ✗ Note over C: Result: NOT EQUAL ✓ The ``map_free_vars`` flag ~~~~~~~~~~~~~~~~~~~~~~~~~~ ``structural_equal(lhs, rhs, map_free_vars=True)`` starts the comparison in "definition region" mode. This is useful for comparing standalone expressions where you want alpha-equivalence at the top level without an enclosing function: .. code-block:: python # With map_free_vars=True: structural_equal(x + 1, y + 1, map_free_vars=True) # → True # With map_free_vars=False (default): structural_equal(x + 1, y + 1) # → False ``"singleton"`` — Singletons ------------------------------ .. code-block:: python @py_class(structural_eq="singleton") class Op(Object): name: str **Meaning**: "There is exactly one instance of this object per logical identity. Pointer equality is the only valid comparison." No content comparison is ever performed. Different pointers are always unequal; same pointer is always equal. .. code-block:: python op_conv = Op.get("nn.conv2d") op_relu = Op.get("nn.relu") structural_equal(op_conv, op_conv) # → True (same pointer) structural_equal(op_conv, op_relu) # → False (different pointers) .. _field-annotations: Field-Level Annotations ----------------------- The ``structural_eq`` parameter on :func:`~tvm_ffi.dataclasses.field` controls how structural equality/hashing treats that specific field. ``structural_eq="ignore"`` — Exclude a field ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. code-block:: python @py_class(structural_eq="tree") class MyNode(Object): value: int span: str = field(structural_eq="ignore") **Meaning**: "This field is not part of the node's structural identity. Skip it during comparison and hashing." Use for: - **Source locations** (``span``) — where the node came from in source code doesn't affect what it means. - **Cached/derived values** — computed from other fields, would be redundant to compare. - **Debug annotations** — names, comments, metadata for human consumption. ``structural_eq="def-recursive"`` / ``"def-non-recursive"`` — Definition region ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. code-block:: python @py_class(structural_eq="tree") class Lambda(Object): params: list[Var] = field(structural_eq="def-recursive") body: Expr **Meaning**: "This field introduces new variable bindings. When comparing or hashing this field, allow new variable correspondences to be established." This is the counterpart to ``"var"``. A ``"var"`` type says "I am a variable"; the ``"def-*"`` flags on a field say "this field is where variables are defined." Together they enable alpha-equivalence: comparing functions up to consistent variable renaming. There are two flavors of definition region, distinguished by what happens when a ``"var"`` reached through the field carries its own sub-fields (for example, a shape annotation in the var's type): - ``"def-recursive"`` (alias: ``"def"``) — the variable's sub-fields stay inside the definition region. Any free variables encountered in those sub-fields are themselves treated as fresh definitions at the same site. One example is **function parameter lists**, where the value var and any shape parameters in its type are co-introduced together at the function boundary. - ``"def-non-recursive"`` — only the immediate variable(s) reached through the field bind. The variable's sub-fields are walked outside the definition region, so any free variables there are *use* references that must resolve against an outer-scope binding. One example is a **normal binding** whose value type references outer-scope shape parameters (a ``let v = expr`` where ``v``'s type refers to vars defined earlier). When the distinction does not matter (no nested free vars under the bound variable), either flavor works and ``"def-recursive"`` is the conventional default — that's why the bare ``"def"`` alias resolves to it. Use for: - **Function parameter lists** — ``"def-recursive"`` so shape parameters in each param's type co-introduce at the same site. - **Normal binding left-hand sides** (let bindings, for-loop iterators) whose value type references outer-scope vars — ``"def-non-recursive"`` so those references don't rebind. - **Any field that introduces names into scope** — pick the flavor that matches the binding form's contract; default to ``"def-recursive"`` when in doubt. .. _sequal-shash: Custom Equality and Hashing: ``__s_equal__`` / ``__s_hash__`` -------------------------------------------------------------- For types where the default field-by-field traversal is insufficient (for example, fields that need to be visited in a specific order, cross-field invariants, or sub-values that need a different ``def_region`` setting than the declarative field flags allow), you can register custom callbacks as **type attributes**: - ``__s_equal__`` — custom structural equality logic. - ``__s_hash__`` — custom structural hashing logic. These are the *only* supported way to override structural comparison. ``structural_equal`` / ``structural_hash`` never consult Python ``__eq__`` / ``__hash__`` — those dunders serve a separate purpose (``==`` and ``hash()``, which default to pointer identity). When either hook is registered, it replaces the default field iteration for that type. All kind-specific machinery (``"dag"`` memoization, ``"var"`` mapping, the pointer shortcut of ``"const-tree"``, etc.) is still managed by the framework — the custom callback only controls *which* sub-values are compared or hashed, *in what order*, and *with what* ``def_region`` flag. Signatures ~~~~~~~~~~ ``__s_equal__``: .. code-block:: text (self, other, eq_cb) -> bool eq_cb(lhs, rhs, def_region_kind: int, field_name: str) -> bool ``__s_hash__``: .. code-block:: text (self, init_hash: int, hash_cb) -> int hash_cb(value, init_hash: int, def_region_kind: int) -> int The ``def_region_kind`` argument on each recursive call mirrors the field-level ``"def-*"`` flags and controls whether the sub-value is compared/hashed inside a definition region: - ``0`` — not in a def region (matches ``None`` on a field). - ``1`` — recursive def region (matches ``"def-recursive"``, alias ``"def"``). - ``2`` — non-recursive def region (matches ``"def-non-recursive"``). For back-compat with the original single-flag API, the callback also accepts a plain ``bool``: ``True`` is treated as ``1`` (recursive) and ``False`` as ``0`` (not in a def region). The Python examples below use ``True`` / ``False`` for that reason; pass an explicit ``2`` (or the ``kTVMFFIDefRegionKindNonRecursive`` enum value from C++) when the non-recursive kind is needed. The ``field_name`` argument on ``eq_cb`` is used only for mismatch path reporting from :py:func:`~tvm_ffi.get_first_structural_mismatch`. Example (Python) ~~~~~~~~~~~~~~~~ .. code-block:: python @py_class(structural_eq="tree") class Lambda(Object): params: list body: Any comment: str # not part of identity, but also not iterated below def __s_equal__(self, other, eq_cb): # params is a definition region; body is not. if not eq_cb(self.params, other.params, True, "params"): return False if not eq_cb(self.body, other.body, False, "body"): return False return True def __s_hash__(self, init_hash, hash_cb): h = hash_cb(self.params, init_hash, True) h = hash_cb(self.body, h, False) return h The two methods must agree: if ``__s_equal__`` considers two instances equal, ``__s_hash__`` must produce the same hash for them. Example (C++) ~~~~~~~~~~~~~ .. code-block:: c++ class MyNodeObj : public Object { public: Array params; Array body; bool SEqual(const MyNodeObj* other, ffi::TypedFunction cmp) const { if (!cmp(params, other->params, /*def_region=*/true, "params")) return false; if (!cmp(body, other->body, /*def_region=*/false, "body")) return false; return true; } int64_t SHash(int64_t init_hash, ffi::TypedFunction hash) const { int64_t h = hash(params, init_hash, /*def_region=*/true); h = hash(body, h, /*def_region=*/false); return h; } static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("params", &MyNodeObj::params) .def_ro("body", &MyNodeObj::body); refl::TypeAttrDef() .def(refl::type_attr::kSEqual, &MyNodeObj::SEqual) .def(refl::type_attr::kSHash, &MyNodeObj::SHash); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my.Node", MyNodeObj, Object); }; See :cpp:var:`tvm::ffi::reflection::type_attr::kSEqual` and :cpp:var:`tvm::ffi::reflection::type_attr::kSHash` in ``include/tvm/ffi/reflection/accessor.h`` for the full reference. All Kinds at a Glance --------------------- The following diagram visualizes the five comparable kinds, arranged by how much structural information they track: .. mermaid:: graph LR UI["singleton
pointer only"] TN["tree
content only"] CTN["const-tree
content + pointer shortcut"] DN["dag
content + sharing"] FV["var
content + binding position"] UI --- TN TN --- CTN TN --- DN TN --- FV style UI fill:#e2e3e5 style TN fill:#d4edda style CTN fill:#d4edda style DN fill:#cce5ff style FV fill:#fff3cd .. list-table:: :header-rows: 1 :widths: 18 18 18 18 18 * - - Content comparison - Pointer shortcut - Tracks sharing - Tracks binding position * - ``"singleton"`` - No - Yes (only) - No - No * - ``"tree"`` - Yes - No - No - No * - ``"const-tree"`` - Yes - Yes (fast path) - No - No * - ``"dag"`` - Yes - No - Yes - No * - ``"var"`` - Yes - No - No - Yes Decision Guide -------------- When defining a new type: .. mermaid:: graph TD Start["New @py_class type"] --> Q1{"Singleton?
(one instance per
logical identity)"} Q1 -->|Yes| UI["structural_eq="singleton""] Q1 -->|No| Q2{"Represents a
variable binding?"} Q2 -->|Yes| FV["structural_eq="var""] Q2 -->|No| Q3{"Pointer sharing
semantically
meaningful?"} Q3 -->|Yes| DN["structural_eq="dag""] Q3 -->|No| Q4{"Immutable AND
no transitive
var children?"} Q4 -->|Yes| CTN["structural_eq="const-tree""] Q4 -->|No| TN["structural_eq="tree""] style UI fill:#e2e3e5 style FV fill:#fff3cd style DN fill:#cce5ff style CTN fill:#d4edda style TN fill:#d4edda For fields: .. mermaid:: graph TD Start["field() parameter"] --> Q1{"Irrelevant to
structural identity?
(span, cache, debug)"} Q1 -->|Yes| IGN["structural_eq="ignore""] Q1 -->|No| Q2{"Introduces new
variable bindings?"} Q2 -->|Yes| DEF["structural_eq="def""] Q2 -->|No| NONE["No flag needed"] style IGN fill:#f8d7da style DEF fill:#fff3cd style NONE fill:#d4edda Worked Example -------------- Putting it all together for a function node with parameters, body, and source location: .. code-block:: python @py_class(structural_eq="tree") class Lambda(Object): params: list[Var] = field(structural_eq="def") body: Expr span: str = field(structural_eq="ignore", default="") @py_class(structural_eq="var") class Var(Object): name: str = field(structural_eq="ignore") @py_class(structural_eq="singleton") class Op(Object): name: str With these annotations, alpha-equivalent functions are structurally equal: .. code-block:: text # These two are structurally equal: fun [x] → x + 1 (span="a.py:1") fun [y] → y + 1 (span="b.py:5") # - params has structural_eq="def" → x maps to y # - body uses that mapping → (x + 1) ≅ (y + 1) # - span has structural_eq="ignore" → locations don't matter And in Python: .. code-block:: python from tvm_ffi import structural_equal, structural_hash x, y = Var("x"), Var("y") f1 = Lambda([x], x + 1, span="a.py:1") f2 = Lambda([y], y + 1, span="b.py:5") assert structural_equal(f1, f2) # alpha-equivalent assert structural_hash(f1) == structural_hash(f2) # same hash tvm-ffi-0.1.12/docs/concepts/tensor.rst000066400000000000000000000416741521067262500177770ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Tensor and DLPack ================= At runtime, TVM-FFI often needs to accept tensors from many sources: * Frameworks (e.g. PyTorch, JAX, PaddlePaddle) via :py:meth:`array_api.array.__dlpack__`; * C/C++ callers passing :c:struct:`DLTensor* `; * Tensors allocated by a library but managed by TVM-FFI itself. TVM-FFI standardizes on **DLPack as the lingua franca**: tensors are built on top of DLPack structs with additional C++ convenience methods and minimal extensions for ownership management. .. tip:: Prefer :cpp:class:`tvm::ffi::TensorView` or :cpp:class:`tvm::ffi::Tensor` in C++ code; they provide safer and more convenient abstractions over raw DLPack structs. This tutorial covers common usage patterns, tensor classes, and how tensors flow across ABI boundaries. Glossary -------- DLPack A cross-library tensor interchange standard defined in the small C header ``dlpack.h``. It defines pure C data structures for describing n-dimensional arrays and their memory layout, including :c:struct:`DLTensor`, :c:struct:`DLManagedTensorVersioned`, :c:struct:`DLDataType`, :c:struct:`DLDevice`, and related types. View (non-owning) A "header" that describes a tensor but does not own its memory. When a consumer receives a view, it must respect that the producer owns the underlying storage and controls its lifetime. The view is valid only while the producer guarantees it remains valid. Managed object (owning) An object that includes lifetime management, using reference counting or a cleanup callback mechanism. This establishes a contract between producer and consumer about when the consumer's ownership ends. .. note:: As a loose analogy, think of **view** vs. **managed** as similar to ``T*`` (raw pointer) vs. ``std::shared_ptr`` (reference-counted pointer) in C++. Common Usage ------------ This section introduces the most important APIs for day-to-day use in C++ and Python. Kernel Signatures ~~~~~~~~~~~~~~~~~ A typical kernel implementation accepts :cpp:class:`TensorView ` parameters, validates metadata (dtype, shape, device), and then accesses the data pointer for computation. Use :c:macro:`TVM_FFI_THROW` to report validation errors (see :doc:`exception_handling`): .. code-block:: cpp #include void MyKernel(tvm::ffi::TensorView input, tvm::ffi::TensorView output) { // Validate dtype & device if (input.dtype() != DLDataType{kDLFloat, 32, 1}) TVM_FFI_THROW(TypeError) << "Expect float32 input, but got " << input.dtype(); if (input.device() != DLDevice{kDLCUDA, 0}) TVM_FFI_THROW(ValueError) << "Expect input on CUDA:0, but got " << input.device(); // Access data pointer float* input_data_ptr = static_cast(input.data_ptr()); float* output_data_ptr = static_cast(output.data_ptr()); Kernel<<<...>>>(..., input_data_ptr, output_data_ptr, ...); } On the C++ side, the following APIs are available to query a tensor's metadata: :cpp:func:`TensorView::shape() ` and :cpp:func:`Tensor::shape() ` Shape array. For kernels exported as TVM-FFI functions, see :doc:`func_module`. :cpp:func:`TensorView::dtype() ` and :cpp:func:`Tensor::dtype() ` element data type :cpp:func:`TensorView::data_ptr() ` and :cpp:func:`Tensor::data_ptr() ` base pointer to the tensor's data :cpp:func:`TensorView::device() ` and :cpp:func:`Tensor::device() ` device type and id :cpp:func:`TensorView::byte_offset() ` and :cpp:func:`Tensor::byte_offset() ` byte offset to the first element :cpp:func:`TensorView::ndim() ` and :cpp:func:`Tensor::ndim() ` number of dimensions (:cpp:func:`ShapeView::size `) :cpp:func:`TensorView::numel() ` and :cpp:func:`Tensor::numel() ` total number of elements (:cpp:func:`ShapeView::Product `) PyTorch Interop ~~~~~~~~~~~~~~~ On the Python side, :py:class:`tvm_ffi.Tensor` is a managed n-dimensional array that: * can be created via :py:func:`tvm_ffi.from_dlpack(ext_tensor, ...) ` to import tensors from external frameworks, e.g., :ref:`PyTorch `, :ref:`JAX `, :ref:`PaddlePaddle `, :ref:`NumPy/CuPy `; * implements the DLPack protocol so it can be passed back to frameworks without copying, e.g., :py:func:`torch.from_dlpack`. The following example demonstrates a typical round-trip pattern: .. code-block:: python import tvm_ffi import torch x_torch = torch.randn(1024, device="cuda") x_tvm_ffi = tvm_ffi.from_dlpack(x_torch, require_contiguous=True) x_torch_again = torch.from_dlpack(x_tvm_ffi) In this example, :py:func:`tvm_ffi.from_dlpack` creates ``x_tvm_ffi``, which views the same memory as ``x_torch``. Similarly, :py:func:`torch.from_dlpack` creates ``x_torch_again``, which shares the underlying buffer with both ``x_tvm_ffi`` and ``x_torch``. No data is copied in either direction. C++ Allocation ~~~~~~~~~~~~~~ TVM-FFI is not a kernel library and is not linked to any specific device memory allocator or runtime. However, it provides standardized allocation entry points for kernel library developers by interfacing with the surrounding framework's allocator - for example, using PyTorch's allocator when running inside a PyTorch environment. **Env Allocator.** Use :cpp:func:`Tensor::FromEnvAlloc() ` along with C API :cpp:func:`TVMFFIEnvTensorAlloc` to allocate a tensor using the framework's allocator. .. code-block:: cpp Tensor tensor = Tensor::FromEnvAlloc( TVMFFIEnvTensorAlloc, /*shape=*/{1, 2, 3}, /*dtype=*/DLDataType({kDLFloat, 32, 1}), /*device=*/DLDevice({kDLCPU, 0}) ); In a PyTorch environment, this is equivalent to :py:func:`torch.empty`. .. warning:: While allocation APIs are available, it is generally **recommended** to avoid allocating tensors inside kernels. Instead, prefer pre-allocating outputs and passing them as :cpp:class:`tvm::ffi::TensorView` parameters. This approach: - avoids memory fragmentation and performance pitfalls, - prevents CUDA graph incompatibilities on GPU, and - allows the outer framework to control allocation policy (pools, device strategies, etc.). **Custom Allocator.** Use :cpp:func:`Tensor::FromNDAlloc(custom_alloc, ...) `, or its advanced variant :cpp:func:`Tensor::FromNDAllocStrided(custom_alloc, ...) `, to allocate a tensor with a user-provided allocation callback. The following example uses ``cudaMalloc``/``cudaFree`` as custom allocators for GPU tensors: .. code-block:: cpp struct CUDANDAlloc { void AllocData(DLTensor* tensor) { size_t data_size = ffi::GetDataSize(*tensor); void* ptr = nullptr; cudaError_t err = cudaMalloc(&ptr, data_size); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err); tensor->data = ptr; } void FreeData(DLTensor* tensor) { if (tensor->data != nullptr) { cudaError_t err = cudaFree(tensor->data); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err); tensor->data = nullptr; } } }; ffi::Tensor cuda_tensor = ffi::Tensor::FromNDAlloc( CUDANDAlloc(), /*shape=*/{3, 4, 5}, /*dtype=*/DLDataType({kDLFloat, 32, 1}), /*device=*/DLDevice({kDLCUDA, 0}) ); C++ Stream Handling ~~~~~~~~~~~~~~~~~~~ Stream context is essential for GPU kernel execution. While CUDA does not have a global context for default streams, frameworks like PyTorch maintain a "current stream" per device (:py:func:`torch.cuda.current_stream`), and kernel libraries must read this stream from the embedding environment. As a hardware-agnostic abstraction layer, TVM-FFI is not linked to any specific stream management library. However, to ensure GPU kernels launch on the correct stream, it provides standardized APIs to obtain the stream context from the host framework (e.g., PyTorch). **Obtain Stream Context.** Use the C API :cpp:func:`TVMFFIEnvGetStream` to obtain the current stream for a given device: .. code-block:: cpp void func(ffi::TensorView input, ...) { ffi::DLDevice device = input.device(); cudaStream_t stream = reinterpret_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); } This is equivalent to the following PyTorch C++ code: .. code-block:: cpp void func(at::Tensor input, ...) { c10::Device device = input.device(); cudaStream_t stream = reinterpret_cast( c10::cuda::getCurrentCUDAStream(device.index()).stream()); } **Auto-Update Stream Context.** When converting framework tensors via :py:func:`tvm_ffi.from_dlpack`, TVM-FFI automatically updates the stream context to match the device of the converted tensor. For example, when converting a PyTorch tensor on ``torch.device('cuda:3')``, TVM-FFI automatically captures the stream from :py:func:`torch.cuda.current_stream(device='cuda:3')`. **Set Stream Context.** Use :py:func:`tvm_ffi.use_torch_stream` or :py:func:`tvm_ffi.use_raw_stream` to manually set the stream context when automatic detection is insufficient. Tensor Classes -------------- This section defines each tensor type in the TVM-FFI C++ API and explains its intended usage. Exact C layout details are covered in :ref:`Tensor Layouts `. .. tip:: On the Python side, only :py:class:`tvm_ffi.Tensor` exists. It strictly follows DLPack semantics for interop and can be converted to PyTorch via :py:func:`torch.from_dlpack`. DLPack Tensors ~~~~~~~~~~~~~~ DLPack tensors come in two main flavors: *Non-owning* object, :c:struct:`DLTensor` The tensor descriptor is a **view** of the underlying data. It describes the device the tensor lives on, its shape, dtype, and data pointer. It does not own the underlying data. *Owning* object, :c:struct:`DLManagedTensorVersioned`, or its legacy counterpart :c:struct:`DLManagedTensor` It is a **managed** variant that wraps a :c:struct:`DLTensor` descriptor with additional fields. Notably, it includes a ``deleter`` callback that releases ownership when the consumer is done with the tensor, and an opaque ``manager_ctx`` handle used by the producer to store additional context. TVM-FFI Tensors ~~~~~~~~~~~~~~~ Similarly, TVM-FFI defines two main tensor types in C++: *Non-owning* object, :cpp:class:`tvm::ffi::TensorView` A thin C++ wrapper around :c:struct:`DLTensor` for inspecting metadata and accessing the data pointer. It is designed for **kernel authors** to inspect metadata and access the underlying data pointer during a call, without taking ownership of the tensor's memory. Being a **view** also means you must ensure the backing tensor remains valid while you use it. *Owning* object, :cpp:class:`tvm::ffi::TensorObj` and :cpp:class:`tvm::ffi::Tensor` :cpp:class:`Tensor `, similar to ``std::shared_ptr``, is the managed class to hold heap-allocated :cpp:class:`TensorObj `. Once the reference count drops to zero (see :ref:`object-reference-counting`), the cleanup logic deallocates the descriptor and releases ownership of the underlying data buffer. .. note:: - For handwritten C++, always use TVM-FFI tensors over DLPack's raw C tensors. - For compiler development, DLPack's raw C tensors are recommended because C is easier to target from codegen. The owning :cpp:class:`Tensor ` is the recommended interface for passing around managed tensors. Use owning tensors when you need one or more of the following: * return a tensor from a function across ABI, which will be converted to :cpp:class:`tvm::ffi::Any`; * allocate an output tensor as the producer, and hand it to a kernel consumer; * store a tensor in a long-lived object. .. admonition:: :cpp:class:`TensorObj ` vs :cpp:class:`Tensor ` :class: hint :cpp:class:`Tensor ` is an intrusive pointer of a heap-allocated :cpp:class:`TensorObj `. As an analogy to ``std::shared_ptr``, think of .. code-block:: cpp using Tensor = std::shared_ptr; You can convert between the two types: - :cpp:func:`Tensor::get() ` converts it to :cpp:class:`TensorObj* `. - :cpp:func:`GetRef\ ` converts a :cpp:class:`TensorObj* ` back to :cpp:class:`Tensor `. .. _layout-and-conversion: Tensor Layouts ~~~~~~~~~~~~~~ :ref:`Figure 1 ` summarizes the layout relationships among DLPack tensors and TVM-FFI tensors. All tensor classes are POD-like; :cpp:class:`tvm::ffi::TensorObj` is also a standard TVM-FFI object, typically heap-allocated and reference-counted. .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tvm-ffi/tensor-layout.png :alt: Layout of DLPack Tensors and TVM-FFI Tensors :align: center :name: fig:layout-tensor Figure 1. Layout specification of DLPack tensors and TVM-FFI tensors. All the tensor types share :c:struct:`DLTensor` as the common descriptor, while carrying different metadata and ownership semantics. As demonstrated in the figure, all tensor classes share :c:struct:`DLTensor` as the common descriptor. In particular, - :c:struct:`DLTensor` and :cpp:class:`TensorView ` share the exact same memory layout. - :c:struct:`DLManagedTensorVersioned` and :cpp:class:`TensorObj ` both have a deleter callback to manage the lifetime of the underlying data buffer, while :c:struct:`DLTensor` and :cpp:class:`TensorView ` do not. - Compared with :cpp:class:`TensorView `, :cpp:class:`TensorObj ` has an extra TVM-FFI object header, making it reference-countable via the standard managed reference :cpp:class:`Tensor `. What Tensor Is Not ~~~~~~~~~~~~~~~~~~ TVM-FFI is not a tensor library. While it provides a unified representation for tensors, it does not include: * kernels (e.g., vector addition, matrix multiplication), * host-device copy or synchronization primitives, * advanced indexing or slicing, or * automatic differentiation or computational graph support. Conversion between :cpp:class:`TVMFFIAny` ----------------------------------------- At the stable C ABI boundary, TVM-FFI passes values using :cpp:class:`Any ` (owning) or :cpp:class:`AnyView ` (non-owning). Tensors have two possible representations: * **Non-owning:** :c:struct:`DLTensor* ` with type index :cpp:enumerator:`TVMFFITypeIndex::kTVMFFIDLTensorPtr` * **Owning:** :cpp:class:`TensorObj* ` with type index :cpp:enumerator:`TVMFFITypeIndex::kTVMFFITensor` When extracting a tensor from :cpp:class:`TVMFFIAny`, check the :cpp:member:`type_index ` to determine the representation before conversion. .. important:: An owning tensor can be converted to a non-owning view, but not vice versa. See :ref:`abi-tensor` for C code examples demonstrating: - Extracting a :c:struct:`DLTensor` pointer from :cpp:class:`TVMFFIAny` - Constructing a :cpp:class:`~tvm::ffi::TensorObj` from DLPack - Exporting a :cpp:class:`~tvm::ffi::TensorObj` to DLPack Further Reading --------------- - :doc:`object_and_class`: The object system that backs :cpp:class:`~tvm::ffi::TensorObj` - :doc:`any`: How tensors are stored in :cpp:class:`~tvm::ffi::Any` containers - :doc:`func_module`: How tensors are passed through TVM-FFI function calls - :doc:`exception_handling`: Throwing errors during tensor validation - :doc:`abi_overview`: Low-level C ABI details for tensor conversion - :doc:`../guides/kernel_library_guide`: Best practices for building kernel libraries with TVM-FFI - :external+dlpack:doc:`DLPack C API `: The underlying tensor interchange standard tvm-ffi-0.1.12/docs/conf.py000066400000000000000000000372311521067262500154060ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Sphinx configuration for the tvm-ffi documentation site.""" # -*- coding: utf-8 -*- from __future__ import annotations import importlib import inspect import os import shutil import subprocess import sys from pathlib import Path import setuptools_scm import sphinx os.environ["TVM_FFI_BUILD_DOCS"] = "1" build_exhale = os.environ.get("BUILD_CPP_DOCS", "0") == "1" build_rust_docs = os.environ.get("BUILD_RUST_DOCS", "0") == "1" # Auto-detect sphinx-autobuild: Check if sphinx-autobuild is in the execution path is_autobuild = any("sphinx-autobuild" in str(arg) for arg in sys.argv) # -- Path constants ------------------------------------------------------- _DOCS_DIR = Path(__file__).resolve().parent _RUST_DIR = _DOCS_DIR.parent / "rust" # -- General configuration ------------------------------------------------ # Determine version without reading pyproject.toml # Always use setuptools_scm (assumed available in docs env) __version__ = setuptools_scm.get_version(root="..") project = "tvm-ffi" author = "Apache TVM FFI contributors" version = __version__ release = __version__ # -- Extensions and extension configurations -------------------------------- extensions = [ "breathe", "myst_parser", "nbsphinx", "autodocsumm", "sphinx_design", "sphinx.ext.autodoc", "sphinx.ext.autosectionlabel", "sphinx.ext.autosummary", "sphinx.ext.intersphinx", "sphinx.ext.mathjax", "sphinx.ext.napoleon", "sphinx.ext.viewcode", "sphinx.ext.ifconfig", "sphinx_autodoc_typehints", "sphinx_copybutton", "sphinx_reredirects", "sphinx_tabs.tabs", "sphinx_toolbox.collapse", "sphinxcontrib.httpdomain", "sphinxcontrib.mermaid", ] if build_exhale: extensions.append("exhale") breathe_default_project = "tvm-ffi" breathe_projects = {"tvm-ffi": "./_build/doxygen/xml"} exhaleDoxygenStdin = """ INPUT = ../include PREDEFINED += TVM_FFI_DLL= TVM_FFI_DLL_EXPORT= TVM_FFI_INLINE= \ TVM_FFI_EXTRA_CXX_API= TVM_FFI_WEAK= TVM_FFI_DOXYGEN_MODE \ TVM_FFI_COLD_CODE= \ TVM_FFI_PREDICT_FALSE(x)=x TVM_FFI_PREDICT_TRUE(x)=x \ __cplusplus=201703 EXCLUDE_SYMBOLS += *details* *TypeTraits* std \ *use_default_type_traits_v* *is_optional_type_v* *operator* \ tvm::ffi::reflection::default_ tvm::ffi::reflection::default_factory EXCLUDE_PATTERNS += */function_details.h */container_details.h ENABLE_PREPROCESSING = YES MACRO_EXPANSION = YES WARNINGS = YES WARN_AS_ERROR = FAIL_ON_WARNINGS_PRINT # if your Doxygen version supports it """ exhaleAfterTitleDescription = """ This page contains the full API index for the C++ API. """ # Setup the exhale extension exhale_args = { "containmentFolder": "reference/cpp/generated", "rootFileName": "index.rst", "doxygenStripFromPath": "../include", "rootFileTitle": "Full API Index", "createTreeView": True, "exhaleExecutesDoxygen": True, "exhaleDoxygenStdin": exhaleDoxygenStdin, "afterTitleDescription": exhaleAfterTitleDescription, } nbsphinx_allow_errors = True cpp_id_attributes = [ "TVM_FFI_DLL", "TVM_FFI_DLL_EXPORT", "TVM_FFI_INLINE", "TVM_FFI_EXTRA_CXX_API", "TVM_FFI_WEAK", "TVM_FFI_COLD_CODE", ] c_id_attributes = [ "TVM_FFI_DLL", "TVM_FFI_DLL_EXPORT", "TVM_FFI_WEAK", ] nbsphinx_execute = "never" autosectionlabel_prefix_document = True nbsphinx_allow_directives = True myst_enable_extensions = [ "dollarmath", "amsmath", "deflist", "colon_fence", "html_image", "linkify", "attrs_block", "substitution", ] myst_heading_anchors = 3 myst_ref_domains = ["std", "py"] myst_all_links_external = False intersphinx_mapping = { "python": ("https://docs.python.org/3.12", None), "typing_extensions": ("https://typing-extensions.readthedocs.io/en/latest", None), "pillow": ("https://pillow.readthedocs.io/en/stable", None), "numpy": ("https://numpy.org/doc/stable", None), "torch": ("https://pytorch.org/docs/2.11", None), "torch-cpp": ("https://docs.pytorch.org/cppdocs", None), "dlpack": ("https://dmlc.github.io/dlpack/latest", None), "data-api": ("https://data-apis.org/array-api/latest", None), "scikit_build_core": ("https://scikit-build-core.readthedocs.io/en/stable/", None), } autosummary_generate = True # actually create stub pages # Map object names -> stub docnames (no ".rst"; relative to your :toctree: dir) autosummary_filename_map = { "tvm_ffi.device": "tvm_ffi.device_function", "tvm_ffi.Device": "tvm_ffi.Device_class", } _STUBS = { "_stubs/cpp_index.rst": "reference/cpp/generated/index.rst", } def _prepare_stub_files() -> None: """Move stub files into place if they do not already exist.""" for src, dst in _STUBS.items(): src_path = _DOCS_DIR / src dst_path = _DOCS_DIR / dst dst_path.parent.mkdir(parents=True, exist_ok=True) if not dst_path.exists(): dst_path.write_text(src_path.read_text(encoding="utf-8"), encoding="utf-8") def _build_rust_docs() -> None: """Build Rust documentation using cargo doc.""" if not build_rust_docs: return print("Building Rust documentation...") try: target_doc = _RUST_DIR / "target" / "doc" # In auto-reload mode (sphinx-autobuild), keep incremental builds # Otherwise (CI/production), do clean rebuild if not is_autobuild and target_doc.exists(): print("Clean rebuild: removing old documentation...") shutil.rmtree(target_doc) # Generate documentation (without dependencies) subprocess.run( ["cargo", "doc", "--no-deps", "--workspace", "--target-dir", "target"], check=True, cwd=_RUST_DIR, env={**os.environ, "RUSTDOCFLAGS": "--cfg docsrs"}, ) print(f"Rust documentation built successfully at {target_doc}") except subprocess.CalledProcessError as e: print(f"Warning: Failed to build Rust documentation: {e}") except FileNotFoundError: print("Warning: cargo not found, skipping Rust documentation build") def _apply_config_overrides(_: object, config: object) -> None: """Apply runtime configuration overrides derived from environment variables.""" config.build_exhale = build_exhale config.build_rust_docs = build_rust_docs def _copy_rust_docs_to_output(app: sphinx.application.Sphinx, exception: Exception | None) -> None: """Copy Rust documentation to the HTML output directory after build completes.""" if exception is not None or not build_rust_docs: return src_dir = _RUST_DIR / "target" / "doc" dst_dir = Path(app.outdir) / "reference" / "rust" / "generated" if src_dir.exists(): if dst_dir.exists(): shutil.rmtree(dst_dir) shutil.copytree(src_dir, dst_dir) print(f"Copied Rust documentation from {src_dir} to {dst_dir}") else: print( f"Warning: Rust documentation source directory not found at {src_dir}. Skipping copy." ) def _mark_exhale_root_orphan( app: sphinx.application.Sphinx, docname: str, source: list[str] ) -> None: """Prepend :orphan: to exhale-generated root so it stays out of the sidebar.""" if docname == "reference/cpp/generated/index": source[0] = ":orphan:\n\n" + source[0] def setup(app: sphinx.application.Sphinx) -> None: """Register custom Sphinx configuration values.""" _prepare_stub_files() _build_rust_docs() app.add_config_value("build_exhale", build_exhale, "env") app.add_config_value("build_rust_docs", build_rust_docs, "env") app.connect("config-inited", _apply_config_overrides) app.connect("source-read", _mark_exhale_root_orphan) app.connect("build-finished", _copy_rust_docs_to_output) app.connect("autodoc-skip-member", _filter_inherited_members) app.connect("autodoc-process-docstring", _link_inherited_members) def _filter_inherited_members(app, what, name, obj, skip, options): # noqa: ANN001, ANN202 if name in _autodoc_always_show: return False if "built-in method " in str(obj): # Skip: `str.maketrans`, `EnumType.from_bytes` return True if getattr(obj, "__objclass__", None) in _py_native_classes: return True return None def _link_inherited_members(app, what, name, obj, options, lines) -> None: # noqa: ANN001 # Only act on members (methods/attributes/properties) if what not in {"method", "attribute", "property"}: return cls = _import_cls(name.rsplit(".", 1)[0]) if cls is None: return member_name = name.rsplit(".", 1)[-1] # just "foo" base = _defining_class(cls, member_name) # If we can't find a base or this class defines it, nothing to do if base is None or base is cls: return # If it comes from builtins we already hide it; no link needed if base in _py_native_classes or getattr(base, "__module__", "") == "builtins": return owner_fq = f"{base.__module__}.{base.__qualname__}".replace("tvm_ffi.core.", "tvm_ffi.") if owner_fq.endswith(".CObject"): owner_fq = owner_fq.removesuffix(".CObject") + ".Object" role = "attr" if what in {"attribute", "property"} else "meth" lines.clear() lines.append( f"*Defined in* :class:`~{owner_fq}` *as {what}* :{role}:`~{owner_fq}.{member_name}`." ) def _defining_class(cls: type | None, attr_name: str) -> type | None: """Find the first class in cls.__mro__ that defines attr_name in its __dict__.""" if not isinstance(cls, type): return None method = getattr(cls, attr_name, None) if method is None: return None for base in reversed(inspect.getmro(cls)): d = getattr(base, "__dict__", {}) if d.get(attr_name, None) is method: return base return None def _import_cls(cls_name: str) -> type | None: """Import and return the class object given its module and class name.""" try: mod, clsname = cls_name.rsplit(".", 1) m = importlib.import_module(mod) return getattr(m, clsname, None) except Exception: return None autodoc_mock_imports = ["torch"] autodoc_default_options = { "members": True, "undoc-members": True, "show-inheritance": True, "inherited-members": False, "member-order": "bysource", } _autodoc_always_show = { "__dlpack__", "__dlpack_device__", "__device_type_name__", "__ffi_init__", "__from_extern_c__", "__from_mlir_packed_safe_call__", "_move", "__move_handle_from__", "__init_handle_by_constructor__", } # If a member method comes from one of these native types, hide it in the docs _py_native_classes: tuple[type, ...] = ( str, tuple, list, dict, set, frozenset, bytes, bytearray, memoryview, int, float, complex, bool, object, ) autodoc_typehints = "description" # or "none" always_use_bars_union = True # Preserve how defaults are written in your source (e.g., DEFAULT_SENTINEL) # Requires Sphinx ≥ 4.0 autodoc_preserve_defaults = True # Ask the extension to include defaults alongside types # 'braces' works well with NumPy-style "Parameters" tables typehints_defaults = "comma" # also accepts: "comma", "braces-after" # Optional: also add stubs for params you didn't list in the docstring always_document_param_types = True # Optional (pairs nicely with NumPy style) napoleon_use_rtype = False # -- Other Options -------------------------------------------------------- templates_path = [] redirects = {} source_suffix = {".rst": "restructuredtext", ".md": "markdown"} language = "en" exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "README.md", "_stubs"] # The name of the Pygments (syntax highlighting) style to use. pygments_style = "sphinx" # A list of ignored prefixes for module index sorting. # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False # -- Options for HTML output ---------------------------------------------- html_theme = "sphinx_book_theme" html_title = project html_copy_source = True html_last_updated_fmt = "" html_favicon = "https://tvm.apache.org/images/logo/tvm-logo-square.png" footer_dropdown = { "name": "ASF", "items": [ ("ASF Homepage", "https://apache.org/"), ("License", "https://www.apache.org/licenses/"), ("Sponsorship", "https://www.apache.org/foundation/sponsorship.html"), ("Security", "https://tvm.apache.org/docs/reference/security.html"), ("Thanks", "https://www.apache.org/foundation/thanks.html"), ("Events", "https://www.apache.org/events/current-event"), ], } footer_copyright = "Copyright © 2025, Apache Software Foundation" footer_note = ( "Apache TVM, Apache, the Apache feather, and the Apache TVM project " + "logo are either trademarks or registered trademarks of the Apache Software Foundation." ) def footer_html() -> str: """Generate HTML for the documentation footer.""" # Create footer HTML with two-line layout # Generate dropdown menu items dropdown_items = "" for item_name, item_url in footer_dropdown["items"]: dropdown_items += f'
  • {item_name}
    • \n' footer_dropdown_html = f"""
    {footer_copyright}
      {dropdown_items}
    {footer_note}
    """ return footer_dropdown_html html_theme_options = { "repository_url": "https://github.com/apache/tvm-ffi", "use_repository_button": True, "show_toc_level": 2, "extra_footer": footer_html(), } html_context = { "display_github": True, "github_user": "apache", "github_version": "main", "conf_py_path": "/docs/", } html_static_path = ["_static"] # Copy Rust documentation to output if enabled html_extra_path = ["reference/rust/generated"] if build_rust_docs else [] html_css_files = ["custom.css"] show_warning_types = True tvm-ffi-0.1.12/docs/dev/000077500000000000000000000000001521067262500146575ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/dev/ci_cd.rst000066400000000000000000000136021521067262500164540ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Reproduce CI/CD =============== This guide explains how to reproduce CI checks and tests locally, and how wheel builds and releases work. All CI/CD workflows are defined under `.github/workflows/ `__. For building the project from source, see :doc:`source_build`. Linters ------- Pre-commit ~~~~~~~~~~ The project uses `pre-commit `__ to run linters and formatters. All hooks are defined in `.pre-commit-config.yaml `__. Install and register the git hooks so they run automatically before each commit: .. code-block:: bash uv tool install pre-commit pre-commit install You can also run hooks manually: .. code-block:: bash # Run all hooks on every file pre-commit run --all-files # Run only on staged files pre-commit run # Run a single hook in isolation pre-commit run ruff-check --all-files pre-commit run clang-format --all-files The main linters per language are: - **Python** -- ``ruff`` (lint + format), ``ty`` (type checking) - **C/C++** -- ``clang-format`` (format), ``clang-tidy`` (lint, see below) - **Cython** -- ``cython-lint`` - **CMake** -- ``cmake-format``, ``cmake-lint`` - **Shell** -- ``shfmt``, ``shellcheck`` If you run into issues with pre-commit: - **Version problems** -- ensure you have pre-commit 2.18.0 or later (``pre-commit --version``). - **Stale cache** -- run ``pre-commit clean`` to clear the hook cache. - **Auto-fixed files** -- most formatting hooks fix issues in place. Review the changes, stage them with ``git add -u``, and commit again. clang-tidy ~~~~~~~~~~ ``clang-tidy`` is run as a separate CI job (not as a pre-commit hook) and only checks C++ files that have changed. To reproduce it locally: .. code-block:: bash # Run clang-tidy on specific files uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=$(nproc) \ include/tvm/ffi/c_api.h src/some_file.cc # Or run on all C++ sources uv run --no-project --with "clang-tidy==21.1.1" \ python tests/lint/clang_tidy_precommit.py \ --build-dir=build-pre-commit \ --jobs=$(nproc) \ ./src/ ./include ./tests .. note:: On macOS, ``clang-tidy`` is resolved through ``xcrun``. The wrapper ``tests/lint/clang_tidy_precommit.py`` handles this automatically. C++ Tests --------- Build and run locally. First, set ``CMAKE_BUILD_PARALLEL_LEVEL`` to speed up the build: .. code-block:: bash export CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) # Linux export CMAKE_BUILD_PARALLEL_LEVEL=$(sysctl -n hw.ncpu) # macOS Then configure, build, and run: .. code-block:: bash # Configure with tests enabled cmake . -B build_test -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Debug # Build the test target cmake --build build_test --clean-first --config Debug --target tvm_ffi_tests # Run tests ctest -V -C Debug --test-dir build_test --output-on-failure .. note:: On Windows, make sure you run the build from a **Developer Command Prompt for VS** or have the MSVC toolchain on your ``PATH``. Python Tests ------------ Reproduce locally with: .. code-block:: bash # Install the project in editable mode with test dependencies uv pip install --reinstall --verbose --group test -e . # Run the full test suite uv run pytest -vvs tests/python Rust Tests ---------- Rust tests live in the ``rust/`` workspace. Run them with: .. code-block:: bash cd rust && cargo test This tests all workspace members (``tvm-ffi``, ``tvm-ffi-sys``, ``tvm-ffi-macros``). .. note:: CI runs Rust tests only after the Python package is installed ( ``uv pip install --group test -e .``), because the Rust FFI bindings link against the built shared library. Make sure the Python package is installed before running ``cargo test``. Build Python Wheels ------------------- CI builds wheels using `cibuildwheel `__ on Linux (x86_64, aarch64), Windows (AMD64), and macOS (arm64). The wheel configuration lives in the ``[tool.cibuildwheel]`` section of ``pyproject.toml``. To build a wheel locally: .. code-block:: bash uv tool install cibuildwheel cibuildwheel --output-dir dist You can restrict the build to a single platform: .. code-block:: bash # Build only for the current platform cibuildwheel --only cp312-macosx_arm64 Use environment variables to control the target platform: .. code-block:: bash # Choose manylinux image (e.g. manylinux2014, manylinux_2_28) CIBW_MANYLINUX_X86_64_IMAGE=manylinux_2_28 cibuildwheel --output-dir dist # Set macOS deployment target CIBW_ENVIRONMENT_MACOS="MACOSX_DEPLOYMENT_TARGET=10.14" cibuildwheel --output-dir dist # Build only specific Python versions CIBW_BUILD="cp312-*" cibuildwheel --output-dir dist .. seealso:: - :doc:`../packaging/python_packaging`: Packaging shared libraries as Python wheels with scikit-build-core. - :doc:`release_process`: Publishing wheels and creating release artifacts. tvm-ffi-0.1.12/docs/dev/doc_build.rst000066400000000000000000000075661521067262500173530ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Build This Doc Site =================== This guide walks through building the TVM FFI documentation locally. Building the docs requires the Python package to be installed first; see :doc:`source_build` for instructions. .. admonition:: Prerequisite :class: hint - `uv `__ manages the Python environment for all docs commands. - Ensure you are in the repository root before running the commands below. - Optional: install ``Doxygen`` if you plan to generate the C++ API reference (see :ref:`build-cpp-docs`). Interactive Build (Auto-Reload) ------------------------------- Rebuilds and serves the documentation locally with live reload: .. code-block:: bash uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/cpp/generated By default, open ``http://127.0.0.1:8000`` in your browser after the initial build completes. One-Off Build ------------- Generates the HTML documentation once, without running a server: .. code-block:: bash uv run --group docs sphinx-build -M html docs docs/_build .. _build-cpp-docs: Build with C++ Docs ------------------- Generating the C++ reference takes longer and requires Doxygen: .. code-block:: bash brew install doxygen # macOS sudo apt install doxygen # Linux Set ``BUILD_CPP_DOCS=1`` on the desired build command to enable the extra step: .. code-block:: bash # Interactive build with auto-rebuild on C++ header changes BUILD_CPP_DOCS=1 uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/cpp/generated --watch include # One-off build BUILD_CPP_DOCS=1 uv run --group docs sphinx-build -M html docs docs/_build Build with Rust Docs -------------------- Generating the Rust reference requires ``cargo`` to be installed: .. code-block:: bash # Install Rust toolchain if not already installed curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh Set ``BUILD_RUST_DOCS=1`` on the desired build command to enable Rust documentation: .. code-block:: bash # Interactive build with auto-rebuild on Rust source changes BUILD_RUST_DOCS=1 uv run --group docs sphinx-autobuild docs docs/_build/html \ --ignore docs/reference/rust/generated --watch rust # One-off build BUILD_RUST_DOCS=1 uv run --group docs sphinx-build -M html docs docs/_build Build All Documentation ----------------------- To build documentation with all language references enabled: .. code-block:: bash # Interactive build BUILD_CPP_DOCS=1 BUILD_RUST_DOCS=1 uv run --group docs sphinx-autobuild \ docs docs/_build/html \ --ignore docs/reference/cpp/generated \ --ignore docs/reference/rust/generated \ --watch include --watch rust # One-off build BUILD_CPP_DOCS=1 BUILD_RUST_DOCS=1 uv run --group docs sphinx-build \ -M html docs docs/_build Cleanup ------- Remove generated artifacts when they are no longer needed: .. code-block:: bash rm -rf docs/_build/ rm -rf docs/reference/python/generated rm -rf docs/reference/cpp/generated rm -rf docs/reference/rust/generated tvm-ffi-0.1.12/docs/dev/release_process.rst000066400000000000000000000160721521067262500205750ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Release Process =============== This guide describes Apache TVM-FFI release workflow for creating a release candidate, staging artifacts on ASF SVN, finalizing the release tag, and publishing release artifacts. .. admonition:: Prerequisite :class: hint The following environment variables need to be set before running the release steps: .. code-block:: bash export FFI_VERSION="v0.1.7-rc0" export FFI_RELEASE_VERSION="v0.1.7" export ASF_USERNAME="your-apache-username" - Tools: ``git``, ``svn``, ``gpg``, ``gtar``, and ``shasum``; - A configured GPG key for signing artifacts. .. important:: This tutorial is based on macOS. Adjust commands accordingly for other operating systems. Pre-release checklist (Step 0) ------------------------------ Before running the steps below, create the release candidate tag and open the release vote. **Step 0.1.** Tag pre-release ``$FFI_VERSION`` on ``__. **Step 0.2.** Open a voting thread: `https://github.com/apache/tvm-ffi/issues/new `__ (`Example `__) Step 1. Create Release Candidate -------------------------------- This builds and signs the source release artifacts under ``tvm-ffi-$FFI_VERSION/``. .. code-block:: bash _make_tarball() { local _ver="$1" local _workdir="tvm-ffi-$_ver-release-files" local _release_dir="tvm-ffi-$_ver" local _tarball="apache-tvm-ffi-$_ver.tar.gz" mkdir -p "$_workdir" "$_release_dir" ( cd "$_workdir" git clone --recursive https://github.com/apache/tvm-ffi.git ffi-release cd ffi-release git checkout "$_ver" rm -rf .DS_Store find . -name ".git*" -print0 | xargs -0 rm -rf cd .. gtar -czvf "$_tarball" -C ffi-release . gpg --armor --output "$_tarball.asc" --detach-sig "$_tarball" shasum -a 512 "$_tarball" > "$_tarball.sha512" ) mv "$_workdir/$_tarball" "$_release_dir/" mv "$_workdir/$_tarball.asc" "$_release_dir/" mv "$_workdir/$_tarball.sha512" "$_release_dir/" rm -rf "$_workdir" } _upload_to_apache_dev_svn() { local _ver="$1" local _asf_username="$2" local _svn_dir="$(pwd)/svn-tvm" local _release_dir="tvm-ffi-$_ver" ( svn co --depth=files "https://dist.apache.org/repos/dist/dev/tvm" $_svn_dir cp -r "${_release_dir}/" "$_svn_dir/$_release_dir" cd "$_svn_dir" svn add "$_release_dir" svn ci --username "$_asf_username" -m "Add TVM-FFI $_ver" ) } _make_tarball "$FFI_VERSION" _upload_to_apache_dev_svn "$FFI_VERSION" "$ASF_USERNAME" Step 2. Conclude Release ------------------------ After the vote passes, retag the release, publish the wheel, bump versions, and trigger the docs release. **Step 2.1.** Conclude voting results: ``__. (`Example `__) **Step 2.2.** Publish PyPI wheel: ``__. (See :doc:`ci_cd` for how wheels are built with cibuildwheel.) **Step 2.3.** Update documentation to latest: ``__. **Step 2.4.** Re-tag the release candidate to the final release version, and bump the version in the source tree: .. code-block:: bash _retag_and_bump_version() { local _ffi_version="$1" local _ffi_release_version="$2" # Configuration variables local _repo_url="git@github.com:apache/tvm-ffi.git" local _work_dir="tvm-ffi-release" local _git_remote="upstream" local _header_file="include/tvm/ffi/c_api.h" # 1. Git clone with remote named "upstream" echo "Cloning repository..." git clone -o "$_git_remote" "$_repo_url" "$_work_dir" cd "$_work_dir" || return 1 git fetch "$_git_remote" --tags # 2. Replace tag and push local _ffi_commit="$(git rev-parse "${_ffi_version}^{commit}")" echo "Creating release tag $_ffi_release_version at commit $_ffi_commit..." git tag -a "$_ffi_release_version" "$_ffi_commit" -m "Release $_ffi_release_version" git push "$_git_remote" "$_ffi_release_version" git push "$_git_remote" --delete "$_ffi_version" # 3. Version bump after the release local _today=$(date +%Y-%m-%d) local _branch_name="${_today}/ver-bump" local _current_patch=$(grep "#define TVM_FFI_VERSION_PATCH" "$_header_file" | awk '{print $3}') local _new_patch=$((_current_patch + 1)) git checkout -b "$_branch_name" echo "Bumping TVM_FFI_VERSION_PATCH from $_current_patch to $_new_patch" sed "s/#define TVM_FFI_VERSION_PATCH $_current_patch/#define TVM_FFI_VERSION_PATCH $_new_patch/" "$_header_file" > "${_header_file}.tmp" && mv "${_header_file}.tmp" "$_header_file" echo "Committing and pushing changes..." git add "$_header_file" git commit -m "chore(release): Version bump after release $_ffi_release_version" git push -u "$_git_remote" "$_branch_name" } _retag_and_bump_version "$FFI_VERSION" "$FFI_RELEASE_VERSION" Step 3. Upload Release Artifacts -------------------------------- After the release is final, copy the RC artifacts into ``dist/release`` with the final version name. .. code-block:: bash _upload_svn() { local _ffi_version="$1" local _ffi_release_version="$2" local _asf_username="$3" local _release_dir="tvm-ffi-$_ffi_version" local _svn_dir="$(pwd)/svn-tvm-release" ( svn co --depth=files "https://dist.apache.org/repos/dist/release/tvm" $_svn_dir mkdir -p "$_svn_dir/tvm-ffi-$_ffi_release_version" cp "${_release_dir}/apache-tvm-ffi-$_ffi_version.tar.gz" "$_svn_dir/tvm-ffi-$_ffi_release_version/apache-tvm-ffi-src-$_ffi_release_version.tar.gz" cp "${_release_dir}/apache-tvm-ffi-$_ffi_version.tar.gz.asc" "$_svn_dir/tvm-ffi-$_ffi_release_version/apache-tvm-ffi-src-$_ffi_release_version.tar.gz.asc" cp "${_release_dir}/apache-tvm-ffi-$_ffi_version.tar.gz.sha512" "$_svn_dir/tvm-ffi-$_ffi_release_version/apache-tvm-ffi-src-$_ffi_release_version.tar.gz.sha512" cd "$_svn_dir" svn add "tvm-ffi-$_ffi_release_version" svn ci --username "$_asf_username" -m "Add TVM FFI $_ffi_release_version" ) } _upload_svn $FFI_VERSION $FFI_RELEASE_VERSION $ASF_USERNAME tvm-ffi-0.1.12/docs/dev/source_build.rst000066400000000000000000000116311521067262500200720ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Build from Source ================= This guide covers two common workflows: - Python package building, which automatically includes building the core C++ library; - C++-only package building without Python. .. admonition:: Prerequisite :class: tip - Python: 3.9 or newer - Compiler: C++17-capable toolchain - Linux: GCC or Clang with C++17 support - macOS: Apple Clang (via Xcode Command Line Tools) - Windows: MSVC (Visual Studio 2019 or 2022, x64) - Build tools: CMake 3.18+; Ninja Build the Python Package ------------------------ Download the source via: .. code-block:: bash git clone --recursive https://github.com/apache/tvm-ffi cd tvm-ffi .. note:: Always clone with ``--recursive`` to pull submodules. If you already cloned without it, run: .. code-block:: bash git submodule update --init --recursive Build the Python package with scikit-build-core, which drives CMake to compile the C++ core and Cython extension: .. code-block:: bash uv pip install --force-reinstall --verbose -e . The ``--force-reinstall`` flag forces a rebuild, and ``-e`` (editable) install means future Python-only changes are reflected immediately without having to rebuild. **CMake flags** can be passed via ``--config-settings cmake.define.=``. For example, to attach debug symbols: .. code-block:: bash uv pip install --force-reinstall --verbose -e . \ --config-settings cmake.define.TVM_FFI_ATTACH_DEBUG_SYMBOLS=ON Available CMake options (see `CMakeLists.txt `__) include: - ``TVM_FFI_ATTACH_DEBUG_SYMBOLS`` -- Attach debug symbols even in release mode (default: ``OFF``). - ``TVM_FFI_USE_LIBBACKTRACE`` -- Enable libbacktrace (default: ``ON``). - ``TVM_FFI_USE_EXTRA_CXX_API`` -- Enable extra C++ API in shared lib (default: ``ON``). - ``TVM_FFI_BACKTRACE_ON_SEGFAULT`` -- Set signal handler to print backtrace on segfault (default: ``ON``). - ``CMAKE_EXPORT_COMPILE_COMMANDS`` -- Generate ``compile_commands.json`` for clangd and other tools (default: ``OFF``). .. warning:: However, changes to C++/Cython always require re-running the install command. Verify the install: .. code-block:: bash uv run python -c "import tvm_ffi; print(tvm_ffi.__version__)" uv run tvm-ffi-config -h .. tip:: Use ``tvm-ffi-config`` to query include and link flags when consuming TVM FFI from external C/C++ projects: .. code-block:: bash tvm-ffi-config --includedir tvm-ffi-config --dlpack-includedir tvm-ffi-config --libfiles # or --libs/--ldflags on Unix Build the C/C++ Library Only ---------------------------- TVM FFI can be used as a standalone C/C++ library without Python. The instruction below should work for Linux, macOS and Windows. .. code-block:: bash cmake . -B build_cpp -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build_cpp --parallel --config RelWithDebInfo --target tvm_ffi_shared cmake --install build_cpp --config RelWithDebInfo --prefix ./dist After installation, you should see: - Headers are installed under ``dist/include/``; - Libraries are installed under ``dist/lib/``. Troubleshooting --------------- - **Rebuilding after C++/Cython changes**: re-run ``uv pip install --force-reinstall -e .``. Editable installs only auto-reflect pure Python changes. - **Submodules missing**: run ``git submodule update --init --recursive`` from the repo root. - **Library not found at import time**: ensure your dynamic loader can find the shared library. If built from source, add the ``lib`` directory to ``LD_LIBRARY_PATH`` (Linux), ``DYLD_LIBRARY_PATH`` (macOS), or ``PATH`` (Windows). - **Wrong generator/build type**: Ninja/Unix Makefiles use ``-DCMAKE_BUILD_TYPE=...``; Visual Studio requires ``--config ...`` at build/ctest time. .. seealso:: - :doc:`ci_cd`: Reproduce linters, unit tests, and wheel builds locally. - :doc:`../get_started/quickstart`: End-to-end walkthrough of building and running a C++/CUDA kernel. - :doc:`../packaging/cpp_tooling`: CMake integration, compiler flags, and library distribution for downstream projects. tvm-ffi-0.1.12/docs/get_started/000077500000000000000000000000001521067262500164065ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/get_started/quickstart.rst000066400000000000000000000331451521067262500213400ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. Quick Start =========== .. note:: All the code in this tutorial is under `examples/quickstart `_ in the repository. This guide walks through shipping a minimal ``add_one`` function that computes ``y = x + 1`` in C++ and CUDA. TVM-FFI's Open ABI and FFI make it possible to **ship one library** for multiple frameworks and languages. We can build a single shared library that works across: - **ML frameworks**, e.g. PyTorch, JAX, PaddlePaddle, NumPy, CuPy, and others; - **Languages**, e.g. C++, Python, Rust, and others; - **Python ABI versions**, e.g. one wheel that supports all Python versions, including free-threaded ones. .. admonition:: Prerequisite :class: hint :name: prerequisite - Python: 3.9 or newer - Compiler: C++17-capable toolchain (GCC/Clang/MSVC) - Optional ML frameworks for testing: NumPy, PyTorch, JAX, CuPy, PaddlePaddle - CUDA: Any modern version (if you want to try the CUDA part) - TVM-FFI installed via: .. code-block:: bash pip install --reinstall --upgrade apache-tvm-ffi Write a Simple ``add_one`` -------------------------- Source Code ~~~~~~~~~~~ Suppose we implement a C++ function ``AddOne`` that performs elementwise ``y = x + 1`` for a 1-D ``float32`` vector. The source code (C++ and CUDA) is: .. hint:: Include the umbrella header to access all the core C++ APIs. .. code-block:: cpp #include .. tabs:: .. group-tab:: C++ .. _cpp_add_one_kernel: .. literalinclude:: ../../examples/quickstart/compile/add_one_cpu.cc :language: cpp :emphasize-lines: 7, 16 :start-after: [example.begin] :end-before: [example.end] .. group-tab:: CUDA .. literalinclude:: ../../examples/quickstart/compile/add_one_cuda.cu :language: cpp :emphasize-lines: 14, 21, 25 :start-after: [example.begin] :end-before: [example.end] The macro :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` exports the C++ function ``AddOne`` as a TVM-FFI-compatible symbol ``__tvm_ffi_add_one_cpu/cuda``. If :c:macro:`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is set to 1, it also exports the function's metadata as a symbol ``__tvm_ffi__metadata_add_one_cpu/cuda`` for type checking and stub generation. The class :cpp:class:`tvm::ffi::TensorView` enables zero-copy interop with tensors from different ML frameworks: - NumPy, CuPy, - PyTorch, JAX, PaddlePaddle, or - any array type that supports the standard :external+data-api:doc:`DLPack protocol `. Finally, :cpp:func:`TVMFFIEnvGetStream` can be used in the CUDA code to launch kernels on the caller's stream. .. seealso:: - :doc:`../guides/export_func_cls`: All three export mechanisms (C symbols, global functions, classes) with complete examples. - :doc:`../guides/kernel_library_guide`: Production-grade CUDA kernel patterns including input validation, device guard, stream handling, and dtype dispatch. - :doc:`../concepts/tensor`: Tensor concepts, :cpp:class:`~tvm::ffi::TensorView` API, and DLPack interop details. .. _sec-cpp-compile-with-tvm-ffi: Compile with TVM-FFI ~~~~~~~~~~~~~~~~~~~~ **Raw command.** Use the following minimal commands to compile the source code: .. tabs:: .. group-tab:: C++ .. literalinclude:: ../../examples/quickstart/raw_compile.sh :language: bash :start-after: [cpp_compile.begin] :end-before: [cpp_compile.end] .. group-tab:: CUDA .. literalinclude:: ../../examples/quickstart/raw_compile.sh :language: bash :start-after: [cuda_compile.begin] :end-before: [cuda_compile.end] These steps produce shared libraries ``add_one_cpu.so`` and ``add_one_cuda.so`` that can be used across languages and frameworks. .. hint:: For a single-file C++/CUDA project, :py:func:`tvm_ffi.cpp.load_inline` minimizes boilerplate for compilation, linking, and loading. **CMake.** CMake is the preferred approach for building across platforms. TVM-FFI integrates with CMake via ``find_package`` as demonstrated below: .. tabs:: .. group-tab:: C++ .. code-block:: cmake # Run `tvm-ffi-config --cmakedir` to set `tvm_ffi_ROOT` find_package(Python COMPONENTS Interpreter REQUIRED) execute_process(COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT) find_package(tvm_ffi CONFIG REQUIRED) # Link C++ target to `tvm_ffi::header` and `tvm_ffi::shared` add_library(add_one_cpu SHARED compile/add_one_cpu.cc) tvm_ffi_configure_target(add_one_cpu) .. group-tab:: CUDA .. code-block:: cmake enable_language(CUDA) # Run `tvm-ffi-config --cmakedir` to set `tvm_ffi_ROOT` find_package(Python COMPONENTS Interpreter REQUIRED) execute_process(COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT) find_package(tvm_ffi CONFIG REQUIRED) # Link CUDA target to `tvm_ffi::header` and `tvm_ffi::shared` add_library(add_one_cuda SHARED compile/add_one_cuda.cu) tvm_ffi_configure_target(add_one_cuda) **Artifact.** The resulting ``add_one_cpu.so`` and ``add_one_cuda.so`` are small libraries that are agnostic to: - Python version/ABI. They are not compiled or linked with Python and depend only on TVM-FFI's stable C ABI; - Languages, including C++, Python, Rust, or any other language that can interop with the C ABI; - ML frameworks, such as PyTorch, JAX, PaddlePaddle, NumPy, CuPy, or any array library that implements the standard :external+data-api:doc:`DLPack protocol `. .. seealso:: :doc:`../packaging/cpp_tooling` for the full build toolchain guide covering CMake integration, raw compiler commands, and cross-platform library distribution. .. _sec-use-across-framework: Ship Across ML Frameworks ------------------------- TVM-FFI's Python package provides :py:func:`tvm_ffi.load_module` to load either ``add_one_cpu.so`` or ``add_one_cuda.so`` into a :py:class:`tvm_ffi.Module`. .. code-block:: python import tvm_ffi mod : tvm_ffi.Module = tvm_ffi.load_module("add_one_cpu.so") func : tvm_ffi.Function = mod.add_one_cpu ``mod.add_one_cpu`` retrieves a callable :py:class:`tvm_ffi.Function` that accepts tensors from host frameworks directly. This is zero-copy, requires no boilerplate code, and adds very little overhead. .. seealso:: :ref:`sec:module` in :doc:`../concepts/func_module` for details on the module system, function retrieval, and the calling convention that enables cross-language interop. We can then use these functions in the following ways: .. _ship-to-pytorch: PyTorch ~~~~~~~ .. literalinclude:: ../../examples/quickstart/load/load_pytorch.py :language: python :start-after: [example.begin] :end-before: [example.end] .. _ship-to-jax: JAX ~~~ Support is provided via `nvidia/jax-tvm-ffi `_. Install it with: .. code-block:: bash pip install jax-tvm-ffi After installation, ``add_one_cuda`` can be registered as a target for JAX's ``ffi_call``. .. code-block:: python # Step 1. Load `build/add_one_cuda.so` import tvm_ffi mod = tvm_ffi.load_module("build/add_one_cuda.so") # Step 2. Register `mod.add_one_cuda` into JAX import jax_tvm_ffi jax_tvm_ffi.register_ffi_target("add_one", mod.add_one_cuda, platform="gpu") # Step 3. Run `mod.add_one_cuda` with JAX import jax import jax.numpy as jnp jax_device, *_ = jax.devices("gpu") x = jnp.array([1, 2, 3, 4, 5], dtype=jnp.float32, device=jax_device) y = jax.ffi.ffi_call( "add_one", # name of the registered function jax.ShapeDtypeStruct(x.shape, x.dtype), # shape and dtype of the output vmap_method="broadcast_all", )(x) print(y) .. _ship-to-paddle: PaddlePaddle ~~~~~~~~~~~~ Since PaddlePaddle 3.3.0, full TVM FFI support is provided. .. literalinclude:: ../../examples/quickstart/load/load_paddle.py :language: python :start-after: [example.begin] :end-before: [example.end] .. _ship-to-numpy: NumPy/CuPy ~~~~~~~~~~ .. literalinclude:: ../../examples/quickstart/load/load_numpy.py :language: python :start-after: [example.begin] :end-before: [example.end] .. literalinclude:: ../../examples/quickstart/load/load_cupy.py :language: python :start-after: [example.begin] :end-before: [example.end] Ship Across Languages --------------------- TVM-FFI's core loading mechanism is ABI-stable and works across language boundaries. A single library can be loaded in any language TVM-FFI supports, without recompiling for different ABIs or languages. .. _ship-to-python: Python ~~~~~~ As shown in the :ref:`previous section`, :py:func:`tvm_ffi.load_module` loads a language- and framework-independent ``add_one_cpu.so`` or ``add_one_cuda.so`` and can be used with any Python array framework that implements the standard :external+data-api:doc:`DLPack protocol `. .. _ship-to-cpp: C++ ~~~ TVM-FFI's C++ API :cpp:func:`tvm::ffi::Module::LoadFromFile` loads ``add_one_cpu.so`` or ``add_one_cuda.so`` and can be used directly from C/C++ without a Python dependency. .. literalinclude:: ../../examples/quickstart/load/load_cpp.cc :language: cpp :start-after: [main.begin] :end-before: [main.end] .. dropdown:: Auxiliary Logics .. literalinclude:: ../../examples/quickstart/load/load_cpp.cc :language: cpp :start-after: [aux.begin] :end-before: [aux.end] Compile and run it with: .. literalinclude:: ../../examples/quickstart/raw_compile.sh :language: bash :start-after: [load_cpp.begin] :end-before: [load_cpp.end] .. note:: Prefer not to load shared libraries? Static linking is also supported. In such cases, use :cpp:func:`tvm::ffi::Function::FromExternC` to create a :cpp:class:`tvm::ffi::Function` from the exported symbol, or directly use :cpp:func:`tvm::ffi::Function::InvokeExternC` to invoke the function. This feature can be useful on iOS, or when the exported module is generated by another DSL compiler targeting the ABI. .. code-block:: cpp // Linked with `add_one_cpu.o` or `add_one_cuda.o` #include // declare reference to the exported symbol extern "C" int __tvm_ffi_add_one_cpu(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); namespace ffi = tvm::ffi; int bundle_add_one(ffi::TensorView x, ffi::TensorView y) { void* closure_handle = nullptr; ffi::Function::InvokeExternC(closure_handle, __tvm_ffi_add_one_cpu, x, y); return 0; } .. _ship-to-rust: Rust ~~~~ TVM-FFI's Rust API ``tvm_ffi::Module::load_from_file`` loads ``add_one_cpu.so`` or ``add_one_cuda.so`` and then retrieves a function ``add_one_cpu`` or ``add_one_cuda`` from it. This mirrors the C++ and Python flows: .. code-block:: rust fn run_add_one(x: &Tensor, y: &Tensor) -> Result<()> { let module = tvm_ffi::Module::load_from_file("add_one_cpu.so")?; let func = module.get_function("add_one_cpu")?; let typed_fn = into_typed_fn!(func, Fn(&Tensor, &Tensor) -> Result<()>); typed_fn(x, y)?; Ok(()) } .. hint:: You can also use the Rust API to target the TVM-FFI ABI. This lets you write the function implementation in Rust and export it to Python/C++ in the same way. Troubleshooting --------------- - ``OSError: cannot open shared object file``: Add an rpath (Linux/macOS) or ensure the DLL is on ``PATH`` (Windows). Example run-path: ``-Wl,-rpath,$(tvm-ffi-config --libdir)``. - ``undefined symbol: __tvm_ffi_add_one_cpu``: Ensure you used :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` and compiled with default symbol visibility (``-fvisibility=hidden`` is fine; the macro ensures export). - ``CUDA error: invalid device function``: Rebuild with the correct ``-arch=sm_XX`` for your GPU, or include multiple ``-gencode`` entries. Further Reading --------------- - :doc:`stable_c_abi`: The stable C ABI layout, calling convention, and end-to-end C examples from both the callee and caller sides. - :doc:`../guides/export_func_cls`: Export C symbols, global functions, and classes across C, C++, and Python. - :doc:`../guides/kernel_library_guide`: Production-grade CUDA kernel patterns (validation, device guard, stream, dtype dispatch). - :doc:`../concepts/func_module`: Calling convention, module system, and global registry concepts. - :doc:`../concepts/tensor`: Tensor representation, DLPack interop, and stream handling. - :doc:`../concepts/abi_overview`: Low-level ABI specification for all TVM-FFI types. - :doc:`../concepts/exception_handling`: Throwing, catching, and propagating exceptions across language boundaries. - :doc:`../packaging/python_packaging`: Packaging extensions as ABI-agnostic Python wheels. - :doc:`../packaging/cpp_tooling`: Build toolchain, CMake integration, and library distribution. - :doc:`../packaging/stubgen`: Generating Python type stubs from C++ metadata. tvm-ffi-0.1.12/docs/get_started/stable_c_abi.rst000066400000000000000000000315051521067262500215330ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. Stable C ABI ============ .. note:: All code used in this guide is under `examples/stable_c_abi `_. .. admonition:: Prerequisite :class: hint - Python: 3.9 or newer (for the ``tvm_ffi.config``/``tvm-ffi-config`` helpers) - Compiler: C11-capable toolchain (GCC/Clang/MSVC) - TVM-FFI installed via .. code-block:: bash pip install --reinstall --upgrade apache-tvm-ffi This guide introduces TVM-FFI's stable C ABI: a single, minimal ABI that represents cross-language calls and is designed for DSL and ML compiler codegen. TVM-FFI is built around the following key idea: .. _tvm_ffi_c_abi: .. admonition:: Key Idea: A Single C ABI for all Functions :class: important Every function call can be represented by a single stable C ABI: .. code-block:: c int tvm_ffi_c_abi( // returns 0 on success; non-zero on failure void* handle, // library handle const TVMFFIAny* args, // inputs: args[0 ... N - 1] int N, // number of inputs TVMFFIAny* result, // output: *result ); where :cpp:class:`TVMFFIAny` is a tagged union of all supported types, e.g. integers, floats, tensors, strings, and more, and can be extended to user-defined types. Built on top of this stable C ABI, TVM-FFI defines a common C ABI protocol for all functions and provides an extensible, performant, and ecosystem-friendly solution. The rest of this guide covers: - The stable C layout and calling convention of ``tvm_ffi_c_abi``; - C examples from both the callee and caller side of this ABI. Stable C Layout --------------- TVM-FFI's :ref:`C ABI ` uses a stable layout for all input and output arguments. Layout of :cpp:class:`TVMFFIAny` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ :cpp:class:`TVMFFIAny` is a fixed-size (128-bit) tagged union that represents all supported types. - First 32 bits: type index indicating which value is stored (supports up to 2^32 types). - Next 32 bits: reserved (used for flags in rare cases, e.g., small-string optimization). - Last 64 bits: payload that is either a 64-bit integer, a 64-bit floating-point number, or a pointer to a heap-allocated object. .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tvm-ffi/stable-c-abi-layout-any.svg :alt: Layout of the 128-bit Any tagged union :name: fig:layout-any Figure 1. Layout spec for the :cpp:class:`TVMFFIAny` tagged union. The following conventions apply when representing values in :cpp:class:`TVMFFIAny`: - Primitive types: the last 64 bits directly store the value, for example: * Integers * Floating-point numbers - Heap-allocated objects: the last 64 bits store a pointer to the actual object, for example: * Managed tensor objects that follow :external+data-api:doc:`DLPack ` (i.e. `DLTensor `_) layout. - Arbitrary objects: the type index identifies the concrete type, and the last 64 bits store a pointer to a reference-counted object in TVM-FFI's object format, for example: * :py:class:`tvm_ffi.Function`, representing all functions, such as Python/C++ functions/lambdas, etc.; * :py:class:`tvm_ffi.Array` and :py:class:`tvm_ffi.Map` (list/dict containers of :cpp:class:`TVMFFIAny` values); * Extending to up to 2^32 types is supported. Function Calling Convention ~~~~~~~~~~~~~~~~~~~~~~~~~~~ Function calls in TVM-FFI share the same calling convention, :ref:`tvm_ffi_c_abi `, as described above. - ``handle: void*``: optional library/closure handle passed to the callee. For exported symbols this is typically ``NULL``; closures may use it to capture context. - ``args: TVMFFIAny*``: pointer to a contiguous array of input arguments. - ``num_args: int``: number of input arguments. - ``result: TVMFFIAny*``: out-parameter that receives the function result (use ``kTVMFFINone`` for "no return value"). .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tvm-ffi/stable-c-abi-layout-func.svg :alt: Layout and calling convention for tvm_ffi_c_abi :name: fig:layout-func Figure 2. Layout and calling convention of :ref:`tvm_ffi_c_abi `, where ``Any`` in this figure refers to :cpp:class:`TVMFFIAny`. Stability and Interoperability ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ **Stability.** The pure C layout and the calling convention are stable across compiler versions and independent of host languages or frameworks. **Cross-language.** TVM-FFI implements this calling convention in multiple languages (C, C++, Python, Rust, ...), enabling code written in one language - or generated by a DSL targeting the ABI - to be called from another language. **Cross-framework.** TVM-FFI uses standard data structures such as :external+data-api:doc:`DLPack tensors ` to represent arrays, so compiled functions can be used from any array framework that implements the DLPack protocol (NumPy, PyTorch, TensorFlow, CuPy, JAX, PaddlePaddle, and others). .. seealso:: - :doc:`../concepts/any`: Full guide on ``Any`` type semantics, ownership, and conversions. - :doc:`../concepts/func_module`: Function and module concepts, including the :ref:`calling convention ` and :ref:`module system `. - :doc:`../concepts/tensor`: Tensor representation and DLPack interop. - :doc:`../concepts/object_and_class`: Object system, type hierarchy, and reference counting. - :doc:`../concepts/abi_overview`: Complete low-level ABI specification for all TVM-FFI types. Stable ABI in C Code -------------------- .. hint:: You can build and run the examples either with raw compiler commands or with CMake. Both approaches are demonstrated below. TVM-FFI's :ref:`C ABI ` is designed with DSL and ML compilers in mind. DSL codegen often targets MLIR, LLVM, or low-level C, where C++ features are unavailable and stable C ABIs are preferred for simplicity and stability. This section shows how to write C code that follows the stable C ABI using two examples: - Callee side: A CPU ``add_one_cpu`` kernel in C that is equivalent to the :ref:`C++ example `. - Caller side: A loader and runner in C that invokes the kernel, a direct C translation of the :ref:`C++ example `. The C code is minimal and dependency-free, so it can serve as a direct reference for DSL compilers that want to expose or invoke kernels through the ABI. Callee: ``add_one_cpu`` Kernel ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Below is a minimal ``add_one_cpu`` kernel in C that follows the stable C ABI in three steps: - **Step 1**. Extract input ``x`` and output ``y`` as DLPack tensors; - **Step 2**. Implement the kernel ``y = x + 1`` on CPU with a simple for-loop; - **Step 3**. Set the output result in ``result``. .. literalinclude:: ../../examples/stable_c_abi/src/add_one_cpu.c :language: c :start-after: [example.begin] :end-before: [example.end] Build it with either approach: .. tabs:: .. group-tab:: Raw command .. literalinclude:: ../../examples/stable_c_abi/raw_compile.sh :language: bash :start-after: [kernel.begin] :end-before: [kernel.end] .. group-tab:: CMake .. code-block:: bash cmake . -B build -DEXAMPLE_NAME="kernel" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo **Compiler codegen.** This C code serves as a direct reference for DSL compilers. To emit a function that follows the stable C ABI, ensure the following: - Symbol naming: define the exported symbol name as ``__tvm_ffi_{func_name}``; - Type checking: check input types via :cpp:member:`TVMFFIAny::type_index`, then marshal inputs from :cpp:class:`TVMFFIAny` to the desired types; - Error handling: return 0 on success, or a non-zero code on failure. When an error occurs, set an error message via :cpp:func:`TVMFFIErrorSetRaisedFromCStr` or :cpp:func:`TVMFFIErrorSetRaisedFromCStrParts`. .. seealso:: - :ref:`sec-exception-handling` for full details on exception propagation across language boundaries. - :doc:`../guides/export_func_cls`: Export C symbols, global functions, and classes from C++ with higher-level macros and reflection helpers. - :doc:`../guides/compiler_integration`: Integrating DSL compilers, graph compilers, and runtime state management with TVM-FFI. **C vs. C++.** Compared to the :ref:`C++ example `, there are a few key differences: - The explicit marshalling in **Step 1** is only needed in C. In C++, templates hide these details. - The C++ macro :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` (used to export ``add_one_cpu``) is not needed in C, since this example directly defines the exported C symbol ``__tvm_ffi_add_one_cpu``. .. hint:: In TVM-FFI's C++ APIs, many invocables (functions, lambdas, functors) are automatically converted into the universal C ABI form by :cpp:class:`tvm::ffi::Function` and :cpp:class:`tvm::ffi::TypedFunction`. Rule of thumb: if an invocable's arguments and result can be converted to/from :cpp:class:`tvm::ffi::Any` (the C++ equivalent of :cpp:class:`TVMFFIAny`), it can be wrapped as a universal C ABI function. Caller: Kernel Loader ~~~~~~~~~~~~~~~~~~~~~ Next, a minimal C loader invokes the ``add_one_cpu`` kernel. It mirrors the :ref:`C++ example ` and performs: - **Step 1**. Load the shared library ``build/add_one_cpu.so`` that contains the kernel; - **Step 2**. Get function ``add_one_cpu`` from the library; - **Step 3**. Invoke the function with two `DLTensor `_ inputs ``x`` and ``y``; .. literalinclude:: ../../examples/stable_c_abi/src/load.c :language: c :start-after: [main.begin] :end-before: [main.end] .. dropdown:: Auxiliary Logics .. literalinclude:: ../../examples/stable_c_abi/src/load.c :language: c :start-after: [aux.begin] :end-before: [aux.end] Build and run the loader with either approach: .. tabs:: .. group-tab:: Raw command .. literalinclude:: ../../examples/stable_c_abi/raw_compile.sh :language: bash :start-after: [load.begin] :end-before: [load.end] .. group-tab:: CMake .. code-block:: bash cmake . -B build -DEXAMPLE_NAME="load" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build/load In C, the idiomatic steps to call a function via the stable C ABI are: - Convert input arguments to the :cpp:class:`TVMFFIAny` type; - Call the target function (e.g., ``add_one_cpu``) via :cpp:func:`TVMFFIFunctionCall`; - Optionally convert the output :cpp:class:`TVMFFIAny` back to the desired type, if the function returns a value. What's Next ----------- **ABI specification.** See the full ABI specification in :doc:`../concepts/abi_overview`. **Convenient compiler target.** The stable C ABI is a simple, portable codegen target for DSL compilers. Emit C that follows this ABI to integrate with TVM-FFI and call the result from multiple languages and frameworks. See :doc:`../concepts/abi_overview`. **Rich and extensible type system.** TVM-FFI supports a rich set of types in the stable C ABI: primitive types (integers, floats), DLPack tensors, strings, built-in reference-counted objects (functions, arrays, maps), and user-defined reference-counted objects. See :doc:`../concepts/object_and_class`. **Export higher-level APIs.** Beyond C symbols, TVM-FFI provides global function registration and class reflection for structured cross-language APIs. See :doc:`../guides/export_func_cls`. **Ship kernels.** For production-grade CUDA kernel patterns including validation, device guard, stream handling, and dtype dispatch, see :doc:`../guides/kernel_library_guide`. **Package and distribute.** Package extensions as Python wheels with type stubs: - :doc:`../packaging/python_packaging`: ABI-agnostic Python wheel builds. - :doc:`../packaging/cpp_tooling`: Build toolchain, CMake integration, and library distribution. - :doc:`../packaging/stubgen`: Generating Python type stubs from C++ metadata. tvm-ffi-0.1.12/docs/guides/000077500000000000000000000000001521067262500153615ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/guides/compiler_integration.md000066400000000000000000000177441521067262500221350ustar00rootroot00000000000000 # Compiler Integration TVM FFI is a standard ABI designed as a standalone module that is independent from compiler or intermediate representation implementations. It specifies a runtime ABI that DSL compilers and languages can integrate with. ## Kernel Language Compilers Kernel languages such as OpenAI Triton, TileLang, Mojo, cuteDSL, Helion, and Hidet usually leverage their own internal compilation mechanisms to build code. To connect these functions to the FFI convention, one can use the following options: - For compilers that generate host functions via codegen (e.g., LLVM), one can generate the symbol `__tvm_ffi_`, where `` is the exported function. Optionally, also generate `__tvm_ffi__metadata_` for reflection. - For kernel generators that generate C++ host code, use {c:macro}`TVM_FFI_DLL_EXPORT_TYPED_FUNC`. This macro automatically exports function metadata when {c:macro}`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is set to 1. - To export documentation strings, use {c:macro}`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC` separately after exporting the function. This enhances tooling support (stub generation, IDE tooltips). Documentation export is also controlled by {c:macro}`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`. The following code snippet shows C code that corresponds to a function performing `add_one_c` under the ABI. It is reasonably straightforward for low-level code generators to replicate this C logic. You can run this code as part of the [quick start example](https://github.com/apache/tvm-ffi/tree/dev/examples/quick_start). ```c #include #include // Helper function to extract DLTensor from TVMFFIAny (can be inlined into generated code) int ReadDLTensorPtr(const TVMFFIAny *value, DLTensor** out) { if (value->type_index == kTVMFFIDLTensorPtr) { *out = (DLTensor*)(value->v_ptr); return 0; } if (value->type_index != kTVMFFITensor) { // Use TVMFFIErrorSetRaisedFromCStr / TVMFFIErrorSetRaisedFromCStrParts to set an error which will // be propagated to the caller TVMFFIErrorSetRaisedFromCStr("ValueError", "Expects a Tensor input"); return -1; } *out = (DLTensor*)((char*)(value->v_obj) + sizeof(TVMFFIObject)); return 0; } // FFI function implementing add_one operation int __tvm_ffi_add_one_c( void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result ) { DLTensor *x, *y; // Extract tensor arguments // return -1 for error, error is set through TVMFFIErrorSetRaisedFromCStr if (ReadDLTensorPtr(&args[0], &x) == -1) return -1; if (ReadDLTensorPtr(&args[1], &y) == -1) return -1; // Get current stream for device synchronization (e.g., CUDA) // not needed for CPU, just keep here for demonstration purpose void* stream = TVMFFIEnvGetStream(x->device.device_type, x->device.device_id); // perform the actual operation for (int i = 0; i < x->shape[0]; ++i) { ((float*)(y->data))[i] = ((float*)(x->data))[i] + 1; } // return 0 for success run return 0; } ``` Some of the key takeaways include: - Prefix the symbol with `__tvm_ffi_` - Call {cpp:func}`TVMFFIEnvGetStream` to get the current environment stream - Use return value for error handling, set error via {cpp:func}`TVMFFIErrorSetRaisedFromCStr` or {cpp:func}`TVMFFIErrorSetRaisedFromCStrParts`. You can also check out the [ABI overview](../concepts/abi_overview.rst) for a more complete guide. ## Graph Compilers Machine learning graph compilers take computational graphs and can integrate with TVM FFI through: - Supporting the `call_tvm_ffi` primitive that calls into `my_func` that follows the ABI: ```python Op.call_tvm_ffi("my_func", *args) ``` - Using the module API to load the modules into context and run. Alternatively, look up global functions that are registered and invoke them. - For ahead-of-time compilation (AOT) with minimum runtime, the AOT compiler can generate direct calls into FFI functions: - Use the TVMFFIFunctionCall API to call into custom {cpp:class}`tvm::ffi::Function`s - If the function exposes a C symbol following the FFI ABI, call it directly. This approach provides a unified mechanism to call into any libraries and other DSLs that expose kernels following the FFI convention, enabling seamless interoperability with various kernel DSLs and libraries. ## Runtime and State Management for Compilers While TVM FFI provides a standard ABI for compiler-generated kernels, many compilers and domain-specific languages (DSLs) require their own **runtime** to manage states like dynamic shapes, workspace memory, or other application-specific data. This runtime can be a separate shared library accessible to all kernels from a specific compiler. ### Recommended Approach for State Management The recommended approach for managing compiler-specific state is to define the state within a **separate shared library**. This library exposes its functionality by registering functions as global `tvm::ffi::Function`s. Here's a breakdown of the process: 1. **Define a Global State**: Create a class or structure to hold your compiler's runtime state. A simple singleton pattern is often used for this. 2. **Register Global Functions**: Use the `TVM_FFI_STATIC_INIT_BLOCK()` macro to register a global function that returns a pointer to your state. For example: ```c++ class GlobalState { ... // your state variables here public: GlobalState* Global() { static auto *inst = new GlobalState(); return inst; } }; TVM_FFI_STATIC_INIT_BLOCK() { using refl = tvm::ffi::reflection; refl.GlobalDef().def("mylang.get_global_state", []()-> void*{ return GlobalState::Global()}); // other runtime APIs can be registered here } ``` This method allows both C++ and Python to access the runtime state through a consistent API. 3. **Access State from Kernels**: Within your compiler-generated kernels, you can use `GetGlobalRequired("mylang.get_global_state")` in C++ or the C equivalent `TVMFFIGetGlobalFunction("mylang.get_global_state", ...)` to get the function and then call it to retrieve the state pointer. ### Distributing the Runtime For a user to use a kernel from your compiler, they must have access to your runtime library. The preferred method is to package the runtime shared library (e.g., `libmylang_runtime.so`) as part of a Python or C++ package. Users must install and import this package before loading any kernels compiled by your system. This approach ensures the state is shared among different kernels. ### Common vs. Custom State It's important to distinguish between compiler-specific state and **common state** managed by TVM FFI. TVM FFI handles common states like **streams** and **memory allocators** through environment functions (e.g., `TVMFFIEnvGetStream`), allowing kernels to access these without managing their own. However, for any unique state required by your compiler, the global function registration approach is the most suitable method. ```{seealso} For creating custom runtime modules that wrap platform-specific driver APIs (e.g., ``cuModuleLoad`` for PTX), see {ref}`sec:custom-modules` in {doc}`../concepts/func_module`. ``` tvm-ffi-0.1.12/docs/guides/cpp_lang_guide.md000066400000000000000000000635031521067262500206520ustar00rootroot00000000000000 # C++ Guide {#cpp-guide} This guide introduces the tvm-ffi C++ API. We provide C++ API on top of the stable C ABI to provide a type-safe and efficient way to work with the tvm-ffi. The C++ API is designed to abstract away the complexity of the C ABI while maintaining full compatibility. The C++ API builds around the following key concepts: - **Any and AnyView**: Type-erased containers that can hold values of any supported type in tvm-ffi. - **Function**: A type-erased "packed" function that can be invoked like normal functions. - **Objects and ObjectRefs**: Reference-counted objects to manage on-heap data types. Code examples in this guide use `EXPECT_EQ` for demonstration purposes, which is a testing framework macro. In actual applications, you would use standard C++ assertions or error handling. You can find runnable code of the examples under tests/cpp/test_example.cc. ## Any and AnyView ```{seealso} For a deep dive into Any including memory layout, ownership semantics, and the type conversion machinery, see {doc}`../concepts/any`. ``` `Any` and `AnyView` are the foundation of tvm-ffi, providing ways to store values that are compatible with the ffi system. The following example shows how we can interact with Any and AnyView. ```cpp #include void ExampleAny() { namespace ffi = tvm::ffi; // Create an Any from various types // EXPECT_EQ is used here for demonstration purposes (testing framework) ffi::Any int_value = 42; ffi::Any float_value = 3.14; ffi::Any string_value = "hello world"; // AnyView provides a lightweight view without ownership ffi::AnyView view = int_value; // we can cast Any/AnyView to a specific type int extracted = view.cast(); EXPECT_EQ(extracted, 42); // If we are not sure about the type // we can use as to get an optional value std::optional maybe_int = view.as(); if (maybe_int.has_value()) { EXPECT_EQ(maybe_int.value(), 42); } // Try cast is another version that will try to run the type // conversion even if the type does not exactly match std::optional maybe_int_try = view.try_cast(); if (maybe_int_try.has_value()) { EXPECT_EQ(maybe_int_try.value(), 42); } } ``` At a high level, we can perform the following operations: - We can store a value into Any, under the hood, Any will record the type of the value by its type_index. - We can fetch a value from Any or AnyView using the `cast` function. - If we are unsure about the type in Any, we can use `as` or `try_cast` function to get an optional value. Under the hood, Any and AnyView store the value via the ABI convention and also manage the reference counting correctly when the stored value is an on-heap object. ## Object and ObjectRef The tvm-ffi object system provides the foundation for all managed, reference-counted objects in the system. It enables type safety, cross-language compatibility, and efficient memory management. The object system is built around three key classes: Object, ObjectPtr, and ObjectRef. The `Object` class is the base class of all heap-allocated objects. It contains a common header that includes the `type_index`, reference counter and deleter for the object. Users do not need to explicitly manage these fields as part of the C++ API. Instead, they are automatically managed through a smart pointer `ObjectPtr` which points to a heap-allocated object instance. The following code shows an example object and the creation of an `ObjectPtr`: ```cpp #include #include class MyIntPairObj : public tvm::ffi::Object { public: int64_t a; int64_t b; MyIntPairObj() = default; MyIntPairObj(int64_t a, int64_t b) : a(a), b(b) {} // Required: declare type information // to register a dynamic type index through the system TVM_FFI_DECLARE_OBJECT_INFO_FINAL("example.MyIntPair", MyIntPairObj, tvm::ffi::Object); }; void ExampleObjectPtr() { namespace ffi = tvm::ffi; // make_object automatically sets up the deleter correctly // This function creates a new ObjectPtr with proper memory management // It handles allocation, initialization, and sets up the reference counting system ffi::ObjectPtr obj = ffi::make_object(100, 200); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(obj->a, 100); EXPECT_EQ(obj->b, 200); } ``` We typically provide a reference class that wraps the ObjectPtr. The `ObjectRef` base class provides the interface and reference counting functionality for these wrapper classes. ```cpp #include #include class MyIntPair : public tvm::ffi::ObjectRef { public: // Constructor explicit MyIntPair(int64_t a, int64_t b) { data_ = tvm::ffi::make_object(a, b); } // Required: define object reference methods // This macro provides the necessary methods for ObjectRef functionality TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MyIntPair, tvm::ffi::ObjectRef, MyIntPairObj); }; void ExampleObjectRef() { namespace ffi = tvm::ffi; MyIntPair pair(100, 200); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(pair->a, 100); EXPECT_EQ(pair->b, 200); } ``` **Note:** The ObjectRef provides a user-friendly interface while ObjectPtr handles the low-level memory management. The ObjectRef acts as a smart pointer wrapper that automatically manages the ObjectPtr lifecycle. The overall implementation pattern is as follows: - **Object Class**: Inherits from `ffi::Object`, stores data and implements the core functionality. - **ObjectPtr**: Smart pointer that manages the Object lifecycle and reference counting. - **Ref Class**: Inherits from `ffi::ObjectRef`, provides a user-friendly interface and automatic memory management. This design ensures efficient memory management while providing a clean API for users. Once we define an ObjectRef class, we can integrate it with the Any, AnyView and Functions. ```cpp #include #include void ExampleObjectRefAny() { namespace ffi = tvm::ffi; MyIntPair pair(100, 200); ffi::Any any = pair; MyIntPair pair2 = any.cast(); // Note: EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(pair2->a, 100); EXPECT_EQ(pair2->b, 200); } ``` Under the hood, ObjectPtr manages the lifecycle of the object through the same mechanism as shared pointers. We designed the object to be intrusive, which means the reference counter and type index metadata are embedded at the header of each object. This design allows us to allocate the control block and object memory together. As we will see in future sections, all of our heap-allocated classes such as Function, on-heap String, Array and Map are managed using subclasses of Object, and the user-facing classes such as Function are ObjectRefs. We provide a collection of built-in object and reference types, which are sufficient for common cases. Developers can also bring new object types as shown in the example of this section. We provide mechanisms to expose these objects to other language bindings such as Python. ## Function The `Function` class provides a type-safe way to create and invoke callable objects through tvm-ffi ABI convention. We can create a `ffi::Function` from an existing typed lambda function. ```cpp #include void ExampleFunctionFromTyped() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fadd1 = ffi::Function::FromTyped( [](const int a) -> int { return a + 1; } ); int b = fadd1(1).cast(); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(b, 2); } ``` Under the hood, tvm-ffi leverages Any and AnyView to create a unified ABI for all functions. The following example demonstrates the low-level way of defining a "packed" function for the same `fadd1`. ```cpp void ExampleFunctionFromPacked() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fadd1 = ffi::Function::FromPacked( [](const ffi::AnyView* args, int32_t num_args, ffi::Any* rv) { // Check that we have exactly one argument TVM_FFI_ICHECK_EQ(num_args, 1); int a = args[0].cast(); *rv = a + 1; } ); int b = fadd1(1).cast(); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(b, 2); } ``` At a high level, `ffi::Function` implements function calling by the following convention: - The arguments are passed through an on-stack array of `ffi::AnyView` - Return values are passed through `ffi::Any` Because the return value is `ffi::Any`, we need to explicitly call `cast` to convert the return value to the desirable type. Importantly, `ffi::Function` itself is a value type that is compatible with tvm-ffi, which means we can pass it as an argument and return values. The following code shows an example of passing a function as an argument and applying it inside. ```cpp void ExampleFunctionPassFunction() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fapply = ffi::Function::FromTyped( [](const ffi::Function f, ffi::Any param) { return f(param.cast()); }); ffi::Function fadd1 = ffi::Function::FromTyped( // [](const int a) -> int { return a + 1; }); int b = fapply(fadd1, 2).cast(); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(b, 3); } ``` This pattern is very powerful because we can construct `ffi::Function` not only from C++, but from any languages that expose to the tvm-ffi ABI. For example, this means we can easily call functions passed in or registered from Python for quick debugging or other purposes. ### Global Function Registry Besides creating functions locally, tvm-ffi provides a global function registry that allows functions to be registered and called across different modules and languages. The following code shows an example ```cpp #include #include void ExampleGlobalFunctionRegistry() { namespace ffi = tvm::ffi; ffi::reflection::GlobalDef().def("xyz.add1", [](const int a) -> int { return a + 1; }); ffi::Function fadd1 = ffi::Function::GetGlobalRequired("xyz.add1"); int b = fadd1(1).cast(); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(b, 2); } ``` You can also access and register global functions from the Python API. ### Exporting as Library Symbol Besides the API that allows registration of functions into the global table, we also provide a macro to export static functions as `TVMFFISafeCallType` symbols in a dynamic library. ```c++ void AddOne(DLTensor* x, DLTensor* y) { // ... implementation omitted ... } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, my_ffi_extension::AddOne); ``` The new `add_one` takes the signature of `TVMFFISafeCallType`, and can loaded and queried through the C++ `ffi::Module` API. When flag `TVM_FFI_DLL_EXPORT_TYPED_FUNC_METADATA` is on, the macro exports both the function and its type metadata, enabling signature validation without calling the function. The metadata contains: - **type_schema**: JSON string describing function signature (return type and argument types) ```cpp ffi::Module mod = ffi::Module::LoadFromFile("path/to/export_lib.so"); // Get the function ffi::Function func = mod->GetFunction("add_one").value(); // Query metadata (type schema information) ffi::Optional metadata = mod->GetFunctionMetadata("add_one"); if (metadata.has_value()) { // Parse JSON metadata for validation // Contains: {"type_schema": "..."} } ``` For functions that need documentation, use the `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC` macro separately: ```cpp #define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1 void ProcessBatch(ffi::TensorView input, ffi::TensorView output) { // ... implementation } // Export the function TVM_FFI_DLL_EXPORT_TYPED_FUNC(process_batch, ProcessBatch); // Export documentation separately (make sure TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is set to 1) TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC( process_batch, R"(Process a batch of inputs and write results to output tensor. Parameters ---------- input : TensorView Input tensor to process output : TensorView Output tensor for results)"); // Query documentation ffi::Optional doc = mod->GetFunctionDoc("process_batch"); ``` ## Error Handling We provide a specific `ffi::Error` type that is also made compatible with the ffi ABI. We also provide a macro `TVM_FFI_THROW` to simplify the error throwing step. ```cpp // file: cpp/test_example.cc #include void FuncThrowError() { namespace ffi = tvm::ffi; TVM_FFI_THROW(TypeError) << "test0"; } void ExampleErrorHandling() { namespace ffi = tvm::ffi; try { FuncThrowError(); } catch (const ffi::Error& e) { EXPECT_EQ(e.kind(), "TypeError"); EXPECT_EQ(e.message(), "test0"); std::cout << e.TracebackMostRecentCallLast() << std::endl; } } ``` The structured error class records kind, message and backtrace that can be mapped to Pythonic style error types and traces. The `TracebackMostRecentCallLast()` call reverses the backtrace and print out follows the Python style, tvm-ffi will try to preserve the backtrace when possible. In the above example, you can see the output as ```text ... more lines omitted File "cpp/test_example.cc", line 106, in ExampleErrorHandling File "cpp/test_example.cc", line 100, in void FuncThrowError() ``` The ffi ABI provides minimal but sufficient mechanisms to propagate these errors across language boundaries. So when we call the function from Python, the Error will be translated into a corresponding Error type. Similarly, when we call a Python callback from C++, the error will be translated into the right error kind and message. ## Tensor For many use cases, we do not need to manage the nd-array/Tensor memory. In such cases, `DLTensor*` can be used as the function arguments. There can be cases for a managed container for multi-dimensional arrays. `ffi::Tensor` is a minimal container to provide such support. Notably, specific logic of device allocations and array operations are non-goals of the FFI. Instead, we provide minimal generic API `ffi::Tensor::FromNDAlloc` to enable flexible customization of Tensor allocation. ```cpp #include #include struct CPUNDAlloc { void AllocData(DLTensor* tensor) { tensor->data = malloc(tvm::ffi::GetDataSize(*tensor)); } void FreeData(DLTensor* tensor) { free(tensor->data); } }; void ExampleTensor() { namespace ffi = tvm::ffi; ffi::Shape shape = {1, 2, 3}; DLDataType dtype = {kDLFloat, 32, 1}; DLDevice device = {kDLCPU, 0}; ffi::Tensor tensor = ffi::Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); // now tensor is a managed tensor } ``` The above example shows how we define `CPUNDAlloc` that customizes `AllocData` and `FreeData` behavior. The CPUNDAlloc struct will be kept alive with the Tensor object. This pattern allows us to implement various Tensor allocations using the same API: - For CUDA allocation, we can change malloc to cudaMalloc - For memory-pool based allocation, we can update `CPUNDAlloc` to keep a strong reference to the pool, so we can keep memory-pool alive when the array is alive. **Working with Shapes** As you may have noticed in the example, we have a `ffi::Shape` container that is used to represent the shapes in nd-array. This container allows us to have compact and efficient representation of managed shapes and we provide quick conversions from standard vector types. ### DLPack Conversion We provide first-class DLPack support to the `ffi::Tensor` that enables efficient exchange through the DLPack Protocol. ```cpp #include void ExampleTensorDLPack() { namespace ffi = tvm::ffi; ffi::Shape shape = {1, 2, 3}; DLDataType dtype = {kDLFloat, 32, 1}; DLDevice device = {kDLCPU, 0}; ffi::Tensor tensor = ffi::Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); // convert to DLManagedTensorVersioned DLManagedTensorVersioned* dlpack = nd.ToDLPackVersioned(); // load back from DLManagedTensorVersioned ffi::Tensor tensor2 = ffi::Tensor::FromDLPackVersioned(dlpack); } ``` These APIs are also available through the C APIs `TVMFFITensorFromDLPackVersioned` and `TVMFFITensorToDLPackVersioned`. ## String and Bytes The tvm-ffi provides first-class support for `String` and `Bytes` types that are efficient, FFI-compatible, and interoperable with standard C++ string types. ```cpp #include void ExampleString() { namespace ffi = tvm::ffi; ffi::String str = "hello world"; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(str.size(), 11); std::string std_str = str; EXPECT_EQ(std_str, "hello world"); } ``` Alternatively, users can always directly use `std::string` in function arguments, conversion will happen automatically. **Rationale:** We need to have separate Bytes and String so they map well to corresponding Python types. `ffi::String` is backed by a possibly managed object that makes it more compatible with the Object system. ## Container Types To enable effective passing and storing of collections of values that are compatible with tvm-ffi, we provide several built-in container types. See [Containers](../concepts/containers.rst) for a conceptual overview. | Type | Header | Mutability | Semantics | | ------ | -------- | ------------ | ----------- | | `Array` | `container/array.h` | Immutable (copy-on-write) | Homogeneous sequence | | `List` | `container/list.h` | Mutable (shared reference) | Homogeneous sequence | | `Tuple` | `container/tuple.h` | Immutable (copy-on-write) | Heterogeneous fixed-size sequence | | `Map` | `container/map.h` | Immutable (copy-on-write) | Homogeneous key-value mapping | | `Dict` | `container/dict.h` | Mutable (shared reference) | Homogeneous key-value mapping | ### Array `Array` provides an array data type that can be used as function arguments. When we use `Array` as an argument of a Function, it will perform runtime checks of the elements to ensure the values match the expected type. ```cpp #include void ExampleArray() { namespace ffi = tvm::ffi; ffi::Array numbers = {1, 2, 3}; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(numbers.size(), 3); EXPECT_EQ(numbers[0], 1); ffi::Function head = ffi::Function::FromTyped([](const ffi::Array a) { return a[0]; }); EXPECT_EQ(head(numbers).cast(), 1); try { // throw an error because 2.2 is not int head(ffi::Array({1, 2.2})); } catch (const ffi::Error& e) { EXPECT_EQ(e.kind(), "TypeError"); } } ``` Under the hood, Array is backed by a reference-counted Object `ArrayObj` that stores a collection of Any values. Note that conversion from Any to `Array` will result in runtime checks of elements because the type index only indicates `ArrayObj` as the backing storage. If you want to defer such checks at the FFI function boundary, consider using `Array` instead. When passing lists and tuples from Python, the values will be converted to `Array` before being passed into the Function. **Performance note:** Repeatedly converting Any to `Array` can incur repeated checking overhead at each element. Consider using `Array` to defer checking or only run conversion once. ### List `List` provides a mutable sequence container with shared reference semantics. Unlike `Array`, mutations happen directly on the underlying shared `ListObj` -- there is no copy-on-write. All handles sharing the same `ListObj` see mutations immediately. ```cpp #include void ExampleList() { namespace ffi = tvm::ffi; ffi::List numbers = {1, 2, 3}; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(numbers.size(), 3); EXPECT_EQ(numbers[0], 1); // Mutate in-place numbers.push_back(4); EXPECT_EQ(numbers.size(), 4); numbers.Set(0, 10); EXPECT_EQ(numbers[0], 10); // Shared reference semantics: both handles see the mutation ffi::List alias = numbers; alias.push_back(5); EXPECT_EQ(numbers.size(), 5); } ``` Under the hood, `List` is backed by a reference-counted `ListObj` that stores a collection of Any values. Like `Array`, conversion from Any to `List` will result in runtime checks of elements. **When to use List vs Array:** Use `Array` when you need an immutable snapshot (e.g., passing a collection through FFI boundaries where the receiver should not mutate). Use `List` when you need to build up or modify a collection in place. ### Tuple `Tuple` provides type-safe fixed-size collections. ```cpp #include void ExampleTuple() { namespace ffi = tvm::ffi; ffi::Tuple tup(42, "hello", true); // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(tup.get<0>(), 42); EXPECT_EQ(tup.get<1>(), "hello"); EXPECT_EQ(tup.get<2>(), true); } ``` Under the hood, Tuple is backed by the same `ArrayObj` as the Array container. This enables zero-cost exchange with input arguments. **Rationale:** This design unifies the conversion rules from Python list/tuple to Array/Tuple. We always need a container representation for tuples to be stored in Any. ### Map `Map` provides a key-value based hashmap container that can accept dict-style parameters. ```cpp #include void ExampleMap() { namespace ffi = tvm::ffi; ffi::Map map0 = {{"Alice", 100}, {"Bob", 95}}; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(map0.size(), 2); EXPECT_EQ(map0.at("Alice"), 100); EXPECT_EQ(map0.count("Alice"), 1); } ``` Under the hood, Map is backed by a reference-counted Object `MapObj` that stores a collection of Any values. The implementation provides a SmallMap variant that stores values as an array and another variant that is based on a hashmap. The Map preserves insertion order like Python dictionaries. Conversion from Any to `Map` will result in runtime checks of its elements because the type index only indicates `MapObj` as the backing storage. If you want to defer such checks at the FFI function boundary, consider using `Map` instead. When passing dictionaries from Python, the values will be converted to `Map` before being passed into the Function. **Performance note:** Repeatedly converting Any to `Map` can incur repeated checking overhead at each element. Consider using `Map` to defer checking or only run conversion once. ### Dict `Dict` provides a mutable key-value mapping with shared reference semantics. Unlike `Map`, mutations happen directly on the underlying shared `DictObj`. All handles sharing the same `DictObj` see mutations immediately. ```cpp #include void ExampleDict() { namespace ffi = tvm::ffi; ffi::Dict scores = {{"Alice", 100}, {"Bob", 95}}; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(scores.size(), 2); EXPECT_EQ(scores.at("Alice"), 100); // Mutate in-place scores.Set("Charlie", 88); EXPECT_EQ(scores.size(), 3); // Shared reference semantics: both handles see the mutation ffi::Dict alias = scores; alias.Set("Dave", 92); EXPECT_EQ(scores.size(), 4); // Erase scores.erase("Bob"); EXPECT_EQ(scores.size(), 3); } ``` Under the hood, `Dict` is backed by a reference-counted `DictObj` that stores a collection of Any key-value pairs. Like `Map`, the `Dict` preserves insertion order. **When to use Dict vs Map:** Use `Map` when you need an immutable snapshot with copy-on-write semantics. Use `Dict` when you need mutable in-place operations. ### Optional `Optional` provides a safe way to handle values that may or may not exist. We specialize Optional for `ffi::String` and Object types to be more compact, using nullptr to indicate non-existence. ```cpp #include void ExampleOptional() { namespace ffi = tvm::ffi; ffi::Optional opt0 = 100; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(opt0.has_value(), true); EXPECT_EQ(opt0.value(), 100); ffi::Optional opt1; EXPECT_EQ(opt1.has_value(), false); EXPECT_EQ(opt1.value_or("default"), "default"); } ``` ### Variant `Variant` provides a type-safe union of different types. ```cpp #include void ExampleVariant() { namespace ffi = tvm::ffi; ffi::Variant var0 = 100; // EXPECT_EQ is used here for demonstration purposes (testing framework) EXPECT_EQ(var0.get(), 100); var0 = ffi::String("hello"); std::optional maybe_str = var0.as(); EXPECT_EQ(maybe_str.value(), "hello"); std::optional maybe_int2 = var0.as(); EXPECT_EQ(maybe_int2.has_value(), false); } ``` Under the hood, Variant is a wrapper around Any that restricts the type to the specific types in the list. tvm-ffi-0.1.12/docs/guides/cubin_launcher.rst000066400000000000000000000447771521067262500211170ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. CUBIN Launcher Guide ==================== This guide demonstrates how to load and launch CUDA kernels from CUBIN (CUDA Binary) modules using TVM-FFI. The CUBIN launcher enables you to execute pre-compiled or runtime-compiled CUDA kernels efficiently through the CUDA Runtime API or Driver API. Overview -------- TVM-FFI provides utilities for loading and launching CUDA kernels from CUBIN modules. The implementation supports both **CUDA Runtime API** (default for CUDA >= 12.8) and **CUDA Driver API**. **Runtime API (CUDA >= 12.8):** - ``cudaLibraryLoadData()`` - Load CUBIN from memory buffer - ``cudaLibraryGetKernel()`` - Get kernel handle by name - ``cudaLaunchKernel()`` - Launch kernel with grid/block dimensions **Driver API:** - ``cuLibraryLoadData()`` - Load CUBIN from memory buffer - ``cuLibraryGetKernel()`` - Get kernel handle by name - ``cuLaunchKernel()`` - Launch kernel with grid/block dimensions **Customization:** By default, the implementation uses the Runtime API if compiled with CUDA >= 12.8, falling back to the Driver API for older versions. You can force the usage of the Driver API (or Runtime API) by defining the macro ``TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API`` (set to ``1`` for Driver API, ``0`` for Runtime API) before including the header. .. warning:: **CMAKE_CUDA_RUNTIME_LIBRARY and Driver API** When using CMake, the default behavior (if ``CMAKE_CUDA_RUNTIME_LIBRARY`` is not set) is to link against the CUDA Runtime Library (``cudart``). TVM-FFI's CMake utility automatically defaults this variable to ``Shared`` if it is undefined. This introduces a dependency on the CUDA runtime version, requiring the system's driver to be compatible with that runtime version. If you intend to use the Driver API only (e.g. by setting ``TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1``) to avoid this runtime dependency: 1. You must explicitly set ``CMAKE_CUDA_RUNTIME_LIBRARY`` to ``None`` in your CMake configuration to prevent linking ``cudart``. 2. You must manually link your target against the CUDA Driver library (usually ``cuda`` on Linux/Windows or `CUDA::cuda_driver` provided by CMake's ``FindCUDAToolkit``). This ensures your application relies solely on the widely compatible CUDA Driver API (``libcuda.so.1``). The implementation is in ``tvm/ffi/extra/cuda/cubin_launcher.h`` and provides: - :cpp:class:`tvm::ffi::CubinModule`: RAII wrapper for loading CUBIN modules from memory - :cpp:class:`tvm::ffi::CubinKernel`: Handle for launching CUDA kernels with specified parameters - :c:macro:`TVM_FFI_EMBED_CUBIN`: Macro for embedding CUBIN data at compile time (legacy / object-linking approach) - :c:macro:`TVM_FFI_EMBED_CUBIN_FROM_BYTES`: Macro for embedding CUBIN data from byte arrays (manual embedding approach) - :c:macro:`TVM_FFI_EMBED_CUBIN_GET_KERNEL`: Macro for retrieving kernels from embedded CUBIN The CUBIN launcher supports: - Loading CUBIN from memory (embedded data or runtime-generated) - Multi-GPU execution using CUDA primary contexts - Kernel parameter management and launch configuration - Integration with NVRTC, Triton, and other CUDA compilation tools **Build Integration:** TVM-FFI provides convenient tools for embedding CUBIN data at build time: - **CMake utilities** (``cmake/Utils/EmbedCubin.cmake``): Functions for compiling CUDA to CUBIN/FATBIN and embedding it into C++ code or linking it. - **Python utility** (``python -m tvm_ffi.utils.embed_cubin``): Command-line tool for embedding CUBIN into object files. - **Python API** (:py:func:`tvm_ffi.cpp.load_inline`): Runtime embedding via ``embed_cubin`` parameter. Python Usage ------------ Basic Workflow ~~~~~~~~~~~~~~ The typical workflow for launching CUBIN kernels from Python involves: 1. **Generate CUBIN**: Compile your CUDA kernel to CUBIN format 2. **Define C++ Wrapper**: Write C++ code to load and launch the kernel 3. **Load Module**: Use :py:func:`tvm_ffi.cpp.load_inline` with ``embed_cubin`` parameter 4. **Call Kernel**: Invoke the kernel function from Python Example: NVRTC Compilation ~~~~~~~~~~~~~~~~~~~~~~~~~~~ Here's a complete example using NVRTC to compile CUDA source at runtime. **Step 1: Compile CUDA source to CUBIN using NVRTC** .. literalinclude:: ../../examples/cubin_launcher/example_nvrtc_cubin.py :language: python :start-after: [cuda_source.begin] :end-before: [cuda_source.end] :dedent: 4 **Step 2: Define C++ wrapper with embedded CUBIN** .. literalinclude:: ../../examples/cubin_launcher/example_nvrtc_cubin.py :language: python :start-after: [cpp_wrapper.begin] :end-before: [cpp_wrapper.end] :dedent: 4 **Key Points:** - The ``embed_cubin`` parameter is a dictionary mapping CUBIN names to their binary data - CUBIN names in ``embed_cubin`` must match names in :c:macro:`TVM_FFI_EMBED_CUBIN` - Use ``cuda_sources`` parameter (instead of ``cpp_sources``) to automatically link with CUDA libraries - The C++ wrapper handles device management, stream handling, and kernel launching Example: Using Triton Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ You can compile Triton kernels to CUBIN and launch them through TVM-FFI. **Step 1: Define and compile Triton kernel** .. literalinclude:: ../../examples/cubin_launcher/example_triton_cubin.py :language: python :start-after: [triton_kernel.begin] :end-before: [triton_kernel.end] :dedent: 4 **Step 2: Define C++ wrapper to launch the Triton kernel** .. literalinclude:: ../../examples/cubin_launcher/example_triton_cubin.py :language: python :start-after: [cpp_wrapper.begin] :end-before: [cpp_wrapper.end] :dedent: 4 .. note:: Triton kernels may require extra dummy parameters in the argument list. Check the compiled kernel's signature to determine the exact parameter count needed. C++ Usage --------- Embedding CUBIN at Compile Time ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The most convenient way to embed CUBIN/FATBIN data in C++ is using the TVM-FFI build utilities. There are three main approaches: 1. **Object Linking (Standard)**: Use CMake utilities to compile and link the CUBIN data. 2. **Header Inclusion (Portable)**: Convert CUBIN to a C header file using ``bin2c``. 3. **C++ Embedding (Modern)**: Use C++23 ``#embed`` (or compiler extensions). **Method 1: Object Linking (Standard)** This approach uses CMake utilities to compile and link the CUBIN data. It works across all supported compilers and handles the low-level details of object file generation and symbol naming. .. literalinclude:: ../../examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/src/lib_embedded.cc :language: cpp :start-after: [example.begin] :end-before: [example.end] **Method 2: Header Inclusion (Portable)** You can use tools like ``bin2c`` to generate a header file containing the byte array and include it. .. literalinclude:: ../../examples/cubin_launcher/embedded_cubin/include_bin2c/src/lib_embedded.cc :language: cpp :start-after: [example.begin] :end-before: [example.end] **Method 3: C++ Embedding (Modern)** Using C++23 ``#embed`` (or compiler extensions like ``#embed`` in Clang/GCC) allows you to include the binary data directly. .. literalinclude:: ../../examples/cubin_launcher/embedded_cubin/cpp_embed/src/lib_embedded.cc :language: cpp :start-after: [example.begin] :end-before: [example.end] **Key Points:** - Use ``static auto kernel`` to cache the kernel lookup for efficiency - Kernel arguments must be pointers to the actual values (use ``&`` for addresses) - :cpp:type:`tvm::ffi::dim3` supports 1D, 2D, or 3D configurations: ``dim3(x)``, ``dim3(x, y)``, ``dim3(x, y, z)`` - ``TVMFFIEnvGetStream`` retrieves the correct CUDA stream for the device - Always check kernel launch results with :c:macro:`TVM_FFI_CHECK_CUDA_ERROR` (which checks CUDA Runtime API or Driver API errors depending on configuration) Loading CUBIN at Runtime ~~~~~~~~~~~~~~~~~~~~~~~~~ You can also load CUBIN modules dynamically from memory: .. literalinclude:: ../../examples/cubin_launcher/dynamic_cubin/src/lib_dynamic.cc :language: cpp :start-after: [example.begin] :end-before: [example.end] Embedding CUBIN with CMake Utilities ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ TVM-FFI provides CMake utility functions that simplify the CUBIN embedding process. This is the recommended approach for CMake-based projects. **Using CMake Utilities:** .. literalinclude:: ../../examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/CMakeLists.txt :language: cmake :start-after: [cmake_example.begin] :end-before: [cmake_example.end] **Available CMake Functions:** - ``add_tvm_ffi_cubin( CUDA )``: Creates an object library that compiles CUDA source to CUBIN format. This is a compatibility wrapper; for CMake >= 3.27, you can use standard ``CUDA_CUBIN_COMPILATION`` property. - ``add_tvm_ffi_fatbin( CUDA )``: Creates an object library that compiles CUDA source to FATBIN format. This is a compatibility wrapper; for CMake >= 3.27, you can use standard ``CUDA_FATBIN_COMPILATION`` property. - ``tvm_ffi_embed_bin_into( SYMBOL BIN )``: Embeds a CUBIN/FATBIN file into an existing object library target. This works by linking the binary data into the target, allowing access via ``TVM_FFI_EMBED_CUBIN()``. - ``target``: The target to embed into (must be an object library or have object files). - ``symbol``: Symbol name to use (must match ``TVM_FFI_EMBED_CUBIN(symbol)``). - ``BIN``: Path to the CUBIN/FATBIN file (e.g., from ``$``). .. note:: When including ``cmake/Utils/EmbedCubin.cmake``, if ``CMAKE_CUDA_RUNTIME_LIBRARY`` is not set, it defaults to ``Shared``. This prevents static linking of cudart, which requires an exact driver version match. If you intend to use the Driver API only (e.g., via ``TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1``), you should explicitly set ``CMAKE_CUDA_RUNTIME_LIBRARY`` to ``None`` in your CMake configuration before including this utility to avoid linking against the CUDA runtime library. And link with CUDA Driver API. Embedding CUBIN with Python Utility ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ For more advanced use cases or non-CMake build systems, you can use the Python command-line utility to embed CUBIN data into existing object files. **Command-Line Usage:** .. code-block:: bash # Step 1: Compile C++ source to object file g++ -c -fPIC -std=c++17 -I/path/to/tvm-ffi/include mycode.cc -o mycode.o # Step 2: Embed CUBIN into the object file python -m tvm_ffi.utils.embed_cubin \ --output-obj mycode_with_cubin.o \ --input-obj mycode.o \ --cubin kernel.cubin \ --name my_kernels # Step 3: Link into final library g++ -o mylib.so -shared mycode_with_cubin.o -lcudart **Python API:** .. code-block:: python from pathlib import Path from tvm_ffi.utils.embed_cubin import embed_cubin embed_cubin( cubin_path=Path("kernel.cubin"), input_obj_path=Path("mycode.o"), output_obj_path=Path("mycode_with_cubin.o"), name="my_kernels", verbose=True # Optional: print detailed progress ) The Python utility performs these steps: 1. Creates intermediate CUBIN object file using ``ld -r -b binary`` 2. Adds ``.note.GNU-stack`` section for security 3. Renames symbols to match TVM-FFI format (``__tvm_ffi__cubin_``) 4. Merges with input object file using relocatable linking 5. Localizes symbols to prevent conflicts when multiple object files use the same name Manual CUBIN Embedding ~~~~~~~~~~~~~~~~~~~~~~ For reference, here's how to manually embed CUBIN using objcopy and ld: **Step 1: Compile CUDA kernel to CUBIN** .. code-block:: bash nvcc --cubin -arch=sm_75 kernel.cu -o kernel.cubin **Step 2: Convert CUBIN to object file** .. code-block:: bash ld -r -b binary -o kernel_data.o kernel.cubin **Step 3: Rename symbols with objcopy** .. code-block:: bash objcopy --rename-section .data=.rodata,alloc,load,readonly,data,contents \ --redefine-sym _binary_kernel_cubin_start=__tvm_ffi__cubin_my_kernels \ --redefine-sym _binary_kernel_cubin_end=__tvm_ffi__cubin_my_kernels_end \ kernel_data.o **Step 4: Link with your library** .. code-block:: bash g++ -o mylib.so -shared mycode.cc kernel_data.o -Wl,-z,noexecstack -lcudart The symbol names must match the name used in :c:macro:`TVM_FFI_EMBED_CUBIN`. **When to Use Each Approach:** - **CMake utilities**: Best for CMake-based projects, provides cleanest integration (recommended) - **Python utility**: Best for custom build systems, Makefile-based projects, or advanced workflows (recommended) - **Manual objcopy**: Low-level approach, useful for understanding the process or debugging (only for customized use cases) Advanced Topics --------------- Multi-GPU Support ~~~~~~~~~~~~~~~~~ The CUBIN launcher automatically handles multi-GPU execution through CUDA primary contexts. Kernels will execute on the device associated with the input tensors: .. code-block:: cpp void MultiGPUExample(tvm::ffi::TensorView x_gpu0, tvm::ffi::TensorView x_gpu1) { static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "process"); // Launch on GPU 0 (device determined by x_gpu0.device()) LaunchOnDevice(kernel, x_gpu0); // Launch on GPU 1 (device determined by x_gpu1.device()) LaunchOnDevice(kernel, x_gpu1); } The :cpp:class:`tvm::ffi::CubinKernel` automatically uses the device context from the input tensors. Kernel Launch Configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ When writing the C++ wrapper, important considerations include: - **Grid/Block Dimensions**: Use :cpp:type:`tvm::ffi::dim3` for 1D, 2D, or 3D configurations - 1D: ``dim3(x)`` → ``(x, 1, 1)`` - 2D: ``dim3(x, y)`` → ``(x, y, 1)`` - 3D: ``dim3(x, y, z)`` → ``(x, y, z)`` - **Kernel Arguments**: Must be pointers to actual values - For device pointers: ``void* ptr = tensor.data_ptr(); args[] = {&ptr}`` - For scalars: ``int n = 42; args[] = {&n}`` - **Stream Management**: Use ``TVMFFIEnvGetStream`` to get the correct CUDA stream for synchronization with DLPack tensors - **Error Checking**: Always use :c:macro:`TVM_FFI_CHECK_CUDA_ERROR` to validate CUDA Runtime API results Dynamic Shared Memory ~~~~~~~~~~~~~~~~~~~~~ To use dynamic shared memory, specify the size in the :cpp:func:`tvm::ffi::CubinKernel::Launch` call: .. code-block:: cpp // Allocate 1KB of dynamic shared memory uint32_t shared_mem_bytes = 1024; TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(args, grid, block, stream, shared_mem_bytes)); Integration with Different Compilers ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The CUBIN launcher works with various CUDA compilation tools: - **NVCC**: Standard NVIDIA compiler, produces highly optimized CUBIN - **NVRTC**: Runtime compilation for JIT scenarios (via :py:mod:`tvm_ffi.cpp.nvrtc`) - **Triton**: High-level DSL that compiles to CUBIN - **Custom compilers**: Any tool that generates valid CUDA CUBIN Complete Examples ----------------- For complete working examples, see the ``examples/cubin_launcher/`` directory: - ``embedded_cubin/`` - Pre-compiled CUBIN embedded at build time - ``dynamic_cubin/`` - CUBIN data passed dynamically at runtime - ``example_nvrtc_cubin.py`` - NVRTC runtime compilation - ``example_triton_cubin.py`` - Triton kernel compilation These examples demonstrate: - Compiling CUDA kernels to CUBIN - Embedding CUBIN in C++ modules - Launching kernels with proper error handling - Testing and verification API Reference ------------- C++ Classes ~~~~~~~~~~~ - :cpp:class:`tvm::ffi::CubinModule`: RAII wrapper for CUBIN module lifecycle - :cpp:func:`tvm::ffi::CubinModule::CubinModule`: Load CUBIN from memory - :cpp:func:`tvm::ffi::CubinModule::GetKernel`: Get kernel by name - :cpp:func:`tvm::ffi::CubinModule::GetKernelWithMaxDynamicSharedMemory`: Get kernel by name with maximum dynamic shared memory set - :cpp:func:`tvm::ffi::CubinModule::operator[]`: Convenient kernel access - :cpp:class:`tvm::ffi::CubinKernel`: Handle for launching kernels - :cpp:func:`tvm::ffi::CubinKernel::Launch`: Launch kernel with specified parameters - :cpp:type:`tvm::ffi::dim3`: 3D dimension structure - ``dim3()``: Default (1, 1, 1) - ``dim3(unsigned int x)``: 1D - ``dim3(unsigned int x, unsigned int y)``: 2D - ``dim3(unsigned int x, unsigned int y, unsigned int z)``: 3D C++ Macros ~~~~~~~~~~ - :c:macro:`TVM_FFI_EMBED_CUBIN`: Declare embedded CUBIN module - :c:macro:`TVM_FFI_EMBED_CUBIN_FROM_BYTES`: Load CUBIN from byte array - :c:macro:`TVM_FFI_EMBED_CUBIN_GET_KERNEL`: Get kernel from embedded module - :c:macro:`TVM_FFI_CHECK_CUDA_ERROR`: Check CUDA Runtime API result Python Functions ~~~~~~~~~~~~~~~~ - :py:func:`tvm_ffi.cpp.nvrtc.nvrtc_compile`: Compile CUDA source to CUBIN - :py:func:`tvm_ffi.cpp.load_inline`: Load inline module with embedded CUBIN Python Utilities ~~~~~~~~~~~~~~~~ - ``python -m tvm_ffi.utils.embed_cubin``: Command-line utility to embed CUBIN into object files - ``--output-obj PATH``: Output combined object file path - ``--input-obj PATH``: Input compiled object file containing C++ code with ``TVM_FFI_EMBED_CUBIN`` - ``--cubin PATH``: Input CUBIN file to embed - ``--name NAME``: Symbol name matching ``TVM_FFI_EMBED_CUBIN(name)`` macro - ``--verbose``: Print detailed command output (optional) - :py:func:`tvm_ffi.utils.embed_cubin.embed_cubin`: Python API for embedding CUBIN - ``cubin_path``: Path to input CUBIN file - ``input_obj_path``: Path to existing object file - ``output_obj_path``: Path to output combined object file - ``name``: Symbol name for the embedded CUBIN - ``verbose``: Enable detailed output (default: False) CMake Functions ~~~~~~~~~~~~~~~ - ``add_tvm_ffi_cubin( CUDA )``: Compile CUDA source to CUBIN - ``add_tvm_ffi_fatbin( CUDA )``: Compile CUDA source to FATBIN - ``tvm_ffi_embed_bin_into( SYMBOL BIN )``: Embed CUBIN/FATBIN into object target tvm-ffi-0.1.12/docs/guides/dataclass_reflection.rst000066400000000000000000000336641521067262500223000ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. .. _dataclass-reflection: Dataclass-Style Reflection ========================== TVM-FFI's reflection system provides Python-dataclass-style features for C++ classes: auto-generated constructors, default values, keyword-only parameters, repr, hashing, comparison, and deep copy. These features are enabled by per-field and per-class traits registered via :cpp:class:`~tvm::ffi::reflection::ObjectDef`. This guide assumes familiarity with :doc:`export_func_cls` and :doc:`../concepts/object_and_class`. Quick Start ----------- Define a C++ object with fields, register it with traits, and use it from Python with full dataclass semantics: .. code-block:: cpp #include namespace ffi = tvm::ffi; class PointObj : public ffi::Object { public: int64_t x; int64_t y; ffi::String label; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my_ext.Point", PointObj, ffi::Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; refl::ObjectDef() .def_rw("x", &PointObj::x) .def_rw("y", &PointObj::y) .def_rw("label", &PointObj::label, refl::default_("")); } No ``refl::init<>()`` call is needed — the reflection system auto-generates a packed ``__ffi_init__`` from the reflected fields: .. code-block:: python import tvm_ffi @tvm_ffi.register_object("my_ext.Point") class Point(tvm_ffi.Object): x: int y: int label: str p1 = Point(1, 2) # positional args p2 = Point(1, 2, label="origin") # keyword arg with default .. _auto-init: Auto-Generated Constructors ---------------------------- When no explicit ``refl::init()`` is registered, ``ObjectDef`` auto-generates a packed constructor (``__ffi_init__``) from the reflected fields. The generated signature follows Python conventions: 1. **Required positional** parameters come first. 2. **Optional positional** parameters (those with defaults) come next. 3. **Keyword-only** parameters follow after a ``*`` separator. Field Traits for Init ~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 30 70 * - Trait - Effect on auto-init * - ``refl::default_(value)`` - Makes the parameter optional with a literal default value. * - ``refl::default_factory(fn)`` - Makes the parameter optional; calls ``fn()`` each time a default is needed. Use for mutable defaults (e.g. ``Array``, ``Dict``). * - ``refl::kw_only(true)`` - Moves the parameter after the ``*`` separator (keyword-only). * - ``refl::init(false)`` - Excludes the field from the constructor entirely. The field must have a default value or be initialized by a base-class constructor. .. code-block:: cpp refl::ObjectDef() .def_rw("batch_size", &ConfigObj::batch_size) .def_rw("lr", &ConfigObj::lr, refl::default_(0.001)) .def_rw("device", &ConfigObj::device, refl::kw_only(true), refl::default_("cpu")) .def_rw("_cache", &ConfigObj::_cache, refl::init(false), refl::default_factory([] { return ffi::Dict(); })); The generated Python signature is: .. code-block:: python def __init__(self, batch_size, lr=0.001, *, device="cpu"): ... # _cache is excluded from __init__ but initialized to Dict() Suppressing Auto-Init ~~~~~~~~~~~~~~~~~~~~~ Pass ``refl::init(false)`` at the class level to suppress auto-init entirely: .. code-block:: cpp refl::ObjectDef(refl::init(false)) .def_rw("x", &InternalObj::x) .def_rw("y", &InternalObj::y); The object will have no ``__ffi_init__`` method. Construction must happen through a custom factory or from C++. Explicit Constructors ~~~~~~~~~~~~~~~~~~~~~ Use ``refl::init()`` to register an explicit typed constructor instead of auto-init: .. code-block:: cpp refl::ObjectDef() .def(refl::init()) .def_ro("a", &IntPairObj::a) .def_ro("b", &IntPairObj::b); This calls ``IntPairObj(int64_t, int64_t)`` directly. Auto-init is automatically suppressed when an explicit constructor is registered. Default Values -------------- Literal Defaults ~~~~~~~~~~~~~~~~ ``refl::default_(value)`` stores a literal default. The value is captured once at registration time: .. code-block:: cpp .def_rw("threshold", &Obj::threshold, refl::default_(0.5)) Factory Defaults ~~~~~~~~~~~~~~~~ ``refl::default_factory(fn)`` calls ``fn()`` each time a default is needed. Use this for mutable containers to avoid aliasing: .. code-block:: cpp .def_rw("items", &Obj::items, refl::default_factory([] { return ffi::List(); })) .. note:: ``refl::default_`` and ``refl::default_factory`` are the preferred names for new code. The original names ``refl::DefaultValue`` and ``refl::DefaultFactory`` are retained for backward compatibility. .. _field-traits-ref: Field Traits Reference ---------------------- All traits are passed as extra arguments to ``def_ro`` or ``def_rw``. Multiple traits can be combined on a single field: .. code-block:: cpp .def_rw("name", &Obj::name, refl::default_("unnamed"), refl::kw_only(true), refl::repr(false)) .. list-table:: :header-rows: 1 :widths: 30 70 * - Trait - Description * - ``refl::init(bool)`` - Include/exclude field from auto-generated ``__ffi_init__``. * - ``refl::kw_only(bool)`` - Mark field as keyword-only in auto-init. * - ``refl::default_(value)`` - Literal default value for the field. * - ``refl::default_factory(fn)`` - Factory function ``() -> Any`` for mutable defaults. * - ``refl::repr(bool)`` - Include/exclude field from repr output. * - ``refl::hash(bool)`` - Include/exclude field from recursive hashing. * - ``refl::compare(bool)`` - Include/exclude field from recursive comparison. * - ``refl::Metadata({...})`` - Attach custom key-value metadata to the field. .. _dataclass-operations: Dataclass Operations -------------------- Once a class is registered with ``ObjectDef``, several dataclass operations are available automatically. These are defined in ``include/tvm/ffi/extra/dataclass.h``: .. code-block:: cpp #include ffi::Any value = ...; ffi::Any copy = ffi::DeepCopy(value); ffi::String repr = ffi::ReprPrint(value); int64_t h = ffi::RecursiveHash(value); bool eq = ffi::RecursiveEq(a, b); All operations use iterative DFS with an explicit stack (no recursion), so they are safe for deep object graphs. Deep Copy ~~~~~~~~~ ``DeepCopy(value)`` recursively copies an object and all reachable objects in its graph. Objects that register a copy constructor via ``ObjectDef`` automatically support deep copy (through the ``__ffi_shallow_copy__`` type attribute). - **Immutable leaves** (returned as-is): primitives, ``String``, ``Bytes``, ``Shape`` - **Recursively copied**: ``Array``, ``List``, ``Map``, ``Dict``, and reflected objects Repr ~~~~ ``ReprPrint(value)`` produces a human-readable string representation using field names and values from reflection metadata. It handles cycles (prints ``...``) and DAG structures (caches repr strings). Exclude a field from repr output with ``refl::repr(false)``: .. code-block:: cpp .def_ro("internal_state", &Obj::internal_state, refl::repr(false)) Hashing ~~~~~~~ ``RecursiveHash(value)`` computes a deterministic recursive hash. The hash is consistent with equality: if ``RecursiveEq(a, b)`` then ``RecursiveHash(a) == RecursiveHash(b)``. Exclude a field from hashing with ``refl::hash(false)``: .. code-block:: cpp .def_ro("cache_key", &Obj::cache_key, refl::hash(false)) Comparison ~~~~~~~~~~ ``RecursiveEq(a, b)`` tests structural equality. Ordering comparisons (``RecursiveLt``, ``RecursiveLe``, ``RecursiveGt``, ``RecursiveGe``) provide lexicographic field-by-field ordering. Exclude a field from comparison with ``refl::compare(false)``: .. code-block:: cpp .def_ro("timestamp", &Obj::timestamp, refl::compare(false)) In Python, these are wired up as ``__eq__``, ``__lt__``, ``__le__``, ``__gt__``, ``__ge__``, and ``__hash__`` on classes created with ``c_class``. Custom Hooks ------------ Override the default behavior for repr, hash, or comparison by registering type-level attributes via :cpp:class:`~tvm::ffi::reflection::TypeAttrDef`: .. code-block:: cpp namespace refl = ffi::reflection; // Custom hash: only hash the "key" field refl::TypeAttrDef().def( refl::type_attr::kHash, [](const Object* self, const Function& fn_hash) -> int64_t { auto* obj = static_cast(self); return fn_hash(AnyView(obj->key)).cast(); }); // Custom equality: only compare "key" fields refl::TypeAttrDef().def( refl::type_attr::kEq, [](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool { auto* a = static_cast(lhs); auto* b = static_cast(rhs); return fn_eq(AnyView(a->key), AnyView(b->key)).cast(); }); // Custom three-way comparison refl::TypeAttrDef().def( refl::type_attr::kCompare, [](const Object* lhs, const Object* rhs, const Function& fn_cmp) -> int32_t { auto* a = static_cast(lhs); auto* b = static_cast(rhs); return fn_cmp(AnyView(a->key), AnyView(b->key)).cast(); }); // Custom repr refl::TypeAttrDef().def( refl::type_attr::kRepr, [](const Object* self, const Function& fn_repr) -> String { auto* obj = static_cast(self); return "MyObj(key=" + fn_repr(AnyView(obj->key)).cast() + ")"; }); .. list-table:: Type attribute hooks :header-rows: 1 :widths: 30 30 40 * - Attribute - Constant - Signature * - ``__ffi_hash__`` - ``type_attr::kHash`` - ``(const Object* self, Function fn_hash) -> int64_t`` * - ``__ffi_eq__`` - ``type_attr::kEq`` - ``(const Object* lhs, const Object* rhs, Function fn_eq) -> bool`` * - ``__ffi_compare__`` - ``type_attr::kCompare`` - ``(const Object* lhs, const Object* rhs, Function fn_cmp) -> int32_t`` * - ``__ffi_repr__`` - ``type_attr::kRepr`` - ``(const Object* self, Function fn_repr) -> String`` The ``fn_hash``, ``fn_eq``, ``fn_cmp``, and ``fn_repr`` callbacks are provided by the framework for recursing into child values. Python ``c_class`` Decorator ----------------------------- On the Python side, ``tvm_ffi.dataclasses.c_class`` provides equivalent functionality to Python's ``@dataclass`` for C++-backed objects: .. code-block:: python from tvm_ffi.dataclasses import c_class @c_class("my_ext.Point") class Point(tvm_ffi.Object): x: int y: int label: str The decorator: 1. Looks up the C++ type info registered by ``ObjectDef``. 2. Matches Python annotations to C++ fields. 3. Generates ``__init__`` from ``__ffi_init__``, respecting ``kw_only``, defaults, and ``init(false)`` settings from C++. 4. Installs ``__copy__``, ``__deepcopy__``, ``__eq__``, ``__hash__``, ``__repr__``, and comparison operators. .. note:: ``@tvm_ffi.register_object`` can also be used, which delegates to ``c_class`` internally for objects with reflected fields. Inheritance ----------- Dataclass traits compose across inheritance. A child class inherits the parent's fields and adds its own: .. code-block:: cpp class ParentObj : public ffi::Object { public: int64_t parent_required; int64_t parent_default; static constexpr bool _type_mutable = true; static constexpr uint32_t _type_child_slots = 1; TVM_FFI_DECLARE_OBJECT_INFO("my_ext.Parent", ParentObj, ffi::Object); }; class ChildObj : public ParentObj { public: int64_t child_required; int64_t child_kw_only; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my_ext.Child", ChildObj, ParentObj); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; refl::ObjectDef() .def_rw("parent_required", &ParentObj::parent_required) .def_rw("parent_default", &ParentObj::parent_default, refl::default_(int64_t{5})); refl::ObjectDef() .def_rw("child_required", &ChildObj::child_required) .def_rw("child_kw_only", &ChildObj::child_kw_only, refl::kw_only(true)); } In Python, the child's auto-init includes all fields: .. code-block:: python # Generated signature: # def __init__(self, parent_required, child_required, parent_default=5, *, child_kw_only): child = Child(1, 2, child_kw_only=3) assert child.parent_required == 1 assert child.child_required == 2 assert child.parent_default == 5 # uses default assert child.child_kw_only == 3 Further Reading --------------- - :doc:`export_func_cls`: Basic function and class export guide - :doc:`../concepts/object_and_class`: Object system fundamentals, type registry, and ABI - :ref:`sec-stubgen`: Auto-generating Python type stubs from reflection metadata - :doc:`cpp_lang_guide`: Full C++ API guide covering ``Any``, ``Function``, containers tvm-ffi-0.1.12/docs/guides/export_func_cls.rst000066400000000000000000000242761521067262500213230ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. Export Functions and Classes ============================ TVM-FFI provides three mechanisms to make functions and classes available across C, C++, and Python. Each targets a different use case: .. list-table:: :header-rows: 1 :widths: 18 42 40 * - Mechanism - When to use - How it works * - :ref:`C Symbols ` - Kernel libraries, compiler codegen - Export ``__tvm_ffi_`` in a shared library; load via :py:func:`tvm_ffi.load_module` * - :ref:`Global Functions ` - Application-level APIs, cross-language callbacks - Register by string name; retrieve from any language * - :ref:`Classes ` - Structured data with fields and methods - Define a C++ :cpp:class:`~tvm::ffi::Object` subclass; use from Python as a dataclass Include the umbrella header to access all C++ APIs used in this guide: .. code-block:: cpp #include Metadata (type signatures, field names, docstrings) is captured automatically and can be turned into Python type hints via stub generation. .. seealso:: - All C++ examples in this guide are under `examples/python_packaging/ `_; ANSI C examples are under `examples/stable_c_abi/ `_. - :doc:`../concepts/func_module`: Calling convention, module system, and global registry concepts. - :doc:`../get_started/stable_c_abi`: Low-level C ABI walkthrough. - :doc:`../packaging/python_packaging`: Packaging extensions as Python wheels. - :doc:`../packaging/stubgen`: Generating Python type stubs from C++ metadata. .. _export-c-symbols: C Symbols --------- C symbols are the most direct export mechanism. A function is compiled into a shared library with a ``__tvm_ffi_`` symbol, then loaded dynamically at runtime. This is the recommended approach for kernel libraries and compiler codegen because it keeps the language boundary thin. .. tip:: For exporting and calling C symbols from pure ANSI C code, see :doc:`../get_started/stable_c_abi`. Export and Look up in C++ ~~~~~~~~~~~~~~~~~~~~~~~~~ **Export.** Use :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` to export a C++ function as a C symbol that follows the :ref:`TVM-FFI calling convention `: .. literalinclude:: ../../examples/python_packaging/src/extension.cc :language: cpp :start-after: [tvm_ffi_abi.begin] :end-before: [tvm_ffi_abi.end] This creates a symbol ``__tvm_ffi_add_two`` in the shared library. The macro handles argument unmarshalling and error propagation automatically. **Look up.** Use :cpp:func:`tvm::ffi::Module::LoadFromFile` to load a shared library and retrieve functions by name: .. code-block:: cpp namespace ffi = tvm::ffi; ffi::Module mod = ffi::Module::LoadFromFile("path/to/library.so"); ffi::Function func = mod->GetFunction("add_two").value(); int result = func(40).cast(); // -> 42 Load from Python ~~~~~~~~~~~~~~~~ Use :py:func:`tvm_ffi.load_module` to load the shared library and call its functions by name: .. code-block:: python import tvm_ffi mod = tvm_ffi.load_module("path/to/library.so") result = mod.add_two(40) # -> 42 .. seealso:: For DSO loading at Python package level, see :ref:`sec-load-the-library` in the Python packaging guide. .. _export-embedded-binary-data: Embedded Binary Data ~~~~~~~~~~~~~~~~~~~~ Shared libraries can embed a binary symbol ``__tvm_ffi__library_bin`` alongside function symbols to support **composite modules** — modules that import custom sub-modules (e.g., a PTX module loaded via ``cuModuleLoad``). The binary layout is: .. code-block:: text [val0: bytes] [val1: bytes] ... - ``nbytes``: total byte count following this header. - ``import_tree``: a CSR sparse array (``> >``) encoding the parent-child relationships among module nodes. - Each ``key`` is a module kind string, or the special value ``_lib`` for the host dynamic library itself. For entries other than ``_lib``, ``val`` contains the serialized bytes of the custom sub-module. - Both ``str`` and ``bytes`` values are length-prefixed: `` ``. When :py:func:`tvm_ffi.load_module` opens a library containing this symbol, it deserializes each sub-module by calling ``ffi.Module.load_from_bytes.``, reconstructs the import tree, and returns the composed module. The custom module class must be available at load time — either by importing its runtime library beforehand or by embedding the class definition in the generated library. See :ref:`Custom Modules ` in :doc:`../concepts/func_module` for details on custom module subclassing. .. _export-global-functions: Global Functions ---------------- Global functions are registered by string name in a shared registry, making them accessible from any language without loading a specific module. This is useful for application-level APIs and cross-language callbacks. Register and Retrieve in C++ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ **Register.** Use :cpp:class:`tvm::ffi::reflection::GlobalDef` inside a :c:macro:`TVM_FFI_STATIC_INIT_BLOCK` to register a function during library initialization: .. literalinclude:: ../../examples/python_packaging/src/extension.cc :language: cpp :start-after: [global_function.begin] :end-before: [global_function.end] The function becomes available by its string name (``my_ffi_extension.add_one``) from any language once the library is loaded. **Retrieve.** Use :cpp:func:`tvm::ffi::Function::GetGlobal` (returns ``Optional``) or :cpp:func:`tvm::ffi::Function::GetGlobalRequired` (throws if missing): .. code-block:: cpp namespace ffi = tvm::ffi; // Optional retrieval ffi::Optional maybe_func = ffi::Function::GetGlobal("my_ffi_extension.add_one"); if (maybe_func.has_value()) { int result = maybe_func.value()(3).cast(); } // Required retrieval (throws if not found) ffi::Function func = ffi::Function::GetGlobalRequired("my_ffi_extension.add_one"); int result = func(3).cast(); // -> 4 Register and Retrieve in Python ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ **Register.** Use the :py:func:`tvm_ffi.register_global_func` decorator: .. code-block:: python import tvm_ffi @tvm_ffi.register_global_func("my_ffi_extension.add_one") def add_one(x: int) -> int: return x + 1 **Retrieve.** Use :py:func:`tvm_ffi.get_global_func` to look up a function by name: .. code-block:: python import tvm_ffi add_one = tvm_ffi.get_global_func("my_ffi_extension.add_one") result = add_one(3) # -> 4 .. seealso:: For packaged extensions, :ref:`stub generation ` produces type-annotated bindings so that users can call global functions directly (e.g. ``my_ffi_extension.add_one(3)``) with full IDE support. .. note:: Global functions can also be retrieved via :cpp:func:`TVMFFIFunctionGetGlobal` in C. .. _export-classes: Classes ------- Any class derived from :cpp:class:`tvm::ffi::Object` can be registered, exported, and instantiated from Python. The reflection helper :cpp:class:`tvm::ffi::reflection::ObjectDef` makes it easy to expose: - **Fields**: immutable via :cpp:func:`ObjectDef::def_ro `, mutable via :cpp:func:`ObjectDef::def_rw ` - **Methods**: instance via :cpp:func:`ObjectDef::def `, static via :cpp:func:`ObjectDef::def_static ` - **Constructors** via ``init()`` Register in C++ ~~~~~~~~~~~~~~~ Define a class that inherits from :cpp:class:`tvm::ffi::Object`, then register it with :cpp:class:`tvm::ffi::reflection::ObjectDef` inside a :c:macro:`TVM_FFI_STATIC_INIT_BLOCK`: .. literalinclude:: ../../examples/python_packaging/src/extension.cc :language: cpp :start-after: [object.begin] :end-before: [object.end] Key elements: - :c:macro:`TVM_FFI_DECLARE_OBJECT_INFO_FINAL` declares the type with a unique string key (``my_ffi_extension.IntPair``), the class name, and parent class. - ``static constexpr bool _type_mutable = true`` allows field modification from Python. Omit this (or set to ``false``) for immutable objects. Use in Python ~~~~~~~~~~~~~ After importing the extension, the class is available with property access and method calls: .. code-block:: python import my_ffi_extension pair = my_ffi_extension.IntPair(1, 2) print(pair.a) # -> 1 print(pair.b) # -> 2 print(pair.sum()) # -> 3 .. seealso:: For packaged extensions, :ref:`stub generation ` produces type-annotated Python classes with full IDE support. Further Reading --------------- - :doc:`../get_started/stable_c_abi`: End-to-end C ABI walkthrough with callee and caller examples - :doc:`../concepts/func_module`: Calling convention, module system, and global registry concepts - :doc:`../concepts/object_and_class`: Object system, type hierarchy, and reference counting - :doc:`../packaging/python_packaging`: Packaging extensions as Python wheels with stub generation - :doc:`../packaging/cpp_tooling`: Build toolchain, CMake integration, and library distribution - :doc:`cpp_lang_guide`: Full C++ API guide covering Any, Function, containers, and more tvm-ffi-0.1.12/docs/guides/kernel_library_guide.rst000066400000000000000000000311311521067262500222730ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. Kernel Library Guide ==================== This guide covers shipping C++/CUDA kernel libraries with TVM-FFI. The resulting libraries are agnostic to Python version and ML framework — a single ``.so`` works with PyTorch, JAX, PaddlePaddle, NumPy, and more. .. seealso:: - :doc:`../get_started/quickstart`: End-to-end walkthrough of a simpler ``add_one`` kernel - :doc:`../packaging/cpp_tooling`: Build toolchain, CMake integration, and library distribution - All example code in this guide is under `examples/kernel_library/ `_. - Production examples: `FlashInfer `_ ships CUDA kernels via TVM-FFI. Anatomy of a Kernel Function ----------------------------- Every TVM-FFI CUDA kernel follows the same sequence: 1. **Validate** inputs (device, dtype, shape, contiguity) 2. **Set device guard** to match the tensor's device 3. **Acquire stream** from the host framework 4. **Dispatch** on dtype and **launch** the kernel Here is a complete ``Scale`` kernel that computes ``y = x * factor``: .. literalinclude:: ../../examples/kernel_library/scale_kernel.cu :language: cpp :start-after: [function.begin] :end-before: [function.end] The CUDA kernel itself is a standard ``__global__`` function: .. literalinclude:: ../../examples/kernel_library/scale_kernel.cu :language: cpp :start-after: [cuda_kernel.begin] :end-before: [cuda_kernel.end] The following subsections break down each step. Input Validation ~~~~~~~~~~~~~~~~ Kernel functions should validate inputs early and fail with clear error messages. A common pattern is to define reusable ``CHECK_*`` macros on top of :c:macro:`TVM_FFI_CHECK` (see :doc:`../concepts/exception_handling`): .. literalinclude:: ../../examples/kernel_library/tvm_ffi_utils.h :language: cpp :start-after: [check_macros.begin] :end-before: [check_macros.end] For **user-facing errors** (bad arguments, unsupported dtypes, shape mismatches), use :c:macro:`TVM_FFI_THROW` or :c:macro:`TVM_FFI_CHECK` with a specific error kind so that callers receive an actionable message: .. code-block:: cpp TVM_FFI_THROW(TypeError) << "Unsupported dtype: " << input.dtype(); TVM_FFI_CHECK(input.numel() > 0, ValueError) << "input must be non-empty"; TVM_FFI_CHECK(input.numel() == output.numel(), ValueError) << "size mismatch"; For **internal invariants** that indicate bugs in the kernel itself, use :c:macro:`TVM_FFI_ICHECK`: .. code-block:: cpp TVM_FFI_ICHECK_GE(n, 0) << "element count must be non-negative"; Device Guard and Stream ~~~~~~~~~~~~~~~~~~~~~~~ Before launching a CUDA kernel, two things must happen: 1. **Set the CUDA device** to match the tensor's device. :cpp:class:`tvm::ffi::CUDADeviceGuard` is an RAII guard that calls ``cudaSetDevice`` on construction and restores the original device on destruction. 2. **Acquire the stream** from the host framework via :cpp:func:`TVMFFIEnvGetStream`. When Python code calls a kernel with PyTorch tensors, TVM-FFI automatically captures PyTorch's current stream for the tensor's device. A small helper keeps this concise: .. literalinclude:: ../../examples/kernel_library/tvm_ffi_utils.h :language: cpp :start-after: [get_stream.begin] :end-before: [get_stream.end] Every kernel function then follows the same two-line pattern: .. code-block:: cpp ffi::CUDADeviceGuard guard(input.device().device_id); cudaStream_t stream = get_cuda_stream(input.device()); See :doc:`../concepts/tensor` for details on stream handling and automatic stream context updates. Dtype Dispatch ~~~~~~~~~~~~~~ Kernels typically support multiple dtypes. Dispatch on :c:struct:`DLDataType` at runtime while instantiating templates at compile time: .. code-block:: cpp constexpr DLDataType dl_float32 = DLDataType{kDLFloat, 32, 1}; constexpr DLDataType dl_float16 = DLDataType{kDLFloat, 16, 1}; if (input.dtype() == dl_float32) { ScaleKernel<<>>( static_cast(output.data_ptr()), ...); } else if (input.dtype() == dl_float16) { ScaleKernel<<>>( static_cast(output.data_ptr()), ...); } else { TVM_FFI_THROW(TypeError) << "Unsupported dtype: " << input.dtype(); } For libraries that support many dtypes, define dispatch macros (see `FlashInfer's tvm_ffi_utils.h `_ for a production example). Export and Load --------------- Export and Build ~~~~~~~~~~~~~~~~ **Export.** Use :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` to create a C symbol that follows the :doc:`TVM-FFI calling convention <../concepts/func_module>`: .. literalinclude:: ../../examples/kernel_library/scale_kernel.cu :language: cpp :start-after: [export.begin] :end-before: [export.end] This creates a symbol ``__tvm_ffi_scale`` in the shared library. **Build.** Compile the kernel into a shared library using GCC/NVCC or CMake (see :doc:`../packaging/cpp_tooling` for full details): .. code-block:: bash nvcc -shared -O3 scale_kernel.cu -o build/scale_kernel.so \ -Xcompiler -fPIC,-fvisibility=hidden \ $(tvm-ffi-config --cxxflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) **Optional arguments.** Wrap any argument type with :cpp:class:`tvm::ffi::Optional` to accept ``None`` from the Python side: .. code-block:: cpp void MyKernel(TensorView output, TensorView input, Optional bias, Optional scale) { if (bias.has_value()) { // use bias.value().data_ptr() } double s = scale.value_or(1.0); } .. code-block:: python mod.my_kernel(y, x, None, None) # no bias, default scale mod.my_kernel(y, x, bias_tensor, 2.0) # with bias and scale Load from Python ~~~~~~~~~~~~~~~~ Use :py:func:`tvm_ffi.load_module` to load the library and call its functions. PyTorch tensors (and other framework tensors) are automatically converted to :cpp:class:`~tvm::ffi::TensorView` at the ABI boundary: .. literalinclude:: ../../examples/kernel_library/load_scale.py :language: python :start-after: [load_and_call.begin] :end-before: [load_and_call.end] See :doc:`../get_started/quickstart` for examples with JAX, PaddlePaddle, NumPy, CuPy, Rust, and pure C++. Tensor Handling --------------- TensorView vs Tensor ~~~~~~~~~~~~~~~~~~~~ TVM-FFI provides two tensor types (see :doc:`../concepts/tensor` for full details): :cpp:class:`~tvm::ffi::TensorView` *(non-owning)* A lightweight view of an existing tensor. **Use this for kernel parameters.** It adds no reference count overhead and works with all framework tensors. :cpp:class:`~tvm::ffi::Tensor` *(owning)* A reference-counted tensor that manages its own lifetime. Use this only when you need to **allocate and return** a tensor from C++. .. important:: Prefer :cpp:class:`~tvm::ffi::TensorView` in kernel signatures. It is more lightweight, supports more use cases (including XLA buffers that only provide views), and avoids unnecessary reference counting. Tensor Metadata ~~~~~~~~~~~~~~~ Both :cpp:class:`~tvm::ffi::TensorView` and :cpp:class:`~tvm::ffi::Tensor` expose identical metadata accessors. These are the methods kernel code uses most: validating inputs, computing launch parameters, and accessing data pointers. **Shape and elements.** :cpp:func:`~tvm::ffi::TensorView::ndim` returns the number of dimensions, :cpp:func:`~tvm::ffi::TensorView::shape` returns the full shape as a :cpp:class:`~tvm::ffi::ShapeView` (a lightweight ``span``-like view of ``int64_t``), and :cpp:func:`~tvm::ffi::TensorView::size` returns the size of a single dimension (supports negative indexing, e.g. ``size(-1)`` for the last dimension). :cpp:func:`~tvm::ffi::TensorView::numel` returns the total element count — use it for computing grid dimensions: .. code-block:: cpp int64_t n = input.numel(); int threads = 256; int blocks = (n + threads - 1) / threads; **Dtype.** :cpp:func:`~tvm::ffi::TensorView::dtype` returns a :c:struct:`DLDataType` with three fields: ``code`` (e.g. ``kDLFloat``, ``kDLBfloat``), ``bits`` (e.g. 16, 32), and ``lanes`` (almost always 1). Compare it against predefined constants to dispatch on dtype: .. code-block:: cpp constexpr DLDataType dl_float32 = DLDataType{kDLFloat, 32, 1}; if (input.dtype() == dl_float32) { ... } **Device.** :cpp:func:`~tvm::ffi::TensorView::device` returns a :c:struct:`DLDevice` with ``device_type`` (e.g. ``kDLCUDA``) and ``device_id``. Use these for validation and to set the device guard: .. code-block:: cpp TVM_FFI_ICHECK_EQ(input.device().device_type, kDLCUDA); ffi::CUDADeviceGuard guard(input.device().device_id); **Data pointer.** :cpp:func:`~tvm::ffi::TensorView::data_ptr` returns ``void*``; cast it to the appropriate typed pointer before passing it to a kernel: .. code-block:: cpp auto* out = static_cast(output.data_ptr()); auto* in = static_cast(input.data_ptr()); **Strides and contiguity.** :cpp:func:`~tvm::ffi::TensorView::strides` returns the stride array as a :cpp:class:`~tvm::ffi::ShapeView`, and :cpp:func:`~tvm::ffi::TensorView::stride` returns a single dimension's stride. :cpp:func:`~tvm::ffi::TensorView::IsContiguous` checks whether the tensor is contiguous in memory. Most kernels require contiguous inputs — the ``CHECK_CONTIGUOUS`` macro shown above enforces this at the top of each function. .. tip:: The API is designed to be familiar to PyTorch developers. ``dim()``, ``sizes()``, ``size(i)``, ``stride(i)``, and ``is_contiguous()`` are all available as aliases of their TVM-FFI counterparts. See :doc:`../concepts/tensor` for the full API reference. Tensor Allocation ~~~~~~~~~~~~~~~~~ **Always pre-allocate output tensors on the Python side** and pass them into the kernel as :cpp:class:`~tvm::ffi::TensorView` parameters. Allocating tensors inside a kernel function is almost never the right choice: - it causes **memory fragmentation** from repeated small allocations, - it **breaks CUDA graph capture**, which requires deterministic memory addresses, and - it **bypasses the framework's allocator** (caching pools, device placement, memory planning). The pre-allocation pattern is straightforward: .. code-block:: python # Python: pre-allocate output y = torch.empty_like(x) mod.scale(y, x, 2.0) .. code-block:: cpp // C++: kernel writes into pre-allocated output void Scale(TensorView output, TensorView input, double factor); If C++-side allocation is truly unavoidable — for example, when the output shape is data-dependent and cannot be determined before the kernel runs — use :cpp:func:`tvm::ffi::Tensor::FromEnvAlloc` to at least reuse the host framework's allocator (e.g., ``torch.empty`` under PyTorch): .. literalinclude:: ../../examples/kernel_library/tvm_ffi_utils.h :language: cpp :start-after: [alloc_tensor.begin] :end-before: [alloc_tensor.end] For custom allocators (e.g., ``cudaMalloc``/``cudaFree``), use :cpp:func:`tvm::ffi::Tensor::FromNDAlloc`. Note that the kernel library must outlive any tensors allocated this way, since the custom deleter lives in the library. See :doc:`../concepts/tensor` for details. Further Reading --------------- - :doc:`../get_started/quickstart`: End-to-end walkthrough shipping ``add_one`` across frameworks and languages - :doc:`../packaging/cpp_tooling`: Build toolchain, CMake integration, GCC/NVCC flags, and library distribution - :doc:`../packaging/python_packaging`: Packaging kernel libraries as Python wheels - :doc:`../concepts/tensor`: Tensor classes, DLPack interop, stream handling, and allocation APIs - :doc:`../concepts/func_module`: Function calling convention, modules, and the global registry - :doc:`../concepts/exception_handling`: Error handling across language boundaries - :doc:`../concepts/abi_overview`: Low-level C ABI details tvm-ffi-0.1.12/docs/guides/python_lang_guide.md000066400000000000000000000276231521067262500214140ustar00rootroot00000000000000 # Python Guide This guide introduces the `tvm_ffi` Python package. At a high level, the `tvm_ffi` Python package provides first-class Python support for - Pythonic classes to represent values in TVM FFI Any ABI. - Mechanisms to call into TVM FFI ABI compatible functions. - Conversion between Python values and `tvm_ffi` values. In this guide, we will run examples that make use of pre-registered testing functions in `tvm_ffi`. If so, we will also briefly copy snippets that show the corresponding C++ behavior. ## Load and Run Module The most common use case of TVM FFI is to load a runnable module and run the corresponding function. You can follow the [quickstart guide](../get_started/quickstart.rst) for details on building the library `build/add_one_cpu.so`. Let's walk through the load and run example again for NumPy ```python import tvm_ffi import numpy as np # Load the compiled module mod = tvm_ffi.load_module("build/add_one_cpu.so") # Create input and output arrays x = np.array([1, 2, 3, 4, 5], dtype=np.float32) y = np.empty_like(x) # Call the function mod.add_one_cpu(x, y) ``` In this case, {py:func}`tvm_ffi.load_module` will return a {py:class}`tvm_ffi.Module` class that contains the exported functions. You can access the functions by their names. ## Tensor `tvm_ffi` provides a managed DLPack-compatible Tensor. ```python import numpy as np import tvm_ffi # Demonstrate DLPack conversion between NumPy and TVM FFI np_data = np.array([1, 2, 3, 4], dtype=np.float32) tvm_array = tvm_ffi.from_dlpack(np_data) # Convert back to NumPy np_result = np.from_dlpack(tvm_array) ``` In most cases, however, you do not have to explicitly create Tensors. The Python interface can take in `torch.Tensor` and `numpy.ndarray` objects and automatically convert them to {py:class}`tvm_ffi.Tensor`. ## Functions and Callbacks {py:class}`tvm_ffi.Function` provides the Python interface for `ffi::Function` in the C++. You can retrieve globally registered functions via {py:func}`tvm_ffi.get_global_func`. ```python import tvm_ffi # testing.echo is defined and registered in C++ # [](ffi::Any x) { return x; } fecho = tvm_ffi.get_global_func("testing.echo") assert fecho(1) == 1 ``` You can pass a Python function as an argument to another FFI function as callbacks. Under the hood, {py:func}`tvm_ffi.convert` is called to convert the Python function into a {py:class}`tvm_ffi.Function`. ```python import tvm_ffi # testing.apply is registered in C++ # [](ffi::Function f, ffi::Any val) { return f(x); } fapply = tvm_ffi.get_global_func("testing.apply") # invoke fapply with lambda callback as f assert fapply(lambda x: x + 1, 1) == 2 ``` This is a very powerful pattern that allows us to inject Python callbacks into the C++ code. You can also register a Python callback as a global function. ```python import tvm_ffi @tvm_ffi.register_global_func("example.add_one") def add_one(a): return a + 1 assert tvm_ffi.get_global_func("example.add_one")(1) == 2 ``` ## Container Types TVM FFI provides five container types that split into **immutable** (copy-on-write) and **mutable** (shared reference) variants. See the [Containers concept page](../concepts/containers.rst) for a detailed overview. | Type | Mutability | Python ABC | Semantics | | ------ | ----------- | ------------ | ----------- | | {py:class}`tvm_ffi.Array` | Immutable | `Sequence[T]` | Homogeneous sequence (copy-on-write) | | {py:class}`tvm_ffi.List` | Mutable | `MutableSequence[T]` | Homogeneous sequence (shared reference) | | {py:class}`tvm_ffi.Map` | Immutable | `Mapping[K, V]` | Homogeneous mapping (copy-on-write) | | {py:class}`tvm_ffi.Dict` | Mutable | `MutableMapping[K, V]` | Homogeneous mapping (shared reference) | ### Array (immutable sequence) When an FFI function takes arguments from lists/tuples, they will be converted into {py:class}`tvm_ffi.Array`. Arrays are **read-only** from Python -- they support indexing and iteration but not mutation. ```python import tvm_ffi # Python lists become Arrays arr = tvm_ffi.convert([1, 2, 3, 4]) assert isinstance(arr, tvm_ffi.Array) assert len(arr) == 4 assert arr[0] == 1 # arr[0] = 10 # TypeError: Array does not support item assignment ``` ### List (mutable sequence) {py:class}`tvm_ffi.List` is a mutable sequence with shared-reference semantics. All handles sharing the same underlying `ListObj` see mutations immediately. ```python import tvm_ffi lst = tvm_ffi.List([1, 2, 3]) assert len(lst) == 3 # Mutate in-place lst.append(4) assert len(lst) == 4 lst[0] = 10 assert lst[0] == 10 ``` ### Map (immutable mapping) Dictionaries will be converted to {py:class}`tvm_ffi.Map`. Maps are **read-only** from Python -- they support key lookup and iteration but not mutation. ```python import tvm_ffi map_obj = tvm_ffi.convert({"a": 1, "b": 2}) assert isinstance(map_obj, tvm_ffi.Map) assert len(map_obj) == 2 assert map_obj["a"] == 1 assert map_obj["b"] == 2 # map_obj["c"] = 3 # TypeError: Map does not support item assignment ``` ### Dict (mutable mapping) {py:class}`tvm_ffi.Dict` is a mutable mapping with shared-reference semantics. All handles sharing the same underlying `DictObj` see mutations immediately. ```python import tvm_ffi d = tvm_ffi.Dict({"a": 1, "b": 2}) d["c"] = 3 assert len(d) == 3 del d["a"] assert len(d) == 2 ``` ### Tuple conversion When a Python tuple is passed to an FFI function, it is converted to {py:class}`tvm_ffi.Array` (the same backing type). The C++ `Tuple` template provides compile-time heterogeneous type safety, but on the Python side both lists and tuples map to `Array`. ```python import tvm_ffi t = tvm_ffi.convert((1, "hello", 3.14)) assert isinstance(t, tvm_ffi.Array) ``` When container values are returned from FFI functions, they are stored in the corresponding container type (Array, List, Map, or Dict). ## Inline Module You can also load a _inline module_ where the C++/CUDA code is directly embedded in the Python script and then compiled on the fly. For example, we can define a simple kernel that adds one to each element of an array as follows: ```python import torch from tvm_ffi import Module import tvm_ffi.cpp # define the cpp source code cpp_source = ''' void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } ''' # compile the cpp source code and load the module mod: Module = tvm_ffi.cpp.load_inline( name='hello', cpp_sources=cpp_source, functions='add_one_cpu' ) # use the function from the loaded module to perform x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) ``` The above code defines a C++ function `add_one_cpu` in Python script, compiles it on the fly and then loads the compiled {py:class}`tvm_ffi.Module` object via {py:func}`tvm_ffi.cpp.load_inline`. You can then call the function `add_one_cpu` from the module as usual. We can also use {py:func}`tvm_ffi.cpp.build_inline` to build the inline module without loading it. The built shared library is returned and can be loaded via {py:func}`tvm_ffi.load_module`. ## Error Handling An FFI function may raise an error. In such cases, the Python package will automatically translate the error to the corresponding error kind in Python ```python import tvm_ffi # defined in C++ # [](String kind, String msg) { throw Error(kind, msg, backtrace); } test_raise_error = tvm_ffi.get_global_func("testing.test_raise_error") test_raise_error("ValueError", "message") ``` The above code shows an example where an error is raised in C++, resulting in the following error trace ```text Traceback (most recent call last): File "example.py", line 7, in test_raise_error("ValueError", "message") ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^ File "python/tvm_ffi/cython/function.pxi", line 927, in tvm_ffi.core.Function.__call__ raise error.py_error() ^^^ File "src/ffi/extra/testing.cc", line 60, in void tvm::ffi::TestRaiseError(tvm::ffi::String, tvm::ffi::String) throw ffi::Error(kind, msg, TVMFFITraceback(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0)); ``` We register common error kinds. You can also register extra error dispatch via the {py:func}`tvm_ffi.register_error` function. ## Advanced: Register Your Own Object For advanced use cases, you may want to register your own objects. This can be achieved through the reflection registry in the TVM-FFI API. First, let's review the C++ side of the code. For this example, you do not need to change the C++ side as this code is pre-shipped with the testing module of the `tvm_ffi` package. ```cpp #include // Step 1: Define the object class (stores the actual data) class TestIntPairObj : public tvm::ffi::Object { public: int64_t a; int64_t b; TestIntPairObj() = default; TestIntPairObj(int64_t a, int64_t b) : a(a), b(b) {} // Required: declare type information TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestIntPair", TestIntPairObj, tvm::ffi::Object); }; // Step 2: Define the reference wrapper (user-facing interface) class TestIntPair : public tvm::ffi::ObjectRef { public: // Constructor explicit TestIntPair(int64_t a, int64_t b) { data_ = tvm::ffi::make_object(a, b); } // Required: define object reference methods TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestIntPair, tvm::ffi::ObjectRef, TestIntPairObj); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; // register the object into the system // register field accessors and a global static function `__ffi_init__` as ffi::Function refl::ObjectDef() .def(refl::init()) .def_ro("a", &TestIntPairObj::a) .def_ro("b", &TestIntPairObj::b); } ``` You can then create wrapper classes for objects that are in the library as follows: ```python import tvm_ffi # Register the class @tvm_ffi.register_object("testing.TestIntPair") class TestIntPair(tvm_ffi.Object): def __init__(self, a, b): # This is a special method to call an FFI function whose return # value exactly initializes the object handle of the object self.__ffi_init__(a, b) test_int_pair = TestIntPair(1, 2) # We can access the fields by name # The properties are populated by the reflection mechanism assert test_int_pair.a == 1 assert test_int_pair.b == 2 ``` Under the hood, we leverage the information registered through the reflection registry to generate efficient field accessors and methods for each class. Importantly, when you have multiple inheritance, you need to call {py:func}`tvm_ffi.register_object` on both the base class and the child class. tvm-ffi-0.1.12/docs/guides/rust_lang_guide.md000066400000000000000000000110621521067262500210560ustar00rootroot00000000000000 # Rust Guide ```{note} The Rust support is currently in an experimental stage. ``` This guide demonstrates how to use TVM FFI from Rust applications. ## Installation ### Prerequisites The Rust support depends on `libtvm_ffi`. First, install the `tvm-ffi` Python package: ```bash pip install -v -e . ``` Confirm that `tvm-ffi-config` is available: ```bash tvm-ffi-config --libdir ``` ### Adding to Your Project Add to your `Cargo.toml`: ```toml [dependencies] tvm-ffi = { path = "path/to/tvm-ffi/rust/tvm-ffi" } ``` For published versions (when available): ```toml [dependencies] tvm-ffi = "0.1.0-alpha.0" ``` ### Environment Setup Set the library path so `libtvm_ffi` can be found at runtime: ```bash export LD_LIBRARY_PATH=$(tvm-ffi-config --libdir):$LD_LIBRARY_PATH ``` ## Basic Usage ### Loading a Module Load a compiled TVM FFI module and call its functions: ```rust use tvm_ffi::{Module, Result}; fn main() -> Result<()> { // Load compiled module let module = Module::load_from_file("build/add_one_cpu.so")?; // Get function by name let add_fn = module.get_function("add_one_cpu")?; Ok(()) } ``` ### Working with Tensors Create and manipulate tensors: ```rust use tvm_ffi::Tensor; // Create a tensor from a slice let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; let tensor = Tensor::from_slice(&data, &[2, 3])?; ``` ### Calling Functions Call functions with tensors: ```rust use tvm_ffi::{Module, Tensor, Result}; fn run_example() -> Result<()> { let module = Module::load_from_file("build/add_one_cpu.so")?; let func = module.get_function("add_one_cpu")?; // Create input and output tensors let input = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[4])?; let output = Tensor::from_slice(&[0.0f32; 4], &[4])?; // Call function func.call_tuple((&input, &output))?; Ok(()) } ``` ## Advanced Topics ### Global Functions Register and access global functions: ```rust use tvm_ffi::Function; // Get global function let func = Function::get_global("my_function")?; // Register a new global function let my_func = Function::from_packed(|args: &[AnyView]| -> Result { // Function implementation Ok(Any::default()) }); Function::register_global("my_custom_func", my_func)?; ``` ### Type-Erased Functions Create functions from Rust closures: ```rust use tvm_ffi::{Function, Any, AnyView, Result}; // From packed closure let func = Function::from_packed(|args: &[AnyView]| -> Result { // Process args and return result Ok(Any::default()) }); // From typed closure let typed_func = Function::from_typed(|x: i64, y: i64| -> Result { Ok(x + y) }); ``` ### Error Handling TVM FFI uses standard Rust `Result` types: ```rust use tvm_ffi::{Error, Module, Result, VALUE_ERROR}; fn may_fail(value: i32) -> Result<()> { // Operations that may fail let module = Module::load_from_file("path.so")?; // Custom errors if value < 0 { return Err(Error::new( VALUE_ERROR, "Value must be non-negative", "" )); } Ok(()) } ``` ## Examples The repository includes a complete example in `rust/tvm-ffi/examples/load_library.rs`. Run it with: ```bash cd rust cargo run --example load_library --features example ``` ## Building the Workspace Build the entire Rust workspace: ```bash cd rust cargo build ``` Run tests: ```bash cargo test ``` ## API Reference For detailed API documentation, see the [Rust API Reference](../reference/rust/index.rst). ## Related Resources - [Quick Start Guide](../get_started/quickstart.rst) - General TVM FFI introduction - [C++ Guide](./cpp_lang_guide.md) - C++ API usage - [Python Guide](./python_lang_guide.md) - Python API usage tvm-ffi-0.1.12/docs/index.rst000066400000000000000000000045271521067262500157520ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Apache TVM FFI Documentation ============================ Welcome to the documentation for TVM FFI. You can get started by reading the get started section, or reading through the guides and concepts sections. Installation ------------ To install TVM-FFI via pip or uv, run: .. code-block:: bash pip install apache-tvm-ffi pip install torch-c-dlpack-ext # compatibility package for torch <= 2.9 Table of Contents ----------------- .. toctree:: :maxdepth: 1 :caption: Get Started get_started/quickstart.rst get_started/stable_c_abi.rst .. toctree:: :maxdepth: 1 :caption: Guides guides/export_func_cls.rst guides/dataclass_reflection.rst guides/kernel_library_guide.rst guides/compiler_integration.md guides/cubin_launcher.rst guides/python_lang_guide.md guides/cpp_lang_guide.md guides/rust_lang_guide.md .. toctree:: :maxdepth: 1 :caption: Concepts concepts/abi_overview.rst concepts/any.rst concepts/containers.rst concepts/object_and_class.rst concepts/tensor.rst concepts/func_module.rst concepts/exception_handling.rst concepts/structural_eq_hash.rst .. toctree:: :maxdepth: 1 :caption: Packaging packaging/python_packaging.rst packaging/stubgen.rst packaging/cpp_tooling.rst .. toctree:: :maxdepth: 1 :caption: Reference reference/python/index.rst reference/cpp/index.rst reference/rust/index.rst .. toctree:: :maxdepth: 1 :caption: Developer Manual dev/source_build.rst dev/doc_build.rst dev/ci_cd.rst dev/release_process.rst tvm-ffi-0.1.12/docs/packaging/000077500000000000000000000000001521067262500160255ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/packaging/cpp_tooling.rst000066400000000000000000000223351521067262500211010ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one .. or more contributor license agreements. See the NOTICE file .. distributed with this work for additional information .. regarding copyright ownership. The ASF licenses this file .. to you under the Apache License, Version 2.0 (the .. "License"); you may not use this file except in compliance .. with the License. You may obtain a copy of the License at .. .. http://www.apache.org/licenses/LICENSE-2.0 .. .. Unless required by applicable law or agreed to in writing, .. software distributed under the License is distributed on an .. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY .. KIND, either express or implied. See the License for the .. specific language governing permissions and limitations .. under the License. C++ Tooling =========== This guide covers the TVM-FFI C++ toolchain: header layout, build integration, library distribution, and editor setup. For an end-to-end walkthrough that builds, loads, and calls a C++/CUDA kernel, see :doc:`../get_started/quickstart`. .. admonition:: Prerequisite :class: hint - Compiler: C++17-capable toolchain (GCC/Clang/MSVC) - CMake: 3.18 or newer - TVM-FFI installed via .. code-block:: bash pip install --reinstall --upgrade apache-tvm-ffi C++ Headers ----------- **Core APIs.** A single umbrella header exposes most of the API surface, including :doc:`functions <../concepts/func_module>`, :doc:`objects <../concepts/object_and_class>`, :doc:`Any/AnyView <../concepts/any>`, :doc:`tensors <../concepts/tensor>`, and :doc:`exception handling <../concepts/exception_handling>`. .. code-block:: cpp #include Extra features live in dedicated headers under `tvm/ffi/extra/ `_ and should be included only when needed: - **Environment API** -- allocator, stream, and device access (see :doc:`../concepts/tensor` for usage): ``#include `` - **Dynamic module loading** (see :ref:`sec:module`): ``#include `` - **CUBIN launcher** (see :doc:`../guides/cubin_launcher`): ``#include `` Build with TVM-FFI ------------------ CMake ~~~~~ TVM-FFI ships CMake utilities and imported targets through its package configuration. The two primary functions are ``tvm_ffi_configure_target`` and ``tvm_ffi_install``, both defined in ``cmake/Utils/Library.cmake``. .. hint:: See `examples/python_packaging/CMakeLists.txt `_ for a complete working example. Configure Target """""""""""""""" ``tvm_ffi_configure_target`` wires a CMake target to TVM-FFI with sensible defaults: it links headers and the shared library, configures debug symbols, and optionally runs :ref:`stub generation ` as a post-build step. .. code-block:: cmake tvm_ffi_configure_target( [LINK_SHARED ON|OFF ] [LINK_HEADER ON|OFF ] [DEBUG_SYMBOL ON|OFF ] [MSVC_FLAGS ON|OFF ] [STUB_INIT ON|OFF ] [STUB_DIR ] [STUB_PKG ] [STUB_PREFIX ] ) :LINK_SHARED: (default: ON) Link against the TVM-FFI shared library ``tvm_ffi::shared``. Set OFF for header-only usage or deferred runtime loading. :LINK_HEADER: (default: ON) Link against the TVM-FFI headers via ``tvm_ffi::header``. Set OFF when you manage include paths and compile definitions manually. :DEBUG_SYMBOL: (default: ON) Enable debug symbol post-processing hooks. On Apple platforms this runs ``dsymutil``. :MSVC_FLAGS: (default: ON) Apply MSVC compatibility flags. :STUB_DIR: (default: "") Output directory for generated Python stubs. When set, stub generation runs as a post-build step. :STUB_INIT: (default: OFF) Allow the stub generator to initialize a new package layout. Requires ``STUB_DIR``. :STUB_PKG: (default: "") Package name passed to the stub generator. Requires ``STUB_DIR`` and ``STUB_INIT=ON``. :STUB_PREFIX: (default: "") Module prefix passed to the stub generator. Requires ``STUB_DIR`` and ``STUB_INIT=ON``. See :ref:`sec-stubgen-cmake` for a detailed explanation of each ``STUB_*`` option and the generation modes they control. Install Target """""""""""""" ``tvm_ffi_install`` installs platform-specific artifacts for a target. On Apple platforms it installs the ``.dSYM`` bundle when present; on other platforms this is currently a no-op. .. code-block:: cmake tvm_ffi_install( [DESTINATION ] ) :DESTINATION: Install destination directory relative to ``CMAKE_INSTALL_PREFIX``. Set ``tvm_ffi_ROOT`` """""""""""""""""""" If ``find_package(tvm_ffi CONFIG REQUIRED)`` fails because CMake cannot locate the package, pass ``tvm_ffi_ROOT`` explicitly: .. code-block:: bash cmake -S . -B build \ -Dtvm_ffi_ROOT="$(tvm-ffi-config --cmakedir)" .. note:: When packaging Python wheels with scikit-build-core, ``tvm_ffi_ROOT`` is discovered automatically from the active Python environment. GCC/NVCC ~~~~~~~~ For quick prototyping or CI scripts without CMake, invoke ``g++`` or ``nvcc`` directly with flags from ``tvm-ffi-config``. The examples below are from the :doc:`Quick Start <../get_started/quickstart>` tutorial: .. tabs:: .. group-tab:: C++ .. literalinclude:: ../../examples/quickstart/raw_compile.sh :language: bash :start-after: [cpp_compile.begin] :end-before: [cpp_compile.end] .. group-tab:: CUDA .. literalinclude:: ../../examples/quickstart/raw_compile.sh :language: bash :start-after: [cuda_compile.begin] :end-before: [cuda_compile.end] The three ``tvm-ffi-config`` flags provide: :``--cxxflags``: Include paths and compile definitions (``-I...``, ``-D...``) :``--ldflags``: Library search paths (``-L...``, ``-Wl,-rpath,...``) :``--libs``: Libraries to link (``-ltvm_ffi``) **RPATH handling.** The resulting shared library links against ``libtvm_ffi.so``, so the dynamic linker must be able to find it at load time: - **Python distribution.** ``import tvm_ffi`` preloads ``libtvm_ffi.so`` into the process before any user library is loaded, so the RPATH requirement is already satisfied without additional linker flags. - **Pure C++ distribution.** You must ensure ``libtvm_ffi.so`` is on the library search path. Either set ``-Wl,-rpath,$(tvm-ffi-config --libdir)`` at link time, or place ``libtvm_ffi.so`` alongside your binary. Library Distribution -------------------- When distributing pre-built shared libraries on Linux, glibc symbol versioning can cause load-time failures on systems with a different glibc version. The standard solution is the `manylinux `_ approach: **build on old glibc, run on new**. **Build environment.** Use a manylinux Docker image: .. code-block:: bash docker pull quay.io/pypa/manylinux2014_x86_64 Build host and device code inside the container. For CUDA: .. code-block:: bash nvcc -shared -Xcompiler -fPIC your_kernel.cu -o kernel.so \ $(tvm-ffi-config --cxxflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) **Verify glibc requirements.** Inspect the minimum glibc version your binary requires: .. code-block:: bash objdump -T your_kernel.so | grep GLIBC_ The ``apache-tvm-ffi`` wheel is already manylinux-compatible, so linking against it inside a manylinux build environment produces portable binaries. Editor Setup ------------ The following configuration enables code completion and diagnostics in VSCode, Cursor, or any editor backed by clangd. **CMake Tools (VSCode/Cursor).** Add these workspace settings so CMake Tools can locate TVM-FFI and generate ``compile_commands.json``: .. code-block:: json { "cmake.buildDirectory": "${workspaceFolder}/build-vscode", "cmake.configureArgs": [ "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON" ], "cmake.configureSettings": { "Python_EXECUTABLE": "${workspaceFolder}/.venv/bin/python3", "tvm_ffi_ROOT": "${workspaceFolder}/.venv/lib/pythonX.Y/site-packages/tvm_ffi/share/cmake/tvm_ffi" } } .. important:: Make sure ``Python_EXECUTABLE`` and ``tvm_ffi_ROOT`` match the virtual environment you intend to use. **clangd.** Create a ``.clangd`` file at the project root pointing to the CMake compilation database. The snippet below also strips NVCC flags that clangd does not recognize: .. code-block:: yaml CompileFlags: CompilationDatabase: build-vscode/ Remove: # for NVCC compatibility - -forward-unknown-to-host-compiler - --generate-code* - -Xcompiler* Further Reading --------------- - :doc:`../get_started/quickstart`: End-to-end walkthrough building and shipping a C++/CUDA kernel - :doc:`../dev/source_build`: Building TVM-FFI from source with CMake flags - :doc:`stubgen`: Generating Python type stubs from C++ reflection metadata - :doc:`python_packaging`: Packaging shared libraries as Python wheels - :doc:`../concepts/func_module`: Defining and exporting TVM-FFI functions - :doc:`../concepts/object_and_class`: Defining C++ classes with cross-language reflection - :doc:`../concepts/exception_handling`: Error handling across language boundaries - :doc:`../concepts/abi_overview`: Low-level C ABI details tvm-ffi-0.1.12/docs/packaging/python_packaging.rst000066400000000000000000000150231521067262500221050ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Python Packaging ================ This guide walks through a small, complete workflow for packaging a TVM-FFI extension as a Python wheel. The goal is to help you wire up a simple extension, produce a wheel, and ship user-friendly typing annotations without needing to know every detail of TVM internals. We cover three checkpoints: - Build a Python wheel; - Export C++ to Python; - Generate Python package stubs. .. note:: All code used in this guide is under `examples/python_packaging `_. .. admonition:: Prerequisite :class: hint - Python: 3.9 or newer (for the ``tvm_ffi.config``/``tvm-ffi-config`` helpers) - Compiler: C11-capable toolchain (GCC/Clang/MSVC) - TVM-FFI installed via .. code-block:: bash pip install --reinstall --upgrade apache-tvm-ffi Build Python Wheel ------------------ Start by defining the Python packaging and build wiring. TVM-FFI provides helpers to build and ship ABI-agnostic Python extensions using standard packaging tools. The steps below set up the build so you can plug in the C++ exports from the next section. The flow below uses :external+scikit_build_core:doc:`scikit-build-core ` to drive a CMake build, but the same ideas apply to setuptools or other :pep:`517` backends. CMake Target ~~~~~~~~~~~~ Assume the source tree contains ``src/extension.cc``. Create a ``CMakeLists.txt`` that creates a shared target ``my_ffi_extension`` and configures it against TVM-FFI. .. literalinclude:: ../../examples/python_packaging/CMakeLists.txt :language: cmake :start-after: [example.cmake.begin] :end-before: [example.cmake.end] Function ``tvm_ffi_configure_target`` sets up TVM-FFI include paths and links against the TVM-FFI library. Additional options for stub generation are covered in :ref:`sec-stubgen`. Function ``tvm_ffi_install`` places necessary information (e.g., debug symbols on macOS) next to the shared library for packaging. Python Build Backend ~~~~~~~~~~~~~~~~~~~~ Define a :pep:`517` build backend in ``pyproject.toml`` with the following steps: - Specify ``apache-tvm-ffi`` as a build requirement, so that CMake can find TVM-FFI; - Configure ``wheel.py-api`` that indicates a Python ABI-agnostic wheel; - Specify the source directory of the package via ``wheel.packages``, and the installation destination via ``wheel.install-dir``. .. literalinclude:: ../../examples/python_packaging/pyproject.toml :language: toml :start-after: [pyproject.build.begin] :end-before: [pyproject.build.end] Once specified, scikit-build-core will invoke CMake and drive the extension build. Wheel Auditing ~~~~~~~~~~~~~~ **Build wheels**. You can build wheels using standard workflows, for example: - `pip workflow `_ or `editable install `_ .. code-block:: bash # editable install pip install -e . # standard wheel build pip wheel -w dist . - `uv workflow `_ .. code-block:: bash uv build --wheel --out-dir dist . - `cibuildwheel `_ for multi-platform build .. code-block:: bash cibuildwheel --output-dir dist **Audit wheels**. In practice, an extra step is usually needed to remove redundant and error-prone shared library dependencies. In our case, because ``libtvm_ffi.so`` (or its platform variants) is guaranteed to be loaded by importing ``tvm_ffi``, we can safely exclude this dependency from the final wheel. .. code-block:: bash # Linux auditwheel repair --exclude libtvm_ffi.so dist/*.whl # macOS delocate-wheel -w dist -v --exclude libtvm_ffi.dylib dist/*.whl # Windows delvewheel repair --exclude tvm_ffi.dll -w dist dist\\*.whl .. _sec-load-the-library: Load the Library ~~~~~~~~~~~~~~~~ Once the wheel is installed, use :py:func:`tvm_ffi.libinfo.load_lib_module` to load the shared library: .. code-block:: python from tvm_ffi.libinfo import load_lib_module LIB = load_lib_module( package="my-ffi-extension", target_name="my_ffi_extension", ) The parameters are: - ``package``: The Python package name as registered with pip (e.g., ``"my-ffi-extension"`` or ``"apache-tvm-ffi"``). This is the name in ``pyproject.toml``, **not** the import name (e.g., ``tvm_ffi``). The function uses ``importlib.metadata.distribution(package)`` internally to locate installed package files. - ``target_name``: The CMake target name (e.g., ``"my_ffi_extension"``). It is used to derive the platform-specific shared library filename: * Linux: ``lib{target_name}.so`` * macOS: ``lib{target_name}.dylib`` * Windows: ``{target_name}.dll`` Once the library is loaded, functions and classes can be exported from C++ and called from Python. See :doc:`../guides/export_func_cls` for the three export mechanisms (C symbols, global functions, and classes) with complete examples. Stub Generation --------------- TVM-FFI provides a stub generation tool ``tvm-ffi-stubgen`` that creates Python type hints from C++ reflection metadata. The tool integrates with CMake and can generate complete stub files automatically, or update existing files using special directive comments. For most projects, enable automatic stub generation in CMake: .. code-block:: cmake tvm_ffi_configure_target(my_ffi_extension STUB_DIR "python" STUB_INIT ON ) This generates ``_ffi_api.py`` and ``__init__.py`` files with proper type hints for all registered global functions and classes. .. seealso:: - :doc:`stubgen`: Complete stub generation guide, including directive-based customization and command-line usage. - :doc:`../dev/ci_cd`: Reproducing wheel builds locally with cibuildwheel. tvm-ffi-0.1.12/docs/packaging/stubgen.rst000066400000000000000000000307021521067262500202300ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. .. _sec-stubgen: Stub Generation =============== TVM-FFI provides ``tvm-ffi-stubgen``, a tool that generates Python type stubs from C++ reflection metadata. It turns registered global functions and classes into proper Python type hints, enabling IDE auto-completion and static type checking. .. admonition:: Prerequisite :class: hint This tutorial uses `examples/python_packaging `_ as a running example. The package registers global functions (``my_ffi_extension.add_one``, ``my_ffi_extension.raise_error``) and an object type (``my_ffi_extension.IntPair``). .. _sec-stubgen-cmake: CMake-based Generation ---------------------- The recommended approach is to use CMake's ``tvm_ffi_configure_target`` function. This runs stub generation automatically after each build. .. code-block:: cmake tvm_ffi_configure_target( STUB_DIR [STUB_INIT ON|OFF] [STUB_PKG ] [STUB_PREFIX ] ) From the example's `CMakeLists.txt `_: .. code-block:: cmake add_library(my_ffi_extension SHARED src/extension.cc) tvm_ffi_configure_target(my_ffi_extension STUB_DIR "./python" STUB_INIT ON) STUB_DIR Required. Directory containing Python source files, relative to ``CMAKE_CURRENT_SOURCE_DIR``. STUB_INIT (default: ``OFF``) Optional. When ``OFF``, fills in existing inline directives without creating new ones. When ``ON``, generates new files and directives. ``STUB_INIT`` is ``OFF`` ~~~~~~~~~~~~~~~~~~~~~~~~ The tool fills in existing inline directives only. It scans files for directive markers and regenerates the content within those blocks: - Class fields and methods (``object/``) .. code-block:: python @tvm_ffi.register_object("my_ffi_extension.IntPair") class IntPair(tvm_ffi.Object): # tvm-ffi-stubgen(begin): object/my_ffi_extension.IntPair a: int b: int if TYPE_CHECKING: def __init__(self, a: int, b: int) -> None: ... def __ffi_init__(self, a: int, b: int) -> None: ... def sum(self, /) -> int: ... # tvm-ffi-stubgen(end) - Global functions (``global/``) .. code-block:: python # tvm-ffi-stubgen(begin): global/my_ffi_extension tvm_ffi.init_ffi_api("my_ffi_extension", __name__) if TYPE_CHECKING: def add_one(_0: int, /) -> int: ... def raise_error(_0: str, /) -> None: ... # tvm-ffi-stubgen(end) - Import sections, ``__all__`` lists, and exports (``import-section``, ``__all__``, ``export/``) See :ref:`sec-stubgen-directives` for a complete reference of all inline directives. .. important:: No new files or directives are created. Use this mode after you have customized the generated files and want to preserve your changes. ``STUB_INIT`` is ``ON`` ~~~~~~~~~~~~~~~~~~~~~~~ The tool generates new files (``_ffi_api.py``, ``__init__.py``) and new directives from scratch. This is the recommended starting point for new projects. Two additional options are available: STUB_PKG (default: ``SKBUILD_PROJECT_NAME`` or CMake target name) Python package name. Determines the directory structure (``//``) and generates the library loading code in ``_ffi_api.py``: .. code-block:: python LIB = tvm_ffi.libinfo.load_lib_module("", "") This uses :py:func:`tvm_ffi.libinfo.load_lib_module` to load the shared library. STUB_PREFIX (default: ``.``) Registry prefix filter. Only functions and classes whose registered names start with this prefix are included. The tool generates: - ``global/`` directives for global functions. In ``//_ffi_api.py``: .. code-block:: python # tvm-ffi-stubgen(begin): global/ ... # tvm-ffi-stubgen(end) Functions in nested namespaces (e.g., ``submodule.func``) are placed in corresponding subdirectories. - ``object/`` directives for each registered class. For a class with type key ``IntPair``: .. code-block:: python @tvm_ffi.register_object("IntPair") class IntPair(_ffi_Object): # tvm-ffi-stubgen(begin): object/IntPair a: int b: int if TYPE_CHECKING: def __init__(self, a: int, b: int) -> None: ... def __ffi_init__(self, a: int, b: int) -> None: ... def sum(self, /) -> int: ... # tvm-ffi-stubgen(end) Classes in nested namespaces are similarly placed in corresponding subdirectories. .. _sec-stubgen-cli: CLI-based Generation -------------------- For standalone usage or custom build systems, use the command-line tool directly. The CLI options correspond to CMake options as follows: .. list-table:: :header-rows: 1 :widths: 30 30 40 * - CLI Option - CMake Option - Description * - ```` - ``STUB_DIR`` - Directory containing Python source files * - ``--init-*`` flags - ``STUB_INIT ON`` - Presence of ``--init-*`` flags enables full generation * - ``--init-pypkg`` - ``STUB_PKG`` - Python package name * - ``--init-prefix`` - ``STUB_PREFIX`` - Registry prefix filter * - ``--init-lib`` - *(target name)* - CMake target name for the shared library Basic Usage ~~~~~~~~~~~ To generate complete stub files from scratch (equivalent to ``STUB_INIT ON``): .. code-block:: bash tvm-ffi-stubgen \ python/my_ffi_extension \ --dlls build/libmy_ffi_extension.so \ --init-pypkg my-ffi-extension \ --init-lib my_ffi_extension \ --init-prefix "my_ffi_extension." To update only existing directives (equivalent to ``STUB_INIT OFF``), omit the ``--init-*`` flags: .. code-block:: bash tvm-ffi-stubgen \ python/my_ffi_extension \ --dlls build/libmy_ffi_extension.so Command Line Options ~~~~~~~~~~~~~~~~~~~~ **Required arguments:** ```` (positional) Directory to generate/update Python stubs. Corresponds to ``STUB_DIR``. ``--dlls`` Shared libraries to load during generation. The stub generator loads these DLLs to discover registered global functions and object types. Without loading the DLLs, the generator cannot know what functions and classes exist. Use semicolons to separate multiple libraries. **Arguments for full generation** (corresponds to ``STUB_INIT ON``): All three are required together. When omitted, the tool operates in directive-only mode (equivalent to ``STUB_INIT OFF``). ``--init-pypkg`` Python package name. Corresponds to ``STUB_PKG``. ``--init-lib`` CMake target name. Used as the second argument in the generated ``load_lib_module("", "")`` call. ``--init-prefix`` Registry prefix filter. Corresponds to ``STUB_PREFIX``. **Optional arguments:** ``--verbose`` Print a unified diff of changes to each file. ``--dry-run`` Preview changes without writing to files. ``--imports`` Additional Python modules to import before generation (semicolon-separated). ``--indent`` Indentation width for generated code (default: 4). For a complete list of options, run ``tvm-ffi-stubgen --help``. .. _sec-stubgen-advanced: Advanced Topics --------------- .. _sec-stubgen-directives: Inline Directives ~~~~~~~~~~~~~~~~~ Inline directives are special comments that mark regions where the tool should generate or update code. They allow you to: - Add custom code alongside generated stubs - Control the exact placement of generated code - Maintain hand-written documentation or additional methods - Integrate with existing codebases The tool preserves content outside directive blocks, only modifying marked regions. **Directive Syntax** Directives use comment markers to delimit generated regions: .. code-block:: python # tvm-ffi-stubgen(begin): / # ... generated code appears here ... # tvm-ffi-stubgen(end) When you run the tool, it: 1. Parses files for directive markers 2. Regenerates content between ``begin`` and ``end`` markers 3. Preserves everything outside the markers **Directive Reference** ``global/`` - Global Functions Marks a section for global function stubs. All functions registered with the given prefix are included. .. code-block:: python # tvm-ffi-stubgen(begin): global/my_ext.arith # fmt: off tvm_ffi.init_ffi_api("my_ext.arith", __name__) if TYPE_CHECKING: def add_one(_0: int, /) -> int: ... def add_two(_0: int, /) -> int: ... # fmt: on # tvm-ffi-stubgen(end) ``object/`` - Object Types Marks a section for class field and method stubs. Place inside a class definition decorated with ``@tvm_ffi.register_object``. .. code-block:: python @tvm_ffi.register_object("my_ffi_extension.IntPair") class IntPair(_ffi_Object): # tvm-ffi-stubgen(begin): object/my_ffi_extension.IntPair # fmt: off a: int b: int if TYPE_CHECKING: def __init__(self, a: int, b: int) -> None: ... def __ffi_init__(self, a: int, b: int) -> None: ... def sum(self, /) -> int: ... # fmt: on # tvm-ffi-stubgen(end) ``import-section`` - Import Section Populates import statements for types used in generated stubs. The tool automatically determines which imports are needed. .. code-block:: python # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from tvm_ffi import init_ffi_api as _FFI_INIT_FUNC from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Mapping, Sequence from tvm_ffi import Device, Object, Tensor, dtype # isort: on # fmt: on # tvm-ffi-stubgen(end) ``export/`` - Export Re-exports names from a submodule, typically used in ``__init__.py`` to aggregate exports. .. code-block:: python # tvm-ffi-stubgen(begin): export/_ffi_api # fmt: off # isort: off from ._ffi_api import * # noqa: F403 from ._ffi_api import __all__ as _ffi_api__all__ if "__all__" not in globals(): __all__ = [] __all__.extend(_ffi_api__all__) # isort: on # fmt: on # tvm-ffi-stubgen(end) ``__all__`` - Export List Populates the ``__all__`` list with all generated class and function names. .. code-block:: python __all__ = [ # tvm-ffi-stubgen(begin): __all__ "LIB", "IntPair", "add_one", # tvm-ffi-stubgen(end) ] ``ty-map`` - Type Mapping Maps a C++ type key to a different Python type path. This is useful when the Python class is defined in a different module than the type key suggests. .. code-block:: python # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep ``import-object`` - Import Object Injects a custom import into generated code. The format is ``;;``. .. code-block:: python # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object This imports ``ffi.Object`` as ``_ffi_Object`` for use in generated code. The second field (``False``) indicates the import is not TYPE_CHECKING-only. ``skip-file`` - Skip File Prevents the tool from modifying the file. Place anywhere in the file. .. code-block:: python # tvm-ffi-stubgen(skip-file) tvm-ffi-0.1.12/docs/reference/000077500000000000000000000000001521067262500160375ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/reference/cpp/000077500000000000000000000000001521067262500166215ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/reference/cpp/index.rst000066400000000000000000000070241521067262500204650ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. C++ API ======= This page contains the API reference for the C++ API. The full API index below can be a bit dense, so we recommend the following tips first: - Please read the :ref:`C++ Guide` for a high-level overview of the C++ API. - The C++ Guide and examples will likely be sufficient to get started with most use cases. - The :ref:`cpp-key-classes` lists the key classes that are most commonly used. - You can go to the Full API Index at the bottom of this page to access the full list of APIs. - We usually group the APIs by files. You can look at the file hierarchy in the full API index and navigate to the specific file to find the APIs in that file. Header Organization ------------------- The C++ APIs are organized into the following folders: .. list-table:: :header-rows: 1 :widths: 30 70 * - Folder - Description * - ``tvm/ffi/`` - Core functionalities that support Function, Any, Object, etc. * - ``tvm/ffi/container/`` - Additional container types such as Array, Map, Shape, Tensor, Variant ... * - ``tvm/ffi/reflection/`` - Reflection support for function and type information registration. * - ``tvm/ffi/extra/`` - Extra APIs that are built on top. .. _cpp-key-classes: Key Classes ----------- .. list-table:: :header-rows: 1 :widths: 30 70 * - Class - Description * - :cpp:class:`tvm::ffi::Function` - Type-erased function that implements the ABI. * - :cpp:class:`tvm::ffi::Any` - Type-erased container for any supported value. * - :cpp:class:`tvm::ffi::AnyView` - Lightweight view of Any without ownership. * - :cpp:class:`tvm::ffi::Object` - Base class for all heap-allocated FFI objects. * - :cpp:class:`tvm::ffi::ObjectRef` - Reference class for objects. * - :cpp:class:`tvm::ffi::Tensor` - Multi-dimensional tensor with DLPack support. * - :cpp:class:`tvm::ffi::TensorView` - Multi-dimensional tensor view with DLPack support. It does not own the underlying data. * - :cpp:class:`tvm::ffi::Shape` - Tensor shape container. * - :cpp:class:`tvm::ffi::Module` - Dynamic library module that can load exported functions. * - :cpp:class:`tvm::ffi::String` - String type for FFI. * - :cpp:class:`tvm::ffi::Bytes` - Byte array type. * - :cpp:class:`tvm::ffi::Array` - Dynamic array container. * - :cpp:class:`tvm::ffi::Tuple` - Heterogeneous tuple container. * - :cpp:class:`tvm::ffi::Map` - Key-value map container. * - :cpp:class:`tvm::ffi::Optional` - Optional value wrapper. * - :cpp:class:`tvm::ffi::Variant` - Type-safe union container. .. _cpp-full-api-index: Full API Index -------------- `Browse the full API index `__ tvm-ffi-0.1.12/docs/reference/python/000077500000000000000000000000001521067262500173605ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/reference/python/index.rst000066400000000000000000000046641521067262500212330ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Python API ========== .. automodule:: tvm_ffi :no-members: .. currentmodule:: tvm_ffi .. contents:: Table of Contents :local: :depth: 2 Object ------ .. autosummary:: :toctree: generated/ Object Tensor ~~~~~~ .. autosummary:: :toctree: generated/ Tensor from_dlpack Shape dtype Device DLDeviceType device Function ~~~~~~~~ .. autosummary:: :toctree: generated/ Function Module ~~~~~~ .. autosummary:: :toctree: generated/ Module system_lib load_module Containers ~~~~~~~~~~ .. autosummary:: :toctree: generated/ Array Map Global Registry --------------- .. autosummary:: :toctree: generated/ register_error register_object register_global_func get_global_func get_global_func_metadata init_ffi_api remove_global_func Stream Context -------------- .. autosummary:: :toctree: generated/ StreamContext use_torch_stream use_raw_stream get_raw_stream C++ Extension ------------- C++ integration helpers for building and loading inline modules. .. autosummary:: :toctree: cpp/generated/ cpp.load_inline cpp.build_inline cpp.load cpp.build libinfo.load_lib_module NVRTC Utilities --------------- .. autosummary:: :toctree: cpp/generated/ cpp.nvrtc.nvrtc_compile Misc ---- .. autosummary:: :toctree: generated/ serialization.from_json_graph_str serialization.to_json_graph_str access_path.AccessKind access_path.AccessPath access_path.AccessStep convert convert_func ObjectConvertible .. (Experimental) Dataclasses .. -------------------------- .. .. autosummary:: .. :toctree: generated/ .. dataclasses.c_class .. dataclasses.field tvm-ffi-0.1.12/docs/reference/rust/000077500000000000000000000000001521067262500170345ustar00rootroot00000000000000tvm-ffi-0.1.12/docs/reference/rust/index.rst000066400000000000000000000030011521067262500206670ustar00rootroot00000000000000.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Rust API ======== This page contains the API reference for the Rust API. The tvm-ffi project provides Rust bindings that allow you to use the TVM FFI from Rust programs. Crates ------ The Rust API is organized into three crates: .. list-table:: :header-rows: 1 :widths: 30 70 * - Crate - Description * - 📦 `tvm-ffi `_ - High-level Rust bindings with safe abstractions for FFI objects, functions, tensors, and containers. * - 📦 `tvm-ffi-sys `_ - Low-level unsafe bindings to the C API. * - 📦 `tvm-ffi-macros `_ - Procedural macros for deriving Object and ObjectRef traits. tvm-ffi-0.1.12/examples/000077500000000000000000000000001521067262500147675ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/abi_overview/000077500000000000000000000000001521067262500174505ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/abi_overview/example_code.c000066400000000000000000000214611521067262500222450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * Example code for TVM-FFI ABI overview. * * Compilation command: * * ```bash * gcc $(tvm-ffi-config --cflags) \ * $(tvm-ffi-config --ldflags) \ * $(tvm-ffi-config --libs) \ * -Wl,-rpath,$(tvm-ffi-config --libdir) \ * -o ./example_code * ``` */ // NOLINTBEGIN(modernize-deprecated-headers,modernize-use-nullptr) #include #include #include #include #include #include int IS_OWNING_ANY = 1; // [Any_AnyView.FromInt_Float.begin] TVMFFIAny Any_AnyView_FromInt(int64_t value) { TVMFFIAny any; any.type_index = kTVMFFIInt; any.zero_padding = 0; any.v_int64 = value; return any; } TVMFFIAny Any_AnyView_FromFloat(double value) { TVMFFIAny any; any.type_index = kTVMFFIFloat; any.zero_padding = 0; any.v_float64 = value; return any; } // [Any_AnyView.FromInt_Float.end] // [Any_AnyView.FromObjectPtr.begin] TVMFFIAny Any_AnyView_FromObjectPtr(TVMFFIObject* obj) { TVMFFIAny any; assert(obj != NULL); any.type_index = kTVMFFIObject; any.zero_padding = 0; any.v_obj = obj; // Increment refcount if it's Any (owning) instead of AnyView (borrowing) if (IS_OWNING_ANY) { TVMFFIObjectIncRef(obj); } return any; } // [Any_AnyView.FromObjectPtr.end] // [Any_AnyView.Destroy.begin] void Any_AnyView_Destroy(TVMFFIAny* any) { if (IS_OWNING_ANY) { // Checks if `any` holds a heap-allocated object, // and if so, decrements the reference count if (any->type_index >= kTVMFFIStaticObjectBegin) { TVMFFIObjectDecRef(any->v_obj); } } *any = (TVMFFIAny){0}; // Clears the `any` struct } // [Any_AnyView.Destroy.end] // [Any_AnyView.GetInt_Float.begin] int64_t Any_AnyView_GetInt(const TVMFFIAny* any) { if (any->type_index == kTVMFFIInt || any->type_index == kTVMFFIBool) { return any->v_int64; } else if (any->type_index == kTVMFFIFloat) { return (int64_t)(any->v_float64); } assert(0); // FAILED to read int return 0; } double Any_AnyView_GetFloat(const TVMFFIAny* any) { if (any->type_index == kTVMFFIInt || any->type_index == kTVMFFIBool) { return (double)(any->v_int64); } else if (any->type_index == kTVMFFIFloat) { return any->v_float64; } assert(0); // FAILED to read float return 0.0; } // [Any_AnyView.GetInt_Float.end] // [Any_AnyView.GetDLTensor.begin] DLTensor* Any_AnyView_GetDLTensor(const TVMFFIAny* value) { if (value->type_index == kTVMFFIDLTensorPtr) { return (DLTensor*)(value->v_ptr); } else if (value->type_index == kTVMFFITensor) { return (DLTensor*)((char*)(value->v_obj) + sizeof(TVMFFIObject)); } assert(0); // FAILED to read DLTensor return NULL; } // [Any_AnyView.GetDLTensor.end] // [Any_AnyView.GetObject.begin] TVMFFIObject* Any_AnyView_GetObject(const TVMFFIAny* value) { if (value->type_index == kTVMFFINone) { return NULL; // Handling nullptr if needed } else if (value->type_index >= kTVMFFIStaticObjectBegin) { return value->v_obj; } assert(0); // FAILED: not a TVM-FFI object return NULL; } // [Any_AnyView.GetObject.end] // [Object.IsInstance.begin] int Object_IsInstance(int32_t sub_type_index, int32_t super_type_index, int32_t super_type_depth) { const TVMFFITypeInfo* sub_type_info = NULL; // Everything is a subclass of object. if (sub_type_index == super_type_index) { return 1; } // Invariance: parent index is always smaller than the child. if (sub_type_index < super_type_index) { return 0; } sub_type_info = TVMFFIGetTypeInfo(sub_type_index); return sub_type_info->type_depth > super_type_depth && sub_type_info->type_ancestors[super_type_depth]->type_index == super_type_index; } // [Object.IsInstance.end] // [Object.MoveFromAny.begin] void Object_MoveFromAny(TVMFFIAny* any, TVMFFIObject** obj) { assert(any->type_index >= kTVMFFIStaticObjectBegin); *obj = any->v_obj; (*any) = (TVMFFIAny){0}; if (!IS_OWNING_ANY) { TVMFFIObjectIncRef(*obj); } } // [Object.MoveFromAny.end] // [Object.Destroy.begin] void Object_Destroy(TVMFFIObject* obj) { assert(obj != NULL); TVMFFIObjectDecRef(obj); } // [Object.Destroy.end] // [Tensor.AccessDLTensor.begin] DLTensor* Tensor_AccessDLTensor(TVMFFIObject* tensor) { assert(tensor != NULL); return (DLTensor*)((char*)tensor + sizeof(TVMFFIObject)); } // [Tensor.AccessDLTensor.end] // [Tensor_FromDLPack.begin] TVMFFIObject* Tensor_FromDLPack(DLManagedTensorVersioned* from) { int err_code = 0; TVMFFIObject* out = NULL; err_code = TVMFFITensorFromDLPackVersioned( // from, // input DLPack tensor /*require_alignment=*/0, // no alignment requirement /*require_contiguous=*/1, // require contiguous tensor (void**)(&out)); assert(err_code == 0); return out; } // [Tensor_FromDLPack.end] // [Tensor_Alloc.begin] TVMFFIObject* Tensor_Alloc(DLTensor* prototype) { int err_code = 0; TVMFFIObject* out = NULL; assert(prototype->data == NULL); err_code = TVMFFIEnvTensorAlloc(prototype, (void**)(&out)); assert(err_code == 0); return out; } // [Tensor_Alloc.end] // [Tensor_ToDLPackVersioned.begin] DLManagedTensorVersioned* Tensor_ToDLPackVersioned(TVMFFIObject* tensor) { int err_code = 0; DLManagedTensorVersioned* out = NULL; err_code = TVMFFITensorToDLPackVersioned(tensor, &out); assert(err_code == 0); return out; } // [Tensor_ToDLPackVersioned.end] // [Function.Construct.begin] TVMFFIObject* Function_Construct(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void* self)) { int err_code; TVMFFIObject* out = NULL; err_code = TVMFFIFunctionCreate(self, safe_call, deleter, (void**)(&out)); assert(err_code == 0); return out; } // [Function.Construct.end] // [Function.Call.begin] int64_t CallFunction(TVMFFIObject* func, int64_t x, int64_t y) { int err_code; TVMFFIAny args[2]; TVMFFIAny result = (TVMFFIAny){0}; args[0] = Any_AnyView_FromInt(x); args[1] = Any_AnyView_FromInt(y); err_code = TVMFFIFunctionCall(func, args, 2, &result); assert(err_code == 0); return Any_AnyView_GetInt(&result); } // [Function.Call.end] // [Function.GetGlobal.begin] TVMFFIObject* Function_RetrieveGlobal(const char* name) { TVMFFIObject* out = NULL; TVMFFIByteArray name_byte_array = {name, strlen(name)}; int err_code = TVMFFIFunctionGetGlobal(&name_byte_array, (void**)(&out)); assert(err_code == 0); return out; } // [Function.GetGlobal.end] // [Function.SetGlobal.begin] void Function_SetGlobal(const char* name, TVMFFIObject* func) { TVMFFIByteArray name_byte_array = {name, strlen(name)}; int err_code = TVMFFIFunctionSetGlobal(&name_byte_array, func, 0); assert(err_code == 0); } // [Function.SetGlobal.end] // [Error.Print.begin] void PrintError(TVMFFIObject* err) { TVMFFIErrorCell* cell = (TVMFFIErrorCell*)((char*)err + sizeof(TVMFFIObject)); fprintf(stderr, "%.*s: %.*s\n", // (int)cell->kind.size, cell->kind.data, // e.g. "ValueError" (int)cell->message.size, cell->message.data); // e.g. "Expected at least 2 arguments" if (cell->backtrace.size) { fprintf(stderr, "Backtrace:\n%.*s\n", (int)cell->backtrace.size, cell->backtrace.data); } } // [Error.Print.end] // [Error.HandleReturnCode.begin] void Error_HandleReturnCode(int rc) { TVMFFIObject* err = NULL; if (rc == -1) { // Move the raised error from TLS (clears TLS slot) TVMFFIErrorMoveFromRaised((void**)(&err)); // now `err` owns the error object if (err != NULL) { PrintError(err); // print the error // IMPORTANT: Release the error object, or gets memory leaks TVMFFIObjectDecRef(err); } } } // [Error.HandleReturnCode.end] // [Error.RaiseException.begin] int Error_RaiseException(void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { TVMFFIErrorSetRaisedFromCStr("ValueError", "Expected at least 2 arguments"); return -1; } // [Error.RaiseException.end] // NOLINTEND(modernize-deprecated-headers,modernize-use-nullptr) int main() { return 0; } tvm-ffi-0.1.12/examples/cubin_launcher/000077500000000000000000000000001521067262500177505ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/README.md000066400000000000000000000150501521067262500212300ustar00rootroot00000000000000 # CUBIN Launcher ## Overview Demonstrates loading and executing CUDA kernels from CUBIN files using TVM-FFI. The `cubin_launcher.h` header wraps CUDA Runtime API to provide lightweight CUBIN module and kernel management. ## Techniques The implementation supports both CUDA Runtime API (CUDA >= 12.8) and Driver API for Library Management. **Runtime API (CUDA >= 12.8):** - **`cudaLibraryLoadData()`** - Load CUBIN from memory buffer - **`cudaLibraryGetKernel()`** - Get kernel handle by name - **`cudaLaunchKernel()`** - Launch kernel with grid/block dimensions **Driver API:** - **`cuLibraryLoadData()`** - Load CUBIN from memory buffer - **`cuLibraryGetKernel()`** - Get kernel handle by name - **`cuLaunchKernel()`** - Launch kernel with grid/block dimensions **Customization:** By default, the implementation uses the Runtime API if compiled with CUDA >= 12.8, falling back to the Driver API for older versions. You can force the usage of the Driver API (or Runtime API) by defining the macro `TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API` (set to `1` for Driver API, `0` for Runtime API) before including the header. Key features: - Multi-GPU support via CUDA primary contexts - RAII-based resource management (CubinModule, CubinKernel) - CUBIN embedding at compile time - Object linking (via `ld` + `objcopy`) - Header inclusion (via `bin2c`) - Modern C++ embedding (via `#embed`) - TVM-FFI integration for tensor argument passing - **Macros:** - `TVM_FFI_EMBED_CUBIN`: Declare symbols for object-linked CUBIN - `TVM_FFI_EMBED_CUBIN_FROM_BYTES`: Load CUBIN from byte array (for `#embed`/`bin2c`) - `TVM_FFI_EMBED_CUBIN_GET_KERNEL`: Helper to retrieve kernels - **Python Integration:** `embed_cubin` parameter in `tvm_ffi.cpp.load_inline` for seamless CUBIN integration - **Runtime Compilation:** `tvm_ffi.cpp.nvrtc` module for runtime CUDA compilation ## Examples ### 1. Embedded CUBIN The `embedded_cubin` directory contains three examples demonstrating different embedding techniques. #### 1.1 Object Linking (Standard) Demonstrates embedding CUBIN data directly into the shared library at build time using the `tvm_ffi_embed_bin_into` CMake utility. This is the most robust method for CMake projects. **Location:** `embedded_cubin/embed_with_tvm_ffi/` **Build and run:** ```bash cd examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi mkdir build && cd build cmake .. make cd .. python main.py ``` #### 1.2 Header Inclusion (Portable) Demonstrates converting the CUBIN to a C header file using `bin2c` and including it in the C++ source. This is highly portable and works with any compiler. **Location:** `embedded_cubin/include_bin2c/` **Build and run:** ```bash cd examples/cubin_launcher/embedded_cubin/include_bin2c mkdir build && cd build cmake .. make cd .. python main.py ``` #### 1.3 C++ Embedding (Modern) Demonstrates using C++23 `#embed` (or compiler extensions in GCC/Clang) to directly include binary data. This is the cleanest approach for modern toolchains. **Location:** `embedded_cubin/cpp_embed/` **Build and run:** ```bash cd examples/cubin_launcher/embedded_cubin/cpp_embed mkdir build && cd build cmake .. make cd .. python main.py ``` ### 2. Dynamic CUBIN Loading Demonstrates loading CUBIN data from a file at runtime using the CUDA Driver API. **Location:** `dynamic_cubin/` **Build and run:** ```bash cd examples/cubin_launcher/dynamic_cubin mkdir build && cd build cmake .. make cd .. python main.py ``` **Key features:** - CUBIN loaded from file at runtime - More flexible - can swap CUBIN files - Useful for JIT-compiled kernels ### 3. Triton Kernel with Embedded CUBIN (Experimental) `example_triton_cubin.py` - Triton kernel compiled to CUBIN and embedded inline using the `embed_cubin` parameter. ```bash # Requires: triton, torch python examples/cubin_launcher/example_triton_cubin.py ``` ### 4. NVRTC with Embedded CUBIN `example_nvrtc_cubin.py` - CUDA source compiled to CUBIN using NVRTC and embedded inline. ```bash # Requires: cuda-python, torch python examples/cubin_launcher/example_nvrtc_cubin.py ``` ## Using Embedded CUBIN with `tvm_ffi.cpp.load_inline` The new `embed_cubin` parameter makes it easy to embed CUBIN binaries into your module: ```python from tvm_ffi import cpp from tvm_ffi.cpp import nvrtc # Compile CUDA source to CUBIN cuda_source = """ extern "C" __global__ void my_kernel(float* data, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) data[idx] *= 2.0f; } """ cubin_bytes = nvrtc.nvrtc_compile(cuda_source) # C++ code using the embedded CUBIN cpp_code = """ #include TVM_FFI_EMBED_CUBIN(my_module); void launch_kernel(TensorView data) { static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_module, "my_kernel"); // ... launch kernel } """ # Load with embedded CUBIN mod = cpp.load_inline( "my_module", cpp_sources=cpp_code, embed_cubin={"my_module": cubin_bytes}, extra_ldflags=["-lcudart"], ) ``` ## Project Structure ### Core Files - `include/tvm/ffi/extra/cuda/cubin_launcher.h` - Header-only C++ library with CUBIN utilities - `python/tvm_ffi/utils/embed_cubin.py` - Python utility for embedding CUBIN into object files - `python/tvm_ffi/cpp/nvrtc.py` - NVRTC compilation utilities - `cmake/Utils/EmbedCubin.cmake` - CMake utilities ### Example Directories **`embedded_cubin/`** - Different CUBIN embedding techniques: - `embed_with_tvm_ffi/` - Standard object linking - `include_bin2c/` - Header inclusion - `cpp_embed/` - Modern C++ `#embed` **`dynamic_cubin/`** - CUBIN loaded at runtime **Additional Examples** (at root level) - `example_triton_cubin.py` - Triton kernel with embedded CUBIN - `example_nvrtc_cubin.py` - NVRTC compilation with embedded CUBIN tvm-ffi-0.1.12/examples/cubin_launcher/benchmark_overhead.py000066400000000000000000000212021521067262500241260ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Benchmark kernel launch overhead: Triton vs TVM-FFI. This script compares the launch overhead between: 1. Triton's native kernel launcher 2. TVM-FFI's CUBIN launcher Both launch the same empty kernel that does nothing, allowing us to measure pure launch overhead without compute time interference. Notes: - Requires `triton` to be installed in the Python environment. """ from __future__ import annotations import platform import sys import time import traceback from typing import Callable import torch import triton import triton.language as tl from tvm_ffi import cpp from tvm_ffi.module import Module def get_cpu_name() -> str: """Get the name of the CPU.""" cpu_name = platform.processor() if not cpu_name or cpu_name == "x86_64": # Fallback: Try to read /proc/cpuinfo for better model string on Linux try: with open("/proc/cpuinfo") as f: # noqa: PTH123 for line in f: if "model name" in line: cpu_name = line.strip().split(":", 1)[1].strip() break except Exception: cpu_name = "Unknown CPU" return cpu_name # Define empty kernel at global scope @triton.jit def empty_kernel(A_ptr, B_ptr, C_ptr, n, BLOCK: tl.constexpr = 128): # noqa # ty: ignore[invalid-parameter-default] """Empty kernel that does nothing - for measuring pure launch overhead.""" pass def print_speed(name: str, time_per_call: float) -> None: """Print benchmark result in a formatted way. Parameters ---------- name : str Name of the benchmark time_per_call : float Time per call in seconds """ time_us = time_per_call * 1e6 print(f" {name:30s}: {time_us:8.3f} μs/call") def benchmark_call(name: str, func: Callable, args: tuple, num_calls: int = 10000) -> float: """Benchmark a function by calling it multiple times. Parameters ---------- name : str Name of the benchmark func : callable Function to benchmark args : tuple Arguments to pass to the function num_calls : int Number of calls to average over Returns ------- float Time per call in seconds """ # Warmup func(*args) torch.cuda.synchronize() # Benchmark start_time = time.time() for _ in range(num_calls): func(*args) torch.cuda.synchronize() end_time = time.time() time_per_call = (end_time - start_time) / num_calls return time_per_call def generate_cubin() -> bytes: """Compile the empty kernel to CUBIN. Returns ------- bytes Compiled CUBIN bytes """ # Trigger kernel compilation by doing a dummy call n = 128 a_dummy = torch.empty(n, dtype=torch.float32, device="cuda") b_dummy = torch.empty(n, dtype=torch.float32, device="cuda") c_dummy = torch.empty(n, dtype=torch.float32, device="cuda") compiled_kernel = empty_kernel[1,](a_dummy, b_dummy, c_dummy, n) # Get CUBIN bytes cubin_bytes = compiled_kernel.kernel return cubin_bytes def load_cubin_module(cubin_bytes: bytes) -> Module: """Load CUBIN kernel through TVM-FFI. Parameters ---------- cubin_bytes : bytes Compiled CUBIN bytes Returns ------- Module TVM-FFI module with launch function """ # Define C++ code inline to load and launch the empty kernel using embedded CUBIN sources = """ #include #include #include #include #include // Embed CUBIN module with name "empty_cubin" TVM_FFI_EMBED_CUBIN(empty_cubin); namespace empty_loader { void LaunchEmpty(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView c) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(empty_cubin, "empty_kernel"); TVM_FFI_CHECK(a.ndim() == 1, ValueError) << "Input a must be 1D tensor"; TVM_FFI_CHECK(b.ndim() == 1, ValueError) << "Input b must be 1D tensor"; TVM_FFI_CHECK(c.ndim() == 1, ValueError) << "Input c must be 1D tensor"; uint32_t n = static_cast(a.size(0)); void* a_ptr = a.data_ptr(); void* b_ptr = b.data_ptr(); void* c_ptr = c.data_ptr(); uint64_t dummy_ptr = 0; // Workaround for Triton extra params: pass dummy addresses for unused parameters void* args[] = {&a_ptr, &b_ptr, &c_ptr, &n, &dummy_ptr, &dummy_ptr}; // Use single thread block with 128 threads tvm::ffi::dim3 grid(1); tvm::ffi::dim3 block(128); DLDevice device = a.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); cudaError_t result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUDA_ERROR(result); } } // namespace empty_loader TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_empty, empty_loader::LaunchEmpty); """ mod = cpp.load_inline( "empty_loader", cuda_sources=sources, embed_cubin={"empty_cubin": cubin_bytes}, ) return mod def run_benchmark(cubin_bytes: bytes, num_calls: int = 10000) -> int: """Run overhead benchmarks for both Triton and TVM-FFI. Parameters ---------- cubin_bytes : bytes Compiled CUBIN bytes num_calls : int Number of calls to average over Returns ------- int 0 on success, non-zero on failure """ # Prepare test tensors n = 128 a = torch.empty(n, dtype=torch.float32, device="cuda") b = torch.empty(n, dtype=torch.float32, device="cuda") c = torch.empty(n, dtype=torch.float32, device="cuda") print(f"\nBenchmarking kernel launch overhead ({num_calls:,} calls)...") print("=" * 60) # Benchmark 1: Triton native launcher def triton_launch() -> None: empty_kernel[1,](a, b, c, n) triton_time = benchmark_call("Triton launch", triton_launch, (), num_calls) # Benchmark 2: TVM-FFI launcher mod = load_cubin_module(cubin_bytes) launch_fn = mod["launch_empty"] def tvm_ffi_launch() -> None: launch_fn(a, b, c) tvm_ffi_time = benchmark_call("TVM-FFI launch", tvm_ffi_launch, (), num_calls) # Summary print_speed("Triton launch", triton_time) print_speed("TVM-FFI launch", tvm_ffi_time) overhead_pct = ((tvm_ffi_time - triton_time) / triton_time) * 100 print(f"\n Overhead: {overhead_pct:+.2f}%") if tvm_ffi_time < triton_time: print(f" TVM-FFI is {triton_time / tvm_ffi_time:.2f}x faster") else: print(f" Triton is {tvm_ffi_time / triton_time:.2f}x faster") print("=" * 60) print( "Note: we did not check dtype/shape constraints in the benchmarks. \n" " Triton usually checks them in Python while TVM-FFI checks them in C++. \n" ) return 0 def main() -> int: """Main benchmark entry point.""" # noqa: D401 print("Kernel Launch Overhead Benchmark: Triton vs TVM-FFI") print("=" * 60) if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f" CPU: {get_cpu_name()}") print(f" CUDA device: {torch.cuda.get_device_name(0)}") print(f" PyTorch version: {torch.__version__}") print(f" Triton version: {triton.__version__}") # Compile empty kernel to CUBIN try: print("\nCompiling empty Triton kernel to CUBIN...") cubin_bytes = generate_cubin() print(f"Compiled CUBIN: {len(cubin_bytes):,} bytes") except Exception as e: print(f"[ERROR] Failed to compile Triton kernel: {e}") traceback.print_exc() return 2 # Run benchmarks try: return run_benchmark(cubin_bytes) except Exception as e: print(f"[ERROR] Failed to run benchmark: {e}") traceback.print_exc() return 3 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/000077500000000000000000000000001521067262500225545ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/CMakeLists.txt000066400000000000000000000055241521067262500253220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.20) project(dynamic_cubin_example LANGUAGES CXX CUDA) # Prefer virtualenv when searching for python set(Python_FIND_VIRTUALENV FIRST) # cmake-lint: disable=C0103 set(CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API OFF CACHE BOOL "Use driver API in cubin launcher" ) # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Find CUDA toolkit find_package(CUDAToolkit REQUIRED) # [cmake_example.begin] # Step 1: Compile kernel.cu to CUBIN using add_tvm_ffi_cubin utility or CUDA_CUBIN_COMPILATION. Use # CMAKE_CUDA_ARCHITECTURES=native to automatically detect the GPU architecture set(CMAKE_CUDA_ARCHITECTURES native) if (CMAKE_VERSION VERSION_LESS "3.27.0") add_tvm_ffi_cubin(kernel_cubin CUDA src/kernel.cu) else () add_library(kernel_cubin OBJECT src/kernel.cu) set_property(TARGET kernel_cubin PROPERTY CUDA_CUBIN_COMPILATION ON) endif () add_custom_target( kernel.cubin COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${CMAKE_CURRENT_BINARY_DIR}/kernel.cubin" DEPENDS kernel_cubin COMMENT "Copy cubin to build dir" ) # Step 2: Build lib_dynamic shared library (loads CUBIN from file at runtime) add_library(lib_dynamic SHARED src/lib_dynamic.cc) include_directories(${CUDAToolkit_INCLUDE_DIRS}) target_link_libraries(lib_dynamic PRIVATE tvm_ffi::header tvm_ffi::shared) add_dependencies(lib_dynamic kernel.cubin) set_target_properties( lib_dynamic PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/" PREFIX "" SUFFIX ".so" ) # Step 3: Link against CUDA Driver API or Runtime API based on config if (CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API) add_compile_definitions(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1) target_link_libraries(lib_dynamic PRIVATE cuda) else () target_link_libraries(lib_dynamic PRIVATE CUDA::cudart) endif () # [cmake_example.end] tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/main.py000066400000000000000000000075771521067262500240720ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example script for dynamic CUBIN loading. This example demonstrates using lib_dynamic.so which accepts CUBIN data dynamically at runtime. """ import sys from pathlib import Path import torch from tvm_ffi import load_module def main() -> int: # noqa: PLR0915 """Test the lib_dynamic.so library with dynamic CUBIN loading.""" print("=" * 60) print("Example: Dynamic CUBIN Loading") print("=" * 60) # Check CUDA availability if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Load the library lib_path = Path(__file__).parent / "build" / "lib_dynamic.so" mod = load_module(str(lib_path)) print(f"Loaded library: {lib_path}") # Read CUBIN file into memory cubin_path = Path(__file__).parent / "build" / "kernel.cubin" cubin_bytes = cubin_path.read_bytes() print(f"Read CUBIN from: {cubin_path} ({len(cubin_bytes)} bytes)") # Set CUBIN data to the module set_cubin = mod["set_cubin"] set_cubin(cubin_bytes) print("CUBIN module loaded successfully") # Get the kernel functions add_one = mod["add_one"] mul_two = mod["mul_two"] print("Loaded functions: add_one, mul_two") # Test add_one kernel print("\n[Test 1] add_one kernel") n = 2048 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.1 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") add_one(x, y) # Verify results expected = x + 1 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test mul_two kernel print("\n[Test 2] mul_two kernel") n = 1024 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.25 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") mul_two(x, y) # Verify results expected = x * 2 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test chained execution print("\n[Test 3] Chained execution: (x + 1) * 2") n = 512 x = torch.full((n,), 5.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Initial value: {x[0].item()}") add_one(x, temp) # temp = x + 1 = 6 mul_two(temp, y) # y = temp * 2 = 12 expected = 12.0 if torch.allclose(y, torch.tensor(expected, device="cuda")): print(f" [PASS] Result: {y[0].item()}") else: print(f" [FAIL] Expected {expected}, got {y[0].item()}") return 1 print("\n[PASS] All tests passed!") return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/src/000077500000000000000000000000001521067262500233435ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/src/kernel.cu000066400000000000000000000033221521067262500251540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/kernel.cu * \brief Simple CUDA kernel for testing cubin_launcher functionality. */ #include // [kernels.begin] /*! * \brief CUDA kernel that adds 1 to each element of an array. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void add_one_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } /*! * \brief CUDA kernel that multiplies each element by 2. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void mul_two_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } // [kernels.end] tvm-ffi-0.1.12/examples/cubin_launcher/dynamic_cubin/src/lib_dynamic.cc000066400000000000000000000115121521067262500261240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/lib_dynamic.cc * \brief TVM-FFI library with dynamic CUBIN loading. * * This library exports TVM-FFI functions to load CUBIN from file and * launch CUDA kernels. */ // [example.begin] #include #include #include #include #include #include #include #include namespace cubin_dynamic { // Global CUBIN module and kernels (loaded dynamically) static std::unique_ptr g_cubin_module; static std::unique_ptr g_add_one_kernel; static std::unique_ptr g_mul_two_kernel; /*! * \brief Set CUBIN module from binary data. * \param cubin CUBIN binary data as Bytes object. */ void SetCubin(const tvm::ffi::Bytes& cubin) { // Load CUBIN module from memory g_cubin_module = std::make_unique(cubin); g_add_one_kernel = std::make_unique((*g_cubin_module)["add_one_cuda"]); g_mul_two_kernel = std::make_unique((*g_cubin_module)["mul_two_cuda"]); } /*! * \brief Launch add_one_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_cubin_module != nullptr, RuntimeError) << "CUBIN module not loaded. Call set_cubin first."; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = g_add_one_kernel->Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } } // namespace cubin_dynamic // [example.end] namespace cubin_dynamic { /*! * \brief Launch mul_two_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void MulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_cubin_module != nullptr, RuntimeError) << "CUBIN module not loaded. Call set_cubin first."; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = g_mul_two_kernel->Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } // Export TVM-FFI functions TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_cubin, cubin_dynamic::SetCubin) TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, cubin_dynamic::AddOne) TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_two, cubin_dynamic::MulTwo) } // namespace cubin_dynamic tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/000077500000000000000000000000001521067262500226615ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/000077500000000000000000000000001521067262500245775ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/CMakeLists.txt000066400000000000000000000061461521067262500273460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.27) project(embedded_cubin_example LANGUAGES CXX CUDA) set(CMAKE_CXX_STANDARD 26) # Check for `#embed` after setting C++ standard, otherwise compilers (e.g. GCC) won't define the # macro. include(CheckCXXSymbolExists) check_cxx_symbol_exists(__cpp_pp_embed "" HAVE_EMBED) if (NOT HAVE_EMBED) message(FATAL_ERROR "Compiler does not support `#embed`." "Please use a newer compiler (e.g. GCC 15+, Clang 19+)" ) endif () set(CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API OFF CACHE BOOL "Use driver API in cubin launcher" ) # Prefer virtualenv when searching for python set(Python_FIND_VIRTUALENV FIRST) # cmake-lint: disable=C0103 # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Find CUDA toolkit find_package(CUDAToolkit REQUIRED) include_directories(${CUDAToolkit_INCLUDE_DIRS}) # [cmake_example.begin] # Step 1: Compile kernel.cu to FATBIN using add_tvm_ffi_fatbin utility or `CUDA_FATBIN_COMPILATION` set(CMAKE_CUDA_ARCHITECTURES 75;80;86;89;90;100;120) add_library(kernel_fatbin OBJECT src/kernel.cu) set_target_properties(kernel_fatbin PROPERTIES CUDA_FATBIN_COMPILATION ON) add_custom_target( kernel_fatbin.fatbin COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${CMAKE_CURRENT_SOURCE_DIR}/src/kernel_fatbin.fatbin" DEPENDS kernel_fatbin COMMENT "Copy fatbin to source dir" ) # Step 2: Build lib_embedded shared library add_library(lib_embedded SHARED src/lib_embedded.cc) add_dependencies(lib_embedded kernel_fatbin.fatbin) target_link_libraries(lib_embedded PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties(lib_embedded PROPERTIES POSITION_INDEPENDENT_CODE ON) # Step 3: Link against CUDA Driver API or Runtime API based on config if (CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API) add_compile_definitions(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1) target_link_libraries(lib_embedded PRIVATE cuda) else () target_link_libraries(lib_embedded PRIVATE CUDA::cudart) endif () set_target_properties( lib_embedded PROPERTIES PREFIX "" SUFFIX ".so" LINKER_LANGUAGE CXX ) # [cmake_example.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/main.py000066400000000000000000000070261521067262500261020ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example script for embedded CUBIN library. This example demonstrates using lib_embedded.so which has CUBIN data embedded directly in the shared library using objcopy. """ import sys from pathlib import Path import torch from tvm_ffi import load_module def main() -> int: """Test the lib_embedded.so library with embedded CUBIN.""" print("=" * 60) print("Example: Embedded CUBIN Library") print("=" * 60) # Check CUDA availability if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Load the library lib_path = Path(__file__).parent / "build" / "lib_embedded.so" mod = load_module(str(lib_path)) print(f"Loaded library: {lib_path}") # Get the functions add_one = mod["add_one"] mul_two = mod["mul_two"] print("Loaded functions: add_one, mul_two") # Test add_one kernel print("\n[Test 1] add_one kernel") n = 1024 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") add_one(x, y) # Verify results expected = x + 1 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test mul_two kernel print("\n[Test 2] mul_two kernel") n = 512 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") mul_two(x, y) # Verify results expected = x * 2 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test chained execution print("\n[Test 3] Chained execution: (x + 1) * 2") n = 256 x = torch.full((n,), 10.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Initial value: {x[0].item()}") add_one(x, temp) # temp = x + 1 = 11 mul_two(temp, y) # y = temp * 2 = 22 expected = 22.0 if torch.allclose(y, torch.tensor(expected, device="cuda")): print(f" [PASS] Result: {y[0].item()}") else: print(f" [FAIL] Expected {expected}, got {y[0].item()}") return 1 print("\n[PASS] All tests passed!") return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/src/000077500000000000000000000000001521067262500253665ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/src/kernel.cu000066400000000000000000000033221521067262500271770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/kernel.cu * \brief Simple CUDA kernel for testing cubin_launcher functionality. */ #include // [kernels.begin] /*! * \brief CUDA kernel that adds 1 to each element of an array. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void add_one_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } /*! * \brief CUDA kernel that multiplies each element by 2. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void mul_two_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } // [kernels.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/cpp_embed/src/lib_embedded.cc000066400000000000000000000103321521067262500302530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/lib_embedded.cc * \brief TVM-FFI library with embedded CUBIN kernels. * * This library exports TVM-FFI functions to launch CUDA kernels from * embedded CUBIN data. */ #include #include #include #include #include // [example.begin] constexpr unsigned char image[]{ // clang >= 20 or gcc >= 14 #embed "kernel_fatbin.fatbin" }; TVM_FFI_EMBED_CUBIN_FROM_BYTES(env, image); // [example.end] namespace cubin_embedded { /*! * \brief Launch add_one_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "add_one_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } } // namespace cubin_embedded namespace cubin_embedded { /*! * \brief Launch mul_two_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void MulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "mul_two_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } // Export TVM-FFI functions TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, cubin_embedded::AddOne) TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_two, cubin_embedded::MulTwo) } // namespace cubin_embedded tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/000077500000000000000000000000001521067262500265025ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/CMakeLists.txt000066400000000000000000000054111521067262500312430ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.20) project(embedded_cubin_example LANGUAGES CXX CUDA) set(CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API OFF CACHE BOOL "Use driver API in cubin launcher" ) # Prefer virtualenv when searching for python set(Python_FIND_VIRTUALENV FIRST) # cmake-lint: disable=C0103 # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Find CUDA toolkit find_package(CUDAToolkit REQUIRED) include_directories(${CUDAToolkit_INCLUDE_DIRS}) # [cmake_example.begin] # Step 1: Compile kernel.cu to FATBIN using add_tvm_ffi_fatbin utility or `CUDA_FATBIN_COMPILATION` set(CMAKE_CUDA_ARCHITECTURES 75;80;86;89;90;100;120) if (CMAKE_VERSION VERSION_LESS "3.27.0") add_tvm_ffi_fatbin(kernel_fatbin CUDA src/kernel.cu) else () add_library(kernel_fatbin OBJECT src/kernel.cu) set_target_properties(kernel_fatbin PROPERTIES CUDA_FATBIN_COMPILATION ON) endif () # Step 2: Build lib_embedded shared library add_library(lib_embedded SHARED src/lib_embedded.cc) target_link_libraries(lib_embedded PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties(lib_embedded PROPERTIES POSITION_INDEPENDENT_CODE ON) # Step 3: Link against CUDA Driver API or Runtime API based on config if (CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API) add_compile_definitions(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1) target_link_libraries(lib_embedded PRIVATE cuda) else () target_link_libraries(lib_embedded PRIVATE CUDA::cudart) endif () # Step 4: Embed CUBIN into shared library just defined, using tvm_ffi_embed_cubin utility This # creates symbols: __tvm_ffi__cubin_env (local) tvm_ffi_embed_bin_into(lib_embedded SYMBOL env BIN "$") set_target_properties( lib_embedded PROPERTIES PREFIX "" SUFFIX ".so" LINKER_LANGUAGE CXX ) # [cmake_example.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/main.py000066400000000000000000000070261521067262500300050ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example script for embedded CUBIN library. This example demonstrates using lib_embedded.so which has CUBIN data embedded directly in the shared library using objcopy. """ import sys from pathlib import Path import torch from tvm_ffi import load_module def main() -> int: """Test the lib_embedded.so library with embedded CUBIN.""" print("=" * 60) print("Example: Embedded CUBIN Library") print("=" * 60) # Check CUDA availability if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Load the library lib_path = Path(__file__).parent / "build" / "lib_embedded.so" mod = load_module(str(lib_path)) print(f"Loaded library: {lib_path}") # Get the functions add_one = mod["add_one"] mul_two = mod["mul_two"] print("Loaded functions: add_one, mul_two") # Test add_one kernel print("\n[Test 1] add_one kernel") n = 1024 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") add_one(x, y) # Verify results expected = x + 1 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test mul_two kernel print("\n[Test 2] mul_two kernel") n = 512 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") mul_two(x, y) # Verify results expected = x * 2 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test chained execution print("\n[Test 3] Chained execution: (x + 1) * 2") n = 256 x = torch.full((n,), 10.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Initial value: {x[0].item()}") add_one(x, temp) # temp = x + 1 = 11 mul_two(temp, y) # y = temp * 2 = 22 expected = 22.0 if torch.allclose(y, torch.tensor(expected, device="cuda")): print(f" [PASS] Result: {y[0].item()}") else: print(f" [FAIL] Expected {expected}, got {y[0].item()}") return 1 print("\n[PASS] All tests passed!") return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/src/000077500000000000000000000000001521067262500272715ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/src/kernel.cu000066400000000000000000000033221521067262500311020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/kernel.cu * \brief Simple CUDA kernel for testing cubin_launcher functionality. */ #include // [kernels.begin] /*! * \brief CUDA kernel that adds 1 to each element of an array. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void add_one_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } /*! * \brief CUDA kernel that multiplies each element by 2. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void mul_two_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } // [kernels.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/embed_with_tvm_ffi/src/lib_embedded.cc000066400000000000000000000103531521067262500321610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/lib_embedded.cc * \brief TVM-FFI library with embedded CUBIN kernels. * * This library exports TVM-FFI functions to launch CUDA kernels from * embedded CUBIN data. */ #include #include #include #include #include // [example.begin] // Embed CUBIN module with name "env" // This creates the necessary symbols and singleton struct for accessing the embedded CUBIN TVM_FFI_EMBED_CUBIN(env); // [example.end] namespace cubin_embedded { /*! * \brief Launch add_one_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "add_one_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } } // namespace cubin_embedded namespace cubin_embedded { /*! * \brief Launch mul_two_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void MulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "mul_two_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } // Export TVM-FFI functions TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, cubin_embedded::AddOne) TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_two, cubin_embedded::MulTwo) } // namespace cubin_embedded tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/000077500000000000000000000000001521067262500253615ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/CMakeLists.txt000066400000000000000000000053271521067262500301300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.27) project(embedded_cubin_example LANGUAGES CXX CUDA) set(CMAKE_CXX_STANDARD 17) set(CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API OFF CACHE BOOL "Use driver API in cubin launcher" ) # Prefer virtualenv when searching for python set(Python_FIND_VIRTUALENV FIRST) # cmake-lint: disable=C0103 # Find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) # Find CUDA toolkit find_package(CUDAToolkit REQUIRED) include_directories(${CUDAToolkit_INCLUDE_DIRS}) # [cmake_example.begin] # Step 1: Compile kernel.cu to FATBIN using add_tvm_ffi_fatbin utility or `CUDA_FATBIN_COMPILATION` set(CMAKE_CUDA_ARCHITECTURES 75;80;86;89;90;100;120) add_library(kernel_fatbin OBJECT src/kernel.cu) set_target_properties(kernel_fatbin PROPERTIES CUDA_FATBIN_COMPILATION ON) add_custom_target( kernel_fatbin.h COMMAND bin2c -c "$" > "${CMAKE_CURRENT_SOURCE_DIR}/src/kernel_fatbin.h" DEPENDS kernel_fatbin COMMENT "Run bin2c for ${kernel_fatbin}" ) # Step 2: Build lib_embedded shared library add_library(lib_embedded SHARED src/lib_embedded.cc) add_dependencies(lib_embedded kernel_fatbin.h) target_link_libraries(lib_embedded PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties(lib_embedded PROPERTIES POSITION_INDEPENDENT_CODE ON) # Step 3: Link against CUDA Driver API or Runtime API based on config if (CMAKE_TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API) add_compile_definitions(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API=1) target_link_libraries(lib_embedded PRIVATE cuda) else () target_link_libraries(lib_embedded PRIVATE CUDA::cudart) endif () set_target_properties( lib_embedded PROPERTIES PREFIX "" SUFFIX ".so" LINKER_LANGUAGE CXX ) # [cmake_example.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/main.py000066400000000000000000000070261521067262500266640ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example script for embedded CUBIN library. This example demonstrates using lib_embedded.so which has CUBIN data embedded directly in the shared library using objcopy. """ import sys from pathlib import Path import torch from tvm_ffi import load_module def main() -> int: """Test the lib_embedded.so library with embedded CUBIN.""" print("=" * 60) print("Example: Embedded CUBIN Library") print("=" * 60) # Check CUDA availability if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Load the library lib_path = Path(__file__).parent / "build" / "lib_embedded.so" mod = load_module(str(lib_path)) print(f"Loaded library: {lib_path}") # Get the functions add_one = mod["add_one"] mul_two = mod["mul_two"] print("Loaded functions: add_one, mul_two") # Test add_one kernel print("\n[Test 1] add_one kernel") n = 1024 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") add_one(x, y) # Verify results expected = x + 1 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test mul_two kernel print("\n[Test 2] mul_two kernel") n = 512 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") mul_two(x, y) # Verify results expected = x * 2 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 1 # Test chained execution print("\n[Test 3] Chained execution: (x + 1) * 2") n = 256 x = torch.full((n,), 10.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Initial value: {x[0].item()}") add_one(x, temp) # temp = x + 1 = 11 mul_two(temp, y) # y = temp * 2 = 22 expected = 22.0 if torch.allclose(y, torch.tensor(expected, device="cuda")): print(f" [PASS] Result: {y[0].item()}") else: print(f" [FAIL] Expected {expected}, got {y[0].item()}") return 1 print("\n[PASS] All tests passed!") return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/src/000077500000000000000000000000001521067262500261505ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/src/.gitignore000066400000000000000000000000201521067262500301300ustar00rootroot00000000000000kernel_fatbin.h tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/src/kernel.cu000066400000000000000000000033221521067262500277610ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/kernel.cu * \brief Simple CUDA kernel for testing cubin_launcher functionality. */ #include // [kernels.begin] /*! * \brief CUDA kernel that adds 1 to each element of an array. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void add_one_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } /*! * \brief CUDA kernel that multiplies each element by 2. * * \param x Input array pointer. * \param y Output array pointer. * \param n Number of elements. */ extern "C" __global__ void mul_two_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } // [kernels.end] tvm-ffi-0.1.12/examples/cubin_launcher/embedded_cubin/include_bin2c/src/lib_embedded.cc000066400000000000000000000102341521067262500310360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file examples/cubin_launcher/src/lib_embedded.cc * \brief TVM-FFI library with embedded CUBIN kernels. * * This library exports TVM-FFI functions to launch CUDA kernels from * embedded CUBIN data. */ #include #include #include #include #include #include "kernel_fatbin.h" // [example.begin] TVM_FFI_EMBED_CUBIN_FROM_BYTES(env, imageBytes); // [example.end] namespace cubin_embedded { /*! * \brief Launch add_one_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "add_one_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } } // namespace cubin_embedded namespace cubin_embedded { /*! * \brief Launch mul_two_cuda kernel on input tensor. * \param x Input tensor (float32, 1D) * \param y Output tensor (float32, 1D, same shape as x) */ void MulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(env, "mul_two_cuda"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); tvm::ffi::cuda_api::StreamHandle stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel tvm::ffi::cuda_api::ResultType result = kernel.Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } // Export TVM-FFI functions TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, cubin_embedded::AddOne) TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_two, cubin_embedded::MulTwo) } // namespace cubin_embedded tvm-ffi-0.1.12/examples/cubin_launcher/example_nvrtc_cubin.py000066400000000000000000000204521521067262500243540ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example script using NVRTC to compile CUDA kernels and embed them inline. This example demonstrates: 1. Compiling CUDA C++ source code to CUBIN using NVRTC 2. Embedding the CUBIN into a C++ module using tvm_ffi.cpp.load_inline 3. Launching the kernel through TVM-FFI Notes: - Requires `cuda-python` to be installed in the Python environment. """ from __future__ import annotations import sys import traceback import torch from tvm_ffi import cpp from tvm_ffi.cpp import nvrtc def generate_cubin() -> bytes: """Define CUDA kernels and compile them to a CUBIN file. The kernels are named `add_one` and `mul_two` and compute y[i] = x[i] + 1 and y[i] = x[i] * 2, respectively. Returns ------- bytes Compiled CUBIN bytes """ # [cuda_source.begin] # Define CUDA kernels cuda_source = """ extern "C" __global__ void add_one(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } extern "C" __global__ void mul_two(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } """ # Compile CUDA source to CUBIN using NVRTC print("Compiling CUDA kernels to CUBIN using NVRTC...") cubin_bytes = nvrtc.nvrtc_compile(cuda_source, name="kernels.cu") print(f"Compiled CUBIN: {len(cubin_bytes)} bytes\n") # [cuda_source.end] return cubin_bytes def use_cubin_kernel(cubin_bytes: bytes) -> int: """Load and test CUBIN kernels through TVM-FFI. Parameters ---------- cubin_bytes : bytes Compiled CUBIN bytes Returns ------- int: 0 on success, non-zero error code on failure """ # [cpp_wrapper.begin] # Define C++ code inline to launch the CUDA kernels using embedded CUBIN sources = """ #include #include #include #include #include // Embed CUBIN module with name "nvrtc_cubin" TVM_FFI_EMBED_CUBIN(nvrtc_cubin); namespace nvrtc_loader { void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(nvrtc_cubin, "add_one"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(args, grid, block, stream)); } void MulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(nvrtc_cubin, "mul_two"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Input and output must have same size"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); // Prepare kernel arguments void* args[] = {reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n)}; // Launch configuration tvm::ffi::dim3 grid((n + 255) / 256); tvm::ffi::dim3 block(256); // Get CUDA stream DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); // Launch kernel TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(args, grid, block, stream)); } } // namespace nvrtc_loader TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, nvrtc_loader::AddOne); TVM_FFI_DLL_EXPORT_TYPED_FUNC(mul_two, nvrtc_loader::MulTwo); """ print("Compiling C++ sources with tvm_ffi.cpp.load_inline...") mod = cpp.load_inline( "nvrtc_loader", cuda_sources=sources, embed_cubin={"nvrtc_cubin": cubin_bytes}, ) print("Successfully compiled and loaded C++ sources with embedded CUBIN\n") # [cpp_wrapper.end] # Get the functions add_one_fn = mod["add_one"] mul_two_fn = mod["mul_two"] # Test add_one kernel print("[Test 1] add_one kernel") n = 1024 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") add_one_fn(x, y) expected = x + 1 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 5 # Test mul_two kernel print("\n[Test 2] mul_two kernel") n = 512 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") mul_two_fn(x, y) expected = x * 2 if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 6 # Test chained execution print("\n[Test 3] Chained execution: (x + 1) * 2") n = 256 x = torch.full((n,), 10.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Initial value: {x[0].item()}") add_one_fn(x, temp) # temp = x + 1 = 11 mul_two_fn(temp, y) # y = temp * 2 = 22 expected = 22.0 if torch.allclose(y, torch.tensor(expected, device="cuda")): print(f" [PASS] Result: {y[0].item()}") else: print(f" [FAIL] Expected {expected}, got {y[0].item()}") return 7 print("\n[PASS] All tests passed!") return 0 def main() -> int: """Compile and launch CUDA kernel through NVRTC -> TVM-FFI.""" print("Example: NVRTC -> CUBIN -> C++ (inline embed) -> TVM-FFI") print("=" * 60) if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Generate CUBIN try: cubin_bytes = generate_cubin() except Exception as e: print(f"[ERROR] Failed to compile CUDA kernels: {e}") traceback.print_exc() return 2 # Use CUBIN kernels try: return use_cubin_kernel(cubin_bytes) except Exception as e: print(f"[ERROR] Failed to use CUBIN kernels: {e}") traceback.print_exc() return 3 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/cubin_launcher/example_triton_cubin.py000066400000000000000000000144121521067262500245360ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Single-file Triton example: define kernel, compile to CUBIN, load via inline C++. This script: 1. Embeds a minimal Triton kernel definition (elementwise square) 2. Compiles it to a CUBIN using the Triton runtime API 3. Defines C++ code inline using tvm_ffi.cpp.load_inline to load the CUBIN 4. Launches the kernel through the TVM-FFI exported function pointer Notes: - Requires `triton` to be installed in the Python environment. """ from __future__ import annotations import sys import traceback import torch import triton import triton.language as tl from tvm_ffi import cpp def generate_cubin() -> bytes: """Define a Triton kernel in-process and compile it to a CUBIN file. The kernel is named `square_kernel` and computes y[i] = x[i] * x[i]. Returns ------- bytes Compiled CUBIN bytes """ # [triton_kernel.begin] # Define the kernel dynamically @triton.jit def square_kernel(X_ptr, Y_ptr, n, BLOCK: tl.constexpr = 1024): # noqa # ty: ignore[invalid-parameter-default] pid = tl.program_id(0) start = pid * BLOCK offsets = start + tl.arange(0, BLOCK) mask = offsets < n x = tl.load(X_ptr + offsets, mask=mask, other=0.0) y = x * x tl.store(Y_ptr + offsets, y, mask=mask) # Trigger kernel compilation by doing a dummy call x_dummy = torch.ones(1024, dtype=torch.float32, device="cuda") y_dummy = torch.empty(1024, dtype=torch.float32, device="cuda") compiled_kernel = square_kernel[1, 1](x_dummy, y_dummy, 1024) # Get CUBIN bytes cubin_bytes = compiled_kernel.kernel # [triton_kernel.end] return cubin_bytes def use_cubin_kernel(cubin_bytes: bytes) -> int: """Load and test Triton CUBIN kernel through TVM-FFI. Parameters ---------- cubin_bytes : bytes Compiled CUBIN bytes Returns ------- int: 0 on success, non-zero error code on failure """ # [cpp_wrapper.begin] # Define C++ code inline to load and launch the Triton kernel using embedded CUBIN sources = """ #include #include #include #include #include // Embed CUBIN module with name "triton_cubin" TVM_FFI_EMBED_CUBIN(triton_cubin); namespace triton_loader { void LaunchSquare(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Get kernel from embedded CUBIN (cached in static variable for efficiency) static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(triton_cubin, "square_kernel"); TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Sizes must match"; uint32_t n = static_cast(x.size(0)); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); uint64_t dummy_ptr = 0; // Workaround for Triton extra params: pass dummy addresses for unused parameters void* args[] = {&x_ptr, &y_ptr, &n, &dummy_ptr, &dummy_ptr}; // Kernel was compiled with .reqntid 128, not 1024 tvm::ffi::dim3 grid((n + 127) / 128); tvm::ffi::dim3 block(128); DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(args, grid, block, stream)); } } // namespace triton_loader TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_square, triton_loader::LaunchSquare); """ print("Compiling C++ sources with tvm_ffi.cpp.load_inline...") # Find CUDA include path mod = cpp.load_inline( "triton_loader", cuda_sources=sources, embed_cubin={"triton_cubin": cubin_bytes}, ) print("Successfully compiled and loaded C++ sources with embedded CUBIN\n") # [cpp_wrapper.end] # Get the launch function launch_fn = mod["launch_square"] # Test kernel: compute square print("[Test] square kernel") n = 4096 x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") print(f" Input shape: {x.shape}, device: {x.device}") launch_fn(x, y) expected = x * x if torch.allclose(y, expected): print(f" [PASS] Verified {n} elements correctly") return 0 else: print(f" [FAIL] Verification failed, max error: {(y - expected).abs().max().item()}") return 5 def main() -> int: """Load and launch Triton kernel through TVM-FFI.""" print("Example: Triton (inline) -> CUBIN -> C++ (inline) -> TVM-FFI") print("=" * 60) if not torch.cuda.is_available(): print("[ERROR] CUDA is not available") return 1 print(f"CUDA device: {torch.cuda.get_device_name(0)}") print(f"PyTorch version: {torch.__version__}\n") # Compile Triton kernel to CUBIN try: print("Compiling Triton kernel to CUBIN...") cubin_bytes = generate_cubin() print(f"Compiled CUBIN: {len(cubin_bytes)} bytes\n") except Exception as e: print(f"[ERROR] Failed to compile Triton kernel: {e}") traceback.print_exc() return 2 # Use CUBIN kernel try: return use_cubin_kernel(cubin_bytes) except Exception as e: print(f"[ERROR] Failed to use CUBIN kernel: {e}") traceback.print_exc() return 3 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/examples/inline_module/000077500000000000000000000000001521067262500176125ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/inline_module/main.py000066400000000000000000000076041521067262500211170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Example: Build and run an inline C++/CUDA tvm-ffi module.""" import torch import tvm_ffi.cpp from tvm_ffi.module import Module def main() -> None: """Build, load, and run inline CPU/CUDA functions.""" mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } void add_one_cuda(tvm::ffi::TensorView x, tvm::ffi::TensorView y); """, cuda_sources=r""" __global__ void AddOneKernel(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1; } } void add_one_cuda(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; int64_t n = x.size(0); int64_t nthread_per_block = 256; int64_t nblock = (n + nthread_per_block - 1) / nthread_per_block; // Obtain the current stream from the environment // it will be set to torch.cuda.current_stream() when calling the function // with torch.Tensors cudaStream_t stream = static_cast( TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); // launch the kernel AddOneKernel<<>>(static_cast(x.data_ptr()), static_cast(y.data_ptr()), n); } """, functions=["add_one_cpu", "add_one_cuda"], ) x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) x_cuda = x.cuda() y_cuda = torch.empty_like(x_cuda) mod.add_one_cuda(x_cuda, y_cuda) torch.testing.assert_close(x_cuda + 1, y_cuda) if __name__ == "__main__": main() tvm-ffi-0.1.12/examples/kernel_library/000077500000000000000000000000001521067262500177735ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/kernel_library/load_scale.py000066400000000000000000000023351521067262500224360ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Load and call a scale kernel.""" # [load_and_call.begin] import torch import tvm_ffi # Load the compiled shared library mod = tvm_ffi.load_module("build/scale_kernel.so") # Pre-allocate input and output tensors in PyTorch x = torch.randn(1024, device="cuda", dtype=torch.float32) y = torch.empty_like(x) # Call the kernel — PyTorch tensors are auto-converted to TensorView mod.scale(y, x, 2.0) assert torch.allclose(y, x * 2.0) # [load_and_call.end] tvm-ffi-0.1.12/examples/kernel_library/scale_kernel.cu000066400000000000000000000047711521067262500227640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include "tvm_ffi_utils.h" // [cuda_kernel.begin] template __global__ void ScaleKernel(T* out, const T* in, T factor, int64_t n) { int64_t i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { out[i] = in[i] * factor; } } // [cuda_kernel.end] // [function.begin] void Scale(TensorView output, TensorView input, double factor) { // --- 1. Validate inputs --- CHECK_INPUT(input); CHECK_INPUT(output); CHECK_DIM(1, input); CHECK_DEVICE(input, output); TVM_FFI_CHECK(input.dtype() == output.dtype(), ValueError) << "input/output dtype mismatch"; TVM_FFI_CHECK(input.numel() == output.numel(), ValueError) << "input/output size mismatch"; // --- 2. Device guard and stream --- ffi::CUDADeviceGuard guard(input.device().device_id); cudaStream_t stream = get_cuda_stream(input.device()); // --- 3. Dispatch on dtype and launch --- int64_t n = input.numel(); int threads = 256; int blocks = (n + threads - 1) / threads; if (input.dtype() == dl_float32) { ScaleKernel<<>>(static_cast(output.data_ptr()), static_cast(input.data_ptr()), static_cast(factor), n); } else if (input.dtype() == dl_float16) { ScaleKernel<<>>(static_cast(output.data_ptr()), static_cast(input.data_ptr()), static_cast(factor), n); } else { TVM_FFI_THROW(TypeError) << "Unsupported dtype: " << input.dtype(); } } // [function.end] // [export.begin] TVM_FFI_DLL_EXPORT_TYPED_FUNC(scale, Scale); // [export.end] tvm-ffi-0.1.12/examples/kernel_library/tvm_ffi_utils.h000066400000000000000000000057171521067262500230300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #ifndef KERNEL_LIBRARY_TVM_FFI_UTILS_H_ #define KERNEL_LIBRARY_TVM_FFI_UTILS_H_ #include #include #include namespace ffi = tvm::ffi; using ffi::Optional; using ffi::Tensor; using ffi::TensorView; // [check_macros.begin] // --- Reusable validation macros --- #define CHECK_CUDA(x) \ TVM_FFI_CHECK((x).device().device_type == kDLCUDA, ValueError) << #x " must be a CUDA tensor" #define CHECK_CONTIGUOUS(x) \ TVM_FFI_CHECK((x).IsContiguous(), ValueError) << #x " must be contiguous" #define CHECK_INPUT(x) \ do { \ CHECK_CUDA(x); \ CHECK_CONTIGUOUS(x); \ } while (0) #define CHECK_DIM(d, x) \ TVM_FFI_CHECK((x).ndim() == (d), ValueError) << #x " must be a " #d "D tensor" #define CHECK_DEVICE(a, b) \ do { \ TVM_FFI_CHECK((a).device().device_type == (b).device().device_type, ValueError) \ << #a " and " #b " must be on the same device type"; \ TVM_FFI_CHECK((a).device().device_id == (b).device().device_id, ValueError) \ << #a " and " #b " must be on the same device"; \ } while (0) // [check_macros.end] // [get_stream.begin] // --- Stream helper --- inline cudaStream_t get_cuda_stream(DLDevice device) { return static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); } // [get_stream.end] // [alloc_tensor.begin] // --- Tensor allocation helper --- inline ffi::Tensor alloc_tensor(const ffi::Shape& shape, DLDataType dtype, DLDevice device) { return ffi::Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, shape, dtype, device); } // [alloc_tensor.end] // [dtype_constants.begin] // --- DLPack dtype constants --- constexpr DLDataType dl_float16 = DLDataType{kDLFloat, 16, 1}; constexpr DLDataType dl_float32 = DLDataType{kDLFloat, 32, 1}; constexpr DLDataType dl_float64 = DLDataType{kDLFloat, 64, 1}; constexpr DLDataType dl_bfloat16 = DLDataType{kDLBfloat, 16, 1}; // [dtype_constants.end] #endif // KERNEL_LIBRARY_TVM_FFI_UTILS_H_ tvm-ffi-0.1.12/examples/python_packaging/000077500000000000000000000000001521067262500203145ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/python_packaging/CMakeLists.txt000066400000000000000000000023041521067262500230530ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.18) project(my_ffi_extension) find_package( Python COMPONENTS Interpreter REQUIRED ) find_package(tvm_ffi CONFIG REQUIRED) # [example.cmake.begin] add_library(my_ffi_extension SHARED src/extension.cc) tvm_ffi_configure_target(my_ffi_extension STUB_DIR "./python" STUB_INIT ON) install(TARGETS my_ffi_extension DESTINATION .) tvm_ffi_install(my_ffi_extension DESTINATION .) # [example.cmake.end] tvm-ffi-0.1.12/examples/python_packaging/README.md000066400000000000000000000035531521067262500216010ustar00rootroot00000000000000 # TVM FFI Packaging Example This is an example project that packages a tvm-ffi based library into a Python ABI-agnostic wheel. This example can also serve as a guideline for general packaging as well. - Source-level build for cross-compilation support in CMake - Registration via global function table ## Install the wheel Use `uv pip` (the same tooling used in CI) to build and install the example wheel: ```bash cd examples/python_packaging uv pip install --reinstall --verbose . ``` ### Note on build and auditwheel Note: When running the auditwheel process, make sure to skip `libtvm_ffi.so` as they are shipped via the tvm_ffi package. ## Run the example After installing the `my_ffi_extension` example package, you can run the following example. ```bash python run_example.py ``` This runs four flows: calling `add_two` via the TVM-FFI ABI, calling `add_one` via the global registry, calling `raise_error` to demonstrate error propagation, and constructing/using the `IntPair` object. tvm-ffi-0.1.12/examples/python_packaging/pyproject.toml000066400000000000000000000036621521067262500232370ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [project] name = "my-ffi-extension" version = "0.1.0" readme = "README.md" license = { text = "Apache 2.0" } classifiers = [ "License :: OSI Approved :: Apache Software License", "Development Status :: 4 - Beta", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", ] keywords = ["machine learning", "inference"] requires-python = ">=3.9" dependencies = ["apache-tvm-ffi"] # [pyproject.build.begin] [build-system] requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"] build-backend = "scikit_build_core.build" [tool.scikit-build] # The wheel is Python ABI-agnostic wheel.py-api = "py3" # The package contains the Python module at `python/my_ffi_extension` wheel.packages = ["python/my_ffi_extension"] # The install dir matches the import name wheel.install-dir = "my_ffi_extension" # Build the extension in-place under "./build-wheel" build-dir = "build-wheel" # Specify minimum CMake version cmake.version = "CMakeLists.txt" # Specify CMake build type cmake.build-type = "Release" # Pass custom CMake definitions [tool.scikit-build.cmake.define] CMAKE_EXPORT_COMPILE_COMMANDS = "ON" # [pyproject.build.end] tvm-ffi-0.1.12/examples/python_packaging/python/000077500000000000000000000000001521067262500216355ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/python_packaging/python/my_ffi_extension/000077500000000000000000000000001521067262500252025ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/python_packaging/python/my_ffi_extension/__init__.py000066400000000000000000000020641521067262500273150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations. """Package my_ffi_extension.""" # tvm-ffi-stubgen(begin): export/_ffi_api # fmt: off # isort: off from ._ffi_api import * # noqa: F403 from ._ffi_api import __all__ as _ffi_api__all__ if "__all__" not in globals(): __all__ = [] __all__.extend(_ffi_api__all__) # isort: on # fmt: on # tvm-ffi-stubgen(end) tvm-ffi-0.1.12/examples/python_packaging/python/my_ffi_extension/_ffi_api.py000066400000000000000000000045421521067262500273150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations. """FFI API bindings for my_ffi_extension.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from tvm_ffi import Object as _ffi_Object, init_ffi_api as _FFI_INIT_FUNC, register_object as _FFI_REG_OBJ from tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB from typing import TYPE_CHECKING if TYPE_CHECKING: from tvm_ffi import Object # isort: on # fmt: on # tvm-ffi-stubgen(end) # tvm-ffi-stubgen(import-object): tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB LIB = _FFI_LOAD_LIB("my_ffi_extension", "my_ffi_extension") # tvm-ffi-stubgen(begin): global/my_ffi_extension # fmt: off _FFI_INIT_FUNC("my_ffi_extension", __name__) if TYPE_CHECKING: def add_one(_0: int, /) -> int: ... def raise_error(_0: str, /) -> None: ... # fmt: on # tvm-ffi-stubgen(end) # tvm-ffi-stubgen(import-object): tvm_ffi.register_object;False;_FFI_REG_OBJ # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object @_FFI_REG_OBJ("my_ffi_extension.IntPair") class IntPair(_ffi_Object): """FFI binding for `my_ffi_extension.IntPair`.""" # tvm-ffi-stubgen(begin): object/my_ffi_extension.IntPair # fmt: off a: int b: int if TYPE_CHECKING: def __init__(self, _0: int, _1: int, /) -> None: ... def __ffi_shallow_copy__(self, /) -> Object: ... @staticmethod def __c_ffi_init__(_0: int, _1: int, /) -> Object: ... def sum(self, /) -> int: ... # fmt: on # tvm-ffi-stubgen(end) __all__ = [ # tvm-ffi-stubgen(begin): __all__ "LIB", "IntPair", "add_one", "raise_error", # tvm-ffi-stubgen(end) ] tvm-ffi-0.1.12/examples/python_packaging/run_example.py000066400000000000000000000037651521067262500232200ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations. # Base logic to load library for extension package """Run functions from the example packaged tvm-ffi extension.""" import traceback import my_ffi_extension def run_add_two() -> None: """Invoke add_two from the extension and print the result.""" print("=========== Example 1: add_two ===========") print(my_ffi_extension.LIB.add_two(1)) def run_add_one() -> None: """Invoke add_one from the extension and print the result.""" print("=========== Example 2: add_one ===========") print(my_ffi_extension.add_one(3)) def run_raise_error() -> None: """Invoke raise_error from the extension to demonstrate error handling.""" print("=========== Example 3: raise_error ===========") try: my_ffi_extension.raise_error("This is an error") except RuntimeError: traceback.print_exc() def run_int_pair() -> None: """Invoke IntPair from the extension to demonstrate object handling.""" print("=========== Example 4: IntPair ===========") pair = my_ffi_extension.IntPair(1, 2) # ty: ignore[too-many-positional-arguments] print(f"a={pair.a}") print(f"b={pair.b}") print(f"sum={pair.sum()}") if __name__ == "__main__": run_add_two() run_add_one() run_raise_error() run_int_pair() tvm-ffi-0.1.12/examples/python_packaging/src/000077500000000000000000000000001521067262500211035ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/python_packaging/src/extension.cc000066400000000000000000000051441521067262500234320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file extension.cc * \brief Example of a tvm-ffi based library that registers various functions. */ #include #include namespace my_ffi_extension { namespace ffi = tvm::ffi; // [tvm_ffi_abi.begin] static int AddTwo(int x) { return x + 2; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_two, AddTwo) // [tvm_ffi_abi.end] // [global_function.begin] static int AddOne(int x) { return x + 1; } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() // .def("my_ffi_extension.add_one", AddOne); } // [global_function.end] /*! * \brief Raise a runtime error to demonstrate error propagation. * * \param msg The error message to raise. * * \code{.py} * import my_ffi_extension * try: * my_ffi_extension.raise_error("boom") * except RuntimeError: * pass * \endcode */ static void RaiseError(const ffi::String& msg) { TVM_FFI_THROW(RuntimeError) << msg; } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() // .def("my_ffi_extension.raise_error", RaiseError); } // [object.begin] class IntPairObj : public ffi::Object { public: int64_t a; int64_t b; IntPairObj(int64_t a, int64_t b) : a(a), b(b) {} int64_t Sum() const { return a + b; } static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL( /*type_key=*/"my_ffi_extension.IntPair", /*class=*/IntPairObj, /*parent_class=*/ffi::Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def(refl::init()) .def_rw("a", &IntPairObj::a, "the first field") .def_rw("b", &IntPairObj::b, "the second field") .def("sum", &IntPairObj::Sum, "IntPairObj::Sum() method"); } // [object.end] } // namespace my_ffi_extension tvm-ffi-0.1.12/examples/quickstart/000077500000000000000000000000001521067262500171615ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/quickstart/CMakeLists.txt000066400000000000000000000064441521067262500217310ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.20) project(quickstart LANGUAGES CXX) option(EXAMPLE_NAME "Which example to build ('compile_cpu'/'compile_cuda'/'load_cpp')" "compile_cpu" ) message(STATUS "Building example: ${EXAMPLE_NAME}") # cmake-lint: disable=C0111 function (set_flat_output_dirs target_name) set(output_dir "${CMAKE_BINARY_DIR}") # Ensure multi-config generators (e.g., MSVC) output to the same flat directory. foreach (config IN ITEMS "" DEBUG RELEASE RELWITHDEBINFO MINSIZEREL) if (config STREQUAL "") set_property(TARGET ${target_name} PROPERTY RUNTIME_OUTPUT_DIRECTORY "${output_dir}") set_property(TARGET ${target_name} PROPERTY LIBRARY_OUTPUT_DIRECTORY "${output_dir}") set_property(TARGET ${target_name} PROPERTY ARCHIVE_OUTPUT_DIRECTORY "${output_dir}") else () set_property( TARGET ${target_name} PROPERTY RUNTIME_OUTPUT_DIRECTORY_${config} "${output_dir}" ) set_property( TARGET ${target_name} PROPERTY LIBRARY_OUTPUT_DIRECTORY_${config} "${output_dir}" ) set_property( TARGET ${target_name} PROPERTY ARCHIVE_OUTPUT_DIRECTORY_${config} "${output_dir}" ) endif () endforeach () endfunction () # Run `tvm_ffi.config --cmakedir` to find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) if (EXAMPLE_NAME STREQUAL "compile_cpu") # Example 1. C++ `add_one` add_library(add_one_cpu SHARED compile/add_one_cpu.cc) target_link_libraries(add_one_cpu PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties(add_one_cpu PROPERTIES PREFIX "" SUFFIX ".so") set_flat_output_dirs(add_one_cpu) elseif (EXAMPLE_NAME STREQUAL "compile_cuda") # Example 2. CUDA `add_one` enable_language(CUDA) add_library(add_one_cuda SHARED compile/add_one_cuda.cu) target_link_libraries(add_one_cuda PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties(add_one_cuda PROPERTIES PREFIX "" SUFFIX ".so") set_flat_output_dirs(add_one_cuda) elseif (EXAMPLE_NAME STREQUAL "load_cpp") # Example 3. Load C++ shared library add_executable(load_cpp load/load_cpp.cc) target_link_libraries(load_cpp PRIVATE tvm_ffi::header tvm_ffi::shared) set_flat_output_dirs(load_cpp) else () message(FATAL_ERROR "Unknown EXAMPLE_NAME option: ${EXAMPLE_NAME}. " "Expected: 'compile_cpu', 'compile_cuda', 'load_cpp'." ) endif () tvm-ffi-0.1.12/examples/quickstart/README.md000066400000000000000000000042341521067262500204430ustar00rootroot00000000000000 # Quick Start Code Example This directory contains all the source code for [tutorial](https://tvm.apache.org/ffi/get_started/quickstart.html). ## Run everything (CPU path) On Linux/macOS: ```bash bash run_all_cpu.sh ``` On Windows: ```batch run_all_cpu.bat ``` ## Compile and Distribute `add_one_*` manually To compile the C++ Example: ```bash cmake . -B build -DEXAMPLE_NAME="compile_cpu" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo ``` This produces `build/add_one_cpu.so`. To compile CUDA Example (Linux with CUDA toolchain available): ```bash cmake . -B build -DEXAMPLE_NAME="compile_cuda" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo ``` ## Load the Distributed `add_one_*` To run library loading examples across ML frameworks (requires CUDA for the CUDA example): ```bash python load/load_pytorch.py python load/load_paddle.py python load/load_numpy.py python load/load_cupy.py ``` To run library loading example in C++: ```bash cmake . -B build -DEXAMPLE_NAME="load_cpp" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build/load_cpp ``` The executable is emitted as `build/load_cpp` (`build/load_cpp.exe` on Windows). For a CUDA end-to-end run, use: ```bash bash run_all_cuda.sh ``` tvm-ffi-0.1.12/examples/quickstart/compile/000077500000000000000000000000001521067262500206115ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/quickstart/compile/add_one_cpu.cc000066400000000000000000000025411521067262500233620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // [example.begin] // File: compile/add_one_cpu.cc #include namespace tvm_ffi_example_cpu { /*! \brief Perform vector add one: y = x + 1 (1-D float32) */ void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { int64_t n = x.size(0); float* x_data = static_cast(x.data_ptr()); float* y_data = static_cast(y.data_ptr()); for (int64_t i = 0; i < n; ++i) { y_data[i] = x_data[i] + 1; } } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, tvm_ffi_example_cpu::AddOne) } // namespace tvm_ffi_example_cpu // [example.end] tvm-ffi-0.1.12/examples/quickstart/compile/add_one_cuda.cu000066400000000000000000000032541521067262500235330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // [example.begin] // File: compile/add_one_cuda.cu #include #include namespace tvm_ffi_example_cuda { __global__ void AddOneKernel(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1; } } void AddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { int64_t n = x.size(0); float* x_data = static_cast(x.data_ptr()); float* y_data = static_cast(y.data_ptr()); int64_t threads = 256; int64_t blocks = (n + threads - 1) / threads; cudaStream_t stream = static_cast(TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); AddOneKernel<<>>(x_data, y_data, n); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cuda, tvm_ffi_example_cuda::AddOne); } // namespace tvm_ffi_example_cuda // [example.end] tvm-ffi-0.1.12/examples/quickstart/load/000077500000000000000000000000001521067262500201005ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/quickstart/load/load_cpp.cc000066400000000000000000000050651521067262500221760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // [main.begin] // File: load/load_cpp.cc #include #include namespace { namespace ffi = tvm::ffi; /*! * \brief Main logics of library loading and function calling. * \param x The input tensor. * \param y The output tensor. */ void Run(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Load shared library `build/add_one_cpu.so` ffi::Module mod = ffi::Module::LoadFromFile("build/add_one_cpu.so"); // Look up `add_one_cpu` function ffi::Function add_one_cpu = mod->GetFunction("add_one_cpu").value(); // Call the function add_one_cpu(x, y); } } // namespace // [main.end] /************* Auxiliary Logics *************/ // [aux.begin] /*! * \brief Allocate a 1D float32 `tvm::ffi::Tensor` on CPU from an braced initializer list. * \param data The input data. * \return The allocated Tensor. */ ffi::Tensor Alloc1DTensor(std::initializer_list data) { struct CPUAllocator { void AllocData(DLTensor* tensor) { tensor->data = std::malloc(tensor->shape[0] * sizeof(float)); } void FreeData(DLTensor* tensor) { std::free(tensor->data); } }; DLDataType f32 = DLDataType({kDLFloat, 32, 1}); DLDevice cpu = DLDevice({kDLCPU, 0}); int64_t n = static_cast(data.size()); ffi::Tensor x = ffi::Tensor::FromNDAlloc(CPUAllocator(), {n}, f32, cpu); float* x_data = static_cast(x.data_ptr()); for (float v : data) { *x_data++ = v; } return x; } int main() { ffi::Tensor x = Alloc1DTensor({1, 2, 3, 4, 5}); ffi::Tensor y = Alloc1DTensor({0, 0, 0, 0, 0}); Run(x, y); std::cout << "[ "; const float* y_data = static_cast(y.data_ptr()); for (int i = 0; i < 5; ++i) { std::cout << y_data[i] << " "; } std::cout << "]" << std::endl; return 0; } // [aux.end] tvm-ffi-0.1.12/examples/quickstart/load/load_cuda.cc000066400000000000000000000071151521067262500223260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // [main.begin] // File: load/load_cuda.cc #include #include namespace { namespace ffi = tvm::ffi; /*! * \brief Main logics of library loading and function calling with CUDA tensors. * \param x The input tensor on CUDA device. * \param y The output tensor on CUDA device. */ void Run(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // Load shared library `build/add_one_cuda.so` ffi::Module mod = ffi::Module::LoadFromFile("build/add_one_cuda.so"); // Look up `add_one_cuda` function ffi::Function add_one_cuda = mod->GetFunction("add_one_cuda").value(); // Call the function with CUDA tensors add_one_cuda(x, y); } } // namespace // [main.end] /************* Auxiliary Logics *************/ // [aux.begin] #include #include #include #include struct CUDANDAlloc { void AllocData(DLTensor* tensor) { size_t data_size = ffi::GetDataSize(*tensor); void* ptr = nullptr; cudaError_t err = cudaMalloc(&ptr, data_size); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err); tensor->data = ptr; } void FreeData(DLTensor* tensor) { if (tensor->data != nullptr) { cudaError_t err = cudaFree(tensor->data); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err); tensor->data = nullptr; } } }; /*! * \brief Allocate a CUDA tensor with the given shape and data type. * \param shape The shape of the tensor. * \param dtype The data type of the tensor. * \param device The CUDA device. * \return The allocated CUDA tensor. */ inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) { return ffi::Tensor::FromNDAlloc(CUDANDAlloc(), shape, dtype, device); } int main() { DLDataType f32_dtype{kDLFloat, 32, 1}; DLDevice cuda_device{kDLCUDA, 0}; constexpr int ARRAY_SIZE = 5; ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); std::vector host_x(ARRAY_SIZE); for (int i = 0; i < ARRAY_SIZE; ++i) { host_x[i] = static_cast(i + 1); } size_t nbytes = host_x.size() * sizeof(float); cudaError_t err = cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMemcpy host to device failed: " << cudaGetErrorString(err); Run(x, y); std::vector host_y(host_x.size()); err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost); TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMemcpy device to host failed: " << cudaGetErrorString(err); std::cout << "[ "; for (float value : host_y) { std::cout << value << " "; } std::cout << "]" << std::endl; return 0; } // [aux.end] tvm-ffi-0.1.12/examples/quickstart/load/load_cupy.py000066400000000000000000000020651521067262500224340ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # fmt: off # ruff: noqa # mypy: ignore-errors # [example.begin] # File: load/load_cupy.py import tvm_ffi mod = tvm_ffi.load_module("build/add_one_cuda.so") import cupy as cp x = cp.array([1, 2, 3, 4, 5], dtype=cp.float32) y = cp.empty_like(x) mod.add_one_cuda(x, y) print(y) # [example.end] tvm-ffi-0.1.12/examples/quickstart/load/load_numpy.py000066400000000000000000000020651521067262500226240ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # fmt: off # ruff: noqa # mypy: ignore-errors # [example.begin] # File: load/load_numpy.py import tvm_ffi mod = tvm_ffi.load_module("build/add_one_cpu.so") import numpy as np x = np.array([1, 2, 3, 4, 5], dtype=np.float32) y = np.empty_like(x) mod.add_one_cpu(x, y) print(y) # [example.end] tvm-ffi-0.1.12/examples/quickstart/load/load_paddle.py000066400000000000000000000021171521067262500227030ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # fmt: off # ruff: noqa # mypy: ignore-errors # [example.begin] # File: load/load_paddle.py import tvm_ffi mod = tvm_ffi.load_module("build/add_one_cuda.so") import paddle x = paddle.tensor([1, 2, 3, 4, 5], dtype=paddle.float32, device="cuda") y = paddle.empty_like(x) mod.add_one_cuda(x, y) print(y) # [example.end] tvm-ffi-0.1.12/examples/quickstart/load/load_pytorch.py000066400000000000000000000022411521067262500231400ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # fmt: off # ruff: noqa # mypy: ignore-errors # [example.begin] # File: load/load_pytorch.py # Step 1. Load `build/add_one_cuda.so` import tvm_ffi mod = tvm_ffi.load_module("build/add_one_cuda.so") # Step 2. Run `mod.add_one_cuda` with PyTorch import torch x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32, device="cuda") y = torch.empty_like(x) mod.add_one_cuda(x, y) print(y) # [example.end] tvm-ffi-0.1.12/examples/quickstart/raw_compile.sh000077500000000000000000000056601521067262500220300ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # shellcheck disable=SC2046 set -ex BUILD_DIR=build mkdir -p $BUILD_DIR # Example 1. Compile C++ `add_one_cpu.cc` to shared library `add_one_cpu.so` # [cpp_compile.begin] g++ -shared -O3 compile/add_one_cpu.cc \ -fPIC -fvisibility=hidden \ $(tvm-ffi-config --cxxflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) \ -o $BUILD_DIR/add_one_cpu.so # [cpp_compile.end] # Example 2. Compile CUDA `add_one_cuda.cu` to shared library `add_one_cuda.so` if command -v nvcc >/dev/null 2>&1; then # [cuda_compile.begin] nvcc -shared -O3 compile/add_one_cuda.cu \ -Xcompiler -fPIC,-fvisibility=hidden \ $(tvm-ffi-config --cxxflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) \ -o $BUILD_DIR/add_one_cuda.so # [cuda_compile.end] fi # Example 3. Load and run `add_one_cpu.so` in C++ if [ -f "$BUILD_DIR/add_one_cpu.so" ]; then # [load_cpp.begin] g++ -fvisibility=hidden -O3 \ load/load_cpp.cc \ $(tvm-ffi-config --cxxflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) \ -Wl,-rpath,$(tvm-ffi-config --libdir) \ -o build/load_cpp build/load_cpp # [load_cpp.end] fi # Example 4. Load and run `add_one_cuda.so` in C++ # Before run this example, make sure you have a CUDA-capable GPU and the CUDA toolkit installed. # See CONTRIBUTING.md to use a pre-built Docker image with CUDA support. # if [ -f "$BUILD_DIR/add_one_cuda.so" ] && command -v nvcc >/dev/null 2>&1; then # # [load_cuda.begin] # g++ -fvisibility=hidden -O3 \ # load/load_cuda.cc \ # $(tvm-ffi-config --cxxflags) \ # $(tvm-ffi-config --ldflags) \ # $(tvm-ffi-config --libs) \ # -I/usr/local/cuda/include \ # -L/usr/local/cuda/lib64 \ # -lcudart \ # -Wl,-rpath,$(tvm-ffi-config --libdir) \ # -Wl,-rpath,/usr/local/cuda/lib64 \ # -o build/load_cuda # build/load_cuda # # [load_cuda.end] # fi tvm-ffi-0.1.12/examples/quickstart/run_all_cpu.bat000066400000000000000000000030601521067262500221530ustar00rootroot00000000000000@echo off rem Licensed to the Apache Software Foundation (ASF) under one rem or more contributor license agreements. See the NOTICE file rem distributed with this work for additional information rem regarding copyright ownership. The ASF licenses this file rem to you under the Apache License, Version 2.0 (the rem "License"); you may not use this file except in compliance rem with the License. You may obtain a copy of the License at rem rem http://www.apache.org/licenses/LICENSE-2.0 rem rem Unless required by applicable law or agreed to in writing, rem software distributed under the License is distributed on an rem "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY rem KIND, either express or implied. See the License for the rem specific language governing permissions and limitations rem under the License. setlocal enabledelayedexpansion cd /d "%~dp0" if exist build rmdir /s /q build rem To compile `compile/add_one_cpu.cc` to shared library `build/add_one_cpu.so` cmake . -B build -DEXAMPLE_NAME="compile_cpu" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo rem To load and run `add_one_cpu.so` in NumPy python load\load_numpy.py rem To load and run `add_one_cpu.so` in C++ for /f "delims=" %%i in ('python -c "from tvm_ffi import libinfo; import pathlib; print(pathlib.Path(libinfo.find_libtvm_ffi()).parent)"') do set "TVM_FFI_LIBDIR=%%i" set "PATH=%TVM_FFI_LIBDIR%;%PATH%" cmake . -B build -DEXAMPLE_NAME="load_cpp" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build\load_cpp.exe endlocal tvm-ffi-0.1.12/examples/quickstart/run_all_cpu.sh000077500000000000000000000023531521067262500220260ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. set -ex # To compile `compile/add_one_cpu.cc` to shared library `build/add_one_cpu.so` cmake . -B build -DEXAMPLE_NAME="compile_cpu" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo # To load and run `add_one_cpu.so` in NumPy python load/load_numpy.py # To load and run `add_one_cpu.so` in C++ cmake . -B build -DEXAMPLE_NAME="load_cpp" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build/load_cpp tvm-ffi-0.1.12/examples/quickstart/run_all_cuda.sh000077500000000000000000000023261521067262500221530ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. set -ex # To compile `compile/add_one_cuda.cu` to shared library `build/add_one_cuda.so` cmake . -B build -DEXAMPLE_NAME="compile_cuda" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo # To load and run `add_one_cuda.so` in PyTorch python load/load_pytorch.py # To load and run `add_one_cuda.so` in CuPy python load/load_cupy.py # To load and run `add_one_cuda.so` in PaddlePaddle python load/load_paddle.py tvm-ffi-0.1.12/examples/stable_c_abi/000077500000000000000000000000001521067262500173565ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/stable_c_abi/CMakeLists.txt000066400000000000000000000063201521067262500221170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. cmake_minimum_required(VERSION 3.20) project(tvm_ffi_example LANGUAGES C) option(EXAMPLE_NAME "Which example to build ('kernel'/'load')" "kernel") message(STATUS "Building example: ${EXAMPLE_NAME}") # cmake-lint: disable=C0111 function (set_flat_output_dirs target_name) set(output_dir "${CMAKE_BINARY_DIR}") foreach (config IN ITEMS "" DEBUG RELEASE RELWITHDEBINFO MINSIZEREL) if (config STREQUAL "") set_property(TARGET ${target_name} PROPERTY RUNTIME_OUTPUT_DIRECTORY "${output_dir}") set_property(TARGET ${target_name} PROPERTY LIBRARY_OUTPUT_DIRECTORY "${output_dir}") set_property(TARGET ${target_name} PROPERTY ARCHIVE_OUTPUT_DIRECTORY "${output_dir}") else () set_property( TARGET ${target_name} PROPERTY RUNTIME_OUTPUT_DIRECTORY_${config} "${output_dir}" ) set_property( TARGET ${target_name} PROPERTY LIBRARY_OUTPUT_DIRECTORY_${config} "${output_dir}" ) set_property( TARGET ${target_name} PROPERTY ARCHIVE_OUTPUT_DIRECTORY_${config} "${output_dir}" ) endif () endforeach () endfunction () # Run `tvm_ffi.config --cmakedir` to find tvm-ffi package find_package( Python COMPONENTS Interpreter REQUIRED ) execute_process( COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT ) find_package(tvm_ffi CONFIG REQUIRED) if (EXAMPLE_NAME STREQUAL "kernel") # Example 1. `add_one_cpu` in C add_library(add_one_cpu SHARED src/add_one_cpu.c) target_link_libraries(add_one_cpu PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties( add_one_cpu PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/" PREFIX "" SUFFIX ".so" C_STANDARD 11 C_STANDARD_REQUIRED YES C_EXTENSIONS NO ) set_flat_output_dirs(add_one_cpu) elseif (EXAMPLE_NAME STREQUAL "load") # Example 2. Load `add_one_cpu` shared library in C add_executable(load src/load.c) target_link_libraries(load PRIVATE tvm_ffi::header tvm_ffi::shared) set_target_properties( load PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/" PREFIX "" C_STANDARD 11 C_STANDARD_REQUIRED YES C_EXTENSIONS NO ) set_flat_output_dirs(load) else () message(FATAL_ERROR "Unknown EXAMPLE_NAME option: ${EXAMPLE_NAME}. " "Expected: 'kernel' or 'load'." ) endif () tvm-ffi-0.1.12/examples/stable_c_abi/README.md000066400000000000000000000032401521067262500206340ustar00rootroot00000000000000 # Stable C ABI Code Example This directory contains all the source code for [tutorial](https://tvm.apache.org/ffi/get_started/stable_c_abi.html). ## Run everything On Linux/macOS: ```bash bash run_all.sh ``` On Windows: ```batch run_all.bat ``` ## Compile and Distribute `add_one_cpu` manually To compile the C Example: ```bash cmake . -B build -DEXAMPLE_NAME="kernel" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo ``` This produces `build/add_one_cpu.so`. ## Load the Distributed `add_one_cpu` manually To run library loading example in C: ```bash cmake . -B build -DEXAMPLE_NAME="load" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build/load ``` The executable is emitted as `build/load` (`build/load.exe` on Windows). tvm-ffi-0.1.12/examples/stable_c_abi/raw_compile.sh000077500000000000000000000032101521067262500222120ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # shellcheck disable=SC2046 set -ex BUILD_DIR=build mkdir -p $BUILD_DIR # Example 1. Compile C++ `add_one_cpu.cc` to shared library `add_one_cpu.so` # [kernel.begin] gcc -shared -O3 -std=c11 src/add_one_cpu.c \ -fPIC -fvisibility=hidden \ $(tvm-ffi-config --cflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) \ -o $BUILD_DIR/add_one_cpu.so # [kernel.end] # Example 2. Load and run `add_one_cpu.so` in C if [ -f "$BUILD_DIR/add_one_cpu.so" ]; then # [load.begin] gcc -fvisibility=hidden -O3 -std=c11 \ src/load.c \ $(tvm-ffi-config --cflags) \ $(tvm-ffi-config --ldflags) \ $(tvm-ffi-config --libs) \ -Wl,-rpath,$(tvm-ffi-config --libdir) \ -o build/load build/load # [load.end] fi tvm-ffi-0.1.12/examples/stable_c_abi/run_all.bat000066400000000000000000000027401521067262500215050ustar00rootroot00000000000000@echo off rem Licensed to the Apache Software Foundation (ASF) under one rem or more contributor license agreements. See the NOTICE file rem distributed with this work for additional information rem regarding copyright ownership. The ASF licenses this file rem to you under the Apache License, Version 2.0 (the rem "License"); you may not use this file except in compliance rem with the License. You may obtain a copy of the License at rem rem http://www.apache.org/licenses/LICENSE-2.0 rem rem Unless required by applicable law or agreed to in writing, rem software distributed under the License is distributed on an rem "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY rem KIND, either express or implied. See the License for the rem specific language governing permissions and limitations rem under the License. setlocal enabledelayedexpansion cd /d "%~dp0" if exist build rmdir /s /q build rem To compile `src/add_one_cpu.c` to shared library `build/add_one_cpu.so` cmake . -B build -DEXAMPLE_NAME="kernel" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo rem To compile `src/load.c` to executable `build/load` for /f "delims=" %%i in ('python -c "from tvm_ffi import libinfo; import pathlib; print(pathlib.Path(libinfo.find_libtvm_ffi()).parent)"') do set "TVM_FFI_LIBDIR=%%i" set "PATH=%TVM_FFI_LIBDIR%;%PATH%" cmake . -B build -DEXAMPLE_NAME="load" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo build\load.exe endlocal tvm-ffi-0.1.12/examples/stable_c_abi/run_all.sh000077500000000000000000000022371521067262500213550ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. set -ex # To compile `src/add_one_cpu.c` to shared library `build/add_one_cpu.so` cmake . -B build -DEXAMPLE_NAME="kernel" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo # To compile `src/load.c` to executable `build/load` cmake . -B build -DEXAMPLE_NAME="load" -DCMAKE_BUILD_TYPE=RelWithDebInfo cmake --build build --config RelWithDebInfo ./build/load tvm-ffi-0.1.12/examples/stable_c_abi/src/000077500000000000000000000000001521067262500201455ustar00rootroot00000000000000tvm-ffi-0.1.12/examples/stable_c_abi/src/add_one_cpu.c000066400000000000000000000046111521067262500225530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOLINTBEGIN(bugprone-reserved-identifier,google-readability-braces-around-statements) #include #include // clang-format off // [example.begin] // File: src/add_one_cpu.cc TVM_FFI_DLL_EXPORT int __tvm_ffi_add_one_cpu(void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { // Step 1. Extract inputs from `Any` // Step 1.1. Extract `x := args[0]` DLTensor* x; if (args[0].type_index == kTVMFFIDLTensorPtr) x = (DLTensor*)(args[0].v_ptr); else if (args[0].type_index == kTVMFFITensor) x = (DLTensor*)(args[0].v_c_str + sizeof(TVMFFIObject)); else { TVMFFIErrorSetRaisedFromCStr("ValueError", "Expects a Tensor input"); return -1; } // Step 1.2. Extract `y := args[1]` DLTensor* y; if (args[1].type_index == kTVMFFIDLTensorPtr) y = (DLTensor*)(args[1].v_ptr); else if (args[1].type_index == kTVMFFITensor) y = (DLTensor*)(args[1].v_c_str + sizeof(TVMFFIObject)); else { TVMFFIErrorSetRaisedFromCStr("ValueError", "Expects a Tensor output"); return -1; } // Step 2. Perform add one: y = x + 1 for (int64_t i = 0; i < x->shape[0]; ++i) { ((float*)y->data)[i] = ((float*)x->data)[i] + 1.0f; } // Step 3. Return error code 0 (success) // // Note that `result` is not set, as the output is passed in via `y` argument, // which is functionally similar to a Python function with signature: // // def add_one(x: Tensor, y: Tensor) -> None: ... return 0; } // [example.end] // clang-format on // NOLINTEND(bugprone-reserved-identifier,google-readability-braces-around-statements) tvm-ffi-0.1.12/examples/stable_c_abi/src/load.c000066400000000000000000000117111521067262500212310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOLINTBEGIN(modernize-deprecated-headers,modernize-use-nullptr,bugprone-assignment-in-if-condition,modernize-loop-convert) // [main.begin] // File: src/load.c #include #include #include // Global functions are looked up during `Initialize` and deallocated during `Finalize` // - global function: "ffi.Module.load_from_file.so" static TVMFFIObjectHandle fn_load_module = NULL; // - global function: "ffi.ModuleGetFunction" static TVMFFIObjectHandle fn_get_function = NULL; int Run(DLTensor* x, DLTensor* y) { int ret_code = 0; TVMFFIAny call_args[3] = {}; TVMFFIAny mod = {.type_index = kTVMFFINone, .v_obj = NULL}; TVMFFIAny func = {.type_index = kTVMFFINone, .v_obj = NULL}; TVMFFIAny none = {.type_index = kTVMFFINone}; // ignore the return value // Step 1. Load module // Equivalent to: // mod = tvm::ffi::Module::LoadFromFile("build/add_one_cpu.so") call_args[0] = (TVMFFIAny){.type_index = kTVMFFIRawStr, .v_c_str = "build/add_one_cpu.so"}; call_args[1] = (TVMFFIAny){.type_index = kTVMFFISmallStr, .v_int64 = 0}; if ((ret_code = TVMFFIFunctionCall(fn_load_module, call_args, 2, &mod))) goto _RAII; // Step 2. Get function `add_one_cpu` from module // Equivalent to: // func = mod->GetFunction("add_one_cpu", /*query_imports=*/false).value() call_args[0] = (TVMFFIAny){.type_index = mod.type_index, .v_obj = mod.v_obj}; call_args[1] = (TVMFFIAny){.type_index = kTVMFFIRawStr, .v_c_str = "add_one_cpu"}; call_args[2] = (TVMFFIAny){.type_index = kTVMFFIBool, .v_int64 = 0}; if ((ret_code = TVMFFIFunctionCall(fn_get_function, call_args, 3, &func))) goto _RAII; // Step 3. Call function `add_one_cpu(x, y)` // Equivalent to: // func(x, y) call_args[0] = (TVMFFIAny){.type_index = kTVMFFIDLTensorPtr, .v_ptr = x}; call_args[1] = (TVMFFIAny){.type_index = kTVMFFIDLTensorPtr, .v_ptr = y}; if ((ret_code = TVMFFIFunctionCall(func.v_ptr, call_args, 2, &none))) goto _RAII; _RAII: if (mod.type_index >= kTVMFFIObject) TVMFFIObjectDecRef(mod.v_obj); if (func.type_index >= kTVMFFIObject) TVMFFIObjectDecRef(func.v_obj); if (none.type_index >= kTVMFFIObject) TVMFFIObjectDecRef(none.v_obj); return ret_code; } // [main.end] /************* Auxiliary Logics *************/ // [aux.begin] static inline int Initialize() { int ret_code = 0; TVMFFIByteArray name_load_module = {.data = "ffi.Module.load_from_file.so", .size = 28}; TVMFFIByteArray name_get_function = {.data = "ffi.ModuleGetFunction", .size = 21}; if ((ret_code = TVMFFIFunctionGetGlobal(&name_load_module, &fn_load_module))) return ret_code; if ((ret_code = TVMFFIFunctionGetGlobal(&name_get_function, &fn_get_function))) return ret_code; return 0; } static inline void Finalize(int ret_code) { TVMFFIObjectHandle err = NULL; TVMFFIErrorCell* cell = NULL; if (fn_load_module) TVMFFIObjectDecRef(fn_load_module); if (fn_get_function) TVMFFIObjectDecRef(fn_get_function); if (ret_code) { TVMFFIErrorMoveFromRaised(&err); cell = (TVMFFIErrorCell*)((char*)(err) + sizeof(TVMFFIObject)); printf("%.*s: %.*s\n", (int)(cell->kind.size), cell->kind.data, (int)(cell->message.size), cell->message.data); } } int main() { int ret_code = 0; float x_data[5] = {1.0, 2.0, 3.0, 4.0, 5.0}; float y_data[5] = {0.0, 0.0, 0.0, 0.0, 0.0}; int64_t shape[1] = {5}; int64_t strides[1] = {1}; DLDataType f32 = {.code = kTVMFFIFloat, .bits = 32, .lanes = 1}; DLDevice cpu = {.device_type = kDLCPU, .device_id = 0}; DLTensor x = {// .data = x_data, .device = cpu, .ndim = 1, .dtype = f32, .shape = shape, .strides = strides, .byte_offset = 0}; DLTensor y = {// .data = y_data, .device = cpu, .ndim = 1, .dtype = f32, .shape = shape, .strides = strides, .byte_offset = 0}; if ((ret_code = Initialize())) goto _RAII; if ((ret_code = Run(&x, &y))) goto _RAII; printf("[ "); for (int i = 0; i < 5; ++i) printf("%f ", y_data[i]); printf("]\n"); _RAII: Finalize(ret_code); return ret_code; } // [aux.end] // NOLINTEND(modernize-deprecated-headers,modernize-use-nullptr,bugprone-assignment-in-if-condition,modernize-loop-convert) tvm-ffi-0.1.12/include/000077500000000000000000000000001521067262500145745ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/000077500000000000000000000000001521067262500154025ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/000077500000000000000000000000001521067262500161465ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/any.h000066400000000000000000000743331521067262500171200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/any.h * \brief Any value support. */ #ifndef TVM_FFI_ANY_H_ #define TVM_FFI_ANY_H_ #include #include #include #include #include namespace tvm { namespace ffi { class Any; namespace details { // Helper to perform // unsafe operations related to object struct AnyUnsafe; } // namespace details /*! * \brief AnyView allows us to take un-managed reference view of any value. */ class AnyView { protected: /*! \brief The underlying backing data of the any object */ TVMFFIAny data_; // Any can see AnyView friend class Any; public: // NOTE: the following functions use style // since they are common functions appearing in FFI. /*! * \brief Reset any view to None */ void reset() { data_.type_index = TypeIndex::kTVMFFINone; // invariance: always set the union padding part to 0 data_.zero_padding = 0; data_.v_int64 = 0; } /*! * \brief Swap this AnyView with another AnyView * \param other The other AnyView */ TVM_FFI_INLINE void swap(AnyView& other) noexcept { std::swap(data_, other.data_); } /*! \return the internal type index */ TVM_FFI_INLINE int32_t type_index() const noexcept { return data_.type_index; } /*! \brief Default constructor */ AnyView() { data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } ~AnyView() = default; // constructors from any view /*! \brief Copy constructor */ AnyView(const AnyView&) = default; /*! \brief Copy assignment operator */ AnyView& operator=(const AnyView&) = default; /*! \brief Move constructor */ AnyView(AnyView&& other) noexcept = default; AnyView& operator=(AnyView&& other) noexcept = default; /*! * \brief Constructor from a general type. * \tparam T The type to convert from. * \param other The value to convert from. */ template ::convert_enabled>> AnyView(const T& other) { // NOLINT(*) TypeTraits::CopyToAnyView(other, &data_); } /*! * \brief Assign from a general type. * \tparam T The type to convert from. * \param other The value to convert from. */ template ::convert_enabled>> TVM_FFI_INLINE AnyView& operator=(const T& other) { // NOLINT(*) // copy-and-swap idiom AnyView(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief Try to see if we can reinterpret the AnyView to as T object. * * \tparam T The type to cast to. * \return The casted value, or std::nullopt if the cast is not possible. * \note This function won't try run type conversion (use try_cast for that purpose). */ template ::convert_enabled>> TVM_FFI_INLINE std::optional as() const { if (TypeTraits::CheckAnyStrict(&data_)) { return TypeTraits::CopyFromAnyViewAfterCheck(&data_); } else { return std::optional(std::nullopt); } } /*! * \brief Shortcut of as Object to cast to a const pointer when T is an Object. * * \tparam T The object type. * \return The requested pointer, returns nullptr if type mismatches. */ template >> TVM_FFI_INLINE const T* as() const { return this->as().value_or(nullptr); } /*! * \brief Cast to a type T. * * \tparam T The type to cast to. * \return The casted value, or throws an exception if the cast is not possible. */ template ::convert_enabled>> TVM_FFI_INLINE T cast() const { std::optional opt = TypeTraits::TryCastFromAnyView(&data_); if (!opt.has_value()) { TVM_FFI_THROW(TypeError) << "Cannot convert from type `" << TypeTraits::GetMismatchTypeInfo(&data_) << "` to `" << TypeTraits::TypeStr() << "`"; } return *std::move(opt); } /*! * \brief Try to cast to a type T, return std::nullopt if the cast is not possible. * * \tparam T The type to cast to. * \return The casted value, or std::nullopt if the cast is not possible. */ template ::convert_enabled>> TVM_FFI_INLINE std::optional try_cast() const { return TypeTraits::TryCastFromAnyView(&data_); } // comparison with nullptr TVM_FFI_INLINE bool operator==(std::nullptr_t) const noexcept { return data_.type_index == TypeIndex::kTVMFFINone; } TVM_FFI_INLINE bool operator!=(std::nullptr_t) const noexcept { return data_.type_index != TypeIndex::kTVMFFINone; } /*! * \brief Get the type key of the Any * \return The type key of the Any */ TVM_FFI_INLINE std::string GetTypeKey() const { return TypeIndexToTypeKey(data_.type_index); } // The following functions are only used for testing purposes /*! * \return The underlying supporting data of any view * \note This function is used only for testing purposes. */ TVM_FFI_INLINE TVMFFIAny CopyToTVMFFIAny() const { return data_; } /*! * \return Create an AnyView from TVMFFIAny * \param data the underlying ffi data. */ TVM_FFI_INLINE static AnyView CopyFromTVMFFIAny(TVMFFIAny data) { AnyView view; view.data_ = data; return view; } }; namespace details { /*! * \brief Helper function to inplace convert any view to any. * \param data The pointer that represents the format as any view. * \param extra_any_bytes Indicate that the data may contain extra bytes following * the TVMFFIAny data structure. This is reserved for future possible optimizations * of small-string and extended any object. */ TVM_FFI_INLINE void InplaceConvertAnyViewToAny(TVMFFIAny* data, [[maybe_unused]] size_t extra_any_bytes = 0) { if (data->type_index >= TVMFFITypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::IncRefObjectHandle(data->v_obj); } else if (data->type_index >= TypeIndex::kTVMFFIRawStr) { if (data->type_index == TypeIndex::kTVMFFIRawStr) { // convert raw string to owned string object String temp(data->v_c_str); TypeTraits::MoveToAny(std::move(temp), data); } else if (data->type_index == TypeIndex::kTVMFFIByteArrayPtr) { // convert byte array to owned bytes object Bytes temp(*static_cast(data->v_ptr)); TypeTraits::MoveToAny(std::move(temp), data); } else if (data->type_index == TypeIndex::kTVMFFIObjectRValueRef) { // convert rvalue ref to owned object Object** obj_addr = static_cast(data->v_ptr); TVM_FFI_ICHECK(obj_addr[0] != nullptr) << "RValueRef already moved"; ObjectRef temp(details::ObjectUnsafe::ObjectPtrFromOwned(obj_addr[0])); // set the rvalue ref to nullptr to avoid double move obj_addr[0] = nullptr; TypeTraits::MoveToAny(std::move(temp), data); } } } } // namespace details /*! * \brief Managed Any that takes strong reference to a value. * * \note Develooper invariance: the TVMFFIAny data_ * in the Any can be safely used in AnyView. */ class Any { protected: /*! \brief The underlying backing data of the any object */ TVMFFIAny data_; public: /*! * \brief Reset any to None */ TVM_FFI_INLINE void reset() { if (data_.type_index >= TVMFFITypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::DecRefObjectHandle(data_.v_obj); } data_.type_index = TVMFFITypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } /*! * \brief Swap this Any with another Any * \param other The other Any */ TVM_FFI_INLINE void swap(Any& other) noexcept { std::swap(data_, other.data_); } /*! \return the internal type index */ TVM_FFI_INLINE int32_t type_index() const noexcept { return data_.type_index; } /*! * \brief Default constructor */ Any() { data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } /*! * \brief Destructor */ ~Any() { this->reset(); } /*! * \brief Constructor from another Any * \param other The other Any */ Any(const Any& other) : data_(other.data_) { if (data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::IncRefObjectHandle(data_.v_obj); } } /*! * \brief Move constructor from another Any * \param other The other Any */ Any(Any&& other) noexcept : data_(other.data_) { other.data_.type_index = TypeIndex::kTVMFFINone; other.data_.zero_padding = 0; other.data_.v_int64 = 0; } /*! * \brief Assign from another Any * \param other The other Any */ TVM_FFI_INLINE Any& operator=(const Any& other) { // copy-and-swap idiom Any(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief Move assign from another Any * \param other The other Any */ TVM_FFI_INLINE Any& operator=(Any&& other) noexcept { // copy-and-swap idiom Any(std::move(other)).swap(*this); // NOLINT(*) return *this; } /*! * \brief Constructor from another AnyView * \param other The other AnyView */ Any(const AnyView& other) : data_(other.data_) { // NOLINT(*) details::InplaceConvertAnyViewToAny(&data_); } /*! * \brief Assign from another AnyView * \param other The other AnyView */ TVM_FFI_INLINE Any& operator=(const AnyView& other) { // copy-and-swap idiom Any(other).swap(*this); // NOLINT(*) return *this; } /*! \brief Any can be converted to AnyView in zero cost. */ operator AnyView() const { // NOLINT(google-explicit-constructor) return AnyView::CopyFromTVMFFIAny(data_); } /*! * \brief Constructor from a general type * \tparam T The value type of the other */ template ::convert_enabled>> Any(T other) { // NOLINT(*) TypeTraits::MoveToAny(std::move(other), &data_); } /*! * \brief Assignment from a general type * \tparam T The value type of the other */ template ::convert_enabled>> TVM_FFI_INLINE Any& operator=(T other) { // NOLINT(*) // copy-and-swap idiom Any(std::move(other)).swap(*this); // NOLINT(*) return *this; } /** * \brief Try to reinterpret the Any as a type T, return std::nullopt if it is not possible. * * \tparam T The type to cast to. * \return The casted value, or std::nullopt if the cast is not possible. * \note This function won't try to run type conversion (use try_cast for that purpose). */ template ::storage_enabled || std::is_same_v>> TVM_FFI_INLINE std::optional as() && { if constexpr (std::is_same_v) { return std::move(*this); } else { if (TypeTraits::CheckAnyStrict(&data_)) { return TypeTraits::MoveFromAnyAfterCheck(&data_); } else { return std::optional(std::nullopt); } } } /** * \brief Try to reinterpret the Any as a type T, return std::nullopt if it is not possible. * * \tparam T The type to cast to. * \return The casted value, or std::nullopt if the cast is not possible. * \note This function won't try to run type conversion (use try_cast for that purpose). */ template ::convert_enabled || std::is_same_v>> TVM_FFI_INLINE std::optional as() const& { if constexpr (std::is_same_v) { return *this; } else { if (TypeTraits::CheckAnyStrict(&data_)) { return TypeTraits::CopyFromAnyViewAfterCheck(&data_); } else { return std::optional(std::nullopt); } } } /*! * \brief Shortcut of as Object to cast to a const pointer when T is an Object. * * \tparam T The object type. * \return The requested pointer, returns nullptr if type mismatches. */ template >> TVM_FFI_INLINE const T* as() const& { return this->as().value_or(nullptr); } /** * \brief Cast to a type T, throw an exception if the cast is not possible. * * \tparam T The type to cast to. */ template ::convert_enabled>> TVM_FFI_INLINE T cast() const& { std::optional opt = TypeTraits::TryCastFromAnyView(&data_); if (!opt.has_value()) { TVM_FFI_THROW(TypeError) << "Cannot convert from type `" << TypeTraits::GetMismatchTypeInfo(&data_) << "` to `" << TypeTraits::TypeStr() << "`"; } return *std::move(opt); } /** * \brief Cast to a type T, throw an exception if the cast is not possible. * * \tparam T The type to cast to. */ template ::storage_enabled>> TVM_FFI_INLINE T cast() && { if (TypeTraits::CheckAnyStrict(&data_)) { return TypeTraits::MoveFromAnyAfterCheck(&data_); } // slow path, try to do fallback convert std::optional opt = TypeTraits::TryCastFromAnyView(&data_); if (!opt.has_value()) { TVM_FFI_THROW(TypeError) << "Cannot convert from type `" << TypeTraits::GetMismatchTypeInfo(&data_) << "` to `" << TypeTraits::TypeStr() << "`"; } return *std::move(opt); } /** * \brief Try to cast to a type T. * * \tparam T The type to cast to. * \return The casted value, or std::nullopt if the cast is not possible. * \note use STL name since it to be more consistent with cast API. */ template ::convert_enabled || std::is_same_v>> TVM_FFI_INLINE std::optional try_cast() const { if constexpr (std::is_same_v) { return *this; } else { return TypeTraits::TryCastFromAnyView(&data_); } } /*! * \brief Check if the two Any are same type and value in shallow comparison. * \param other The other Any * \return True if the two Any are same type and value, false otherwise. */ TVM_FFI_INLINE bool same_as(const Any& other) const noexcept { return data_.type_index == other.data_.type_index && data_.zero_padding == other.data_.zero_padding && data_.v_int64 == other.data_.v_int64; } /*! * \brief Check if any and ObjectRef are same type and value in shallow comparison. * \param other The other ObjectRef * \return True if the two Any are same type and value, false otherwise. */ TVM_FFI_INLINE bool same_as(const ObjectRef& other) const noexcept { if (other.get() != nullptr) { return (data_.type_index == other->type_index() && reinterpret_cast(data_.v_obj) == other.get()); } else { return data_.type_index == TypeIndex::kTVMFFINone; } } TVM_FFI_INLINE bool operator==(std::nullptr_t) const noexcept { return data_.type_index == TypeIndex::kTVMFFINone; } TVM_FFI_INLINE bool operator!=(std::nullptr_t) const noexcept { return data_.type_index != TypeIndex::kTVMFFINone; } /*! * \brief Get the type key of the Any * \return The type key of the Any */ TVM_FFI_INLINE std::string GetTypeKey() const { return TypeIndexToTypeKey(data_.type_index); } friend struct details::AnyUnsafe; friend struct AnyHash; friend struct AnyEqual; }; // layout assert to ensure we can freely cast between the two types static_assert(sizeof(AnyView) == sizeof(TVMFFIAny)); static_assert(sizeof(Any) == sizeof(TVMFFIAny)); // AnyView is a non-owning, POD-shaped view of TVMFFIAny. Keeping it // trivially copyable lets the C++ ABI pass it through registers and // match the C ABI for struct passing. Any user-defined copy/move ctor // or non-trivial destructor on AnyView would silently regress this // calling convention — catch it here at compile time. static_assert(std::is_trivially_copyable_v, "AnyView must be trivially copyable."); namespace details { template struct Type2Str { static std::string v() { return TypeTraitsNoCR::TypeStr(); } }; template <> struct Type2Str { static std::string v() { return "Any"; } }; template <> struct Type2Str { static std::string v() { return "Any"; } }; template <> struct Type2Str { static std::string v() { return "Any"; } }; template <> struct Type2Str { static std::string v() { return "AnyView"; } }; template <> struct Type2Str { static std::string v() { return "AnyView"; } }; template <> struct Type2Str { static std::string v() { return "AnyView"; } }; template <> struct Type2Str { static std::string v() { return "void"; } }; // Extra unsafe method to help any manipulation struct AnyUnsafe : public ObjectUnsafe { // FFI related operations TVM_FFI_INLINE static TVMFFIAny MoveAnyToTVMFFIAny(Any&& any) { TVMFFIAny result = any.data_; any.data_.type_index = TypeIndex::kTVMFFINone; any.data_.zero_padding = 0; any.data_.v_int64 = 0; return result; } TVM_FFI_INLINE static Any MoveTVMFFIAnyToAny(TVMFFIAny* data) { Any any; any.data_ = *data; data->type_index = TypeIndex::kTVMFFINone; data->zero_padding = 0; data->v_int64 = 0; return any; } template TVM_FFI_INLINE static bool CheckAnyStrict(const Any& ref) { return TypeTraits::CheckAnyStrict(&(ref.data_)); } template TVM_FFI_INLINE static T CopyFromAnyViewAfterCheck(const Any& ref) { if constexpr (!std::is_same_v) { return TypeTraits::CopyFromAnyViewAfterCheck(&(ref.data_)); } else { return ref; } } template TVM_FFI_INLINE static T MoveFromAnyAfterCheck(Any&& ref) { if constexpr (!std::is_same_v) { return TypeTraits::MoveFromAnyAfterCheck(&(ref.data_)); } else { return std::move(ref); } } TVM_FFI_INLINE static Object* ObjectPtrFromAnyAfterCheck(const Any& ref) { return reinterpret_cast(ref.data_.v_obj); } TVM_FFI_INLINE static const TVMFFIAny* TVMFFIAnyPtrFromAny(const Any& ref) { return &(ref.data_); } template TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const Any& ref) { return TypeTraits::GetMismatchTypeInfo(&(ref.data_)); } }; } // namespace details /*! \brief String-aware Any equal functor */ struct AnyHash { public: /*! * \brief Calculate the hash code of an Any * \param a The given Any * \return Hash code of a, string hash for strings and pointer address otherwise. */ uint64_t operator()(const Any& src) const { if (src.data_.type_index == TypeIndex::kTVMFFISmallStr) { // for small string, we use the same type key hash as normal string // so heap allocated string and on stack string will have the same hash return details::StableHashCombine(TypeIndex::kTVMFFIStr, details::StableHashSmallStrBytes(&src.data_)); } else if (src.data_.type_index == TypeIndex::kTVMFFISmallBytes) { // use byte the same type key as bytes return details::StableHashCombine(TypeIndex::kTVMFFIBytes, details::StableHashSmallStrBytes(&src.data_)); } else if (src.data_.type_index == TypeIndex::kTVMFFIStr || src.data_.type_index == TypeIndex::kTVMFFIBytes) { const details::BytesObjBase* src_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(src); return details::StableHashCombine(src.data_.type_index, details::StableHashBytes(src_str->data, src_str->size)); } else { if (src.data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { static const TVMFFITypeAttrColumn* custom_hash_column = GetAnyHashTypeAttrColumn(); if (custom_hash_column != nullptr) { int32_t offset = src.data_.type_index - custom_hash_column->begin_index; if (offset >= 0 && offset < custom_hash_column->size) { const TVMFFIAny& custom_any_hash = custom_hash_column->data[offset]; if (custom_any_hash.type_index != TypeIndex::kTVMFFINone) { return details::StableHashCombine(src.data_.type_index, CallCustomAnyHash(custom_any_hash, src)); } } } } return details::StableHashCombine(src.data_.type_index, src.data_.v_uint64); } } private: /*! * \brief Get the type attribute column for any hash. * \return The type attribute column for any hash. */ static const TVMFFITypeAttrColumn* GetAnyHashTypeAttrColumn() { constexpr const char* kAttrName = "__any_hash__"; TVMFFIByteArray attr_name = TVMFFIByteArray{kAttrName, std::char_traits::length(kAttrName)}; return TVMFFIGetTypeAttrColumn(&attr_name); } /*! * \brief Call the custom any hash function registered in type attribute column. * \param custom_any_hash The custom any hash function object or function pointer. * \param src The source Any. * \return The hash value. */ static uint64_t CallCustomAnyHash(const TVMFFIAny& custom_any_hash, const Any& src) { // NOTE: we explicitly use low-level ABI here since we do not want to have dep on function.h // it also keeps the logic simple if (custom_any_hash.type_index == TypeIndex::kTVMFFIOpaquePtr) { // we allow this attribute to be a function pointer for fast path using FCustomAnyHashPtr = int64_t (*)(const Any&); FCustomAnyHashPtr hash_func = reinterpret_cast(custom_any_hash.v_ptr); TVM_FFI_ICHECK_NOTNULL(hash_func); return static_cast( (*hash_func)(src)); // NOLINT(clang-analyzer-core.CallAndMessage) } else { // alternatively it can be a ffi.Function object. TVM_FFI_ICHECK_EQ(custom_any_hash.type_index, TypeIndex::kTVMFFIFunction); TVMFFIAny arg = src.data_; TVMFFIAny result; result.type_index = TypeIndex::kTVMFFINone; result.zero_padding = 0; result.v_int64 = 0; TVMFFIFunctionCell* func_cell = TVMFFIFunctionGetCellPtr(custom_any_hash.v_obj); if (func_cell->cpp_call != nullptr) { // Fast path: invoke C++ ABI call directly when available. using FCppCall = void (*)(const void*, const TVMFFIAny*, int32_t, TVMFFIAny*); reinterpret_cast(func_cell->cpp_call)(custom_any_hash.v_obj, &arg, 1, &result); } else { if (func_cell->safe_call(custom_any_hash.v_obj, &arg, 1, &result) != 0) { throw details::MoveFromSafeCallRaised(); } } Any result_any = details::AnyUnsafe::MoveTVMFFIAnyToAny(&result); TVM_FFI_ICHECK_EQ(result_any.type_index(), TypeIndex::kTVMFFIInt); return static_cast(result_any.data_.v_int64); } } }; /*! \brief String-aware Any hash functor */ struct AnyEqual { public: /*! * \brief Check if the two Any are equal * \param lhs left operand. * \param rhs right operand * \return String equality if both are strings, pointer address equality otherwise. */ bool operator()(const Any& lhs, const Any& rhs) const { // header with type index const int64_t* lhs_as_int64 = reinterpret_cast(&lhs.data_); const int64_t* rhs_as_int64 = reinterpret_cast(&rhs.data_); static_assert(sizeof(TVMFFIAny) == 16); // fast path, check byte equality if (lhs_as_int64[0] == rhs_as_int64[0] && lhs_as_int64[1] == rhs_as_int64[1]) { return true; } // common false case type index match, in this case we only need to pay attention to string // equality if (lhs.data_.type_index == rhs.data_.type_index) { // specialy handle string hash if (lhs.data_.type_index == TypeIndex::kTVMFFIStr || lhs.data_.type_index == TypeIndex::kTVMFFIBytes) { const details::BytesObjBase* lhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); const details::BytesObjBase* rhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs_str->data, rhs_str->data, lhs_str->size, rhs_str->size); } if (lhs.data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { static const TVMFFITypeAttrColumn* custom_equal_column = GetAnyEqualTypeAttrColumn(); if (custom_equal_column != nullptr) { int32_t offset = lhs.data_.type_index - custom_equal_column->begin_index; if (offset >= 0 && offset < custom_equal_column->size) { const TVMFFIAny& custom_any_equal = custom_equal_column->data[offset]; if (custom_any_equal.type_index != TypeIndex::kTVMFFINone) { return CallCustomAnyEqual(custom_any_equal, lhs, rhs); } } } } return false; } else { // type_index mismatch, if index is not string, return false if (lhs.data_.type_index != kTVMFFIStr && lhs.data_.type_index != kTVMFFISmallStr && lhs.data_.type_index != kTVMFFISmallBytes && lhs.data_.type_index != kTVMFFIBytes) { return false; } // small string and normal string comparison if (lhs.data_.type_index == kTVMFFIStr && rhs.data_.type_index == kTVMFFISmallStr) { const details::BytesObjBase* lhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); return Bytes::memequal(lhs_str->data, rhs.data_.v_bytes, lhs_str->size, rhs.data_.small_str_len); } if (lhs.data_.type_index == kTVMFFISmallStr && rhs.data_.type_index == kTVMFFIStr) { const details::BytesObjBase* rhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs.data_.v_bytes, rhs_str->data, lhs.data_.small_str_len, rhs_str->size); } if (lhs.data_.type_index == kTVMFFIBytes && rhs.data_.type_index == kTVMFFISmallBytes) { const details::BytesObjBase* lhs_bytes = details::AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); return Bytes::memequal(lhs_bytes->data, rhs.data_.v_bytes, lhs_bytes->size, rhs.data_.small_str_len); } if (lhs.data_.type_index == kTVMFFISmallBytes && rhs.data_.type_index == kTVMFFIBytes) { const details::BytesObjBase* rhs_bytes = details::AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs.data_.v_bytes, rhs_bytes->data, lhs.data_.small_str_len, rhs_bytes->size); } return false; } } private: /*! * \brief Get the type attribute column for any equal. * \return The type attribute column for any equal. */ static const TVMFFITypeAttrColumn* GetAnyEqualTypeAttrColumn() { constexpr const char* kAttrName = "__any_equal__"; TVMFFIByteArray attr_name = TVMFFIByteArray{kAttrName, std::char_traits::length(kAttrName)}; return TVMFFIGetTypeAttrColumn(&attr_name); } /*! * \brief Call the custom any equal function registered in type attribute column. * \param custom_any_equal The custom any equal function object or function pointer. * \param lhs The left-hand side Any. * \param rhs The right-hand side Any. * \return The equality result. */ static bool CallCustomAnyEqual(const TVMFFIAny& custom_any_equal, const Any& lhs, const Any& rhs) { // NOTE: we explicitly use low-level ABI here since we do not want to have dep on function.h // it also keeps the logic simple if (custom_any_equal.type_index == TypeIndex::kTVMFFIOpaquePtr) { // we allow this attribute to be a function pointer for fast path using FCustomAnyEqualPtr = bool (*)(const Any&, const Any&); FCustomAnyEqualPtr equal_func = reinterpret_cast(custom_any_equal.v_ptr); TVM_FFI_ICHECK_NOTNULL(equal_func); return (*equal_func)(lhs, rhs); // NOLINT(clang-analyzer-core.CallAndMessage) } else { // alternatively it can be a ffi.Function object. TVM_FFI_ICHECK_EQ(custom_any_equal.type_index, TypeIndex::kTVMFFIFunction); TVMFFIAny args[2] = {lhs.data_, rhs.data_}; TVMFFIAny result; result.type_index = TypeIndex::kTVMFFINone; result.zero_padding = 0; result.v_int64 = 0; TVMFFIFunctionCell* func_cell = TVMFFIFunctionGetCellPtr(custom_any_equal.v_obj); if (func_cell->cpp_call != nullptr) { // Fast path: invoke C++ ABI call directly when available. using FCppCall = void (*)(const void*, const TVMFFIAny*, int32_t, TVMFFIAny*); reinterpret_cast(func_cell->cpp_call)(custom_any_equal.v_obj, args, 2, &result); } else { if (func_cell->safe_call(custom_any_equal.v_obj, args, 2, &result) != 0) { throw details::MoveFromSafeCallRaised(); } } Any result_any = details::AnyUnsafe::MoveTVMFFIAnyToAny(&result); TVM_FFI_ICHECK(result_any.type_index() == TypeIndex::kTVMFFIBool); return result_any.data_.v_int64 != 0; } } }; } // namespace ffi // Expose to the tvm namespace for usability // Rationale: no ambiguity even in root using tvm::ffi::Any; using tvm::ffi::AnyView; } // namespace tvm #endif // TVM_FFI_ANY_H_ tvm-ffi-0.1.12/include/tvm/ffi/base_details.h000066400000000000000000000302151521067262500207370ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/base_details.h * \brief Internal detail utils that can be used by files in tvm/ffi. * \note details headers are for internal use only * and not to be directly used by user. */ #ifndef TVM_FFI_BASE_DETAILS_H_ #define TVM_FFI_BASE_DETAILS_H_ #include #include #include #include #if defined(_MSC_VER) #ifndef WIN32_LEAN_AND_MEAN #define WIN32_LEAN_AND_MEAN #endif #ifndef NOMINMAX #define NOMINMAX #endif #include #if (defined(_M_ARM64) || defined(_ARM64_) || defined(_M_ARM64EC)) && \ !defined(_InlineInterlockedAdd64) #define _InlineInterlockedAdd64 InterlockedAdd64 #endif #ifdef ERROR #undef ERROR #endif #endif /// \cond Doxygen_Suppress #if defined(_MSC_VER) #define TVM_FFI_INLINE [[msvc::forceinline]] inline #else #define TVM_FFI_INLINE [[gnu::always_inline]] inline #endif /*! * \brief Macro helper to force a function not to be inlined. * It is only used in places that we know not inlining is good, * e.g. some logging functions. */ #if defined(_MSC_VER) #define TVM_FFI_NO_INLINE [[msvc::noinline]] #else #define TVM_FFI_NO_INLINE [[gnu::noinline]] #endif #if defined(_MSC_VER) #define TVM_FFI_UNREACHABLE() __assume(false) #else #define TVM_FFI_UNREACHABLE() __builtin_unreachable() #endif /*! * \brief Mark a function as cold so the toolchain places it in a * separate cold region of `.text`. Apply to functions that only * run on error / setup / teardown paths. * * On GCC and Clang, expands to `[[gnu::cold]]`, which emits the * function into a per-TU `.text.unlikely` section. The default GNU * linker script gathers `.text.unlikely.*` into a contiguous slot * inside `.text`, so cold-marked functions cluster away from hot * code without any additional CMake flags. On MSVC the attribute * does not exist and the macro is a no-op. */ #if defined(__GNUC__) || defined(__clang__) #define TVM_FFI_COLD_CODE [[gnu::cold]] #else #define TVM_FFI_COLD_CODE #endif /*! * \brief Branch-prediction / layout hint that the condition is unlikely * to be true. Use on error-checking branches to keep the hot * fall-through contiguous and push the error-handling block to * the function tail. * * if (TVM_FFI_PREDICT_FALSE(rc != 0)) { ...error... } * * On GCC/Clang, expands to `__builtin_expect((cond), 0)`. On MSVC, * expands to `(cond)` (no equivalent builtin; modern MSVC does its own * profile-driven block reordering). */ #if defined(__GNUC__) || defined(__clang__) #define TVM_FFI_PREDICT_FALSE(cond) (__builtin_expect(static_cast(cond), 0)) #define TVM_FFI_PREDICT_TRUE(cond) (__builtin_expect(static_cast(cond), 1)) #else #define TVM_FFI_PREDICT_FALSE(cond) (cond) #define TVM_FFI_PREDICT_TRUE(cond) (cond) #endif /*! * \brief Translates into __builtin_assume / __assume / __attribute__((assume)). * * Use ONLY when the external invariant guarantees cond. The compiler * will remove all paths inconsistent with cond. This is not an * assertion or check -- using on a wrong cond will result in * undefined behavior. cond must be side-effect-free. */ #if defined(__clang__) #define TVM_FFI_UNSAFE_ASSUME(cond) __builtin_assume(cond) #elif defined(__GNUC__) // GCC 13+ supports __attribute__((assume(...))); fall back to the void-cast // no-op for older GCC where __builtin_assume is absent. #if __GNUC__ >= 13 #define TVM_FFI_UNSAFE_ASSUME(cond) __attribute__((assume(cond))) #else #define TVM_FFI_UNSAFE_ASSUME(cond) static_cast(0) #endif #elif defined(_MSC_VER) #define TVM_FFI_UNSAFE_ASSUME(cond) __assume(cond) #else #define TVM_FFI_UNSAFE_ASSUME(cond) static_cast(0) #endif #define TVM_FFI_STR_CONCAT_(__x, __y) __x##__y #define TVM_FFI_STR_CONCAT(__x, __y) TVM_FFI_STR_CONCAT_(__x, __y) #if defined(__GNUC__) || defined(__clang__) #define TVM_FFI_FUNC_SIG __PRETTY_FUNCTION__ #elif defined(_MSC_VER) #define TVM_FFI_FUNC_SIG __FUNCSIG__ #else #define TVM_FFI_FUNC_SIG __func__ #endif /// \endcond #if defined(TVM_FFI_DOXYGEN_MODE) /*! * \brief Macro that defines a block that will be called during static initialization. * * \code{.cpp} * TVM_FFI_STATIC_INIT_BLOCK() { * RegisterFunctions(); * } * \endcode */ #define TVM_FFI_STATIC_INIT_BLOCK() #elif defined(__GNUC__) // gcc and clang: attribute constructor #define TVM_FFI_STATIC_INIT_BLOCK_DEF_(FnName) __attribute__((constructor)) static void FnName() #define TVM_FFI_STATIC_INIT_BLOCK() \ TVM_FFI_STATIC_INIT_BLOCK_DEF_(TVM_FFI_STR_CONCAT(__TVMFFIStaticInitFunc, __COUNTER__)) #else // other compilers: use the variable trick #define TVM_FFI_STATIC_INIT_BLOCK_DEF_(FnName, RegVar) \ static void FnName(); \ [[maybe_unused]] static inline int RegVar = []() { \ FnName(); \ return 0; \ }(); \ static void FnName() #define TVM_FFI_STATIC_INIT_BLOCK() \ TVM_FFI_STATIC_INIT_BLOCK_DEF_(TVM_FFI_STR_CONCAT(__TVMFFIStaticInitFunc, __COUNTER__), \ TVM_FFI_STR_CONCAT(__TVMFFIStaticInitReg, __COUNTER__)) #endif /// \cond Doxygen_Suppress /* * \brief Define the default copy/move constructor and assign operator * \param TypeName The class typename. */ #define TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TypeName) \ TypeName(const TypeName& other) = default; /* NOLINT(bugprone-macro-parentheses) */ \ TypeName(TypeName&& other) noexcept = default; /* NOLINT(bugprone-macro-parentheses) */ \ TypeName& operator=(const TypeName& other) = default; /* NOLINT(bugprone-macro-parentheses) */ \ TypeName& operator=(TypeName&& other) noexcept = default; /* NOLINT(bugprone-macro-parentheses)*/ /*! * \brief marks the begining of a C call that logs exception */ #define TVM_FFI_LOG_EXCEPTION_CALL_BEGIN() \ try { \ (void)0 /*! * \brief Marks the end of a C call that logs exception */ #define TVM_FFI_LOG_EXCEPTION_CALL_END(Name) \ } \ catch (const std::exception& err) { \ std::cerr << "Exception caught during " << #Name << ":\n" << err.what() << std::endl; \ exit(-1); \ } /*! * \brief Clear the padding parts so we can safely use v_int64 for hash * and equality check even when the value stored is a pointer. * * This macro is used to clear the padding parts for hash and equality check * in 32bit platform. */ #define TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result) \ if constexpr (sizeof(void*) != sizeof(int64_t)) { \ (result)->v_int64 = 0; \ } namespace tvm { namespace ffi { namespace details { // for each iterator struct for_each_dispatcher { template static void run(std::index_sequence, const F& f, Args&&... args) { // NOLINT(*) (f(I, std::forward(args)), ...); } }; template void for_each(const F& f, Args&&... args) { // NOLINT(*) for_each_dispatcher::run(std::index_sequence_for{}, f, std::forward(args)...); } /*! * \brief hash an object and combines uint64_t key with previous keys * * This hash function is stable across platforms. * * \param key The left operand. * \param value The right operand. * \return the combined result. */ template , bool> = true> TVM_FFI_INLINE uint64_t StableHashCombine(uint64_t key, const T& value) { // XXX: do not use std::hash in this function. This hash must be stable // across different platforms and std::hash is implementation dependent. return key ^ (uint64_t(value) + 0x9e3779b9 + (key << 6) + (key >> 2)); } /*! * \brief Hash the binary bytes * \param data The data pointer * \param size The size of the bytes. * \return the hash value. */ TVM_FFI_INLINE uint64_t StableHashBytes(const void* data_ptr, size_t size) { // NOLINTBEGIN(clang-analyzer-security.ArrayBound) const char* data = reinterpret_cast(data_ptr); const constexpr uint64_t kMultiplier = 1099511628211ULL; const constexpr uint64_t kMod = 2147483647ULL; union Union { uint8_t a[8]; uint64_t b; } u; static_assert(sizeof(Union) == sizeof(uint64_t), "sizeof(Union) != sizeof(uint64_t)"); const char* it = data; const char* end = it + size; uint64_t result = 0; if constexpr (TVM_FFI_IO_NO_ENDIAN_SWAP) { // if alignment requirement is met, directly use load if (reinterpret_cast(it) % 8 == 0) { for (; it + 8 <= end; it += 8) { u.b = *reinterpret_cast(it); result = (result * kMultiplier + u.b) % kMod; } } else { // unaligned version for (; it + 8 <= end; it += 8) { u.a[0] = it[0]; u.a[1] = it[1]; u.a[2] = it[2]; u.a[3] = it[3]; u.a[4] = it[4]; u.a[5] = it[5]; u.a[6] = it[6]; u.a[7] = it[7]; result = (result * kMultiplier + u.b) % kMod; } } } else { // need endian swap for (; it + 8 <= end; it += 8) { u.a[0] = it[7]; u.a[1] = it[6]; u.a[2] = it[5]; u.a[3] = it[4]; u.a[4] = it[3]; u.a[5] = it[2]; u.a[6] = it[1]; u.a[7] = it[0]; result = (result * kMultiplier + u.b) % kMod; } } if (it < end) { u.b = 0; uint8_t* a = u.a; if (it + 4 <= end) { a[0] = it[0]; a[1] = it[1]; a[2] = it[2]; a[3] = it[3]; it += 4; a += 4; } if (it + 2 <= end) { a[0] = it[0]; a[1] = it[1]; it += 2; a += 2; } if (it + 1 <= end) { a[0] = it[0]; } if constexpr (!TVM_FFI_IO_NO_ENDIAN_SWAP) { std::swap(u.a[0], u.a[7]); std::swap(u.a[1], u.a[6]); std::swap(u.a[2], u.a[5]); std::swap(u.a[3], u.a[4]); } result = (result * kMultiplier + u.b) % kMod; } // NOLINTEND(clang-analyzer-security.ArrayBound) return result; } /*! * \brief Same as StableHashBytes, but for small string data. * \param data The data pointer * \return the hash value. */ TVM_FFI_INLINE uint64_t StableHashSmallStrBytes(const TVMFFIAny* data) { if constexpr (TVM_FFI_IO_NO_ENDIAN_SWAP) { // fast path, no endian swap, simply hash as uint64_t const constexpr uint64_t kMod = 2147483647ULL; return data->v_uint64 % kMod; } return StableHashBytes(reinterpret_cast(data), sizeof(data->v_uint64)); } /*! * \brief Helper to generate a JSON-based type schema for a given type. * \tparam T The type to generate the schema for. Assuming `T` is not * const-qualified or reference-qualified. */ template struct TypeSchemaImpl; /*! * \brief Helper to generate a JSON-based type schema for a given type. * \tparam T The type to generate the schema for. * \note This type removes const and reference qualifiers from `T` before * passing it to `TypeSchemaImpl`. */ template using TypeSchema = TypeSchemaImpl>>; } // namespace details } // namespace ffi } // namespace tvm /// \endcond #endif // TVM_FFI_BASE_DETAILS_H_ tvm-ffi-0.1.12/include/tvm/ffi/c_api.h000066400000000000000000001547531521067262500174110ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOLINTBEGIN(modernize-use-using,bugprone-reserved-identifier,modernize-deprecated-headers) /* * \file tvm/ffi/c_api.h * \brief This file defines the C convention of the FFI convention */ #ifndef TVM_FFI_C_API_H_ #define TVM_FFI_C_API_H_ #include #include // Macros to do weak linking #ifdef _MSC_VER #define TVM_FFI_WEAK __declspec(selectany) #else #define TVM_FFI_WEAK __attribute__((weak)) #endif // Defines two macros // TVM_FFI_DLL: marks the function as a DLL export/import // depending on whether TVM_FFI_EXPORTS is defined // TVM_FFI_DLL_EXPORT: always marks the function as a DLL export #if !defined(TVM_FFI_DLL) && defined(__EMSCRIPTEN__) #include #define TVM_FFI_DLL EMSCRIPTEN_KEEPALIVE #define TVM_FFI_DLL_EXPORT EMSCRIPTEN_KEEPALIVE #endif #if !defined(TVM_FFI_DLL) && defined(_MSC_VER) #ifdef TVM_FFI_EXPORTS #define TVM_FFI_DLL __declspec(dllexport) #else #define TVM_FFI_DLL __declspec(dllimport) #endif #define TVM_FFI_DLL_EXPORT __declspec(dllexport) #endif #ifndef TVM_FFI_DLL #define TVM_FFI_DLL __attribute__((visibility("default"))) #define TVM_FFI_DLL_EXPORT __attribute__((visibility("default"))) #endif // NOLINTBEGIN(modernize-macro-to-enum) /*! \brief TVM FFI major version. */ #define TVM_FFI_VERSION_MAJOR 0 /*! \brief TVM FFI minor version. */ #define TVM_FFI_VERSION_MINOR 1 /*! \brief TVM FFI patch version. */ #define TVM_FFI_VERSION_PATCH 12 // NOLINTEND(modernize-macro-to-enum) #ifdef __cplusplus extern "C" { #endif /*! * \brief TVM FFI version. */ typedef struct { /*! \brief TVM FFI major version. */ uint32_t major; /*! \brief TVM FFI minor version. */ uint32_t minor; /*! \brief TVM FFI patch version. */ uint32_t patch; } TVMFFIVersion; // [TVMFFITypeIndex.begin] #ifdef __cplusplus enum TVMFFITypeIndex : int32_t { #else typedef enum { #endif /* * \brief The root type of all FFI objects. * * We include it so TypeIndex captures all possible runtime values. * `kTVMFFIAny` code will never appear in Any::type_index. * However, it may appear in field annotations during reflection. */ kTVMFFIAny = -1, // [Section] On-stack POD and special types: [0, kTVMFFIStaticObjectBegin) // N.B. `kTVMFFIRawStr` is a string backed by a `\0`-terminated char array, // which is not owned by TVMFFIAny. It is required that the following // invariant holds: // - `Any::type_index` is never `kTVMFFIRawStr` // - `AnyView::type_index` can be `kTVMFFIRawStr` // /*! \brief None/nullptr value */ kTVMFFINone = 0, /*! \brief POD int value */ kTVMFFIInt = 1, /*! \brief POD bool value */ kTVMFFIBool = 2, /*! \brief POD float value */ kTVMFFIFloat = 3, /*! \brief Opaque pointer object */ kTVMFFIOpaquePtr = 4, /*! \brief DLDataType */ kTVMFFIDataType = 5, /*! \brief DLDevice */ kTVMFFIDevice = 6, /*! \brief DLTensor* */ kTVMFFIDLTensorPtr = 7, /*! \brief const char* */ kTVMFFIRawStr = 8, /*! \brief TVMFFIByteArray* */ kTVMFFIByteArrayPtr = 9, /*! \brief R-value reference to ObjectRef */ kTVMFFIObjectRValueRef = 10, /*! \brief Small string on stack */ kTVMFFISmallStr = 11, /*! \brief Small bytes on stack */ kTVMFFISmallBytes = 12, /*! \brief Start of statically defined objects. */ kTVMFFIStaticObjectBegin = 64, /*! * \brief Object, all objects starts with TVMFFIObject as its header. * \note We will also add other fields */ kTVMFFIObject = 64, /*! * \brief String object, layout = { TVMFFIObject, TVMFFIByteArray, ... } */ kTVMFFIStr = 65, /*! * \brief Bytes object, layout = { TVMFFIObject, TVMFFIByteArray, ... } */ kTVMFFIBytes = 66, /*! \brief Error object. */ kTVMFFIError = 67, /*! \brief Function object. */ kTVMFFIFunction = 68, /*! * \brief Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... } */ kTVMFFIShape = 69, /*! * \brief Tensor object, layout = { TVMFFIObject, DLTensor, ... } */ kTVMFFITensor = 70, /*! \brief Array object. */ kTVMFFIArray = 71, /*! \brief Map object. */ kTVMFFIMap = 72, /*! \brief Runtime dynamic loaded module object. */ kTVMFFIModule = 73, /*! * \brief Opaque python object. * * This is a special type index to indicate we are storing an opaque PyObject. * Such object may interact with callback functions that are registered to support * python-related operations. * * We only translate the objects that we do not recognize into this type index. * * \sa TVMFFIObjectCreateOpaque */ kTVMFFIOpaquePyObject = 74, /*! \brief List object. */ kTVMFFIList = 75, /*! \brief Dict object. */ kTVMFFIDict = 76, //---------------------------------------------------------------- // more complex objects //---------------------------------------------------------------- kTVMFFIStaticObjectEnd, // [Section] Dynamic Boxed: [kTVMFFIDynObjectBegin, +oo) /*! \brief Start of type indices that are allocated at runtime. */ kTVMFFIDynObjectBegin = 128 #ifdef __cplusplus }; #else } TVMFFITypeIndex; #endif // [TVMFFITypeIndex.end] /*! \brief Handle to Object from C API's pov */ typedef void* TVMFFIObjectHandle; /*! * \brief bitmask of the object deleter flag. */ #ifdef __cplusplus enum TVMFFIObjectDeleterFlagBitMask : int32_t { #else typedef enum { #endif /*! * \brief deleter action when strong reference count becomes zero. * Need to call destructor of the object but not free the memory block. */ kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0, /*! * \brief deleter action when weak reference count becomes zero. * Need to free the memory block. */ kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1, /*! * \brief deleter action when both strong and weak reference counts become zero. * \note This is the most common case. */ kTVMFFIObjectDeleterFlagBitMaskBoth = (kTVMFFIObjectDeleterFlagBitMaskStrong | kTVMFFIObjectDeleterFlagBitMaskWeak), #ifdef __cplusplus }; #else } TVMFFIObjectDeleterFlagBitMask; #endif // [TVMFFIObject.begin] /*! * \brief C-based type of all FFI object header that allocates on heap. */ typedef struct { /*! * \brief Combined strong and weak reference counter of the object. * * Strong ref counter is packed into the lower 32 bits. * Weak ref counter is packed into the upper 32 bits. * * It is equivalent to { uint32_t strong_ref_count, uint32_t weak_ref_count } * in little-endian structure: * * - strong_ref_count: `combined_ref_count & 0xFFFFFFFF` * - weak_ref_count: `(combined_ref_count >> 32) & 0xFFFFFFFF` * * Rationale: atomic ops on strong ref counter remains the same as +1/-1, * this combined ref counter allows us to use u64 atomic once * instead of a separate atomic read of weak counter during deletion. * * The ref counter goes first to align ABI with most intrusive ptr designs. * It is also likely more efficient as rc operations can be quite common. */ uint64_t combined_ref_count; /*! * \brief type index of the object. * \note The type index of Object and Any are shared in FFI. */ int32_t type_index; /*! \brief Extra padding to ensure 8 bytes alignment. */ uint32_t __padding; #if !defined(TVM_FFI_DOXYGEN_MODE) union { #endif /*! * \brief Deleter to be invoked when strong reference counter goes to zero. * \param self The self object handle. * \param flags The flags to indicate deletion behavior. * \sa TVMFFIObjectDeleterFlagBitMask */ void (*deleter)(void* self, int flags); /*! * \brief auxilary field to TVMFFIObject is always 8 bytes aligned. * \note This helps us to ensure cross platform compatibility. */ int64_t __ensure_align; #if !defined(TVM_FFI_DOXYGEN_MODE) }; #endif } TVMFFIObject; // [TVMFFIObject.end] // [TVMFFIAny.begin] /*! * \brief C-based type of all on stack Any value. * * Any value can hold on stack values like int, * as well as reference counted pointers to object. */ typedef struct { /*! * \brief type index of the object. * \note The type index of Object and Any are shared in FFI. */ int32_t type_index; #if !defined(TVM_FFI_DOXYGEN_MODE) union { // 4 bytes #endif /*! \brief padding, must set to zero for values other than small string. */ uint32_t zero_padding; /*! * \brief Length of small string, with a max value of 7. * * We keep small str to start at next 4 bytes to ensure alignment * when accessing the small str content. */ uint32_t small_str_len; #if !defined(TVM_FFI_DOXYGEN_MODE) }; #endif #if !defined(TVM_FFI_DOXYGEN_MODE) union { // 8 bytes #endif /*! \brief integers */ int64_t v_int64; /*! \brief floating-point numbers */ double v_float64; /*! \brief typeless pointers */ void* v_ptr; /*! \brief raw C-string */ const char* v_c_str; /*! \brief ref counted objects */ TVMFFIObject* v_obj; /*! \brief data type */ DLDataType v_dtype; /*! \brief device */ DLDevice v_device; /*! \brief small string */ char v_bytes[8]; /*! \brief uint64 repr mainly used for hashing */ uint64_t v_uint64; #if !defined(TVM_FFI_DOXYGEN_MODE) }; #endif } TVMFFIAny; // [TVMFFIAny.end] // [TVMFFIByteArray.begin] /*! * \brief Byte array data structure used by String and Bytes. * * String and Bytes object layout = { TVMFFIObject, TVMFFIByteArray, ... } * * \note This byte array data structure layout differs in 32/64 bit platforms. * as size_t equals to the size of the pointer, use this convetion to * be consistent with std::string and also avoid need to calculate padding * for the size field on 32-bit platforms. * The FFI binding should be careful when treating this ABI. */ typedef struct { /*! \brief The data pointer. */ const char* data; /*! \brief The size of the data. */ size_t size; } TVMFFIByteArray; // [TVMFFIByteArray.end] /*! * \brief Shape cell used in shape object following header. */ typedef struct { /*! \brief The data pointer. */ const int64_t* data; /*! \brief The size of the data. */ size_t size; } TVMFFIShapeCell; // [TVMFFISeqCell.begin] /*! * \brief Sequence cell used by sequence-like containers. * * ArrayObj and ListObj both inherit from this cell. */ #ifdef __cplusplus struct TVMFFISeqCell { #else typedef struct { #endif /*! \brief Data pointer to the first element of the sequence. */ void* data; /*! \brief Number of elements used. */ int64_t size; /*! \brief Number of elements allocated. */ int64_t capacity; /*! * \brief Optional deleter for the data buffer. * * When non-null, data was allocated separately from the object * (e.g. ListObj heap buffer) and data_deleter is called to free it. * * When nullptr, data lives inside the object allocation itself * (e.g. ArrayObj inplace storage via make_inplace_array_object) * and is freed together with the object. */ void (*data_deleter)(void*); #ifdef __cplusplus }; #else } TVMFFISeqCell; #endif // [TVMFFISeqCell.end] /*! * \brief Mode to update the backtrace of the error. */ #ifdef __cplusplus enum TVMFFIBacktraceUpdateMode : int32_t { #else typedef enum { #endif kTVMFFIBacktraceUpdateModeReplace = 0, kTVMFFIBacktraceUpdateModeAppend = 1, #ifdef __cplusplus }; #else } TVMFFIBacktraceUpdateMode; #endif // [TVMFFIErrorCell.begin] /*! * \brief Error cell used in error object following header. */ typedef struct { /*! \brief The kind of the error. */ TVMFFIByteArray kind; /*! \brief The message of the error. */ TVMFFIByteArray message; /*! * \brief The backtrace of the error. * * The backtrace is in the order of recent call first from the top of the stack * to the bottom of the stack. This order makes it helpful for appending * the extra backtrace to the end as we go up when error is propagated. * * When printing out, we encourage reverse the order of lines to make it * align with python style. */ TVMFFIByteArray backtrace; /*! * \brief Function handle to update the backtrace of the error. * \param self The self object handle. * \param backtrace The backtrace to update. * \param update_mode The mode to update the backtrace, * can be either kTVMFFIBacktraceUpdateModeReplace, kTVMFFIBacktraceUpdateModeAppend. */ void (*update_backtrace)(TVMFFIObjectHandle self, const TVMFFIByteArray* backtrace, int32_t update_mode); /*! * \brief Optional cause error chain that caused this error to be raised. * \note This handle is owned by the ErrorCell. */ TVMFFIObjectHandle cause_chain; /*! * \brief Optional extra context that can be used to record additional info about the error. * \note This handle is owned by the ErrorCell. */ TVMFFIObjectHandle extra_context; } TVMFFIErrorCell; // [TVMFFIErrorCell.end] // [TVMFFISafeCallType.begin] /*! * \brief Type that defines C-style safe call convention * * Safe call explicitly catches exception on function boundary. * * \param handle The function handle * \param num_args Number of input arguments * \param args The input arguments to the call. * \param result Store output result. * * IMPORTANT: caller must initialize result->type_index to be kTVMFFINone, * or any other value smaller than kTVMFFIStaticObjectBegin. * * \return The call returns 0 if call is successful. * It returns non-zero value if there is an error. * * Possible return error of the API functions: * * 0: success * * -1: error happens, can be retrieved by TVMFFIErrorMoveFromRaised * * \note We decided to leverage TVMFFIErrorMoveFromRaised and TVMFFIErrorSetRaised * for C function error propagation. This design choice, while * introducing a dependency for TLS runtime, simplifies error * propgation in chains of calls in compiler codegen. * As we do not need to propagate error through argument but simply * set them in the runtime environment. * * \sa TVMFFIErrorMoveFromRaised * \sa TVMFFIErrorSetRaised * \sa TVMFFIErrorSetRaisedFromCStr * \sa TVMFFIErrorSetRaisedFromCStrParts */ typedef int (*TVMFFISafeCallType)(void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result); // [TVMFFISafeCallType.end] // [TVMFFIFunctionCell.begin] /*! * \brief Object cell for function object following header. */ typedef struct { /*! \brief A C API compatible call with exception catching. */ TVMFFISafeCallType safe_call; /*! * \brief A function pointer to an underlying cpp call. * * The signature is the same as TVMFFISafeCallType except the return type is void, * and the function throws exception directly instead of returning error code. * We use void* here to avoid depending on c++ compiler. * * This pointer should be set to NULL for functions that are not originally created in cpp. * * \note The caller must assume the same cpp exception catching abi when using this pointer. * When used across FFI boundaries, always use safe_call. */ void* cpp_call; } TVMFFIFunctionCell; // [TVMFFIFunctionCell.end] /*! * \brief Object cell for opaque object following header. */ typedef struct { /*! \brief The handle of the opaque object, for python it is PyObject* */ void* handle; } TVMFFIOpaqueObjectCell; //----------------------------------------------------------------------- // Section: Version API //----------------------------------------------------------------------- /*! * \brief Get the TVM FFI version from the current C ABI. * * This function is always stable across all versions of the C ABI. * * \param out_version The output version. */ TVM_FFI_DLL void TVMFFIGetVersion(TVMFFIVersion* out_version); //------------------------------------------------------------ // Section: Basic object API //------------------------------------------------------------ /*! * \brief Increase the strong reference count of an object handle * \param obj The object handle. * \note Internally we increase the reference counter of the object. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIObjectIncRef(TVMFFIObjectHandle obj); /*! * \brief Free an object handle by decreasing strong reference * \param obj The object handle. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIObjectDecRef(TVMFFIObjectHandle obj); /*! * \brief Create an Opaque object by passing in handle, type_index and deleter. * * The opaque object's lifetime is managed as an Object, so it can be retained * and released like other objects. * When the opaque object is kTVMFFIOpaquePyObject, it can be converted back to * the python type when returned or passed as arguments to a python function. * * We can support ffi::Function that interacts with these objects, * most likely callback registered from python. * * For language bindings, we only convert types that we do not recognize into this type. * On the C++ side, the most common way to represent such OpaqueObject is to simply * use ffi::ObjectRef or ffi::Any. * * \param handle The resource handle of the opaque object. * \param type_index The type index of the object. * \param deleter deleter to recycle * \param out The output of the opaque object. * \return 0 when success, nonzero when failure happens * * \note The caller must ensure the type_index is a valid opaque object type index. * \sa kTVMFFIOpaquePyObject */ TVM_FFI_DLL int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index, void (*deleter)(void* handle), TVMFFIObjectHandle* out); /*! * \brief Convert type key to type index. * \param type_key The key of the type. * \param out_tindex the corresponding type index. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFITypeKeyToIndex(const TVMFFIByteArray* type_key, int32_t* out_tindex); //----------------------------------------------------------------------- // Section: Basic function calling API for function implementation //----------------------------------------------------------------------- /*! * \brief Create a FFIFunc by passing in callbacks from a C callback. * The registered function can then be retrieved by the backend using its name. * \param self The resource handle of the C callback. * \param safe_call The C callback implementation. * \param deleter The deleter to recycle. * \param out The output of the function. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIFunctionCreate(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void* self), TVMFFIObjectHandle* out); /*! * \brief Get a global function registered in the system. * \param name The name of the function. * \param out The result function pointer, NULL if it does not exist. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIFunctionGetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle* out); /*! * \brief Convert an AnyView to an owned Any. * \param any_view The AnyView to convert. * \param out The output Any, must be an empty object. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIAnyViewToOwnedAny(const TVMFFIAny* any_view, TVMFFIAny* out); /*! * \brief Call a FFIFunc by passing in arguments. * \param func The resource handle of the C callback. * \param args The input arguments to the call. * \param num_args The number of input arguments. * \param result The output result, caller must ensure result->type_index is set to kTVMFFINone. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIFunctionCall(TVMFFIObjectHandle func, TVMFFIAny* args, int32_t num_args, TVMFFIAny* result); /*! * \brief Move the last error from the environment to the result. * \param result The result error. * \note This function clears the error stored in the TLS. */ TVM_FFI_DLL void TVMFFIErrorMoveFromRaised(TVMFFIObjectHandle* result); /*! * \brief Set a raised error in TLS, which can be fetched by TVMFFIErrorMoveFromRaised. * \param error The error object handle */ TVM_FFI_DLL void TVMFFIErrorSetRaised(TVMFFIObjectHandle error); /*! * \brief Set a raised error in TLS, which can be fetched by TVMFFIErrorMoveFromRaised. * \param kind The kind of the error. * \param message The error message. * \note This is a convenient method for the C API side to set an error directly from a string. */ TVM_FFI_DLL void TVMFFIErrorSetRaisedFromCStr(const char* kind, const char* message); /*! * \brief Set a raised error in TLS, which can be fetched by TVMFFIErrorMoveFromRaised. * * Rationale: This function can be used by compilers to create error messages by * concatenating multiple parts of the error message, which can reduce the * storage size for common parts such as function signatures. * * For example, the following are possible error messages from a kernel DSL * * - Argument 1 mismatch in `matmul(x: Tensor, y: Tensor, z: Tensor)`, dtype mismatch * - Argument 2 mismatch in `matmul(x: Tensor, y: Tensor, z: Tensor)`, shape[0] mismatch * - Argument 2 mismatch in `matmul(x: Tensor, y: Tensor, z: Tensor)`, shape[1] mismatch * * Storing each part of the error message as a separate global string can cause quite * a bit of duplication, especially considering the kinds of error reports we may have. * Instead, compilers can store error messages in parts, where items like * `matmul(x: Tensor, y: Tensor, z: Tensor)` can be reused across multiple error messages. * This API simplifies error reporting for such cases. * * \param kind The kind of the error. * \param message_parts The error message parts, each part can be NULL and will be skipped. * \param num_parts The number of error message parts. */ TVM_FFI_DLL void TVMFFIErrorSetRaisedFromCStrParts(const char* kind, const char** message_parts, int32_t num_parts); /*! * \brief Create an initial error object. * \param kind The kind of the error. * \param message The error message. * \param backtrace The backtrace of the error. * \param out The output error object handle. * \return 0 on success, nonzero on failure(likely MemoryError) * * \note This function is different from other functions as it is used in the error handling loop. * So we do not follow normal error handling patterns. When error happens it will not set * the error in TLS (since TLS error setting also involves creating an Error object). * Instead, caller should simply report MemoryError to the logger. */ TVM_FFI_DLL int TVMFFIErrorCreate(const TVMFFIByteArray* kind, const TVMFFIByteArray* message, const TVMFFIByteArray* backtrace, TVMFFIObjectHandle* out); /*! * \brief Create an initial error object with cause chain and extra context. * \param kind The kind of the error. * \param message The error message. * \param backtrace The backtrace of the error. * \param cause_chain The cause error chain that caused this error to be raised. * \param extra_context The extra context that can be used to record additional information. * \param out The output error object handle. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIErrorCreateWithCauseAndExtraContext( const TVMFFIByteArray* kind, const TVMFFIByteArray* message, const TVMFFIByteArray* backtrace, TVMFFIObjectHandle cause_chain, TVMFFIObjectHandle extra_context, TVMFFIObjectHandle* out); //------------------------------------------------------------ // Section: DLPack support APIs //------------------------------------------------------------ /*! * \brief Produce a managed Tensor from a DLPack tensor. * \param from The source DLPack tensor. * \param require_alignment The minimum alignment required of the data + byte_offset. * \param require_contiguous Boolean flag indicating if we need to check for contiguity. * \param out The output Tensor handle. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITensorFromDLPack(DLManagedTensor* from, int32_t require_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out); /*! * \brief Produce a DLManagedTensor from the array that shares data memory with the array. * \param from The source array. * \param out The DLManagedTensor handle. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITensorToDLPack(TVMFFIObjectHandle from, DLManagedTensor** out); /*! * \brief Produce a managed Tensor from a DLPack tensor. * \param from The source DLPack tensor. * \param require_alignment The minimum alignment required of the data + byte_offset. * \param require_contiguous Boolean flag indicating if we need to check for contiguity. * \param out The output Tensor handle. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITensorFromDLPackVersioned(DLManagedTensorVersioned* from, int32_t require_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out); /*! * \brief Produce a DLManagedTensor from the array that shares data memory with the array. * \param from The source array. * \param out The DLManagedTensor handle. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITensorToDLPackVersioned(TVMFFIObjectHandle from, DLManagedTensorVersioned** out); /*! * \brief Create a Tensor view from source using metadata in the prototype while retaining the * source tensor. * \param source The source tensor whose data memory will be shared by the view. * \param prototype The prototype DLTensor that contains the metadata for the view. * \param out The output Tensor handle. * \return 0 on success, nonzero on failure. * \note This function is unsafe and the caller must ensure the prototype is valid and that * the prototype's data pointer points to memory owned by the source tensor. The callee * allocates shape and strides arrays in the output tensor and copies them from the prototype. */ TVM_FFI_DLL int TVMFFITensorCreateUnsafeView(TVMFFIObjectHandle source, const DLTensor* prototype, TVMFFIObjectHandle* out); //--------------------------------------------------------------- // Section: string/bytes support APIs. // These APIs are used to simplify the string/bytes construction //--------------------------------------------------------------- /*! * \brief Reinterpret the content of TVMFFIByteArray to String. * \param input The TVMFFIByteArray to convert. * \param out The output String owned by the caller, maybe a SmallStr or a Str object. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIStringFromByteArray(const TVMFFIByteArray* input, TVMFFIAny* out); /*! * \brief Reinterpret the content of TVMFFIByteArray to Bytes. * \param input The TVMFFIByteArray to convert. * \param out The output Bytes owned by the caller, maybe a SmallBytes or a Bytes object. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIBytesFromByteArray(const TVMFFIByteArray* input, TVMFFIAny* out); //--------------------------------------------------------------- // Section: dtype string support APIs. // These APIs are used to simplify the dtype printings during FFI //--------------------------------------------------------------- /*! * \brief Convert a string to a DLDataType. * \param str The string to convert. * \param out The output DLDataType. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType* out); /*! * \brief Convert a DLDataType to a string. * \param dtype The DLDataType to convert. * \param out The output string. * \return 0 on success, nonzero on failure. * \note out is a String object that needs to be freed by the caller via TVMFFIObjectDecRef. The content of string can be accessed via TVMFFIObjectGetByteArrayPtr. * \note The input dtype is a pointer to the DLDataType to avoid ABI compatibility issues. */ TVM_FFI_DLL int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIAny* out); //------------------------------------------------------------ // Section: Type reflection support APIs // // The reflec //------------------------------------------------------------ /*! * \brief Getter that can take the address of a field and set the result. * \param field The raw address of the field. * \param result Stores the result. * \return 0 on success, nonzero on failure. */ typedef int (*TVMFFIFieldGetter)(void* field, TVMFFIAny* result); /*! * \brief Getter that can take the address of a field and set it to a value. * \param field The raw address of the field. * \param value The value to set. * \return 0 on success, nonzero on failure. */ typedef int (*TVMFFIFieldSetter)(void* field, const TVMFFIAny* value); /*! * \brief Function that creates a new instance of the type. * \param result The new object handle * \return 0 on success, nonzero on failure. */ typedef int (*TVMFFIObjectCreator)(TVMFFIObjectHandle* result); /*! * \brief bitmask of the field. */ #ifdef __cplusplus enum TVMFFIFieldFlagBitMask : int32_t { #else typedef enum { #endif /*! \brief The field is writable. */ kTVMFFIFieldFlagBitMaskWritable = 1 << 0, /*! \brief The field has default value. */ kTVMFFIFieldFlagBitMaskHasDefault = 1 << 1, /*! \brief The field is a static method. */ kTVMFFIFieldFlagBitMaskIsStaticMethod = 1 << 2, /*! * \brief The field should be ignored when performing structural eq/hash * * This is an optional meta-data for structural eq/hash. */ kTVMFFIFieldFlagBitMaskSEqHashIgnore = 1 << 3, /*! * \brief The field enters a recursive def region. * * This is an optional meta-data for structural eq/hash. * * \sa TVMFFIDefRegionKind for the def-region semantics. */ kTVMFFIFieldFlagBitMaskSEqHashDefRecursive = 1 << 4, /*! * \brief The default_value_or_factory is a callable factory function () -> Any. * * When this flag is set along with kTVMFFIFieldFlagBitMaskHasDefault, * the default_value_or_factory field contains a Function that should be * called with no arguments to produce the default value, rather than * being used directly as the default value. */ kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5, /*! * \brief The field is excluded from repr output. * * When set, the field will not appear in the generic reflection-based repr. * By default this flag is off (meaning the field is included in repr). */ kTVMFFIFieldFlagBitMaskReprOff = 1 << 6, /*! * \brief The field is excluded from recursive comparison. * * When set, the field will not participate in generic reflection-based * recursive comparison (RecursiveEq, RecursiveLt, etc.). * By default this flag is off (meaning the field is included in comparison). */ kTVMFFIFieldFlagBitMaskCompareOff = 1 << 7, /*! * \brief The field is excluded from recursive hashing. * * When set, the field will not participate in generic reflection-based * recursive hashing (RecursiveHash). * By default this flag is off (meaning the field is included in hashing). */ kTVMFFIFieldFlagBitMaskHashOff = 1 << 8, /*! * \brief The field is excluded from auto-generated ``__ffi_init__``. * * When set, the field will not appear as a parameter of the reflection-based * auto-generated constructor. The field must either have a default value * or be initialized by the creator (default constructor / UnsafeInit). * By default this flag is off (meaning the field is included in init). */ kTVMFFIFieldFlagBitMaskInitOff = 1 << 9, /*! * \brief The field is keyword-only in the auto-generated ``__ffi_init__``. * * When set, the field can only be provided via the KWARGS calling convention * (not as a positional argument) in the auto-generated constructor. * This flag is only meaningful when kTVMFFIFieldFlagBitMaskInitOff is *not* set. * By default this flag is off (meaning the field accepts positional arguments). */ kTVMFFIFieldFlagBitMaskKwOnly = 1 << 10, /*! * \brief The setter field is a TVMFFIObjectHandle pointing to a FunctionObj. * * When this flag is set, the ``setter`` member of TVMFFIFieldInfo is not a * TVMFFIFieldSetter function pointer but instead a TVMFFIObjectHandle * pointing to a FunctionObj. The FunctionObj is called with two arguments: * ``(field_addr_as_OpaquePtr, value_as_AnyView)``. */ kTVMFFIFieldFlagBitSetterIsFunctionObj = 1 << 11, /*! * \brief The field enters a non-recursive def region. * * This is an optional meta-data for structural eq/hash. * * \sa TVMFFIDefRegionKind for the def-region semantics. * * \note Bit 1 << 12 is used here because bits 1 << 5 .. 1 << 11 are * already taken by other field flags above. */ kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive = 1 << 12, #ifdef __cplusplus }; #else } TVMFFIFieldFlagBitMask; #endif /*! * \brief Optional meta-data for structural eq/hash. * * This meta-data is only useful when we want to leverage the information * to perform richer semantics aware structural comparison and hash. * It can be safely ignored if such information is not needed. * * The meta-data record comparison method in tree node and DAG node. * * \code{.cpp} * x = VarNode() * v0 = AddNode(x, 1) * v1 = AddNode(x, 1) * v2 = AddNode(v0, v0) * v3 = AddNode(v1, v0) * \endcode * * Consider the construct sequence of AddNode below, * if AddNode is treated as a tree node, then v2 and v3 * structural equals to each other, but if AddNode is * treated as a DAG node, then v2 and v3 does not * structural equals to each other. */ #ifdef __cplusplus enum TVMFFISEqHashKind : int32_t { #else typedef enum { #endif /*! \brief Do not support structural eq/hash. */ kTVMFFISEqHashKindUnsupported = 0, /*! * \brief The object be compared as a tree node. */ kTVMFFISEqHashKindTreeNode = 1, /*! * \brief The object is treated as a free variable that can be mapped * to another free variable in the definition region. */ kTVMFFISEqHashKindFreeVar = 2, /*! * \brief The field should be compared as a DAG node. */ kTVMFFISEqHashKindDAGNode = 3, /*! * \brief The object is treated as a constant tree node. * * Same as tree node, but the object does not contain free var * as any of its nested children. * * That means we can use pointer equality for equality. */ kTVMFFISEqHashKindConstTreeNode = 4, /*! * \brief One can simply use pointer equality for equality. * * This is useful for "singleton"-style object that can * is only an unique copy of each value. */ kTVMFFISEqHashKindUniqueInstance = 5, #ifdef __cplusplus }; #else } TVMFFISEqHashKind; #endif /*! * \brief Kind of def region a structural-equal / structural-hash callback is * currently in when visiting a field. */ #ifdef __cplusplus enum TVMFFIDefRegionKind : int32_t { #else typedef enum { #endif /*! * \brief Not in a def region. * * Free vars reachable through this field are uses; they must already * be bound by an enclosing def region or equality / hashing falls * back to pointer identity. */ kTVMFFIDefRegionKindNone = 0, /*! * \brief In a recursive def region. * * When we see a free var for the first time, we define the var, and * the sub-fields of the var (e.g. its struct_info / type_annotation / * shape) are also still in the def region — any free vars discovered * inside those sub-fields are themselves treated as fresh defs at the * same site. * * One example is function parameter lists: the value var and any * shape parameters in its type are co-introduced at the same binding * site. */ kTVMFFIDefRegionKindRecursive = 1, /*! * \brief In a non-recursive def region. * * When we see a free var for the first time, we define the var, but * the sub-fields of the var are NOT in the def region — they are * treated as use references that must resolve against an outer * binding. Free vars found in those sub-fields therefore do not * rebind; if they are not already bound, equality fails. * * One example is a normal binding whose value type contains shape * parameters: the value var is introduced fresh, but its shape * parameters reference vars defined in an outer scope. */ kTVMFFIDefRegionKindNonRecursive = 2, #ifdef __cplusplus }; #else } TVMFFIDefRegionKind; #endif /*! * \brief Information support for optional object reflection. */ typedef struct { /*! \brief The name of the field. */ TVMFFIByteArray name; /*! \brief The docstring about the field. */ TVMFFIByteArray doc; /*! \brief The structured metadata of the field in JSON string. */ TVMFFIByteArray metadata; /*! * \brief bitmask flags of the field. */ int64_t flags; /*! \brief The size of the field. */ int64_t size; /*! \brief The alignment of the field. */ int64_t alignment; /*! \brief The offset of the field. */ int64_t offset; /*! \brief The getter to access the field. */ TVMFFIFieldGetter getter; /*! * \brief The setter to access the field. * * When kTVMFFIFieldFlagBitSetterIsFunctionObj is NOT set (default), * this is a TVMFFIFieldSetter function pointer (cast to void*). * When kTVMFFIFieldFlagBitSetterIsFunctionObj IS set, * this is a TVMFFIObjectHandle pointing to a FunctionObj. * * \note The setter is set even if the field is readonly for serialization. */ void* setter; /*! * \brief The default value or default factory of the field. * * When flags has kTVMFFIFieldFlagBitMaskHasDefault set: * - If kTVMFFIFieldFlagBitMaskDefaultFromFactory is NOT set, * this holds the static default value as AnyView. * - If kTVMFFIFieldFlagBitMaskDefaultFromFactory IS set, * this holds a Function (() -> Any) that produces the default. */ TVMFFIAny default_value_or_factory; /*! * \brief The compile-time static type index of the field. * * This reflects the type declared at compile time, which is only trustworthy * for statically-inferrable types. It does NOT necessarily match the runtime * type. For example, a field declared as `Any` will have * `field_static_type_index == kTVMFFIAny` even if it holds an `int` at runtime, * and `Array` will report `kTVMFFIArray` even though the elements may be * `Array` at runtime. * * \warning Do NOT use this field to determine the actual type of a value at * runtime. It is purely a compile-time hint derived from the C++ field * declaration. The actual runtime type must be obtained from the value's * own `type_index`. * * \warning When the static type is a generic container (e.g. `Array`), * this field only tells you it is an Array — it says nothing about the * element types actually stored inside. Similarly, `kTVMFFIObject` only * means "some ObjectRef" without any subtype information. * * \note This is used by the serializer to inline POD field values directly * (avoiding node-graph overhead for None, Bool, Int, Float, DataType), * while routing other types through the node graph. */ int32_t field_static_type_index; } TVMFFIFieldInfo; /*! * \brief Method information that can appear in reflection table. */ typedef struct { /*! \brief The name of the field. */ TVMFFIByteArray name; /*! \brief The docstring about the method. */ TVMFFIByteArray doc; // Rationale: We separate the docstring from the metadata since docstrings // can be unstructured and sometimes large, while metadata can be focused // on storing structured information. /*! \brief Optional structured metadata of the method in JSON string. */ TVMFFIByteArray metadata; /*! \brief bitmask flags of the method. */ int64_t flags; /*! * \brief The method wrapped as ffi::Function, stored as AnyView. * \note The first argument to the method is always the self for instance methods. */ TVMFFIAny method; } TVMFFIMethodInfo; /*! * \brief Extra information of object type that can be used for reflection. * * \note This information is optional and can be used to enable reflection based * creation of the object. */ typedef struct { /*! \brief The docstring about the object. */ TVMFFIByteArray doc; /*! * \brief An optional function that can create a new empty instance of the type. * * When known_fixed_size is non-zero, creator can be called * with nullptr passed to optional_bytes. * * \note Caller must call setter for each field to initialize the object for * the final object to be in valid state. * * \note This field is optional to enable reflection based creation. */ TVMFFIObjectCreator creator; /*! * \brief Total size of the object struct, if it is fixed and known. * * This field is set optional and set to 0 if not registered. */ int32_t total_size; /*! * \brief Optional meta-data for structural eq/hash. */ TVMFFISEqHashKind structural_eq_hash_kind; } TVMFFITypeMetadata; /*! * \brief One column of a type–attribute table: extra attributes keyed by runtime type index. * * TypeAttr is designed to support an open set of possible attributes that can be * registered and queried across different downstream use cases. * * Conceptually, TypeAttr is a dynamic variant of TypeTraits as seen in C++/Rust. * It behaves like c++ type_traits, so column[T] does not contain attributes from base classes. * * Typical use cases for TypeAttr include: * - Storing extra magic trait functions that can customize behaviors like hash/eq/repr. * - Storing extra metadata about type data structures (e.g., mutability) and properties of * op/operator structures that can be used by compiler transformations (e.g., passes). * * This column covers type indices in the range [begin_index, begin_index + size). * For a given type_index, the corresponding entry is data[type_index - begin_index] * when begin_index <= type_index < begin_index + size; otherwise the entry is * not present (treat as None/null). * * \sa TVMFFIRegisterTypeAttr */ typedef struct { /*! \brief The data of the column, indexed by (type_index - begin_index). */ const TVMFFIAny* data; /*! * \brief The number of elements in the data array. * * The column covers type indices in the range [begin_index, begin_index + size). * For a given type_index, the corresponding entry is data[type_index - begin_index] * when begin_index <= type_index < begin_index + size. */ int32_t size; /*! * \brief The starting type index of the column data. * A value of 0 means the data array starts at type_index 0. */ int32_t begin_index; } TVMFFITypeAttrColumn; /*! * \brief Runtime type information for object type checking. */ #ifdef __cplusplus struct TVMFFITypeInfo { #else typedef struct TVMFFITypeInfo { #endif /*! *\brief The runtime type index, * It can be allocated during runtime if the type is dynamic. */ int32_t type_index; /*! \brief number of parent types in the type hierachy. */ int32_t type_depth; /*! \brief the unique type key to identify the type. */ TVMFFIByteArray type_key; /*! * \brief type_ancestors[depth] stores the type_index of the acenstors at depth level * \note To keep things simple, we do not allow multiple inheritance so the * hieracy stays as a tree */ const struct TVMFFITypeInfo** type_ancestors; // The following fields are used for reflection /*! \brief Cached hash value of the type key, used for consistent structural hashing. */ uint64_t type_key_hash; /*! \brief number of reflection accessible fields. */ int32_t num_fields; /*! \brief number of reflection acccesible methods. */ int32_t num_methods; /*! \brief The reflection field information. */ const TVMFFIFieldInfo* fields; /*! \brief The reflection method. */ const TVMFFIMethodInfo* methods; /*! \brief The extra information of the type. */ const TVMFFITypeMetadata* metadata; #ifdef __cplusplus }; #else } TVMFFITypeInfo; #endif /*! * \brief Register the function to runtime's global table. * The registered function can then be retrieved by the backend using its name. * \param name The name of the function. * \param f The function to be registered. * \param allow_override Whether to allow overriding an already registered function. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIFunctionSetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle f, int allow_override); /*! * \brief Register the function to runtime's global table with method info. * This is the same as TVMFFIFunctionSetGlobal but with method info that can provide extra * metadata used in the runtime. * \param method_info The method info to be registered. * \param allow_override Whether to allow overriding an already registered function. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIFunctionSetGlobalFromMethodInfo(const TVMFFIMethodInfo* method_info, int allow_override); /*! * \brief Register type field information for runtime reflection. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITypeRegisterField(int32_t type_index, const TVMFFIFieldInfo* info); /*! * \brief Register type method information for runtime reflection. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITypeRegisterMethod(int32_t type_index, const TVMFFIMethodInfo* info); /*! * \brief Register type creator information for runtime reflection. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITypeRegisterMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata); /*! * \brief Register extra type attributes that can be looked up during runtime. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFITypeRegisterAttr(int32_t type_index, const TVMFFIByteArray* attr_name, const TVMFFIAny* attr_value); /*! * \brief Get the type attribute column by name. * \return The pointer to the type attribute column. * \return NULL if the attribute was not registered in the system. */ TVM_FFI_DLL const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray* attr_name); //------------------------------------------------------------ // Section: Backend noexcept functions for internal use // // These functions are used internally and do not throw error // instead the error will be logged and abort the process // These are function are being called in startup or exit time // so exception handling do not apply //------------------------------------------------------------ /*! * \brief Get stack backtrace in a string. * * The backtrace is in the order of recent call first from the top of the stack * to the bottom of the stack. This order makes it helpful for appending * the extra backtrace as we unwind the stack. * * When printing out, we encourage reverse the order of lines to make it * align with python style. * * \param filename The current file name. * \param lineno The current line number * \param func The current function * \param cross_ffi_boundary Whether the backtrace is crossing the ffi boundary * or we should stop at the ffi boundary when detected * \return The backtrace string * * \note filename/func can be nullptr, then this info is skipped, they are useful * for cases when debug symbols are not available. */ TVM_FFI_DLL const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const char* func, int cross_ffi_boundary); /*! * \brief Initialize the type info during runtime. * * When the function is first called for a type, * it will register the type to the type table in the runtime. * If the static_tindex is non-negative, the function will * allocate a runtime type index. * Otherwise, we will populate the type table and return the static index. * * \param type_key The type key. * \param type_depth The type depth. * \param static_type_index Static type index if any, can be -1, which means this is a dynamic index * \param num_child_slots Number of slots reserved for its children. * \param child_slots_can_overflow Whether to allow child to overflow the slots. * \param parent_type_index Parent type index, pass in -1 if it is root. * * \return The allocated type index. */ TVM_FFI_DLL int32_t TVMFFITypeGetOrAllocIndex(const TVMFFIByteArray* type_key, int32_t static_type_index, int32_t type_depth, int32_t num_child_slots, int32_t child_slots_can_overflow, int32_t parent_type_index); /*! * \brief Get dynamic type info by type index. * \return The type info. */ TVM_FFI_DLL const TVMFFITypeInfo* TVMFFIGetTypeInfo(int32_t type_index); // ---------------------------------------------------------------------------- // Static handle initialization and deinitialization API // ---------------------------------------------------------------------------- /*! * \brief Initialize a handle once in a thread-safe manner. * * This function checks if *handle_addr is nullptr, * and if so, calls the initialization function * and stores the result in *handle_addr. * * This function is thread-safe and is meant to be used by DSLs that, * unlike C++, may not have static initialization support. * * \param handle_addr The address of the handle to be initialized. * \param init_func The initialization function to be called once to create the result handle. * \return 0 on success, nonzero on failure. * * \note If init_func encounters an error, it should call TVMFFIErrorSetRaisedFromCStr * to set the error and return nonzero, which will then be propagated to the * caller of TVMFFIHandleInitOnce. */ TVM_FFI_DLL int TVMFFIHandleInitOnce(void** handle_addr, int (*init_func)(void** result)); /*! * \brief Deinitialize a handle once in a thread-safe manner. * * This function checks if *handle_addr is not nullptr, and if so, * calls the deinitialization function * and sets *handle_addr to nullptr. * * This function is thread-safe and is meant to be used by DSLs that, * unlike C++, may not have static deinitialization support. * * \param handle_addr The address of the handle to be deinitialized. * \param deinit_func The deinitialization function to be called if *handle_addr is not nullptr. * \return 0 on success, nonzero on failure. */ TVM_FFI_DLL int TVMFFIHandleDeinitOnce(void** handle_addr, int (*deinit_func)(void* handle)); #ifdef __cplusplus } // TVM_FFI_EXTERN_C #endif //--------------------------------------------------------------- // The following API defines static object attribute accessors // for language bindings. // // They are defined in C++ inline functions for cleaner code. // Note that they only have to do with address offset computation. // So they can always be reimplemented in bindings when c++ is // not available or when binding only wants to refer to the dll. //---------------------------------------------------------------- #ifdef __cplusplus /*! * \brief Get the type index of an object. * \param obj The object handle. * \return The type index. */ inline int32_t TVMFFIObjectGetTypeIndex(TVMFFIObjectHandle obj) { return static_cast(obj)->type_index; } /*! * \brief Get the content of a small string in bytearray format. * \param value The value to get the content of the small string in bytearray format. * \return The content of the small string in bytearray format. */ inline TVMFFIByteArray TVMFFISmallBytesGetContentByteArray(const TVMFFIAny* value) { return TVMFFIByteArray{value->v_bytes, static_cast(value->small_str_len)}; } /*! * \brief Get the data pointer of a bytearray from a string or bytes object. * \param obj The object handle. * \return The data pointer. */ inline TVMFFIByteArray* TVMFFIBytesGetByteArrayPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Get the data pointer of a ErrorInfo from an Error object. * \param obj The object handle. * \return The cell pointer. */ inline TVMFFIErrorCell* TVMFFIErrorGetCellPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Get the data pointer of a function cell from a function object. * \param obj The object handle. * \return The cell pointer. */ inline TVMFFIFunctionCell* TVMFFIFunctionGetCellPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Get the data pointer of a opaque object cell from a opaque object. * \param obj The object handle. * \return The cell pointer. */ inline TVMFFIOpaqueObjectCell* TVMFFIOpaqueObjectGetCellPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Get the data pointer of a shape array from a shape object. * \param obj The object handle. * \return The cell pointer. */ inline TVMFFIShapeCell* TVMFFIShapeGetCellPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Get the DLTensor pointer from an Tensor object. * \param obj The object handle. * \return The DLTensor pointer. */ inline DLTensor* TVMFFITensorGetDLTensorPtr(TVMFFIObjectHandle obj) { return reinterpret_cast(reinterpret_cast(obj) + sizeof(TVMFFIObject)); } /*! * \brief Create a DLDevice from a device type and device id. * \param device_type The device type. * \param device_id The device id. * \return The DLDevice. */ inline DLDevice TVMFFIDLDeviceFromIntPair(int32_t device_type, int32_t device_id) { return DLDevice{static_cast(device_type), device_id}; } #endif // __cplusplus #endif // TVM_FFI_C_API_H_ // NOLINTEND(modernize-use-using,bugprone-reserved-identifier,modernize-deprecated-headers) tvm-ffi-0.1.12/include/tvm/ffi/cast.h000066400000000000000000000052401521067262500172520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/cast.h * \brief Extra value casting helpers */ #ifndef TVM_FFI_CAST_H_ #define TVM_FFI_CAST_H_ #include #include #include namespace tvm { namespace ffi { /*! * \brief Get a reference type from a raw object ptr type * * It is always important to get a reference type * if we want to return a value as reference or keep * the object alive beyond the scope of the function. * * \param ptr The object pointer * \tparam RefType The reference type * \tparam ObjectType The object type * \return The corresponding RefType */ template inline RefType GetRef(const ObjectType* ptr) { using ContainerType = typename RefType::ContainerType; static_assert(std::is_base_of_v, "Can only cast to the ref of same container type"); if constexpr (is_optional_type_v || RefType::_type_is_nullable) { if (ptr == nullptr) { return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); } } else { TVM_FFI_ICHECK_NOTNULL(ptr); } return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned( const_cast(static_cast(ptr)))); } /*! * \brief Get an object ptr type from a raw object ptr. * * \param ptr The object pointer * \tparam BaseType The reference type * \tparam ObjectType The object type * \return The corresponding RefType */ template inline ObjectPtr GetObjectPtr(ObjectType* ptr) { static_assert(std::is_base_of_v, "Can only cast to the ref of same container type"); return details::ObjectUnsafe::ObjectPtrFromUnowned(ptr); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_CAST_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/000077500000000000000000000000001521067262500201305ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/container/array.h000066400000000000000000000722711521067262500214300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/array.h * \brief Array type. * * tvm::ffi::Array is an erased type that contains a list of content */ #ifndef TVM_FFI_CONTAINER_ARRAY_H_ #define TVM_FFI_CONTAINER_ARRAY_H_ #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief Array node content in array */ class ArrayObj : public SeqBaseObj { public: /*! * \brief Constructs a container and copy from another * \param cap The capacity of the container * \param from Source of the copy * \return Ref-counted ArrayObj requested */ static ObjectPtr CopyFrom(int64_t cap, ArrayObj* from) { int64_t size = from->TVMFFISeqCell::size; if (size > cap) { TVM_FFI_THROW(ValueError) << "Not enough capacity"; } ObjectPtr p = ArrayObj::Empty(cap); Any* write = p->MutableBegin(); Any* read = from->MutableBegin(); // To ensure exception safety, size is only incremented after the initialization succeeds for (int64_t& i = p->TVMFFISeqCell::size = 0; i < size; ++i) { new (write++) Any(*read++); } return p; } /*! * \brief Constructs a container and move from another * \param cap The capacity of the container * \param from Source of the move * \return Ref-counted ArrayObj requested */ static ObjectPtr MoveFrom(int64_t cap, ArrayObj* from) { int64_t size = from->TVMFFISeqCell::size; if (size > cap) { TVM_FFI_THROW(RuntimeError) << "Not enough capacity"; } ObjectPtr p = ArrayObj::Empty(cap); Any* write = p->MutableBegin(); Any* read = from->MutableBegin(); // To ensure exception safety, size is only incremented after the initialization succeeds for (int64_t& i = p->TVMFFISeqCell::size = 0; i < size; ++i) { new (write++) Any(std::move(*read++)); } from->TVMFFISeqCell::size = 0; return p; } /*! * \brief Constructs a container with n elements. Each element is a copy of val * \param n The size of the container * \param val The init value * \return Ref-counted ArrayObj requested */ static ObjectPtr CreateRepeated(int64_t n, const Any& val) { ObjectPtr p = ArrayObj::Empty(n); Any* itr = p->MutableBegin(); for (int64_t& i = p->TVMFFISeqCell::size = 0; i < n; ++i) { new (itr++) Any(val); } return p; } /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIArray; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIArray, ArrayObj, Object); /// \endcond private: /*! \return Size of initialized memory, used by InplaceArrayBase. */ size_t GetSize() const { return TVMFFISeqCell::size; } /*! * \brief Create an ArrayObj with the given capacity. * \param n Required capacity * \return Ref-counted ArrayObj requested */ static ObjectPtr Empty(int64_t n = kInitSize) { ObjectPtr p = make_inplace_array_object(n); p->TVMFFISeqCell::capacity = n; p->TVMFFISeqCell::size = 0; p->data = reinterpret_cast(p.get()) + sizeof(ArrayObj); p->data_deleter = nullptr; return p; } /*! * \brief Inplace-initialize the elements starting idx from [first, last) * \param idx The starting point * \param first Begin of iterator * \param last End of iterator * \tparam IterType The type of iterator * \return Self */ template ArrayObj* InitRange(int64_t idx, IterType first, IterType last) { Any* itr = MutableBegin() + idx; for (; first != last; ++first) { Any ref = *first; new (itr++) Any(std::move(ref)); } return this; } /*! \brief Initial size of ArrayObj */ static constexpr int64_t kInitSize = 4; /*! \brief Expansion factor of the Array */ static constexpr int64_t kIncFactor = 2; // Reference class template friend class Array; template friend class Tuple; template friend struct TypeTraits; // To specialize make_object friend ObjectPtr make_object<>(); }; /*! \brief Helper struct for type-checking * * is_valid_iterator::value will be true if IterType can * be dereferenced into a type that can be stored in an Array, and * false otherwise. */ template struct is_valid_iterator : std::bool_constant< std::is_same_v< T, std::remove_cv_t())>>> || std::is_base_of_v< T, std::remove_cv_t())>>>> { }; template struct is_valid_iterator, IterType> : is_valid_iterator {}; template struct is_valid_iterator : std::true_type {}; /*! * \brief Check whether IterType is valid iterator for T. * \tparam T The type. * \tparam IterType The type of iterator. */ template inline constexpr bool is_valid_iterator_v = is_valid_iterator::value; /*! * \brief Array, container representing a contiguous sequence of ObjectRefs. * * Array implements in-place copy-on-write semantics. * * As in typical copy-on-write, a method which would typically mutate the array * instead opaquely copies the underlying container, and then acts on its copy. * * If the array has reference count equal to one, we directly update the * container in place without copying. This is optimization is sound because * when the reference count is equal to one this reference is guranteed to be * the sole pointer to the container. * * * operator[] only provides const access, use Set to mutate the content. * \tparam T The content Value type, must be compatible with tvm::ffi::Any */ template >> class Array : public ObjectRef { public: /*! \brief The value type of the array */ using value_type = T; // constructors /*! * \brief Construct an Array with UnsafeInit */ explicit Array(UnsafeInit tag) : ObjectRef(tag) {} /*! * \brief default constructor */ Array() { data_ = ArrayObj::Empty(); } // NOLINT(modernize-use-equals-default) /*! * \brief Move constructor * \param other The other array */ Array(Array&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief Copy constructor * \param other The other array */ Array(const Array& other) : ObjectRef(other.data_) {} // NOLINT(google-explicit-constructor) /*! * \brief Constructor from another array * \param other The other array * \tparam U The value type of the other array */ template >> Array(Array&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief Constructor from another array * \param other The other array * \tparam U The value type of the other array */ template >> Array(const Array& other) // NOLINT(google-explicit-constructor) : ObjectRef(other.data_) {} /*! * \brief Move assignment from another array * \param other The other array */ TVM_FFI_INLINE Array& operator=(Array&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Assignment from another array * \param other The other array */ TVM_FFI_INLINE Array& operator=(const Array& other) { data_ = other.data_; return *this; } /*! * \brief Move assignment from another array * \param other The other array * \tparam U The value type of the other array */ template >> TVM_FFI_INLINE Array& operator=(Array&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Assignment from another array * \param other The other array * \tparam U The value type of the other array */ template >> TVM_FFI_INLINE Array& operator=(const Array& other) { data_ = other.data_; return *this; } /*! * \brief Constructor from pointer * \param n the container pointer */ explicit Array(ObjectPtr n) : ObjectRef(std::move(n)) {} /*! * \brief Constructor from iterator * \param first begin of iterator * \param last end of iterator * \tparam IterType The type of iterator */ template Array(IterType first, IterType last) { // NOLINT(performance-unnecessary-value-param) static_assert(is_valid_iterator_v, "IterType cannot be inserted into a tvm::Array"); Assign(first, last); } /*! * \brief constructor from initializer list * \param init The initializer list */ Array(std::initializer_list init) { // NOLINT(*) Assign(init.begin(), init.end()); } /*! * \brief constructor from vector * \param init The vector */ Array(const std::vector& init) { // NOLINT(*) Assign(init.begin(), init.end()); } /*! * \brief Constructs a container with n elements. Each element is a copy of val * \param n The size of the container * \param val The init value */ explicit Array(const size_t n, const T& val) { data_ = ArrayObj::CreateRepeated(n, val); } public: // iterators /// \cond Doxygen_Suppress struct ValueConverter { using ResultType = T; /*! * \brief Convert any to T * \param n The any value to convert * \return The converted value */ static T convert(const Any& n) { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(n); } }; /// \endcond /*! \brief The iterator type of the array */ using iterator = details::IterAdapter; /*! \brief The reverse iterator type of the array */ using reverse_iterator = details::ReverseIterAdapter; /*! \return begin iterator */ iterator begin() const { return iterator(GetArrayObj()->begin()); } /*! \return end iterator */ iterator end() const { return iterator(GetArrayObj()->end()); } /*! \return rbegin iterator */ reverse_iterator rbegin() const { // ArrayObj::end() is never nullptr return reverse_iterator(GetArrayObj()->end() - 1); } /*! \return rend iterator */ reverse_iterator rend() const { // ArrayObj::begin() is never nullptr return reverse_iterator(GetArrayObj()->begin() - 1); } public: // const methods in std::vector /*! * \brief Immutably read i-th element from array. * \param i The index * \return the i-th element. */ const T operator[](int64_t i) const { ArrayObj* p = GetArrayObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null array"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->at(i)); } /*! \return The size of the array */ size_t size() const { ArrayObj* p = GetArrayObj(); return p == nullptr ? 0 : p->size(); } /*! \return The capacity of the array */ size_t capacity() const { ArrayObj* p = GetArrayObj(); return p == nullptr ? 0 : p->SeqBaseObj::capacity(); } /*! \return Whether array is empty */ bool empty() const { return size() == 0; } /*! \return The first element of the array */ T front() const { ArrayObj* p = GetArrayObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null array"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->front()); } /*! \return The last element of the array */ T back() const { ArrayObj* p = GetArrayObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null array"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->back()); } public: // mutation in std::vector, implements copy-on-write /*! * \brief push a new item to the back of the list * \param item The item to be pushed. */ void push_back(const T& item) { ArrayObj* p = CopyOnWrite(1); p->EmplaceInit(p->TVMFFISeqCell::size++, item); } /*! * \brief Emplace a new element at the back of the array * \param args The arguments to construct the new element */ template void emplace_back(Args&&... args) { ArrayObj* p = CopyOnWrite(1); p->EmplaceInit(p->TVMFFISeqCell::size++, std::forward(args)...); } /*! * \brief Insert an element into the given position * \param position An iterator pointing to the insertion point * \param val The element to insert */ void insert(iterator position, const T& val) { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot insert a null array"; } int64_t idx = std::distance(begin(), position); CopyOnWrite(1)->insert(idx, Any(val)); } /*! * \brief Insert a range of elements into the given position * \param position An iterator pointing to the insertion point * \param first The begin iterator of the range * \param last The end iterator of the range */ template void insert(iterator position, IterType first, IterType last) { static_assert(is_valid_iterator_v, "IterType cannot be inserted into a tvm::Array"); if (first == last) return; if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot insert a null array"; } int64_t idx = std::distance(begin(), position); int64_t numel = std::distance(first, last); CopyOnWrite(numel)->insert(idx, first, last); } /*! \brief Remove the last item of the list */ void pop_back() { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot pop_back a null array"; } CopyOnWrite()->pop_back(); } /*! * \brief Erase an element on the given position * \param position An iterator pointing to the element to be erased */ void erase(iterator position) { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot erase a null array"; } int64_t idx = std::distance(begin(), position); CopyOnWrite()->erase(idx); } /*! * \brief Erase a given range of elements * \param first The begin iterator of the range * \param last The end iterator of the range */ void erase(iterator first, iterator last) { if (first == last) return; if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot erase a null array"; } int64_t st = std::distance(begin(), first); int64_t ed = std::distance(begin(), last); CopyOnWrite()->erase(st, ed); } /*! * \brief Resize the array. * \param n The new size. */ void resize(int64_t n) { if (n < 0) { TVM_FFI_THROW(ValueError) << "cannot resize an Array to negative size"; } if (data_ == nullptr) { SwitchContainer(n); return; } int64_t cur_size = GetArrayObj()->TVMFFISeqCell::size; if (cur_size < n) { CopyOnWrite(n - cur_size)->resize(n); } else if (cur_size > n) { CopyOnWrite()->resize(n); } } /*! * \brief Make sure the list has the capacity of at least n * \param n lower bound of the capacity */ void reserve(int64_t n) { if (data_ == nullptr || n > static_cast(GetArrayObj()->SeqBaseObj::capacity())) { SwitchContainer(n); } } /*! \brief Release reference to all the elements */ void clear() { if (data_ != nullptr) { ArrayObj* p = CopyOnWrite(); p->clear(); } } /// \cond Doxygen_Suppress template static size_t CalcCapacityImpl() { return 0; } template static size_t CalcCapacityImpl(Array value, Args... args) { return value.size() + CalcCapacityImpl(args...); } template static size_t CalcCapacityImpl(T value, Args... args) { return 1 + CalcCapacityImpl(args...); } template static void AgregateImpl(Array& dest) {} // NOLINT(*) template static void AgregateImpl(Array& dest, Array value, Args... args) { // NOLINT(*) dest.insert(dest.end(), value.begin(), value.end()); AgregateImpl(dest, args...); } template static void AgregateImpl(Array& dest, T value, Args... args) { // NOLINT(*) dest.push_back(value); AgregateImpl(dest, args...); } /// \endcond public: // Array's own methods /*! * \brief set i-th element of the array. * \param i The index * \param value The value to be setted. */ void Set(int64_t i, T value) { CopyOnWrite()->SetItem(i, std::move(value)); } /*! \return The underlying ArrayObj */ ArrayObj* GetArrayObj() const { return static_cast(data_.get()); } /*! * \brief Helper function to apply a map function onto the array. * * \param fmap The transformation function T -> U. * * \tparam F The type of the mutation function. * * \tparam U The type of the returned array, inferred from the * return type of F. If overridden by the user, must be something * that is convertible from the return type of F. * * \note This function performs copy on write optimization. If * `fmap` returns an object of type `T`, and all elements of the * array are mapped to themselves, then the returned array will be * the same as the original, and reference counts of the elements in * the array will not be incremented. * * \return The transformed array. */ template > Array Map(F fmap) const { return Array(MapHelper(data_, fmap)); } /*! * \brief Helper function to apply fmutate to mutate an array. * \param fmutate The transformation function T -> T. * \tparam F the type of the mutation function. * \note This function performs copy on write optimization. */ template >>> void MutateByApply(F fmutate) { data_ = MapHelper(std::move(data_), fmutate); } /*! * \brief reset the array to content from iterator. * \param first begin of iterator * \param last end of iterator * \tparam IterType The type of iterator */ template void Assign(IterType first, IterType last) { // NOLINT(performance-unnecessary-value-param) int64_t cap = std::distance(first, last); if (cap < 0) { TVM_FFI_THROW(ValueError) << "cannot construct an Array of negative size"; } ArrayObj* p = GetArrayObj(); if (p != nullptr && data_.unique() && p->TVMFFISeqCell::capacity >= cap) { // do not have to make new space p->clear(); } else { // create new space data_ = ArrayObj::Empty(cap); p = GetArrayObj(); } // To ensure exception safety, size is only incremented after the initialization succeeds Any* itr = p->MutableBegin(); for (int64_t& i = p->TVMFFISeqCell::size = 0; i < cap; ++i, ++first, ++itr) { new (itr) Any(*first); } } /*! * \brief Copy on write semantics * Do nothing if current handle is the unique copy of the array. * Otherwise make a new copy of the array to ensure the current handle * hold a unique copy. * * \return Handle to the internal node container(which ganrantees to be unique) */ ArrayObj* CopyOnWrite() { if (data_ == nullptr) { return SwitchContainer(ArrayObj::kInitSize); } if (!data_.unique()) { return SwitchContainer(capacity()); } return static_cast(data_.get()); } /*! \brief specify container node */ using ContainerType = ArrayObj; /*! * \brief Agregate arguments into a single Array * \param args sequence of T or Array elements * \return Agregated Array */ template static Array Agregate(Args... args) { Array result; result.reserve(CalcCapacityImpl(args...)); AgregateImpl(result, args...); return result; } private: /*! * \brief Implement copy-on-write semantics, and ensures capacity is enough for extra elements. * \param reserve_extra Number of extra slots needed * \return ArrayObj pointer to the unique copy */ ArrayObj* CopyOnWrite(int64_t reserve_extra) { ArrayObj* p = GetArrayObj(); if (p == nullptr) { // necessary to get around the constexpr address issue before c++17 const int64_t kInitSize = ArrayObj::kInitSize; return SwitchContainer(std::max(kInitSize, reserve_extra)); } if (p->TVMFFISeqCell::capacity >= p->TVMFFISeqCell::size + reserve_extra) { return CopyOnWrite(); } int64_t cap = p->TVMFFISeqCell::capacity * ArrayObj::kIncFactor; cap = std::max(cap, p->TVMFFISeqCell::size + reserve_extra); return SwitchContainer(cap); } /*! * \brief Move or copy the ArrayObj to new address with the given capacity * \param capacity The capacity requirement of the new address */ ArrayObj* SwitchContainer(int64_t capacity) { if (data_ == nullptr) { data_ = ArrayObj::Empty(capacity); } else if (data_.unique()) { data_ = ArrayObj::MoveFrom(capacity, GetArrayObj()); } else { data_ = ArrayObj::CopyFrom(capacity, GetArrayObj()); } return static_cast(data_.get()); } /*! \brief Helper method for mutate/map * * A helper function used internally by both `Array::Map` and * `Array::MutateInPlace`. Given an array of data, apply the * mapping function to each element, returning the collected array. * Applies both mutate-in-place and copy-on-write optimizations, if * possible. * * \param data A pointer to the ArrayObj containing input data. * Passed by value to allow for mutate-in-place optimizations. * * \param fmap The mapping function * * \tparam F The type of the mutation function. * * \tparam U The output type of the mutation function. Inferred * from the callable type given. Must inherit from ObjectRef. * * \return The mapped array. Depending on whether mutate-in-place * or copy-on-write optimizations were applicable, may be the same * underlying array as the `data` parameter. */ template > static ObjectPtr MapHelper(ObjectPtr data, F fmap) { if (data == nullptr) { return nullptr; } TVM_FFI_ICHECK(data->IsInstance()); constexpr bool is_same_output_type = std::is_same_v; if constexpr (is_same_output_type) { if (data.unique()) { // Mutate-in-place path. Only allowed if the output type U is // the same as type T, we have a mutable this*, and there are // no other shared copies of the array. auto arr = static_cast(data.get()); for (auto it = arr->MutableBegin(); it != arr->MutableEnd(); it++) { T value = details::AnyUnsafe::CopyFromAnyViewAfterCheck(*it); // reset the original value to nullptr, to ensure unique ownership it->reset(); T mapped = fmap(std::move(value)); *it = std::move(mapped); } return data; } } constexpr bool compatible_types = is_valid_iterator_v || is_valid_iterator_v; ObjectPtr output = nullptr; auto arr = static_cast(data.get()); auto it = arr->begin(); if constexpr (compatible_types) { // Copy-on-write path, if the output Array might be // represented by the same underlying array as the existing // Array. Typically, this is for functions that map `T` to // `T`, but can also apply to functions that map `T` to // `Optional`, or that map `T` to a subclass or superclass of // `T`. bool all_identical = true; for (; it != arr->end(); it++) { U mapped = fmap(details::AnyUnsafe::CopyFromAnyViewAfterCheck(*it)); if (!(*it).same_as(mapped)) { // At least one mapped element is different than the // original. Therefore, prepare the output array, // consisting of any previous elements that had mapped to // themselves (if any), and the element that didn't map to // itself. // // We cannot use `U()` as the default object, as `U` may be // a non-nullable type. Since the default `Any()` // will be overwritten before returning, all objects will be // of type `U` for the calling scope. all_identical = false; output = ArrayObj::CreateRepeated(static_cast(arr->size()), Any()); output->InitRange(0, arr->begin(), it); output->SetItem(it - arr->begin(), std::move(mapped)); it++; break; } } if (all_identical) { return data; } } else { // Path for incompatible types. The constexpr check for // compatible types isn't strictly necessary, as the first // (*it).same_as(mapped) would return false, but we might as well // avoid it altogether. // // We cannot use `U()` as the default object, as `U` may be a // non-nullable type. Since the default `Any()` will be // overwritten before returning, all objects will be of type `U` // for the calling scope. output = ArrayObj::CreateRepeated(static_cast(arr->size()), Any()); } // Normal path for incompatible types, or post-copy path for // copy-on-write instances. // // If the types are incompatible, then at this point `output` is // empty, and `it` points to the first element of the input. // // If the types were compatible, then at this point `output` // contains zero or more elements that mapped to themselves // followed by the first element that does not map to itself, and // `it` points to the element just after the first element that // does not map to itself. Because at least one element has been // changed, we no longer have the opportunity to avoid a copy, so // we don't need to check the result. // // In both cases, `it` points to the next element to be processed, // so we can either start or resume the iteration from that point, // with no further checks on the result. for (; it != arr->end(); it++) { U mapped = fmap(details::AnyUnsafe::CopyFromAnyViewAfterCheck(*it)); output->SetItem(it - arr->begin(), std::move(mapped)); } return output; } template friend class Array; }; /*! * \brief Concat two Arrays. * \param lhs first Array to be concatenated. * \param rhs second Array to be concatenated. * \return The concatenated Array. Original Arrays are kept unchanged. */ template || TypeTraits::convert_enabled>> inline Array Concat(Array lhs, const Array& rhs) { for (const auto& x : rhs) { lhs.push_back(x); } return std::move(lhs); } /*! * \brief Specialize make_object * \return The empty array object. */ template <> inline ObjectPtr make_object() { return ArrayObj::Empty(); } // Traits for Array template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public SeqTypeTraitsBase>, Array, T> { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIArray; static constexpr int32_t kPrimaryTypeIndex = TypeIndex::kTVMFFIArray; static constexpr int32_t kOtherTypeIndex = TypeIndex::kTVMFFIList; static constexpr const char* kTypeName = "Array"; static constexpr const char* kStaticTypeKey = StaticTypeKey::kTVMFFIArray; TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << kStaticTypeKey << R"(","args":[)"; oss << details::TypeSchema::v(); oss << "]}"; return oss.str(); } }; namespace details { template inline constexpr bool type_contains_v, Array> = type_contains_v; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_ARRAY_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/container_details.h000066400000000000000000000156161521067262500240010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/container_details.h * \brief Common utilities for typed container types. */ #ifndef TVM_FFI_CONTAINER_CONTAINER_DETAILS_H_ #define TVM_FFI_CONTAINER_CONTAINER_DETAILS_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace details { /*! * \brief iterator adapter that adapts TIter to return another type. * \tparam Converter a struct that contains converting function * \tparam TIter the content iterator type. */ template class IterAdapter { public: using difference_type = typename std::iterator_traits::difference_type; using value_type = typename Converter::ResultType; using pointer = const typename Converter::ResultType*; using reference = const typename Converter::ResultType; using iterator_category = typename std::iterator_traits::iterator_category; explicit IterAdapter(TIter iter) : iter_(iter) {} IterAdapter& operator++() { ++iter_; return *this; } IterAdapter& operator--() { --iter_; return *this; } IterAdapter operator++(int) { IterAdapter copy = *this; ++iter_; return copy; } IterAdapter operator--(int) { IterAdapter copy = *this; --iter_; return copy; } IterAdapter operator+(difference_type offset) const { return IterAdapter(iter_ + offset); } IterAdapter operator-(difference_type offset) const { return IterAdapter(iter_ - offset); } IterAdapter& operator+=(difference_type offset) { iter_ += offset; return *this; } IterAdapter& operator-=(difference_type offset) { iter_ -= offset; return *this; } template inline std::enable_if_t, typename T::difference_type> operator-(const IterAdapter& rhs) const { return iter_ - rhs.iter_; } bool operator==(IterAdapter other) const { return iter_ == other.iter_; } bool operator!=(IterAdapter other) const { return !(*this == other); } reference operator*() const { return Converter::convert(*iter_); } private: TIter iter_; }; /*! * \brief iterator adapter that adapts TIter to return another type. * \tparam Converter a struct that contains converting function * \tparam TIter the content iterator type. */ template class ReverseIterAdapter { public: using difference_type = typename std::iterator_traits::difference_type; using value_type = typename Converter::ResultType; using pointer = const typename Converter::ResultType*; using reference = const typename Converter::ResultType; using iterator_category = typename std::iterator_traits::iterator_category; explicit ReverseIterAdapter(TIter iter) : iter_(iter) {} ReverseIterAdapter& operator++() { --iter_; return *this; } ReverseIterAdapter& operator--() { ++iter_; return *this; } ReverseIterAdapter operator++(int) { ReverseIterAdapter copy = *this; --iter_; return copy; } ReverseIterAdapter operator--(int) { ReverseIterAdapter copy = *this; ++iter_; return copy; } ReverseIterAdapter operator+(difference_type offset) const { return ReverseIterAdapter(iter_ - offset); } template inline std::enable_if_t, typename T::difference_type> operator-(const ReverseIterAdapter& rhs) const { return rhs.iter_ - iter_; } bool operator==(ReverseIterAdapter other) const { return iter_ == other.iter_; } bool operator!=(ReverseIterAdapter other) const { return !(*this == other); } reference operator*() const { return Converter::convert(*iter_); } private: TIter iter_; }; /*! * \brief Check if T is compatible with Any. * * \tparam T The type to check. * \return True if T is compatible with Any, false otherwise. */ template inline constexpr bool storage_enabled_v = std::is_same_v || TypeTraits::storage_enabled; /*! * \brief Check if all T are compatible with Any. * * \tparam T The type to check. * \return True if T is compatible with Any, false otherwise. */ template inline constexpr bool all_storage_enabled_v = (storage_enabled_v && ...); /*! * \brief Check if all T are compatible with Any. * * \tparam T The type to check. * \return True if T is compatible with Any, false otherwise. */ template inline constexpr bool all_object_ref_v = (std::is_base_of_v && ...); /** * \brief Check if Any storage of Derived can always be directly used as Base. * * \tparam Base The base type. * \tparam Derived The derived type. * \return True if Derived's storage can be used as Base's storage, false otherwise. */ template inline constexpr bool type_contains_v = std::is_base_of_v || std::is_same_v; // special case for Any template inline constexpr bool type_contains_v = true; /*! * \brief Create a string of the container type. * \tparam V The types of the elements in the container. * \param name The name of the container type. * \return A string of the container type. */ template std::string ContainerTypeStr(const char* name) { std::stringstream ss; // helper to construct concated string of TypeStr class TypeStrHelper { public: TypeStrHelper(std::stringstream& stream) : stream_(stream) {} // NOLINT(*) TypeStrHelper& operator<<(const std::string& str) { if (counter_ > 0) { stream_ << ", "; } stream_ << str; counter_++; return *this; } private: std::stringstream& stream_; // NOLINT(*) int counter_ = 0; }; TypeStrHelper helper(ss); ss << name << '<'; (helper << ... << Type2Str::v()); ss << '>'; return ss.str(); } } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_CONTAINER_DETAILS_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/dict.h000066400000000000000000000267121521067262500212340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/dict.h * \brief Mutable dictionary container type. * * All handles sharing the same DictObj see mutations immediately. */ #ifndef TVM_FFI_CONTAINER_DICT_H_ #define TVM_FFI_CONTAINER_DICT_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief Dict object — mutable map with shared reference semantics. */ class DictObj : public MapBaseObj { public: /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIDict; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIDict, DictObj, Object); /// \endcond protected: template friend class Dict; }; static_assert(sizeof(DictObj) == sizeof(MapBaseObj), "DictObj must match MapBaseObj layout"); /*! * \brief Mutable dictionary container with shared reference semantics. * * Mutations happen directly on the underlying shared DictObj. * All handles sharing the same DictObj see mutations immediately. * * \tparam K The key type. * \tparam V The value type. */ template && details::storage_enabled_v>> class Dict : public ObjectRef { public: /*! \brief The key type of the dict */ using key_type = K; /*! \brief The mapped type of the dict */ using mapped_type = V; /*! \brief The iterator type of the dict */ class iterator; /*! * \brief Construct a Dict with UnsafeInit */ explicit Dict(UnsafeInit tag) : ObjectRef(tag) {} /*! * \brief default constructor */ Dict() { data_ = DictObj::Empty(); } /*! * \brief move constructor * \param other source */ Dict(Dict&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief copy constructor * \param other source */ Dict(const Dict& other) // NOLINT(google-explicit-constructor) : ObjectRef(other.data_) {} /*! * \brief Move constructor * \param other The other dict * \tparam KU The key type of the other dict * \tparam VU The mapped type of the other dict */ template && details::type_contains_v>> Dict(Dict&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief Copy constructor * \param other The other dict * \tparam KU The key type of the other dict * \tparam VU The mapped type of the other dict */ template && details::type_contains_v>> // NOLINTNEXTLINE(google-explicit-constructor) Dict(const Dict& other) : ObjectRef(other.data_) {} /*! * \brief Move assignment * \param other The other dict */ Dict& operator=(Dict&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Copy assignment * \param other The other dict */ Dict& operator=(const Dict& other) { data_ = other.data_; return *this; } /*! * \brief Move assignment * \param other The other dict * \tparam KU The key type of the other dict * \tparam VU The mapped type of the other dict */ template && details::type_contains_v>> Dict& operator=(Dict&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Copy assignment * \param other The other dict * \tparam KU The key type of the other dict * \tparam VU The mapped type of the other dict */ template && details::type_contains_v>> Dict& operator=(const Dict& other) { data_ = other.data_; return *this; } /*! * \brief constructor from pointer * \param n the container pointer */ explicit Dict(ObjectPtr n) : ObjectRef(n) {} /*! * \brief constructor from iterator * \param begin begin of iterator * \param end end of iterator * \tparam IterType The type of iterator */ template Dict(IterType begin, IterType end) { data_ = DictObj::CreateFromRange(begin, end); } /*! * \brief constructor from initializer list * \param init The initalizer list */ Dict(std::initializer_list> init) { data_ = DictObj::CreateFromRange(init.begin(), init.end()); } /*! * \brief constructor from unordered_map * \param init The unordered_map */ template Dict(const std::unordered_map& init) { // NOLINT(*) data_ = DictObj::CreateFromRange(init.begin(), init.end()); } /*! * \brief Read element from dict. * \param key The key * \return the corresponding element. */ V at(const K& key) const { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(GetDictObj()->at(key)); } /*! * \brief Read element from dict. * \param key The key * \return the corresponding element. */ V operator[](const K& key) const { return this->at(key); } /*! \return The size of the dict */ size_t size() const { DictObj* n = GetDictObj(); return n == nullptr ? 0 : n->size(); } /*! \return The number of elements of the key */ size_t count(const K& key) const { DictObj* n = GetDictObj(); return n == nullptr ? 0 : n->count(key); } /*! \return whether dict is empty */ bool empty() const { return size() == 0; } /*! \brief Release reference to all the elements */ void clear() { DictObj* n = GetDictObj(); if (n != nullptr) { n->clear(); } } /*! * \brief Set a key-value pair in the Dict (mutates in-place). * \param key The index key. * \param value The value to be set. */ void Set(const K& key, const V& value) { EnsureDictObj(); ObjectPtr new_container = MapBaseObj::InsertMaybeReHash(DictObj::KVType(key, value), data_); if (new_container != nullptr) { static_cast(data_.get())->InplaceSwitchTo(std::move(new_container)); } } /*! \return begin iterator */ iterator begin() const { return iterator(GetDictObj()->begin()); } /*! \return end iterator */ iterator end() const { return iterator(GetDictObj()->end()); } /*! \return find the key and returns the associated iterator */ iterator find(const K& key) const { return iterator(GetDictObj()->find(key)); } /*! \return The value associated with the key, std::nullopt if not found */ std::optional Get(const K& key) const { DictObj::iterator iter = GetDictObj()->find(key); if (iter == GetDictObj()->end()) { return std::nullopt; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(iter->second); } /*! * \brief Erase the entry associated with the key (mutates in-place) * \param key The key */ void erase(const K& key) { DictObj* n = GetDictObj(); if (n != nullptr) { n->erase(key); } } /*! \brief specify container node */ using ContainerType = DictObj; /// \cond Doxygen_Suppress /*! \brief Iterator of the hash map */ class iterator { public: using iterator_category = std::bidirectional_iterator_tag; using difference_type = int64_t; using value_type = const std::pair; using pointer = value_type*; using reference = value_type; iterator() : itr() {} /*! \brief Compare iterators */ bool operator==(const iterator& other) const { return itr == other.itr; } /*! \brief Compare iterators */ bool operator!=(const iterator& other) const { return itr != other.itr; } /*! \brief De-reference iterators is not allowed */ pointer operator->() const = delete; /*! \brief De-reference iterators */ reference operator*() const { auto& kv = *itr; return std::make_pair(details::AnyUnsafe::CopyFromAnyViewAfterCheck(kv.first), details::AnyUnsafe::CopyFromAnyViewAfterCheck(kv.second)); } /*! \brief Prefix self increment, e.g. ++iter */ iterator& operator++() { ++itr; return *this; } /*! \brief Suffix self increment */ iterator operator++(int) { iterator copy = *this; ++(*this); return copy; } /*! \brief Prefix self decrement, e.g. --iter */ iterator& operator--() { --itr; return *this; } /*! \brief Suffix self decrement */ iterator operator--(int) { iterator copy = *this; --(*this); return copy; } private: iterator(const DictObj::iterator& itr) // NOLINT(*) : itr(itr) {} template friend class Dict; DictObj::iterator itr; }; /// \endcond private: /*! \brief Return data_ as type of pointer of DictObj */ DictObj* GetDictObj() const { return static_cast(data_.get()); } /*! \brief Ensure we have a valid DictObj */ void EnsureDictObj() { if (data_ == nullptr) { data_ = DictObj::Empty(); } } template friend class Dict; }; // Traits for Dict template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public MapTypeTraitsBase>, Dict, K, V> { static constexpr int32_t kPrimaryTypeIndex = TypeIndex::kTVMFFIDict; static constexpr int32_t kOtherTypeIndex = TypeIndex::kTVMFFIMap; static constexpr const char* kTypeName = "Dict"; TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << StaticTypeKey::kTVMFFIDict << R"(","args":[)"; oss << details::TypeSchema::v() << ","; oss << details::TypeSchema::v(); oss << "]}"; return oss.str(); } }; namespace details { template inline constexpr bool type_contains_v, Dict> = type_contains_v && type_contains_v; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_DICT_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/list.h000066400000000000000000000405001521067262500212530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/list.h * \brief Mutable list type. * * tvm::ffi::List is an erased mutable sequence container. */ #ifndef TVM_FFI_CONTAINER_LIST_H_ #define TVM_FFI_CONTAINER_LIST_H_ #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief List node content in list */ class ListObj : public SeqBaseObj { public: /*! * \brief Constructs a container with n elements. Each element is a copy of val * \param n The size of the container * \param val The init value * \return Ref-counted ListObj requested */ static ObjectPtr CreateRepeated(int64_t n, const Any& val) { ObjectPtr p = ListObj::Empty(n); Any* itr = p->MutableBegin(); for (int64_t& i = p->TVMFFISeqCell::size = 0; i < n; ++i) { new (itr++) Any(val); } return p; } /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIList; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIList, ListObj, Object); /// \endcond private: /*! * \brief Ensure the list has at least n slots. * \param n The lower bound of required capacity. * * \note Leak-safety: Any's move constructor is noexcept, so the * move loop below cannot throw and leave the buffer in a * partially-constructed state. */ void Reserve(int64_t n) { if (n <= TVMFFISeqCell::capacity) { return; } Any* old_data = MutableBegin(); Any* new_data = static_cast(::operator new(sizeof(Any) * static_cast(n))); for (int64_t i = 0; i < TVMFFISeqCell::size; ++i) { new (new_data + i) Any(std::move(old_data[i])); } for (int64_t j = 0; j < TVMFFISeqCell::size; ++j) { (old_data + j)->Any::~Any(); } data_deleter(data); data = new_data; TVMFFISeqCell::capacity = n; } /*! * \brief Create an empty ListObj with the given capacity. * \param n Required capacity * \return Ref-counted ListObj requested */ static ObjectPtr Empty(int64_t n = kInitSize) { if (n < 0) { TVM_FFI_THROW(ValueError) << "cannot construct a List of negative size"; } ObjectPtr p = make_object(); p->TVMFFISeqCell::capacity = n; p->TVMFFISeqCell::size = 0; p->data = n == 0 ? nullptr : static_cast(::operator new(sizeof(Any) * n)); p->data_deleter = RawDataDeleter; return p; } static void RawDataDeleter(void* data) { ::operator delete(data); } /*! \brief Initial size of ListObj */ static constexpr int64_t kInitSize = 4; /*! \brief Expansion factor of the List */ static constexpr int64_t kIncFactor = 2; template friend class List; template friend struct TypeTraits; }; /*! * \brief List, container representing a mutable contiguous sequence of ObjectRefs. * * Unlike Array, List is mutable and does not implement copy-on-write. * Mutations happen directly on the underlying shared ListObj. * * \note Thread safety: List is **not** thread-safe. Concurrent reads and writes * from multiple threads require external synchronization. * * \warning Because List elements are stored as `Any` (which may hold `ObjectRef` * pointers), it is possible to create reference cycles (e.g. a List that * contains itself). Such cycles will **not** be collected by the reference- * counting mechanism alone; avoid them in long-lived data structures. * * \tparam T The content Value type, must be compatible with tvm::ffi::Any */ template >> class List : public ObjectRef { public: /*! \brief The value type of the list */ using value_type = T; /*! \brief Construct a List with UnsafeInit */ explicit List(UnsafeInit tag) : ObjectRef(tag) {} /*! \brief default constructor */ List() { data_ = ListObj::Empty(0); } // NOLINT(modernize-use-equals-default) /*! \brief Move constructor */ List(List&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! \brief Copy constructor */ List(const List& other) : ObjectRef(other.data_) {} // NOLINT(google-explicit-constructor) /*! * \brief Constructor from another list * \tparam U The value type of the other list */ template >> List(List&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief Constructor from another list * \tparam U The value type of the other list */ template >> List(const List& other) // NOLINT(google-explicit-constructor) : ObjectRef(other.data_) {} /*! * \brief Move assignment from another list. * \param other The other list. */ TVM_FFI_INLINE List& operator=(List&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Assignment from another list. * \param other The other list. */ TVM_FFI_INLINE List& operator=(const List& other) { data_ = other.data_; return *this; } /*! * \brief Move assignment from another list. * \param other The other list. * \tparam U The value type of the other list. */ template >> TVM_FFI_INLINE List& operator=(List&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Assignment from another list. * \param other The other list. * \tparam U The value type of the other list. */ template >> TVM_FFI_INLINE List& operator=(const List& other) { data_ = other.data_; return *this; } /*! \brief Constructor from pointer */ explicit List(ObjectPtr n) : ObjectRef(std::move(n)) {} /*! * \brief Constructor from iterator * \param first begin of iterator * \param last end of iterator * \tparam IterType The type of iterator */ template List(IterType first, IterType last) { // NOLINT(performance-unnecessary-value-param) static_assert(is_valid_iterator_v, "IterType cannot be inserted into a tvm::List"); Assign(first, last); } /*! \brief constructor from initializer list */ List(std::initializer_list init) { // NOLINT(*) Assign(init.begin(), init.end()); } /*! \brief constructor from vector */ List(const std::vector& init) { // NOLINT(*) Assign(init.begin(), init.end()); } /*! * \brief Constructs a container with n elements. Each element is a copy of val * \param n The size of the container * \param val The init value */ explicit List(const size_t n, const T& val) { data_ = ListObj::CreateRepeated(n, val); } public: // iterators /// \cond Doxygen_Suppress struct ValueConverter { using ResultType = T; static T convert(const Any& n) { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(n); } }; /// \endcond /*! \brief The iterator type of the list */ using iterator = details::IterAdapter; /*! \brief The reverse iterator type of the list */ using reverse_iterator = details::ReverseIterAdapter; /*! \return begin iterator */ iterator begin() const { return iterator(GetListObj()->begin()); } /*! \return end iterator */ iterator end() const { return iterator(GetListObj()->end()); } /*! \return rbegin iterator */ reverse_iterator rbegin() const { return reverse_iterator(GetListObj()->end() - 1); } /*! \return rend iterator */ reverse_iterator rend() const { return reverse_iterator(GetListObj()->begin() - 1); } public: // const methods in std::vector /*! * \brief Immutably read i-th element from list. * \param i The index * \return the i-th element. */ T operator[](int64_t i) const { ListObj* p = GetListObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null list"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->at(i)); } /*! \return The size of the list */ size_t size() const { ListObj* p = GetListObj(); return p == nullptr ? 0 : p->size(); } /*! \return The capacity of the list */ size_t capacity() const { ListObj* p = GetListObj(); return p == nullptr ? 0 : p->SeqBaseObj::capacity(); } /*! \return Whether list is empty */ bool empty() const { return size() == 0; } /*! \return The first element of the list */ T front() const { ListObj* p = GetListObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null list"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->front()); } /*! \return The last element of the list */ T back() const { ListObj* p = GetListObj(); if (p == nullptr) { TVM_FFI_THROW(IndexError) << "cannot index a null list"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(p->back()); } public: // mutation in std::vector /*! * \brief push a new item to the back of the list * \param item The item to be pushed. */ void push_back(const T& item) { ListObj* p = EnsureCapacity(1); p->EmplaceInit(p->TVMFFISeqCell::size++, item); } /*! * \brief Emplace a new element at the back of the list * \param args The arguments to construct the new element */ template void emplace_back(Args&&... args) { ListObj* p = EnsureCapacity(1); p->EmplaceInit(p->TVMFFISeqCell::size++, std::forward(args)...); } /*! * \brief Insert an element into the given position * \param position An iterator pointing to the insertion point * \param val The element to insert */ void insert(iterator position, const T& val) { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot insert to a null list"; } int64_t idx = std::distance(begin(), position); EnsureCapacity(1)->insert(idx, Any(val)); } /*! * \brief Insert a range of elements into the given position * \param position An iterator pointing to the insertion point * \param first The begin iterator of the range * \param last The end iterator of the range */ template void insert(iterator position, IterType first, IterType last) { static_assert(is_valid_iterator_v, "IterType cannot be inserted into a tvm::List"); if (first == last) return; if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot insert to a null list"; } int64_t idx = std::distance(begin(), position); int64_t numel = std::distance(first, last); EnsureCapacity(numel)->insert(idx, first, last); } /*! \brief Remove the last item of the list */ void pop_back() { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot pop_back a null list"; } GetListObj()->pop_back(); } /*! * \brief Erase an element on the given position * \param position An iterator pointing to the element to be erased */ void erase(iterator position) { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot erase a null list"; } int64_t idx = std::distance(begin(), position); GetListObj()->erase(idx); } /*! * \brief Erase a given range of elements * \param first The begin iterator of the range * \param last The end iterator of the range */ void erase(iterator first, iterator last) { if (first == last) return; if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "cannot erase a null list"; } int64_t st = std::distance(begin(), first); int64_t ed = std::distance(begin(), last); GetListObj()->erase(st, ed); } /*! * \brief Resize the list. * \param n The new size. */ void resize(int64_t n) { if (n < 0) { TVM_FFI_THROW(ValueError) << "cannot resize a List to negative size"; } EnsureCapacity(std::max(0, n - static_cast(size())))->resize(n); } /*! * \brief Make sure the list has the capacity of at least n * \param n lower bound of the capacity */ void reserve(int64_t n) { EnsureListObj()->Reserve(n); } /*! \brief Release reference to all the elements */ void clear() { if (data_ != nullptr) { GetListObj()->clear(); } } public: // List's own methods /*! * \brief set i-th element of the list. * \param i The index * \param value The value to be set. */ void Set(int64_t i, T value) { EnsureListObj()->SetItem(i, std::move(value)); } /*! \return The underlying ListObj */ ListObj* GetListObj() const { return static_cast(data_.get()); } /*! * \brief reset the list to content from iterator. * \param first begin of iterator * \param last end of iterator * \tparam IterType The type of iterator */ template void Assign(IterType first, IterType last) { // NOLINT(performance-unnecessary-value-param) int64_t cap = std::distance(first, last); if (cap < 0) { TVM_FFI_THROW(ValueError) << "cannot construct a List of negative size"; } ListObj* p = EnsureListObj(); p->Reserve(cap); p->clear(); Any* itr = p->MutableBegin(); for (int64_t& i = p->TVMFFISeqCell::size = 0; i < cap; ++i, ++first, ++itr) { new (itr) Any(*first); } } /*! \brief specify container node */ using ContainerType = ListObj; private: /*! * \brief Ensure the list object exists and has room for reserve_extra new entries. * \param reserve_extra Number of extra slots needed * \return ListObj pointer */ ListObj* EnsureCapacity(int64_t reserve_extra) { ListObj* p = EnsureListObj(); if (p->TVMFFISeqCell::capacity >= p->TVMFFISeqCell::size + reserve_extra) { return p; } int64_t cap = p->TVMFFISeqCell::capacity * ListObj::kIncFactor; cap = std::max(cap, p->TVMFFISeqCell::size + reserve_extra); p->Reserve(cap); return p; } /*! * \brief Ensure this list has a container object. * \return ListObj pointer */ ListObj* EnsureListObj() { if (data_ == nullptr) { data_ = ListObj::Empty(); } return static_cast(data_.get()); } template friend class List; }; // Traits for List template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public SeqTypeTraitsBase>, List, T> { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIList; static constexpr int32_t kPrimaryTypeIndex = TypeIndex::kTVMFFIList; static constexpr int32_t kOtherTypeIndex = TypeIndex::kTVMFFIArray; static constexpr const char* kTypeName = "List"; static constexpr const char* kStaticTypeKey = StaticTypeKey::kTVMFFIList; TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << kStaticTypeKey << R"(","args":[)"; oss << details::TypeSchema::v(); oss << "]}"; return oss.str(); } }; namespace details { template inline constexpr bool type_contains_v, List> = type_contains_v; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_LIST_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/map.h000066400000000000000000000277221521067262500210700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/map.h * \brief Immutable Map container type. */ #ifndef TVM_FFI_CONTAINER_MAP_H_ #define TVM_FFI_CONTAINER_MAP_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief Map object */ class MapObj : public MapBaseObj { public: /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIMap; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIMap, MapObj, Object); /// \endcond protected: template friend class Map; }; /*! * \brief Map container of NodeRef->NodeRef in DSL graph. * Map implements copy on write semantics, which means map is mutable * but copy will happen when array is referenced in more than two places. * * operator[] only provide const acces, use Set to mutate the content. * \tparam K The key NodeRef type. * \tparam V The value NodeRef type. */ template && details::storage_enabled_v>> class Map : public ObjectRef { public: /*! \brief The key type of the map */ using key_type = K; /*! \brief The mapped type of the map */ using mapped_type = V; /*! \brief The iterator type of the map */ class iterator; /*! * \brief Construct an Map with UnsafeInit */ explicit Map(UnsafeInit tag) : ObjectRef(tag) {} /*! * \brief default constructor */ Map() { data_ = MapObj::Empty(); } /*! * \brief move constructor * \param other source */ Map(Map&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief copy constructor * \param other source */ Map(const Map& other) // NOLINT(google-explicit-constructor) : ObjectRef(other.data_) {} /*! * \brief Move constructor * \param other The other map * \tparam KU The key type of the other map * \tparam VU The mapped type of the other map */ template && details::type_contains_v>> Map(Map&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other.data_)) {} /*! * \brief Copy constructor * \param other The other map * \tparam KU The key type of the other map * \tparam VU The mapped type of the other map */ template && details::type_contains_v>> Map(const Map& other) : ObjectRef(other.data_) {} // NOLINT(google-explicit-constructor) /*! * \brief Move assignment * \param other The other map */ Map& operator=(Map&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Copy assignment * \param other The other map */ Map& operator=(const Map& other) { data_ = other.data_; return *this; } /*! * \brief Move assignment * \param other The other map * \tparam KU The key type of the other map * \tparam VU The mapped type of the other map */ template && details::type_contains_v>> Map& operator=(Map&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Copy assignment * \param other The other map * \tparam KU The key type of the other map * \tparam VU The mapped type of the other map */ template && details::type_contains_v>> Map& operator=(const Map& other) { data_ = other.data_; return *this; } /*! * \brief constructor from pointer * \param n the container pointer */ explicit Map(ObjectPtr n) : ObjectRef(n) {} /*! * \brief constructor from iterator * \param begin begin of iterator * \param end end of iterator * \tparam IterType The type of iterator */ template Map(IterType begin, IterType end) { data_ = MapObj::CreateFromRange(begin, end); } /*! * \brief constructor from initializer list * \param init The initalizer list */ Map(std::initializer_list> init) { data_ = MapObj::CreateFromRange(init.begin(), init.end()); } /*! * \brief constructor from unordered_map * \param init The unordered_map */ template Map(const std::unordered_map& init) { // NOLINT(*) data_ = MapObj::CreateFromRange(init.begin(), init.end()); } /*! * \brief Read element from map. * \param key The key * \return the corresonding element. */ const V at(const K& key) const { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(GetMapObj()->at(key)); } /*! * \brief Read element from map. * \param key The key * \return the corresonding element. */ const V operator[](const K& key) const { return this->at(key); } /*! \return The size of the array */ size_t size() const { MapObj* n = GetMapObj(); return n == nullptr ? 0 : n->size(); } /*! \return The number of elements of the key */ size_t count(const K& key) const { MapObj* n = GetMapObj(); return n == nullptr ? 0 : GetMapObj()->count(key); } /*! \return whether array is empty */ bool empty() const { return size() == 0; } /*! \brief Release reference to all the elements */ void clear() { MapObj* n = GetMapObj(); if (n != nullptr) { data_ = MapObj::Empty(); } } /*! * \brief set the Map. * \param key The index key. * \param value The value to be setted. */ void Set(const K& key, const V& value) { CopyOnWrite(); ObjectPtr new_data = MapObj::InsertMaybeReHash(MapObj::KVType(key, value), data_); if (new_data != nullptr) { data_ = std::move(new_data); } } /*! \return begin iterator */ iterator begin() const { return iterator(GetMapObj()->begin()); } /*! \return end iterator */ iterator end() const { return iterator(GetMapObj()->end()); } /*! \return find the key and returns the associated iterator */ iterator find(const K& key) const { return iterator(GetMapObj()->find(key)); } /*! \return The value associated with the key, std::nullopt if not found */ std::optional Get(const K& key) const { MapObj::iterator iter = GetMapObj()->find(key); if (iter == GetMapObj()->end()) { return std::nullopt; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(iter->second); } /*! * \brief Erase the entry associated with the key * \param key The key */ void erase(const K& key) { CopyOnWrite()->erase(key); } /*! * \brief copy on write semantics * Do nothing if current handle is the unique copy of the array. * Otherwise make a new copy of the array to ensure the current handle * hold a unique copy. * * \return Handle to the internal node container(which guarantees to be unique) */ MapObj* CopyOnWrite() { if (data_.get() == nullptr) { data_ = MapObj::Empty(); } else if (!data_.unique()) { data_ = MapObj::CopyFrom(GetMapObj()); } return GetMapObj(); } /*! \brief specify container node */ using ContainerType = MapObj; /// \cond Doxygen_Suppress /*! \brief Iterator of the hash map */ class iterator { public: using iterator_category = std::bidirectional_iterator_tag; using difference_type = int64_t; using value_type = const std::pair; using pointer = value_type*; using reference = value_type; iterator() : itr() {} /*! \brief Compare iterators */ bool operator==(const iterator& other) const { return itr == other.itr; } /*! \brief Compare iterators */ bool operator!=(const iterator& other) const { return itr != other.itr; } /*! \brief De-reference iterators is not allowed */ pointer operator->() const = delete; /*! \brief De-reference iterators */ reference operator*() const { auto& kv = *itr; return std::make_pair(details::AnyUnsafe::CopyFromAnyViewAfterCheck(kv.first), details::AnyUnsafe::CopyFromAnyViewAfterCheck(kv.second)); } /*! \brief Prefix self increment, e.g. ++iter */ iterator& operator++() { ++itr; return *this; } /*! \brief Suffix self increment */ iterator operator++(int) { iterator copy = *this; ++(*this); return copy; } /*! \brief Prefix self decrement, e.g. --iter */ iterator& operator--() { --itr; return *this; } /*! \brief Suffix self decrement */ iterator operator--(int) { iterator copy = *this; --(*this); return copy; } private: iterator(const MapObj::iterator& itr) // NOLINT(*) : itr(itr) {} template friend class Map; MapObj::iterator itr; }; /// \endcond private: /*! \brief Return data_ as type of pointer of MapObj */ MapObj* GetMapObj() const { return static_cast(data_.get()); } template friend class Map; }; /*! * \brief Merge two Maps. * \param lhs the first Map to merge. * \param rhs the second Map to merge. * @return The merged Array. Original Maps are kept unchanged. */ template && details::storage_enabled_v>> inline Map Merge(Map lhs, const Map& rhs) { for (const auto& p : rhs) { lhs.Set(p.first, p.second); } return std::move(lhs); } // Traits for Map template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public MapTypeTraitsBase>, Map, K, V> { static constexpr int32_t kPrimaryTypeIndex = TypeIndex::kTVMFFIMap; static constexpr int32_t kOtherTypeIndex = TypeIndex::kTVMFFIDict; static constexpr const char* kTypeName = "Map"; TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << StaticTypeKey::kTVMFFIMap << R"(","args":[)"; oss << details::TypeSchema::v() << ","; oss << details::TypeSchema::v(); oss << "]}"; return oss.str(); } }; namespace details { template inline constexpr bool type_contains_v, Map> = type_contains_v && type_contains_v; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_MAP_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/map_base.h000066400000000000000000001766241521067262500220700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/map_base.h * \brief Base class for Map container types. */ #ifndef TVM_FFI_CONTAINER_MAP_BASE_H_ #define TVM_FFI_CONTAINER_MAP_BASE_H_ #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /// \cond Doxygen_Suppress #if TVM_FFI_DEBUG_WITH_ABI_CHANGE #define TVM_FFI_MAP_FAIL_IF_CHANGED() \ TVM_FFI_ICHECK(state_marker == self->state_marker) << "Concurrent modification of the Map"; #else #define TVM_FFI_MAP_FAIL_IF_CHANGED() #endif // TVM_FFI_DEBUG_WITH_ABI_CHANGE /// \endcond /*! * \brief Shared content of all specializations of hash map * * MapBaseObj is transparent to the FFI type system (no type index), * following the same pattern as BytesObjBase. */ class MapBaseObj : public Object { public: /*! \brief Type of the keys in the hash map */ using key_type = Any; /*! \brief Type of the values in the hash map */ using mapped_type = Any; /*! \brief Type of value stored in the hash map */ using KVType = std::pair; /// \cond Doxygen_Suppress /*! \brief Type of raw storage of the key-value pair in the hash map */ struct KVRawStorageType { TVMFFIAny first; TVMFFIAny second; }; /// \endcond /*! \brief Iterator class */ class iterator; static_assert(std::is_standard_layout_v, "KVType is not standard layout"); static_assert(sizeof(KVType) == 32, "sizeof(KVType) incorrect"); /*! * \brief Number of elements in the MapBaseObj * \return The result */ size_t size() const { return size_; } /*! * \brief Count the number of times a key exists in the hash map * \param key The indexing key * \return The result, 0 or 1 */ size_t count(const key_type& key) const; /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The const reference to the value */ const mapped_type& at(const key_type& key) const; /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The mutable reference to the value */ mapped_type& at(const key_type& key); /*! \return begin iterator */ iterator begin() const; /*! \return end iterator */ iterator end() const; /*! * \brief Index value associated with a key * \param key The indexing key * \return The iterator of the entry associated with the key, end iterator if not exists */ iterator find(const key_type& key) const; /*! * \brief Erase the entry associated with the iterator * \param position The iterator */ void erase(const iterator& position); /*! * \brief Erase the entry associated with the key, do nothing if not exists * \param key The indexing key */ void erase(const key_type& key) { erase(find(key)); } /*! * \brief clear all entries */ void clear(); /// \cond Doxygen_Suppress class iterator { public: using iterator_category = std::forward_iterator_tag; using difference_type = int64_t; using value_type = KVType; using pointer = KVType*; using reference = KVType&; /*! \brief Default constructor */ #if TVM_FFI_DEBUG_WITH_ABI_CHANGE iterator() : state_marker(0), index(0), self(nullptr) {} #else iterator() : index(0), self(nullptr) {} #endif // TVM_FFI_DEBUG_WITH_ABI_CHANGE /*! \brief Compare iterators */ bool operator==(const iterator& other) const { TVM_FFI_MAP_FAIL_IF_CHANGED() return index == other.index && self == other.self; } /*! \brief Compare iterators */ bool operator!=(const iterator& other) const { return !(*this == other); } /*! \brief De-reference iterators */ pointer operator->() const; /*! \brief De-reference iterators */ reference operator*() const { TVM_FFI_MAP_FAIL_IF_CHANGED() return *((*this).operator->()); } /*! \brief Prefix self increment, e.g. ++iter */ iterator& operator++(); /*! \brief Prefix self decrement, e.g. --iter */ iterator& operator--(); /*! \brief Suffix self increment */ iterator operator++(int) { TVM_FFI_MAP_FAIL_IF_CHANGED() iterator copy = *this; ++(*this); return copy; } /*! \brief Suffix self decrement */ iterator operator--(int) { TVM_FFI_MAP_FAIL_IF_CHANGED() iterator copy = *this; --(*this); return copy; } protected: #if TVM_FFI_DEBUG_WITH_ABI_CHANGE uint64_t state_marker; /*! \brief Construct by value */ iterator(uint64_t index, const MapBaseObj* self) : state_marker(self->state_marker), index(index), self(self) {} #else iterator(uint64_t index, const MapBaseObj* self) : index(index), self(self) {} #endif // TVM_FFI_DEBUG_WITH_ABI_CHANGE /*! \brief The position on the array */ uint64_t index; /*! \brief The container it points to */ const MapBaseObj* self; friend class DenseMapBaseObj; friend class SmallMapBaseObj; }; /// \endcond protected: #if TVM_FFI_DEBUG_WITH_ABI_CHANGE uint64_t state_marker; #endif // TVM_FFI_DEBUG_WITH_ABI_CHANGE /*! * \brief Inplace switch the storage contents to the other map * * The other map will be consumed after this operation * The current content will be reset after this operation * * \param other The other map */ inline void InplaceSwitchTo(ObjectPtr&& other); /*! * \brief Create an empty container * \return The object created */ template static inline ObjectPtr Empty(); /*! * \brief Create the map using contents from the given iterators. * \param first Begin of iterator * \param last End of iterator * \tparam MapObjType The type of map object * \tparam IterType The type of iterator * \return ObjectPtr to the map created */ template static inline ObjectPtr CreateFromRange(IterType first, IterType last); /*! * \brief InsertMaybeReHash an entry into the given hash map * \param kv The entry to be inserted * \param map The reference to the map container * \tparam MapObjType The type of map object * \return A new container if re-hashing happens, nullptr otherwise */ template static inline ObjectPtr InsertMaybeReHash(KVType&& kv, const ObjectPtr& map); /*! * \brief Create an empty container with elements copying from another SmallMapBaseObj * \param from The source container * \tparam MapObjType The type of map object * \return The object created */ template static inline ObjectPtr CopyFrom(MapBaseObj* from); /*! * \brief data pointer to the data region of the map. * \note For immutable inplace small map we do not need data_, * but we keep it here for future compact with mutable container. */ void* data_; /*! \brief number of entries in the container */ uint64_t size_; /*! \brief number of slots */ uint64_t slots_; /*! * \brief Small layout tag mask * \note The most significant bit is used to indicate the small map layout. */ static constexpr uint64_t kSmallTagMask = static_cast(1) << 63; /*! * \brief Check if the map is a small map * \return True if the map is a small map */ bool IsSmallMap() const { return (slots_ & kSmallTagMask) != 0ull; } /*! * \brief Optional data deleter when data is allocated separately * and its deletion is not managed by MapBaseObj::deleter_. */ void (*data_deleter_)(void*) = nullptr; // Reference class friend class SmallMapBaseObj; friend class DenseMapBaseObj; template friend class Dict; template friend struct TypeTraits; }; /*! \brief A specialization of hash map that implements the idea of array-based hash map. * Another reference implementation can be found [1]. * * A. Overview * * DenseMapBaseObj did several improvements over traditional separate chaining hash, * in terms of cache locality, memory footprints and data organization. * * A1. Implicit linked list. For better cache locality, instead of using linked list * explicitly for each bucket, we store list data into a single array that spans contiguously * in memory, and then carefully design access patterns to make sure most of them fall into * a single cache line. * * A2. 1-byte metadata. There is only 1 byte overhead for each slot in the array to indexing and * traversal. This can be divided in 3 parts. * 1) Reserved code: (0b11111111)_2 indicates a slot is empty; (0b11111110)_2 indicates protected, * which means the slot is empty but not allowed to be written. * 2) If not empty or protected, the highest bit is used to indicate whether data in the slot is * head of a linked list. * 3) The rest 7 bits are used as the "next pointer" (i.e. pointer to the next element). On 64-bit * architecture, an ordinary pointer can take up to 8 bytes, which is not acceptable overhead when * dealing with 16-byte ObjectRef pairs. Based on a commonly noticed fact that the lists are * relatively short (length <= 3) in hash maps, we follow [1]'s idea that only allows the pointer to * be one of the 126 possible values, i.e. if the next element of i-th slot is (i + x)-th element, * then x must be one of the 126 pre-defined values. * * A3. Data blocking. We organize the array in the way that every 16 elements forms a data block. * The 16-byte metadata of those 16 elements are stored together, followed by the real data, i.e. * 16 key-value pairs. * * B. Implementation details * * B1. Power-of-2 table size and Fibonacci Hashing. We use power-of-two as table size to avoid * modulo for more efficient arithmetics. To make the hash-to-slot mapping distribute more evenly, * we use the Fibonacci Hashing [2] trick. * * B2. Traverse a linked list in the array. * 1) List head. Assume Fibonacci Hashing maps a given key to slot i, if metadata at slot i * indicates that it is list head, then we found the head; otherwise the list is empty. No probing * is done in this procedure. 2) Next element. To find the next element of a non-empty slot i, we * look at the last 7 bits of the metadata at slot i. If they are all zeros, then it is the end of * list; otherwise, we know that the next element is (i + candidates[the-last-7-bits]). * * B3. InsertMaybeReHash an element. Following B2, we first traverse the linked list to see if this * element is in the linked list, and if not, we put it at the end by probing the next empty * position in one of the 126 candidate positions. If the linked list does not even exist, but the * slot for list head has been occupied by another linked list, we should find this intruder another * place. * * B4. Quadratic probing with triangle numbers. In open address hashing, it is provable that probing * with triangle numbers can traverse power-of-2-sized table [3]. In our algorithm, we follow the * suggestion in [1] that also use triangle numbers for "next pointer" as well as sparing for list * head. * * [1] https://github.com/skarupke/flat_hash_map * [2] https://programmingpraxis.com/2018/06/19/fibonacci-hash/ * [3] https://fgiesen.wordpress.com/2015/02/22/triangular-numbers-mod-2n/ */ class DenseMapBaseObj : public MapBaseObj { private: /*! \brief The number of elements in a memory block */ static constexpr int kBlockCap = 16; /*! \brief Maximum load factor of the hash map */ static constexpr double kMaxLoadFactor = 0.99; /*! \brief Binary representation of the metadata of an empty slot */ static constexpr uint8_t kEmptySlot = static_cast(0b11111111); /*! \brief Binary representation of the metadata of a protected slot */ static constexpr uint8_t kProtectedSlot = static_cast(0b11111110); /*! \brief Number of probing choices available */ static constexpr int kNumJumpDists = 126; /*! \brief Index indicator to indicate an invalid index */ static constexpr uint64_t kInvalidIndex = std::numeric_limits::max(); /*! \brief Head of the implicit linked list */ struct ListNode; /*! \brief item type of the dense map, including a kv data and prev/next pointer */ struct ItemType { KVType data; uint64_t prev = kInvalidIndex; uint64_t next = kInvalidIndex; explicit ItemType(KVType&& data) : data(std::move(data)) {} explicit ItemType(key_type key, mapped_type value) : data(std::move(key), std::move(value)) {} }; /*! \brief POD type of a block of memory */ struct Block { uint8_t bytes[kBlockCap + kBlockCap * sizeof(ItemType)]; }; static_assert(sizeof(Block) == kBlockCap * (sizeof(ItemType) + 1), "sizeof(Block) incorrect"); static_assert(std::is_standard_layout_v, "Block is not standard layout"); /*! * \brief Deleter for the Block * \param data The pointer to the Block */ static void BlockDeleter(void* data) { delete[] static_cast(data); } public: using MapBaseObj::iterator; /*! * \brief Return the number of usable slots for Dense layout (MSB clear => identity). * \return The number of usable slots */ uint64_t NumSlots() const { return slots_; } /*! * \brief Destroy the DenseMapBaseObj */ ~DenseMapBaseObj() { this->Reset(); } /*! \return The number of elements of the key */ size_t count(const key_type& key) const { return !Search(key).IsNone(); } /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The const reference to the value */ const mapped_type& at(const key_type& key) const { return At(key); } /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The mutable reference to the value */ mapped_type& at(const key_type& key) { return At(key); } /*! * \brief Index value associated with a key * \param key The indexing key * \return The iterator of the entry associated with the key, end iterator if not exists */ iterator find(const key_type& key) const { ListNode node = Search(key); return node.IsNone() ? end() : iterator(node.index, this); } /*! * \brief Erase the entry associated with the iterator * \param position The iterator */ void erase(const iterator& position) { uint64_t index = position.index; if (position.self != nullptr && index <= this->NumSlots()) { Erase(ListNode(index, this)); } } /*! \return begin iterator */ iterator begin() const { return iterator(iter_list_head_, this); } /*! \return end iterator */ iterator end() const { return iterator(kInvalidIndex, this); } /*! * \brief clear all entries */ void clear() { uint64_t n_blocks = CalcNumBlocks(this->NumSlots()); for (uint64_t bi = 0; bi < n_blocks; ++bi) { uint8_t* meta_ptr = GetBlock(bi)->bytes; ItemType* data_ptr = reinterpret_cast(GetBlock(bi)->bytes + kBlockCap); for (int j = 0; j < kBlockCap; ++j, ++meta_ptr, ++data_ptr) { uint8_t& meta = *meta_ptr; if (meta != kProtectedSlot && meta != kEmptySlot) { meta = kEmptySlot; data_ptr->ItemType::~ItemType(); } } } size_ = 0; iter_list_head_ = kInvalidIndex; iter_list_tail_ = kInvalidIndex; } private: Block* GetBlock(size_t index) const { return static_cast(data_) + index; } /*! * \brief Unlink the entry from iterator list * \param node The node to be unlinked * \note This function is usually used before deletion, * and it does not change data content of the node. */ void IterListUnlink(ListNode node) { // update head and tail of iterator list if needed if (node.Item().prev == kInvalidIndex) { iter_list_head_ = node.Item().next; } else { ListNode prev_node(node.Item().prev, this); prev_node.Item().next = node.Item().next; } if (node.Item().next == kInvalidIndex) { iter_list_tail_ = node.Item().prev; } else { ListNode next_node(node.Item().next, this); next_node.Item().prev = node.Item().prev; } } /*! * \brief Insert the entry into tail of iterator list * \param node The node to be inserted * \note this function does not change data content of the node. */ void IterListPushBack(ListNode node) { node.Item().prev = iter_list_tail_; node.Item().next = kInvalidIndex; if (iter_list_tail_ != kInvalidIndex) { ListNode prev_node(iter_list_tail_, this); prev_node.Item().next = node.index; } if (iter_list_head_ == kInvalidIndex) { iter_list_head_ = node.index; } iter_list_tail_ = node.index; } /*! * \brief Replace node src by dst in the iter list * \param src The source node * \param dst The destination node, must be empty * \note This function does not change data content of the nodes, * which needs to be updated by the caller. */ void IterListReplaceNodeBy(ListNode src, ListNode dst) { // set link correctly on the dst dst.Item().prev = src.Item().prev; dst.Item().next = src.Item().next; // update prev and next of dst if (dst.Item().prev == kInvalidIndex) { iter_list_head_ = dst.index; } else { ListNode prev_node(dst.Item().prev, this); prev_node.Item().next = dst.index; } if (dst.Item().next == kInvalidIndex) { iter_list_tail_ = dst.index; } else { ListNode next_node(dst.Item().next, this); next_node.Item().prev = dst.index; } } /*! * \brief Search for the given key * \param key The key * \return ListNode that associated with the key */ ListNode Search(const key_type& key) const { if (this->size_ == 0) { return ListNode(); } for (ListNode iter = GetListHead(AnyHash()(key)); !iter.IsNone(); iter.MoveToNext(this)) { if (AnyEqual()(key, iter.Key())) { return iter; } } return ListNode(); } /*! * \brief Search for the given key, throw exception if not exists * \param key The key * \return ListNode that associated with the key */ mapped_type& At(const key_type& key) const { ListNode iter = Search(key); if (iter.IsNone()) { TVM_FFI_THROW(KeyError) << "key is not in Map"; } return iter.Val(); } /*! * \brief Try to insert a key, or do nothing if already exists * \param key The indexing key * \param result The linked-list entry found or just constructed * \return A boolean, indicating if actual insertion happens */ bool TryInsert(const key_type& key, ListNode* result) { if (slots_ == 0) { return false; } // required that `iter` to be the head of a linked list through which we can iterator ListNode iter = IndexFromHash(AnyHash()(key)); // `iter` can be: 1) empty; 2) body of an irrelevant list; 3) head of the relevant list // Case 1: empty if (iter.IsEmpty()) { iter.NewHead(ItemType(key, Any(nullptr))); this->size_ += 1; *result = iter; return true; } // Case 2: body of an irrelevant list if (!iter.IsHead()) { // we move the elements around and construct the single-element linked list return IsFull() ? false : TrySpareListHead(iter, key, result); } // Case 3: head of the relevant list // we iterate through the linked list until the end // make sure `iter` is the previous element of `next` ListNode next = iter; do { // find equal item, do not insert if (AnyEqual()(key, next.Key())) { // we plan to take next, so we need to unlink it from iterator list IterListUnlink(next); *result = next; return true; } // make sure `iter` is the previous element of `next` iter = next; } while (next.MoveToNext(this)); // `iter` is the tail of the linked list // always check capacity before insertion if (IsFull()) { return false; } // find the next empty slot uint8_t jump; if (!iter.GetNextEmpty(this, &jump, result)) { return false; } result->NewTail(ItemType(key, Any(nullptr))); // link `iter` to `empty`, and move forward iter.SetJump(jump); this->size_ += 1; return true; } /*! * \brief Spare an entry to be the head of a linked list. * As described in B3, during insertion, it is possible that the entire linked list does not * exist, but the slot of its head has been occupied by other linked lists. In this case, we need * to spare the slot by moving away the elements to another valid empty one to make insertion * possible. * \param target The given entry to be spared * \param key The indexing key * \param result The linked-list entry constructed as the head * \return A boolean, if actual insertion happens */ bool TrySpareListHead(ListNode target, const key_type& key, ListNode* result) { // `target` is not the head of the linked list // move the original item of `target` (if any) // and construct new item on the position `target` // To make `target` empty, we // 1) find `w` the previous element of `target` in the linked list // 2) copy the linked list starting from `r = target` // 3) paste them after `w` // read from the linked list after `r` ListNode r = target; // write to the tail of `w` ListNode w = target.FindPrev(this); // after `target` is moved, we disallow writing to the slot bool is_first = true; uint8_t r_meta, jump; ListNode empty; do { // `jump` describes how `w` is jumped to `empty` // rehash if there is no empty space after `w` if (!w.GetNextEmpty(this, &jump, &empty)) { return false; } // move `r` to `empty` // first move the data over empty.NewTail(ItemType(std::move(r.Data()))); // then move link list chain of r to empty // this needs to happen after NewTail so empty's prev/next get updated IterListReplaceNodeBy(r, empty); // explicit call destructor to destroy the item in `r` r.DestructData(); // clear the metadata of `r` r_meta = r.Meta(); if (is_first) { is_first = false; r.SetProtected(); } else { r.SetEmpty(); } // link `w` to `empty`, and move forward w.SetJump(jump); w = empty; // move `r` forward as well } while (r.MoveToNext(this, r_meta)); // finally we have done moving the linked list // fill data_ into `target` target.NewHead(ItemType(key, Any(nullptr))); this->size_ += 1; *result = target; return true; } /*! * \brief Remove a ListNode * \param iter The node to be removed */ void Erase(const ListNode& iter) { this->size_ -= 1; if (!iter.HasNext()) { // `iter` is the last if (!iter.IsHead()) { // cut the link if there is any iter.FindPrev(this).SetJump(0); } // unlink the node from iterator list IterListUnlink(iter); // IMPORTANT: must explicit call destructor `iter` to avoid memory leak // This is because we need to recycle iter's data iter.DestructData(); // set the meta data to be empty iter.SetEmpty(); } else { ListNode last = iter, prev = iter; for (last.MoveToNext(this); last.HasNext(); prev = last, last.MoveToNext(this)) { } // needs to first unlink iter from the list IterListUnlink(iter); // move data from last to iter iter.Data() = std::move(last.Data()); // Move link chain of iter to last as we stores last node to the new iter loc. IterListReplaceNodeBy(last, iter); // IMPORTANT: must explicit call destructor `last` to avoid memory leak // likely we don't need this in this particular case because Any move behavior // keep it here to be safe so code do not depend on specific move behavior of KVType last.DestructData(); // set the meta data to be empty last.SetEmpty(); prev.SetJump(0); } } /*! \brief Clear the container to empty, release all entries and memory acquired */ void Reset() { clear(); ReleaseMemory(); } /*! * \brief Inplace switch the storage contents to the other map * * The other map will be consumed after this operation * The current content will be reset after this operation * * \param other The other map */ void InplaceSwitchTo(ObjectPtr&& other) { this->Reset(); MapBaseObj* other_map = static_cast(other.get()); // to simplify implementation, inplace switch from another dense map // since all current switch usecases are pushing elements to map // so we don't need to handle small -> dense switch TVM_FFI_ICHECK(!other_map->IsSmallMap()); DenseMapBaseObj* other_dense_map = static_cast(other_map); DenseMapBaseObj* this_dense_map = this; this_dense_map->size_ = other_dense_map->size_; this_dense_map->slots_ = other_dense_map->slots_; this_dense_map->data_ = other_dense_map->data_; this_dense_map->data_deleter_ = other_dense_map->data_deleter_; this_dense_map->fib_shift_ = other_dense_map->fib_shift_; this_dense_map->iter_list_head_ = other_dense_map->iter_list_head_; this_dense_map->iter_list_tail_ = other_dense_map->iter_list_tail_; other_dense_map->data_ = nullptr; other_dense_map->data_deleter_ = nullptr; other_dense_map->size_ = 0; other_dense_map->slots_ = 0; } /*! * \brief Release the memory acquired by the container without deleting its entries stored inside */ void ReleaseMemory() { if (data_ != nullptr) { TVM_FFI_ICHECK(data_deleter_ != nullptr); data_deleter_(data_); } data_ = nullptr; data_deleter_ = nullptr; slots_ = 0; size_ = 0; fib_shift_ = 63; } /*! * \brief Create an empty container * \param fib_shift The fib shift provided * \param n_slots Number of slots required, should be power-of-two * \tparam MapObjType The type of map object * \return The object created */ template static ObjectPtr Empty(uint32_t fib_shift, uint64_t n_slots) { // Ensure even slot count (power-of-two expected by callers; this guard // makes the method robust if a non-even value slips through). ObjectPtr p = make_object(); uint64_t n_blocks = CalcNumBlocks(n_slots); Block* block = new Block[n_blocks]; p->data_ = block; // assign block deleter so even if we take re-alloc data // in another shared-lib that may have different malloc/free behavior // it will still be safe. p->data_deleter_ = BlockDeleter; p->SetSlotsAndDenseLayoutTag(n_slots); p->size_ = 0; p->fib_shift_ = fib_shift; p->iter_list_head_ = kInvalidIndex; p->iter_list_tail_ = kInvalidIndex; // manually set the type index to specialized subclass // data layout is ensured to be the same except for the type index static_assert(std::is_base_of_v); static_assert(sizeof(MapObjType) == sizeof(MapBaseObj)); p->header_.type_index = MapObjType::RuntimeTypeIndex(); for (uint64_t i = 0; i < n_blocks; ++i, ++block) { std::fill(block->bytes, block->bytes + kBlockCap, kEmptySlot); } return p; } /*! * \brief Create an empty container with elements copying from another DenseMapBaseObj * \param from The source container * \tparam MapObjType The type of map object * \return The object created */ template static ObjectPtr CopyFrom(DenseMapBaseObj* from) { ObjectPtr p = make_object(); uint64_t n_blocks = CalcNumBlocks(from->NumSlots()); p->data_ = new Block[n_blocks]; // assign block deleter so even if we take re-alloc data // in another shared-lib that may have different malloc/free behavior // it will still be safe. p->data_deleter_ = BlockDeleter; p->SetSlotsAndDenseLayoutTag(from->NumSlots()); p->size_ = from->size_; p->fib_shift_ = from->fib_shift_; p->iter_list_head_ = from->iter_list_head_; p->iter_list_tail_ = from->iter_list_tail_; // manually set the type index to specialized subclass // data layout is ensured to be the same except for the type index static_assert(std::is_base_of_v); static_assert(sizeof(MapObjType) == sizeof(MapBaseObj)); p->header_.type_index = MapObjType::RuntimeTypeIndex(); for (uint64_t bi = 0; bi < n_blocks; ++bi) { uint8_t* meta_ptr_from = from->GetBlock(bi)->bytes; ItemType* data_ptr_from = reinterpret_cast(from->GetBlock(bi)->bytes + kBlockCap); uint8_t* meta_ptr_to = p->GetBlock(bi)->bytes; ItemType* data_ptr_to = reinterpret_cast(p->GetBlock(bi)->bytes + kBlockCap); for (int j = 0; j < kBlockCap; ++j, ++meta_ptr_from, ++data_ptr_from, ++meta_ptr_to, ++data_ptr_to) { uint8_t& meta = *meta_ptr_to = *meta_ptr_from; TVM_FFI_ICHECK(meta != kProtectedSlot); if (meta != kEmptySlot) { new (data_ptr_to) ItemType(*data_ptr_from); } } } return p; } /*! * \brief InsertMaybeReHash an entry into the given hash map * \param kv The entry to be inserted * \param map The pointer to the map, can be changed if re-hashing happens */ template static ObjectPtr InsertMaybeReHash(KVType&& kv, const ObjectPtr& map) { DenseMapBaseObj* map_node = static_cast(map.get()); ListNode iter; // Try to insert. If succeed, we simply return if (map_node->TryInsert(kv.first, &iter)) { iter.Val() = std::move(kv.second); // update the iter list relation map_node->IterListPushBack(iter); return ObjectPtr(nullptr); } TVM_FFI_ICHECK(!map_node->IsSmallMap()); // Otherwise, start rehash ObjectPtr p = Empty(map_node->fib_shift_ - 1, map_node->NumSlots() * 2); // need to insert in the same order as the original map for (uint64_t index = map_node->iter_list_head_; index != kInvalidIndex;) { ListNode node(index, map_node); // now try move src_data into the new map, note that src may still not // be fully consumed into the call, but destructor will be called. ObjectPtr rehashed = InsertMaybeReHash(std::move(node.Data()), p); if (rehashed != nullptr) p = std::move(rehashed); // Important, needs to explicit call destructor in case move did remove // node's internal item index = node.Item().next; // IMPORTANT: must explicit call destructor `node` to avoid memory leak // We must call node.DestructData() here. // This is because std::move() arguments in IterMaybeReHash may or may not // explicitly move out the node.Data() // Remove this call will cause memory leak very likely. node.DestructData(); } { ObjectPtr rehashed = InsertMaybeReHash(std::move(kv), p); if (rehashed != nullptr) p = std::move(rehashed); } map_node->ReleaseMemory(); return p; } /*! * \brief Check whether the hash table is full * \return A boolean indicating whether hash table is full */ bool IsFull() const { // NOLINTNEXTLINE(bugprone-narrowing-conversions) return (size_ + 1) > static_cast(NumSlots()) * kMaxLoadFactor; } /*! * \brief Increment the pointer * \param index The pointer to be incremented * \return The increased pointer */ uint64_t IncItr(uint64_t index) const { // keep at the end of iterator if (index == kInvalidIndex) { return index; } ListNode node(index, this); return node.Item().next; } /*! * \brief Decrement the pointer * \param index The pointer to be decremented * \return The decreased pointer */ uint64_t DecItr(uint64_t index) const { // this is the end iterator, we need to return tail. if (index == kInvalidIndex) { return iter_list_tail_; } // circle around the iterator list, which is OK ListNode node(index, this); return node.Item().prev; } /*! * \brief De-reference the pointer * \param index The pointer to be dereferenced * \return The result */ KVType* DeRefItr(uint64_t index) const { return &ListNode(index, this).Data(); } /*! \brief Construct from hash code */ ListNode IndexFromHash(uint64_t hash_value) const { return ListNode(FibHash(hash_value, fib_shift_), this); } /*! \brief Construct from hash code if the position is head of list */ ListNode GetListHead(uint64_t hash_value) const { ListNode node = IndexFromHash(hash_value); return node.IsHead() ? node : ListNode(); } /*! \brief Construct the number of blocks in the hash table */ static uint64_t CalcNumBlocks(uint64_t n_slots) { return (n_slots + kBlockCap - 1) / kBlockCap; } /*! * \brief Calculate the power-of-2 table size given the lower-bound of required capacity. * \param cap The lower-bound of the required capacity * \param fib_shift The result shift for Fibonacci Hashing * \param n_slots The result number of slots */ static void CalcTableSize(uint64_t cap, uint32_t* fib_shift, uint64_t* n_slots) { uint32_t shift = 64; uint64_t slots = 1; for (uint64_t c = cap; c; c >>= 1) { shift -= 1; slots <<= 1; } TVM_FFI_ICHECK_GT(slots, cap); if (slots < cap * 2) { *fib_shift = shift - 1; *n_slots = slots << 1; } else { *fib_shift = shift; *n_slots = slots; } } /*! * \brief Fibonacci Hashing, maps a hash code to an index in a power-of-2-sized table. * See also: https://programmingpraxis.com/2018/06/19/fibonacci-hash/. * \param hash_value The raw hash value * \param fib_shift The shift in Fibonacci Hashing * \return An index calculated using Fibonacci Hashing */ static uint64_t FibHash(uint64_t hash_value, uint32_t fib_shift) { constexpr uint64_t coeff = 11400714819323198485ull; return (coeff * hash_value) >> fib_shift; } /*! \brief The implicit in-place linked list used to index a chain */ struct ListNode { /*! \brief Construct None */ ListNode() : index(0), block(nullptr) {} /*! \brief Construct from position */ ListNode(uint64_t index, const DenseMapBaseObj* self) : index(index), block(self->GetBlock(index / kBlockCap)) {} /*! \brief Metadata on the entry */ uint8_t& Meta() const { return *(block->bytes + index % kBlockCap); } /*! \brief Data on the entry */ ItemType& Item() const { return *(reinterpret_cast(block->bytes + kBlockCap + (index % kBlockCap) * sizeof(ItemType))); } /*! \brief Data on the entry */ KVType& Data() const { return Item().data; } /*! \brief Key on the entry */ key_type& Key() const { return Data().first; } /*! \brief Value on the entry */ mapped_type& Val() const { return Data().second; } /*! \brief If the entry is head of linked list */ bool IsHead() const { return (Meta() & 0b10000000) == 0b00000000; } /*! \brief If the entry is none */ bool IsNone() const { return block == nullptr; } /*! \brief If the entry is empty slot */ bool IsEmpty() const { return Meta() == kEmptySlot; } /*! \brief If the entry is protected slot */ bool IsProtected() const { return Meta() == kProtectedSlot; } /*! \brief Set the entry to be empty */ void SetEmpty() const { Meta() = kEmptySlot; } /*! \brief Destruct the item in the entry */ void DestructData() const { // explicit call destructor to destroy the item // Favor this over ~KVType as MSVC may not support ~KVType (need the original name) (&Data())->first.Any::~Any(); (&Data())->second.Any::~Any(); } /*! \brief Set the entry to be protected */ void SetProtected() const { Meta() = kProtectedSlot; } /*! \brief Set the entry's jump to its next entry */ void SetJump(uint8_t jump) const { (Meta() &= 0b10000000) |= jump; } /*! \brief Construct a head of linked list in-place */ void NewHead(ItemType v) const { Meta() = 0b00000000; new (&Item()) ItemType(std::move(v)); } /*! \brief Construct a tail of linked list in-place */ void NewTail(ItemType v) const { Meta() = 0b10000000; new (&Item()) ItemType(std::move(v)); } /*! \brief If the entry has next entry on the linked list */ bool HasNext() const { return NextProbeLocation(Meta() & 0b01111111) != 0; } /*! \brief Move the entry to the next entry on the linked list */ bool MoveToNext(const DenseMapBaseObj* self, uint8_t meta) { uint64_t offset = NextProbeLocation(meta & 0b01111111); if (offset == 0) { index = 0; block = nullptr; return false; } // the probing will go to next position and round back to stay within the // correct range of the slots index = (index + offset) % self->NumSlots(); block = self->GetBlock(index / kBlockCap); return true; } /*! \brief Move the entry to the next entry on the linked list */ bool MoveToNext(const DenseMapBaseObj* self) { return MoveToNext(self, Meta()); } /*! \brief Get the previous entry on the linked list */ ListNode FindPrev(const DenseMapBaseObj* self) const { // start from the head of the linked list, which must exist ListNode next = self->IndexFromHash(AnyHash()(Key())); // `prev` is always the previous item of `next` ListNode prev = next; for (next.MoveToNext(self); index != next.index; prev = next, next.MoveToNext(self)) { } return prev; } /*! \brief Get the next empty jump */ bool GetNextEmpty(const DenseMapBaseObj* self, uint8_t* jump, ListNode* result) const { for (uint8_t idx = 1; idx < kNumJumpDists; ++idx) { // the probing will go to next position and round back to stay within the // correct range of the slots ListNode candidate((index + NextProbeLocation(idx)) % self->NumSlots(), self); if (candidate.IsEmpty()) { *jump = idx; *result = candidate; return true; } } return false; } /*! \brief Index on the real array */ uint64_t index; /*! \brief Pointer to the actual block */ Block* block; }; protected: /*! \brief fib shift in Fibonacci Hashing */ uint32_t fib_shift_; /*! \brief the head of iterator list */ uint64_t iter_list_head_ = kInvalidIndex; /*! \brief the tail of iterator list */ uint64_t iter_list_tail_ = kInvalidIndex; static uint64_t NextProbeLocation(size_t index) { /* clang-format off */ /*! \brief Candidates of probing distance */ static const uint64_t kNextProbeLocation[kNumJumpDists] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // Quadratic probing with triangle numbers. See also: // 1) https://en.wikipedia.org/wiki/Quadratic_probing // 2) https://fgiesen.wordpress.com/2015/02/22/triangular-numbers-mod-2n/ // 3) https://github.com/skarupke/flat_hash_map 21, 28, 36, 45, 55, 66, 78, 91, 105, 120, 136, 153, 171, 190, 210, 231, 253, 276, 300, 325, 351, 378, 406, 435, 465, 496, 528, 561, 595, 630, 666, 703, 741, 780, 820, 861, 903, 946, 990, 1035, 1081, 1128, 1176, 1225, 1275, 1326, 1378, 1431, 1485, 1540, 1596, 1653, 1711, 1770, 1830, 1891, 1953, 2016, 2080, 2145, 2211, 2278, 2346, 2415, 2485, 2556, 2628, // larger triangle numbers 8515, 19110, 42778, 96141, 216153, 486591, 1092981, 2458653, 5532801, 12442566, 27993903, 62983476, 141717030, 318844378, 717352503, 1614057336, 3631522476, 8170957530, 18384510628, 41364789378, 93070452520, 209408356380, 471168559170, 1060128894105, 2385289465695, 5366898840628, 12075518705635, 27169915244790, 61132312065111, 137547689707000, 309482283181501, 696335127828753, 1566753995631385, 3525196511162271, 7931691992677701, 17846306936293605, 40154190677507445, 90346928918121501, 203280589587557251, 457381325854679626, 1029107982097042876, 2315492959180353330, 5209859154120846435, }; /* clang-format on */ return kNextProbeLocation[index]; } friend class MapBaseObj; friend class SmallMapBaseObj; private: /*! * \brief Set the number of slots and attach tags bit. * \param n The number of slots */ void SetSlotsAndDenseLayoutTag(uint64_t n) { TVM_FFI_ICHECK(((n & kSmallTagMask) == 0ull)) << "DenseMap expects MSB clear"; slots_ = n; } }; /*! \brief A specialization of small-sized hash map */ class SmallMapBaseObj : public MapBaseObj { private: static constexpr uint64_t kInitSize = 2; static constexpr uint64_t kMaxSize = 4; public: using MapBaseObj::iterator; using MapBaseObj::KVType; // Return the number of usable slots for Small layout (mask off tag). /*! * \brief Return the number of usable slots for Small layout (mask off tag). * \return The number of usable slots */ uint64_t NumSlots() const { return slots_ & ~kSmallTagMask; } ~SmallMapBaseObj() { // in destructor, need to check if the map was inplace switched to a dense map. MapBaseObj* base_map = static_cast(this); if (base_map->IsSmallMap()) { this->Reset(); } else { // this map was inplace switched to a dense map. DenseMapBaseObj* this_dense_map = static_cast(base_map); this_dense_map->Reset(); } } /*! * \brief clear all entries */ void clear() { KVType* begin = static_cast(data_); for (uint64_t index = 0; index < size_; ++index) { // call destructor to destroy the item in `begin + index` // Explicit call Any::~Any() to destroy the Any object // Favor this over ~KVType as MSVC may not support ~KVType (need the original name) (begin + index)->first.Any::~Any(); (begin + index)->second.Any::~Any(); } size_ = 0; } /*! * \brief Count the number of times a key exists in the SmallMapBaseObj * \param key The indexing key * \return The result, 0 or 1 */ size_t count(const key_type& key) const { return find(key).index < size_; } /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The const reference to the value */ const mapped_type& at(const key_type& key) const { iterator itr = find(key); if (itr.index >= size_) { TVM_FFI_THROW(KeyError) << "key is not in Map"; } return itr->second; } /*! * \brief Index value associated with a key, throw exception if the key does not exist * \param key The indexing key * \return The mutable reference to the value */ mapped_type& at(const key_type& key) { iterator itr = find(key); if (itr.index >= size_) { TVM_FFI_THROW(KeyError) << "key is not in Map"; } return itr->second; } /*! \return begin iterator */ iterator begin() const { return iterator(0, this); } /*! \return end iterator */ iterator end() const { return iterator(size_, this); } /*! * \brief Index value associated with a key * \param key The indexing key * \return The iterator of the entry associated with the key, end iterator if not exists */ iterator find(const key_type& key) const { KVType* ptr = static_cast(data_); for (uint64_t i = 0; i < size_; ++i, ++ptr) { if (AnyEqual()(ptr->first, key)) { return iterator(i, this); } } return iterator(size_, this); } /*! * \brief Erase the entry associated with the iterator * \param position The iterator */ void erase(const iterator& position) { Erase(position.index); } private: /*! * \brief Set the number of slots and attach tags bit. * \param n The number of slots */ void SetSlotsAndSmallLayoutTag(uint64_t n) { slots_ = (n & ~kSmallTagMask) | kSmallTagMask; } /* * \brief Reset original small Storage so it can be ready for switch. */ void Reset() { this->clear(); if (data_deleter_ != nullptr) { data_deleter_(data_); // after deletion we have zero slots slots_ = 0; data_ = nullptr; data_deleter_ = nullptr; } } /*! * \brief Inplace deleter from data * \param data The data */ static void InplaceSmallMapDeleterFromData(void* data) { details::ObjectUnsafe::ObjectPtrFromOwned( reinterpret_cast(reinterpret_cast(data) - sizeof(SmallMapBaseObj))) .reset(); } /*! * \brief Inplace switch the storage contents to the other map * * The other map will be consumed after this operation * The current content will be reset after this operation * * \param other The other map */ void InplaceSwitchTo(ObjectPtr&& other) { // invariant this map have not been inplace switched to a dense map TVM_FFI_ICHECK(this->IsSmallMap()); this->Reset(); MapBaseObj* other_map = static_cast(other.get()); if (other_map->IsSmallMap()) { SmallMapBaseObj* other_small_map = static_cast(other_map); SmallMapBaseObj* this_small_map = this; this_small_map->size_ = other_small_map->size_; this_small_map->slots_ = other_small_map->slots_; this_small_map->data_ = other_small_map->data_; if (other_small_map->data_deleter_ != nullptr) { this_small_map->data_deleter_ = other_small_map->data_deleter_; other_small_map->data_deleter_ = nullptr; } else { // we switch to the inplace deleter from the data this_small_map->data_deleter_ = InplaceSmallMapDeleterFromData; // move out other from the ptr so the deletion can only be triggered // via InplaceSmallMapDeleterFromData details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(other)); } other_small_map->data_ = nullptr; other_small_map->size_ = 0; other_small_map->slots_ = 0; } else { // Reinterpret this SmallMapBaseObj's memory as DenseMapBaseObj to write fields at the // correct offsets. This is raw memory manipulation: SmallMapBaseObj's allocation is // guaranteed large enough (see static_assert in Empty()), and all member access // compiles to fixed-offset stores with no virtual dispatch involved. // The destructor will also cross check and apply the correct deletion. // As a result, we can inplace switch the container storage to dense map DenseMapBaseObj* other_dense_map = static_cast(other_map); DenseMapBaseObj* this_dense_map = reinterpret_cast(this); this_dense_map->size_ = other_dense_map->size_; this_dense_map->slots_ = other_dense_map->slots_; this_dense_map->data_ = other_dense_map->data_; this_dense_map->data_deleter_ = other_dense_map->data_deleter_; this_dense_map->fib_shift_ = other_dense_map->fib_shift_; this_dense_map->iter_list_head_ = other_dense_map->iter_list_head_; this_dense_map->iter_list_tail_ = other_dense_map->iter_list_tail_; other_dense_map->data_ = nullptr; other_dense_map->data_deleter_ = nullptr; other_dense_map->size_ = 0; other_dense_map->slots_ = 0; } } /*! * \brief Remove a position in SmallMapBaseObj * \param index The position to be removed */ void Erase(const uint64_t index) { if (index >= size_) { return; } KVType* begin = static_cast(data_); // call destructor to destroy the item in `begin + index` // Explicit call Any::~Any() to destroy the Any object // Favor this over ~KVType as MSVC may not support ~KVType (need the original name) (begin + index)->first.Any::~Any(); (begin + index)->second.Any::~Any(); // IMPORTANT: We do direct raw memmove to bring later items to the current position // to preserve the order of insertion. // This works because direct memory copy preserves the Any's move semantics. if (index + 1 < size_) { std::memmove(reinterpret_cast(begin + index), reinterpret_cast(begin + index + 1), (size_ - index - 1) * sizeof(KVType)); } size_ -= 1; } /*! * \brief Create an empty container * \param n Number of empty slots * \return The object created * \tparam MapObjType The type of map object */ template static ObjectPtr Empty(uint64_t n = kInitSize) { // We always allocate a SmallMapObj to be large enough so it can inplace switch to DenseMapObj static_assert(alignof(SmallMapBaseObj) % alignof(KVType) == 0); static_assert(sizeof(SmallMapBaseObj) + kInitSize * sizeof(KVType) >= sizeof(DenseMapBaseObj)); n = std::max(n, static_cast(kInitSize)); // allocate a SmallMapBaseObj with enough space to inplace store n elements and switch to // DenseMapBaseObj ObjectPtr p = ffi::make_inplace_array_object(n); // data_ is after the SmallMapBaseObj header p->data_ = reinterpret_cast(p.get()) + sizeof(SmallMapBaseObj); p->size_ = 0; p->SetSlotsAndSmallLayoutTag(n); // manually set the type index to specialized subclass // data layout is ensured to be the same except for the type index static_assert(std::is_base_of_v); static_assert(sizeof(MapObjType) == sizeof(MapBaseObj)); p->header_.type_index = MapObjType::RuntimeTypeIndex(); return p; } /*! * \brief Create an empty container initialized with a given range * \param n Number of empty slots * \param first begin of iterator * \param last end of iterator * \tparam MapObjType The type of map object * \tparam IterType The type of iterator * \return The object created */ template static ObjectPtr CreateFromRange(uint64_t n, IterType first, IterType last) { ObjectPtr p = Empty(n); SmallMapBaseObj* small_map_base_obj = static_cast(p.get()); KVType* ptr = static_cast(small_map_base_obj->data_); for (; first != last; ++first, ++small_map_base_obj->size_) { new (ptr++) KVType(*first); } return p; } /*! * \brief Create an empty container with elements copying from another SmallMapBaseObj * \param from The source container * \return The object created */ template static ObjectPtr CopyFrom(SmallMapBaseObj* from) { KVType* first = static_cast(from->data_); KVType* last = first + from->size_; return CreateFromRange(from->size_, first, last); } /*! * \brief InsertMaybeReHash an entry into the given hash map * \param kv The entry to be inserted * \param map The pointer to the map, can be changed if re-hashing happens */ template static ObjectPtr InsertMaybeReHash(KVType&& kv, const ObjectPtr& map) { SmallMapBaseObj* map_node = static_cast(map.get()); iterator itr = map_node->find(kv.first); if (itr.index < map_node->size_) { itr->second = kv.second; return ObjectPtr(nullptr); } if (map_node->size_ < map_node->NumSlots()) { KVType* ptr = static_cast(map_node->data_) + map_node->size_; new (ptr) KVType(std::move(kv)); ++map_node->size_; return ObjectPtr(nullptr); } uint64_t next_size = std::max(map_node->NumSlots() * 2, kInitSize); next_size = std::min(next_size, kMaxSize); TVM_FFI_ICHECK_GT(next_size, map_node->NumSlots()); ObjectPtr new_map = CreateFromRange(next_size, map_node->begin(), map_node->end()); ObjectPtr rehashed = InsertMaybeReHash(std::move(kv), new_map); if (rehashed != nullptr) new_map = std::move(rehashed); return new_map; } /*! * \brief Increment the pointer * \param index The pointer to be incremented * \return The increased pointer */ uint64_t IncItr(uint64_t index) const { return index + 1 < size_ ? index + 1 : size_; } /*! * \brief Decrement the pointer * \param index The pointer to be decremented * \return The decreased pointer */ uint64_t DecItr(uint64_t index) const { return index > 0 ? index - 1 : size_; } /*! * \brief De-reference the pointer * \param index The pointer to be dereferenced * \return The result */ KVType* DeRefItr(uint64_t index) const { return static_cast(data_) + index; } protected: friend class MapBaseObj; friend class DenseMapBaseObj; }; /// \cond #define TVM_FFI_DISPATCH_MAP(base, var, body) \ { \ using TSmall = SmallMapBaseObj*; \ using TDense = DenseMapBaseObj*; \ if ((base)->IsSmallMap()) { \ TSmall var = static_cast((base)); \ body; \ } else { \ TDense var = static_cast((base)); \ body; \ } \ } #define TVM_FFI_DISPATCH_MAP_CONST(base, var, body) \ { \ using TSmall = const SmallMapBaseObj*; \ using TDense = const DenseMapBaseObj*; \ if ((base)->IsSmallMap()) { \ TSmall var = static_cast((base)); \ body; \ } else { \ TDense var = static_cast((base)); \ body; \ } \ } inline MapBaseObj::iterator::pointer MapBaseObj::iterator::operator->() const { TVM_FFI_MAP_FAIL_IF_CHANGED() TVM_FFI_DISPATCH_MAP_CONST(self, p, { return p->DeRefItr(index); }); } inline MapBaseObj::iterator& MapBaseObj::iterator::operator++() { TVM_FFI_MAP_FAIL_IF_CHANGED() TVM_FFI_DISPATCH_MAP_CONST(self, p, { index = p->IncItr(index); return *this; }); } inline MapBaseObj::iterator& MapBaseObj::iterator::operator--() { TVM_FFI_MAP_FAIL_IF_CHANGED() TVM_FFI_DISPATCH_MAP_CONST(self, p, { index = p->DecItr(index); return *this; }); } inline size_t MapBaseObj::count(const key_type& key) const { TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->count(key); }); } inline const MapBaseObj::mapped_type& MapBaseObj::at(const MapBaseObj::key_type& key) const { TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->at(key); }); } inline MapBaseObj::mapped_type& MapBaseObj::at(const MapBaseObj::key_type& key) { TVM_FFI_DISPATCH_MAP(this, p, { return p->at(key); }); } inline MapBaseObj::iterator MapBaseObj::begin() const { TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->begin(); }); } inline MapBaseObj::iterator MapBaseObj::end() const { TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->end(); }); } inline MapBaseObj::iterator MapBaseObj::find(const MapBaseObj::key_type& key) const { TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->find(key); }); } inline void MapBaseObj::erase(const MapBaseObj::iterator& position) { TVM_FFI_DISPATCH_MAP(this, p, { p->erase(position); }); } inline void MapBaseObj::clear() { TVM_FFI_DISPATCH_MAP(this, p, { p->clear(); }); } inline void MapBaseObj::InplaceSwitchTo(ObjectPtr&& other) { TVM_FFI_DISPATCH_MAP(this, p, { p->InplaceSwitchTo(std::move(other)); }); } /// \endcond #undef TVM_FFI_DISPATCH_MAP #undef TVM_FFI_DISPATCH_MAP_CONST template inline ObjectPtr MapBaseObj::Empty() { return SmallMapBaseObj::Empty(); } template inline ObjectPtr MapBaseObj::CopyFrom(MapBaseObj* from) { if (from->IsSmallMap()) { return SmallMapBaseObj::CopyFrom(static_cast(from)); } else { return DenseMapBaseObj::CopyFrom(static_cast(from)); } } template inline ObjectPtr MapBaseObj::CreateFromRange(IterType first, IterType last) { int64_t _cap = std::distance(first, last); if (_cap < 0) { return SmallMapBaseObj::Empty(); } uint64_t cap = static_cast(_cap); if (cap < SmallMapBaseObj::kMaxSize) { if (cap < 2) { return SmallMapBaseObj::CreateFromRange(cap, first, last); } // need to insert to avoid duplicate keys ObjectPtr obj = SmallMapBaseObj::Empty(cap); for (; first != last; ++first) { KVType kv(*first); ObjectPtr rehashed = SmallMapBaseObj::InsertMaybeReHash(std::move(kv), obj); if (rehashed != nullptr) obj = std::move(rehashed); } return obj; } else { uint32_t fib_shift; uint64_t n_slots; DenseMapBaseObj::CalcTableSize(cap, &fib_shift, &n_slots); ObjectPtr obj = DenseMapBaseObj::Empty(fib_shift, n_slots); for (; first != last; ++first) { KVType kv(*first); ObjectPtr rehashed = DenseMapBaseObj::InsertMaybeReHash(std::move(kv), obj); if (rehashed != nullptr) obj = std::move(rehashed); } return obj; } } template inline ObjectPtr MapBaseObj::InsertMaybeReHash(KVType&& kv, const ObjectPtr& map) { MapBaseObj* base = static_cast(map.get()); #if TVM_FFI_DEBUG_WITH_ABI_CHANGE base->state_marker++; #endif // TVM_FFI_DEBUG_WITH_ABI_CHANGE if (base->IsSmallMap()) { SmallMapBaseObj* sm = static_cast(base); if (sm->NumSlots() < SmallMapBaseObj::kMaxSize) { return SmallMapBaseObj::InsertMaybeReHash(std::move(kv), map); } else if (sm->NumSlots() == SmallMapBaseObj::kMaxSize) { if (base->size_ < sm->NumSlots()) { return SmallMapBaseObj::InsertMaybeReHash(std::move(kv), map); } else { ObjectPtr new_map = MapBaseObj::CreateFromRange(base->begin(), base->end()); ObjectPtr rehashed = DenseMapBaseObj::InsertMaybeReHash(std::move(kv), new_map); if (rehashed != nullptr) new_map = std::move(rehashed); return new_map; } } } else { return DenseMapBaseObj::InsertMaybeReHash(std::move(kv), map); } return ObjectPtr(nullptr); } /// \cond Doxygen_Suppress /*! * \brief Specialize make_object to be deleted for make_object and * make_object only. */ template <> inline ObjectPtr make_object<>() = delete; /// \endcond /*! * \brief CRTP base for map type-traits (Map, Dict). * * \tparam Derived Must expose: * - `static constexpr int32_t kPrimaryTypeIndex` — the canonical FFI type index * - `static constexpr int32_t kOtherTypeIndex` — an alternative accepted type index * - `static constexpr const char* kTypeName` — human-readable name for diagnostics */ template struct MapTypeTraitsBase : public ObjectRefTypeTraitsBase { using Base = ObjectRefTypeTraitsBase; using Base::CopyFromAnyViewAfterCheck; TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex) return false; if constexpr (std::is_same_v && std::is_same_v) { return true; } else { const MapBaseObj* n = reinterpret_cast(src->v_obj); for (const auto& kv : *n) { if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.first)) return false; } if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.second)) return false; } } return true; } } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex && src->type_index != Derived::kOtherTypeIndex) { return TypeTraitsBase::GetMismatchTypeInfo(src); } if constexpr (!std::is_same_v || !std::is_same_v) { const MapBaseObj* n = reinterpret_cast(src->v_obj); for (const auto& kv : *n) { if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.first) && !kv.first.try_cast().has_value()) { return std::string(Derived::kTypeName) + "[some key is " + details::AnyUnsafe::GetMismatchTypeInfo(kv.first) + ", V]"; } } if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.second) && !kv.second.try_cast().has_value()) { return std::string(Derived::kTypeName) + "[K, some value is " + details::AnyUnsafe::GetMismatchTypeInfo(kv.second) + "]"; } } } } TVM_FFI_THROW(InternalError) << "Cannot reach here"; TVM_FFI_UNREACHABLE(); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex && src->type_index != Derived::kOtherTypeIndex) { return std::nullopt; } const MapBaseObj* n = reinterpret_cast(src->v_obj); if constexpr (!std::is_same_v || !std::is_same_v) { bool storage_check = [&]() { for (const auto& kv : *n) { if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.first)) return false; } if constexpr (!std::is_same_v) { if (!details::AnyUnsafe::CheckAnyStrict(kv.second)) return false; } } return true; }(); // fast path: if storage check passes and type is primary, return directly. if (storage_check && src->type_index == Derived::kPrimaryTypeIndex) { return CopyFromAnyViewAfterCheck(src); } // slow path: create a new map and convert each key-value pair. MapRef ret; for (const auto& kv : *n) { auto k = kv.first.try_cast(); auto v = kv.second.try_cast(); if (!k.has_value() || !v.has_value()) return std::nullopt; ret.Set(*std::move(k), *std::move(v)); } return ret; } else { if (src->type_index == Derived::kPrimaryTypeIndex) { return CopyFromAnyViewAfterCheck(src); } // cross-type conversion for Any,Any: create new MapRef, copy all entries. MapRef ret; for (const auto& kv : *n) { ret.Set(kv.first, kv.second); } return ret; } } TVM_FFI_INLINE static std::string TypeStr() { return std::string(Derived::kTypeName) + "<" + details::Type2Str::v() + ", " + details::Type2Str::v() + ">"; } private: MapTypeTraitsBase() = default; friend Derived; }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_MAP_BASE_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/seq_base.h000066400000000000000000000261531521067262500220720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/seq_base.h * \brief Base class for sequence containers (Array, List). */ #ifndef TVM_FFI_CONTAINER_SEQ_BASE_H_ #define TVM_FFI_CONTAINER_SEQ_BASE_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Base class for sequence containers (ArrayObj, ListObj). * * SeqBaseObj is transparent to the FFI type system (no type index), * following the same pattern as BytesObjBase. */ class SeqBaseObj : public Object, protected TVMFFISeqCell { public: SeqBaseObj() { data = nullptr; TVMFFISeqCell::size = 0; TVMFFISeqCell::capacity = 0; data_deleter = nullptr; } ~SeqBaseObj() { Any* begin = MutableBegin(); for (int64_t i = 0; i < TVMFFISeqCell::size; ++i) { (begin + i)->Any::~Any(); } if (data_deleter != nullptr) { data_deleter(data); } } /*! \return The size of the sequence */ size_t size() const { return static_cast(TVMFFISeqCell::size); } /*! \return The capacity of the sequence */ size_t capacity() const { return static_cast(TVMFFISeqCell::capacity); } /*! \return Whether the sequence is empty */ bool empty() const { return TVMFFISeqCell::size == 0; } /*! * \brief Read i-th element from the sequence. * \param i The index * \return the i-th element. */ const Any& at(int64_t i) const { return this->operator[](i); } /*! * \brief Read i-th element from the sequence. * \param i The index * \return the i-th element. */ const Any& operator[](int64_t i) const { if (i < 0 || i >= TVMFFISeqCell::size) { TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << TVMFFISeqCell::size; } return static_cast(data)[i]; } /*! \return The first element */ const Any& front() const { if (TVMFFISeqCell::size == 0) { TVM_FFI_THROW(IndexError) << "front() on empty sequence"; } return static_cast(data)[0]; } /*! \return The last element */ const Any& back() const { if (TVMFFISeqCell::size == 0) { TVM_FFI_THROW(IndexError) << "back() on empty sequence"; } return static_cast(data)[TVMFFISeqCell::size - 1]; } /*! \return begin constant iterator */ const Any* begin() const { return static_cast(data); } /*! \return end constant iterator */ const Any* end() const { return begin() + TVMFFISeqCell::size; } /*! \brief Release reference to all the elements */ void clear() { Any* itr = MutableEnd(); while (TVMFFISeqCell::size > 0) { (--itr)->Any::~Any(); --TVMFFISeqCell::size; } } /*! * \brief Set i-th element of the sequence in-place * \param i The index * \param item The value to be set */ void SetItem(int64_t i, Any item) { if (i < 0 || i >= TVMFFISeqCell::size) { TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << TVMFFISeqCell::size; } static_cast(data)[i] = std::move(item); } /*! \brief Remove the last element */ void pop_back() { if (TVMFFISeqCell::size == 0) { TVM_FFI_THROW(IndexError) << "pop_back on empty sequence"; } ShrinkBy(1); } /*! * \brief Erase element at position idx * \param idx The index to erase */ void erase(int64_t idx) { if (idx < 0 || idx >= TVMFFISeqCell::size) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds " << TVMFFISeqCell::size; } MoveElementsLeft(idx, idx + 1, TVMFFISeqCell::size); ShrinkBy(1); } /*! * \brief Erase elements in half-open range [first, last) * \param first Start index (inclusive) * \param last End index (exclusive) */ void erase(int64_t first, int64_t last) { if (first == last) return; if (first < 0 || last > TVMFFISeqCell::size || first >= last) { TVM_FFI_THROW(IndexError) << "Erase range [" << first << ", " << last << ") out of bounds " << TVMFFISeqCell::size; } MoveElementsLeft(first, last, TVMFFISeqCell::size); ShrinkBy(last - first); } /*! * \brief Insert element at position idx * \param idx The index to insert at * \param item The value to insert * \note Caller must ensure capacity >= size + 1 */ void insert(int64_t idx, Any item) { int64_t sz = TVMFFISeqCell::size; if (idx < 0 || idx > sz) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds [0, " << sz << "]"; } EnlargeBy(1); MoveElementsRight(idx + 1, idx, sz); MutableBegin()[idx] = std::move(item); } /*! * \brief Insert elements from iterator range at position idx * \param idx The index to insert at * \param first Begin of iterator * \param last End of iterator * \tparam IterType The type of iterator * \note Caller must ensure capacity >= size + distance(first, last) */ template void insert(int64_t idx, IterType first, IterType last) { int64_t count = std::distance(first, last); if (count == 0) return; int64_t sz = TVMFFISeqCell::size; if (idx < 0 || idx > sz) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds [0, " << sz << "]"; } EnlargeBy(count); MoveElementsRight(idx + count, idx, sz); Any* dst = MutableBegin() + idx; for (; first != last; ++first, ++dst) { *dst = Any(*first); } } /*! \brief Reverse the elements in-place */ void Reverse() { std::reverse(MutableBegin(), MutableBegin() + TVMFFISeqCell::size); } /*! * \brief Resize the sequence * \param n The new size * \note Caller must ensure capacity >= n when growing */ void resize(int64_t n) { if (n < 0) { TVM_FFI_THROW(ValueError) << "Cannot resize to negative size"; } int64_t old_size = TVMFFISeqCell::size; if (old_size < n) { EnlargeBy(n - old_size); } else if (old_size > n) { ShrinkBy(old_size - n); } } protected: /// \cond Doxygen_Suppress Any* MutableBegin() const { return static_cast(this->data); } Any* MutableEnd() const { return MutableBegin() + TVMFFISeqCell::size; } template void EmplaceInit(size_t idx, Args&&... args) { Any* itr = MutableBegin() + idx; new (itr) Any(std::forward(args)...); } void EnlargeBy(int64_t delta, const Any& val = Any()) { Any* itr = MutableEnd(); while (delta-- > 0) { new (itr++) Any(val); ++TVMFFISeqCell::size; } } void ShrinkBy(int64_t delta) { Any* itr = MutableEnd(); while (delta-- > 0) { (--itr)->Any::~Any(); --TVMFFISeqCell::size; } } void MoveElementsLeft(int64_t dst, int64_t src_begin, int64_t src_end) { Any* begin = MutableBegin(); std::move(begin + src_begin, begin + src_end, begin + dst); } void MoveElementsRight(int64_t dst, int64_t src_begin, int64_t src_end) { Any* begin = MutableBegin(); std::move_backward(begin + src_begin, begin + src_end, begin + dst + (src_end - src_begin)); } /// \endcond }; /*! * \brief CRTP base for sequence type-traits (Array, List). * * \tparam Derived Must expose: * - `static constexpr int32_t kPrimaryTypeIndex` — the canonical FFI type index * - `static constexpr int32_t kOtherTypeIndex` — an alternative accepted type index * - `static constexpr const char* kTypeName` — human-readable name for diagnostics */ template struct SeqTypeTraitsBase : public ObjectRefTypeTraitsBase { using Base = ObjectRefTypeTraitsBase; using Base::CopyFromAnyViewAfterCheck; TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex) return false; if constexpr (std::is_same_v) { return true; } else { const SeqBaseObj* n = reinterpret_cast(src->v_obj); for (const Any& any_v : *n) { if (!details::AnyUnsafe::CheckAnyStrict(any_v)) return false; } return true; } } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex && src->type_index != Derived::kOtherTypeIndex) { return TypeTraitsBase::GetMismatchTypeInfo(src); } if constexpr (!std::is_same_v) { const SeqBaseObj* n = reinterpret_cast(src->v_obj); for (size_t i = 0; i < n->size(); i++) { const Any& any_v = n->at(static_cast(i)); if (details::AnyUnsafe::CheckAnyStrict(any_v)) continue; if (any_v.try_cast()) continue; return std::string(Derived::kTypeName) + "[index " + std::to_string(i) + ": " + details::AnyUnsafe::GetMismatchTypeInfo(any_v) + "]"; } } TVM_FFI_THROW(InternalError) << "Cannot reach here"; TVM_FFI_UNREACHABLE(); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index != Derived::kPrimaryTypeIndex && src->type_index != Derived::kOtherTypeIndex) { return std::nullopt; } const SeqBaseObj* n = reinterpret_cast(src->v_obj); if constexpr (!std::is_same_v) { bool storage_check = [&]() { for (const Any& any_v : *n) { if (!details::AnyUnsafe::CheckAnyStrict(any_v)) return false; } return true; }(); if (storage_check && src->type_index == Derived::kPrimaryTypeIndex) { return CopyFromAnyViewAfterCheck(src); } SeqRef result; result.reserve(static_cast(n->size())); for (const Any& any_v : *n) { if (auto opt_v = any_v.try_cast()) { result.push_back(*std::move(opt_v)); } else { return std::nullopt; } } return result; } else { if (src->type_index == Derived::kPrimaryTypeIndex) { return CopyFromAnyViewAfterCheck(src); } SeqRef result; result.reserve(static_cast(n->size())); for (const Any& any_v : *n) { result.push_back(any_v); } return result; } } TVM_FFI_INLINE static std::string TypeStr() { return std::string(Derived::kTypeName) + "<" + details::Type2Str::v() + ">"; } private: SeqTypeTraitsBase() = default; friend Derived; }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_SEQ_BASE_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/shape.h000066400000000000000000000237211521067262500214060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/shape.h * \brief Container to store shape of an Tensor. */ #ifndef TVM_FFI_CONTAINER_SHAPE_H_ #define TVM_FFI_CONTAINER_SHAPE_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Lightweight view non-owning class for shape. */ class ShapeView { public: /*! \brief Default constructor. */ ShapeView() : cell_{nullptr, 0} {} /*! \brief Copy constructor. */ ShapeView(const ShapeView& other) = default; /*! \brief Copy assignment operator. */ ShapeView& operator=(const ShapeView& other) = default; /*! \brief Move constructor. */ ShapeView(ShapeView&& other) = default; /*! \brief Move assignment operator. */ ShapeView& operator=(ShapeView&& other) = default; /*! \brief Constructor from data and size. */ ShapeView(const int64_t* data, size_t size) : cell_{data, size} {} /*! \brief Constructor from initializer list. */ ShapeView(const std::initializer_list& other) : cell_{other.begin(), other.size()} {} /*! \brief Get the data pointer. */ const int64_t* data() const { return cell_.data; } /*! \brief Get the size of the shape. */ size_t size() const { return cell_.size; } /*! \brief Get the product of the shape. */ int64_t Product() const { int64_t product = 1; for (size_t i = 0; i < cell_.size; ++i) { product *= cell_.data[i]; } return product; } /*! \brief Get the i-th element of the shape. */ int64_t operator[](size_t idx) const { return cell_.data[idx]; } /*! \return begin iterator */ const int64_t* begin() const { return cell_.data; } /*! \return end iterator */ const int64_t* end() const { return cell_.data + cell_.size; } /*! \return Whether shape tuple is empty */ bool empty() const { return size() == 0; } /*! \return The first element of the shape tuple */ int64_t front() const { return this->at(0); } /*! \return The last element of the shape tuple */ int64_t back() const { return this->at(this->size() - 1); } /*! \brief Get the i-th element of the shape. */ int64_t at(size_t idx) const { if (idx >= this->size()) { TVM_FFI_THROW(IndexError) << "indexing " << idx << " on a Shape of size " << this->size(); } return cell_.data[idx]; } private: TVMFFIShapeCell cell_; }; /*! \brief An object representing a shape tuple. */ class ShapeObj : public Object, public TVMFFIShapeCell { public: /*! \brief The type of shape index element. */ using index_type = int64_t; /*! \brief Get "numel", meaning the number of elements of an array if the array has this shape */ int64_t Product() const { int64_t product = 1; for (size_t i = 0; i < this->size; ++i) { product *= this->data[i]; } return product; } /// \cond Doxygen_Suppress static constexpr const uint32_t _type_index = TypeIndex::kTVMFFIShape; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIShape, ShapeObj, Object); /// \endcond }; namespace details { class ShapeObjStdImpl : public ShapeObj { public: explicit ShapeObjStdImpl(std::vector other) : data_{std::move(other)} { this->data = data_.data(); this->size = static_cast(data_.size()); } private: std::vector data_; }; TVM_FFI_INLINE ObjectPtr MakeEmptyShape(size_t length, int64_t** mutable_data) { ObjectPtr p = make_inplace_array_object(length); static_assert(alignof(ShapeObj) % alignof(int64_t) == 0); static_assert(sizeof(ShapeObj) % alignof(int64_t) == 0); int64_t* data = reinterpret_cast(reinterpret_cast(p.get()) + sizeof(ShapeObj)); if (mutable_data) { *mutable_data = data; } p->data = data; p->size = length; return p; } // inplace shape allocation template TVM_FFI_INLINE ObjectPtr MakeInplaceShape(IterType begin, IterType end) { size_t length = std::distance(begin, end); int64_t* mutable_data; ObjectPtr p = MakeEmptyShape(length, &mutable_data); std::copy(begin, end, mutable_data); return p; } /*! * \brief Get the product of a shape. * \param shape The input shape. * \param out_strides The output strides. * \return The product of the shape. */ TVM_FFI_INLINE void FillStridesFromShape(ShapeView shape, int64_t* out_strides) { int64_t stride = 1; for (int64_t i = static_cast(shape.size()) - 1; i >= 0; --i) { out_strides[i] = stride; stride *= shape[i]; } } /*! * \brief Make a strides shape from a shape view. * \param shape The input shape. * \return The shape. */ TVM_FFI_INLINE ObjectPtr MakeStridesFromShape(ShapeView shape) { int64_t* strides_data; ObjectPtr strides = details::MakeEmptyShape(shape.size(), &strides_data); FillStridesFromShape(shape, strides_data); return strides; } } // namespace details /*! * \brief Managed reference to shape object. * * When possible, use ShapeView instead of Shape to reduce memory allocation. * Use Shape when you need to have a managed reference to on-heap allocated shape. * * \sa ShapeView */ class Shape : public ObjectRef { public: /*! \brief The type of shape index element. */ using index_type = ShapeObj::index_type; /*! \brief Default constructor */ Shape() : ObjectRef(details::MakeEmptyShape(0, nullptr)) {} /*! * \brief Constructor from iterator * \param begin begin of iterator * \param end end of iterator * \tparam IterType The type of iterator */ template Shape(IterType begin, IterType end) : Shape(details::MakeInplaceShape(begin, end)) {} /** * \brief Constructor from Array * \param shape The Array * * \note This constructor will copy the data content. */ Shape(Array shape) // NOLINT(*) : Shape(shape.begin(), shape.end()) {} /*! * \brief constructor from initializer list * \param shape The initializer list */ Shape(std::initializer_list shape) : Shape(shape.begin(), shape.end()) {} /*! * \brief constructor from int64_t [N] * * \param other a int64_t array. */ Shape(std::vector other) // NOLINT(*) : ObjectRef(make_object(std::move(other))) {} /*! * \brief constructor from shape view. * \param other The shape view. */ Shape(ShapeView other) : Shape(other.begin(), other.end()) {} // NOLINT(*) /*! * \brief Create a strides from a shape. * \param shape The input shape. * \return The strides. */ static Shape StridesFromShape(ShapeView shape) { return Shape(details::MakeStridesFromShape(shape)); } /*! * \brief Convert to shape view. * \return The shape view. */ operator ShapeView() const { return ShapeView(data(), size()); } // NOLINT(*) /*! * \brief Return the data pointer * * \return const index_type* data pointer */ const int64_t* data() const { return get()->data; } /*! * \brief Return the size of the shape tuple * * \return size_t shape tuple size */ size_t size() const { return get()->size; } /*! * \brief Immutably read i-th element from the shape tuple. * \param idx The index * \return the i-th element. */ int64_t operator[](size_t idx) const { return this->data()[idx]; } /*! * \brief Immutably read i-th element from the shape tuple. * \param idx The index * \return the i-th element. */ int64_t at(size_t idx) const { if (idx >= this->size()) { TVM_FFI_THROW(IndexError) << "indexing " << idx << " on a Shape of size " << this->size(); } return this->operator[](idx); } /*! \return Whether shape tuple is empty */ bool empty() const { return size() == 0; } /*! \return The first element of the shape tuple */ int64_t front() const { return this->at(0); } /*! \return The last element of the shape tuple */ int64_t back() const { return this->at(this->size() - 1); } /*! \return begin iterator */ const int64_t* begin() const { return get()->data; } /*! \return end iterator */ const int64_t* end() const { return (get()->data + size()); } /*! \return The product of the shape tuple */ int64_t Product() const { return get()->Product(); } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Shape, ObjectRef, ShapeObj); /// \endcond private: explicit Shape(ObjectPtr ptr) : ObjectRef(std::move(ptr)) {} }; inline std::ostream& operator<<(std::ostream& os, const Shape& shape) { os << '['; for (size_t i = 0; i < shape.size(); ++i) { if (i != 0) { os << ", "; } os << shape[i]; } os << ']'; return os; } // Shape template <> inline constexpr bool use_default_type_traits_v = false; // Allow auto conversion from Array to Shape, but not from Shape to Array template <> struct TypeTraits : public ObjectRefWithFallbackTraitsBase> { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIShape; TVM_FFI_INLINE static Shape ConvertFallbackValue(Array src) { return Shape(std::move(src)); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_SHAPE_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/tensor.h000066400000000000000000001043731521067262500216230ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/tensor.h * \brief Container to store a Tensor. */ #ifndef TVM_FFI_CONTAINER_TENSOR_H_ #define TVM_FFI_CONTAINER_TENSOR_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { class Tensor; /*! * \brief Check if the device uses direct address, where address of data indicate alignment. * \param device The input device. * \return True if the device uses direct address, false otherwise. */ inline bool IsDirectAddressDevice(const DLDevice& device) { return device.device_type <= kDLCUDAHost || device.device_type == kDLCUDAManaged || device.device_type == kDLROCM || device.device_type == kDLROCMHost; } /*! * \brief check if a DLTensor is contiguous. * \param arr The input DLTensor. * \return The check result. */ inline bool IsContiguous(const DLTensor& arr) { if (arr.strides == nullptr) return true; // An empty tensor (numel == 0) is trivially contiguous regardless of strides, // matching NumPy/PyTorch semantics. for (int32_t i = 0; i < arr.ndim; ++i) { if (arr.shape[i] == 0) return true; } int64_t expected_stride = 1; for (int32_t i = arr.ndim; i != 0; --i) { int32_t k = i - 1; if (arr.shape[k] == 1) { // Skip stride check if shape[k] is 1, where the dimension is contiguous // regardless of the value of stride. // // For example, PyTorch will normalize stride to 1 if shape is 1 when exporting // to DLPack. // More context: https://github.com/pytorch/pytorch/pull/83158 continue; } if (arr.strides[k] != expected_stride) return false; expected_stride *= arr.shape[k]; } return true; } /** * \brief Check if the data in the DLTensor is aligned to the given alignment. * \param arr The input DLTensor. * \param alignment The alignment to check. * \return True if the data is aligned to the given alignment, false otherwise. */ inline bool IsAligned(const DLTensor& arr, size_t alignment) { if (IsDirectAddressDevice(arr.device)) { return (reinterpret_cast(static_cast(arr.data) + arr.byte_offset) % alignment == 0); } else { return arr.byte_offset % alignment == 0; } } /*! * \brief return the total number of bytes needed to store packed data * * \param numel the number of elements in the array * \param dtype the data type of the array * \return the total number of bytes needed to store packed data */ inline size_t GetDataSize(size_t numel, DLDataType dtype) { // compatible handling sub-byte uint1(bool), which usually stored as uint8_t // TODO(tqchen): revisit and switch to kDLBool if (dtype.code == kDLUInt && dtype.bits == 1 && dtype.lanes == 1) { return numel; } // for other sub-byte types, packing is preferred // Use uint64_t to avoid overflow on 32-bit platforms (WASM) for large allocations. return static_cast((static_cast(numel) * dtype.bits * dtype.lanes + 7) / 8); } /*! * \brief return the size of data the DLTensor holds, in terms of number of bytes * * \param arr the input DLTensor * \return number of bytes of data in the DLTensor. */ inline size_t GetDataSize(const DLTensor& arr) { size_t size = 1; for (int i = 0; i < arr.ndim; ++i) { size *= static_cast(arr.shape[i]); } return GetDataSize(size, arr.dtype); } /*! \brief An object representing a Tensor. */ class TensorObj : public Object, public DLTensor { public: /// \cond Doxygen_Suppress static constexpr const uint32_t _type_index = TypeIndex::kTVMFFITensor; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFITensor, TensorObj, Object); /// \endcond /*! * \brief Move a Tensor to a DLPack managed tensor. * \return The converted DLPack managed tensor. */ DLManagedTensor* ToDLPack() const { TensorObj* self = const_cast(this); DLManagedTensor* ret = new DLManagedTensor(); ret->dl_tensor = *static_cast(self); ret->manager_ctx = self; ret->deleter = DLManagedTensorDeleter; details::ObjectUnsafe::IncRefObjectHandle(self); return ret; } /*! * \brief Move a Tensor to a DLPack managed tensor. * \return The converted DLPack managed tensor. */ DLManagedTensorVersioned* ToDLPackVersioned() const { TensorObj* self = const_cast(this); DLManagedTensorVersioned* ret = new DLManagedTensorVersioned(); ret->version.major = DLPACK_MAJOR_VERSION; ret->version.minor = DLPACK_MINOR_VERSION; ret->dl_tensor = *static_cast(self); ret->manager_ctx = self; ret->deleter = DLManagedTensorDeleter; details::ObjectUnsafe::IncRefObjectHandle(self); return ret; } protected: /*! * \brief Deleter for DLManagedTensor. * \param tensor The DLManagedTensor to be deleted. */ template static void DLManagedTensorDeleter(TDLManagedTensor* tensor) { TensorObj* obj = static_cast(tensor->manager_ctx); details::ObjectUnsafe::DecRefObjectHandle(obj); delete tensor; } friend class Tensor; }; namespace details { /*! *\brief Helper class to create an TensorObj from an NDAllocator * * The underlying allocator needs to be implemented by user. */ template class TensorObjFromNDAlloc : public TensorObj { public: using Self = TensorObjFromNDAlloc; template TensorObjFromNDAlloc(TNDAlloc alloc, ffi::ShapeView shape, DLDataType dtype, DLDevice device, ExtraArgs&&... extra_args) : alloc_(std::move(alloc)) { this->device = device; this->ndim = static_cast(shape.size()); this->dtype = dtype; this->byte_offset = 0; // inplace alloc shape and strides after data structure this->shape = reinterpret_cast(reinterpret_cast(this) + sizeof(Self)); this->strides = this->shape + shape.size(); std::copy(shape.begin(), shape.end(), this->shape); details::FillStridesFromShape(shape, this->strides); // call allocator to alloc data alloc_.AllocData(static_cast(this), std::forward(extra_args)...); } template TensorObjFromNDAlloc(TNDAlloc alloc, const DLTensor& prototype, ExtraArgs&&... extra_args) : alloc_(std::move(alloc)) { *static_cast(this) = prototype; this->shape = reinterpret_cast(reinterpret_cast(this) + sizeof(Self)); this->strides = this->shape + prototype.ndim; TVM_FFI_ICHECK_NOTNULL(prototype.strides); std::copy(prototype.shape, prototype.shape + prototype.ndim, this->shape); std::copy(prototype.strides, prototype.strides + prototype.ndim, this->strides); // call allocator to alloc data alloc_.AllocData(static_cast(this), std::forward(extra_args)...); } ~TensorObjFromNDAlloc() { alloc_.FreeData(static_cast(this)); } private: TNDAlloc alloc_; }; /*! \brief helper class to import from DLPack legacy DLManagedTensor */ template class TensorObjFromDLPack : public TensorObj { public: using Self = TensorObjFromDLPack; explicit TensorObjFromDLPack(TDLPackManagedTensor* tensor, bool extra_strides_at_tail) : tensor_(tensor) { *static_cast(this) = tensor_->dl_tensor; if (extra_strides_at_tail) { this->strides = reinterpret_cast(reinterpret_cast(this) + sizeof(Self)); details::FillStridesFromShape(ShapeView(tensor_->dl_tensor.shape, tensor_->dl_tensor.ndim), this->strides); } } ~TensorObjFromDLPack() { // run DLPack deleter if needed. if (tensor_->deleter != nullptr) { (*tensor_->deleter)(tensor_); } } private: TDLPackManagedTensor* tensor_; }; } // namespace details /*! * \brief Managed Tensor (n-dimensional array). * The tensor is backed by reference counted blocks. * * \note This class can be subclassed to implement downstream customized * Tensor types that are backed by the same TensorObj storage type. */ class Tensor : public ObjectRef { public: /*! * \brief Default constructor. */ Tensor() = default; /*! * \brief Constructor from a ObjectPtr. * \param n The ObjectPtr. */ explicit Tensor(::tvm::ffi::ObjectPtr n) : ObjectRef(std::move(n)) {} /*! * \brief Constructor from a UnsafeInit tag. * \param tag The UnsafeInit tag. */ explicit Tensor(::tvm::ffi::UnsafeInit tag) : ObjectRef(tag) {} /// \cond Doxygen_Suppress TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(Tensor) /// \endcond /*! * \brief Get the data pointer of the Tensor. * \return The data pointer of the Tensor. */ void* data_ptr() const { return get()->data; } /*! * \brief Get the device of the Tensor. * \return The device of the Tensor. */ DLDevice device() const { return get()->device; } /*! * \brief Get the number of dimensions in the Tensor. * \return The number of dimensions in the Tensor. */ int32_t ndim() const { return get()->ndim; } /*! * \brief Get the data type of the Tensor. * \return The data type of the Tensor. */ DLDataType dtype() const { return get()->dtype; } /*! * \brief Get the shape of the Tensor. * \return The shape of the Tensor. */ ShapeView shape() const { const TensorObj* obj = get(); return tvm::ffi::ShapeView(obj->shape, obj->ndim); } /*! * \brief Get the strides of the Tensor. * \return The strides of the Tensor. */ ShapeView strides() const { const TensorObj* obj = get(); TVM_FFI_ICHECK(obj->strides != nullptr || obj->ndim == 0); return ShapeView(obj->strides, obj->ndim); } /*! * \brief Get the size of the idx-th dimension. If the idx is negative, * it gets the size of last idx-th dimension. * \param idx The index of the size. * \return The size of the idx-th dimension. */ int64_t size(int64_t idx) const { const TensorObj* ptr = get(); int64_t adjusted_idx = idx >= 0 ? idx : (ptr->ndim + idx); if (adjusted_idx < 0 || adjusted_idx >= ptr->ndim) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds for tensor with " << ptr->ndim << " dimensions"; } return ptr->shape[adjusted_idx]; } /*! * \brief Get the stride of the idx-th dimension. If the idx is negative, * it gets the stride of last idx-th dimension. * \param idx The index of the stride. * \return The stride of the idx-th dimension. */ int64_t stride(int64_t idx) const { const TensorObj* ptr = get(); int64_t adjusted_idx = idx >= 0 ? idx : (ptr->ndim + idx); if (adjusted_idx < 0 || adjusted_idx >= ptr->ndim) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds for tensor with " << ptr->ndim << " dimensions"; } return ptr->strides[adjusted_idx]; } /*! * \brief Get the number of elements in the Tensor. * \return The number of elements in the Tensor. */ int64_t numel() const { return this->shape().Product(); } /*! * \brief Get the byte offset of the Tensor. * \return The byte offset of the Tensor. */ uint64_t byte_offset() const { return get()->byte_offset; } /*! * \brief Check if the Tensor is contiguous. * \return True if the Tensor is contiguous, false otherwise. */ bool IsContiguous() const { return tvm::ffi::IsContiguous(*get()); } /*! * \brief Check if the Tensor data is aligned to the given alignment. * \param alignment The alignment to check. * \return True if the Tensor data is aligned to the given alignment, false otherwise. */ bool IsAligned(size_t alignment) const { return tvm::ffi::IsAligned(*get(), alignment); } /*! * \brief Create a new Tensor as a strided view of the current Tensor. * \param shape The shape of the new Tensor. * \param strides The strides of the new Tensor. * \param element_offset The element offset of the new Tensor in the unit of dtype elements. * \return The new Tensor. * \note element_offset is in the unit of dtype elements not bytes. */ Tensor as_strided(ShapeView shape, ShapeView strides, std::optional element_offset = std::nullopt) const { DLTensor prototype; prototype = *static_cast(get()); prototype.shape = const_cast(shape.data()); prototype.ndim = static_cast(shape.size()); prototype.strides = const_cast(strides.data()); int64_t elem_offset_as_i64 = element_offset.value_or(0); TVM_FFI_ICHECK_GE(elem_offset_as_i64, 0); prototype.byte_offset += GetDataSize(static_cast(elem_offset_as_i64), prototype.dtype); if (prototype.byte_offset != 0 && IsDirectAddressDevice(prototype.device)) { // If the device supports direct address, we can just add the byte offset to the data pointer. prototype.data = reinterpret_cast(reinterpret_cast(prototype.data) + prototype.byte_offset); prototype.byte_offset = 0; } TVMFFIObjectHandle out; Object* obj_handle = const_cast(get()); TVM_FFI_CHECK_SAFE_CALL(TVMFFITensorCreateUnsafeView(obj_handle, &prototype, &out)); return Tensor( details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(out))); } /*! * \brief Create a Tensor from a NDAllocator. * * \note When building a kernel library, we always recommend use FromEnvAlloc when possible to * allocate intermediate Tensors. When a loaded module returns an allocated tensor to the caller, * we need to keep the module alive before the returned tensors get freed, because its * deleter is defined within the module. FromNDAlloc can be used by C++ applications and runtimes * to create Tensors. * * Example usage: * \code{.cpp} * // CPU Allocator * struct CPUNDAlloc { * void AllocData(DLTensor* tensor) { tensor->data = malloc(ffi::GetDataSize(*tensor)); } * void FreeData(DLTensor* tensor) { free(tensor->data); } * }; * * // CUDA Allocator * struct CUDANDAlloc { * void AllocData(DLTensor* tensor) { * size_t data_size = ffi::GetDataSize(*tensor); * void* ptr = nullptr; * cudaError_t err = cudaMalloc(&ptr, data_size); * TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err); * tensor->data = ptr; * } * void FreeData(DLTensor* tensor) { * if (tensor->data != nullptr) { * cudaError_t err = cudaFree(tensor->data); * TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err); * tensor->data = nullptr; * } * } * }; * * // NVSHMEM Allocator * struct NVSHMEMNDAlloc { * void AllocData(DLTensor* tensor) { * size_t size = tvm::ffi::GetDataSize(*tensor); * tensor->data = nvshmem_malloc(size); * TVM_FFI_ICHECK_NE(tensor->data, nullptr) << "nvshmem_malloc failed. size: " << size; * } * void FreeData(DLTensor* tensor) { nvshmem_free(tensor->data); } * }; * * // Allocator usage * ffi::Tensor cpu_tensor = ffi::Tensor::FromNDAlloc(CPUNDAlloc(), ...); * ffi::Tensor cuda_tensor = ffi::Tensor::FromNDAlloc(CUDANDAlloc(), ...); * ffi::Tensor nvshmem_tensor = ffi::Tensor::FromNDAlloc(NVSHMEMNDAlloc(), ...); * \endcode * * \param alloc The NDAllocator. * \param shape The shape of the Tensor. * \param dtype The data type of the Tensor. * \param device The device of the Tensor. * \param extra_args Extra arguments to be forwarded to TNDAlloc. * \return The created Tensor. * \tparam TNDAlloc The type of the NDAllocator, impelments Alloc and Free. * \tparam ExtraArgs Extra arguments to be passed to Alloc. */ template static Tensor FromNDAlloc(TNDAlloc alloc, ffi::ShapeView shape, DLDataType dtype, DLDevice device, ExtraArgs&&... extra_args) { // inplace alloc shape and strides after data structure (as a result why multiply 2) size_t num_extra_i64_at_tail = shape.size() * 2; return Tensor(make_inplace_array_object, int64_t>( num_extra_i64_at_tail, alloc, shape, dtype, device, std::forward(extra_args)...)); } /*! * \brief A variant of FromNDAlloc that allows explicit passing a strides. * * \note This function needs to ensure that strides are well-defined * with respect to the allocated compact shape. * * \param alloc The NDAllocator. * \param shape The shape of the Tensor. * \param strides The strides of the Tensor. * \param dtype The data type of the Tensor. * \param device The device of the Tensor. * \param extra_args Extra arguments to be forwarded to TNDAlloc. * \return The created Tensor. * \tparam TNDAlloc The type of the NDAllocator, impelments Alloc and Free. * \tparam ExtraArgs Extra arguments to be passed to Alloc. */ template static Tensor FromNDAllocStrided(TNDAlloc alloc, ffi::ShapeView shape, ffi::ShapeView strides, DLDataType dtype, DLDevice device, ExtraArgs&&... extra_args) { TVM_FFI_CHECK(shape.size() == strides.size(), ValueError) << "shape and strides must have the same size."; // inplace alloc shape and strides after data structure (as a result why multiply 2) size_t num_extra_i64_at_tail = shape.size() * 2; DLTensor prototype; prototype.data = nullptr; prototype.device = device; prototype.dtype = dtype; prototype.shape = const_cast(shape.data()); prototype.ndim = static_cast(shape.size()); prototype.strides = const_cast(strides.data()); prototype.byte_offset = 0; return Tensor(make_inplace_array_object, int64_t>( num_extra_i64_at_tail, alloc, prototype, std::forward(extra_args)...)); } /*! * \brief Create a Tensor from the TVMFFIEnvTensorAlloc API * * This function can be used together with TVMFFIEnvSetDLPackManagedTensorAllocator * in the extra/c_env_api.h to create a Tensor from the thread-local environment allocator. * We explicitly pass TVMFFIEnvTensorAlloc to maintain explicit dependency on extra/c_env_api.h * * \code{.cpp} * ffi::Tensor tensor = ffi::Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, shape, dtype, device); * \endcode * * \param env_alloc TVMFFIEnvTensorAlloc function pointer. * \param shape The shape of the Tensor. * \param dtype The data type of the Tensor. * \param device The device of the Tensor. * \return The created Tensor. * * \sa TVMFFIEnvTensorAlloc */ static Tensor FromEnvAlloc(int (*env_alloc)(DLTensor*, TVMFFIObjectHandle*), ffi::ShapeView shape, DLDataType dtype, DLDevice device) { TVMFFIObjectHandle out; DLTensor prototype{}; prototype.device = device; prototype.dtype = dtype; prototype.shape = const_cast(shape.data()); prototype.ndim = static_cast(shape.size()); TVM_FFI_CHECK_SAFE_CALL(env_alloc(&prototype, &out)); return Tensor( details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(out))); } /*! * \brief Create a Tensor from a DLPack managed tensor, pre v1.0 API. * \param tensor The input DLPack managed tensor. * \param require_alignment The minimum alignment requored of the data + byte_offset. * \param require_contiguous Boolean flag indicating if we need to check for contiguity. * \note This function will not run any checks on flags. * \return The created Tensor. */ static Tensor FromDLPack(DLManagedTensor* tensor, size_t require_alignment = 0, bool require_contiguous = false) { if (require_alignment != 0 && !ffi::IsAligned(tensor->dl_tensor, require_alignment)) { TVM_FFI_THROW(RuntimeError) << "FromDLPack: Data is not aligned to " << require_alignment << " bytes."; } if (require_contiguous && !ffi::IsContiguous(tensor->dl_tensor)) { TVM_FFI_THROW(RuntimeError) << "FromDLPack: Tensor is not contiguous."; } if (tensor->dl_tensor.strides != nullptr || tensor->dl_tensor.ndim == 0) { return Tensor(make_object>( tensor, /*extra_strides_at_tail=*/false)); } else { return Tensor( make_inplace_array_object, int64_t>( tensor->dl_tensor.ndim, tensor, /*extra_strides_at_tail=*/true)); } } /*! * \brief Create a Tensor from a DLPack managed tensor, post v1.0 API. * \param tensor The input DLPack managed tensor. * \param require_alignment The minimum alignment requored of the data + byte_offset. * \param require_contiguous Boolean flag indicating if we need to check for contiguity. * \return The created Tensor. */ static Tensor FromDLPackVersioned(DLManagedTensorVersioned* tensor, size_t require_alignment = 0, bool require_contiguous = false) { if (require_alignment != 0 && !ffi::IsAligned(tensor->dl_tensor, require_alignment)) { TVM_FFI_THROW(RuntimeError) << "FromDLPack: Data is not aligned to " << require_alignment << " bytes."; } if (require_contiguous && !ffi::IsContiguous(tensor->dl_tensor)) { TVM_FFI_THROW(RuntimeError) << "FromDLPack: Tensor is not contiguous."; } if (tensor->flags & DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED) { TVM_FFI_THROW(RuntimeError) << "Subbyte type padded is not yet supported"; } if (tensor->dl_tensor.strides != nullptr || tensor->dl_tensor.ndim == 0) { return Tensor(make_object>( tensor, /*extra_strides_at_tail=*/false)); } else { return Tensor( make_inplace_array_object, int64_t>(tensor->dl_tensor.ndim, tensor, /*extra_strides_at_tail=*/true)); } } /*! * \brief Convert the Tensor to a DLPack managed tensor. * \return The converted DLPack managed tensor. */ DLManagedTensor* ToDLPack() const { return get_mutable()->ToDLPack(); } /*! * \brief Convert the Tensor to a DLPack managed tensor. * \return The converted DLPack managed tensor. */ DLManagedTensorVersioned* ToDLPackVersioned() const { return get_mutable()->ToDLPackVersioned(); } /*! * \brief Get the underlying DLTensor pointer. * \return The underlying DLTensor pointer. */ const DLTensor* GetDLTensorPtr() const { return get(); } /// \cond Doxygen_Suppress [[maybe_unused]] static constexpr bool _type_is_nullable = true; using ContainerType = TensorObj; /// \endcond // the following code are convenient APIs redirections created to provide aten-style api /*! * \brief This functions redirects to ndim(). * \return The number of dimensions in the Tensor. */ inline int32_t dim() { return ndim(); } /*! * \brief This functions redirects to shape(). * \return The shape of the Tensor. */ inline ShapeView sizes() const { return shape(); } /*! * \brief This functions redirects to IsContiguous(). * \return True if the Tensor is contiguous, false otherwise. */ inline bool is_contiguous() const { return IsContiguous(); } protected: /*! * \brief Get const internal container pointer. * \return a const container pointer. */ const TensorObj* get() const { return static_cast(ObjectRef::get()); } /*! * \brief Get mutable internal container pointer. * \return a mutable container pointer. */ TensorObj* get_mutable() const { return const_cast(get()); } }; /*! * \brief A non-owning view of a Tensor. * * This class stores a light-weight non-owning view of a Tensor. * This is useful for accessing tensor data without retaining a strong reference to the Tensor. * Since the caller may not always be able to pass in a strong referenced tensor. * * It is the user's responsibility to ensure * that the underlying tensor data outlives the `TensorView`. * This responsibility extends to all data pointed to by the underlying DLTensor. * This includes not only the tensor elements in data but also the memory for shape and strides. * * When exposing a function that expects only expects a TensorView, we recommend using * ffi::TensorView as the argument type instead of ffi::Tensor. */ class TensorView { public: /*! * \brief Create a TensorView from a Tensor. * \param tensor The input Tensor. */ TensorView(const Tensor& tensor) { // NOLINT(*) TVM_FFI_ICHECK(tensor.defined()); tensor_ = *tensor.GetDLTensorPtr(); } // NOLINT(*) /*! * \brief Create a TensorView from a DLTensor. * \param tensor The input DLTensor. */ TensorView(const DLTensor* tensor) { // NOLINT(*) TVM_FFI_ICHECK(tensor != nullptr); tensor_ = *tensor; } /*! * \brief Copy constructor. * \param tensor The input TensorView. */ TensorView(const TensorView& tensor) = default; /*! * \brief Move constructor. * \param tensor The input TensorView. */ TensorView(TensorView&& tensor) = default; /*! * \brief Copy assignment operator. * \param tensor The input TensorView. * \return The created TensorView. */ TensorView& operator=(const TensorView& tensor) = default; /*! * \brief Move assignment operator. * \param tensor The input TensorView. * \return The created TensorView. */ TensorView& operator=(TensorView&& tensor) = default; /*! * \brief Assignment operator from a Tensor. * \param tensor The input Tensor. * \return The created TensorView. */ TensorView& operator=(const Tensor& tensor) { TVM_FFI_ICHECK(tensor.defined()); tensor_ = *tensor.GetDLTensorPtr(); return *this; } // explicitly delete move constructor TensorView(Tensor&& tensor) = delete; // NOLINT(*) // delete move assignment operator from owned tensor TensorView& operator=(Tensor&& tensor) = delete; /*! * \brief Get the data pointer of the Tensor. * \return The data pointer of the Tensor. */ void* data_ptr() const { return tensor_.data; } /*! * \brief Get the device of the Tensor. * \return The device of the Tensor. */ DLDevice device() const { return tensor_.device; } /*! * \brief Get the number of dimensions in the Tensor. * \return The number of dimensions in the Tensor. */ int32_t ndim() const { return tensor_.ndim; } /*! * \brief Get the data type of the Tensor. * \return The data type of the Tensor. */ DLDataType dtype() const { return tensor_.dtype; } /*! * \brief Get the shape of the Tensor. * \return The shape of the Tensor. */ ShapeView shape() const { return ShapeView(tensor_.shape, tensor_.ndim); } /*! * \brief Get the number of elements in the Tensor. * \return The number of elements in the Tensor. */ int64_t numel() const { return this->shape().Product(); } /*! * \brief Get the strides of the Tensor. * \return The strides of the Tensor. */ ShapeView strides() const { TVM_FFI_ICHECK(tensor_.strides != nullptr || tensor_.ndim == 0); return ShapeView(tensor_.strides, tensor_.ndim); } /*! * \brief Get the size of the idx-th dimension. If the idx is negative, * it gets the size of last idx-th dimension. * \param idx The index of the size. * \return The size of the idx-th dimension. */ int64_t size(int64_t idx) const { int64_t adjusted_idx = idx >= 0 ? idx : (tensor_.ndim + idx); if (adjusted_idx < 0 || adjusted_idx >= tensor_.ndim) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds for tensor with " << tensor_.ndim << " dimensions"; } return tensor_.shape[adjusted_idx]; } /*! * \brief Get the stride of the idx-th dimension. If the idx is negative, * it gets the stride of last idx-th dimension. * \param idx The index of the stride. * \return The stride of the idx-th dimension. */ int64_t stride(int64_t idx) const { int64_t adjusted_idx = idx >= 0 ? idx : (tensor_.ndim + idx); if (adjusted_idx < 0 || adjusted_idx >= tensor_.ndim) { TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds for tensor with " << tensor_.ndim << " dimensions"; } return tensor_.strides[adjusted_idx]; } /*! * \brief Get the byte offset of the Tensor. * \return The byte offset of the Tensor. */ uint64_t byte_offset() const { return tensor_.byte_offset; } /*! * \brief Check if the Tensor is contiguous. * \return True if the Tensor is contiguous, false otherwise. */ bool IsContiguous() const { return tvm::ffi::IsContiguous(tensor_); } // the following code are convenient APIs redirections created to provide aten-style api /*! * \brief This functions redirects to ndim(). * \return The number of dimensions in the Tensor. */ int32_t dim() const { return ndim(); } /*! * \brief This functions redirects to shape(). * \return The shape of the Tensor. */ ShapeView sizes() const { return shape(); } /*! * \brief This functions redirects to IsContiguous(). * \return True if the Tensor is contiguous, false otherwise. */ bool is_contiguous() const { return IsContiguous(); } /*! * \brief Create a new TensorView as a strided view of the current TensorView. * * Use this function with extreme caution. The user must ensure that the shape and strides * arrays, as well as the data pointer, remain valid for the lifetime of the returned TensorView. * * One common anti-pattern is to create temporary shape/strides arrays and pass them in, but * then deallocate the temporary arrays immediately after the call to as_strided, which * causes the returned TensorView to point to invalid memory. * * \param shape The shape of the new TensorView. * \param strides The strides of the new TensorView. * \param element_offset The element offset of the new TensorView in units of dtype elements, not * bytes. * \return The new TensorView. * * \note The caller must ensure that the shape and strides arrays remain valid for the lifetime * of the returned TensorView. */ TensorView as_strided(ShapeView shape, ShapeView strides, std::optional element_offset = std::nullopt) const { DLTensor prototype = tensor_; prototype.shape = const_cast(shape.data()); prototype.ndim = static_cast(shape.size()); prototype.strides = const_cast(strides.data()); TVM_FFI_ICHECK_EQ(shape.size(), strides.size()); int64_t elem_offset_as_i64 = element_offset.value_or(0); TVM_FFI_ICHECK_GE(elem_offset_as_i64, 0); prototype.byte_offset += GetDataSize(static_cast(elem_offset_as_i64), prototype.dtype); if (prototype.byte_offset != 0 && IsDirectAddressDevice(prototype.device)) { // If the device supports direct address, we can just add the byte offset to the data pointer. prototype.data = reinterpret_cast(reinterpret_cast(prototype.data) + prototype.byte_offset); prototype.byte_offset = 0; } return TensorView(&prototype); } private: DLTensor tensor_; template friend struct TypeTraits; }; /*! * \brief Get the data size of the Tensor. * \param tensor The input Tensor. * \return The data size of the Tensor. */ inline size_t GetDataSize(const Tensor& tensor) { return GetDataSize(tensor.numel(), tensor.dtype()); } /*! * \brief Get the data size of the TensorView. * \param tensor The input TensorView. * \return The data size of the TensorView. */ inline size_t GetDataSize(const TensorView& tensor) { return GetDataSize(tensor.numel(), tensor.dtype()); } // TensorView type, allow implicit casting from DLTensor* // NOTE: we deliberately do not support MoveToAny and MoveFromAny since it does not retain ownership template <> struct TypeTraits : public TypeTraitsBase { static constexpr bool storage_enabled = false; static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDLTensorPtr; TVM_FFI_INLINE static void CopyToAnyView(const TensorView& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIDLTensorPtr; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_ptr = const_cast(&(src.tensor_)); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIDLTensorPtr; } TVM_FFI_INLINE static TensorView CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { return TensorView(static_cast(src->v_ptr)); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIDLTensorPtr) { return TensorView(static_cast(src->v_ptr)); } else if (src->type_index == TypeIndex::kTVMFFITensor) { return TensorView(TVMFFITensorGetDLTensorPtr(src->v_obj)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIDLTensorPtr; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIDLTensorPtr) + R"("})"; } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_TENSOR_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/tuple.h000066400000000000000000000322431521067262500214360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/tuple.h * \brief Typed tuple like std::tuple backed by ArrayObj container. */ #ifndef TVM_FFI_CONTAINER_TUPLE_H_ #define TVM_FFI_CONTAINER_TUPLE_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Typed tuple like std::tuple backed by ArrayObj container. * * Tuple implements in-place copy-on-write semantics. * * \tparam Types The types of the tuple elements */ template class Tuple : public ObjectRef { public: static_assert(details::all_storage_enabled_v, "All types used in Tuple<...> must be compatible with Any"); /*! \brief Default constructor */ Tuple() : ObjectRef(MakeDefaultTupleNode()) {} /*! * \brief Constructor with UnsafeInit */ explicit Tuple(UnsafeInit tag) : ObjectRef(tag) {} /*! \brief Copy constructor */ Tuple(const Tuple& other) : ObjectRef(other) {} /*! \brief Move constructor */ Tuple(Tuple&& other) noexcept : ObjectRef(std::move(other)) {} /*! * \brief Constructor from another tuple * \param other The other tuple * \tparam UTypes The types of the other tuple * \tparam The enable_if_t type */ template && ...), int>> Tuple(const Tuple& other) : ObjectRef(other) {} // NOLINT(google-explicit-constructor) /*! * \brief Constructor from another tuple * \param other The other tuple * \tparam UTypes The types of the other tuple * \tparam The enable_if_t type */ template && ...), int>> Tuple(Tuple&& other) // NOLINT(google-explicit-constructor) : ObjectRef(std::move(other)) {} /*! * \brief Constructor from arguments * \param args The arguments * \tparam UTypes The types of the other tuple */ template , Tuple> && ...))>> explicit Tuple(UTypes&&... args) : ObjectRef(MakeTupleNode(std::forward(args)...)) {} /*! * \brief Assignment from another tuple * \param other The other tuple * \tparam The enable_if_t type */ TVM_FFI_INLINE Tuple& operator=(const Tuple& other) { data_ = other.data_; return *this; } /*! * \brief Assignment from another tuple * \param other The other tuple * \tparam The enable_if_t type */ TVM_FFI_INLINE Tuple& operator=(Tuple&& other) noexcept { data_ = std::move(other.data_); return *this; } /*! * \brief Assignment from another tuple * \param other The other tuple * \tparam UTypes The types of the other tuple * \tparam The enable_if_t type */ template && ...)>> TVM_FFI_INLINE Tuple& operator=(const Tuple& other) { data_ = other.data_; return *this; } /*! * \brief Assignment from another tuple * \param other The other tuple * \tparam UTypes The types of the other tuple * \tparam The enable_if_t type */ template && ...)>> TVM_FFI_INLINE Tuple& operator=(Tuple&& other) { data_ = std::move(other.data_); return *this; } /*! * \brief Get I-th element of the tuple * * \tparam I The index of the element to get * \return The I-th element of the tuple * \note We use stl style since get usually is like a getter. */ template auto get() const& { static_assert(I < sizeof...(Types), "Tuple index out of bounds"); using ReturnType = std::tuple_element_t>; const Any* ptr = GetArrayObj()->begin() + I; return details::AnyUnsafe::CopyFromAnyViewAfterCheck(*ptr); } /*! * \brief Move out I-th element of the tuple * * \tparam I The index of the element to get * \return The I-th element of the tuple * \note We use stl style since get usually is like a getter. */ template auto get() && { if (!this->unique()) { // fallback to const& version if not unique return std::as_const(*this).template get(); } static_assert(I < sizeof...(Types), "Tuple index out of bounds"); using ReturnType = std::tuple_element_t>; Any* ptr = GetArrayObj()->MutableBegin() + I; return details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(*ptr)); } /*! * \brief Set I-th element of the tuple * * \param item The item to set * \tparam I The index of the element to set * \tparam U The type of the item * * \note This function will perform copy on write if underlying * container is not uniquely owned. * We use CamelCase since Set can cause copy on write * and is more complicated than simple field setter. */ template void Set(U&& item) { static_assert(I < sizeof...(Types), "Tuple index out of bounds"); using T = std::tuple_element_t>; this->CopyIfNotUnique(); Any* ptr = GetArrayObj()->MutableBegin() + I; *ptr = T(std::forward(item)); } /*! \brief specify container node */ using ContainerType = ArrayObj; private: static ObjectPtr MakeDefaultTupleNode() { ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); // increase size after each new to ensure exception safety ((new (itr++) Any(Types()), p->TVMFFISeqCell::size++), ...); return p; } template static ObjectPtr MakeTupleNode(UTypes&&... args) { ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); // increase size after each new to ensure exception safety ((new (itr++) Any(Types(std::forward(args))), p->TVMFFISeqCell::size++), ...); return p; } /*! \brief Copy on write */ void CopyIfNotUnique() { if (!data_.unique()) { ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); const Any* read = GetArrayObj()->begin(); // increase size after each new to ensure exception safety for (size_t i = 0; i < sizeof...(Types); ++i) { new (itr++) Any(*read++); p->TVMFFISeqCell::size++; } data_ = std::move(p); } } /*! \return The underlying ArrayObj */ ArrayObj* GetArrayObj() const { return static_cast(data_.get()); } template friend class Tuple; }; template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public ObjectRefTypeTraitsBase> { using ObjectRefTypeTraitsBase>::CopyFromAnyViewAfterCheck; TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) { return TypeTraitsBase::GetMismatchTypeInfo(src); } const ArrayObj* n = reinterpret_cast(src->v_obj); if (n->size() != sizeof...(Types)) { return "Array[size=" + std::to_string(n->size()) + "]"; } return GetMismatchTypeInfoHelper<0, Types...>(n->begin()); } template TVM_FFI_INLINE static std::string GetMismatchTypeInfoHelper(const Any* arr) { if constexpr (!std::is_same_v) { const Any& any_v = arr[I]; if (!details::AnyUnsafe::CheckAnyStrict(any_v) && !(any_v.try_cast().has_value())) { // now report the accurate mismatch information return "Array[index " + std::to_string(I) + ": " + details::AnyUnsafe::GetMismatchTypeInfo(any_v) + "]"; } } if constexpr (sizeof...(Rest) > 0) { return GetMismatchTypeInfoHelper(arr); } TVM_FFI_THROW(InternalError) << "Cannot reach here"; TVM_FFI_UNREACHABLE(); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) return false; const ArrayObj* n = reinterpret_cast(src->v_obj); if (n->size() != sizeof...(Types)) return false; const TVMFFIAny* ffi_any_arr = reinterpret_cast(n->begin()); return CheckAnyStrictHelper<0, Types...>(ffi_any_arr); } template TVM_FFI_INLINE static bool CheckAnyStrictHelper(const TVMFFIAny* src_arr) { if constexpr (!std::is_same_v) { if (!TypeTraits::CheckAnyStrict(src_arr + I)) { return false; } } if constexpr (sizeof...(Rest) > 0) { return CheckAnyStrictHelper(src_arr); } return true; } TVM_FFI_INLINE static std::optional> TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) return std::nullopt; const ArrayObj* n = reinterpret_cast(src->v_obj); if (n->size() != sizeof...(Types)) return std::nullopt; // fast path, storage is already in the right type if (CheckAnyStrict(src)) { return CopyFromAnyViewAfterCheck(src); } // slow path, try to convert to each type to match the tuple storage need. Array arr = TypeTraits>::CopyFromAnyViewAfterCheck(src); Any* ptr = arr.CopyOnWrite()->MutableBegin(); if (TryConvertElements<0, Types...>(ptr)) { return details::ObjectUnsafe::ObjectRefFromObjectPtr>( details::ObjectUnsafe::ObjectPtrFromObjectRef(arr)); } return std::nullopt; } template TVM_FFI_INLINE static bool TryConvertElements(Any* arr) { if constexpr (!std::is_same_v) { if (auto opt_convert = arr[I].try_cast()) { arr[I] = *std::move(opt_convert); } else { return false; } } if constexpr (sizeof...(Rest) > 0) { return TryConvertElements(std::move(arr)); } else { return true; } } TVM_FFI_INLINE static std::string TypeStr() { return details::ContainerTypeStr("Tuple"); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":"Tuple","args":[)"; const char* sep = ""; ((oss << sep << details::TypeSchema::v(), sep = ","), ...); oss << "]}"; return oss.str(); } }; namespace details { template inline constexpr bool type_contains_v, Tuple> = (type_contains_v && ...); } // namespace details /// \cond Doxygen_Suppress /// NOTE: ADL friendly get functions /// Example usage: { using std::get; get<0>(t); } /// ADL will find the right get function /** * \brief get I-th element of the tuple * \tparam I The index of the element to get * \param t The tuple * \return The I-th element of the tuple */ template inline constexpr auto get(const Tuple& t) -> std::tuple_element_t> { return t.template get(); } /** * \brief get I-th element of the tuple * \tparam I The index of the element to get * \param t The tuple (rvalue) * \return The I-th element of the tuple */ template inline constexpr auto get(Tuple&& t) -> std::tuple_element_t> { return std::move(t).template get(); } /// NOTE: C++17 deduction guide template Tuple(UTypes&&...) -> Tuple>...>; /// \endcond } // namespace ffi } // namespace tvm namespace std { template struct tuple_size<::tvm::ffi::Tuple> : public std::integral_constant {}; template struct tuple_element> { using type = std::tuple_element_t>; }; } // namespace std #endif // TVM_FFI_CONTAINER_TUPLE_H_ tvm-ffi-0.1.12/include/tvm/ffi/container/variant.h000066400000000000000000000246341521067262500217560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/container/variant.h * \brief Runtime variant container types. */ #ifndef TVM_FFI_CONTAINER_VARIANT_H_ #define TVM_FFI_CONTAINER_VARIANT_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { namespace details { /*! * \brief Base class for Variant. * * \tparam all_storage_object Whether all types are derived from ObjectRef. */ template class VariantBase { public: TVM_FFI_INLINE bool same_as(const VariantBase& other) const { return data_.same_as(other.data_); } protected: template explicit VariantBase(T other) : data_(std::move(other)) {} TVM_FFI_INLINE void SetData(Any other_data) { data_ = std::move(other_data); } TVM_FFI_INLINE Any MoveToAny() && { return std::move(data_); } TVM_FFI_INLINE AnyView ToAnyView() const { return data_.operator AnyView(); } Any data_; }; // Specialization for all object ref case, backed by ObjectRef. template <> class VariantBase : public ObjectRef { protected: template explicit VariantBase(const T& other) : ObjectRef(other) {} template , VariantBase>>> explicit VariantBase(T&& other) : ObjectRef(std::forward(other)) {} explicit VariantBase(UnsafeInit tag) : ObjectRef(tag) {} explicit VariantBase(Any other) : ObjectRef(details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(other))) {} TVM_FFI_INLINE void SetData(ObjectPtr other) { data_ = std::move(other); } TVM_FFI_INLINE Any MoveToAny() && { return Any(ObjectRef(std::move(data_))); } TVM_FFI_INLINE AnyView ToAnyView() const { TVMFFIAny any_data; if (data_ == nullptr) { any_data.type_index = TypeIndex::kTVMFFINone; any_data.zero_padding = 0; any_data.v_int64 = 0; } else { TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(&any_data); any_data.type_index = data_->type_index(); any_data.zero_padding = 0; any_data.v_obj = details::ObjectUnsafe::TVMFFIObjectPtrFromObjectPtr(data_); } return AnyView::CopyFromTVMFFIAny(any_data); } }; } // namespace details /*! * \brief A typed variant container. * * When all values are ObjectRef, Variant is backed by ObjectRef, * otherwise it is backed by Any. */ template class Variant : public details::VariantBase> { public: /// \cond Doxygen_Suppress using TParent = details::VariantBase>; static_assert(details::all_storage_enabled_v, "All types used in Variant<...> must be compatible with Any"); /* * \brief Helper utility to check if the type can be contained in the variant */ template static constexpr bool variant_contains_v = (details::type_contains_v || ...); /* \brief Helper utility for SFINAE if the type is part of the variant */ template using enable_if_variant_contains_t = std::enable_if_t>; /// \endcond /*! * \brief Constructor from another variant * \param other The other variant */ Variant(const Variant& other) : TParent(other.data_) {} /*! * \brief Constructor from another variant * \param other The other variant */ Variant(Variant&& other) noexcept : TParent(std::move(other.data_)) {} /*! * \brief Assignment from another variant * \param other The other variant */ TVM_FFI_INLINE Variant& operator=(const Variant& other) { this->SetData(other.data_); return *this; } /*! * \brief Assignment from another variant * \param other The other variant */ TVM_FFI_INLINE Variant& operator=(Variant&& other) noexcept { this->SetData(std::move(other.data_)); return *this; } /*! * \brief Constructor from another variant * \param other The other variant */ template > Variant(T other) : TParent(std::move(other)) {} // NOLINT(*) /*! * \brief Assignment from another variant * \param other The other variant */ template > TVM_FFI_INLINE Variant& operator=(T other) { return operator=(Variant(std::move(other))); } /*! * \brief Try to cast to a type T, return std::nullopt if the cast is not possible. * \return The casted value, or std::nullopt if the cast is not possible. * \tparam T The type to cast to. */ template > TVM_FFI_INLINE std::optional as() const { return this->TParent::ToAnyView().template as(); } /*! * \brief Shortcut of as Object to cast to a const pointer when T is an Object. * * \tparam T The object type. * \return The requested pointer, returns nullptr if type mismatches. */ template >> TVM_FFI_INLINE const T* as() const { return this->TParent::ToAnyView().template as().value_or(nullptr); } /*! * \brief Get the value of the variant in type T, throws an exception if cast fails. * \return The value of the variant * \tparam T The type to get. */ template > TVM_FFI_INLINE T get() const& { return this->TParent::ToAnyView().template cast(); } /*! * \brief Get the value of the variant in type T, throws an exception if cast fails. * \return The value of the variant * \tparam T The type to get. */ template > TVM_FFI_INLINE T get() && { return std::move(*this).TParent::MoveToAny().template cast(); } /*! * \brief Get the type key of the variant * \return The type key of the variant */ TVM_FFI_INLINE std::string GetTypeKey() const { return this->TParent::ToAnyView().GetTypeKey(); } private: friend struct TypeTraits>; friend struct ObjectPtrHash; friend struct ObjectPtrEqual; // constructor from any explicit Variant(Any data) : TParent(std::move(data)) {} /*! * \brief Get the object pointer from the variant * \note This function is only available if all types used in Variant<...> are derived from * ObjectRef */ TVM_FFI_INLINE Object* GetObjectPtrForHashEqual() const { constexpr bool all_object_v = (std::is_base_of_v && ...); static_assert(all_object_v, "All types used in Variant<...> must be derived from ObjectRef " "to enable ObjectPtrHash/ObjectPtrEqual"); return this->data_.get(); } // rexpose to friend class using TParent::MoveToAny; using TParent::ToAnyView; }; template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public TypeTraitsBase { TVM_FFI_INLINE static void CopyToAnyView(const Variant& src, TVMFFIAny* result) { *result = src.ToAnyView().CopyToTVMFFIAny(); } TVM_FFI_INLINE static void MoveToAny(Variant src, TVMFFIAny* result) { *result = details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(src).MoveToAny()); } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { return TypeTraitsBase::GetMismatchTypeInfo(src); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return (TypeTraits::CheckAnyStrict(src) || ...); } TVM_FFI_INLINE static Variant CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { return Variant(Any(AnyView::CopyFromTVMFFIAny(*src))); } TVM_FFI_INLINE static Variant MoveFromAnyAfterCheck(TVMFFIAny* src) { return Variant(details::AnyUnsafe::MoveTVMFFIAnyToAny(src)); } TVM_FFI_INLINE static std::optional> TryCastFromAnyView(const TVMFFIAny* src) { // fast path, storage is already in the right type if (CheckAnyStrict(src)) { return CopyFromAnyViewAfterCheck(src); } // More expensive path, try to convert to each type, in order of declaration return TryVariantTypes(src); } template TVM_FFI_INLINE static std::optional> TryVariantTypes(const TVMFFIAny* src) { if (auto opt_convert = TypeTraits::TryCastFromAnyView(src)) { return Variant(*std::move(opt_convert)); } if constexpr (sizeof...(Rest) > 0) { return TryVariantTypes(src); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return details::ContainerTypeStr("Variant"); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":"Variant","args":[)"; const char* sep = ""; ((oss << sep << details::TypeSchema::v(), sep = ","), ...); oss << "]}"; return oss.str(); } }; template TVM_FFI_INLINE size_t ObjectPtrHash::operator()(const Variant& a) const { return std::hash()(a.GetObjectPtrForHashEqual()); } template TVM_FFI_INLINE bool ObjectPtrEqual::operator()(const Variant& a, const Variant& b) const { return a.GetObjectPtrForHashEqual() == b.GetObjectPtrForHashEqual(); } namespace details { template inline constexpr bool type_contains_v, T> = (type_contains_v || ...); } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_CONTAINER_VARIANT_H_ tvm-ffi-0.1.12/include/tvm/ffi/dtype.h000066400000000000000000000136141521067262500174510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/dtype.h * \brief Data type handling. */ #ifndef TVM_FFI_DTYPE_H_ #define TVM_FFI_DTYPE_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Extension code beyond the DLDataType. * * This class is always consistent with the DLPack. */ enum DLExtDataTypeCode { kDLExtCustomBegin = 129 }; namespace details { /* * \brief Convert a DLDataTypeCode to a string. * \param os The output stream. * \param type_code The DLDataTypeCode to convert. */ inline const char* DLDataTypeCodeAsCStr(DLDataTypeCode type_code) { // NOLINT(*) switch (static_cast(type_code)) { case kDLInt: { return "int"; } case kDLUInt: { return "uint"; } case kDLFloat: { return "float"; } case kDLOpaqueHandle: { return "handle"; } case kDLBfloat: { return "bfloat"; } case kDLBool: { return "bool"; } case kDLFloat8_e3m4: { return "float8_e3m4"; } case kDLFloat8_e4m3: { return "float8_e4m3"; } case kDLFloat8_e4m3b11fnuz: { return "float8_e4m3b11fnuz"; } case kDLFloat8_e4m3fn: { return "float8_e4m3fn"; } case kDLFloat8_e4m3fnuz: { return "float8_e4m3fnuz"; } case kDLFloat8_e5m2: { return "float8_e5m2"; } case kDLFloat8_e5m2fnuz: { return "float8_e5m2fnuz"; } case kDLFloat8_e8m0fnu: { return "float8_e8m0fnu"; } case kDLFloat6_e2m3fn: { return "float6_e2m3fn"; } case kDLFloat6_e3m2fn: { return "float6_e3m2fn"; } case kDLFloat4_e2m1fn: { return "float4_e2m1fn"; } default: { if (static_cast(type_code) >= static_cast(DLExtDataTypeCode::kDLExtCustomBegin)) { return "custom"; } else { TVM_FFI_THROW(ValueError) << "DLDataType contains unknown type_code=" << static_cast(type_code); } TVM_FFI_UNREACHABLE(); } } } } // namespace details /*! * \brief Convert a string to a DLDataType. * \param str The string to convert. * \return The DLDataType. */ inline DLDataType StringToDLDataType(const String& str) { DLDataType out; TVMFFIByteArray data{str.data(), str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFIDataTypeFromString(&data, &out)); return out; } /*! * \brief Convert a DLDataType to a string. * \param dtype The DLDataType to convert. * \return The string. */ inline String DLDataTypeToString(DLDataType dtype) { TVMFFIAny out; TVM_FFI_CHECK_SAFE_CALL(TVMFFIDataTypeToString(&dtype, &out)); return TypeTraits::MoveFromAnyAfterCheck(&out); } // DLDataType template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDataType; TVM_FFI_INLINE static void CopyToAnyView(const DLDataType& src, TVMFFIAny* result) { // clear padding part to ensure the equality check can always check the v_uint64 part result->v_uint64 = 0; result->type_index = TypeIndex::kTVMFFIDataType; result->zero_padding = 0; result->v_dtype = src; } TVM_FFI_INLINE static void MoveToAny(DLDataType src, TVMFFIAny* result) { // clear padding part to ensure the equality check can always check the v_uint64 part result->v_uint64 = 0; result->type_index = TypeIndex::kTVMFFIDataType; result->zero_padding = 0; result->v_dtype = src; } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIDataType; } TVM_FFI_INLINE static DLDataType CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIDataType); return src->v_dtype; } TVM_FFI_INLINE static DLDataType MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type — move is just copy. return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIDataType) { return src->v_dtype; } // enable string to dtype auto conversion if (auto opt_str = TypeTraits::TryCastFromAnyView(src)) { return StringToDLDataType(*opt_str); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return ffi::StaticTypeKey::kTVMFFIDataType; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(ffi::StaticTypeKey::kTVMFFIDataType) + R"("})"; } }; } // namespace ffi } // namespace tvm // define DLDataType comparison and printing in root namespace inline std::ostream& operator<<(std::ostream& os, DLDataType dtype) { // NOLINT(*) return os << tvm::ffi::DLDataTypeToString(dtype); } inline bool operator==(const DLDataType& lhs, const DLDataType& rhs) { return lhs.code == rhs.code && lhs.bits == rhs.bits && lhs.lanes == rhs.lanes; } inline bool operator!=(const DLDataType& lhs, const DLDataType& rhs) { return !(lhs == rhs); } #endif // TVM_FFI_DTYPE_H_ tvm-ffi-0.1.12/include/tvm/ffi/endian.h000066400000000000000000000056601521067262500175640ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file tvm/ffi/endian.h * \brief Endian detection and handling */ #ifndef TVM_FFI_ENDIAN_H_ #define TVM_FFI_ENDIAN_H_ #include #include #ifndef TVM_FFI_IO_USE_LITTLE_ENDIAN #define TVM_FFI_IO_USE_LITTLE_ENDIAN 1 #endif #ifdef TVM_FFI_CMAKE_LITTLE_ENDIAN // If compiled with CMake, use CMake's endian detection logic #define TVM_FFI_LITTLE_ENDIAN TVM_FFI_CMAKE_LITTLE_ENDIAN #else #if defined(__APPLE__) || defined(_WIN32) #define TVM_FFI_LITTLE_ENDIAN 1 #elif defined(__GLIBC__) || defined(__GNU_LIBRARY__) || defined(__ANDROID__) || \ defined(__RISCV__) || defined(__MUSL__) #include #define TVM_FFI_LITTLE_ENDIAN (__BYTE_ORDER == __LITTLE_ENDIAN) #elif defined(__FreeBSD__) || defined(__OpenBSD__) #include #define TVM_FFI_LITTLE_ENDIAN (_BYTE_ORDER == _LITTLE_ENDIAN) #elif defined(__QNX__) #include #define TVM_FFI_LITTLE_ENDIAN (BYTE_ORDER == LITTLE_ENDIAN) #elif defined(__EMSCRIPTEN__) || defined(__hexagon__) #define TVM_FFI_LITTLE_ENDIAN 1 #elif defined(__sun) || defined(sun) #include #if defined(_LITTLE_ENDIAN) #define TVM_FFI_LITTLE_ENDIAN 1 #else #define TVM_FFI_LITTLE_ENDIAN 0 #endif #else #error "Unable to determine endianness of your machine; use CMake to compile" #endif #endif /*! \brief whether serialize using little endian */ #define TVM_FFI_IO_NO_ENDIAN_SWAP (TVM_FFI_LITTLE_ENDIAN == TVM_FFI_IO_USE_LITTLE_ENDIAN) namespace tvm { namespace ffi { /*! * \brief A generic inplace byte swapping function. * \param data The data pointer. * \param elem_bytes The number of bytes of the data elements * \param num_elems Number of elements in the data. * \note Always try pass in constant elem_bytes to enable * compiler optimization */ inline void ByteSwap(void* data, size_t elem_bytes, size_t num_elems) { for (size_t i = 0; i < num_elems; ++i) { uint8_t* bptr = reinterpret_cast(data) + elem_bytes * i; for (size_t j = 0; j < elem_bytes / 2; ++j) { uint8_t v = bptr[elem_bytes - 1 - j]; bptr[elem_bytes - 1 - j] = bptr[j]; bptr[j] = v; } } } } // namespace ffi } // namespace tvm #endif // TVM_FFI_ENDIAN_H_ tvm-ffi-0.1.12/include/tvm/ffi/enum.h000066400000000000000000000127121521067262500172660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/enum.h * \brief Base class for FFI-registered enum types. */ #ifndef TVM_FFI_ENUM_H_ #define TVM_FFI_ENUM_H_ #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { class Enum; /*! * \brief Base class for FFI-registered enums. * * Each registered variant is a unique, process-wide singleton with a * dense ordinal (``_value``) and string ``_name``. Subclasses may add * *declared fields* — part of the variant's schema, set at registration * time via ``reflection::EnumDef``. Separately, any consumer may * attach *extensible attributes* (per-variant metadata stored outside * the variant's fields) via ``EnumDef::set_attr`` or the Python * ``Enum.def_attr`` surface, without modifying ``EnumClsObj``. * * \sa reflection::EnumDef */ class EnumObj : public Object { public: /*! \brief Declared field: dense ordinal assigned at registration time (0-indexed per class). */ int64_t _value; /*! \brief Declared field: instance name (e.g., ``"Add"`` for ``Op.Add``). */ String _name; EnumObj() = default; /*! * \brief Construct an EnumObj with an explicit ordinal and name. * \param value The dense ordinal (0-indexed per enum class). * \param name The instance name key. */ EnumObj(int64_t value, String name) : _value(value), _name(std::move(name)) {} /*! * \brief Look up the registered singleton for ``EnumClsObj`` by name. * * Reads from the per-class ``reflection::type_attr::kEnumEntries`` * registry populated by ``reflection::EnumDef``. Instances * are unique per ``(type_key, name)`` pair for the life of the process, * so the returned ``Enum`` compares equal (by pointer) to every other * lookup of the same name. Throws ``RuntimeError`` if no instance with * the given name is registered for ``EnumClsObj``. * * \tparam EnumClsObj An ``Object`` subclass deriving from ``EnumObj``. * \param name The instance name to look up (e.g., ``"Add"``). * \return The registered ``Enum`` singleton. */ template static Enum Get(const String& name); /// \cond Doxygen_Suppress static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance; TVM_FFI_DECLARE_OBJECT_INFO("ffi.Enum", EnumObj, Object); /// \endcond private: /*! * \brief Return the process-wide ``__ffi_enum_entries__`` column pointer. * * The column is registered at library init via ``EnsureTypeAttrColumn`` * and the struct its pointer refers to is stable for the lifetime of the * process, so we cache the lookup in a function-local static. */ static const TVMFFITypeAttrColumn* GetEnumEntriesColumn() { constexpr TVMFFIByteArray kAttrName = reflection::AsByteArray(reflection::type_attr::kEnumEntries); static const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&kAttrName); return column; } }; /*! * \brief ObjectRef wrapper for ``EnumObj``. * * Holds a shared reference to a registered singleton. Two ``Enum`` * values compare structurally equal if and only if they point at the * same underlying object (see ``kTVMFFISEqHashKindUniqueInstance``), * which — given the register-once registry — is equivalent to sharing * the same ``(type_key, name)`` pair. * * \sa EnumObj * \sa reflection::EnumDef */ class Enum : public ObjectRef { public: /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Enum, ObjectRef, EnumObj); /// \endcond }; template inline Enum EnumObj::Get(const String& name) { static_assert(std::is_base_of_v, "EnumObj::Get requires T to be a subclass of EnumObj"); const TVMFFITypeAttrColumn* column = GetEnumEntriesColumn(); int32_t type_index = EnumClsObj::RuntimeTypeIndex(); if (column != nullptr) { int32_t offset = type_index - column->begin_index; if (offset >= 0 && offset < column->size) { const TVMFFIAny* stored = &column->data[offset]; if (stored->type_index != kTVMFFINone) { Dict entries = AnyView::CopyFromTVMFFIAny(*stored).cast>(); auto it = entries.find(name); if (it != entries.end()) { return (*it).second; } } } } TVM_FFI_THROW(RuntimeError) << "Enum `" << EnumClsObj::_type_key << "` has no instance named `" << name << "`"; TVM_FFI_UNREACHABLE(); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_ENUM_H_ tvm-ffi-0.1.12/include/tvm/ffi/error.h000066400000000000000000000453261521067262500174620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file tvm/ffi/error.h * \brief Error handling component. */ #ifndef TVM_FFI_ERROR_H_ #define TVM_FFI_ERROR_H_ #include #include #include #include #include #include #include #include #include #include #include #include /*! * \brief Macro defines whether we enable libbacktrace */ #ifndef TVM_FFI_USE_LIBBACKTRACE #define TVM_FFI_USE_LIBBACKTRACE 1 #endif /*! * \brief Macro defines whether to install signal handler * and print backtrace during segfault */ #ifndef TVM_FFI_BACKTRACE_ON_SEGFAULT #define TVM_FFI_BACKTRACE_ON_SEGFAULT 1 #endif #ifndef TVM_FFI_ALWAYS_LOG_BEFORE_THROW #define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 0 #endif namespace tvm { namespace ffi { /*! * \brief Error object class. */ class ErrorObj : public Object, public TVMFFIErrorCell { public: ErrorObj() { this->cause_chain = nullptr; this->extra_context = nullptr; } ~ErrorObj() { if (this->cause_chain != nullptr) { details::ObjectUnsafe::DecRefObjectHandle(this->cause_chain); } if (this->extra_context != nullptr) { details::ObjectUnsafe::DecRefObjectHandle(this->extra_context); } } /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIError; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIError, ErrorObj, Object); /// \endcond }; namespace details { class ErrorObjFromStd : public ErrorObj { public: ErrorObjFromStd(std::string kind, std::string message, std::string backtrace) : kind_data_(std::move(kind)), message_data_(std::move(message)), backtrace_data_(std::move(backtrace)) { this->kind = TVMFFIByteArray{kind_data_.data(), kind_data_.length()}; this->message = TVMFFIByteArray{message_data_.data(), message_data_.length()}; this->backtrace = TVMFFIByteArray{backtrace_data_.data(), backtrace_data_.length()}; this->update_backtrace = UpdateBacktrace; } private: /*! * \brief Update the backtrace of the error object. * \param backtrace The backtrace to update. * \param update_mode The mode to update the backtrace, * can be either kTVMFFIBacktraceUpdateModeReplace, kTVMFFIBacktraceUpdateModeAppend. */ static void UpdateBacktrace(TVMFFIObjectHandle self, const TVMFFIByteArray* backtrace_str, int32_t update_mode) { ErrorObjFromStd* obj = static_cast(self); if (update_mode == kTVMFFIBacktraceUpdateModeReplace) { obj->backtrace_data_.resize(backtrace_str->size); std::memcpy(obj->backtrace_data_.data(), backtrace_str->data, backtrace_str->size); obj->backtrace = TVMFFIByteArray{obj->backtrace_data_.data(), obj->backtrace_data_.length()}; } else { obj->backtrace_data_.append(backtrace_str->data, backtrace_str->size); obj->backtrace = TVMFFIByteArray{obj->backtrace_data_.data(), obj->backtrace_data_.length()}; } } std::string kind_data_; std::string message_data_; std::string backtrace_data_; }; } // namespace details /*! * \brief Managed reference to ErrorObj * \sa Error Object */ class Error : public ObjectRef, public std::exception { public: /*! * \brief Constructor * \param kind The kind of the error. * \param message The message of the error. * \param backtrace The backtrace of the error. */ Error(std::string kind, std::string message, std::string backtrace) { data_ = make_object(std::move(kind), std::move(message), std::move(backtrace)); } /*! * \brief Constructor * \param kind The kind of the error. * \param message The message of the error. * \param backtrace The backtrace of the error. * \param cause_chain The cause chain of the error. * \param extra_context The extra context of the error. */ Error(std::string kind, std::string message, std::string backtrace, std::optional cause_chain, std::optional extra_context) { ObjectPtr error_obj = make_object( std::move(kind), std::move(message), std::move(backtrace)); if (cause_chain.has_value()) { error_obj->cause_chain = details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(*std::move(cause_chain)); } if (extra_context.has_value()) { error_obj->extra_context = details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(*std::move(extra_context)); } data_ = std::move(error_obj); } /*! * \brief Constructor * \param kind The kind of the error. * \param message The message of the error. * \param backtrace The backtrace of the error. */ Error(std::string kind, std::string message, const TVMFFIByteArray* backtrace) : Error(std::move(kind), std::move(message), std::string(backtrace->data, backtrace->size)) {} /*! * \brief Get the kind of the error object. * \return The kind of the error object. */ std::string kind() const { ErrorObj* obj = static_cast(data_.get()); return std::string(obj->kind.data, obj->kind.size); } /*! * \brief Get the message of the error object. * \return The message of the error object. */ std::string message() const { ErrorObj* obj = static_cast(data_.get()); return std::string(obj->message.data, obj->message.size); } /*! * \brief Get the cause chain of the error object. * \return The cause chain of the error object. */ std::optional cause_chain() const { ErrorObj* obj = static_cast(data_.get()); if (obj->cause_chain != nullptr) { return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned( static_cast(obj->cause_chain))); } else { return std::nullopt; } } /*! * \brief Get the extra context of the error object. * \return The extra context of the error object. */ std::optional extra_context() const { ErrorObj* obj = static_cast(data_.get()); if (obj->extra_context != nullptr) { return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned( static_cast(obj->extra_context))); } else { return std::nullopt; } } /*! * \brief Get the backtrace of the error object. * \return The backtrace of the error object. * \note Consider use TracebackMostRecentCallLast for pythonic style traceback. * * \sa TracebackMostRecentCallLast */ std::string backtrace() const { ErrorObj* obj = static_cast(data_.get()); return std::string(obj->backtrace.data, obj->backtrace.size); } /*! * \brief Get the traceback in the order of most recent call last. * * \return The traceback of the error object. */ std::string TracebackMostRecentCallLast() const { // add placeholder for the first line std::vector line_breakers = {-1}; ErrorObj* obj = static_cast(data_.get()); for (size_t i = 0; i < obj->backtrace.size; i++) { if (obj->backtrace.data[i] == '\n') { line_breakers.push_back(static_cast(i)); } } std::string result; result.reserve(obj->backtrace.size); for (size_t i = line_breakers.size() - 1; i > 0; --i) { int64_t line_start = line_breakers[i - 1] + 1; int64_t line_end = line_breakers[i]; if (line_start == line_end) continue; result.append(obj->backtrace.data + line_start, line_end - line_start); result.append("\n"); } return result; } /*! * \brief Update the backtrace of the error object. * \param backtrace_str The backtrace to update. * \param update_mode The mode to update the backtrace, * can be either kTVMFFIBacktraceUpdateModeReplace, kTVMFFIBacktraceUpdateModeAppend. */ void UpdateBacktrace(const TVMFFIByteArray* backtrace_str, int32_t update_mode) { ErrorObj* obj = static_cast(data_.get()); obj->update_backtrace(obj, backtrace_str, update_mode); } /*! * \brief Get the full message of the error, including kind, message and traceback. * \return The full message of the error object. */ std::string FullMessage() const { ErrorObj* obj = static_cast(data_.get()); return (std::string("Traceback (most recent call last):\n") + TracebackMostRecentCallLast() + std::string(obj->kind.data, obj->kind.size) + std::string(": ") + std::string(obj->message.data, obj->message.size) + '\n'); } /*! * \brief Get the error message * \return The error message * \note To get the full message including kind and traceback, use FullMessage() instead. */ const char* what() const noexcept(true) override { ErrorObj* obj = static_cast(data_.get()); return obj->message.data; } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Error, ObjectRef, ErrorObj); /// \endcond }; /*! * \brief Error object for EnvErrorAlreadySet * * This error can be thrown by EnvCheckSignals to indicate * that there is an error set in the frontend environment(e.g. * python interpreter). * * \code{.cpp} * void ExampleLongRunningFunction() { * if (TVMFFIEnvCheckSignals() != 0) { * throw ::tvm::ffi::EnvErrorAlreadySet(); * } * // do work here * } * \endcode */ inline Error EnvErrorAlreadySet() { return Error("EnvErrorAlreadySet", "", ""); } namespace details { /*! * \brief Move the last raised safe-call error from TLS. * \return The raised error object. */ TVM_FFI_INLINE Error MoveFromSafeCallRaised() { TVMFFIObjectHandle handle; TVMFFIErrorMoveFromRaised(&handle); return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(handle))); } /*! * \brief Set a raised safe-call error into TLS. * \param error The error to be raised. */ TVM_FFI_INLINE void SetSafeCallRaised(const Error& error) { TVMFFIErrorSetRaised(details::ObjectUnsafe::TVMFFIObjectPtrFromObjectRef(error)); } class ErrorBuilder { public: TVM_FFI_COLD_CODE explicit ErrorBuilder(std::string kind, std::string backtrace, bool log_before_throw) : kind_(std::move(kind)), backtrace_(std::move(backtrace)), log_before_throw_(log_before_throw) {} TVM_FFI_COLD_CODE explicit ErrorBuilder(std::string kind, const TVMFFIByteArray* backtrace, bool log_before_throw) : ErrorBuilder(std::move(kind), std::string(backtrace->data, backtrace->size), log_before_throw) {} TVM_FFI_COLD_CODE explicit ErrorBuilder(std::string kind, const TVMFFIByteArray* backtrace, bool log_before_throw, std::optional cause_chain, std::optional extra_context) : ErrorBuilder(std::move(kind), backtrace, log_before_throw) { cause_chain_ = std::move(cause_chain); extra_context_ = std::move(extra_context); } // MSVC disable warning in error builder as it is exepected #ifdef _MSC_VER #pragma warning(push) #pragma warning(disable : 4722) #endif // avoid inline to reduce binary size, error throw path do not need to be fast [[noreturn]] TVM_FFI_COLD_CODE ~ErrorBuilder() noexcept(false) { ::tvm::ffi::Error error(std::move(kind_), stream_.str(), std::move(backtrace_), std::move(cause_chain_), std::move(extra_context_)); if (log_before_throw_) { std::cerr << error.FullMessage(); } throw error; } #ifdef _MSC_VER #pragma warning(pop) #endif std::ostringstream& stream() { return stream_; } protected: std::string kind_; std::ostringstream stream_; std::string backtrace_; bool log_before_throw_; std::optional cause_chain_; std::optional extra_context_; }; } // namespace details /*! * \brief Helper macro to throw an error with backtrace and message * * \code{.cpp} * void ThrowError() { * TVM_FFI_THROW(RuntimeError) << "error message"; * } * \endcode */ #define TVM_FFI_THROW(ErrorKind) \ ::tvm::ffi::details::ErrorBuilder(#ErrorKind, \ TVMFFIBacktrace(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0), \ TVM_FFI_ALWAYS_LOG_BEFORE_THROW) \ .stream() /*! * \brief Explicitly log error in stderr and then throw the error. * * \note This is only necessary on startup functions where we know error * cannot be caught, and it is better to have a clear log message. * In most cases, we should use use TVM_FFI_THROW. */ #define TVM_FFI_LOG_AND_THROW(ErrorKind) \ ::tvm::ffi::details::ErrorBuilder( \ #ErrorKind, TVMFFIBacktrace(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0), true) \ .stream() // Glog style checks with TVM_FFI prefix // NOTE: we explicitly avoid glog style generic macros (LOG/CHECK) in tvm ffi // to avoid potential conflict of downstream users who might have their own GLOG style macros namespace details { template TVM_FFI_INLINE std::unique_ptr LogCheckFormat(const X& x, const Y& y) { std::ostringstream os; os << " (" << x << " vs. " << y << ") "; // CHECK_XX(x, y) requires x and y can be serialized to // string. Use CHECK(x OP y) otherwise. return std::make_unique(os.str()); } #define TVM_FFI_CHECK_FUNC(name, op) \ template \ TVM_FFI_INLINE std::unique_ptr LogCheck##name(const X& x, const Y& y) { \ if (x op y) return nullptr; \ return LogCheckFormat(x, y); \ } \ TVM_FFI_INLINE std::unique_ptr LogCheck##name(int x, int y) { \ return LogCheck##name(x, y); \ } // Inline _Pragma in macros does not work reliably on old version of MSVC and // GCC. We wrap all comparisons in a function so that we can use #pragma to // silence bad comparison warnings. #if defined(__GNUC__) || defined(__clang__) // GCC and Clang #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wsign-compare" #elif defined(_MSC_VER) // MSVC #pragma warning(push) #pragma warning(disable : 4389) // '==' : signed/unsigned mismatch #endif TVM_FFI_CHECK_FUNC(_LT, <) TVM_FFI_CHECK_FUNC(_GT, >) TVM_FFI_CHECK_FUNC(_LE, <=) TVM_FFI_CHECK_FUNC(_GE, >=) TVM_FFI_CHECK_FUNC(_EQ, ==) TVM_FFI_CHECK_FUNC(_NE, !=) #if defined(__GNUC__) || defined(__clang__) // GCC and Clang #pragma GCC diagnostic pop #elif defined(_MSC_VER) // MSVC #pragma warning(pop) #endif } // namespace details #define TVM_FFI_CHECK_BINARY_OP(name, op, x, y, ErrorKind) \ if (auto __tvm_ffi_log_err = /* NOLINT(bugprone-reserved-identifier) */ \ ::tvm::ffi::details::LogCheck##name(x, y)) \ TVM_FFI_THROW(ErrorKind) << "Check failed: " << #x " " #op " " #y << *__tvm_ffi_log_err << ": " #define TVM_FFI_CHECK(cond, ErrorKind) \ if (TVM_FFI_PREDICT_FALSE(!(cond))) \ TVM_FFI_THROW(ErrorKind) << "Check failed: (" #cond << ") is false: " #define TVM_FFI_CHECK_LT(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_LT, <, x, y, ErrorKind) #define TVM_FFI_CHECK_GT(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_GT, >, x, y, ErrorKind) #define TVM_FFI_CHECK_LE(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_LE, <=, x, y, ErrorKind) #define TVM_FFI_CHECK_GE(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_GE, >=, x, y, ErrorKind) #define TVM_FFI_CHECK_EQ(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_EQ, ==, x, y, ErrorKind) #define TVM_FFI_CHECK_NE(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_NE, !=, x, y, ErrorKind) #define TVM_FFI_CHECK_NOTNULL(x, ErrorKind) \ ((x) == nullptr ? TVM_FFI_THROW(ErrorKind) << "Check not null: " #x << ' ', \ (x) : (x)) // NOLINT(*) #define TVM_FFI_ICHECK(x) TVM_FFI_CHECK(x, InternalError) #define TVM_FFI_ICHECK_LT(x, y) TVM_FFI_CHECK_LT(x, y, InternalError) #define TVM_FFI_ICHECK_GT(x, y) TVM_FFI_CHECK_GT(x, y, InternalError) #define TVM_FFI_ICHECK_LE(x, y) TVM_FFI_CHECK_LE(x, y, InternalError) #define TVM_FFI_ICHECK_GE(x, y) TVM_FFI_CHECK_GE(x, y, InternalError) #define TVM_FFI_ICHECK_EQ(x, y) TVM_FFI_CHECK_EQ(x, y, InternalError) #define TVM_FFI_ICHECK_NE(x, y) TVM_FFI_CHECK_NE(x, y, InternalError) #define TVM_FFI_ICHECK_NOTNULL(x) TVM_FFI_CHECK_NOTNULL(x, InternalError) // Debug checks: same as ICHECK but stripped when NDEBUG is defined (release builds). #ifndef NDEBUG #define TVM_FFI_DCHECK(x) TVM_FFI_ICHECK(x) #define TVM_FFI_DCHECK_LT(x, y) TVM_FFI_ICHECK_LT(x, y) #define TVM_FFI_DCHECK_GT(x, y) TVM_FFI_ICHECK_GT(x, y) #define TVM_FFI_DCHECK_LE(x, y) TVM_FFI_ICHECK_LE(x, y) #define TVM_FFI_DCHECK_GE(x, y) TVM_FFI_ICHECK_GE(x, y) #define TVM_FFI_DCHECK_EQ(x, y) TVM_FFI_ICHECK_EQ(x, y) #define TVM_FFI_DCHECK_NE(x, y) TVM_FFI_ICHECK_NE(x, y) #define TVM_FFI_DCHECK_NOTNULL(x) TVM_FFI_ICHECK_NOTNULL(x) #else #define TVM_FFI_DCHECK(x) \ while (false) TVM_FFI_ICHECK(x) #define TVM_FFI_DCHECK_LT(x, y) \ while (false) TVM_FFI_ICHECK_LT(x, y) #define TVM_FFI_DCHECK_GT(x, y) \ while (false) TVM_FFI_ICHECK_GT(x, y) #define TVM_FFI_DCHECK_LE(x, y) \ while (false) TVM_FFI_ICHECK_LE(x, y) #define TVM_FFI_DCHECK_GE(x, y) \ while (false) TVM_FFI_ICHECK_GE(x, y) #define TVM_FFI_DCHECK_EQ(x, y) \ while (false) TVM_FFI_ICHECK_EQ(x, y) #define TVM_FFI_DCHECK_NE(x, y) \ while (false) TVM_FFI_ICHECK_NE(x, y) #define TVM_FFI_DCHECK_NOTNULL(x) (x) #endif // NDEBUG } // namespace ffi } // namespace tvm #endif // TVM_FFI_ERROR_H_ tvm-ffi-0.1.12/include/tvm/ffi/expected.h000066400000000000000000000207131521067262500201230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/expected.h * \brief Runtime Expected container type for exception-free error handling. */ #ifndef TVM_FFI_EXPECTED_H_ #define TVM_FFI_EXPECTED_H_ #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Wrapper to explicitly construct an Expected in the error state. * \tparam E The error type, must derive from Error. */ template class Unexpected { static_assert(std::is_base_of_v>, "Unexpected requires E to be Error or a subclass of Error."); public: /*! \brief Construct from an error value. */ explicit Unexpected(E error) : error_(std::move(error)) {} /*! \brief Access the stored error. */ const E& error() const& noexcept { return error_; } /*! \brief Access the stored error. */ E& error() & noexcept { return error_; } /*! \brief Access the stored error (rvalue). */ const E&& error() const&& noexcept { return std::move(error_); } /*! \brief Access the stored error (rvalue). */ E&& error() && noexcept { return std::move(error_); } private: E error_; }; #ifndef TVM_FFI_DOXYGEN_MODE template Unexpected(E) -> Unexpected; #endif /*! * \brief Expected provides exception-free error handling for FFI functions. * * Expected is similar to Rust's Result or C++23's std::expected. * It can hold either a success value of type T or an error of type Error. * * \tparam T The success type. Must be Any-compatible and cannot be Error. * * Usage: * \code * Expected divide(int a, int b) { * if (b == 0) { * return Error("ValueError", "Division by zero"); * } * return a / b; * } * * Expected result = divide(10, 2); * if (result.is_ok()) { * int value = result.value(); * } else { * Error err = result.error(); * } * \endcode */ template class Expected { public: static_assert(!std::is_same_v, "Expected is not allowed. Use Error directly."); /*! * \brief Implicit constructor from a success value. * \param value The success value. */ // NOLINTNEXTLINE(google-explicit-constructor,runtime/explicit) Expected(T value) : data_(Any(std::move(value))) {} /*! * \brief Implicit constructor from an error. * \param error The error value. */ // NOLINTNEXTLINE(google-explicit-constructor,runtime/explicit) Expected(Error error) : data_(Any(std::move(error))) {} /*! \brief Implicit constructor from an Unexpected wrapper. */ template >>> // NOLINTNEXTLINE(google-explicit-constructor,runtime/explicit) Expected(Unexpected unexpected) : data_(Any(std::move(unexpected).error())) {} /*! \brief Returns true if the Expected contains a success value. */ TVM_FFI_INLINE bool is_ok() const noexcept { return data_.type_index() != TypeIndex::kTVMFFIError; } /*! \brief Returns true if the Expected contains an error. */ TVM_FFI_INLINE bool is_err() const noexcept { return data_.type_index() == TypeIndex::kTVMFFIError; } /*! \brief Alias for is_ok(). */ TVM_FFI_INLINE bool has_value() const noexcept { return is_ok(); } /*! \brief Returns the success value, or throws the contained error. */ TVM_FFI_INLINE T value() const& { if (TVM_FFI_PREDICT_TRUE(is_ok())) { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); } throw details::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); } /*! \brief Returns the success value (moved out), or throws the contained error. */ TVM_FFI_INLINE T value() && { if (TVM_FFI_PREDICT_TRUE(is_ok())) { return details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(data_)); } throw details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(data_)); } /*! \brief Returns the contained error, or throws RuntimeError if is_ok(). */ TVM_FFI_INLINE Error error() const& { // No branch hint: error() is itself a cold path — callers only invoke it // after observing !is_ok(), so the branch direction here doesn't matter. if (is_ok()) { TVM_FFI_THROW(RuntimeError) << "Bad expected access: contains value, not error"; } return details::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); } /*! \brief Returns the contained error (moved out), or throws RuntimeError if is_ok(). */ TVM_FFI_INLINE Error error() && { // No branch hint: error() is itself a cold path — callers only invoke it // after observing !is_ok(), so the branch direction here doesn't matter. if (is_ok()) { TVM_FFI_THROW(RuntimeError) << "Bad expected access: contains value, not error"; } return details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(data_)); } /*! * \brief Returns the success value, or \p default_value if the Expected holds an error. */ template > TVM_FFI_INLINE T value_or(U&& default_value) const& { if (TVM_FFI_PREDICT_TRUE(is_ok())) { return details::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); } return T(std::forward(default_value)); } /*! * \brief Returns the success value (moved out), or \p default_value if the Expected holds an * error. */ template > TVM_FFI_INLINE T value_or(U&& default_value) && { if (TVM_FFI_PREDICT_TRUE(is_ok())) { return details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(data_)); } return T(std::forward(default_value)); } private: Any data_; // Invariant: holds a T (type_index != kTVMFFIError) or an Error. }; // TypeTraits specialization for Expected template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public TypeTraitsBase { TVM_FFI_INLINE static void CopyToAnyView(const Expected& src, TVMFFIAny* result) { if (src.is_err()) { TypeTraits::CopyToAnyView(src.error(), result); } else { TypeTraits::CopyToAnyView(src.value(), result); } } TVM_FFI_INLINE static void MoveToAny(Expected src, TVMFFIAny* result) { if (src.is_err()) { TypeTraits::MoveToAny(std::move(src).error(), result); } else { TypeTraits::MoveToAny(std::move(src).value(), result); } } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return TypeTraits::CheckAnyStrict(src) || TypeTraits::CheckAnyStrict(src); } TVM_FFI_INLINE static Expected CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { if (TypeTraits::CheckAnyStrict(src)) { return TypeTraits::CopyFromAnyViewAfterCheck(src); } return TypeTraits::CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static Expected MoveFromAnyAfterCheck(TVMFFIAny* src) { if (TypeTraits::CheckAnyStrict(src)) { return TypeTraits::MoveFromAnyAfterCheck(src); } return TypeTraits::MoveFromAnyAfterCheck(src); } TVM_FFI_INLINE static std::optional> TryCastFromAnyView(const TVMFFIAny* src) { if (auto opt = TypeTraits::TryCastFromAnyView(src)) { return Expected(*std::move(opt)); } if (auto opt_err = TypeTraits::TryCastFromAnyView(src)) { return Expected(*std::move(opt_err)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return "Expected<" + TypeTraits::TypeStr() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"Expected","args":[)" + details::TypeSchema::v() + R"(,{"type":"ffi.Error"}]})"; } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXPECTED_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/000077500000000000000000000000001521067262500172715ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/extra/base.h000066400000000000000000000033251521067262500203570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/base.h * \brief Base header for Extra API. * * The extra APIs contains a minmal set of extra APIs that are not * required to support essential core functionality. */ #ifndef TVM_FFI_EXTRA_BASE_H_ #define TVM_FFI_EXTRA_BASE_H_ #include /*! * \brief Marks the API as extra c++ api that is defined in cc files. * * They are implemented in cc files to reduce compile-time overhead. * The input/output only uses POD/Any/ObjectRef for ABI stability. * However, these extra APIs may have an issue across MSVC/Itanium ABI, * * Related features are also available through reflection based function * that is fully based on C API * * The project aims to minimize the number of extra C++ APIs to keep things * lightweight and restrict the use to non-core functionalities. */ #ifndef TVM_FFI_EXTRA_CXX_API #define TVM_FFI_EXTRA_CXX_API TVM_FFI_DLL #endif #endif // TVM_FFI_EXTRA_BASE_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/base64.h000066400000000000000000000132521521067262500205310ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * * \file tvm/ffi/extra/base64.h * \brief Base64 encoding and decoding utilities */ #ifndef TVM_FFI_EXTRA_BASE64_H_ #define TVM_FFI_EXTRA_BASE64_H_ #include #include namespace tvm { namespace ffi { /*! * \brief Encode a byte array into a base64 string * \param bytes The byte array to encode * \return The base64 encoded string */ inline String Base64Encode(TVMFFIByteArray bytes) { // encoding every 3 bytes into 4 characters constexpr const char kEncodeTable[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; std::string encoded; encoded.reserve(4 * (bytes.size + 2) / 3); for (size_t i = 0; i < (bytes.size / 3) * 3; i += 3) { int32_t buf[3]; buf[0] = static_cast(static_cast(bytes.data[i])); buf[1] = static_cast(static_cast(bytes.data[i + 1])); buf[2] = static_cast(static_cast(bytes.data[i + 2])); encoded.push_back(kEncodeTable[buf[0] >> 2]); encoded.push_back(kEncodeTable[((buf[0] << 4) | (buf[1] >> 4)) & 0x3F]); encoded.push_back(kEncodeTable[((buf[1] << 2) | (buf[2] >> 6)) & 0x3F]); encoded.push_back(kEncodeTable[buf[2] & 0x3F]); } if (bytes.size % 3 == 1) { int32_t buf[1] = {static_cast(static_cast(bytes.data[bytes.size - 1]))}; encoded.push_back(kEncodeTable[buf[0] >> 2]); encoded.push_back(kEncodeTable[(buf[0] << 4) & 0x3F]); encoded.push_back('='); encoded.push_back('='); } else if (bytes.size % 3 == 2) { int32_t buf[2] = {static_cast(static_cast(bytes.data[bytes.size - 2])), static_cast(static_cast(bytes.data[bytes.size - 1]))}; encoded.push_back(kEncodeTable[buf[0] >> 2]); encoded.push_back(kEncodeTable[((buf[0] << 4) | (buf[1] >> 4)) & 0x3F]); encoded.push_back(kEncodeTable[(buf[1] << 2) & 0x3F]); encoded.push_back('='); } return String(encoded); } /*! * \brief Encode a bytes object into a base64 string * \param data The bytes object to encode * \return The base64 encoded string */ inline String Base64Encode(const Bytes& data) { return Base64Encode(TVMFFIByteArray{data.data(), data.size()}); } /*! * \brief Decode a base64 string into a byte array * \param bytes The bytes to be decoded * \return The decoded byte array */ inline Bytes Base64Decode(TVMFFIByteArray bytes) { constexpr const char kDecodeTable[] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 62, // '+' 0, 0, 0, 63, // '/' 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, // '0'-'9' 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, // 'A'-'Z' 0, 0, 0, 0, 0, 0, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, // 'a'-'z' }; std::string decoded; decoded.reserve(bytes.size * 3 / 4); if (bytes.size == 0) return Bytes(); TVM_FFI_ICHECK(bytes.size % 4 == 0) << "invalid base64 encoding"; // leverage this property to simplify decoding static_assert('=' < sizeof(kDecodeTable) && kDecodeTable[static_cast('=')] == 0); // base64 is always multiple of 4 bytes for (size_t i = 0; i < bytes.size; i += 4) { // decode every 4 characters into 24bits, each character contains 6 bits // note that = is also decoded as 0, which is safe to skip int32_t buf[4] = { static_cast(static_cast(bytes.data[i])), static_cast(static_cast(bytes.data[i + 1])), static_cast(static_cast(bytes.data[i + 2])), static_cast(static_cast(bytes.data[i + 3])), }; int32_t value_i24 = (static_cast(kDecodeTable[buf[0]]) << 18) | (static_cast(kDecodeTable[buf[1]]) << 12) | (static_cast(kDecodeTable[buf[2]]) << 6) | static_cast(kDecodeTable[buf[3]]); // unpack 24bits into 3 bytes, each contains 8 bits decoded.push_back(static_cast((value_i24 >> 16) & 0xFF)); if (buf[2] != '=') { decoded.push_back(static_cast((value_i24 >> 8) & 0xFF)); } if (buf[3] != '=') { decoded.push_back(static_cast(value_i24 & 0xFF)); } } return Bytes(decoded); } /*! * \brief Decode a base64 string into a byte array * \param data The base64 encoded string to decode * \return The decoded byte array */ inline Bytes Base64Decode(const String& data) { return Base64Decode(TVMFFIByteArray{data.data(), data.size()}); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_BASE64_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/c_env_api.h000066400000000000000000000140161521067262500213670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOLINTBEGIN(modernize-use-using) /*! * \file tvm/ffi/extra/c_env_api.h * \brief Extra environment API. */ #ifndef TVM_FFI_EXTRA_C_ENV_API_H_ #define TVM_FFI_EXTRA_C_ENV_API_H_ #include #ifdef __cplusplus extern "C" { #endif // ---------------------------------------------------------------------------- // Stream context // Focusing on minimalistic thread-local context recording stream being used. // We explicitly not handle allocation/de-allocation of stream here. // ---------------------------------------------------------------------------- /*! * \brief The type of the stream handle. */ typedef void* TVMFFIStreamHandle; /*! * \brief FFI function to set the current stream for a device * * \param device_type The type of the device. * \param device_id The id of the device. * \param stream The stream to set. * \param opt_out_original_stream Output original stream if the address is not nullptr. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIEnvSetStream(int32_t device_type, int32_t device_id, TVMFFIStreamHandle stream, TVMFFIStreamHandle* opt_out_original_stream); /*! * \brief FFI function to get the current stream for a device * * \param device_type The type of the device. * \param device_id The id of the device. * \return The current stream of the device. */ TVM_FFI_DLL TVMFFIStreamHandle TVMFFIEnvGetStream(int32_t device_type, int32_t device_id); /*! * \brief Set the current DLPackManagedTensorAllocator in thread-local(TLS) context * * \param allocator The allocator to set. * \param write_to_global_context Whether to also set the allocator to the global context. * \param opt_out_original_allocator Output original TLS allocator if the address is not nullptr. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIEnvSetDLPackManagedTensorAllocator( DLPackManagedTensorAllocator allocator, int write_to_global_context, DLPackManagedTensorAllocator* opt_out_original_allocator); /*! * \brief FFI function get the current DLPackManagedTensorAllocator stored in context. * * This function first queries the global context, and if not found, * queries the thread-local context. * * \return The current setted DLPackManagedTensorAllocator */ TVM_FFI_DLL DLPackManagedTensorAllocator TVMFFIEnvGetDLPackManagedTensorAllocator(); /*! * \brief Allocate a tensor from the allocator set in thread-local(TLS) context. * * This function redirects to one of environment allocator. As of now, we only * support the DLPackManagedTensorAllocator set in thread-local(TLS) context. * * \param prototype The prototype DLTensor, only the dtype, ndim, shape, * and device fields are used, other fields are ignored. * \param out The output tensor in kTVMFFITensor type. * \return 0 when success, nonzero when failure happens * \sa TVMFFIEnvSetDLPackManagedTensorAllocator */ TVM_FFI_DLL int TVMFFIEnvTensorAlloc(DLTensor* prototype, TVMFFIObjectHandle* out); /*! * \brief Check if there are any signals raised in the surrounding env. * \return 0 when success, nonzero when failure happens * \note Under python this function redirects to PyErr_CheckSignals */ TVM_FFI_DLL int TVMFFIEnvCheckSignals(); /*! * \brief Register a symbol into the from the surrounding env such as python * \param name The name of the symbol. * \param symbol The symbol to register. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIEnvRegisterCAPI(const char* name, void* symbol); // ---------------------------------------------------------------------------- // Module symbol management in callee side // ---------------------------------------------------------------------------- /*! * \brief FFI function to lookup a function from a module's imports. * * This is a helper function that is used by generated code. * * \param library_ctx The library context module handle. * \param func_name The name of the function. * \param out The result function. * \note The returned function is a weak reference that is cached/owned by the module. * \return 0 when no error is thrown, -1 when failure happens */ TVM_FFI_DLL int TVMFFIEnvModLookupFromImports(TVMFFIObjectHandle library_ctx, const char* func_name, TVMFFIObjectHandle* out); /*! * \brief Register a symbol value that will be initialized when a library with the symbol is loaded. * * This function can be used to make context functions to be available in the library * module that wants to avoid an explicit link dependency * * \param name The name of the symbol. * \param symbol The symbol to register. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIEnvModRegisterContextSymbol(const char* name, void* symbol); /*! * \brief Register a symbol that will be initialized when a system library is loaded. * * \param name The name of the symbol. * \param symbol The symbol to register. * \return 0 when success, nonzero when failure happens */ TVM_FFI_DLL int TVMFFIEnvModRegisterSystemLibSymbol(const char* name, void* symbol); #ifdef __cplusplus } // extern "C" #endif #endif // TVM_FFI_EXTRA_C_ENV_API_H_ // NOLINTEND(modernize-use-using) tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/000077500000000000000000000000001521067262500202055ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/base.h000066400000000000000000000063031521067262500212720ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/cuda/base.h * \brief CUDA base utilities. */ #ifndef TVM_FFI_EXTRA_CUDA_BASE_H_ #define TVM_FFI_EXTRA_CUDA_BASE_H_ #include #include namespace tvm { namespace ffi { /*! * \brief Macro for checking CUDA runtime API errors. * * This macro checks the return value of CUDA runtime API calls and throws * a RuntimeError with detailed error information if the call fails. * * \param stmt The CUDA runtime API call to check. */ #define TVM_FFI_CHECK_CUDA_ERROR(stmt) \ do { \ cudaError_t __err = (stmt); \ if (__err != cudaSuccess) { \ const char* __err_name = cudaGetErrorName(__err); \ const char* __err_str = cudaGetErrorString(__err); \ TVM_FFI_THROW(RuntimeError) << "CUDA Runtime Error: " << __err_name << " (" \ << static_cast(__err) << "): " << __err_str; \ } \ } while (0) /*! * \brief A simple 3D dimension type for CUDA kernel launch configuration. * * This struct mimics the behavior of dim3 from CUDA Runtime API and provides * a compatible interface for kernel launch configuration. It can be constructed * from 1, 2, or 3 dimensions. */ struct dim3 { /*! \brief X dimension (number of blocks in x-direction or threads in x-direction) */ unsigned int x; /*! \brief Y dimension (number of blocks in y-direction or threads in y-direction) */ unsigned int y; /*! \brief Z dimension (number of blocks in z-direction or threads in z-direction) */ unsigned int z; /*! \brief Default constructor initializes to (1, 1, 1) */ dim3() : x(1), y(1), z(1) {} /*! \brief Construct with x dimension, y and z default to 1 */ explicit dim3(unsigned int x_) : x(x_), y(1), z(1) {} /*! \brief Construct with x and y dimensions, z defaults to 1 */ dim3(unsigned int x_, unsigned int y_) : x(x_), y(y_), z(1) {} /*! \brief Construct with all three dimensions */ dim3(unsigned int x_, unsigned int y_, unsigned int z_) : x(x_), y(y_), z(z_) {} }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_CUDA_BASE_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/cubin_launcher.h000066400000000000000000000547451521067262500233560ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/cuda/cubin_launcher.h * \brief CUDA CUBIN launcher utility for loading and executing CUDA kernels. * * This header provides a lightweight C++ wrapper around CUDA Runtime API * for loading CUBIN modules and launching kernels. It supports: * - Loading CUBIN from memory (embedded data) * - Multi-GPU execution using CUDA primary contexts * - Kernel parameter management and launch configuration */ #ifndef TVM_FFI_EXTRA_CUDA_CUBIN_LAUNCHER_H_ #define TVM_FFI_EXTRA_CUDA_CUBIN_LAUNCHER_H_ #include // NOLINT(clang-diagnostic-error) #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Macro to embed a CUBIN module with static initialization. * * This macro declares external symbols for embedded CUBIN data and creates * a singleton struct to manage the CubinModule instance. The CUBIN data * symbols should be named `__tvm_ffi__cubin_` and `__tvm_ffi__cubin__end`, * typically created using objcopy and ld. * * \par Creating Embedded CUBIN with TVM-FFI Utilities * TVM-FFI provides utilities to simplify CUBIN embedding. You have two options: * * \par Option 1: CMake Utility (Recommended) * Use the `tvm_ffi_embed_cubin` CMake function: * \code{.cmake} * # Find tvm_ffi package (provides tvm_ffi_embed_cubin utility) * find_package(tvm_ffi CONFIG REQUIRED) * find_package(CUDAToolkit REQUIRED) * * # Compile CUDA kernel to CUBIN * tvm_ffi_generate_cubin( * OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/kernel.cubin * SOURCE src/kernel.cu * ARCH native # or sm_75, sm_80, etc. * ) * * # Embed CUBIN into C++ object file * tvm_ffi_embed_cubin( * OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/mycode_with_cubin.o * SOURCE src/mycode.cc * CUBIN ${CMAKE_CURRENT_BINARY_DIR}/kernel.cubin * NAME my_kernels # Must match TVM_FFI_EMBED_CUBIN(my_kernels) in code * ) * * # Link into shared library * add_library(mylib SHARED ${CMAKE_CURRENT_BINARY_DIR}/mycode_with_cubin.o) * target_link_libraries(mylib PRIVATE tvm_ffi_header CUDA::cudart) * \endcode * * \par Option 2: Python Utility * Use the `tvm_ffi.utils.embed_cubin` command-line tool: * \code{.bash} * # Step 1: Compile CUDA kernel to CUBIN * nvcc --cubin -arch=sm_75 kernel.cu -o kernel.cubin * * # Step 2: Compile C++ source to object file * g++ -c -fPIC -std=c++17 -I/path/to/tvm-ffi/include mycode.cc -o mycode.o * * # Step 3: Embed CUBIN using Python utility * python -m tvm_ffi.utils.embed_cubin \ * --output-obj mycode_with_cubin.o \ * --input-obj mycode.o \ * --cubin kernel.cubin \ * --name my_kernels * * # Step 4: Link into shared library * g++ -o mylib.so -shared mycode_with_cubin.o -lcudart * \endcode * * The utilities automatically handle: * - Symbol renaming to __tvm_ffi__cubin_ format * - Adding .note.GNU-stack section for security * - Symbol localization to prevent conflicts * * \par Usage in C++ Code * In your C++ source file, use the embedded CUBIN: * \code{.cpp} * #include * * // Declare the embedded CUBIN module (name must match CMake NAME parameter) * TVM_FFI_EMBED_CUBIN(my_kernels); * * void MyFunction() { * // Get kernel from embedded CUBIN (cached in static variable for efficiency) * static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "my_kernel"); * // Use kernel... * } * \endcode * * \note CMake Setup: To use the utilities, add to your CMakeLists.txt: * \code{.cmake} * find_package(tvm_ffi CONFIG REQUIRED) # Provides tvm_ffi_embed_cubin utility * \endcode * * \par Option 3: Python Integration with load_inline * When using `tvm_ffi.cpp.load_inline()` with the `embed_cubin` parameter, * the CUBIN data is automatically embedded using the Python utility internally: * \code{.py} * from tvm_ffi import cpp * from tvm_ffi.cpp import nvrtc * * # Compile CUDA source to CUBIN * cubin_bytes = nvrtc.nvrtc_compile(cuda_source) * * # Load with embedded CUBIN - automatically handles embedding * mod = cpp.load_inline( * "my_module", * cuda_sources=cpp_code, * embed_cubin={"my_kernels": cubin_bytes}, * extra_ldflags=["-lcudart"] * ) * \endcode * * \param name The identifier for this embedded CUBIN module (must match the * symbol names created with objcopy or the key in embed_cubin dict). * * \see TVM_FFI_EMBED_CUBIN_GET_KERNEL * \see CubinModule * \see CubinKernel */ #define TVM_FFI_EMBED_CUBIN(name) \ extern "C" const char __tvm_ffi__cubin_##name[]; \ extern "C" const char __tvm_ffi__cubin_##name##_end[]; \ namespace { \ struct EmbedCubinModule_##name { \ tvm::ffi::CubinModule mod{__tvm_ffi__cubin_##name}; \ static EmbedCubinModule_##name* Global() { \ static EmbedCubinModule_##name inst; \ return &inst; \ } \ }; \ } /* anonymous namespace */ /*! * \brief Macro to load a CUBIN module from a byte array. * * This macro creates a singleton struct to manage the CubinModule instance * initialized from a byte array (e.g. from `#embed ` or bin2c output). * * \par Usage Example * \code{.cpp} * constexpr unsigned char image[] = { ... }; * TVM_FFI_EMBED_CUBIN_FROM_BYTES(my_kernels, image); * * void MyFunc() { * static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "kernel_name"); * } * \endcode * * \param name The identifier for this embedded CUBIN module. * \param imageBytes The byte array containing the CUBIN/FATBIN data. */ #define TVM_FFI_EMBED_CUBIN_FROM_BYTES(name, imageBytes) \ namespace { \ struct EmbedCubinModule_##name { \ tvm::ffi::CubinModule mod{imageBytes}; \ static EmbedCubinModule_##name* Global() { \ static EmbedCubinModule_##name inst; \ return &inst; \ } \ }; \ } /* anonymous namespace */ /*! * \brief Macro to get a kernel from an embedded CUBIN module. * * This macro retrieves a kernel by name from a previously declared embedded * CUBIN module (using TVM_FFI_EMBED_CUBIN). The result is a CubinKernel object * that can be used to launch the kernel with specified parameters. * * \par Performance Tip * It's recommended to store the result in a static variable to avoid repeated * kernel lookups, which improves performance: * \code{.cpp} * static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "kernel_name"); * \endcode * * \par Complete Example * \code{.cpp} * // Declare embedded CUBIN module * TVM_FFI_EMBED_CUBIN(my_kernels); * * void LaunchKernel(tvm::ffi::TensorView input, tvm::ffi::TensorView output) { * // Get kernel (cached in static variable for efficiency) * static auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "add_one"); * * // Prepare kernel arguments * void* in_ptr = input.data_ptr(); * void* out_ptr = output.data_ptr(); * int64_t n = input.size(0); * void* args[] = {&in_ptr, &out_ptr, &n}; * * // Configure launch * tvm::ffi::dim3 grid((n + 255) / 256); * tvm::ffi::dim3 block(256); * * // Get stream and launch * DLDevice device = input.device(); * cudaStream_t stream = static_cast( * TVMFFIEnvGetStream(device.device_type, device.device_id)); * * cudaError_t result = kernel.Launch(args, grid, block, stream); * TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); * } * \endcode * * \param name The identifier of the embedded CUBIN module (must match the name * used in TVM_FFI_EMBED_CUBIN). * \param kernel_name The name of the kernel function as it appears in the CUBIN * (typically the function name for `extern "C"` kernels). * \return A CubinKernel object for the specified kernel. * * \see TVM_FFI_EMBED_CUBIN * \see CubinKernel::Launch */ #define TVM_FFI_EMBED_CUBIN_GET_KERNEL(name, kernel_name) \ (EmbedCubinModule_##name::Global()->mod[kernel_name]) // Forward declaration class CubinKernel; /*! * \brief CUDA CUBIN module loader and manager. * * This class provides a RAII wrapper around CUDA Runtime API's library management. * It loads a CUBIN module from memory and manages the library handle automatically. * The library is unloaded when the CubinModule object is destroyed. * * \par Features * - Load CUBIN from memory (embedded data or runtime-generated) * - Automatic resource management (RAII pattern) * - Multi-GPU execution using CUDA primary contexts * - Retrieve multiple kernels from the same module * * \par Example Usage * \code{.cpp} * // Load CUBIN from memory * tvm::ffi::Bytes cubin_data = ...; * tvm::ffi::CubinModule module(cubin_data); * * // Get kernels by name * tvm::ffi::CubinKernel kernel1 = module["add_one"]; * tvm::ffi::CubinKernel kernel2 = module.GetKernel("mul_two"); * * // Launch kernels * void* args[] = {...}; * tvm::ffi::dim3 grid(32), block(256); * cudaStream_t stream = ...; * kernel1.Launch(args, grid, block, stream); * \endcode * * \note This class is movable but not copyable. * \see TVM_FFI_EMBED_CUBIN for embedding CUBIN at compile time * \see CubinKernel for kernel launching */ class CubinModule { public: /*! * \brief Load CUBIN module from memory. * * \param bytes CUBIN binary data as a Bytes object. */ explicit CubinModule(const Bytes& bytes) { TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(cuda_api::LoadLibrary(&library_, bytes.data())); } /*! * \brief Load CUBIN module from raw memory buffer. * * \param code Pointer to CUBIN binary data. * \note The `code` buffer points to an ELF image. */ explicit CubinModule(const char* code) { TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(cuda_api::LoadLibrary(&library_, code)); } /*! * \brief Load CUBIN module from raw memory buffer. * * \param code Pointer to CUBIN binary data. * \note The `code` buffer points to an ELF image. */ explicit CubinModule(const unsigned char* code) { TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(cuda_api::LoadLibrary(&library_, code)); } /*! \brief Destructor unloads the library */ ~CubinModule() { if (library_ != nullptr) { cuda_api::UnloadLibrary(library_); } } /*! * \brief Get a kernel function from the module by name. * * \param name Name of the kernel function. * \return CubinKernel object representing the loaded kernel. */ CubinKernel GetKernel(const char* name); /*! * \brief Get a kernel function from the module by name with maximum dynamic shared memory. * * \param name Name of the kernel function. * \param dynamic_smem_max Maximum dynamic shared memory in bytes to set for this kernel. * -1 (default) means maximum available dynamic shared memory * (device max - static shared memory used by kernel). * \return CubinKernel object representing the loaded kernel. */ CubinKernel GetKernelWithMaxDynamicSharedMemory(const char* name, int64_t dynamic_smem_max); /*! * \brief Operator[] for convenient kernel access. * * It's equivalent to calling GetKernel(name, -1). * * \param name Name of the kernel function. * \return CubinKernel object representing the loaded kernel. */ CubinKernel operator[](const char* name); /*! \brief Get the underlying cudaLibrary_t handle */ cuda_api::LibraryHandle GetHandle() const { return library_; } // Non-copyable CubinModule(const CubinModule&) = delete; CubinModule& operator=(const CubinModule&) = delete; /*! * \brief Move constructor for CubinModule. * * Transfers ownership of the CUDA library handle from another CubinModule instance. * * \param other The source CubinModule to move from (will be left in an empty state). */ CubinModule(CubinModule&& other) noexcept : library_(other.library_) { other.library_ = nullptr; } /*! * \brief Move assignment operator for CubinModule. * * Transfers ownership of the CUDA library handle from another CubinModule instance. * Cleans up any existing library handle in this instance before taking ownership. * * \param other The source CubinModule to move from (will be left in an empty state). * \return Reference to this CubinModule. */ CubinModule& operator=(CubinModule&& other) noexcept { if (this != &other) { if (library_ != nullptr) { cuda_api::UnloadLibrary(library_); } library_ = other.library_; other.library_ = nullptr; } return *this; } private: cuda_api::LibraryHandle library_ = nullptr; }; /*! * \brief CUDA kernel handle for launching kernels. * * This class represents a loaded CUDA kernel function and provides * methods to launch it with specified grid/block dimensions, arguments, * and stream configuration. Obtained from CubinModule by kernel name. * * \par Usage Pattern * \code{.cpp} * // Get kernel from module * tvm::ffi::CubinKernel kernel = module["kernel_name"]; * * // Prepare arguments (must be pointers to actual values) * void* data_ptr = tensor.data_ptr(); * int64_t size = tensor.size(0); * void* args[] = {&data_ptr, &size}; * * // Configure launch dimensions * tvm::ffi::dim3 grid(32); // 32 blocks * tvm::ffi::dim3 block(256); // 256 threads per block * * // Launch on stream * cudaStream_t stream = ...; * cudaError_t result = kernel.Launch(args, grid, block, stream); * TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); * \endcode * * \note This class is movable but not copyable. * \see CubinModule for loading CUBIN and getting kernels * \see dim3 for grid/block dimension specification */ class CubinKernel { public: /*! * \brief Construct a CubinKernel from a library and kernel name. * * \param library The cudaLibrary_t handle. * \param name Name of the kernel function. */ CubinKernel(cuda_api::LibraryHandle library, const char* name) { TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(cuda_api::GetKernel(&kernel_, library, name)); } /*! \brief Destructor (kernel handle doesn't need explicit cleanup) */ ~CubinKernel() = default; /*! * \brief Launch the kernel with specified parameters. * * This function launches the kernel on the current CUDA context/device using * the CUDA Runtime API. The kernel executes asynchronously on the specified stream. * * \par Argument Preparation * The `args` array must contain pointers to the actual argument values, not the * values themselves. For example: * \code{.cpp} * void* data_ptr = tensor.data_ptr(); * int64_t size = 100; * void* args[] = {&data_ptr, &size}; // Note: addresses of the variables * \endcode * * \par Launch Configuration * Grid and block dimensions determine the kernel's parallelism: * - Grid: Number of thread blocks (can be 1D, 2D, or 3D) * - Block: Number of threads per block (can be 1D, 2D, or 3D) * - Total threads = grid.x * grid.y * grid.z * block.x * block.y * block.z * * \par Error Checking * Always check the returned cudaError_t: * \code{.cpp} * TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(kernel.Launch(args, grid, block, stream)); * \endcode * * \param args Array of pointers to kernel arguments (must point to actual values). * \param grid Grid dimensions (number of blocks in x, y, z). * \param block Block dimensions (threads per block in x, y, z). * \param stream CUDA stream to launch the kernel on (use 0 for default stream). * \param dyn_smem_bytes Dynamic shared memory size in bytes (default: 0). * \return cudaError_t error code from cudaLaunchKernel (cudaSuccess on success). * * \note The kernel executes asynchronously. Use cudaStreamSynchronize() or * cudaDeviceSynchronize() to wait for completion if needed. */ cuda_api::ResultType Launch(void** args, dim3 grid, dim3 block, cuda_api::StreamHandle stream, uint32_t dyn_smem_bytes = 0) { return cuda_api::LaunchKernel(kernel_, args, grid, block, stream, dyn_smem_bytes); } /*! * \brief Launch the kernel using extended launch API with a pre-built config. * * This enables features like cluster dimensions (SM90+) that require * cuLaunchKernelEx / cudaLaunchKernelExC. * * \param args Array of pointers to kernel arguments. * \param config The launch configuration (populated by ConstructLaunchConfig). * \return Result code. */ cuda_api::ResultType LaunchEx(void** args, const cuda_api::LaunchConfig& config) { return cuda_api::LaunchKernelEx(kernel_, args, config); } /*! \brief Get the underlying cudaKernel_t handle */ cuda_api::KernelHandle GetHandle() const { return kernel_; } // Non-copyable CubinKernel(const CubinKernel&) = delete; CubinKernel& operator=(const CubinKernel&) = delete; /*! * \brief Move constructor for CubinKernel. * * Transfers ownership of the CUDA kernel handle from another CubinKernel instance. * * \param other The source CubinKernel to move from (will be left in an empty state). */ CubinKernel(CubinKernel&& other) noexcept : kernel_(other.kernel_) { other.kernel_ = nullptr; } /*! * \brief Move assignment operator for CubinKernel. * * Transfers ownership of the CUDA kernel handle from another CubinKernel instance. * * \param other The source CubinKernel to move from (will be left in an empty state). * \return Reference to this CubinKernel. */ CubinKernel& operator=(CubinKernel&& other) noexcept { if (this != &other) { kernel_ = other.kernel_; other.kernel_ = nullptr; } return *this; } private: /*! * \brief Set maximum dynamic shared memory for this kernel across all devices. * * This method configures the maximum dynamic shared memory that can be allocated * when launching this kernel. It must be called after the kernel is loaded. * * \param dynamic_smem_max Maximum dynamic shared memory in bytes to set. * -1 (default) means maximum available dynamic shared memory, * which is computed as (device max shared memory - static shared memory). * For -1, the method queries the kernel's static shared memory usage * and sets the attribute to the remaining available shared memory. * * \note This sets the maximum cap but doesn't force allocation. The actual dynamic * shared memory used is controlled by the dyn_smem_bytes parameter in Launch(). * \note This method attempts to set the attribute for all available devices and will * only throw an error if it fails for ALL devices. */ void SetMaxDynamicSharedMemory(int64_t dynamic_smem_max = -1) { int device_count = 0; cuda_api::ResultType err = cuda_api::GetDeviceCount(&device_count); if (err != cuda_api::kSuccess || device_count == 0) { return; // No devices available, nothing to configure } bool any_success = false; for (int device_id = 0; device_id < device_count; ++device_id) { auto device = cuda_api::GetDeviceHandle(device_id); // Query device's maximum shared memory per block int max_shared_mem = 0; err = cuda_api::GetDeviceAttribute( &max_shared_mem, /* CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK/cudaDevAttrMaxSharedMemoryPerBlock */ cuda_api::DeviceAttrType(8), device); if (err != cuda_api::kSuccess) { continue; // Skip this device if we can't get its attribute } int shared_mem_to_set; if (dynamic_smem_max == -1) { int static_shared; err = cuda_api::GetKernelSharedMem(kernel_, static_shared, device); if (err != cuda_api::kSuccess) { continue; // Skip this device if we can't get kernel attributes } // Calculate available dynamic shared memory: // device max shared memory - static shared memory used by kernel int64_t max_shared = static_cast(max_shared_mem); int64_t available = max_shared - static_shared; shared_mem_to_set = (available > 0) ? static_cast(available) : 0; } else { shared_mem_to_set = static_cast(dynamic_smem_max); } // Set the maximum dynamic shared memory size for this device err = cuda_api::SetKernelMaxDynamicSharedMem(kernel_, shared_mem_to_set, device); if (err == cuda_api::kSuccess) { any_success = true; } // Don't error out for individual device failures - user may only use some GPUs } // Only error out if setting failed for ALL devices if (!any_success && device_count > 0) { TVM_FFI_THROW(RuntimeError) << "Failed to set dynamic shared memory attribute for any device"; } } cuda_api::KernelHandle kernel_ = nullptr; friend class CubinModule; }; // Implementation of CubinModule methods that return CubinKernel inline CubinKernel CubinModule::GetKernelWithMaxDynamicSharedMemory(const char* name, int64_t dynamic_smem_max = -1) { auto kernel = CubinKernel(library_, name); kernel.SetMaxDynamicSharedMemory(dynamic_smem_max); return kernel; } inline CubinKernel CubinModule::GetKernel(const char* name) { auto kernel = CubinKernel(library_, name); return kernel; } inline CubinKernel CubinModule::operator[](const char* name) { return GetKernel(name); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_CUDA_CUBIN_LAUNCHER_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/device_guard.h000066400000000000000000000045201521067262500230000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/cuda/device_guard.h * \brief Device guard structs. */ #ifndef TVM_FFI_EXTRA_CUDA_DEVICE_GUARD_H_ #define TVM_FFI_EXTRA_CUDA_DEVICE_GUARD_H_ #include namespace tvm { namespace ffi { /*! * \brief CUDA Device Guard. On construction, it calls `cudaGetDevice` to set the CUDA device to * target index, and stores the original current device index. And on destruction, it sets the * current CUDA device back to original device index. * * Example usage: * \code{.cpp} * void kernel(ffi::TensorView x) { * ffi::CUDADeviceGuard guard(x.device().device_id); * ... * } * \endcode */ struct CUDADeviceGuard { CUDADeviceGuard() = delete; /*! * \brief Constructor from a device index, and store the original device index. * \param device_index The device index to guard. */ explicit CUDADeviceGuard(int device_index) { target_device_index_ = device_index; TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&original_device_index_)); if (target_device_index_ != original_device_index_) { TVM_FFI_CHECK_CUDA_ERROR(cudaSetDevice(device_index)); } } /*! * \brief Destructor to set the current device index back to original one if different. */ ~CUDADeviceGuard() noexcept(false) { if (original_device_index_ != target_device_index_) { TVM_FFI_CHECK_CUDA_ERROR(cudaSetDevice(original_device_index_)); } } private: int original_device_index_; int target_device_index_; }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_CUDA_DEVICE_GUARD_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/internal/000077500000000000000000000000001521067262500220215ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/extra/cuda/internal/unified_api.h000066400000000000000000000242071521067262500244530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #ifndef TVM_FFI_EXTRA_CUDA_INTERNAL_UNIFIED_API_H_ #define TVM_FFI_EXTRA_CUDA_INTERNAL_UNIFIED_API_H_ #include #include #include // =========================================================================== // Section 1: Configuration & Version Checks // =========================================================================== // We only use unified API for cubin launcher for now // this name is intentional to avoid confusion of other API usages #ifndef TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API #if CUDART_VERSION >= 12080 // Use Runtime API by default if possible (CUDA >= 12.8) #define TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API 0 #else // if CUDART_VERSION < 12080 #define TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API 1 #endif #else // if defined(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API) // User explicitly defined the macro, check compatibility #if (!(TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API)) && (CUDART_VERSION < 12080) #define _STRINGIFY(x) #x #define STR(x) _STRINGIFY(x) static_assert(false, "Runtime API only supported for CUDA >= 12.8, got CUDA Runtime version: " STR( CUDART_VERSION)); #undef STR #undef _STRINGIFY #endif #endif namespace tvm { namespace ffi { namespace cuda_api { // =========================================================================== // Section 2: Type Definitions & Macros // =========================================================================== #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API // Driver API Types using StreamHandle = CUstream; using DeviceHandle = CUdevice; using LibraryHandle = CUlibrary; using KernelHandle = CUkernel; using LaunchConfig = CUlaunchConfig; using ResultType = CUresult; using LaunchAttrType = CUlaunchAttribute; using DeviceAttrType = CUdevice_attribute; constexpr ResultType kSuccess = CUDA_SUCCESS; // Driver API Functions #define _TVM_FFI_CUDA_FUNC(name) cu##name // NOLINT(bugprone-reserved-identifier) #else using StreamHandle = cudaStream_t; using DeviceHandle = int; using LibraryHandle = cudaLibrary_t; using KernelHandle = cudaKernel_t; using LaunchConfig = cudaLaunchConfig_t; using ResultType = cudaError_t; using LaunchAttrType = cudaLaunchAttribute; using DeviceAttrType = cudaDeviceAttr; constexpr ResultType kSuccess = cudaSuccess; // Runtime API Functions #define _TVM_FFI_CUDA_FUNC(name) cuda##name #endif // =========================================================================== // Section 3: Error Handling // =========================================================================== // Helper to get error name/string based on API inline void GetErrorString(ResultType err, const char** name, const char** str) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API cuGetErrorName(err, name); cuGetErrorString(err, str); #else *name = cudaGetErrorName(err); *str = cudaGetErrorString(err); #endif } // this macro is only used to check cuda errors in cubin launcher where // we might switch between driver and runtime API. #define TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(stmt) \ do { \ ::tvm::ffi::cuda_api::ResultType __err = (stmt); \ if (__err != ::tvm::ffi::cuda_api::kSuccess) { \ const char *__err_name, *__err_str; \ ::tvm::ffi::cuda_api::GetErrorString(__err, &__err_name, &__err_str); \ TVM_FFI_THROW(RuntimeError) << "CUDA Error: " << __err_name << " (" \ << static_cast(__err) << "): " << __err_str; \ } \ } while (0) // =========================================================================== // Section 4: Unified API Wrappers // =========================================================================== inline ResultType LoadLibrary(LibraryHandle* library, const void* image) { return _TVM_FFI_CUDA_FUNC(LibraryLoadData)(library, image, nullptr, nullptr, 0, nullptr, nullptr, 0); } inline ResultType UnloadLibrary(LibraryHandle library) { return _TVM_FFI_CUDA_FUNC(LibraryUnload)(library); } inline ResultType GetKernel(KernelHandle* kernel, LibraryHandle library, const char* name) { return _TVM_FFI_CUDA_FUNC(LibraryGetKernel)(kernel, library, name); } inline DeviceHandle GetDeviceHandle(int device_id) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API CUdevice dev; // Note: We use CHECK here because this conversion usually shouldn't fail if ID is valid // and we need to return a value. TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(cuDeviceGet(&dev, device_id)); return dev; #else return device_id; #endif } inline ResultType LaunchKernel(KernelHandle kernel, void** args, tvm::ffi::dim3 grid, tvm::ffi::dim3 block, StreamHandle stream, uint32_t dyn_smem_bytes = 0) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API return cuLaunchKernel(reinterpret_cast(kernel), grid.x, grid.y, grid.z, block.x, block.y, block.z, dyn_smem_bytes, stream, args, nullptr); #else return cudaLaunchKernel(reinterpret_cast(kernel), {grid.x, grid.y, grid.z}, {block.x, block.y, block.z}, args, dyn_smem_bytes, stream); #endif } inline ResultType GetKernelSharedMem(KernelHandle kernel, int& out, DeviceHandle device) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API return cuKernelGetAttribute(&out, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, kernel, device); #else cudaFuncAttributes func_attr; cudaError_t err = cudaFuncGetAttributes(&func_attr, kernel); if (err == cudaSuccess) { out = func_attr.sharedSizeBytes; } return err; #endif } inline ResultType SetKernelMaxDynamicSharedMem(KernelHandle kernel, int shmem, DeviceHandle device) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API return cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem, kernel, device); #else return cudaKernelSetAttributeForDevice(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem, device); #endif } /*! * \brief Launch a kernel using the extended launch API with launch attributes. * * This enables features like cluster dimensions (SM90+) that require * cuLaunchKernelEx / cudaLaunchKernelExC. * * \param kernel The kernel handle. * \param args Array of pointers to kernel arguments. * \param config The launch configuration (grid, block, smem, stream, attributes). * \return Result code. */ inline ResultType LaunchKernelEx(KernelHandle kernel, void** args, const LaunchConfig& config) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API return cuLaunchKernelEx(&config, reinterpret_cast(kernel), args, nullptr); #else return cudaLaunchKernelExC(&config, reinterpret_cast(kernel), args); #endif } /*! * \brief Construct a launch configuration with optional cluster dimensions. * * \param kernel The kernel handle. * \param stream The CUDA stream. * \param smem_size Dynamic shared memory size in bytes. * \param grid Grid dimensions. * \param block Block dimensions. * \param cluster_dim Cluster dimension (1 = no clustering, >1 enables cluster launch). * \param[out] config The launch configuration to populate. * \param[out] attr Storage for a launch attribute (must outlive the launch call). * \return kSuccess. */ inline ResultType ConstructLaunchConfig(KernelHandle kernel, StreamHandle stream, uint32_t smem_size, tvm::ffi::dim3 grid, tvm::ffi::dim3 block, int cluster_dim, LaunchConfig& config, LaunchAttrType& attr) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API config.gridDimX = grid.x; config.gridDimY = grid.y; config.gridDimZ = grid.z; config.blockDimX = block.x; config.blockDimY = block.y; config.blockDimZ = block.z; config.sharedMemBytes = smem_size; config.hStream = stream; config.numAttrs = 0; config.attrs = nullptr; if (cluster_dim > 1) { attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; attr.value.clusterDim.x = static_cast(cluster_dim); attr.value.clusterDim.y = 1; attr.value.clusterDim.z = 1; config.attrs = &attr; config.numAttrs = 1; } #else config.gridDim = {grid.x, grid.y, grid.z}; config.blockDim = {block.x, block.y, block.z}; config.dynamicSmemBytes = smem_size; config.stream = stream; config.numAttrs = 0; config.attrs = nullptr; if (cluster_dim > 1) { attr.id = cudaLaunchAttributeClusterDimension; attr.val.clusterDim = {static_cast(cluster_dim), 1, 1}; config.attrs = &attr; config.numAttrs = 1; } #endif return kSuccess; } // Additional wrappers for device operations used in CubinLauncher inline ResultType GetDeviceCount(int* count) { #if TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API return cuDeviceGetCount(count); #else return cudaGetDeviceCount(count); #endif } inline ResultType GetDeviceAttribute(int* value, DeviceAttrType attr, DeviceHandle device) { return _TVM_FFI_CUDA_FUNC(DeviceGetAttribute)(value, attr, device); } } // namespace cuda_api } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_CUDA_INTERNAL_UNIFIED_API_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/dataclass.h000066400000000000000000000114411521067262500214020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/dataclass.h * \brief Reflection-based dataclass operations: deep copy, repr, hash, compare. */ #ifndef TVM_FFI_EXTRA_DATACLASS_H_ #define TVM_FFI_EXTRA_DATACLASS_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { /** * \brief Deep copy an ffi::Any value. * * Recursively copies the value and all reachable objects in its object graph. * Copy-constructible types with `ObjectDef` registration automatically support deep copy. * Primitive types, strings, bytes, and Shape are returned as-is (they are immutable). * Arrays, Lists, Maps, and Dicts are recursively deep copied. * Objects without copy support cause a runtime error. * * \param value The value to deep copy. * \return The deep copied value. */ TVM_FFI_EXTRA_CXX_API Any DeepCopy(const Any& value); /** * \brief Produce a human-readable repr string for an ffi::Any value. * * Recursively formats the value using reflection metadata. * Handles cycles (prints "...") and DAGs (caches repr strings). * Custom __ffi_repr__ hooks are supported via the reflection type attribute system. * * \param value The value to repr-print. * \return The repr string. */ TVM_FFI_EXTRA_CXX_API String ReprPrint(const Any& value); /** * \brief Compute a deterministic recursive hash of an ffi::Any value. * * Recursively hashes the value and all reachable objects using reflection. * Consistent with RecursiveEq: RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b). * Custom __ffi_hash__ hooks are supported. * * \param value The value to hash. * \return The hash as int64_t (for FFI compatibility). */ TVM_FFI_EXTRA_CXX_API int64_t RecursiveHash(const Any& value); /** * \brief Recursive structural equality comparison. * \param lhs Left-hand side value. * \param rhs Right-hand side value. * \return true if the two values are structurally equal. */ TVM_FFI_EXTRA_CXX_API bool RecursiveEq(const Any& lhs, const Any& rhs); /** * \brief Recursive structural less-than comparison. * \param lhs Left-hand side value. * \param rhs Right-hand side value. * \return true if lhs is structurally less than rhs. */ TVM_FFI_EXTRA_CXX_API bool RecursiveLt(const Any& lhs, const Any& rhs); /** * \brief Recursive structural less-than-or-equal comparison. * \param lhs Left-hand side value. * \param rhs Right-hand side value. * \return true if lhs is structurally less than or equal to rhs. */ TVM_FFI_EXTRA_CXX_API bool RecursiveLe(const Any& lhs, const Any& rhs); /** * \brief Recursive structural greater-than comparison. * \param lhs Left-hand side value. * \param rhs Right-hand side value. * \return true if lhs is structurally greater than rhs. */ TVM_FFI_EXTRA_CXX_API bool RecursiveGt(const Any& lhs, const Any& rhs); /** * \brief Recursive structural greater-than-or-equal comparison. * \param lhs Left-hand side value. * \param rhs Right-hand side value. * \return true if lhs is structurally greater than or equal to rhs. */ TVM_FFI_EXTRA_CXX_API bool RecursiveGe(const Any& lhs, const Any& rhs); // std::ostream overloads /*! \brief Stream an ffi::Any using its canonical repr form. */ inline std::ostream& operator<<(std::ostream& os, const Any& value) { return os << ReprPrint(value); } /*! \brief Stream an ffi::ObjectRef using its canonical repr form. */ inline std::ostream& operator<<(std::ostream& os, const ObjectRef& value) { return os << ReprPrint(Any(value)); } /*! \brief Stream an ffi::Variant<...> using its canonical repr form. */ template inline std::ostream& operator<<(std::ostream& os, const Variant& value) { return os << ReprPrint(Any(value)); } /*! \brief Stream an ffi::Optional using its canonical repr form. */ template inline std::ostream& operator<<(std::ostream& os, const Optional& value) { return os << ReprPrint(Any(value)); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_DATACLASS_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/dtype.h000066400000000000000000000125161521067262500205740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/dtype.h * \brief Type traits to map C++ types to DLPack dtypes. */ #ifndef TVM_FFI_EXTRA_DTYPE_H_ #define TVM_FFI_EXTRA_DTYPE_H_ #include #include // Common for both CUDA and HIP struct __half; // CUDA struct __nv_fp8_e4m3; struct __nv_bfloat16; struct __nv_fp8_e5m2; struct __nv_fp8_e8m0; struct __nv_fp4_e2m1; struct __nv_fp4x2_e2m1; // HIP struct __hip_bfloat16; struct hip_bfloat16; // i don't know why this is a struct instead of alias... struct __hip_fp8_e4m3; struct __hip_fp8_e4m3_fnuz; struct __hip_fp8_e5m2; struct __hip_fp8_e5m2_fnuz; struct __hip_fp4_e2m1; struct __hip_fp4x2_e2m1; namespace tvm_ffi { /// \cond Doxygen_Suppress template struct dtype_trait {}; namespace details::dtypes { template struct integer_trait { static constexpr DLDataType value = { /* code = */ std::is_signed_v ? kDLInt : kDLUInt, /* bits = */ static_cast(sizeof(T) * 8), /* lanes = */ 1, }; }; template struct float_trait { static constexpr DLDataType value = { /* code = */ kDLFloat, /* bits = */ static_cast(sizeof(T) * 8), /* lanes = */ 1, }; }; } // namespace details::dtypes template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::integer_trait {}; template <> struct dtype_trait : details::dtypes::float_trait {}; template <> struct dtype_trait : details::dtypes::float_trait {}; // Specialization for bool template <> struct dtype_trait { static constexpr DLDataType value = {DLDataTypeCode::kDLBool, 8, 1}; }; // Specializations for CUDA template <> struct dtype_trait<__half> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat, 16, 1}; }; template <> struct dtype_trait<__nv_bfloat16> { static constexpr DLDataType value = {DLDataTypeCode::kDLBfloat, 16, 1}; }; template <> struct dtype_trait<__nv_fp8_e4m3> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e4m3fn, 8, 1}; }; template <> struct dtype_trait<__nv_fp8_e5m2> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e5m2, 8, 1}; }; template <> struct dtype_trait<__nv_fp8_e8m0> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e8m0fnu, 8, 1}; }; template <> struct dtype_trait<__nv_fp4_e2m1> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat4_e2m1fn, 4, 1}; }; template <> struct dtype_trait<__nv_fp4x2_e2m1> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat4_e2m1fn, 4, 2}; }; // Specializations for HIP template <> struct dtype_trait<__hip_bfloat16> { static constexpr DLDataType value = {DLDataTypeCode::kDLBfloat, 16, 1}; }; template <> struct dtype_trait { static constexpr DLDataType value = {DLDataTypeCode::kDLBfloat, 16, 1}; }; template <> struct dtype_trait<__hip_fp8_e4m3> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e4m3fn, 8, 1}; }; template <> struct dtype_trait<__hip_fp8_e4m3_fnuz> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e4m3fnuz, 8, 1}; }; template <> struct dtype_trait<__hip_fp8_e5m2> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e5m2, 8, 1}; }; template <> struct dtype_trait<__hip_fp8_e5m2_fnuz> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat8_e5m2fnuz, 8, 1}; }; template <> struct dtype_trait<__hip_fp4_e2m1> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat4_e2m1fn, 4, 1}; }; template <> struct dtype_trait<__hip_fp4x2_e2m1> { static constexpr DLDataType value = {DLDataTypeCode::kDLFloat4_e2m1fn, 4, 2}; }; /// \endcond } // namespace tvm_ffi #endif // TVM_FFI_EXTRA_DTYPE_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/json.h000066400000000000000000000052431521067262500204170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/json.h * \brief Minimal lightweight JSON parsing and serialization utilities */ #ifndef TVM_FFI_EXTRA_JSON_H_ #define TVM_FFI_EXTRA_JSON_H_ #include #include #include #include namespace tvm { namespace ffi { namespace json { /*! * \brief alias Any as json Value. * * To keep things lightweight, we simply reuse the ffi::Any system. */ using Value = Any; /*! * \brief alias Map as json Object. * \note We use Map instead of Map to avoid * the overhead of key checking when doing as conversion, * the check will be performed at runtime when we read each key */ using Object = ffi::Map; /*! \brief alias Array as json Array. */ using Array = ffi::Array; /*! * \brief Parse a JSON string into an Any value. * * Besides the standard JSON syntax, this function also supports: * - Infinity/NaN as JavaScript syntax * - int64 integer value * * If error_msg is not nullptr, the error message will be written to it * and no exception will be thrown when parsing fails. * * \param json_str The JSON string to parse. * \param error_msg The output error message, can be nullptr. * * \return The parsed Any value. */ TVM_FFI_EXTRA_CXX_API json::Value Parse(const String& json_str, String* error_msg = nullptr); /*! * \brief Serialize an Any value into a JSON string. * * \param value The Any value to serialize. * \param indent The number of spaces to indent the output. * If not specified, the output will be compact. * \return The output JSON string. */ TVM_FFI_EXTRA_CXX_API String Stringify(const json::Value& value, Optional indent = std::nullopt); } // namespace json } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_JSON_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/module.h000066400000000000000000000253231521067262500207340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/module.h * \brief A managed dynamic module in the TVM FFI. */ #ifndef TVM_FFI_EXTRA_MODULE_H_ #define TVM_FFI_EXTRA_MODULE_H_ #include #include #include #include #include namespace tvm { namespace ffi { // forward declare Module class Module; /*! * \brief A module that can dynamically load ffi::Functions or exportable source code. * \sa Module */ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object { public: /*! * \return The per module type key. * \note This key is used to for serializing custom modules. */ virtual const char* kind() const = 0; /*! * \brief Get the property mask of the module. * \return The property mask of the module. * * \sa Module::ModulePropertyMask */ virtual int GetPropertyMask() const { return 0b000; } /*! * \brief Get a ffi::Function from the module. * \param name The name of the function. * \return The function. */ virtual Optional GetFunction(const String& name) = 0; /*! * \brief Returns true if this module has a definition for a function of \p name. * * Note that even if this function returns true the corresponding \p GetFunction result * may be nullptr if the function is not yet callable without further compilation. * * The default implementation just checks if \p GetFunction is non-null. * \param name The name of the function. * \return True if the module implements the function, false otherwise. */ virtual bool ImplementsFunction(const String& name) { return GetFunction(name).defined(); } /*! * \brief Get the docstring of the function, if available. * \param name The name of the function. * \return The documentation string if available, nullopt otherwise. * * \sa GetFunctionMetadata, TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC */ virtual Optional GetFunctionDoc(const String& name) { return std::nullopt; } // Rationale: We separate the docstring from the metadata since docstrings // can be unstructured and sometimes large, while metadata can be focused // on storing structured information. /*! * \brief Get the metadata of the function, if available. * \param name The name of the function. * \return The metadata as JSON string if available, nullopt otherwise. * * \code{.cpp} * Module mod = Module::LoadFromFile("lib.so"); * Optional metadata = mod->GetFunctionMetadata("my_func"); * if (metadata.has_value()) { * // Parse JSON: {"type_schema": "..."} * validate_signature(*metadata); * } * \endcode * * \sa GetFunctionDoc, TVM_FFI_DLL_EXPORT_TYPED_FUNC */ virtual Optional GetFunctionMetadata(const String& name) { return std::nullopt; } /*! * \brief Write the current module to file with given format (for further compilation). * * \param file_name The file to be saved to. * \param format The format of the file. * * \note This function is mainly used by modules that */ virtual void WriteToFile(const String& file_name, const String& format) const { TVM_FFI_THROW(RuntimeError) << "Module[" << kind() << "] does not support WriteToFile"; } /*! * \brief Get the possible write formats of the module, when available. * \return Possible write formats when available. */ virtual Array GetWriteFormats() const { return Array(); } /*! * \brief Serialize the the module to bytes. * \return The serialized module. */ virtual Bytes SaveToBytes() const { TVM_FFI_THROW(RuntimeError) << "Module[" << kind() << "] does not support SaveToBytes"; TVM_FFI_UNREACHABLE(); } /*! * \brief Get the source code of module, when available. * \param format Format of the source code, can be empty by default. * \return Possible source code when available, or empty string if not available. */ virtual String InspectSource(const String& format) const { return String(); } /*! * \brief Import another module. * \param other The module to import. */ virtual void ImportModule(const Module& other); /*! * \brief Clear all imported modules. */ virtual void ClearImports(); /*! * \brief Overloaded function to optionally query from imports. * \param name The name of the function. * \param query_imports Whether to query imported modules. * \return The function. */ Optional GetFunction(const String& name, bool query_imports); /*! * \brief Overloaded function to optionally query from imports. * \param name The name of the function. * \param query_imports Whether to query imported modules. * \return True if the module implements the function, false otherwise. */ bool ImplementsFunction(const String& name, bool query_imports); /*! * \brief Get the function docstring of the function if available. * \param name The name of the function. * \param query_imports Whether to also query modules imported by this module. * \return The documentation string if available, nullopt otherwise. * * \sa GetFunctionMetadata */ Optional GetFunctionDoc(const String& name, bool query_imports); /*! * \brief Get the function metadata of the function if available. * \param name The name of the function. * \param query_imports Whether to also query modules imported by this module. * \return The metadata as JSON string if available, nullopt otherwise. * * \sa GetFunctionDoc */ Optional GetFunctionMetadata(const String& name, bool query_imports); /*! * \brief Get the imports of the module. * \return The imports of the module. * \note Note the signature is not part of the public API. */ const Array& imports() const { return this->imports_; } struct InternalUnsafe; /// \cond Doxygen_Suppress static constexpr const int32_t _type_index = TypeIndex::kTVMFFIModule; static constexpr const bool _type_mutable = true; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIModule, ModuleObj, Object); /// \endcond protected: friend struct InternalUnsafe; /*! * \brief The modules that this module depends on. * \note Use ObjectRef to avoid circular dep on Module. */ Array imports_; private: /*! * \brief cache used by TVMFFIModuleLookupFromImports */ Map import_lookup_cache_; }; /*! * \brief Reference to module object. * * When invoking a function on a ModuleObj, such as GetFunction, * use operator-> to get the ModuleObj pointer and invoke the member functions. * * \code{.cpp} * ffi::Module mod = ffi::Module::LoadFromFile("path/to/module.so"); * ffi::Function func = mod->GetFunction(name); * \endcode * * \sa ModuleObj which contains most of the function implementations. */ class Module : public ObjectRef { public: /*! * \brief Property of ffi::Module */ enum ModulePropertyMask : int { /*! * \brief The module can be serialized to bytes. * * This prooperty indicates that module implements SaveToBytes. * The system also registers a GlobalDef function * `ffi.Module.load_from_bytes.` with signature (Bytes) -> Module. */ kBinarySerializable = 0b001, /*! * \brief The module can directly get runnable functions. * * This property indicates that module implements GetFunction that returns * runnable ffi::Functions. */ kRunnable = 0b010, /*! * \brief The module can be exported to a object file or source file that then be compiled. * * This property indicates that module implements WriteToFile with a given format * that can be queried by GetLibExportFormat. * * Examples include modules that can be exported to .o, .cc, .cu files. * * Such modules can be exported, compiled and loaded back as a dynamic library module. */ kCompilationExportable = 0b100 }; /*! * \brief Constructor from ObjectPtr. * \param ptr The object pointer. */ explicit Module(const ObjectPtr& ptr) : ObjectRef(ptr) { TVM_FFI_ICHECK(ptr != nullptr); } /*! * \brief Load a module from file. * \param file_name The name of the host function module. * \note This function won't load the import relationship. * Re-create import relationship by calling Import. */ TVM_FFI_EXTRA_CXX_API static Module LoadFromFile(const String& file_name); /*! * \brief Query context symbols that is registered via TVMEnvRegisterSymbols. * \param callback The callback to be called with the symbol name and address. * \note This helper can be used to implement custom Module that needs to access context symbols. */ TVM_FFI_EXTRA_CXX_API static void VisitContextSymbols( const ffi::TypedFunction& callback); /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Module, ObjectRef, ModuleObj); /// \endcond }; /* * \brief Symbols for library module. */ namespace symbol { /*!\ brief symbol prefix for tvm ffi related function symbols */ constexpr const char* tvm_ffi_symbol_prefix = "__tvm_ffi_"; // Special symbols have one extra _ prefix to avoid conflict with user symbols /*! * \brief Default entry function of a library module is tvm_ffi_symbol_prefix + "main" */ constexpr const char* tvm_ffi_main = "__tvm_ffi_main"; /*! \brief Global variable to store context pointer for a library module. */ constexpr const char* tvm_ffi_library_ctx = "__tvm_ffi__library_ctx"; /*! \brief Global variable to store binary data alongside a library module. */ constexpr const char* tvm_ffi_library_bin = "__tvm_ffi__library_bin"; /*! \brief Optional metadata prefix of a symbol. */ constexpr const char* tvm_ffi_metadata_prefix = "__tvm_ffi__metadata_"; /*! \brief Optional documentation prefix of a symbol. */ constexpr const char* tvm_ffi_doc_prefix = "__tvm_ffi__doc_"; } // namespace symbol } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_MODULE_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/serialization.h000066400000000000000000000050071521067262500223210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/serialization.h * \brief Reflection-based serialization utilities */ #ifndef TVM_FFI_EXTRA_SERIALIZATION_H_ #define TVM_FFI_EXTRA_SERIALIZATION_H_ #include #include namespace tvm { namespace ffi { /** * \brief Serialize ffi::Any to a JSON that stores the object graph. * * The JSON graph structure is stored as follows: * * ``` * { * "root_index": , // Index of root node in nodes array * "nodes": [, ...], // Array of serialized nodes * "metadata": // Optional metadata * } * ``` * * Each node has the format: `{"type": "", "data": }` * For object types and strings, the data may contain indices to other nodes. * For object fields whose static type is known as a primitive type, it is stored directly, * otherwise, it is stored as a reference to the nodes array by an index. * * This function preserves the type and multiple references to the same object, * which is useful for debugging and serialization. * * \param value The ffi::Any value to serialize. * \param metadata Extra metadata attached to "metadata" field of the JSON object. * \return The serialized JSON value. */ TVM_FFI_EXTRA_CXX_API json::Value ToJSONGraph(const Any& value, const Any& metadata = Any(nullptr)); /** * \brief Deserialize a JSON that stores the object graph to an ffi::Any value. * * This function can be used to implement deserialization * and debugging. * * \param value The JSON value to deserialize. * \return The deserialized object graph. */ TVM_FFI_EXTRA_CXX_API Any FromJSONGraph(const json::Value& value); } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_SERIALIZATION_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/stl.h000066400000000000000000000531501521067262500202500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/stl.h * \brief STL container support. * \note This file is an extra extension of TVM FFI, * which provides support for STL containers in C++ exported functions. * * Whenever possible, prefer using tvm/ffi/container/ implementations, * such as `tvm::ffi::Array` and `tvm::ffi::Tuple`, over STL containers. * * Native ffi objects comes with stable data layout and can be directly accessed * through compiled languages (Rust) and DSLs(via LLVM) with raw pointer access * for better performance and compatibility. */ #ifndef TVM_FFI_EXTRA_STL_H_ #define TVM_FFI_EXTRA_STL_H_ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "tvm/ffi/function.h" namespace tvm { namespace ffi { namespace details { struct STLTypeMismatch : public std::exception { const char* what() const noexcept override { return "STL type mismatch"; } }; struct STLTypeTrait : public TypeTraitsBase { public: static constexpr bool storage_enabled = false; protected: /// NOTE: we always copy STL types into an Object first, then move the ObjectPtr to Any. template TVM_FFI_INLINE static void MoveToAnyImpl(ObjectPtr&& src, TVMFFIAny* result) { TVMFFIObject* obj_ptr = ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(src)); result->type_index = obj_ptr->type_index; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_obj = obj_ptr; } /// NOTE: we always construct STL types from an Object first, then copy from the ObjectPtr in Any. template TVM_FFI_INLINE static ObjectPtr CopyFromAnyImpl(const TVMFFIAny* src) { return ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj); } /// NOTE: STL types are not natively movable from Any, so we always make a new copy. template TVM_FFI_INLINE static T ConstructFromAny(const Any& value) { using TypeTrait = TypeTraits; if constexpr (std::is_same_v) { return value; } else { /// NOTE: Not all type support `CheckArgStrict` (e.g. std::string), /// so we use `TryCast` instead (without any prior check). auto opt = TypeTrait::TryCastFromAnyView(AnyUnsafe::TVMFFIAnyPtrFromAny(value)); if (!opt.has_value()) { throw STLTypeMismatch{}; } return std::move(*opt); } } }; struct ListTemplate {}; struct MapTemplate {}; } // namespace details template <> struct TypeTraits : public details::STLTypeTrait { public: static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIArray; private: template TVM_FFI_INLINE static ObjectPtr CopyToTupleImpl(std::index_sequence, Tuple&& src) { auto array = ArrayObj::Empty(static_cast(sizeof...(Is))); auto dst = array->MutableBegin(); // increase size after each new to ensure exception safety std::apply( [&](auto&&... elems) { ((::new (dst++) Any(std::forward(elems)), array->TVMFFISeqCell::size++), ...); }, std::forward(src)); return array; } template TVM_FFI_INLINE static ObjectPtr CopyToArrayImpl(Iter src, std::size_t size) { auto array = ArrayObj::Empty(static_cast(size)); auto dst = array->MutableBegin(); // increase size after each new to ensure exception safety for (std::size_t i = 0; i < size; ++i) { ::new (dst++) Any(*(src++)); array->TVMFFISeqCell::size++; } return array; } protected: template TVM_FFI_INLINE static ObjectPtr CopyToTuple(const Tuple& src) { return CopyToTupleImpl(std::make_index_sequence>{}, src); } template TVM_FFI_INLINE static ObjectPtr MoveToTuple(Tuple&& src) { return CopyToTupleImpl(std::make_index_sequence>{}, std::forward(src)); } template TVM_FFI_INLINE static ObjectPtr CopyToArray(const Range& src) { return CopyToArrayImpl(std::begin(src), std::size(src)); } template TVM_FFI_INLINE static ObjectPtr MoveToArray(Range&& src) { return CopyToArrayImpl(std::make_move_iterator(std::begin(src)), std::size(src)); } }; template <> struct TypeTraits : public details::STLTypeTrait { public: static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIMap; protected: template TVM_FFI_INLINE static ObjectPtr CopyToMap(const MapType& src) { return MapObj::CreateFromRange(std::begin(src), std::end(src)); } template TVM_FFI_INLINE static ObjectPtr MoveToMap(MapType&& src) { return MapObj::CreateFromRange(std::make_move_iterator(std::begin(src)), std::make_move_iterator(std::end(src))); } template TVM_FFI_INLINE static MapType ConstructMap(const TVMFFIAny* src) { using KeyType = typename MapType::key_type; using ValueType = typename MapType::mapped_type; auto result = MapType{}; auto map_obj = CopyFromAnyImpl(src); if constexpr (CanReserve) { result.reserve(map_obj->size()); } for (const auto& [key, value] : *map_obj) { result.try_emplace(ConstructFromAny(key), ConstructFromAny(value)); } return result; } }; template struct TypeTraits> : public TypeTraits { private: using Self = std::array; TVM_FFI_INLINE static bool CheckAnyFast(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) return false; const ArrayObj& n = *reinterpret_cast(src->v_obj); return n.TVMFFISeqCell::size == Nm; } public: static_assert(Nm > 0, "Zero-length std::array is not supported."); TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return MoveToAnyImpl(CopyToArray(src), result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return MoveToAnyImpl(MoveToArray(std::move(src)), result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (!CheckAnyFast(src)) return std::nullopt; try { auto array = CopyFromAnyImpl(src); auto begin = array->MutableBegin(); Self result; // no initialization to avoid overhead for (std::size_t i = 0; i < Nm; ++i) { result[i] = ConstructFromAny(begin[i]); } return result; } catch (const details::STLTypeMismatch&) { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return "std::array<" + details::Type2Str::v() + ", " + std::to_string(Nm) + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::array","args":[)" + details::TypeSchema::v() + "," + std::to_string(Nm) + "]}"; } }; template struct TypeTraits> : public TypeTraits { private: using Self = std::vector; TVM_FFI_INLINE static bool CheckAnyFast(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIArray || src->type_index == TypeIndex::kTVMFFIList; } public: TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return MoveToAnyImpl(CopyToArray(src), result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return MoveToAnyImpl(MoveToArray(std::move(src)), result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (!CheckAnyFast(src)) return std::nullopt; try { const SeqBaseObj* seq = reinterpret_cast(src->v_obj); auto result = Self{}; int64_t length = static_cast(seq->size()); result.reserve(length); for (int64_t i = 0; i < length; ++i) { result.emplace_back(ConstructFromAny(seq->at(i))); } return result; } catch (const details::STLTypeMismatch&) { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return "std::vector<" + details::Type2Str::v() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::vector","args":[)" + details::TypeSchema::v() + "]}"; } }; template struct TypeTraits> : public TypeTraitsBase { public: using Self = std::optional; TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { if (src.has_value()) { TypeTraits::CopyToAnyView(*src, result); } else { TypeTraits::CopyToAnyView(nullptr, result); } } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { if (src.has_value()) { TypeTraits::MoveToAny(std::move(*src), result); } else { TypeTraits::MoveToAny(nullptr, result); } } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return true; return TypeTraits::CheckAnyStrict(src); } TVM_FFI_INLINE static Self CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return Self{std::nullopt}; return TypeTraits::CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static Self MoveFromAnyAfterCheck(TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return Self{std::nullopt}; return TypeTraits::MoveFromAnyAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return Self{std::nullopt}; auto result = std::optional{}; if (std::optional opt = TypeTraits::TryCastFromAnyView(src)) { /// NOTE: std::optional is just what we want (Self). result.emplace(std::move(opt)); } else { result.reset(); // failed to cast, indicate failure } return result; } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { return TypeTraits::GetMismatchTypeInfo(src); } TVM_FFI_INLINE static std::string TypeStr() { return "std::optional<" + TypeTraits::TypeStr() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::optional","args":[)" + details::TypeSchema::v() + "]}"; } }; template struct TypeTraits> : public TypeTraitsBase { private: using Self = std::variant; static constexpr std::size_t Nm = sizeof...(Args); template TVM_FFI_INLINE static Self CopyUnsafeAux(const TVMFFIAny* src) { if constexpr (Is >= Nm) { TVM_FFI_ICHECK(false) << "Unreachable: variant TryCast failed."; throw; // unreachable } else { using ElemType = std::variant_alternative_t; if (TypeTraits::CheckAnyStrict(src)) { return Self{std::in_place_index, TypeTraits::CopyFromAnyViewAfterCheck(src)}; } else { return CopyUnsafeAux(src); } } } template TVM_FFI_INLINE static Self MoveUnsafeAux(const TVMFFIAny* src) { if constexpr (Is >= Nm) { TVM_FFI_ICHECK(false) << "Unreachable: variant TryCast failed."; throw; // unreachable } else { using ElemType = std::variant_alternative_t; if (TypeTraits::CheckAnyStrict(src)) { return Self{std::in_place_index, TypeTraits::MoveFromAnyAfterCheck(src)}; } else { return MoveUnsafeAux(src); } } } template TVM_FFI_INLINE static std::optional TryCastAux(const TVMFFIAny* src) { if constexpr (Is >= Nm) { return std::nullopt; } else { using ElemType = std::variant_alternative_t; if (auto opt = TypeTraits::TryCastFromAnyView(src)) { return Self{std::in_place_index, std::move(*opt)}; } else { return TryCastAux(src); } } } public: TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return std::visit( [&](const auto& value) { using ValueType = std::decay_t; TypeTraits::CopyToAnyView(value, result); }, src); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return std::visit( [&](auto&& value) { using ValueType = std::decay_t; TypeTraits::MoveToAny(std::forward(value), result); }, std::move(src)); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return (TypeTraits::CheckAnyStrict(src) || ...); } TVM_FFI_INLINE static Self CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { // find the first possible type to copy return CopyUnsafeAux(src); } TVM_FFI_INLINE static Self MoveFromAnyAfterCheck(TVMFFIAny* src) { // find the first possible type to move return MoveUnsafeAux(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { // try to find the first possible type to copy return TryCastAux(src); } TVM_FFI_INLINE static std::string TypeStr() { std::ostringstream os; os << "std::variant<"; const char* sep = ""; ((os << sep << details::Type2Str::v(), sep = ", "), ...); os << ">"; return std::move(os).str(); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream os; os << R"({"type":"std::variant","args":[)"; const char* sep = ""; ((os << sep << details::TypeSchema::v(), sep = ", "), ...); os << "]}"; return std::move(os).str(); } }; template struct TypeTraits> : public TypeTraits { private: using Self = std::tuple; static constexpr std::size_t Nm = sizeof...(Args); static_assert(Nm > 0, "Zero-length std::tuple is not supported."); TVM_FFI_INLINE static bool CheckAnyFast(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) return false; const ArrayObj& n = *reinterpret_cast(src->v_obj); return n.TVMFFISeqCell::size == Nm; } template TVM_FFI_INLINE static Self ConstructTupleAux(std::index_sequence, const ArrayObj& n) { return Self{ConstructFromAny>(n[Is])...}; } public: static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIArray; TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return MoveToAnyImpl(CopyToTuple(src), result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return MoveToAnyImpl(MoveToTuple(std::move(src)), result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index != TypeIndex::kTVMFFIArray) return false; const ArrayObj& n = *reinterpret_cast(src->v_obj); // check static size first if (n.TVMFFISeqCell::size != Nm) return false; // then check element type return CheckSubTypeAux(std::make_index_sequence{}, n); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (!CheckAnyFast(src)) return std::nullopt; try { auto array = CopyFromAnyImpl(src); return ConstructTupleAux(std::make_index_sequence{}, *array); } catch (const details::STLTypeMismatch&) { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { std::ostringstream os; os << "std::tuple<"; const char* sep = ""; ((os << sep << details::Type2Str::v(), sep = ", "), ...); os << ">"; return std::move(os).str(); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream os; os << R"({"type":"std::tuple","args":[)"; const char* sep = ""; ((os << sep << details::TypeSchema::v(), sep = ", "), ...); os << "]}"; return std::move(os).str(); } }; template struct TypeTraits> : public TypeTraits { private: using Self = std::map; TVM_FFI_INLINE static bool CheckAnyFast(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIMap; } public: TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return MoveToAnyImpl(CopyToMap(src), result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return MoveToAnyImpl(MoveToMap(std::move(src)), result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (!CheckAnyFast(src)) return std::nullopt; try { return ConstructMap(src); } catch (const details::STLTypeMismatch&) { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return "std::map<" + details::Type2Str::v() + ", " + details::Type2Str::v() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::map","args":[)" + details::TypeSchema::v() + "," + details::TypeSchema::v() + "]}"; } }; template struct TypeTraits> : public TypeTraits { private: using Self = std::unordered_map; TVM_FFI_INLINE static bool CheckAnyFast(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIMap; } public: TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return MoveToAnyImpl(CopyToMap(src), result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return MoveToAnyImpl(MoveToMap(std::move(src)), result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (!CheckAnyFast(src)) return std::nullopt; try { return ConstructMap(src); } catch (const details::STLTypeMismatch&) { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return "std::unordered_map<" + details::Type2Str::v() + ", " + details::Type2Str::v() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::unordered_map","args":[)" + details::TypeSchema::v() + "," + details::TypeSchema::v() + "]}"; } }; template struct TypeTraits> : TypeTraitsBase { private: using Self = std::function; using Function = TypedFunction; using ProxyTrait = TypeTraits; TVM_FFI_INLINE static Self GetFunc(Function&& f) { return [fn = std::move(f)](Args... args) -> Ret { return fn(std::forward(args)...); }; } public: static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIFunction; static constexpr bool storage_enabled = false; TVM_FFI_INLINE static void CopyToAnyView(const Self& src, TVMFFIAny* result) { return ProxyTrait::MoveToAny(Function{src}, result); } TVM_FFI_INLINE static void MoveToAny(Self&& src, TVMFFIAny* result) { return ProxyTrait::MoveToAny(Function{std::move(src)}, result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { auto opt = ProxyTrait::TryCastFromAnyView(src); if (opt.has_value()) { return GetFunc(std::move(*opt)); } else { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { std::ostringstream os; os << "std::function<" << details::Type2Str::v() << "("; const char* sep = ""; ((os << sep << details::Type2Str::v(), sep = ", "), ...); os << ")>"; return std::move(os).str(); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream os; os << R"({"type":"std::function","args":[)" << details::TypeSchema::v() << ",["; const char* sep = ""; ((os << sep << details::TypeSchema::v(), sep = ", "), ...); os << "]]}"; return std::move(os).str(); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_STL_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/structural_equal.h000066400000000000000000000061301521067262500230410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/structural_equal.h * \brief Structural equal implementation */ #ifndef TVM_FFI_EXTRA_STRUCTURAL_EQUAL_H_ #define TVM_FFI_EXTRA_STRUCTURAL_EQUAL_H_ #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Structural equality comparators */ class StructuralEqual { public: /** * \brief Compare two Any values for structural equality. * \param lhs The left hand side Any object. * \param rhs The right hand side Any object. * \param map_free_vars Whether to map free variables. * \param skip_tensor_content Whether to skip comparingn darray data content, * useful for cases where we don't care about parameters content * \return True if the two Any values are structurally equal, false otherwise. */ TVM_FFI_EXTRA_CXX_API static bool Equal(const Any& lhs, const Any& rhs, bool map_free_vars = false, bool skip_tensor_content = false); /** * \brief Get the first mismatch AccessPath pair when running * structural equal comparison between two Any values. * * \param lhs The left hand side Any object. * \param rhs The right hand side Any object. * \param map_free_vars Whether to map free variables. * \param skip_tensor_content Whether to skip comparing tensor data content, * useful for cases where we don't care about parameters content * \return If comparison fails, return the first mismatch AccessPath pair, * otherwise return std::nullopt. */ TVM_FFI_EXTRA_CXX_API static Optional GetFirstMismatch( const Any& lhs, const Any& rhs, bool map_free_vars = false, bool skip_tensor_content = false); /* * \brief Compare two Any values for structural equality. * \param lhs The left hand side Any object. * \param rhs The right hand side Any object. * \return True if the two Any values are structurally equal, false otherwise. */ TVM_FFI_INLINE bool operator()(const Any& lhs, const Any& rhs) const { return Equal(lhs, rhs, false, true); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_STRUCTURAL_EQUAL_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/structural_hash.h000066400000000000000000000036061521067262500226620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/structural_hash.h * \brief Structural hash */ #ifndef TVM_FFI_EXTRA_STRUCTURAL_HASH_H_ #define TVM_FFI_EXTRA_STRUCTURAL_HASH_H_ #include #include namespace tvm { namespace ffi { /*! * \brief Structural hash */ class StructuralHash { public: /*! * \brief Hash an Any value. * \param value The Any value to hash. * \param map_free_vars Whether to map free variables. * \param skip_tensor_content Whether to skip comparingn darray data content, * useful for cases where we don't care about parameters content. * \return The hash value. */ TVM_FFI_EXTRA_CXX_API static uint64_t Hash(const Any& value, bool map_free_vars = false, bool skip_tensor_content = false); /*! * \brief Hash an Any value. * \param value The Any value to hash. * \return The hash value. */ TVM_FFI_INLINE uint64_t operator()(const Any& value) const { return Hash(value); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_STRUCTURAL_HASH_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/structural_key.h000066400000000000000000000061571521067262500225330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/structural_key.h * \brief Structural-key wrapper that caches structural hash. */ #ifndef TVM_FFI_EXTRA_STRUCTURAL_KEY_H_ #define TVM_FFI_EXTRA_STRUCTURAL_KEY_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Object node that contains a key and its cached structural hash. */ class StructuralKeyObj : public Object { public: /*! \brief The key value. */ Any key; /*! \brief Cached structural hash of \p key, encoded as int64_t. */ int64_t hash_i64{0}; // Default constructor to support reflection-based initialization. StructuralKeyObj() = default; /*! * \brief Construct a StructuralKeyObj from a key and cache its structural hash. * \param key The key value. */ explicit StructuralKeyObj(Any key) : key(std::move(key)), hash_i64(static_cast(StructuralHash::Hash(this->key))) {} /// \cond Doxygen_Suppress TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ffi.StructuralKey", StructuralKeyObj, Object); /// \endcond }; /*! * \brief ObjectRef wrapper for StructuralKeyObj. * * StructuralKey can be used to ensure that we use structural equality and hash for the wrapped key. */ class StructuralKey : public ObjectRef { public: /*! * \brief Construct a StructuralKey from a key. * \param key The key value. */ explicit StructuralKey(Any key) : ObjectRef(make_object(std::move(key))) {} bool operator==(const StructuralKey& other) const { if (this->same_as(other)) { return true; } if (this->get()->hash_i64 != other->hash_i64) { return false; } return StructuralEqual::Equal(this->get()->key, other->key); } bool operator!=(const StructuralKey& other) const { return !(*this == other); } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralKey, ObjectRef, StructuralKeyObj); /// \endcond }; } // namespace ffi } // namespace tvm namespace std { template <> struct hash { size_t operator()(const tvm::ffi::StructuralKey& key) const { return static_cast(static_cast(key->hash_i64)); } }; } // namespace std #endif // TVM_FFI_EXTRA_STRUCTURAL_KEY_H_ tvm-ffi-0.1.12/include/tvm/ffi/extra/visit_error_context.h000066400000000000000000000243551521067262500235660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/extra/visit_error_context.h * \brief VisitErrorContext: typed payload for Error::extra_context that records the * chain of ObjectRefs visited during a recursive visit when an error is thrown. */ #ifndef TVM_FFI_EXTRA_VISIT_ERROR_CONTEXT_H_ #define TVM_FFI_EXTRA_VISIT_ERROR_CONTEXT_H_ #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Object class for VisitErrorContext. * * \sa VisitErrorContext */ class VisitErrorContextObj : public Object { public: /*! * \brief Visit records that get populated, which include the object visit * path pattern in innermost-first order. Best-effort — not exhaustive. */ List reverse_visit_pattern; /*! * \brief Pre-existing Error::extra_context payload before we placed the * VisitErrorContext. */ Optional prev_error_context; /// \cond Doxygen_Suppress static constexpr bool _type_mutable = true; static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUnsupported; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ffi.VisitErrorContext", VisitErrorContextObj, Object); /// \endcond }; /*! * \brief Typed payload attached to Error::extra_context to support * visit-context-aware error reporting. * * The VisitErrorContext captures the reverse_visit_pattern — * the chain of nodes visited before an error was thrown — so callers * can translate it via FindAccessPaths into a structured access path * for richer error messages. * * Typical usage: * * 1. A recursive visit is instrumented with * TVM_FFI_VISIT_BEGIN / _END(node). The deepest detection * point throws via TVM_FFI_VISIT_THROW(kind, node), which * seeds the context with the throw site; enclosing BEGIN/END pairs * append parent nodes on rethrow. * * 2. The root catch handler retrieves the context via * TryGetFromError(err), then resolves the chain into one or more * reflection::AccessPath instances via FindAccessPaths(root, ctx). * * 3. The caller uses the AccessPath to enrich the error message * with structured position info (e.g., ".body[2].cond.lhs"). */ class VisitErrorContext : public ObjectRef { public: /*! \brief Get the VisitErrorContext attached to err's extra_context. * \param err The error to inspect. * \return The attached VisitErrorContext, or NullOpt if absent. */ TVM_FFI_COLD_CODE static Optional TryGetFromError(const Error& err) { std::optional extra_context = err.extra_context(); if (extra_context) { return extra_context->as(); } return std::nullopt; } /*! \brief Find all access paths that match the pattern specified in the * VisitErrorContext. * \param root The root ObjectRef to search from. * \param visit_context The VisitErrorContext to match against. * \param allow_prefix_match If true, also report paths where only a * prefix of the pattern was matched (i.e., * the algorithm descended through some * matched records but could not find further * matches before reaching a leaf). Default * false — only full pattern matches are * reported. * \return Array of matched access paths. */ TVM_FFI_COLD_CODE TVM_FFI_EXTRA_CXX_API static Array FindAccessPaths( const ObjectRef& root, const VisitErrorContext& visit_context, bool allow_prefix_match = false); /// \cond Doxygen_Suppress explicit VisitErrorContext(ObjectPtr n) : ObjectRef(std::move(n)) {} TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitErrorContext, ObjectRef, VisitErrorContextObj); /// \endcond }; /*! * \brief Begin a visit try block. * * Must be paired with TVM_FFI_VISIT_END(node) at the end of the * visit body. Expands to an open `try {` — a mismatched _END macro is a * compile error (unclosed try block). * * \code{.cpp} * void MyVisitor::VisitNode(const ObjectRef& node) { * TVM_FFI_VISIT_BEGIN(); * DispatchVisit(node); * TVM_FFI_VISIT_END(node); * } * \endcode */ #define TVM_FFI_VISIT_BEGIN() try { /*! * \brief End a visit try block and catch+re-throw any Error, * appending node to the VisitErrorContext on the way up. * * Must be paired with TVM_FFI_VISIT_BEGIN() above the visit body. * * \param node The ObjectRef at the current visit level (appended to the * error context's reverse_visit_pattern on exception). */ #define TVM_FFI_VISIT_END(node) \ } \ catch (::tvm::ffi::Error & _tvm_ffi_visit_err_) { \ ::tvm::ffi::details::UpdateVisitErrorContext(_tvm_ffi_visit_err_, (node)); \ throw; \ } /*! * \brief Throw an error from inside a visit, with `node` recorded * as the innermost frame of the resulting VisitErrorContext. * * Use this when the bad spot is somewhere the BEGIN/END pair does not * already record — typically a child field of the currently-visited node, * or a helper called from a visit that has no BEGIN/END of its own. The * throw site is seeded as the innermost frame; enclosing * TVM_FFI_VISIT_BEGIN/END pairs continue to append their nodes * on rethrow as the stack unwinds. The macro mirrors TVM_FFI_THROW — * it returns an ostream you stream a message into. * * If `node` here is the same as the enclosing END's node (a redundant * throw at the same level), FindAccessPaths normalizes the consecutive * duplicate during matching, so user code does not need to guard against * it — but in that case the throw would have been recorded by END * anyway and a plain TVM_FFI_THROW would suffice. * * \code{.cpp} * // Visiting a TPair node; pin the bad subfield (.lhs) as the throw * // site so the resulting AccessPath ends at .lhs, not at the TPair * // node itself. The surrounding END appends `node` as the next frame. * void Visitor::Visit(const ObjectRef& node) { * TVM_FFI_VISIT_BEGIN(); * if (auto pair = node.as()) { * if (!IsValid(pair.value()->lhs)) { * TVM_FFI_VISIT_THROW(ValueError, pair.value()->lhs) * << "invalid lhs"; * } * } * TVM_FFI_VISIT_END(node); * } * \endcode * * \param ErrorKind The kind of error to throw (e.g. TypeError, ValueError). * \param node The ObjectRef at the throw site (innermost frame). */ #define TVM_FFI_VISIT_THROW(ErrorKind, node) \ ::tvm::ffi::details::ErrorBuilder( \ #ErrorKind, TVMFFIBacktrace(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0), \ TVM_FFI_ALWAYS_LOG_BEFORE_THROW, ::std::nullopt, \ ::std::optional<::tvm::ffi::ObjectRef>(::tvm::ffi::details::MakeVisitErrorContext(node))) \ .stream() namespace details { /*! * \brief Build a fresh VisitErrorContext seeded with `node` as the * innermost (and currently only) frame of reverse_visit_pattern. * * Used by TVM_FFI_VISIT_THROW to attach the throw-site * node to the Error's extra_context at construction time. */ TVM_FFI_COLD_CODE inline VisitErrorContext MakeVisitErrorContext(const ObjectRef& node) { ObjectPtr obj = make_object(); obj->reverse_visit_pattern = List{node}; return VisitErrorContext(std::move(obj)); } /*! * \brief Implementation helper for TVM_FFI_VISIT_END(node). * Calling convention may change; do not call directly from user code. */ TVM_FFI_COLD_CODE inline void UpdateVisitErrorContext(Error& err, const ObjectRef& node) { // NOLINT(*) // NOTE: This function mutates the ErrorObj in place via ObjectUnsafe. // Expected to run only inside the exception throw chain, where the Error // is single-owned by this thread. The tradeoff avoids reallocating a // fresh Error per catch frame; the immutability invariant returns once // the unwind window closes. std::optional extra_context = err.extra_context(); if (extra_context) { Optional visit_context = extra_context->as(); if (visit_context) { visit_context.value()->reverse_visit_pattern.push_back(node); return; } } // Build a fresh VisitErrorContext, preserving any pre-existing payload. ObjectPtr new_context = make_object(); new_context->reverse_visit_pattern = List{node}; if (extra_context) new_context->prev_error_context = *extra_context; ErrorObj* error_obj = static_cast(details::ObjectUnsafe::RawObjectPtrFromObjectRef(err)); if (error_obj->extra_context != nullptr) { details::ObjectUnsafe::DecRefObjectHandle(error_obj->extra_context); } error_obj->extra_context = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(new_context)); } } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_VISIT_ERROR_CONTEXT_H_ tvm-ffi-0.1.12/include/tvm/ffi/function.h000066400000000000000000001212401521067262500201440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/function.h * \brief A managed function in the TVM FFI. */ #ifndef TVM_FFI_FUNCTION_H_ #define TVM_FFI_FUNCTION_H_ /*! * \brief Controls whether DLL exports should include metadata. * * When set to 1, exported functions will include additional metadata. * When set to 0 (default), exports are minimal without metadata. */ #ifndef TVM_FFI_DLL_EXPORT_INCLUDE_METADATA #define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 0 #endif #if TVM_FFI_DLL_EXPORT_INCLUDE_METADATA #include #endif // TVM_FFI_DLL_EXPORT_INCLUDE_METADATA #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /** * Helper macro to construct a safe call * * \brief Marks the beginning of the safe call that catches exception explicitly * \sa TVM_FFI_SAFE_CALL_END * * \code{.cpp} * int TVMFFICStyleFunction() { * TVM_FFI_SAFE_CALL_BEGIN(); * // c++ code region here * TVM_FFI_SAFE_CALL_END(); * } * \endcode */ #define TVM_FFI_SAFE_CALL_BEGIN() \ try { \ (void)0 /*! * \brief Marks the end of safe call. */ #define TVM_FFI_SAFE_CALL_END() \ return 0; \ } \ catch (const ::tvm::ffi::Error& err) { \ ::tvm::ffi::details::SetSafeCallRaised(err); \ return -1; \ } \ catch (const std::exception& ex) { \ ::tvm::ffi::details::SetSafeCallRaised(::tvm::ffi::Error("InternalError", ex.what(), "")); \ return -1; \ } \ TVM_FFI_UNREACHABLE() /*! * \brief Macro to check a call to TVMFFISafeCallType and raise exception if error happens. * \param func The function to check. * * \code{.cpp} * // calls TVMFFIFunctionCall and raises exception if error happens * TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_arr, &type_index)); * \endcode */ #define TVM_FFI_CHECK_SAFE_CALL(func) \ { \ int ret_code = (func); \ if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) { \ throw ::tvm::ffi::details::MoveFromSafeCallRaised(); \ } \ } /*! * \brief Object container class that backs ffi::Function * \note Do not use this class directly, use ffi::Function */ class FunctionObj : public Object, public TVMFFIFunctionCell { public: /*! \brief Typedef for C++ style calling signature that comes with exception propagation */ using FCall = void (*)(const FunctionObj*, const AnyView*, int32_t, Any*); using TVMFFIFunctionCell::cpp_call; using TVMFFIFunctionCell::safe_call; /*! * \brief Call the function in packed format. * \param args The arguments * \param num_args The number of arguments * \param result The return value. */ TVM_FFI_INLINE void CallPacked(const AnyView* args, int32_t num_args, Any* result) const { // if cpp_call is set, use it to call the function, otherwise, redirect to safe_call // use conditional expression here so the select is branchless FCall call_ptr = this->cpp_call ? reinterpret_cast(this->cpp_call) : CppCallDedirectToSafeCall; (*call_ptr)(this, args, num_args, result); } /// \cond Doxygen_Suppress static constexpr const uint32_t _type_index = TypeIndex::kTVMFFIFunction; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIFunction, FunctionObj, Object); /// \endcond protected: /*! \brief Make default constructor protected. */ FunctionObj() {} friend class Function; private: static void CppCallDedirectToSafeCall(const FunctionObj* func, const AnyView* args, int32_t num_args, Any* rv) { FunctionObj* self = static_cast(const_cast(func)); TVM_FFI_CHECK_SAFE_CALL(self->safe_call(self, reinterpret_cast(args), num_args, reinterpret_cast(rv))); } }; namespace details { /*! * \brief Derived object class for constructing FunctionObj backed by a TCallable * * This is a helper class that implements the function call interface. * Invariance: TCallable cannot be const or reference type. */ template class FunctionObjImpl : public FunctionObj { public: static_assert(std::is_same_v>>, "TCallable of FunctionObjImpl cannot be const or reference type"); /*! \brief The type of derived object class */ using TSelf = FunctionObjImpl; /*! * \brief Derived object class for constructing ffi::FunctionObj. * \param args The arguments to construct TCallable */ template explicit FunctionObjImpl(Args&&... args) : callable_(std::forward(args)...) { this->safe_call = SafeCall; this->cpp_call = reinterpret_cast(CppCall); } FunctionObjImpl(const FunctionObjImpl&) = delete; FunctionObjImpl& operator=(const FunctionObjImpl&) = delete; TCallable* GetCallable() { return &callable_; } private: // implementation of call static void CppCall(const FunctionObj* func, const AnyView* args, int32_t num_args, Any* result) { (static_cast(func))->callable_(args, num_args, result); } /// \cond Doxygen_Suppress // Implementing safe call style static int SafeCall(void* func, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { TVM_FFI_SAFE_CALL_BEGIN(); TVM_FFI_ICHECK_LT(result->type_index, TypeIndex::kTVMFFIStaticObjectBegin); FunctionObj* self = static_cast(func); reinterpret_cast(self->cpp_call)(self, reinterpret_cast(args), num_args, reinterpret_cast(result)); TVM_FFI_SAFE_CALL_END(); } /// \endcond /*! \brief Type-erased filed for storing callable object*/ mutable TCallable callable_; }; /*! * \brief FunctionObj specialization for raw C style callback where handle and deleter are null. */ class ExternCFunctionObjNullHandleImpl : public FunctionObj { public: explicit ExternCFunctionObjNullHandleImpl(TVMFFISafeCallType safe_call) { this->safe_call = safe_call; this->cpp_call = nullptr; } }; /*! * \brief FunctionObj specialization that leverages C-style callback definitions. */ class ExternCFunctionObjImpl : public FunctionObj { public: ExternCFunctionObjImpl(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void* self)) : self_(self), safe_call_(safe_call), deleter_(deleter) { this->safe_call = SafeCall; this->cpp_call = nullptr; } ~ExternCFunctionObjImpl() { if (deleter_) deleter_(self_); } private: static int32_t SafeCall(void* func, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* rv) { ExternCFunctionObjImpl* self = reinterpret_cast(func); return self->safe_call_(self->self_, args, num_args, rv); } void* self_; TVMFFISafeCallType safe_call_; void (*deleter_)(void* self); }; // Helper class to set packed arguments class PackedArgsSetter { public: explicit PackedArgsSetter(AnyView* args) : args_(args) {} // NOTE: setter needs to be very carefully designed // such that we do not have temp variable conversion(eg. convert from lvalue to rvalue) // that is why we need T&& and std::forward here template TVM_FFI_INLINE void operator()(size_t i, T&& value) const { args_[i].operator=(std::forward(value)); } private: AnyView* args_; }; } // namespace details /*! * \brief Represents arguments packed in AnyView array * \note This class represent packed arguments to ffi::Function */ class PackedArgs { public: /*! * \brief Constructor * \param data The arguments * \param size The number of arguments */ PackedArgs(const AnyView* data, int32_t size) : data_(data), size_(size) {} /*! \return size of the arguments */ int size() const { return size_; } /*! \return The arguments */ const AnyView* data() const { return data_; } /*! * \brief Slice the arguments * \param begin The begin index * \param end The end index * \return The sliced arguments */ PackedArgs Slice(int begin, int end = -1) const { if (end == -1) { end = size_; } return PackedArgs(data_ + begin, end - begin); } /*! * \brief Get i-th argument * \param i the index. * \return the ith argument. */ AnyView operator[](int i) const { return data_[i]; } /*! * \brief Fill the arguments into the AnyView array * \param data The AnyView array to store the packed arguments * \param args The arguments to be packed * \note Caller must ensure all args are alive during lifetime of data. * A common pitfall is to pass in local variables that are immediately * destroyed after calling Fill. */ template TVM_FFI_INLINE static void Fill(AnyView* data, Args&&... args) { details::for_each(details::PackedArgsSetter(data), std::forward(args)...); } private: /*! \brief The arguments */ const AnyView* data_; /*! \brief The number of arguments */ int32_t size_; }; /*! * \brief ffi::Function is a type-erased function. * The arguments are passed by "packed format" via AnyView */ class Function : public ObjectRef { public: /*! \brief Constructor from null */ Function(std::nullptr_t) : ObjectRef(nullptr) {} // NOLINT(*) /*! * \brief Constructing a packed function from a callable type * whose signature is consistent with `ffi::Function` * \param packed_call The packed function signature * \note legacy purpose, should change to Function::FromPacked for mostfuture use. */ template , Function>>> explicit Function(TCallable&& packed_call) { *this = FromPacked(std::forward(packed_call)); } /*! * \brief Constructing a packed function from a callable type * whose signature is consistent with `ffi::Function` * \param packed_call The packed function signature */ template static Function FromPacked(TCallable&& packed_call) { static_assert( std::is_convertible_v> || std::is_convertible_v>, "tvm::ffi::Function::FromPacked requires input function signature to match packed func " "format"); if constexpr (std::is_convertible_v>) { return FromPackedInternal( [packed_call = std::forward(packed_call)]( const AnyView* args, int32_t num_args, Any* rv) mutable -> void { packed_call(PackedArgs{args, num_args}, rv); }); } else { return FromPackedInternal(std::forward(packed_call)); } } /*! * \brief Constructing a packed function from a callable type * whose signature is consistent with `ffi::Function`. * It will create the Callable object with the given arguments, * and return the inplace constructed Function along with * the pointer to the callable object. The lifetime of the callable * object is managed by the returned Function. * \param args The arguments to construct TCallable * \return A tuple of (Function, TCallable*) */ template static auto FromPackedInplace(Args&&... args) { // We must ensure TCallable is a value type (decay_t) that can hold the callable object static_assert(std::is_same_v>); static_assert(std::is_invocable_v); using ObjType = details::FunctionObjImpl; Function func; auto obj_ptr = make_object(std::forward(args)...); auto* call_ptr = obj_ptr->GetCallable(); func.data_ = std::move(obj_ptr); return std::make_tuple(std::move(func), call_ptr); } /*! * \brief Create ffi::Function from a C style callbacks. * * self and deleter can be nullptr if the function do not need closure support * and corresponds to a raw function pointer. * * \param self Resource handle to the function * \param safe_call The safe_call definition in C. * \param deleter The deleter to release the resource of self. * \return The created function. */ static Function FromExternC(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void* self)) { // the other function coems from a different library Function func; if (self == nullptr && deleter == nullptr) { func.data_ = make_object(safe_call); } else { func.data_ = make_object(self, safe_call, deleter); } return func; } /*! * \brief Get global function by name * \param name The function name * \return The global function. * \note This function will return std::nullopt if the function is not found. */ static std::optional GetGlobal(std::string_view name) { TVMFFIObjectHandle handle; TVMFFIByteArray name_arr{name.data(), name.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFIFunctionGetGlobal(&name_arr, &handle)); if (handle != nullptr) { return Function( details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(handle))); } else { return std::nullopt; } } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will return std::nullopt if the function is not found. */ static std::optional GetGlobal(const std::string& name) { return GetGlobal(std::string_view(name.data(), name.length())); } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will return std::nullopt if the function is not found. */ static std::optional GetGlobal(const String& name) { return GetGlobal(std::string_view(name.data(), name.length())); } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will return std::nullopt if the function is not found. */ static std::optional GetGlobal(const char* name) { return GetGlobal(std::string_view(name)); } /*! * \brief Get global function by name and throw an error if it is not found. * \param name The name of the function * \return The global function * \note This function will throw an error if the function is not found. */ static Function GetGlobalRequired(std::string_view name) { std::optional res = GetGlobal(name); if (!res.has_value()) { TVM_FFI_THROW(ValueError) << "Function " << name << " not found"; } return *res; } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will throw an error if the function is not found. */ static Function GetGlobalRequired(const std::string& name) { return GetGlobalRequired(std::string_view(name.data(), name.length())); } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will throw an error if the function is not found. */ static Function GetGlobalRequired(const String& name) { return GetGlobalRequired(std::string_view(name.data(), name.length())); } /*! * \brief Get global function by name * \param name The name of the function * \return The global function * \note This function will throw an error if the function is not found. */ static Function GetGlobalRequired(const char* name) { return GetGlobalRequired(std::string_view(name)); } /*! * \brief Set global function by name * \param name The name of the function * \param func The function * \param override Whether to override when there is duplication. */ static void SetGlobal(std::string_view name, Function func, // NOLINT(performance-unnecessary-value-param) bool override = false) { TVMFFIByteArray name_arr{name.data(), name.size()}; TVM_FFI_CHECK_SAFE_CALL( TVMFFIFunctionSetGlobal(&name_arr, details::ObjectUnsafe::GetHeader(func.get()), override)); } /*! * \brief List all global names * \return A vector of all global names * \note This function do not depend on Array so core do not have container dep. */ static std::vector ListGlobalNames() { Function fname_functor = GetGlobalRequired("ffi.FunctionListGlobalNamesFunctor")().cast(); std::vector names; int len = fname_functor(-1).cast(); names.reserve(len); for (int i = 0; i < len; ++i) { names.push_back(fname_functor(i).cast()); } return names; } /** * \brief Remove a global function by name * \param name The name of the function */ static void RemoveGlobal(const String& name) { static Function fremove = GetGlobalRequired("ffi.FunctionRemoveGlobal"); fremove(name); } /*! * \brief Constructing a packed function from a normal function. * * \param callable the internal container of packed function. */ template static Function FromTyped(TCallable&& callable) { using FuncInfo = details::FunctionInfo>; // Callable is always captured by value here to avoid possible dangling reference auto call_packed = [callable = std::forward(callable)]( const AnyView* args, int32_t num_args, Any* rv) mutable -> void { details::unpack_call( std::make_index_sequence{}, nullptr, callable, args, num_args, rv); }; return FromPackedInternal(std::move(call_packed)); } /*! * \brief Constructing a packed function from a normal function. * * \param callable the internal container of packed function. * \param name optional name attacked to the function. */ template static Function FromTyped(TCallable&& callable, std::string name) { using FuncInfo = details::FunctionInfo>; // Callable is always captured by value here to avoid possible dangling reference auto call_packed = [callable = std::forward(callable), name = std::move(name)]( const AnyView* args, int32_t num_args, Any* rv) mutable -> void { details::unpack_call( std::make_index_sequence{}, &name, callable, args, num_args, rv); }; return FromPackedInternal(std::move(call_packed)); } /*! * \brief Directly invoke an extern "C" function that follows the TVM FFI SafeCall convention. * * This function can be useful to turn an existing exported symbol into a typed function. * * \code{.cpp} * // An extern "C" function, matching TVMFFISafeCallType * extern "C" int __tvm_ffi_add( * void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny*result * ); * // redirect an existing symbol into a typed function * inline int add(int a, int b) { * return tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add, a, b).cast(); * } * \endcode * * \tparam Args The types of the arguments to the extern function. * \param handle The handle argument, for exported symbols this is usually nullptr. * \param safe_call The function pointer to the extern "C" function. * \param args The arguments to pass to the function. * \return The return value, wrapped in a tvm::ffi::Any. */ template TVM_FFI_INLINE static Any InvokeExternC(void* handle, TVMFFISafeCallType safe_call, Args&&... args) { const int kNumArgs = sizeof...(Args); const int kArraySize = kNumArgs > 0 ? kNumArgs : 1; AnyView args_pack[kArraySize]; PackedArgs::Fill(args_pack, std::forward(args)...); Any result; TVM_FFI_CHECK_SAFE_CALL(safe_call(handle, reinterpret_cast(args_pack), kNumArgs, reinterpret_cast(&result))); return result; } /*! * \brief Call function by directly passing in unpacked arguments. * * \param args Arguments to be passed. * \tparam Args arguments to be passed. * * \code{.cpp} * // Example code on how to call packed function * void CallFFIFunction(tvm::ffi::Function f) { * // call like normal functions by pass in arguments * // return value is automatically converted back * int rvalue = f(1, 2.0); * } * \endcode */ template TVM_FFI_INLINE Any operator()(Args&&... args) const { const int kNumArgs = sizeof...(Args); const int kArraySize = kNumArgs > 0 ? kNumArgs : 1; AnyView args_pack[kArraySize]; PackedArgs::Fill(args_pack, std::forward(args)...); Any result; static_cast(data_.get())->CallPacked(args_pack, kNumArgs, &result); return result; } /*! * \brief Call the function in packed format. * \param args The arguments * \param num_args The number of arguments * \param result The return value. */ TVM_FFI_INLINE void CallPacked(const AnyView* args, int32_t num_args, Any* result) const { static_cast(data_.get())->CallPacked(args, num_args, result); } /*! * \brief Call the function in packed format. * \param args The arguments * \param result The return value. */ TVM_FFI_INLINE void CallPacked(PackedArgs args, Any* result) const { static_cast(data_.get())->CallPacked(args.data(), args.size(), result); } /*! * \brief Call the function and return Expected for exception-free error handling. * \tparam T The expected return type (default: Any). * \param args The arguments to pass to the function. * \return Expected containing either the result or an error. * * This method provides exception-free calling by catching all exceptions * and returning them as Error values in the Expected type. * * \code * Function func = Function::GetGlobal("risky_function"); * Expected result = func.CallExpected(arg1, arg2); * if (result.is_ok()) { * int value = result.value(); * } else { * Error err = result.error(); * } * \endcode */ template TVM_FFI_INLINE Expected CallExpected(Args&&... args) const { constexpr size_t kNumArgs = sizeof...(Args); AnyView args_pack[kNumArgs > 0 ? kNumArgs : 1]; PackedArgs::Fill(args_pack, std::forward(args)...); Any result; FunctionObj* func_obj = static_cast(data_.get()); // Use safe_call path to catch exceptions int ret_code = func_obj->safe_call(func_obj, reinterpret_cast(args_pack), kNumArgs, reinterpret_cast(&result)); if (ret_code == 0) { if constexpr (std::is_same_v) { return std::move(result); } else { // Try T first (fast path), then Error if (auto val = result.template try_cast()) { return *std::move(val); } if (auto err = result.template try_cast()) { return Unexpected(std::move(*err)); } return Unexpected(Error("TypeError", "CallExpected: result type mismatch, expected " + TypeTraits::TypeStr() + ", but got " + result.GetTypeKey(), "")); } } else { return Unexpected(details::MoveFromSafeCallRaised()); } } /*! \return Whether the packed function is nullptr */ TVM_FFI_INLINE bool operator==(std::nullptr_t) const { return data_ == nullptr; } /*! \return Whether the packed function is not nullptr */ TVM_FFI_INLINE bool operator!=(std::nullptr_t) const { return data_ != nullptr; } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Function, ObjectRef, FunctionObj); /// \endcond class Registry; private: /*! * \brief Constructing a packed function from a callable type * whose signature is consistent with `ffi::Function` * \param packed_call The packed function signature */ template static Function FromPackedInternal(TCallable&& packed_call) { // We must make TCallable a value type (decay_t) that can hold the callable object using ObjType = typename details::FunctionObjImpl>; Function func; func.data_ = make_object(std::forward(packed_call)); return func; } }; /*! * \brief Please refer to \ref TypedFunctionAnchor "TypedFunction" */ template class TypedFunction; /*! * \anchor TypedFunctionAnchor * \brief A ffi::Function wrapper to provide typed function signature. * It is backed by a ffi::Function internally. * * TypedFunction enables compile time type checking. * TypedFunction works with the runtime system: * - It can be passed as an argument of ffi::Function. * - It can be assigned to ffi::Any. * - It can be directly converted to a type-erased ffi::Function. * * Developers should prefer TypedFunction over ffi::Function in C++ code * as it enables compile time checking. * We can construct a TypedFunction from a lambda function * with the same signature. * * \code{.cpp} * // user defined lambda function. * auto addone = [](int x)->int { return x + 1; }; * // We can directly convert * // lambda function to TypedFunction * TypedFunction ftyped(addone); * // invoke the function. * int y = ftyped(1); * // Can be directly converted to ffi::Function * ffi::Function packed = ftype; * \endcode * \tparam R The return value of the function. * \tparam Args The argument signature of the function. */ template class TypedFunction { public: /*! \brief short hand for this function type */ using TSelf = TypedFunction; /*! \brief default constructor */ TypedFunction() = default; /*! \brief constructor from null */ TypedFunction(std::nullptr_t null) {} // NOLINT(*) /*! * \brief constructor from a function * \param packed The function */ TypedFunction(Function packed) : packed_(std::move(packed)) {} // NOLINT(*) /*! * \brief construct from a lambda function with the same signature. * * Example usage: * \code{.cpp} * auto typed_lambda = [](int x)->int { return x + 1; } * // construct from packed function * TypedFunction ftyped(typed_lambda, "add_one"); * // call the typed version. * CHECK_EQ(ftyped(1), 2); * \endcode * * \param typed_lambda typed lambda function. * \param name the name of the lambda function. * \tparam FLambda the type of the lambda function. */ template >>> TypedFunction(FLambda&& typed_lambda, std::string name) { packed_ = Function::FromTyped(std::forward(typed_lambda), std::move(name)); } /*! * \brief construct from a lambda function with the same signature. * * This version does not take a name. It is highly recommend you use the * version that takes a name for the lambda. * * Example usage: * \code{.cpp} * auto typed_lambda = [](int x)->int { return x + 1; } * // construct from packed function * TypedFunction ftyped(typed_lambda); * // call the typed version. * CHECK_EQ(ftyped(1), 2); * \endcode * * \param typed_lambda typed lambda function. * \tparam FLambda the type of the lambda function. */ template > && !std::is_same_v, TSelf>>> TypedFunction(FLambda&& typed_lambda) { // NOLINT(google-explicit-constructor) packed_ = Function::FromTyped(std::forward(typed_lambda)); } /*! * \brief copy assignment operator from typed lambda * * Example usage: * \code{.cpp} * // construct from packed function * TypedFunction ftyped; * ftyped = [](int x) { return x + 1; } * // call the typed version. * CHECK_EQ(ftyped(1), 2); * \endcode * * \param typed_lambda typed lambda function. * \tparam FLambda the type of the lambda function. * \returns reference to self. */ template > && !std::is_same_v, TSelf>>> TSelf& operator=(FLambda&& typed_lambda) { packed_ = Function::FromTyped(std::forward(typed_lambda)); return *this; } /*! * \brief copy assignment operator from ffi::Function. * \param packed The packed function. * \returns reference to self. */ TSelf& operator=(Function packed) { packed_ = std::move(packed); return *this; } /*! * \brief Invoke the operator. * \param args The arguments * \returns The return value. */ TVM_FFI_INLINE R operator()(Args... args) const { // NOLINT(performance-unnecessary-value-param) if constexpr (std::is_same_v) { packed_(std::forward(args)...); } else { Any res = packed_(std::forward(args)...); if constexpr (std::is_same_v) { return res; } else { return std::move(res).cast(); } } } /*! * \brief convert to ffi::Function * \return the internal ffi::Function */ operator Function() const { return packed(); } // NOLINT(google-explicit-constructor) /*! * \return reference the internal ffi::Function */ const Function& packed() const& { return packed_; } /*! * \return r-value reference the internal ffi::Function */ constexpr Function&& packed() && { return std::move(packed_); } /*! \return Whether the packed function is nullptr */ bool operator==(std::nullptr_t null) const { return packed_ == nullptr; } /*! \return Whether the packed function is not nullptr */ bool operator!=(std::nullptr_t null) const { return packed_ != nullptr; } /*! * \brief Get the type schema of `TypedFunction` in json format. * \return The type schema of the function in json format. */ static std::string TypeSchema() { return details::FuncFunctorImpl::TypeSchema(); } private: /*! \brief The internal packed function */ Function packed_; }; template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIFunction; TVM_FFI_INLINE static void CopyToAnyView(const TypedFunction& src, TVMFFIAny* result) { TypeTraits::CopyToAnyView(src.packed(), result); } TVM_FFI_INLINE static void MoveToAny(TypedFunction src, TVMFFIAny* result) { // Move from rvalue to trigger TypedFunction rvalue packed() overload TypeTraits::MoveToAny(std::move(src).packed(), result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIFunction; } TVM_FFI_INLINE static TypedFunction CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { return TypedFunction(TypeTraits::CopyFromAnyViewAfterCheck(src)); } TVM_FFI_INLINE static TypedFunction MoveFromAnyAfterCheck(TVMFFIAny* src) { return TypedFunction(TypeTraits::MoveFromAnyAfterCheck(src)); } TVM_FFI_INLINE static std::optional> TryCastFromAnyView( const TVMFFIAny* src) { std::optional opt = TypeTraits::TryCastFromAnyView(src); if (opt.has_value()) { return TypedFunction(*std::move(opt)); } else { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return details::FunctionInfo::Sig(); } TVM_FFI_INLINE static std::string TypeSchema() { return TypedFunction::TypeSchema(); } }; /*! * \brief helper function to get type index from key */ inline int32_t TypeKeyToIndex(std::string_view type_key) { int32_t type_index; TVMFFIByteArray type_key_array = {type_key.data(), type_key.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_array, &type_index)); return type_index; } /// \cond Doxygen_Suppress // Internal implementation macros used by TVM_FFI_DLL_EXPORT_TYPED_FUNC and related macros. // These should not be used directly; use the public macros instead. // Internal implementation macro that generates the C ABI wrapper function #define TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function) \ extern "C" { \ TVM_FFI_DLL_EXPORT int __tvm_ffi_##ExportName(void* self, const TVMFFIAny* args, \ int32_t num_args, TVMFFIAny* result) { \ TVM_FFI_SAFE_CALL_BEGIN(); \ using FuncInfo = ::tvm::ffi::details::FunctionInfo; \ static std::string name = #ExportName; \ ::tvm::ffi::details::unpack_call( \ std::make_index_sequence{}, &name, Function, \ reinterpret_cast(args), num_args, \ reinterpret_cast<::tvm::ffi::Any*>(result)); \ TVM_FFI_SAFE_CALL_END(); \ } \ } /// \endcond /*! * \brief Export typed function as a SafeCallType symbol that follows the FFI ABI. * * This macro exports the function and automatically exports metadata when * TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is defined. * * \param ExportName The symbol name to be exported. * \param Function The typed function. * * \sa ffi::TypedFunction, TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC * * \code{.cpp} * int AddOne_(int x) { * return x + 1; * } * // Expose the function as "AddOne" * TVM_FFI_DLL_EXPORT_TYPED_FUNC(AddOne, AddOne_); * // Expose the function as "SubOne" * TVM_FFI_DLL_EXPORT_TYPED_FUNC(SubOne, [](int x) { * return x - 1; * }); * \endcode * * \note The final symbol names are: * - `__tvm_ffi_` (function) * - `__tvm_ffi__metadata_` (metadata - only when * TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is defined) */ #if TVM_FFI_DLL_EXPORT_INCLUDE_METADATA // Implementation note: we specifically use TVMFFIStringFromByteArray // so the returned string metadata is allocated in the libtvm_ffi and long lived. #define TVM_FFI_DLL_EXPORT_TYPED_FUNC(ExportName, Function) \ TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function) \ extern "C" { \ TVM_FFI_DLL_EXPORT int __tvm_ffi__metadata_##ExportName(void* self, const TVMFFIAny* args, \ int32_t num_args, TVMFFIAny* result) { \ TVM_FFI_SAFE_CALL_BEGIN(); \ using FuncInfo = ::tvm::ffi::details::FunctionInfo; \ std::ostringstream os; \ os << R"({"type_schema":)" \ << ::tvm::ffi::EscapeStringJSON(::tvm::ffi::String(FuncInfo::TypeSchema())) << R"(})"; \ std::string data = os.str(); \ TVMFFIByteArray data_array{data.data(), data.size()}; \ return TVMFFIStringFromByteArray(&data_array, result); \ TVM_FFI_SAFE_CALL_END(); \ } \ } #else #define TVM_FFI_DLL_EXPORT_TYPED_FUNC(ExportName, Function) \ TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function) #endif /*! * \brief Export documentation string for a typed function. * * This macro exports a documentation string associated with a function export name. * The docstring can be used by stub generators and documentation tools. * This macro only exports the docstring; it does not export the function itself. * * \param ExportName The symbol name that the docstring is associated with. * \param DocString The documentation string (C string literal). * * \sa ffi::TypedFunction, TVM_FFI_DLL_EXPORT_TYPED_FUNC * * \code{.cpp} * int Add(int a, int b) { * return a + b; * } * * TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, Add); * TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC( * add, * R"(Add two integers and return the sum. * * Parameters * ---------- * a : int * First integer * b : int * Second integer * * Returns * ------- * result : int * Sum of a and b)"); * \endcode * * \note The exported symbol name is `__tvm_ffi__doc_` (docstring getter function). * This symbol is only exported when TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is defined. */ #if TVM_FFI_DLL_EXPORT_INCLUDE_METADATA // Implementation note: we specifically use TVMFFIStringFromByteArray // so the returned string metadata is allocated in the libtvm_ffi and long lived. #define TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(ExportName, DocString) \ extern "C" { \ TVM_FFI_DLL_EXPORT int __tvm_ffi__doc_##ExportName(void* self, const TVMFFIAny* args, \ int32_t num_args, TVMFFIAny* result) { \ TVM_FFI_SAFE_CALL_BEGIN(); \ std::string_view data(DocString); \ TVMFFIByteArray data_array{data.data(), data.size()}; \ return TVMFFIStringFromByteArray(&data_array, result); \ TVM_FFI_SAFE_CALL_END(); \ } \ } #else #define TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(ExportName, DocString) #endif } // namespace ffi } // namespace tvm #endif // TVM_FFI_FUNCTION_H_ tvm-ffi-0.1.12/include/tvm/ffi/function_details.h000066400000000000000000000243431521067262500216570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/function_details.h * \brief Implements the funciton signature reflection */ #ifndef TVM_FFI_FUNCTION_DETAILS_H_ #define TVM_FFI_FUNCTION_DETAILS_H_ #include #include #include #include #include #include #include namespace tvm { namespace ffi { // Forward declaration for Expected template class Expected; namespace details { template struct Arg2Str { template TVM_FFI_INLINE static void Apply(std::ostream& os) { using Arg = std::tuple_element_t; if constexpr (i != 0) { os << ", "; } os << i << ": " << Type2Str::v(); } template TVM_FFI_INLINE static void Run(std::ostream& os, std::index_sequence) { using TExpander = int[]; (void)TExpander{0, (Apply(os), 0)...}; } }; /// NOTE: We only support `T`, `const T`, `const T&` and `T&&` as argument types. template static constexpr bool ArgTypeSupported = (!std::is_reference_v) || (std::is_const_v> && std::is_lvalue_reference_v) || (!std::is_const_v> && std::is_rvalue_reference_v); template static constexpr bool ArgSupported = (ArgTypeSupported && (std::is_same_v>, Any> || std::is_same_v>, AnyView> || TypeTraitsNoCR::convert_enabled)); // NOTE: return type can only support non-reference managed returns template static constexpr bool RetSupported = (std::is_same_v || std::is_void_v || TypeTraits::convert_enabled); template struct FuncFunctorImpl { using FType = R(Args...); using ArgType = std::tuple; using RetType = R; /*! \brief total number of arguments*/ static constexpr size_t num_args = sizeof...(Args); // MSVC is not that friendly to in-template nested bool evaluation #ifndef _MSC_VER /*! \brief Whether this function can be converted to ffi::Function via FromTyped */ static constexpr bool unpacked_supported = (ArgSupported && ...) && (RetSupported); #endif TVM_FFI_INLINE static std::string Sig() { using IdxSeq = std::make_index_sequence; std::ostringstream ss; ss << "("; Arg2Str>::Run(ss, IdxSeq{}); ss << ") -> " << Type2Str::v(); return ss.str(); } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << StaticTypeKey::kTVMFFIFunction << R"(","args":[)"; oss << details::TypeSchema::v(); ((oss << "," << details::TypeSchema::v()), ...); oss << "]}"; return oss.str(); } }; template struct FunctionInfoHelper; template struct FunctionInfoHelper : FuncFunctorImpl {}; template struct FunctionInfoHelper : FuncFunctorImpl {}; /*! * \brief Template class to get function signature of a function or functor. * \tparam T The function/functor type. * \note We need a decltype redirection because this helps lambda types. */ template struct FunctionInfo : FunctionInfoHelper {}; template struct FunctionInfo : FuncFunctorImpl {}; template struct FunctionInfo : FuncFunctorImpl {}; template struct FunctionInfo : FuncFunctorImpl {}; // Support pointer-to-member functions used in reflection (e.g. &Class::method) template struct FunctionInfo>> : FuncFunctorImpl {}; template struct FunctionInfo>> : FuncFunctorImpl {}; template struct FunctionInfo>> : FuncFunctorImpl {}; template struct FunctionInfo>> : FuncFunctorImpl {}; /*! \brief Using static function to output typed function signature */ using FGetFuncSignature = std::string (*)(); /*! * \brief Auxilary argument value with context for error reporting * \tparam Type The expected type of the argument. * \note We use a template class with non-template operator conversion * instead of a non-template class with template operator conversion. * This is because template operator conversion doesn't play well with * classes with template constructors. * In this case, it may lead to some unintended compiler errors. * An example of class can be `std::optional`. */ template class ArgValueWithContext { public: using TypeWithoutCR = std::remove_const_t>; /*! * \brief move constructor from another return value. * \param args The argument list * \param arg_index In a function call, this argument is at index arg_index (0-indexed). * \param optional_name Name of the function being called. Can be nullptr if the function is not. * \param f_sig Pointer to static function outputting signature of the function being called. * named. */ TVM_FFI_INLINE ArgValueWithContext(const AnyView* args, int32_t arg_index, const std::string* optional_name, FGetFuncSignature f_sig) : args_(args), arg_index_(arg_index), optional_name_(optional_name), f_sig_(f_sig) {} TVM_FFI_INLINE operator TypeWithoutCR() { // NOLINT(google-explicit-constructor) if constexpr (std::is_same_v) { return args_[arg_index_]; } else if constexpr (std::is_same_v) { return Any(args_[arg_index_]); } else { std::optional opt = args_[arg_index_].template try_cast(); if (!opt.has_value()) { TVMFFIAny any_data = args_[arg_index_].CopyToTVMFFIAny(); TVM_FFI_THROW(TypeError) << "Mismatched type on argument #" << arg_index_ << " when calling: `" << (optional_name_ == nullptr ? "" : *optional_name_) << (f_sig_ == nullptr ? "" : (*f_sig_)()) << "`. Expected `" << Type2Str::v() << "` but got `" << TypeTraits::GetMismatchTypeInfo(&any_data) << '`'; } return *std::move(opt); } } private: const AnyView* args_; int32_t arg_index_; const std::string* optional_name_; FGetFuncSignature f_sig_; }; template TVM_FFI_INLINE void unpack_call(std::index_sequence, const std::string* optional_name, const F& f, [[maybe_unused]] const AnyView* args, [[maybe_unused]] int32_t num_args, [[maybe_unused]] Any* rv) { using FuncInfo = FunctionInfo; using PackedArgs = typename FuncInfo::ArgType; FGetFuncSignature f_sig = FuncInfo::Sig; // somehow MSVC does not support the static constexpr member in this case, function is fine #ifndef _MSC_VER static_assert(FuncInfo::unpacked_supported, "The function signature do not support unpacked"); #endif constexpr size_t nargs = sizeof...(Is); if (nargs != num_args) { TVM_FFI_THROW(TypeError) << "Mismatched number of arguments when calling: `" << (optional_name == nullptr ? "" : *optional_name) << (f_sig == nullptr ? "" : (*f_sig)()) << "`. Expected " << nargs << " but got " << num_args << " arguments"; } // use index sequence to do recursive-less unpacking if constexpr (std::is_same_v) { f(ArgValueWithContext>{args, Is, optional_name, f_sig}...); } else { *rv = R(f(ArgValueWithContext>{args, Is, optional_name, f_sig}...)); } // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks } template struct TypeSchemaImpl { static std::string v() { using U = std::remove_const_t>; return TypeTraits::TypeSchema(); } }; template <> struct TypeSchemaImpl { static std::string v() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFINone) + R"("})"; } }; template <> struct TypeSchemaImpl { static std::string v() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIAny) + R"("})"; } }; template <> struct TypeSchemaImpl { static std::string v() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIAny) + R"("})"; } }; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_FUNCTION_DETAILS_H_ tvm-ffi-0.1.12/include/tvm/ffi/memory.h000066400000000000000000000240261521067262500176330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/memory.h * \brief Runtime memory management to allocate on heap object. */ #ifndef TVM_FFI_MEMORY_H_ #define TVM_FFI_MEMORY_H_ #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief Deleter function for obeject */ using FObjectDeleter = void (*)(void* obj, int flags); // Detail implementations after this // // The current design allows swapping the // allocator pattern when necessary. // // Possible future allocator optimizations: // - Arena allocator that gives ownership of memory to arena (deleter = nullptr) // - Thread-local object pools: one pool per size and alignment requirement. // - Can specialize by type of object to give the specific allocator to each object. namespace details { /*! * \brief Allocate aligned memory. * \param size The size. * \tparam align The alignment, must be a power of 2. * \return The pointer to the allocated memory. */ template TVM_FFI_INLINE void* AlignedAlloc(size_t size) { static_assert(align != 0 && (align & (align - 1)) == 0, "align must be a power of 2"); #ifdef _MSC_VER // MSVC have to use _aligned_malloc if (void* ptr = _aligned_malloc(size, align)) { return ptr; } throw std::bad_alloc(); #else if constexpr (align <= alignof(std::max_align_t)) { // malloc guarantees alignment of std::max_align_t if (void* ptr = std::malloc(size)) { return ptr; } throw std::bad_alloc(); } else { void* ptr; // for other alignments, use posix_memalign if (posix_memalign(&ptr, align, size) != 0) { throw std::bad_alloc(); } return ptr; } #endif } /*! * \brief Free aligned memory. * \param data The pointer to the memory to free. */ TVM_FFI_INLINE void AlignedFree(void* data) { #ifdef _MSC_VER // MSVC have to use _aligned_free _aligned_free(data); #else std::free(data); #endif } /*! * \brief Base class of object allocators that implements make. * Use curiously recurring template pattern. * * \tparam Derived The derived class. */ template class ObjAllocatorBase { public: /*! * \brief Make a new object using the allocator. * \tparam T The type to be allocated. * \tparam Args The constructor signature. * \param args The arguments. */ template ObjectPtr make_object(Args&&... args) { using Handler = typename Derived::template Handler; static_assert(std::is_base_of_v, "make can only be used to create Object"); T* ptr = Handler::New(static_cast(this), std::forward(args)...); TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr); ffi_ptr->combined_ref_count = kCombinedRefCountBothOne; ffi_ptr->type_index = T::RuntimeTypeIndex(); ffi_ptr->__padding = 0; ffi_ptr->deleter = Handler::Deleter(); return details::ObjectUnsafe::ObjectPtrFromOwned(ptr); } /*! * \tparam ArrayType The type to be allocated. * \tparam ElemType The type of array element. * \tparam Args The constructor signature. * \param num_elems The number of array elements. * \param args The arguments. */ template ObjectPtr make_inplace_array(size_t num_elems, Args&&... args) { using Handler = typename Derived::template ArrayHandler; static_assert(std::is_base_of_v, "make_inplace_array can only be used to create Object"); ArrayType* ptr = Handler::New(static_cast(this), num_elems, std::forward(args)...); TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr); ffi_ptr->combined_ref_count = kCombinedRefCountBothOne; ffi_ptr->type_index = ArrayType::RuntimeTypeIndex(); ffi_ptr->__padding = 0; ffi_ptr->deleter = Handler::Deleter(); return details::ObjectUnsafe::ObjectPtrFromOwned(ptr); } private: ObjAllocatorBase() = default; friend Derived; }; // Simple allocator that uses new/delete. class SimpleObjAllocator : public ObjAllocatorBase { public: template class Handler { public: template static T* New(SimpleObjAllocator*, Args&&... args) { // NOTE: the first argument is not needed for SimpleObjAllocator // It is reserved for special allocators that needs to recycle // the object to itself (e.g. in the case of object pool). // // In the case of an object pool, an allocator needs to create // a special chunk memory that hides reference to the allocator // and call allocator's release function in the deleter. // NOTE2: Use inplace new to allocate // This is used to get rid of warning when deleting a virtual // class with non-virtual destructor. // We are fine here as we captured the right deleter during construction. // This is also the right way to get storage type for an object pool. void* data = AlignedAlloc(sizeof(T)); new (data) T(std::forward(args)...); return reinterpret_cast(data); } static FObjectDeleter Deleter() { return Deleter_; } private: static void Deleter_(void* objptr, int flags) { T* tptr = details::ObjectUnsafe::RawObjectPtrFromUnowned(static_cast(objptr)); if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) { // It is important to do tptr->T::~T(), // so that we explicitly call the specific destructor // instead of tptr->~T(), which could mean the intention // call a virtual destructor(which may not be available and is not required). tptr->T::~T(); } if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) { AlignedFree(static_cast(tptr)); } } }; // Array handler that uses new/delete. template class ArrayHandler { public: template static ArrayType* New(SimpleObjAllocator*, size_t num_elems, Args&&... args) { // NOTE: the first argument is not needed for ArrayObjAllocator // It is reserved for special allocators that needs to recycle // the object to itself (e.g. in the case of object pool). // // In the case of an object pool, an allocator needs to create // a special chunk memory that hides reference to the allocator // and call allocator's release function in the deleter. // NOTE2: Use inplace new to allocate // This is used to get rid of warning when deleting a virtual // class with non-virtual destructor. // We are fine here as we captured the right deleter during construction. // This is also the right way to get storage type for an object pool. // for now only support elements that aligns with array header. static_assert( alignof(ArrayType) % alignof(ElemType) == 0 && sizeof(ArrayType) % alignof(ElemType) == 0, "element alignment constraint"); size_t size = sizeof(ArrayType) + sizeof(ElemType) * num_elems; // round up to the nearest multiple of align constexpr size_t align = alignof(ArrayType); // C++ standard always guarantees that alignof operator returns a power of 2 size_t aligned_size = (size + (align - 1)) & ~(align - 1); void* data = AlignedAlloc(aligned_size); new (data) ArrayType(std::forward(args)...); return reinterpret_cast(data); } static FObjectDeleter Deleter() { return Deleter_; } private: static void Deleter_(void* objptr, int flags) { ArrayType* tptr = details::ObjectUnsafe::RawObjectPtrFromUnowned( static_cast(objptr)); if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) { // It is important to do tptr->ArrayType::~ArrayType(), // so that we explicitly call the specific destructor // instead of tptr->~ArrayType(), which could mean the intention // call a virtual destructor(which may not be available and is not required). tptr->ArrayType::~ArrayType(); } if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) { AlignedFree(static_cast(tptr)); } } }; }; } // namespace details /*! * \brief Allocate an object * \param args arguments to the constructor. * \tparam T the node type. * \return The ObjectPtr to the allocated object. */ template inline ObjectPtr make_object(Args&&... args) { return details::SimpleObjAllocator().make_object(std::forward(args)...); } /*! * \brief Allocate an Object with additional ElemType[num_elems] that are stored right after. * \param num_elems The number of elements in the array. * \param args arguments to the constructor. * \tparam ArrayType the array type. * \tparam ElemType the element type. * \return The ObjectPtr to the allocated array. */ template inline ObjectPtr make_inplace_array_object(size_t num_elems, Args&&... args) { return details::SimpleObjAllocator().make_inplace_array( num_elems, std::forward(args)...); } } // namespace ffi } // namespace tvm #endif // TVM_FFI_MEMORY_H_ tvm-ffi-0.1.12/include/tvm/ffi/object.h000066400000000000000000001330301521067262500175650ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/object.h * \brief A managed object in the TVM FFI. */ #ifndef TVM_FFI_OBJECT_H_ #define TVM_FFI_OBJECT_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief TypeIndex enum, alias of TVMFFITypeIndex. */ using TypeIndex = TVMFFITypeIndex; /*! * \brief TypeInfo, alias of TVMFFITypeInfo. */ using TypeInfo = TVMFFITypeInfo; /*! * \brief Helper tag to explicitly request unsafe initialization. * * Constructing an ObjectRefType with UnsafeInit{} will set the data_ member to nullptr. * * When initializing Object fields, ObjectRef fields can be set to UnsafeInit. * This enables the "construct with UnsafeInit then set all fields" pattern * when the object does not have a default constructor. * * Used for initialization in controlled scenarios where such unsafe * initialization is known to be safe. * * Each ObjectRefType should have a constructor that takes an UnsafeInit tag. * * \note As the name suggests, do not use it in normal code paths. */ struct UnsafeInit {}; /*! * \brief Known type keys for pre-defined types. */ struct StaticTypeKey { /*! \brief The type key for Any */ static constexpr const char* kTVMFFIAny = "Any"; /*! \brief The type key for None */ static constexpr const char* kTVMFFINone = "None"; /*! \brief The type key for bool */ static constexpr const char* kTVMFFIBool = "bool"; /*! \brief The type key for int */ static constexpr const char* kTVMFFIInt = "int"; /*! \brief The type key for float */ static constexpr const char* kTVMFFIFloat = "float"; /*! \brief The type key for void* */ static constexpr const char* kTVMFFIOpaquePtr = "void*"; /*! \brief The type key for DataType */ static constexpr const char* kTVMFFIDataType = "DataType"; /*! \brief The type key for Device */ static constexpr const char* kTVMFFIDevice = "Device"; /*! \brief The type key for DLTensor* */ static constexpr const char* kTVMFFIDLTensorPtr = "DLTensor*"; /*! \brief The type key for const char* */ static constexpr const char* kTVMFFIRawStr = "const char*"; /*! \brief The type key for TVMFFIByteArray* */ static constexpr const char* kTVMFFIByteArrayPtr = "TVMFFIByteArray*"; /*! \brief The type key for ObjectRValueRef */ static constexpr const char* kTVMFFIObjectRValueRef = "ObjectRValueRef"; /*! \brief The type key for SmallStr */ static constexpr const char* kTVMFFISmallStr = "ffi.SmallStr"; /*! \brief The type key for SmallBytes */ static constexpr const char* kTVMFFISmallBytes = "ffi.SmallBytes"; /*! \brief The type key for Error */ static constexpr const char* kTVMFFIError = "ffi.Error"; /*! \brief The type key for Bytes */ static constexpr const char* kTVMFFIBytes = "ffi.Bytes"; /*! \brief The type key for String */ static constexpr const char* kTVMFFIStr = "ffi.String"; /*! \brief The type key for Shape */ static constexpr const char* kTVMFFIShape = "ffi.Shape"; /*! \brief The type key for Tensor */ static constexpr const char* kTVMFFITensor = "ffi.Tensor"; /*! \brief The type key for Object */ static constexpr const char* kTVMFFIObject = "ffi.Object"; /*! \brief The type key for Function */ static constexpr const char* kTVMFFIFunction = "ffi.Function"; /*! \brief The type key for Array */ static constexpr const char* kTVMFFIArray = "ffi.Array"; /*! \brief The type key for List */ static constexpr const char* kTVMFFIList = "ffi.List"; /*! \brief The type key for Map */ static constexpr const char* kTVMFFIMap = "ffi.Map"; /*! \brief The type key for Module */ static constexpr const char* kTVMFFIModule = "ffi.Module"; /*! \brief The type key for Dict */ static constexpr const char* kTVMFFIDict = "ffi.Dict"; /*! \brief The type key for OpaquePyObject */ static constexpr const char* kTVMFFIOpaquePyObject = "ffi.OpaquePyObject"; }; /*! * \brief Get type key from type index * \param type_index The input type index * \return the type key */ inline std::string TypeIndexToTypeKey(int32_t type_index) { const TypeInfo* type_info = TVMFFIGetTypeInfo(type_index); return std::string(type_info->type_key.data, type_info->type_key.size); } namespace details { // Helper to perform // unsafe operations related to object struct ObjectUnsafe; /*! \brief One counter for weak reference. */ constexpr uint64_t kCombinedRefCountWeakOne = static_cast(1) << 32; /*! \brief One counter for strong reference. */ constexpr uint64_t kCombinedRefCountStrongOne = 1; /*! \brief Both reference counts. */ constexpr uint64_t kCombinedRefCountBothOne = kCombinedRefCountWeakOne | kCombinedRefCountStrongOne; /*! \brief Mask to get the lower 32 bits of the combined reference count. */ constexpr uint64_t kCombinedRefCountMaskUInt32 = (static_cast(1) << 32) - 1; /*! * Check if the type_index is an instance of TargetObjectType. * * \tparam TargetType The target object type to be checked. * * \param object_type_index The type index to be checked, caller * ensures that the index is already within the object index range. * * \return Whether the target type is true. */ template TVM_FFI_INLINE bool IsObjectInstance(int32_t object_type_index); } // namespace details /*! * \brief Base class of all object containers. * * Sub-class of objects should declare the following static constexpr fields: * * - _type_index: * Static type index of the object, if assigned to TypeIndex::kTVMFFIDynObject * the type index will be assigned during runtime. * Runtime type index can be accessed by ObjectType::TypeIndex(); * - _type_key: * The unique string identifier of the type. * - _type_final: * Whether the type is terminal type(there is no subclass of the type in the object system). * This field is automatically set by macro TVM_FFI_DECLARE_OBJECT_INFO_FINAL * It is still OK to sub-class a terminal object type T and construct it using make_object. * But IsInstance check will only show that the object type is T(instead of the sub-class). * - _type_mutable: * Whether we would like to expose cast to non-constant pointer * ObjectType* from Any/AnyView. By default, we set to false so it is not exposed. * * The following two fields are necessary for base classes that can be sub-classed. * * - _type_child_slots: * Number of reserved type index slots for child classes. * Used for runtime optimization for type checking in IsInstance. * If an object's type_index is within range of [type_index, type_index + _type_child_slots] * Then the object can be quickly decided as sub-class of the current object class. * If not, a fallback mechanism is used to check the global type table. * Recommendation: set to estimate number of children needed. * * - _type_child_slots_can_overflow: * Whether we can add additional child classes even if the number of child classes * exceeds the _type_child_slots. A fallback mechanism to check type table will be used. * Recommendation: set to false for optimal runtime speed if we know exact number of children. * * Two macros are used to declare helper functions in the object: * - Use TVM_FFI_DECLARE_OBJECT_INFO for object classes that can be sub-classed. * - Use TVM_FFI_DECLARE_OBJECT_INFO_FINAL for object classes that cannot be sub-classed. * * New objects can be created using make_object function. * Which will automatically populate the type_index and deleter of the object. */ class Object { protected: /*! \brief header field that is the common prefix of all objects */ TVMFFIObject header_; public: Object() { header_.combined_ref_count = 0; header_.type_index = 0; header_.__padding = 0; header_.__ensure_align = 0; } /*! * Check if the object is an instance of TargetType. * \tparam TargetType The target type to be checked. * \return Whether the target type is true. */ template bool IsInstance() const { return details::IsObjectInstance(header_.type_index); } /*! \return The internal runtime type index of the object. */ int32_t type_index() const { return header_.type_index; } /*! * \return the type key of the object. * \note this operation is expensive, can be used for error reporting. */ std::string GetTypeKey() const { // the function checks that the info exists const TypeInfo* type_info = TVMFFIGetTypeInfo(header_.type_index); return std::string(type_info->type_key.data, type_info->type_key.size); } /*! * \return A hash value of the return of GetTypeKey. */ uint64_t GetTypeKeyHash() const { // the function checks that the info exists const TypeInfo* type_info = TVMFFIGetTypeInfo(header_.type_index); return type_info->type_key_hash; } /*! * \brief Get the type key of the corresponding index from runtime. * \param tindex The type index. * \return the result. */ static std::string TypeIndex2Key(int32_t tindex) { const TypeInfo* type_info = TVMFFIGetTypeInfo(tindex); return std::string(type_info->type_key.data, type_info->type_key.size); } /*! * \return Whether the object.use_count() == 1. */ bool unique() const { return use_count() == 1; } /*! * \return The usage count of the cell. * \note We use STL style naming to be consistent with known API in shared_ptr. */ uint64_t use_count() const { // only need relaxed load of counters #ifdef _MSC_VER return ((reinterpret_cast( &header_.combined_ref_count))[0] // NOLINT(*) ) & kCombinedRefCountMaskUInt32; #else return __atomic_load_n(&(header_.combined_ref_count), __ATOMIC_RELAXED) & kCombinedRefCountMaskUInt32; #endif } //---------------------------------------------------------------------------- // The following fields are configuration flags for subclasses of object //---------------------------------------------------------------------------- /*! \brief The type key of the class */ static constexpr const char* _type_key = StaticTypeKey::kTVMFFIObject; /*! \brief Whether the class is final */ static constexpr bool _type_final = false; /*! \brief Whether allow mutable access to fields */ static constexpr bool _type_mutable = false; /*! \brief The number of child slots of the class to pre-allocate to this type */ static constexpr uint32_t _type_child_slots = 0; /*! * \brief Whether allow additional children beyond pre-specified by _type_child_slots */ static constexpr bool _type_child_slots_can_overflow = true; /*! \brief The static type index of the class */ static constexpr int32_t _type_index = TypeIndex::kTVMFFIObject; /*! \brief The static depth of the class in the object hierarchy */ static constexpr int32_t _type_depth = 0; /*! \brief The structural equality and hash kind of the type */ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUnsupported; // The following functions are provided by macro // TVM_FFI_DECLARE_OBJECT_INFO and TVM_FFI_DECLARE_OBJECT_INFO_FINAL /*! * \brief Get the runtime allocated type index of the type * \note Getting this information may need dynamic calls into a global table. */ static int32_t RuntimeTypeIndex() { return TypeIndex::kTVMFFIObject; } /*! * \brief Internal function to get or allocate a runtime index. */ static int32_t _GetOrAllocRuntimeTypeIndex() { // NOLINT(bugprone-reserved-identifier) return TypeIndex::kTVMFFIObject; } private: // exposing detailed constants to here static constexpr uint64_t kCombinedRefCountMaskUInt32 = details::kCombinedRefCountMaskUInt32; static constexpr uint64_t kCombinedRefCountStrongOne = details::kCombinedRefCountStrongOne; static constexpr uint64_t kCombinedRefCountWeakOne = details::kCombinedRefCountWeakOne; static constexpr uint64_t kCombinedRefCountBothOne = details::kCombinedRefCountBothOne; /*! \brief increase strong reference count, the caller must already hold a strong reference */ void IncRef() { #ifdef _MSC_VER _InterlockedIncrement64( reinterpret_cast(&header_.combined_ref_count)); // NOLINT(*) #else __atomic_fetch_add(&(header_.combined_ref_count), 1, __ATOMIC_RELAXED); #endif } /*! * \brief Try to lock the object to increase the strong reference count, * the caller must already hold a strong reference. * \return whether the lock call is successful and object is still alive. */ bool TryPromoteWeakPtr() { #ifdef _MSC_VER uint64_t old_count = (reinterpret_cast(&header_.combined_ref_count))[0]; // NOLINT(*) while ((old_count & kCombinedRefCountMaskUInt32) != 0) { uint64_t new_count = old_count + kCombinedRefCountStrongOne; uint64_t old_count_loaded = _InterlockedCompareExchange64( reinterpret_cast(&header_.combined_ref_count), new_count, old_count); if (old_count == old_count_loaded) { return true; } old_count = old_count_loaded; } return false; #else uint64_t old_count = __atomic_load_n(&(header_.combined_ref_count), __ATOMIC_RELAXED); while ((old_count & kCombinedRefCountMaskUInt32) != 0) { // must do CAS to ensure that we are the only one that increases the reference count // avoid condition when two threads tries to promote weak to strong at same time // or when strong deletion happens between the load and the CAS uint64_t new_count = old_count + kCombinedRefCountStrongOne; if (__atomic_compare_exchange_n(&(header_.combined_ref_count), &old_count, new_count, true, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) { return true; } } return false; #endif } /*! \brief increase weak reference count */ void IncWeakRef() { #ifdef _MSC_VER _InlineInterlockedAdd64( reinterpret_cast(&header_.combined_ref_count), // NOLINT(*) kCombinedRefCountWeakOne); #else __atomic_fetch_add(&(header_.combined_ref_count), kCombinedRefCountWeakOne, __ATOMIC_RELAXED); #endif } /*! \brief decrease strong reference count and delete the object */ void DecRef() { #ifdef _MSC_VER // use simpler impl in windows to ensure correctness uint64_t count_before_sub = _InterlockedDecrement64( // reinterpret_cast(&header_.combined_ref_count) // NOLINT(*) ) + 1; if (count_before_sub == kCombinedRefCountBothOne) { // NOLINT(*) // fast path: both reference counts will go to zero if (header_.deleter != nullptr) { // full barrrier is implicit in InterlockedDecrement header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskBoth); } } else if ((count_before_sub & kCombinedRefCountMaskUInt32) == kCombinedRefCountStrongOne) { // strong reference count becomes zero, we need to first do strong deletion // then decrease weak reference count // full barrrier is implicit in InterlockedAdd if (header_.deleter != nullptr) { header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskStrong); } // decrease weak reference count if (_InlineInterlockedAdd64( // reinterpret_cast(&header_.combined_ref_count), -kCombinedRefCountWeakOne) == 0) { // NOLINT(*) if (header_.deleter != nullptr) { // full barrrier is implicit in InterlockedAdd header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskWeak); } } } #else // first do a release, note we only need to acquire for deleter // optimization: we only need one atomic to tell the common case // where both reference counts are zero uint64_t count_before_sub = __atomic_fetch_sub(&(header_.combined_ref_count), kCombinedRefCountStrongOne, __ATOMIC_RELEASE); if (count_before_sub == kCombinedRefCountBothOne) { // common case, we need to delete both the object and the memory block // only acquire when we need to call deleter __atomic_thread_fence(__ATOMIC_ACQUIRE); if (header_.deleter != nullptr) { // call deleter once header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskBoth); } } else if ((count_before_sub & kCombinedRefCountMaskUInt32) == kCombinedRefCountStrongOne) { // strong count is already zero // Slower path: there is still a weak reference left __atomic_thread_fence(__ATOMIC_ACQUIRE); // call destructor first, then decrease weak reference count if (header_.deleter != nullptr) { header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskStrong); } // now decrease weak reference count if (__atomic_fetch_sub(&(header_.combined_ref_count), kCombinedRefCountWeakOne, __ATOMIC_RELEASE) == kCombinedRefCountWeakOne) { __atomic_thread_fence(__ATOMIC_ACQUIRE); if (header_.deleter != nullptr) { header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskWeak); } } } #endif } /*! \brief decrease weak reference count */ void DecWeakRef() { #ifdef _MSC_VER if (_InlineInterlockedAdd64( // reinterpret_cast(&header_.combined_ref_count), // NOLINT(*) -kCombinedRefCountWeakOne) == 0) { if (header_.deleter != nullptr) { header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskWeak); } } #else // now decrease weak reference count if (__atomic_fetch_sub(&(header_.combined_ref_count), kCombinedRefCountWeakOne, __ATOMIC_RELEASE) == kCombinedRefCountWeakOne) { __atomic_thread_fence(__ATOMIC_ACQUIRE); if (header_.deleter != nullptr) { header_.deleter(&(this->header_), kTVMFFIObjectDeleterFlagBitMaskWeak); } } #endif } // friend classes template friend class ObjectPtr; template friend class WeakObjectPtr; friend struct tvm::ffi::details::ObjectUnsafe; }; /*! * \brief A custom smart pointer for Object. * \tparam T the content data type. * \sa make_object */ template class ObjectPtr { public: /*! \brief default constructor */ ObjectPtr() = default; /*! \brief default constructor */ ObjectPtr(std::nullptr_t) {} // NOLINT(*) /*! * \brief copy constructor * \param other The value to be moved */ ObjectPtr(const ObjectPtr& other) // NOLINT(*) : ObjectPtr(other.data_) {} /*! * \brief copy constructor * \param other The value to be moved */ template ObjectPtr(const ObjectPtr& other) // NOLINT(*) : ObjectPtr(other.data_) { static_assert(std::is_base_of_v, "can only assign of child class ObjectPtr to parent"); } /*! * \brief move constructor * \param other The value to be moved */ ObjectPtr(ObjectPtr&& other) // NOLINT(*) : data_(other.data_) { other.data_ = nullptr; } /*! * \brief move constructor * \param other The value to be moved */ template ObjectPtr(ObjectPtr&& other) // NOLINT(*) : data_(other.data_) { static_assert(std::is_base_of_v, "can only assign of child class ObjectPtr to parent"); other.data_ = nullptr; } /*! \brief destructor */ ~ObjectPtr() { this->reset(); } /*! * \brief Swap this array with another Object * \param other The other Object */ void swap(ObjectPtr& other) { // NOLINT(*) std::swap(data_, other.data_); } /*! * \return Get the content of the pointer */ T* get() const { return static_cast(data_); } /*! * \return The pointer */ T* operator->() const { return get(); } /*! * \return The reference */ T& operator*() const { // NOLINT(*) return *get(); } /*! * \brief copy assignment * \param other The value to be assigned. * \return reference to self. */ ObjectPtr& operator=(const ObjectPtr& other) { // NOLINT(*) // takes in plane operator to enable copy elison. // copy-and-swap idiom ObjectPtr(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief move assignment * \param other The value to be assigned. * \return reference to self. */ ObjectPtr& operator=(ObjectPtr&& other) { // NOLINT(*) // copy-and-swap idiom ObjectPtr(std::move(other)).swap(*this); // NOLINT(*) return *this; } /*! * \brief nullptr check * \return result of comparison of internal pointer with nullptr. */ explicit operator bool() const { return get() != nullptr; } /*! \brief reset the content of ptr to be nullptr */ void reset() { if (data_ != nullptr) { data_->DecRef(); data_ = nullptr; } } /*! \return The use count of the ptr, for debug purposes */ int use_count() const { return data_ != nullptr ? data_->use_count() : 0; } /*! \return whether the reference is unique */ bool unique() const { return data_ != nullptr && data_->use_count() == 1; } /*! \return Whether two ObjectPtr do not equal each other */ bool operator==(const ObjectPtr& other) const { return data_ == other.data_; } /*! \return Whether two ObjectPtr equals each other */ bool operator!=(const ObjectPtr& other) const { return data_ != other.data_; } /*! \return Whether the pointer is nullptr */ bool operator==(std::nullptr_t) const { return data_ == nullptr; } /*! \return Whether the pointer is not nullptr */ bool operator!=(std::nullptr_t) const { return data_ != nullptr; } private: /*! \brief internal pointer field */ Object* data_{nullptr}; /*! * \brief constructor from Object * \param data The data pointer */ explicit ObjectPtr(Object* data) : data_(data) { if (data_ != nullptr) { data_->IncRef(); } } // friend classes friend class Object; friend class ObjectRef; friend struct ObjectPtrHash; template friend class ObjectPtr; template friend class WeakObjectPtr; friend struct tvm::ffi::details::ObjectUnsafe; }; /*! * \brief A custom smart pointer for Object. * \tparam T the content data type. * \sa make_object */ template class WeakObjectPtr { public: /*! \brief default constructor */ WeakObjectPtr() = default; /*! \brief default constructor */ WeakObjectPtr(std::nullptr_t) {} // NOLINT(*) /*! * \brief copy constructor * \param other The value to be moved */ WeakObjectPtr(const WeakObjectPtr& other) // NOLINT(*) : WeakObjectPtr(other.data_) {} /*! * \brief copy constructor * \param other The value to be moved */ WeakObjectPtr(const ObjectPtr& other) // NOLINT(*) : WeakObjectPtr(other.get()) {} /*! * \brief copy constructor * \param other The value to be moved */ template WeakObjectPtr(const WeakObjectPtr& other) // NOLINT(*) : WeakObjectPtr(other.data_) { static_assert(std::is_base_of_v, "can only assign of child class ObjectPtr to parent"); } /*! * \brief copy constructor * \param other The value to be moved */ template WeakObjectPtr(const ObjectPtr& other) // NOLINT(*) : WeakObjectPtr(other.data_) { static_assert(std::is_base_of_v, "can only assign of child class ObjectPtr to parent"); } /*! * \brief move constructor * \param other The value to be moved */ WeakObjectPtr(WeakObjectPtr&& other) // NOLINT(*) : data_(other.data_) { other.data_ = nullptr; } /*! * \brief move constructor * \param other The value to be moved */ template WeakObjectPtr(WeakObjectPtr&& other) // NOLINT(*) : data_(other.data_) { static_assert(std::is_base_of_v, "can only assign of child class ObjectPtr to parent"); other.data_ = nullptr; } /*! \brief destructor */ ~WeakObjectPtr() { this->reset(); } /*! * \brief Swap this array with another Object * \param other The other Object */ void swap(WeakObjectPtr& other) { // NOLINT(*) std::swap(data_, other.data_); } /*! * \brief copy assignment * \param other The value to be assigned. * \return reference to self. */ WeakObjectPtr& operator=(const WeakObjectPtr& other) { // NOLINT(*) // takes in plane operator to enable copy elison. // copy-and-swap idiom WeakObjectPtr(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief move assignment * \param other The value to be assigned. * \return reference to self. */ WeakObjectPtr& operator=(WeakObjectPtr&& other) { // NOLINT(*) // copy-and-swap idiom WeakObjectPtr(std::move(other)).swap(*this); // NOLINT(*) return *this; } /*! \return The internal object pointer if the object is still alive, otherwise nullptr */ ObjectPtr lock() const { if (data_ != nullptr && data_->TryPromoteWeakPtr()) { ObjectPtr ret; // we already increase the reference count, so we don't need to do it again ret.data_ = data_; return ret; } return nullptr; } /*! \brief reset the content of ptr to be nullptr */ void reset() { if (data_ != nullptr) { data_->DecWeakRef(); data_ = nullptr; } } /*! \return The use count of the ptr, for debug purposes */ int use_count() const { return data_ != nullptr ? data_->use_count() : 0; } /*! \return whether the pointer is nullptr */ bool expired() const { return data_ == nullptr || data_->use_count() == 0; } private: /*! \brief internal pointer field */ Object* data_{nullptr}; /*! * \brief constructor from Object * \param data The data pointer */ explicit WeakObjectPtr(Object* data) : data_(data) { if (data_ != nullptr) { data_->IncWeakRef(); } } template friend class WeakObjectPtr; friend struct tvm::ffi::details::ObjectUnsafe; }; /*! * \brief Optional data type in FFI. * \tparam T The underlying type of the optional. * * \note Compared to std::optional, Optional * akes less storage as it used nullptr to represent nullopt. */ template class Optional; /*! \brief Base class of all object reference */ class ObjectRef { public: /*! \brief default constructor */ ObjectRef() = default; /*! \brief copy constructor */ ObjectRef(const ObjectRef& other) = default; /*! \brief move constructor */ ObjectRef(ObjectRef&& other) noexcept : data_(std::move(other.data_)) { other.data_ = nullptr; } /*! \brief copy assignment */ ObjectRef& operator=(const ObjectRef& other) = default; /*! \brief move assignment */ ObjectRef& operator=(ObjectRef&& other) noexcept { data_ = std::move(other.data_); other.data_ = nullptr; return *this; } /*! \brief Constructor from existing object ptr */ explicit ObjectRef(ObjectPtr data) : data_(std::move(data)) {} /*! \brief Constructor from UnsafeInit */ explicit ObjectRef(UnsafeInit) : data_(nullptr) {} /*! * \brief Comparator * \param other Another object ref. * \return the compare result. */ bool same_as(const ObjectRef& other) const { return data_ == other.data_; } /*! * \brief Comparator * \param other Another object ref. * \return the compare result. */ bool operator==(const ObjectRef& other) const { return data_ == other.data_; } /*! * \brief Comparator * \param other Another object ref. * \return the compare result. */ bool operator!=(const ObjectRef& other) const { return data_ != other.data_; } /*! * \brief Comparator * \param other Another object ref by address. * \return the compare result. */ bool operator<(const ObjectRef& other) const { return data_.get() < other.data_.get(); } /*! * \return whether the object is defined. */ bool defined() const { return data_ != nullptr; } /*! \return the internal object pointer */ const Object* get() const { return data_.get(); } /*! \return the internal object pointer */ const Object* operator->() const { return get(); } /*! \return whether the reference is unique */ bool unique() const { return data_.unique(); } /*! \return The use count of the ptr, for debug purposes */ int use_count() const { return data_.use_count(); } /*! * \brief Try to downcast the internal Object to a * raw pointer of a corresponding type. * * The function will return a nullptr if the cast failed. * * if (const AddNode *ptr = node_ref.as()) { * // This is an add node * } * * \tparam ObjectType the target type, must be a subtype of Object * \return The pointer to the requested type. */ template >> const ObjectType* as() const { if (data_ != nullptr && data_->IsInstance()) { return static_cast(data_.get()); } else { return nullptr; } } /*! * \brief Try to downcast the ObjectRef to Optional of the requested type. * * The function will return a std::nullopt if the cast or if the pointer is nullptr. * * \tparam ObjectRefType the target type, must be a subtype of ObjectRef' * \return The optional value of the requested type. */ template >> TVM_FFI_INLINE std::optional as() const { if (data_ != nullptr) { if (data_->IsInstance()) { ObjectRefType ref(UnsafeInit{}); ref.data_ = data_; return ref; } else { return std::nullopt; } } else { return std::nullopt; } } /*! * \brief Get the type index of the ObjectRef * \return The type index of the ObjectRef */ int32_t type_index() const { return data_ != nullptr ? data_->type_index() : TypeIndex::kTVMFFINone; } /*! * \brief Get the type key of the ObjectRef * \return The type key of the ObjectRef */ std::string GetTypeKey() const { return data_ != nullptr ? data_->GetTypeKey() : StaticTypeKey::kTVMFFINone; } /*! \brief type indicate the container type. */ using ContainerType = Object; /*! \brief Whether the reference can point to nullptr */ static constexpr bool _type_is_nullable = true; protected: /*! \brief Internal pointer that backs the reference. */ ObjectPtr data_; /*! \return return a mutable internal ptr, can be used by sub-classes. */ Object* get_mutable() const { return data_.get(); } // friend classes. friend struct ObjectPtrHash; friend struct tvm::ffi::details::ObjectUnsafe; }; // forward delcare variant template class Variant; /*! \brief ObjectRef hash functor */ struct ObjectPtrHash { size_t operator()(const ObjectRef& a) const { return operator()(a.data_); } template size_t operator()(const ObjectPtr& a) const { return std::hash()(a.get()); } template TVM_FFI_INLINE size_t operator()(const Variant& a) const; }; /*! \brief ObjectRef equal functor */ struct ObjectPtrEqual { bool operator()(const ObjectRef& a, const ObjectRef& b) const { return a.same_as(b); } template bool operator()(const ObjectPtr& a, const ObjectPtr& b) const { return a == b; } template TVM_FFI_INLINE bool operator()(const Variant& a, const Variant& b) const; }; /*! * \brief Helper macro to declare object information with static type index. * * For each custom object, you need to call tvm::ffi::reflection::ObjectDef() * once in your cc file to register the type index with the runtime. * Alternatively, you can call TypeName::_GetOrAllocRuntimeTypeIndex() once. * * \param TypeKey The type key of the current type. * \param TypeName The name of the current type. * \param ParentType The name of the ParentType * * \see tvm::ffi::reflection::ObjectDef */ #define TVM_FFI_DECLARE_OBJECT_INFO_STATIC(TypeKey, TypeName, ParentType) \ static constexpr int32_t _type_depth = ParentType::_type_depth + 1; \ static int32_t _GetOrAllocRuntimeTypeIndex() { \ static_assert(!ParentType::_type_final, "ParentType marked as final"); \ static_assert(TypeName::_type_child_slots == 0 || ParentType::_type_child_slots == 0 || \ TypeName::_type_child_slots < ParentType::_type_child_slots, \ "Need to set _type_child_slots when parent specifies it."); \ TVMFFIByteArray type_key{TypeName::_type_key, \ std::char_traits::length(TypeName::_type_key)}; \ static int32_t tindex [[maybe_unused]] = TVMFFITypeGetOrAllocIndex( \ &type_key, TypeName::_type_index, TypeName::_type_depth, TypeName::_type_child_slots, \ TypeName::_type_child_slots_can_overflow, ParentType::_GetOrAllocRuntimeTypeIndex()); \ return TypeName::_type_index; \ } \ static int32_t RuntimeTypeIndex() { return TypeName::_type_index; } \ static constexpr const char* _type_key = TypeKey /*! * \brief Helper macro to declare object information with type key already defined in class. * * \param TypeName The name of the current type. * \param ParentType The name of the ParentType */ #define TVM_FFI_DECLARE_OBJECT_INFO_PREDEFINED_TYPE_KEY(TypeName, ParentType) \ static constexpr int32_t _type_depth = ParentType::_type_depth + 1; \ static int32_t _GetOrAllocRuntimeTypeIndex() { \ static_assert(!ParentType::_type_final, "ParentType marked as final"); \ static_assert(TypeName::_type_child_slots == 0 || ParentType::_type_child_slots == 0 || \ TypeName::_type_child_slots < ParentType::_type_child_slots, \ "Need to set _type_child_slots when parent specifies it."); \ TVMFFIByteArray type_key{TypeName::_type_key, \ std::char_traits::length(TypeName::_type_key)}; \ static int32_t tindex = TVMFFITypeGetOrAllocIndex( \ &type_key, -1, TypeName::_type_depth, TypeName::_type_child_slots, \ TypeName::_type_child_slots_can_overflow, ParentType::_GetOrAllocRuntimeTypeIndex()); \ return tindex; \ } \ static int32_t RuntimeTypeIndex() { return _GetOrAllocRuntimeTypeIndex(); } /*! * \brief Helper macro to declare object information with dynamic type index. * * For each custom object, you need to call tvm::ffi::reflection::ObjectDef() * once in your cc file to register the type index with the runtime. * Alternatively, you can call TypeName::_GetOrAllocRuntimeTypeIndex() once. * * \param TypeKey The type key of the current type. * \param TypeName The name of the current type. * \param ParentType The name of the ParentType * \sa tvm::ffi::reflection::ObjectDef */ #define TVM_FFI_DECLARE_OBJECT_INFO(TypeKey, TypeName, ParentType) \ static constexpr const char* _type_key = TypeKey; \ TVM_FFI_DECLARE_OBJECT_INFO_PREDEFINED_TYPE_KEY(TypeName, ParentType) /*! * \brief Helper macro to declare object information with dynamic type index and is final. * * For each custom object, you need to call tvm::ffi::reflection::ObjectDef() * once in your cc file to register the type index with the runtime. * Alternatively, you can call TypeName::_GetOrAllocRuntimeTypeIndex() once. * * \param TypeKey The type key of the current type. * \param TypeName The name of the current type. * \param ParentType The name of the ParentType * \sa tvm::ffi::reflection::ObjectDef */ #define TVM_FFI_DECLARE_OBJECT_INFO_FINAL(TypeKey, TypeName, ParentType) \ static const constexpr int _type_child_slots [[maybe_unused]] = 0; \ static const constexpr bool _type_final [[maybe_unused]] = true; \ TVM_FFI_DECLARE_OBJECT_INFO(TypeKey, TypeName, ParentType) /*! * \brief Define object reference methods. * * \param TypeName The object type name * \param ParentType The parent type of the objectref * \param ObjectName The type name of the object. * * \note This macro also defines the default constructor that puts the ObjectRef * in undefined state initially. */ #define TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypeName, ParentType, ObjectName) \ TypeName() = default; \ explicit TypeName(::tvm::ffi::ObjectPtr n) : ParentType(std::move(n)) {} \ explicit TypeName(::tvm::ffi::UnsafeInit tag) : ParentType(tag) {} \ TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TypeName) \ using __PtrType = std::conditional_t<(ObjectName::_type_mutable), \ ObjectName*, /* NOLINT(bugprone-macro-parentheses) */ \ const ObjectName*>; \ __PtrType operator->() const { return static_cast<__PtrType>(data_.get()); } \ __PtrType get() const { return static_cast<__PtrType>(data_.get()); } \ [[maybe_unused]] static constexpr bool _type_is_nullable = true; \ using ContainerType = ObjectName /*! * \brief Define object reference methods do not have undefined state. * * \param TypeName The object type name * \param ParentType The parent type of the objectref * \param ObjectName The type name of the object. */ #define TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypeName, ParentType, ObjectName) \ explicit TypeName(::tvm::ffi::UnsafeInit tag) : ParentType(tag) {} \ TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TypeName) \ using __PtrType = std::conditional_t<(ObjectName::_type_mutable), \ ObjectName*, /* NOLINT(bugprone-macro-parentheses) */ \ const ObjectName*>; \ __PtrType operator->() const { return static_cast<__PtrType>(data_.get()); } \ __PtrType get() const { return static_cast<__PtrType>(data_.get()); } \ [[maybe_unused]] static constexpr bool _type_is_nullable = false; \ using ContainerType = ObjectName namespace details { template TVM_FFI_INLINE bool IsObjectInstance(int32_t object_type_index) { static_assert(std::is_base_of_v); // Everything is a subclass of object. if constexpr (std::is_same_v) { return true; } else if constexpr (TargetType::_type_final) { // if the target type is a final type // then we only need to check the equivalence. return object_type_index == TargetType::RuntimeTypeIndex(); } else { // Explicitly enclose in else to eliminate this branch early in compilation. // if target type is a non-leaf type // Check if type index falls into the range of reserved slots. int32_t target_type_index = TargetType::RuntimeTypeIndex(); int32_t begin = target_type_index; // The condition will be optimized by constant-folding. if constexpr (TargetType::_type_child_slots != 0) { // total_slots = child_slots + 1 (including self) int32_t end = begin + TargetType::_type_child_slots + 1; if (object_type_index >= begin && object_type_index < end) return true; } else { if (object_type_index == begin) return true; } if constexpr (TargetType::_type_child_slots_can_overflow) { // Invariance: parent index is always smaller than the child. if (object_type_index < target_type_index) return false; // Do a runtime lookup of type information // the function checks that the info exists const TypeInfo* type_info = TVMFFIGetTypeInfo(object_type_index); return (type_info->type_depth > TargetType::_type_depth && type_info->type_ancestors[TargetType::_type_depth]->type_index == target_type_index); } else { return false; } } } /*! * \brief Namespace to internally manipulate object class. * \note These functions are only supposed to be used by internal * implementations and not external users of the tvm::ffi */ struct ObjectUnsafe { // NOTE: get ffi header from an object TVM_FFI_INLINE static TVMFFIObject* GetHeader(const Object* src) { return const_cast(&(src->header_)); } // Suppress -Winvalid-offsetof: we intentionally use offsetof on non-standard-layout types // to avoid undefined behavior from null pointer arithmetic that sanitizers flag. #if defined(__clang__) || defined(__GNUC__) #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Winvalid-offsetof" #endif template TVM_FFI_INLINE static int64_t GetObjectOffsetToSubclass() { return static_cast(__builtin_offsetof(Class, header_)) - static_cast(__builtin_offsetof(Object, header_)); } #if defined(__clang__) || defined(__GNUC__) #pragma GCC diagnostic pop #endif template TVM_FFI_INLINE static T ObjectRefFromObjectPtr(const ObjectPtr& ptr) { T ref(UnsafeInit{}); ref.data_ = ptr; return ref; } template TVM_FFI_INLINE static T ObjectRefFromObjectPtr(ObjectPtr&& ptr) { T ref(UnsafeInit{}); ref.data_ = std::move(ptr); return ref; } template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromObjectRef(const ObjectRef& ref) { if constexpr (std::is_same_v) { return ref.data_; } else { return tvm::ffi::ObjectPtr(ref.data_.data_); } } template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromObjectRef(ObjectRef&& ref) { if constexpr (std::is_same_v) { return std::move(ref.data_); } else { ObjectPtr result; result.data_ = std::move(ref.data_.data_); ref.data_.data_ = nullptr; return result; } } template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromOwned(Object* raw_ptr) { tvm::ffi::ObjectPtr ptr; ptr.data_ = raw_ptr; return ptr; } template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromOwned(TVMFFIObject* obj_ptr) { return ObjectPtrFromOwned(reinterpret_cast(obj_ptr)); } template TVM_FFI_INLINE static T* RawObjectPtrFromUnowned(TVMFFIObject* obj_ptr) { // NOTE: this is important to first cast to Object* // then cast back to T* because objptr and tptr may not be the same // depending on how sub-class allocates the space. return static_cast(reinterpret_cast(obj_ptr)); } // Create ObjectPtr from unowned ptr template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromUnowned(Object* raw_ptr) { return tvm::ffi::ObjectPtr(raw_ptr); } template TVM_FFI_INLINE static ObjectPtr ObjectPtrFromUnowned(TVMFFIObject* obj_ptr) { return tvm::ffi::ObjectPtr(reinterpret_cast(obj_ptr)); } TVM_FFI_INLINE static void DecRefObjectHandle(TVMFFIObjectHandle handle) { if (handle) reinterpret_cast(handle)->DecRef(); } TVM_FFI_INLINE static void IncRefObjectHandle(TVMFFIObjectHandle handle) { reinterpret_cast(handle)->IncRef(); } TVM_FFI_INLINE static Object* RawObjectPtrFromObjectRef(const ObjectRef& src) { return src.data_.data_; } TVM_FFI_INLINE static TVMFFIObject* TVMFFIObjectPtrFromObjectRef(const ObjectRef& src) { return GetHeader(src.data_.data_); } template TVM_FFI_INLINE static TVMFFIObject* TVMFFIObjectPtrFromObjectPtr(const ObjectPtr& src) { return GetHeader(src.data_); } template TVM_FFI_INLINE static TVMFFIObject* MoveObjectPtrToTVMFFIObjectPtr(ObjectPtr&& src) { Object* obj_ptr = src.data_; src.data_ = nullptr; return GetHeader(obj_ptr); } TVM_FFI_INLINE static TVMFFIObject* MoveObjectRefToTVMFFIObjectPtr(ObjectRef&& src) { Object* obj_ptr = src.data_.data_; src.data_.data_ = nullptr; return GetHeader(obj_ptr); } }; } // namespace details } // namespace ffi } // namespace tvm #endif // TVM_FFI_OBJECT_H_ tvm-ffi-0.1.12/include/tvm/ffi/optional.h000066400000000000000000000317021521067262500201470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/optional.h * \brief Runtime Optional container types. * \note Optional specializes for T is ObjectRef and used nullptr to indicate nullopt. */ #ifndef TVM_FFI_OPTIONAL_H_ #define TVM_FFI_OPTIONAL_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { // Note: We place optional in tvm/ffi instead of tvm/ffi/container // because optional itself is an inherent core component of the FFI system. /// \cond Doxygen_Suppress template inline constexpr bool is_optional_type_v = false; template inline constexpr bool is_optional_type_v> = true; // we can safely used ptr based optional for ObjectRef types // that do not have additional data members and virtual functions. template inline constexpr bool use_ptr_based_optional_v = (std::is_base_of_v && !is_optional_type_v); /// \endcond // Specialization for non-ObjectRef types. // simply fallback to std::optional template class Optional && !std::is_same_v && !std::is_same_v>> { public: // default constructors. Optional() = default; // NOLINTBEGIN(google-explicit-constructor) Optional(const Optional& other) : data_(other.data_) {} Optional(Optional&& other) noexcept : data_(std::move(other.data_)) {} Optional(std::optional other) : data_(std::move(other)) {} Optional(std::nullopt_t) {} Optional(T other) : data_(std::move(other)) {} // NOLINTEND(google-explicit-constructor) TVM_FFI_INLINE Optional& operator=(const Optional& other) { data_ = other.data_; return *this; } TVM_FFI_INLINE Optional& operator=(Optional&& other) noexcept { data_ = std::move(other.data_); return *this; } TVM_FFI_INLINE Optional& operator=(T other) { data_ = std::move(other); return *this; } TVM_FFI_INLINE Optional& operator=(std::nullopt_t) { data_ = std::nullopt; return *this; } TVM_FFI_INLINE const T& value() const& { if (!data_.has_value()) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return *data_; } TVM_FFI_INLINE T&& value() && { if (!data_.has_value()) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return *std::move(data_); } template > TVM_FFI_INLINE T value_or(U&& default_value) const { return data_.value_or(std::forward(default_value)); } TVM_FFI_INLINE explicit operator bool() const noexcept { return data_.has_value(); } TVM_FFI_INLINE bool has_value() const noexcept { return data_.has_value(); } TVM_FFI_INLINE bool operator==(const Optional& other) const { return data_ == other.data_; } TVM_FFI_INLINE bool operator!=(const Optional& other) const { return data_ != other.data_; } template TVM_FFI_INLINE bool operator==(const U& other) const { return data_ == other; } template TVM_FFI_INLINE bool operator!=(const U& other) const { return data_ != other; } // NOLINTBEGIN(bugprone-unchecked-optional-access) /*! * \brief Direct access to the value. * \return the xvalue reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE T&& operator*() && noexcept { return *std::move(data_); } /*! * \brief Direct access to the value. * \return the const reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE const T& operator*() const& noexcept { return *data_; } // NOLINTEND(bugprone-unchecked-optional-access) private: std::optional data_; }; // Specialization for String type, use nullptr to indicate nullopt template class Optional || std::is_same_v>> { public: // default constructors. Optional() = default; // NOLINTBEGIN(google-explicit-constructor) Optional(const Optional& other) : data_(other.data_) {} Optional(Optional&& other) : data_(std::move(other.data_)) {} Optional(std::nullopt_t) {} Optional(T other) : data_(std::move(other)) {} // NOLINTEND(google-explicit-constructor) TVM_FFI_INLINE Optional& operator=(const Optional& other) { data_ = other.data_; return *this; } TVM_FFI_INLINE Optional& operator=(Optional&& other) { data_ = std::move(other.data_); return *this; } TVM_FFI_INLINE Optional& operator=(T other) { data_ = std::move(other); return *this; } TVM_FFI_INLINE Optional& operator=(std::nullopt_t) { T(details::BytesBaseCell(std::nullopt)).swap(data_); return *this; } TVM_FFI_INLINE const T& value() const& { if (data_.data_ == std::nullopt) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return data_; } TVM_FFI_INLINE String&& value() && { if (data_.data_ == std::nullopt) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return std::move(data_); } template TVM_FFI_INLINE T value_or(U&& default_value) const { if (data_.data_ == std::nullopt) { return std::forward(default_value); } return data_; } TVM_FFI_INLINE explicit operator bool() const noexcept { return data_.data_ != std::nullopt; } TVM_FFI_INLINE bool has_value() const noexcept { return data_.data_ != std::nullopt; } TVM_FFI_INLINE bool operator==(const Optional& other) const { if (data_.data_ == std::nullopt) { return other.data_.data_ == std::nullopt; } if (other.data_.data_ == std::nullopt) { return false; } return data_ == other.data_; } TVM_FFI_INLINE bool operator!=(const Optional& other) const { return !(*this == other); } template TVM_FFI_INLINE bool operator==(const U& other) const { if constexpr (std::is_same_v) { return data_.data_ == std::nullopt; } else { if (data_.data_ == std::nullopt) { return false; } return data_ == other; } } template TVM_FFI_INLINE bool operator!=(const U& other) const { if constexpr (std::is_same_v) { return data_.data_ != std::nullopt; } else { if (data_.data_ == std::nullopt) { return true; } return data_ != other; } } /*! * \brief Direct access to the value. * \return the xvalue reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE T&& operator*() && noexcept { return std::move(data_); } /*! * \brief Direct access to the value. * \return the const reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE const T& operator*() const& noexcept { return data_; } private: // this is a private initializer T data_{details::BytesBaseCell(std::nullopt)}; }; // Specialization for ObjectRef types. // nullptr is treated as std::nullopt. template class Optional>> : public ObjectRef { public: using ContainerType = typename T::ContainerType; Optional() = default; // NOLINTBEGIN(google-explicit-constructor) Optional(const Optional& other) : ObjectRef(other) {} Optional(Optional&& other) noexcept : ObjectRef(std::move(other)) {} explicit Optional(ffi::UnsafeInit tag) : ObjectRef(tag) {} Optional(std::nullopt_t) {} Optional(std::optional other) { if (other.has_value()) { *this = *std::move(other); } } Optional(T other) : ObjectRef(std::move(other)) {} // NOLINTEND(google-explicit-constructor) TVM_FFI_INLINE Optional& operator=(T other) { ObjectRef::operator=(std::move(other)); return *this; } TVM_FFI_INLINE Optional& operator=(const Optional& other) { data_ = other.data_; return *this; } TVM_FFI_INLINE Optional& operator=(std::nullptr_t) { data_ = nullptr; return *this; } TVM_FFI_INLINE Optional& operator=(Optional&& other) { data_ = std::move(other.data_); return *this; } TVM_FFI_INLINE T value() const& { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return details::ObjectUnsafe::ObjectRefFromObjectPtr(data_); } TVM_FFI_INLINE T value() && { if (data_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "Back optional access"; } return details::ObjectUnsafe::ObjectRefFromObjectPtr(std::move(data_)); } template > TVM_FFI_INLINE T value_or(U&& default_value) const { return data_ != nullptr ? details::ObjectUnsafe::ObjectRefFromObjectPtr(data_) : T(std::forward(default_value)); } TVM_FFI_INLINE explicit operator bool() const { return data_ != nullptr; } TVM_FFI_INLINE bool has_value() const { return data_ != nullptr; } /*! * \brief Direct access to the value. * \return the const reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE T operator*() const& noexcept { return details::ObjectUnsafe::ObjectRefFromObjectPtr(data_); } /*! * \brief Direct access to the value. * \return the const reference to the stored value. * \note only use this function after checking has_value() */ TVM_FFI_INLINE T operator*() && noexcept { return details::ObjectUnsafe::ObjectRefFromObjectPtr(std::move(data_)); } TVM_FFI_INLINE bool operator==(std::nullptr_t) const noexcept { return !has_value(); } TVM_FFI_INLINE bool operator!=(std::nullptr_t) const noexcept { return has_value(); } // operator overloadings TVM_FFI_INLINE auto operator==(const Optional& other) const { // support case where sub-class returns a symbolic ref type. return EQToOptional(other); } TVM_FFI_INLINE auto operator!=(const Optional& other) const { return NEToOptional(other); } TVM_FFI_INLINE auto operator==(const std::optional& other) const { // support case where sub-class returns a symbolic ref type. return EQToOptional(other); } TVM_FFI_INLINE auto operator!=(const std::optional& other) const { return NEToOptional(other); } TVM_FFI_INLINE auto operator==(const T& other) const { using RetType = decltype(value() == other); if (same_as(other)) return RetType(true); if (has_value()) return operator*() == other; return RetType(false); } TVM_FFI_INLINE auto operator!=(const T& other) const { return !(*this == other); } template TVM_FFI_INLINE auto operator==(const U& other) const { using RetType = decltype(value() == other); if (!has_value()) return RetType(false); return operator*() == other; } template TVM_FFI_INLINE auto operator!=(const U& other) const { using RetType = decltype(value() != other); if (!has_value()) return RetType(true); return operator*() != other; } /*! * \return The internal object pointer with container type of T. * \note This function do not perform not-null checking. */ TVM_FFI_INLINE const ContainerType* get() const { return static_cast(data_.get()); } private: template TVM_FFI_INLINE auto EQToOptional(const U& other) const { // support case where sub-class returns a symbolic ref type. using RetType = decltype(operator*() == *other); if (same_as(other)) return RetType(true); if (has_value() && other.has_value()) { return operator*() == *other; } else { // one of them is nullptr. return RetType(false); } } template TVM_FFI_INLINE auto NEToOptional(const U& other) const { // support case where sub-class returns a symbolic ref type. using RetType = decltype(operator*() != *other); if (same_as(other)) return RetType(false); if (has_value() && other.has_value()) { return operator*() != *other; } else { // one of them is nullptr. return RetType(true); } } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_OPTIONAL_H_ tvm-ffi-0.1.12/include/tvm/ffi/reflection/000077500000000000000000000000001521067262500203005ustar00rootroot00000000000000tvm-ffi-0.1.12/include/tvm/ffi/reflection/access_path.h000066400000000000000000000324361521067262500227360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/registry.h * \brief Registry of reflection metadata. */ #ifndef TVM_FFI_REFLECTION_ACCESS_PATH_H_ #define TVM_FFI_REFLECTION_ACCESS_PATH_H_ #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace reflection { /*! * \brief The kind of the access pattern. */ enum class AccessKind : int32_t { /*! \brief Object attribute access. */ kAttr = 0, /*! \brief Array item access. */ kArrayItem = 1, /*! \brief Map item access. */ kMapItem = 2, // the following two are used for error reporting when // the supposed access field is not available /*! \brief Object attribute missing access. */ kAttrMissing = 3, /*! \brief Array item missing access. */ kArrayItemMissing = 4, /*! \brief Map item missing access. */ kMapItemMissing = 5, }; class AccessStep; /*! * \brief Represent a single step in object field, map key, array index access. */ class AccessStepObj : public Object { public: /*! * \brief The kind of the access pattern. */ AccessKind kind; /*! * \brief The access key * \note for array access, it will always be integer * for field access, it will be string */ Any key; // default constructor to enable auto-serialization AccessStepObj() = default; /*! * \brief Constructor * \param kind The kind of the access step. * \param key The key of the access step. */ AccessStepObj(AccessKind kind, Any key) : kind(kind), key(std::move(key)) {} /*! * \brief Deep check if two steps are equal. * \param other The other step to compare with. * \return True if the two steps are equal, false otherwise. */ inline bool StepEqual(const AccessStep& other) const; /// \cond Doxygen_Suppress static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ffi.reflection.AccessStep", AccessStepObj, Object); /// \endcond }; /*! * \brief ObjectRef class of AccessStepObj. * * \sa AccessStepObj */ class AccessStep : public ObjectRef { public: /*! * \brief Constructor * \param kind The kind of the access step. * \param key The key of the access step. */ AccessStep(AccessKind kind, Any key) : ObjectRef(make_object(kind, std::move(key))) {} /*! * \brief Create an access step for a object attribute access. * \param field_name The name of the field to access. * \return The access step. */ static AccessStep Attr(String field_name) { return AccessStep(AccessKind::kAttr, std::move(field_name)); } /*! * \brief Create an access step for a object attribute missing access. * \param field_name The name of the field to access. * \return The access step. */ static AccessStep AttrMissing(String field_name) { return AccessStep(AccessKind::kAttrMissing, std::move(field_name)); } /*! * \brief Create an access step for a array item access. * \param index The index of the array item to access. * \return The access step. */ static AccessStep ArrayItem(int64_t index) { return AccessStep(AccessKind::kArrayItem, index); } /*! * \brief Create an access step for a array item missing access. * \param index The index of the array item to access. * \return The access step. */ static AccessStep ArrayItemMissing(int64_t index) { return AccessStep(AccessKind::kArrayItemMissing, index); } /*! * \brief Create an access step for a map item access. * \param key The key of the map item to access. * \return The access step. */ static AccessStep MapItem(Any key) { return AccessStep(AccessKind::kMapItem, std::move(key)); } /*! * \brief Create an access step for a map item missing access. * \param key The key of the map item to access. * \return The access step. */ static AccessStep MapItemMissing(Any key = nullptr) { return AccessStep(AccessKind::kMapItemMissing, std::move(key)); } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AccessStep, ObjectRef, AccessStepObj); /// \endcond }; inline bool AccessStepObj::StepEqual(const AccessStep& other) const { return this->kind == other->kind && AnyEqual()(this->key, other->key); } // forward declaration class AccessPath; /*! * \brief ObjectRef class of AccessPathObj. * * \sa AccessPathObj */ class AccessPathObj : public Object { public: /*! * \brief The parent of the access path. * * This parent-pointing tree structure is more space efficient when * representing multiple paths that share a common prefix. * * \note Empty for root. */ Optional parent; /*! * \brief The current of the access path. * \note Empty for root. */ Optional step; /*! * \brief The current depth of the access path, 0 for root */ int32_t depth; // default constructor to enable auto-serialization AccessPathObj() = default; /*! * \brief Constructor for the access path. * \param parent The parent of the access path. * \param step The current step of the access path. * \param depth The current depth of the access path. */ AccessPathObj(Optional parent, Optional step, int32_t depth) : parent(std::move(parent)), step(std::move(step)), depth(depth) {} /*! * \brief Get the parent of the access path. * \return The parent of the access path. */ inline Optional GetParent() const; /*! * \brief Extend the access path with a new step. * \param step The step to extend the access path with. * \return The extended access path. */ inline AccessPath Extend(AccessStep step) const; /*! * \brief Extend the access path with an object attribute access. * \param field_name The name of the field to access. * \return The extended access path. */ inline AccessPath Attr(String field_name) const; /*! * \brief Extend the access path with an object attribute missing access. * \param field_name The name of the field to access. * \return The extended access path. */ inline AccessPath AttrMissing(String field_name) const; /*! * \brief Extend the access path with an array item access. * \param index The index of the array item to access. * \return The extended access path. */ inline AccessPath ArrayItem(int64_t index) const; /*! * \brief Extend the access path with an array item missing access. * \param index The index of the array item to access. * \return The extended access path. */ inline AccessPath ArrayItemMissing(int64_t index) const; /*! * \brief Extend the access path with a map item access. * \param key The key of the map item to access. * \return The extended access path. */ inline AccessPath MapItem(Any key) const; /*! * \brief Extend the access path with a map item missing access. * \param key The key of the map item to access. * \return The extended access path. */ inline AccessPath MapItemMissing(Any key) const; /*! * \brief Get the array of steps that corresponds to the access path. * \return The array of steps that corresponds to the access path. */ inline Array ToSteps() const; /*! * \brief Check if two paths are equal by deep comparing the steps. * \param other The other path to compare with. * \return True if the two paths are equal, false otherwise. */ inline bool PathEqual(const AccessPath& other) const; /*! * \brief Check if this path is a prefix of another path. * \param other The other path to compare with. * \return True if this path is a prefix of the other path, false otherwise. */ inline bool IsPrefixOf(const AccessPath& other) const; /// \cond Doxygen_Suppress static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ffi.reflection.AccessPath", AccessPathObj, Object); /// \endcond private: static bool PathEqual(const AccessPathObj* lhs, const AccessPathObj* rhs) { // fast path for same pointer if (lhs == rhs) return true; if (lhs->depth != rhs->depth) return false; // do deep equality checks while (lhs->parent.has_value()) { TVM_FFI_ICHECK(rhs->parent.has_value()); TVM_FFI_ICHECK(lhs->step.has_value()); TVM_FFI_ICHECK(rhs->step.has_value()); if (!(*lhs->step)->StepEqual(*(rhs->step))) { return false; } lhs = static_cast(lhs->parent.get()); rhs = static_cast(rhs->parent.get()); // fast path for same pointer if (lhs == rhs) return true; TVM_FFI_ICHECK(lhs != nullptr); TVM_FFI_ICHECK(rhs != nullptr); } return true; } }; /*! * \brief ObjectRef class of AccessPath. * * \sa AccessPathObj */ class AccessPath : public ObjectRef { public: /*! * \brief Create an access path from an iterator range of steps. * \param begin The beginning of the iterator range. * \param end The end of the iterator range. * \return The access path. */ template // NOLINTNEXTLINE(performance-unnecessary-value-param) static AccessPath FromSteps(Iter begin, Iter end) { AccessPath path = AccessPath::Root(); for (Iter it = begin; it != end; ++it) { path = path->Extend(*it); } return path; } /*! * \brief Create an access path from an array of steps. * \param steps The array of steps. * \return The access path. */ static AccessPath FromSteps(const Array& steps) { AccessPath path = AccessPath::Root(); for (AccessStep step : steps) { path = path->Extend(step); } return path; } /*! * \brief Create a root access path. * \return The root access path. */ static AccessPath Root() { return AccessPath(make_object(std::nullopt, std::nullopt, 0)); } /// \cond Doxygen_Suppress TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AccessPath, ObjectRef, AccessPathObj); /// \endcond private: friend class AccessPathObj; explicit AccessPath(ObjectPtr ptr) : ObjectRef(std::move(ptr)) {} }; /*! * \brief The pair of access paths. */ using AccessPathPair = Tuple; inline Optional AccessPathObj::GetParent() const { if (auto opt_parent = this->parent.as()) { return opt_parent; } return std::nullopt; } inline AccessPath AccessPathObj::Extend(AccessStep step) const { return AccessPath( make_object(GetRef(this), std::move(step), this->depth + 1)); } inline AccessPath AccessPathObj::Attr(String field_name) const { return this->Extend(AccessStep::Attr(std::move(field_name))); } inline AccessPath AccessPathObj::AttrMissing(String field_name) const { return this->Extend(AccessStep::AttrMissing(std::move(field_name))); } inline AccessPath AccessPathObj::ArrayItem(int64_t index) const { return this->Extend(AccessStep::ArrayItem(index)); } inline AccessPath AccessPathObj::ArrayItemMissing(int64_t index) const { return this->Extend(AccessStep::ArrayItemMissing(index)); } inline AccessPath AccessPathObj::MapItem(Any key) const { return this->Extend(AccessStep::MapItem(std::move(key))); } inline AccessPath AccessPathObj::MapItemMissing(Any key) const { return this->Extend(AccessStep::MapItemMissing(std::move(key))); } inline Array AccessPathObj::ToSteps() const { std::vector reverse_steps; reverse_steps.reserve(this->depth); const AccessPathObj* current = this; while (current->parent.has_value()) { TVM_FFI_ICHECK(current->step.has_value()); reverse_steps.push_back(*(current->step)); current = static_cast(current->parent.get()); TVM_FFI_ICHECK(current != nullptr); } return Array(reverse_steps.rbegin(), reverse_steps.rend()); } inline bool AccessPathObj::PathEqual(const AccessPath& other) const { return PathEqual(this, other.get()); } inline bool AccessPathObj::IsPrefixOf(const AccessPath& other) const { if (this->depth > other->depth) { return false; } const AccessPathObj* rhs_path = other.get(); while (rhs_path->depth > this->depth) { TVM_FFI_ICHECK(rhs_path->parent.has_value()); rhs_path = static_cast(rhs_path->parent.get()); } return PathEqual(this, rhs_path); } } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_REFLECTION_ACCESS_PATH_H_ tvm-ffi-0.1.12/include/tvm/ffi/reflection/accessor.h000066400000000000000000000504501521067262500222570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/accessor.h * \brief Reflection-based accessor for object fields and methods. */ #ifndef TVM_FFI_REFLECTION_ACCESSOR_H_ #define TVM_FFI_REFLECTION_ACCESSOR_H_ #include #include #include #include #include namespace tvm { namespace ffi { namespace reflection { /*! * \brief helper function to get reflection field info by type key and field name */ inline const TVMFFIFieldInfo* GetFieldInfo(std::string_view type_key, const char* field_name) { int32_t type_index; TVMFFIByteArray type_key_array = {type_key.data(), type_key.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_array, &type_index)); const TypeInfo* info = TVMFFIGetTypeInfo(type_index); for (int32_t i = 0; i < info->num_fields; ++i) { if (std::strncmp(info->fields[i].name.data, field_name, info->fields[i].name.size) == 0) { return &(info->fields[i]); } } TVM_FFI_THROW(RuntimeError) << "Cannot find field `" << field_name << "` in " << type_key; TVM_FFI_UNREACHABLE(); } /*! * \brief Call the field setter, dispatching between function pointer and FunctionObj. * * When kTVMFFIFieldFlagBitSetterIsFunctionObj is off, invokes the setter as a * TVMFFIFieldSetter function pointer. When on, calls via TVMFFIFunctionCall * with arguments (field_addr as OpaquePtr, value). * * \param field_info The field info containing the setter. * \param field_addr The address of the field in the object. * \param value The value to set (as a TVMFFIAny pointer). * \return 0 on success, nonzero on failure. */ inline int CallFieldSetter(const TVMFFIFieldInfo* field_info, void* field_addr, const TVMFFIAny* value) { if (!(field_info->flags & kTVMFFIFieldFlagBitSetterIsFunctionObj)) { auto setter = reinterpret_cast(field_info->setter); return setter(field_addr, value); } else { TVMFFIAny args[2]{}; args[0].type_index = kTVMFFIOpaquePtr; args[0].v_ptr = field_addr; args[1] = *value; TVMFFIAny result{}; result.type_index = kTVMFFINone; return TVMFFIFunctionCall(static_cast(field_info->setter), args, 2, &result); } } /*! * \brief helper wrapper class to obtain a getter. */ class FieldGetter { public: /*! * \brief Constructor * \param field_info The field info. */ explicit FieldGetter(const TVMFFIFieldInfo* field_info) : field_info_(field_info) {} /*! * \brief Constructor * \param type_key The type key. * \param field_name The name of the field. */ explicit FieldGetter(std::string_view type_key, const char* field_name) : FieldGetter(GetFieldInfo(type_key, field_name)) {} /*! * \brief Get the value of the field * \param obj_ptr The object pointer. * \return The value of the field. */ Any operator()(const Object* obj_ptr) const { Any result; const void* addr = reinterpret_cast(obj_ptr) + field_info_->offset; TVM_FFI_CHECK_SAFE_CALL( field_info_->getter(const_cast(addr), reinterpret_cast(&result))); return result; } Any operator()(const ObjectPtr& obj_ptr) const { return operator()(obj_ptr.get()); } Any operator()(const ObjectRef& obj) const { return operator()(obj.get()); } private: const TVMFFIFieldInfo* field_info_; }; /*! * \brief helper wrapper class to obtain a setter. */ class FieldSetter { public: /*! * \brief Constructor * \param field_info The field info. */ explicit FieldSetter(const TVMFFIFieldInfo* field_info) : field_info_(field_info) {} /*! * \brief Constructor * \param type_key The type key. * \param field_name The name of the field. */ explicit FieldSetter(std::string_view type_key, const char* field_name) : FieldSetter(GetFieldInfo(type_key, field_name)) {} /*! * \brief Set the value of the field * \param obj_ptr The object pointer. * \param value The value to be set. */ void operator()(const Object* obj_ptr, AnyView value) const { const void* addr = reinterpret_cast(obj_ptr) + field_info_->offset; TVM_FFI_CHECK_SAFE_CALL(CallFieldSetter(field_info_, const_cast(addr), reinterpret_cast(&value))); } void operator()(const ObjectPtr& obj_ptr, AnyView value) const { operator()(obj_ptr.get(), value); } void operator()(const ObjectRef& obj, AnyView value) const { operator()(obj.get(), value); } private: const TVMFFIFieldInfo* field_info_; }; /*! * \brief Helper class to get type attribute column. */ class TypeAttrColumn { public: /*! * \brief Constructor * \param attr_name The name of the type attribute. */ explicit TypeAttrColumn(std::string_view attr_name) { TVMFFIByteArray attr_name_array = {attr_name.data(), attr_name.size()}; column_ = TVMFFIGetTypeAttrColumn(&attr_name_array); if (column_ == nullptr) { TVM_FFI_THROW(RuntimeError) << "Cannot find type attribute " << attr_name; } } /*! * \brief Get the type attribute column by type index. * \param type_index The type index. * \return The type attribute column. */ AnyView operator[](int32_t type_index) const { int32_t offset = type_index - column_->begin_index; if (offset < 0 || offset >= column_->size) { return AnyView(); } const AnyView* any_view_data = reinterpret_cast(column_->data); return any_view_data[offset]; } private: const TVMFFITypeAttrColumn* column_; }; /*! * \brief helper function to get reflection method info by type key and method name * * \param type_key The type key. * \param method_name The name of the method. * \return The method info. */ inline const TVMFFIMethodInfo* GetMethodInfo(std::string_view type_key, const char* method_name) { int32_t type_index; TVMFFIByteArray type_key_array = {type_key.data(), type_key.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_array, &type_index)); const TypeInfo* info = TVMFFIGetTypeInfo(type_index); for (int32_t i = 0; i < info->num_methods; ++i) { if (std::strncmp(info->methods[i].name.data, method_name, info->methods[i].name.size) == 0) { return &(info->methods[i]); } } TVM_FFI_THROW(RuntimeError) << "Cannot find method " << method_name << " in " << type_key; TVM_FFI_UNREACHABLE(); } /*! * \brief helper function to get reflection method function by method info * * \param type_key The type key. * \param method_name The name of the method. * \return The method function. */ inline Function GetMethod(std::string_view type_key, const char* method_name) { const TVMFFIMethodInfo* info = GetMethodInfo(type_key, method_name); return AnyView::CopyFromTVMFFIAny(info->method).cast(); } /*! * \brief Set a field to its default value, calling the factory if applicable. * * When kTVMFFIFieldFlagBitMaskDefaultFromFactory is set, extracts the * Function from default_value_or_factory, calls it with no arguments, * and uses the result. Otherwise, passes default_value_or_factory directly * to the setter. * * \param field_info The field info (must have kTVMFFIFieldFlagBitMaskHasDefault set). * \param field_addr The address of the field in the object. */ inline void SetFieldToDefault(const TVMFFIFieldInfo* field_info, void* field_addr) { if (field_info->flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory) { Function factory = AnyView::CopyFromTVMFFIAny(field_info->default_value_or_factory).cast(); Any default_val = factory(); CallFieldSetter(field_info, field_addr, reinterpret_cast(&default_val)); } else { CallFieldSetter(field_info, field_addr, &(field_info->default_value_or_factory)); } } /*! * \brief Visit each field info of the type info and run callback. * * \tparam Callback The callback function type. * * \param type_info The type info. * \param callback The callback function. * * \note This function calls both the child and parent type info. */ template inline void ForEachFieldInfo(const TypeInfo* type_info, Callback callback) { using ResultType = decltype(callback(type_info->fields)); static_assert(std::is_same_v, "Callback must return void"); // iterate through acenstors in parent to child order // skip the first one since it is always the root object for (int i = 1; i < type_info->type_depth; ++i) { const TVMFFITypeInfo* parent_info = type_info->type_ancestors[i]; for (int j = 0; j < parent_info->num_fields; ++j) { callback(parent_info->fields + j); } } for (int i = 0; i < type_info->num_fields; ++i) { callback(type_info->fields + i); } } /*! * \brief Visit each field info of the type info and run callback which returns bool for early stop. * * \tparam Callback The callback function type, which returns bool for early stop. * * \param type_info The type info. * \param callback_with_early_stop The callback function. * \return true if any of early stop is triggered. * * \note This function calls both the child and parent type info and can be used for searching. */ template inline bool ForEachFieldInfoWithEarlyStop(const TypeInfo* type_info, Callback callback_with_early_stop) { // iterate through acenstors in parent to child order // skip the first one since it is always the root object for (int i = 1; i < type_info->type_depth; ++i) { const TVMFFITypeInfo* parent_info = type_info->type_ancestors[i]; for (int j = 0; j < parent_info->num_fields; ++j) { if (callback_with_early_stop(parent_info->fields + j)) return true; } } for (int i = 0; i < type_info->num_fields; ++i) { if (callback_with_early_stop(type_info->fields + i)) return true; } return false; } /*! * \brief Type attribute names used by the reflection system. * * Each constant names a TypeAttrColumn — a sparse, type-indexed slot that * stores a ``Function`` or any other Any-compliant value. * * - ``TypeAttrDef.def(type_attr::kFoo, ...)`` to register a hook, * - ``TypeAttrColumn(type_attr::kFoo)[type_index]`` to look one up, * */ namespace type_attr { /*! * \brief Zero-arg object allocator. * * Allocates a zero-initialised object of the registered type. For C++ types * this wraps ``metadata->creator``; for Python ``@py_class`` types it is a * ``calloc``-based allocator. * * Signature: ``() -> TSelf``, where ``TSelf`` is a subclass of ObjectRef. */ inline constexpr const char* kNew = "__ffi_new__"; /*! * \brief Packed constructor — creates **and** initialises a new object. * * Used by ``@c_class`` auto-generated ``__init__``: the Python side calls * ``self.__init_handle_by_constructor__(ffi_init, *args)`` which invokes * this function to produce a fully initialised object. * * For C++ types with ``refl::init()``, the function is the explicit * init; otherwise ``ffi._RegisterFFIInit`` generates a default one using reflection. * * Signature: ``(*args, **kwargs) -> TSelf``, where ``TSelf`` is a subclass of ObjectRef. * * Keyword arguments are packed as ``[KWARGS, key0, val0, key1, val1, ...]``. */ inline constexpr const char* kInit = "__ffi_init__"; /*! * \brief Convert ``AnyView`` to a specific reflected ``TSelf`` type. * * Registered via ``TypeAttrDef.def(kConvert, &FFIConvertFromAnyViewToObjectRef)`` * for every type that calls ``.ref()``. Used by the Python type converter * to marshal values into the correct ``TSelf`` subclass. * * Signature: ``(AnyView src) -> TSelf``, where ``TSelf`` is a subclass of ObjectRef. */ inline constexpr const char* kConvert = "__ffi_convert__"; /*! * \brief Shallow-copy factory. * * Allocates a new object and copies all reflected field values from the * source. Used by Python ``copy.copy()`` via ``__copy__`` and by * ``copy.replace()`` via ``__replace__``. * * Signature: ``(TSelf self) -> TSelf``, where ``TSelf`` is a subclass of ObjectRef. */ inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__"; /*! * \brief Custom recursive repr hook. * * If registered, ``RecursiveRepr`` (Python ``__repr__``) calls this instead * of the default field-by-field formatting. The hook receives a callback * ``fn_repr`` to recursively format nested values, avoiding infinite loops. * * Signature: ``(TSelf self, fn_repr: FnRepr) -> String``, where ``TSelf`` is a subclass of * ObjectRef, and ``FnRepr: (AnyView value) -> String`` formats a nested value. */ inline constexpr const char* kRepr = "__ffi_repr__"; /*! * \brief Custom recursive hash hook. * * If registered, ``RecursiveHash`` (Python ``hash()``) calls this instead * of the default field-by-field hashing. The hook receives a callback * ``fn_hash`` to recursively hash nested values, with cycle detection. * * Signature: ``(TSelf self, fn_hash: FnHash) -> int64``, where ``TSelf`` is a subclass of * ObjectRef, and ``FnHash: (AnyView value) -> int64`` hashes a nested value. */ inline constexpr const char* kHash = "__ffi_hash__"; /*! * \brief Custom recursive equality hook. * * If registered, ``RecursiveEq`` (Python ``==``) calls this instead of the * default field-by-field comparison. Falls back to ``kCompare`` if only the * compare hook is present. The hook receives a callback ``fn_eq`` to * recursively compare nested values. * * Signature: ``(TSelf lhs, TSelf rhs, fn_eq: FnEq) -> bool``, where ``TSelf`` is a subclass of * ObjectRef, and ``FnEq: (AnyView lhs, AnyView rhs) -> bool`` compares nested values. */ inline constexpr const char* kEq = "__ffi_eq__"; /*! * \brief Custom recursive three-way comparison hook. * * If registered, ``RecursiveCompare`` (Python ``<``, ``<=``, ``>``, ``>=``) * calls this. Also used as fallback for ``RecursiveEq`` when ``kEq`` is not * registered. The hook receives a callback ``fn_cmp`` to recursively * compare nested values. * * Signature: ``(TSelf lhs, TSelf rhs, fn_cmp: FnCmp) -> int32``, where ``TSelf`` is a subclass of * ObjectRef, and ``FnCmp: (TSelf lhs, TSelf rhs) -> int32`` returns < 0, == 0, or > 0. */ inline constexpr const char* kCompare = "__ffi_compare__"; /*! * \brief Custom Any-level hash for use as ``Map``/``Dict`` key. * * Unlike ``kHash`` (which operates on ``ObjectRef`` and is recursive), this * hook operates at the ``Any`` level — it is called by ``AnyHash`` whenever * an ``Any`` holding an object of this type is used as a container key. * * Can be either a raw C function pointer (fast path, no boxing overhead) or * a ``Function`` object. * * Raw pointer signature: ``int64_t (*)(const Any& src)`` * * Function signature: ``(Any src) -> int64`` */ inline constexpr const char* kAnyHash = "__any_hash__"; /*! * \brief Custom Any-level equality for use as ``Map``/``Dict`` key. * * Unlike ``kEq`` (which operates on ``ObjectRef``), this hook operates at the * ``Any`` level — called by ``AnyEqual`` for container key comparison. * * Can be either a raw C function pointer (fast path) or a ``Function`` object. * * Raw pointer signature: ``bool (*)(const Any& lhs, const Any& rhs)`` * * Function signature: ``(Any lhs, Any rhs) -> bool`` */ inline constexpr const char* kAnyEqual = "__any_equal__"; /*! * \brief Custom structural hash hook (used by ``StructuralHash``). * * Unlike ``kHash`` (which is for ``RecursiveHash`` / Python ``hash()``), * this is for the ``StructuralHash`` system that supports def/use region * semantics (``map_free_vars``). The hook receives the current accumulated * hash and a callback to recursively hash sub-values. * * Signature: ``(TSelf self, int64 init_hash, Function hash_cb) -> int64``, where * ``TSelf`` is a subclass of ``ObjectRef``. * * ``hash_cb(AnyView val, int64 init_hash, bool def_region) -> int64`` * recursively hashes a sub-value; ``def_region`` controls free-variable * mapping. */ inline constexpr const char* kSHash = "__s_hash__"; /*! * \brief Custom structural equality hook (used by ``StructuralEqual``). * * Unlike ``kEq`` (which is for ``RecursiveEq`` / Python ``==``), this is for * the ``StructuralEqual`` system that supports def/use region semantics and * alpha-equivalence of bound variables. The hook receives a callback to * recursively compare sub-values. * * Signature: ``(TSelf lhs, TSelf rhs, Function eq_cb) -> bool``, where ``TSelf`` is a * subclass of ``ObjectRef``. * * ``eq_cb(AnyView lhs, AnyView rhs, bool def_region, AnyView field_name) -> bool`` * recursively compares sub-values; ``def_region`` controls free-variable * mapping; ``field_name`` is used for mismatch path reporting. */ inline constexpr const char* kSEqual = "__s_equal__"; /*! * \brief Serialize object data to a JSON-compatible ``Map``. * * If registered, ``ToJSONGraph`` calls this instead of the default * field-by-field serialization. Allows types with non-trivial internal * state (e.g. ``TInt`` storing a plain ``int64_t``) to define a compact * custom JSON representation. * * Signature: ``(TSelf self) -> Map``, where ``TSelf`` is a subclass of * ``ObjectRef``. */ inline constexpr const char* kDataToJson = "__data_to_json__"; /*! * \brief Deserialize object data from a JSON-compatible ``Map``. * * Counterpart to ``kDataToJson``. If registered, ``FromJSONGraph`` calls * this to reconstruct the object from its serialized ``Map`` representation * instead of using field-by-field deserialization. * * Signature: ``(Map json_data) -> TSelf``, where ``TSelf`` is a subclass of * ``ObjectRef``. */ inline constexpr const char* kDataFromJson = "__data_from_json__"; /*! * \brief Per-class enum entry registry. * * Maps each variant's name to its registered singleton for an * ``EnumObj`` subclass. Populated by ``refl::EnumDef("Name")`` on * the C++ side and by ``Enum`` subclass declarations on the Python * side; both languages share the same underlying storage. * * Value type: ``Dict``. */ inline constexpr const char* kEnumEntries = "__ffi_enum_entries__"; /*! * \brief Per-class column holding extensible attributes for enum variants. * * The outer dict is keyed by extensible-attribute name; each value is a * list indexed by the variant's ordinal (``EnumObj::_value``). Written * by ``refl::EnumDef::set_attr(...)`` on the C++ side and by the * ``EnumAttrMap`` returned from Python ``Enum.def_attr(...)``. * * Value type: ``Dict>``. */ inline constexpr const char* kEnumAttrs = "__ffi_enum_attrs__"; /*! * \brief Per-class payload-to-variant index for enums. * * Parallel to ``kEnumEntries`` (name → variant) but keyed by the * user-visible payload carried on each variant — i.e. the ``value`` * field on Python ``IntEnum`` / ``StrEnum`` subclasses (``int`` or * ``str``) or the equivalent payload field on a C++ ``EnumObj`` * subclass. Populated by the creator of each variant (Python or C++) * when the variant has a payload; absent or partially populated * otherwise. Consumed by FFI converters to resolve a raw payload * (``int``/``str``) to its singleton variant in O(1). * * Value type: ``Dict``. */ inline constexpr const char* kEnumValueEntries = "__ffi_enum_value_entries__"; } // namespace type_attr /*! * \brief C++17 constexpr helper: convert a ``const char*`` to ``TVMFFIByteArray``. * \param s A null-terminated string literal. * \return A ``TVMFFIByteArray`` whose ``data`` points to *s* and ``size`` equals ``strlen(s)``. */ inline constexpr TVMFFIByteArray AsByteArray(const char* s) { return {s, std::char_traits::length(s)}; } } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_REFLECTION_ACCESSOR_H_ tvm-ffi-0.1.12/include/tvm/ffi/reflection/creator.h000066400000000000000000000166041521067262500221170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/creator.h * \brief Reflection-based creator to create objects from type key and fields. */ #ifndef TVM_FFI_REFLECTION_CREATOR_H_ #define TVM_FFI_REFLECTION_CREATOR_H_ #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Create an empty object via the type's native creator or ``__ffi_new__`` type attr. * * Falls back to the ``__ffi_new__`` type attribute (used by Python-defined types) * when the native ``metadata->creator`` is NULL. * * \param type_info The type info for the object to create. * \return An owned ObjectPtr to the newly allocated (zero-initialized) object. * \throws RuntimeError if neither creator nor __ffi_new__ is available. */ inline ObjectPtr CreateEmptyObject(const TVMFFITypeInfo* type_info) { // Fast path: native C++ creator if (type_info->metadata != nullptr && type_info->metadata->creator != nullptr) { TVMFFIObjectHandle handle; TVM_FFI_CHECK_SAFE_CALL(type_info->metadata->creator(&handle)); return details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(handle)); } // Fallback: __ffi_new__ type attr (Python-defined types) constexpr TVMFFIByteArray kFFINewAttrName = reflection::AsByteArray(reflection::type_attr::kNew); const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&kFFINewAttrName); if (column != nullptr) { int32_t offset = type_info->type_index - column->begin_index; if (offset >= 0 && offset < column->size) { AnyView attr_view = AnyView::CopyFromTVMFFIAny(column->data[offset]); if (auto opt_func = attr_view.try_cast()) { ObjectRef obj_ref = (*opt_func)().cast(); return details::ObjectUnsafe::ObjectPtrFromObjectRef(std::move(obj_ref)); } } } TVM_FFI_THROW(RuntimeError) << "Type `" << TypeIndexToTypeKey(type_info->type_index) << "` does not support reflection creation" << " (no native creator or __ffi_new__ type attr)"; TVM_FFI_UNREACHABLE(); } /*! * \brief Check whether a type supports reflection creation. * * Returns true if the type has a native creator or a ``__ffi_new__`` type attr. * * \param type_info The type info to check. * \return true if CreateEmptyObject would succeed. */ inline bool HasCreator(const TVMFFITypeInfo* type_info) { if (type_info->metadata != nullptr && type_info->metadata->creator != nullptr) { return true; } constexpr TVMFFIByteArray kFFINewAttrName = reflection::AsByteArray(reflection::type_attr::kNew); const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&kFFINewAttrName); if (column != nullptr) { int32_t offset = type_info->type_index - column->begin_index; if (offset >= 0 && offset < column->size && column->data[offset].type_index == TypeIndex::kTVMFFIFunction) { return true; } } return false; } namespace reflection { namespace details { /*! * \brief Convert an AnyView to a specific reflected object type. * * \tparam TObjectRef The object reference type to convert to. * \param input The AnyView to convert. * \return The converted object. */ template TObjectRef FFIConvertFromAnyViewToObjectRef(AnyView input) { TVMFFIAny input_pod = input.CopyToTVMFFIAny(); if (auto opt = TypeTraits::TryCastFromAnyView(&input_pod)) { return *std::move(opt); } TVM_FFI_THROW(TypeError) << "Cannot cast from `" << TypeIndexToTypeKey(input_pod.type_index) << "` to `" << TypeTraits::TypeStr() << "`"; } } // namespace details /*! * \brief helper wrapper class of TVMFFITypeInfo to create object based on reflection. */ class ObjectCreator { public: /*! * \brief Constructor * \param type_key The type key. */ explicit ObjectCreator(std::string_view type_key) : ObjectCreator(TVMFFIGetTypeInfo(TypeKeyToIndex(type_key))) {} /*! * \brief Constructor * \param type_info The type info. */ explicit ObjectCreator(const TVMFFITypeInfo* type_info) : type_info_(type_info) { if (!HasCreator(type_info)) { TVM_FFI_THROW(RuntimeError) << "Type `" << TypeIndexToTypeKey(type_info->type_index) << "` does not support default constructor, " << "as a result cannot be created via reflection"; } } /** * \brief Create an object from a map of fields. * \param fields The fields of the object. * \return The created object. */ Any operator()(const Map& fields) const { ObjectPtr ptr = CreateEmptyObject(type_info_); size_t match_field_count = 0; ForEachFieldInfo(type_info_, [&](const TVMFFIFieldInfo* field_info) { String field_name(field_info->name); void* field_addr = reinterpret_cast(ptr.get()) + field_info->offset; if (fields.count(field_name) != 0) { Any field_value = fields[field_name]; CallFieldSetter(field_info, field_addr, reinterpret_cast(&field_value)); ++match_field_count; } else if (field_info->flags & kTVMFFIFieldFlagBitMaskHasDefault) { SetFieldToDefault(field_info, field_addr); } else { TVM_FFI_THROW(TypeError) << "Required field `" << String(field_info->name.data, field_info->name.size) << "` not set in type `" << String(type_info_->type_key.data, type_info_->type_key.size) << "`"; } }); if (match_field_count == fields.size()) return ObjectRef(ptr); // report error that checks if contains extra fields that are not in the type auto check_field_name = [&](const String& field_name) { bool found = false; ForEachFieldInfoWithEarlyStop(type_info_, [&](const TVMFFIFieldInfo* field_info) { if (field_name.compare(field_info->name) == 0) { found = true; return true; } return false; }); return found; }; for (const auto& [field_name, _] : fields) { if (!check_field_name(field_name)) { TVM_FFI_THROW(TypeError) << "Type `" << String(type_info_->type_key.data, type_info_->type_key.size) << "` does not have field `" << field_name << "`"; } } TVM_FFI_UNREACHABLE(); } private: const TVMFFITypeInfo* type_info_; }; } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_REFLECTION_CREATOR_H_ tvm-ffi-0.1.12/include/tvm/ffi/reflection/enum_def.h000066400000000000000000000143611521067262500222400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/enum_def.h * \brief Builder for registering enum instances on ``EnumObj`` subclasses. */ #ifndef TVM_FFI_REFLECTION_ENUM_DEF_H_ #define TVM_FFI_REFLECTION_ENUM_DEF_H_ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace reflection { /*! * \brief Builder that registers a single enum instance on ``EnumClsObj``. * * Each ``EnumDef("Name")`` call allocates a fresh dense ordinal * (``= len(existing entries)``), constructs a variant with ``_value`` and * ``_name`` populated, and writes it into the per-class registry stored in * the ``type_attr::kEnumEntries`` TypeAttr column. Subsequent * ``.set_attr(...)`` calls write *extensible attributes* — per-variant * metadata attached outside the variant's declared fields — into the * per-class ``type_attr::kEnumAttrs`` column. Python bindings of the * same ``type_key`` see every C++-registered variant and every extensible * attribute through the matching ``Enum.def_attr`` surface. * * \tparam EnumClsObj An ``Object`` subclass deriving from ``EnumObj``. * * \code{.cpp} * namespace refl = ::tvm::ffi::reflection; * refl::EnumDef("Add").set_attr("has_side_effects", false); * refl::EnumDef("Mul").set_attr("has_side_effects", false); * \endcode */ template >> class EnumDef : public ReflectionDefBase { public: /*! * \brief Register a new instance named ``instance_name`` on ``EnumClsObj``. * \param instance_name The instance's string name (e.g., ``"Add"``). */ explicit EnumDef(const char* instance_name) : type_index_(EnumClsObj::RuntimeTypeIndex()), name_(instance_name) { Dict entries = EnsureEntriesDict(); String name_str(name_); if (entries.count(name_str) != 0) { TVM_FFI_THROW(RuntimeError) << "Duplicate enum entry `" << name_ << "` for type `" << EnumClsObj::_type_key << "`"; } ordinal_ = static_cast(entries.size()); ObjectPtr obj = make_object(); obj->_value = ordinal_; obj->_name = name_str; instance_ = Enum(ObjectPtr(std::move(obj))); entries.Set(name_str, instance_); // Ensure the attrs dict exists so later ``set_attr`` calls can mutate it. EnsureAttrsDict(); } /*! * \brief Write an *extensible attribute* for this enum variant. * * Writes land in the per-class ``type_attr::kEnumAttrs`` column and * are visible to every binder of the same ``type_key`` — including * Python readers via ``Enum.def_attr`` / ``Enum.attr_dict``. Distinct * from declared fields on ``EnumClsObj``: declared fields are part of * the variant's schema and set during construction, whereas * extensible attributes live outside the variant object and may be * attached by any consumer at any time. * * \tparam T The value type. * \param attr_name The extensible-attribute name (e.g., * ``"has_side_effects"``). * \param value The value to store for this variant's ordinal. * \return Reference to this builder for chaining. */ template EnumDef& set_attr(const char* attr_name, T value) { Dict> attrs = EnsureAttrsDict(); String attr_key(attr_name); List column; auto it = attrs.find(attr_key); if (it == attrs.end()) { column = List(); attrs.Set(attr_key, column); } else { column = (*it).second; } while (static_cast(column.size()) <= ordinal_) { column.push_back(Any(nullptr)); } column.Set(ordinal_, Any(std::move(value))); return *this; } /*! \brief Return the registered instance (for tests / advanced callers). */ Enum instance() const { return instance_; } /*! \brief Return the ordinal assigned to this instance. */ int64_t ordinal() const { return ordinal_; } private: Dict EnsureEntriesDict() { return EnsureDict>(type_attr::kEnumEntries); } Dict> EnsureAttrsDict() { return EnsureDict>>(type_attr::kEnumAttrs); } template DictT EnsureDict(const char* attr_name) { TVMFFIByteArray name_array = {attr_name, std::char_traits::length(attr_name)}; const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&name_array); if (column != nullptr) { int32_t offset = type_index_ - column->begin_index; if (offset >= 0 && offset < column->size) { const TVMFFIAny* stored = &column->data[offset]; if (stored->type_index != kTVMFFINone) { return AnyView::CopyFromTVMFFIAny(*stored).cast(); } } } DictT fresh; TVMFFIAny value_any = AnyView(fresh).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any)); return fresh; } int32_t type_index_; const char* name_; int64_t ordinal_; Enum instance_; }; } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_REFLECTION_ENUM_DEF_H_ tvm-ffi-0.1.12/include/tvm/ffi/reflection/overload.h000066400000000000000000000462651521067262500223010ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/overload.h * \brief Registry of reflection metadata, supporting function overloading. */ #ifndef TVM_FFI_EXTRA_OVERLOAD_H #define TVM_FFI_EXTRA_OVERLOAD_H #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace details { struct OverloadBase { public: // Try Call function pointer type, return true if matched and called using FnPtr = bool (*)(OverloadBase*, const AnyView*, int32_t, Any*); explicit OverloadBase(int32_t num_args, std::optional name) : num_args_(num_args), name_(name ? std::move(*name) : ""), name_ptr_(name ? &this->name_ : nullptr) {} virtual void Register(std::unique_ptr overload) = 0; virtual FnPtr GetTryCallPtr() = 0; virtual void GetMismatchMessage(std::ostringstream& os, const AnyView* args, int32_t num_args) = 0; virtual ~OverloadBase() = default; OverloadBase(const OverloadBase&) = delete; OverloadBase& operator=(const OverloadBase&) = delete; public: static constexpr int32_t kAllMatched = -1; // a fast cache for last matched arg index // on 64-bit platform, this is packed in the same 8 byte with num_args_ int32_t last_mismatch_index_{kAllMatched}; // some constant helper args const int32_t num_args_; const std::string name_; const std::string* const name_ptr_; }; template struct CaptureTupleAux; template struct CaptureTupleAux> { using type = std::tuple>...>; }; template struct TypedOverload : OverloadBase { public: static_assert(std::is_same_v>, "Callable must be value type"); using FuncInfo = details::FunctionInfo; using PackedArgs = typename FuncInfo::ArgType; using Ret = typename FuncInfo::RetType; using CaptureTuple = typename CaptureTupleAux::type; using OverloadBase::name_; using OverloadBase::name_ptr_; using typename OverloadBase::FnPtr; static constexpr auto kNumArgs = FuncInfo::num_args; static constexpr auto kSeq = std::make_index_sequence{}; explicit TypedOverload(const Callable& f, std::optional name = std::nullopt) : OverloadBase(kNumArgs, std::move(name)), f_(f) {} explicit TypedOverload(Callable&& f, std::optional name = std::nullopt) : OverloadBase(kNumArgs, std::move(name)), f_(std::move(f)) {} bool TryCall(const AnyView* args, int32_t num_args, Any* rv) { if (num_args != kNumArgs) return false; CaptureTuple captures{}; if (!TrySetAux(kSeq, captures, args)) return false; // now all captures are set if constexpr (std::is_same_v) { CallAux(kSeq, captures); return true; } else { *rv = CallAux(kSeq, captures); return true; } } void Register(std::unique_ptr overload) override { TVM_FFI_ICHECK(false) << "This should never be called."; } FnPtr GetTryCallPtr() final { // lambda without a capture can be converted to function pointer return [](OverloadBase* base, const AnyView* args, int32_t num_args, Any* rv) -> bool { return static_cast*>(base)->TryCall(args, num_args, rv); }; } void GetMismatchMessage(std::ostringstream& os, const AnyView* args, int32_t num_args) final { FGetFuncSignature f_sig = FuncInfo::Sig; if (num_args != kNumArgs) { os << "Mismatched number of arguments when calling: `" << name_ << " " << (f_sig == nullptr ? "" : (*f_sig)()) << "`. Expected " << kNumArgs << " arguments"; } else { GetMismatchMessageAux<0>(os, args, num_args); } } private: template void GetMismatchMessageAux(std::ostringstream& os, const AnyView* args, [[maybe_unused]] int32_t num_args) { if constexpr (I < kNumArgs) { if (this->last_mismatch_index_ == static_cast(I)) { TVMFFIAny any_data = args[I].CopyToTVMFFIAny(); FGetFuncSignature f_sig = FuncInfo::Sig; using Type = std::decay_t>; os << "Mismatched type on argument #" << I << " when calling: `" << name_ << " " << (f_sig == nullptr ? "" : (*f_sig)()) << "`. Expected `" << Type2Str::v() << "` but got `" << TypeTraits::GetMismatchTypeInfo(&any_data) << '`'; } else { GetMismatchMessageAux(os, args, num_args); } } // end of recursion } template Ret CallAux(std::index_sequence, CaptureTuple& tuple) { /// NOTE: this works for T, const T, const T&, T&& argument types // NOLINTNEXTLINE(bugprone-unchecked-optional-access) return f_(static_cast>(std::move(*std::get(tuple)))...); } template bool TrySetAux(std::index_sequence, CaptureTuple& tuple, const AnyView* args) { return (TrySetOne(tuple, args) && ...); } template bool TrySetOne(CaptureTuple& tuple, const AnyView* args) { using Type = std::decay_t>; auto& capture = std::get(tuple); if constexpr (std::is_same_v) { capture = args[I]; return true; } else if constexpr (std::is_same_v) { capture = Any(args[I]); return true; } else { capture = args[I].template try_cast(); if (capture.has_value()) return true; // slow path: record the last mismatch index this->last_mismatch_index_ = static_cast(I); return false; } } protected: Callable f_; }; template inline auto CreateNewOverload(Callable&& f, std::string name) { using Type = TypedOverload>; return std::make_unique(std::forward(f), std::move(name)); } template struct OverloadedFunction : TypedOverload { public: using TypedBase = TypedOverload; using OverloadBase::name_; using OverloadBase::name_ptr_; using TypedBase::GetTryCallPtr; using TypedBase::kNumArgs; using TypedBase::kSeq; using TypedBase::TypedBase; // constructors using typename OverloadBase::FnPtr; using typename TypedBase::Ret; void Register(std::unique_ptr overload) final { const auto fptr = overload->GetTryCallPtr(); overloads_.emplace_back(std::move(overload), fptr); } void operator()(const AnyView* args, int32_t num_args, Any* rv) { // fast path: only add a little overhead when no overloads if (overloads_.size() == 0) { return unpack_call(kSeq, name_ptr_, f_, args, num_args, rv); } // this can be inlined by compiler, don't worry if (this->TryCall(args, num_args, rv)) return; // virtual calls cannot be inlined, so we fast check the num_args first // we also de-virtualize the fptr to reduce one more indirection for (const auto& [overload, fptr] : overloads_) { if (overload->num_args_ != num_args) continue; if (fptr(overload.get(), args, num_args, rv)) return; } this->HandleOverloadFailure(args, num_args); } private: void HandleOverloadFailure(const AnyView* args, int32_t num_args) { std::ostringstream oss; int32_t i = 0; oss << "Overload #" << i++ << ": "; this->GetMismatchMessage(oss, args, num_args); for (const auto& [overload, _] : overloads_) { oss << "\nOverload #" << i++ << ": "; overload->GetMismatchMessage(oss, args, num_args); } TVM_FFI_THROW(TypeError) << "No matching overload found when calling: `" << name_ << "` with " << num_args << " arguments:\n" << std::move(oss).str(); } using TypedBase::f_; std::vector, FnPtr>> overloads_; }; } // namespace details /*! \brief Reflection namespace */ namespace reflection { /*! * \brief Helper to register Object's reflection metadata. * \tparam Class The class type. * * \code * namespace refl = tvm::ffi::reflection; * refl::ObjectDef().def_ro("my_field", &MyClass::my_field); * \endcode */ template class OverloadObjectDef : private ObjectDef { public: /*! \brief The super class */ using Super = ObjectDef; /*! * \brief Constructor * \tparam ExtraArgs The extra arguments. * \param extra_args The extra arguments. */ template explicit OverloadObjectDef(ExtraArgs&&... extra_args) : Super(std::forward(extra_args)...) {} /*! * \brief Define a readonly field. * * \tparam Class The class type. * \tparam T The field type. * \tparam Extra The extra arguments. * * \param name The name of the field. * \param field_ptr The pointer to the field. * \param extra The extra arguments that can be docstring or default value. * * \return The reflection definition. */ template TVM_FFI_INLINE OverloadObjectDef& def_ro(const char* name, T BaseClass::* field_ptr, Extra&&... extra) { /// NOTE: we don't allow properties to be overloaded Super::def_ro(name, field_ptr, std::forward(extra)...); return *this; } /*! * \brief Define a read-write field. * * \tparam Class The class type. * \tparam T The field type. * \tparam Extra The extra arguments. * * \param name The name of the field. * \param field_ptr The pointer to the field. * \param extra The extra arguments that can be docstring or default value. * * \return The reflection definition. */ template TVM_FFI_INLINE OverloadObjectDef& def_rw(const char* name, T BaseClass::* field_ptr, Extra&&... extra) { /// NOTE: we don't allow properties to be overloaded Super::def_rw(name, field_ptr, std::forward(extra)...); return *this; } /*! * \brief Define a method. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the method. * \param func The function to be registered. * \param extra The extra arguments that can be docstring. * * \return The reflection definition. */ template TVM_FFI_INLINE OverloadObjectDef& def(const char* name, Func&& func, Extra&&... extra) { RegisterMethod(name, false, std::forward(func), std::forward(extra)...); return *this; } /*! * \brief Define a static method. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the method. * \param func The function to be registered. * \param extra The extra arguments that can be docstring. * * \return The reflection definition. */ template TVM_FFI_INLINE OverloadObjectDef& def_static(const char* name, Func&& func, Extra&&... extra) { RegisterMethod(name, true, std::forward(func), std::forward(extra)...); return *this; } /*! * \brief Register a constructor for this object type. * * This method registers a static `__init__` method that constructs an instance * of the object with the specified argument types. The constructor can be invoked * from Python or other FFI bindings. * * \tparam Args The argument types for the constructor. * \tparam Extra Additional arguments (e.g., docstring). * * \param init_func An instance of `init` specifying constructor signature. * \param extra Optional additional metadata such as docstring. * * \return Reference to this `ObjectDef` for method chaining. * * Example: * \code * refl::ObjectDef() * .def(refl::init(), "Constructor docstring"); * \endcode */ template TVM_FFI_INLINE OverloadObjectDef& def([[maybe_unused]] init init_func, Extra&&... extra) { this->has_explicit_init_ = true; bool is_first = registered_fields_.find(kInitMethodName) == registered_fields_.end(); // Register as method (enables overloading). RegisterMethod(kInitMethodName, true, &init::template execute, std::forward(extra)...); // On first init registration, also mirror into __ffi_init__ TypeAttrColumn. // The Function wraps OverloadedFunction, so subsequent overloads added // via Register() are automatically visible through the TypeAttrColumn. if (is_first) { const TVMFFITypeInfo* tinfo = TVMFFIGetTypeInfo(type_index_); constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit); for (int32_t i = 0; i < tinfo->num_methods; ++i) { if (tinfo->methods[i].name.size == attr_name.size && std::strncmp(tinfo->methods[i].name.data, attr_name.data, attr_name.size) == 0) { TVM_FFI_CHECK_SAFE_CALL( TVMFFITypeRegisterAttr(type_index_, &attr_name, &tinfo->methods[i].method)); break; } } } return *this; } private: using ReflectionDefBase::ApplyExtraInfoTrait; using ReflectionDefBase::WrapFunction; using Super::kInitMethodName; using Super::type_index_; using Super::type_key_; template static auto GetOverloadMethod(std::string name, Func&& func) { using WrapFn = decltype(WrapFunction(std::forward(func))); using OverloadFn = ::tvm::ffi::details::OverloadedFunction>; return ffi::Function::FromPackedInplace(WrapFunction(std::forward(func)), std::move(name)); } template static auto NewOverload(std::string name, Func&& func) { return ::tvm::ffi::details::CreateNewOverload(WrapFunction(std::forward(func)), std::move(name)); } template void RegisterExtraInfo(ExtraArgs&&... extra_args) { TVMFFITypeMetadata info; info.total_size = sizeof(Class); info.structural_eq_hash_kind = Class::_type_s_eq_hash_kind; info.creator = nullptr; info.doc = TVMFFIByteArray{nullptr, 0}; if constexpr (std::is_default_constructible_v) { info.creator = ReflectionDefBase::ObjectCreatorDefault; } else if constexpr (std::is_constructible_v) { info.creator = ReflectionDefBase::ObjectCreatorUnsafeInit; } // apply extra info traits ((ApplyExtraInfoTrait(&info, std::forward(extra_args)), ...)); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMetadata(type_index_, &info)); } template void RegisterField(const char* name, T BaseClass::* field_ptr, bool writable, ExtraArgs&&... extra_args) { static_assert(std::is_base_of_v, "BaseClass must be a base class of Class"); FieldInfoBuilder info; info.name = TVMFFIByteArray{name, std::char_traits::length(name)}; info.field_static_type_index = TypeToFieldStaticTypeIndex::value; // store byte offset and setter, getter // so the same setter can be reused for all the same type info.offset = GetFieldByteOffsetToObject(field_ptr); info.size = sizeof(T); info.alignment = alignof(T); info.flags = 0; if (writable) { info.flags |= kTVMFFIFieldFlagBitMaskWritable; } info.getter = ReflectionDefBase::FieldGetter; info.setter = reinterpret_cast(ReflectionDefBase::FieldSetter); // initialize default value to nullptr info.default_value_or_factory = AnyView(nullptr).CopyToTVMFFIAny(); info.doc = TVMFFIByteArray{nullptr, 0}; info.metadata_.emplace_back("type_schema", ::tvm::ffi::details::TypeSchema::v()); // apply field info traits ((ApplyFieldInfoTrait(&info, std::forward(extra_args)), ...)); // call register std::string metadata_str = Metadata::ToJSON(info.metadata_); info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterField(type_index_, &info)); } // register a method template void RegisterMethod(const char* name, bool is_static, Func&& func, Extra&&... extra) { using FuncInfo = ::tvm::ffi::details::FunctionInfo>; MethodInfoBuilder info; info.name = TVMFFIByteArray{name, std::char_traits::length(name)}; info.doc = TVMFFIByteArray{nullptr, 0}; info.flags = 0; if (is_static) { info.flags |= kTVMFFIFieldFlagBitMaskIsStaticMethod; } auto method_name = std::string(type_key_) + "." + name; // if an overload method exists, register to existing overload function if (const auto overload_it = registered_fields_.find(name); overload_it != registered_fields_.end()) { ::tvm::ffi::details::OverloadBase* overload_ptr = overload_it->second; return overload_ptr->Register(NewOverload(std::move(method_name), std::forward(func))); } // first time registering overload method auto [method, overload_ptr] = GetOverloadMethod(std::move(method_name), std::forward(func)); registered_fields_.try_emplace(name, overload_ptr); info.method = AnyView(method).CopyToTVMFFIAny(); info.metadata_.emplace_back("type_schema", FuncInfo::TypeSchema()); // apply method info traits ((ApplyMethodInfoTrait(&info, std::forward(extra)), ...)); std::string metadata_str = Metadata::ToJSON(info.metadata_); info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index_, &info)); } std::unordered_map registered_fields_; }; } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_OVERLOAD_H tvm-ffi-0.1.12/include/tvm/ffi/reflection/registry.h000066400000000000000000001172221521067262500223260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/reflection/registry.h * \brief Registry of reflection metadata. */ #ifndef TVM_FFI_REFLECTION_REGISTRY_H_ #define TVM_FFI_REFLECTION_REGISTRY_H_ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! \brief Reflection namespace */ namespace reflection { /*! * \brief Types of temporary metadata hold in FieldInfoBuilder and MethodInfoBuilder, * before they are filled into final C metadata */ using _MetadataType = std::vector>; // NOLINT(bugprone-reserved-identifier) /*! * \brief Builder for TVMFFIFieldInfo * \sa TVMFFIFieldInfo */ struct FieldInfoBuilder : public TVMFFIFieldInfo { /*! \brief Temporary metadata info to be filled into TVMFFIFieldInfo::metadata */ _MetadataType metadata_; }; /*! * \brief Builder for TVMFFIMethodInfo * \sa TVMFFIMethodInfo */ struct MethodInfoBuilder : public TVMFFIMethodInfo { /*! \brief Temporary metadata info to be filled into TVMFFIMethodInfo::metadata */ _MetadataType metadata_; }; /*! * \brief Trait that can be used to set information attached to a field or a method. * \sa DefaultValue, AttachFieldFlag */ struct InfoTrait {}; /*! \brief User-supplied metadata attached to a field or a method */ class Metadata : public InfoTrait { public: /*! * \brief Constructor * \param dict The initial dictionary */ Metadata(std::initializer_list> dict) : dict_(dict) {} /*! * \brief Move metadata into `FieldInfoBuilder` * \param info The field info builder. */ inline void Apply(FieldInfoBuilder* info) const { this->Apply(&info->metadata_); } /*! * \brief Move metadata into `MethodInfoBuilder` * \param info The method info builder. */ inline void Apply(MethodInfoBuilder* info) const { this->Apply(&info->metadata_); } private: friend class GlobalDef; template friend class ObjectDef; template friend class OverloadObjectDef; /*! * \brief Move metadata into a vector of key-value pairs. * \param out The output vector. */ inline void Apply(_MetadataType* out) const { std::copy(std::make_move_iterator(dict_.begin()), std::make_move_iterator(dict_.end()), std::back_inserter(*out)); } /*! \brief Convert the metadata to JSON string */ static std::string ToJSON(const _MetadataType& metadata) { using ::tvm::ffi::details::StringObj; std::ostringstream os; os << "{"; bool first = true; for (const auto& [key, value] : metadata) { if (!first) { os << ","; } os << "\"" << key << "\":"; if (std::optional v = value.as()) { os << *v; } else if (std::optional v = value.as()) { os << (*v ? "true" : "false"); } else if (std::optional v = value.as()) { String escaped = EscapeStringJSON(*v); os << escaped.c_str(); } else { TVM_FFI_LOG_AND_THROW(TypeError) << "Metadata can be only int, bool or string, but on key `" << key << "`, the type is " << value.GetTypeKey(); } first = false; } os << "}"; return os.str(); } std::vector> dict_; }; /*! * \brief Trait that can be used to set field default value */ class DefaultValue : public InfoTrait { public: /*! * \brief Constructor * \param value The value to be set */ explicit DefaultValue(Any value) : value_(std::move(value)) {} /*! * \brief Apply the default value to the field info * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { info->default_value_or_factory = AnyView(value_).CopyToTVMFFIAny(); info->flags |= kTVMFFIFieldFlagBitMaskHasDefault; } private: Any value_; }; /*! \brief Lowercase alias for DefaultValue, mirrors Python ``default``. */ class default_value : public DefaultValue { using DefaultValue::DefaultValue; }; /*! * \brief Trait that can be used to set field default factory. * * A default factory is a callable () -> Any that is invoked each time * a default value is needed, producing a fresh value. This is important * for mutable defaults (e.g., Array, Map) to avoid aliasing. */ class DefaultFactory : public InfoTrait { public: /*! * \brief Constructor * \param factory The factory function to be called to produce default values. */ explicit DefaultFactory(Function factory) : factory_(std::move(factory)) {} /*! * \brief Apply the default factory to the field info * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { info->default_value_or_factory = AnyView(factory_).CopyToTVMFFIAny(); info->flags |= kTVMFFIFieldFlagBitMaskHasDefault; info->flags |= kTVMFFIFieldFlagBitMaskDefaultFromFactory; } private: Function factory_; }; /*! \brief Lowercase alias for DefaultFactory, mirrors Python ``default_factory``. */ class default_factory : public DefaultFactory { using DefaultFactory::DefaultFactory; }; /*! * \brief Trait that can be used to attach field flag */ class AttachFieldFlag : public InfoTrait { public: /*! * \brief Attach a field flag to the field * \param flag The flag to be set */ explicit AttachFieldFlag(int32_t flag) : flag_(flag) {} /*! * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashDefRecursive * * The field enters a recursive def region: free vars discovered both at * the field's value and inside that value's sub-fields bind as fresh * defs at the same site. Use for "function-style" bindings. */ TVM_FFI_INLINE static AttachFieldFlag SEqHashDefRecursive() { return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashDefRecursive); } /*! * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive * * The field enters a non-recursive def region: only the immediate free * var at the field's value binds; free vars in its sub-fields are uses * that must already be bound by an outer def region. Use for "let-style" * bindings whose sub-fields reference outer-scope vars. */ TVM_FFI_INLINE static AttachFieldFlag SEqHashDefNonRecursive() { return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive); } /*! * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashIgnore */ TVM_FFI_INLINE static AttachFieldFlag SEqHashIgnore() { return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashIgnore); } /*! * \brief Apply the field flag to the field info * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { info->flags |= flag_; } private: int32_t flag_; }; /*! * \brief Trait that controls whether a field appears in repr output. * * By default, all fields appear in repr. Use `repr(false)` to exclude a field. */ class repr : public InfoTrait { public: /*! * \brief Constructor. * \param show Whether the field should appear in repr output. */ explicit repr(bool show) : show_(show) {} /*! * \brief Apply the repr flag to the field info. * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { if (!show_) { info->flags |= kTVMFFIFieldFlagBitMaskReprOff; } } private: bool show_; }; /*! * \brief InfoTrait to control whether a field participates in recursive comparison. * * Usage: `refl::compare(false)` marks a field to be excluded from * RecursiveEq/RecursiveLt/etc. */ class compare : public InfoTrait { public: /*! * \brief Constructor. * \param include Whether the field should participate in comparison. */ explicit compare(bool include) : include_(include) {} /*! * \brief Apply the compare flag to the field info. * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { if (!include_) { info->flags |= kTVMFFIFieldFlagBitMaskCompareOff; } } private: bool include_; }; /*! * \brief InfoTrait to control whether a field participates in recursive hashing. * * Usage: `refl::hash(false)` marks a field to be excluded from * RecursiveHash. */ class hash : public InfoTrait { public: /*! * \brief Constructor. * \param include Whether the field should participate in hashing. */ explicit hash(bool include) : include_(include) {} /*! * \brief Apply the hash flag to the field info. * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { if (!include_) { info->flags |= kTVMFFIFieldFlagBitMaskHashOff; } } private: bool include_; }; // Forward-declare `init` so that the deduction guide can reference it. template struct init; /*! * \brief InfoTrait to mark a field as keyword-only in the auto-generated init. * * Usage: `refl::kw_only(true)` marks a field so it can only be provided via * the KWARGS calling convention in the auto-generated ``__ffi_init__``. */ class kw_only : public InfoTrait { public: /*! * \brief Constructor. * \param is_kw_only Whether the field is keyword-only. */ explicit kw_only(bool is_kw_only) : kw_only_(is_kw_only) {} /*! * \brief Apply the kw_only flag to the field info. * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { if (kw_only_) { info->flags |= kTVMFFIFieldFlagBitMaskKwOnly; } } private: bool kw_only_; }; /*! * \brief Get the byte offset of a class member field. * * \tparam The original class. * \tparam T the field type. * * \param field_ptr A class member pointer * \returns The byteoffset */ template TVM_FFI_INLINE int64_t GetFieldByteOffsetToObject(T Class::* field_ptr) { int64_t field_offset_to_class = reinterpret_cast(&(static_cast(nullptr)->*field_ptr)); return field_offset_to_class - ::tvm::ffi::details::ObjectUnsafe::GetObjectOffsetToSubclass(); } /// \cond Doxygen_Suppress class ReflectionDefBase { protected: template static int FieldGetter(void* field, TVMFFIAny* result) { TVM_FFI_SAFE_CALL_BEGIN(); *result = ::tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(Any(*reinterpret_cast(field))); TVM_FFI_SAFE_CALL_END(); } template static int FieldSetter(void* field, const TVMFFIAny* value) { TVM_FFI_SAFE_CALL_BEGIN(); if constexpr (std::is_same_v) { *reinterpret_cast(field) = AnyView::CopyFromTVMFFIAny(*value); } else { *reinterpret_cast(field) = AnyView::CopyFromTVMFFIAny(*value).cast(); } TVM_FFI_SAFE_CALL_END(); } template static int ObjectCreatorDefault(TVMFFIObjectHandle* result) { TVM_FFI_SAFE_CALL_BEGIN(); ObjectPtr obj = make_object(); *result = ::tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(obj)); TVM_FFI_SAFE_CALL_END(); } template static int ObjectCreatorUnsafeInit(TVMFFIObjectHandle* result) { TVM_FFI_SAFE_CALL_BEGIN(); ObjectPtr obj = make_object(UnsafeInit{}); *result = ::tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(obj)); TVM_FFI_SAFE_CALL_END(); } template TVM_FFI_INLINE static void ApplyFieldInfoTrait(FieldInfoBuilder* info, const T& value) { if constexpr (std::is_base_of_v>) { value.Apply(info); } if constexpr (std::is_same_v, char*>) { info->doc = TVMFFIByteArray{value, std::char_traits::length(value)}; } } template TVM_FFI_INLINE static void ApplyMethodInfoTrait(MethodInfoBuilder* info, const T& value) { if constexpr (std::is_base_of_v>) { value.Apply(info); } if constexpr (std::is_same_v, char*>) { info->doc = TVMFFIByteArray{value, std::char_traits::length(value)}; } } template TVM_FFI_INLINE static void ApplyExtraInfoTrait(TVMFFITypeMetadata* info, const T& value) { if constexpr (std::is_same_v, char*>) { info->doc = TVMFFIByteArray{value, std::char_traits::length(value)}; } } template TVM_FFI_INLINE static Function GetMethod(std::string name, Func&& func) { return ffi::Function::FromTyped(WrapFunction(std::forward(func)), std::move(name)); } template TVM_FFI_INLINE static Func&& WrapFunction(Func&& func) { return std::forward(func); } template TVM_FFI_INLINE static auto WrapFunction(R (Class::*func)(Args...)) { static_assert(std::is_base_of_v || std::is_base_of_v, "Class must be derived from ObjectRef or Object"); if constexpr (std::is_base_of_v) { return [func](Class target, Args... params) -> R { // call method pointer return (target.*func)(std::forward(params)...); }; } if constexpr (std::is_base_of_v) { return [func](const Class* target, Args... params) -> R { // call method pointer return (const_cast(target)->*func)(std::forward(params)...); }; } } template TVM_FFI_INLINE static auto WrapFunction(R (Class::*func)(Args...) const) { static_assert(std::is_base_of_v || std::is_base_of_v, "Class must be derived from ObjectRef or Object"); if constexpr (std::is_base_of_v) { return [func](const Class& target, Args... params) -> R { // call method pointer return (target.*func)(std::forward(params)...); }; } if constexpr (std::is_base_of_v) { return [func](const Class* target, Args... params) -> R { // call method pointer return (target->*func)(std::forward(params)...); }; } } }; /// \endcond /*! * \brief GlobalDef helper to register a global function. * * \code{.cpp} * namespace refl = tvm::ffi::reflection; * refl::GlobalDef().def("my_ffi_extension.my_function", MyFunction); * \endcode */ class GlobalDef : public ReflectionDefBase { public: /*! * \brief Define a global function. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the function. * \param func The function to be registered. * \param extra The extra arguments that can be docstring or subclass of InfoTrait. * * \return The reflection definition. */ template GlobalDef& def(const char* name, Func&& func, Extra&&... extra) { using FuncInfo = ::tvm::ffi::details::FunctionInfo>; RegisterFunc(name, ffi::Function::FromTyped(std::forward(func), std::string(name)), FuncInfo::TypeSchema(), std::forward(extra)...); return *this; } /*! * \brief Define a global function in ffi::PackedArgs format. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the function. * \param func The function to be registered. * \param extra The extra arguments that can be docstring or subclass of InfoTrait. * * \return The reflection definition. */ template GlobalDef& def_packed(const char* name, Func func, Extra&&... extra) { RegisterFunc(name, ffi::Function::FromPacked(func), ::tvm::ffi::details::TypeSchemaImpl::v(), std::forward(extra)...); return *this; } /*! * \brief Expose a class method as a global function. * * An argument will be added to the first position if the function is not static. * * \tparam Class The class type. * \tparam Func The function type. * * \param name The name of the method. * \param func The function to be registered. * \param extra The extra arguments that can be docstring. * * \return The reflection definition. */ template GlobalDef& def_method(const char* name, Func&& func, Extra&&... extra) { using FuncInfo = ::tvm::ffi::details::FunctionInfo>; RegisterFunc(name, GetMethod(std::string(name), std::forward(func)), FuncInfo::TypeSchema(), std::forward(extra)...); return *this; } private: template // NOLINTNEXTLINE(performance-unnecessary-value-param) void RegisterFunc(const char* name, ffi::Function func, String type_schema, Extra&&... extra) { MethodInfoBuilder info; info.name = TVMFFIByteArray{name, std::char_traits::length(name)}; info.doc = TVMFFIByteArray{nullptr, 0}; info.flags = 0; info.method = AnyView(func).CopyToTVMFFIAny(); info.metadata_.emplace_back("type_schema", type_schema); ((ApplyMethodInfoTrait(&info, std::forward(extra)), ...)); std::string metadata_str = Metadata::ToJSON(info.metadata_); info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFIFunctionSetGlobalFromMethodInfo(&info, 0)); } }; /*! * \brief Helper class to register a constructor method for object types, * and (in its zero-argument specialization) an InfoTrait to control * whether a field participates in the auto-generated init. * * \tparam Args The argument types for the constructor. * * When used with one or more template arguments, ``init`` is passed * to ``ObjectDef::def()`` to register an ``__init__`` method that constructs * an object with the specified argument types: * * \code{.cpp} * refl::ObjectDef() * .def(refl::init()); * \endcode * * When used without template arguments (via CTAD), ``init(false)`` marks a * field to be excluded from the auto-generated ``__ffi_init__`` constructor, * or suppresses auto-init entirely when passed to ``ObjectDef``'s constructor: * * \code{.cpp} * // Per-field opt-out: * refl::ObjectDef() * .def_rw("hidden", &Foo::hidden, refl::init(false), refl::default_value(42)); * * // Class-level opt-out: * refl::ObjectDef(refl::init(false)) * .def_rw("x", &Bar::x); * \endcode * * \note The object type is automatically deduced from the ``ObjectDef`` context. */ template struct init { // Allow ObjectDef to access the execute function template friend class ObjectDef; template friend class OverloadObjectDef; /*! * \brief Constructor */ constexpr init() noexcept = default; private: /*! * \brief Execute the constructor * \tparam Class The class type. * \param args The arguments to be passed to the constructor. * \return The constructed object wrapped in an `ObjectRef`. */ template static inline ObjectRef execute(Args&&... args) { return ObjectRef(ffi::make_object(std::forward(args)...)); } }; /*! * \brief Specialization of ``init`` for zero template arguments. * * This specialization doubles as an ``InfoTrait`` that controls whether a field * participates in the auto-generated ``__ffi_init__``. * * - ``init(false)`` as a field trait: excludes the field from the auto-generated init. * - ``init(false)`` as an ``ObjectDef`` constructor argument: suppresses auto-init entirely. * - ``init<>()`` passed to ``ObjectDef::def()``: registers a zero-argument constructor. */ template <> struct init<> : public InfoTrait { // Allow ObjectDef to access the execute function and include_ flag. template friend class ObjectDef; template friend class OverloadObjectDef; /*! * \brief Default constructor (for zero-argument constructor registration). */ constexpr init() noexcept = default; /*! * \brief Constructor for field/class-level init control. * \param include Whether the field should be included in the auto-generated init. */ explicit init(bool include) : include_(include) {} /*! * \brief Apply the init flag to the field info. * \param info The field info. */ TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { if (!include_) { info->flags |= kTVMFFIFieldFlagBitMaskInitOff; } } private: bool include_ = true; /*! * \brief Execute the zero-argument constructor. * \tparam Class The class type. * \return The constructed object wrapped in an `ObjectRef`. */ template static inline ObjectRef execute() { return ObjectRef(ffi::make_object()); } }; /*! \brief CTAD deduction guide: ``init(false)`` deduces to ``init<>``. */ #if !defined(TVM_FFI_DOXYGEN_MODE) init(bool) -> init<>; #endif /*! * \brief Helper to register Object's reflection metadata. * \tparam Class The class type. * * \code{.cpp} * namespace refl = tvm::ffi::reflection; * refl::ObjectDef().def_ro("my_field", &MyClass::my_field); * \endcode */ template class ObjectDef : public ReflectionDefBase { public: /*! * \brief Constructor * \tparam ExtraArgs The extra arguments. * \param extra_args The extra arguments. */ template explicit ObjectDef(ExtraArgs&&... extra_args) : type_index_(Class::_GetOrAllocRuntimeTypeIndex()), type_key_(Class::_type_key) { (MaybeSuppressAutoInit(extra_args), ...); RegisterExtraInfo(std::forward(extra_args)...); } /*! \brief Non-copyable / non-movable. */ ObjectDef(const ObjectDef&) = delete; ObjectDef& operator=(const ObjectDef&) = delete; ObjectDef(ObjectDef&&) = delete; ObjectDef& operator=(ObjectDef&&) = delete; /*! * \brief Destructor, which also potentially registers `__ffi_new__`, `__ffi_init__`, * `__ffi_shallow_copy__`. */ ~ObjectDef() noexcept(false) { const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(type_index_); // Step 1. Register `__ffi_shallow_copy__` <== copy constructor (if it exists and is public) if constexpr (std::is_copy_constructible_v) { Function fn = Function::FromTyped( [](const Class* self) { return ObjectRef(ffi::make_object(*self)); }, std::string(type_key_) + "." + type_attr::kShallowCopy); TVMFFIByteArray attr = AsByteArray(type_attr::kShallowCopy); TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any)); } // Step 2. Register `__ffi_new__` <== info->metadata->creator // Also, `__ffi_init__` if no explicit init is defined. if (info->metadata != nullptr && info->metadata->creator != nullptr) { Function fn = Function::FromTyped( [creator = info->metadata->creator]() -> ObjectRef { TVMFFIObjectHandle handle; TVM_FFI_CHECK_SAFE_CALL(creator(&handle)); return ObjectRef(::tvm::ffi::details::ObjectUnsafe::ObjectPtrFromOwned( static_cast(handle))); }, std::string(type_key_) + "." + type_attr::kNew); TVMFFIByteArray attr = AsByteArray(type_attr::kNew); TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any)); // Step 3. Register `__ffi_init__` if no explicit init is defined. // Use Function::GetGlobal to look up the registration function, which lives in // dataclass.cc and may not be loaded yet for early (builtin) types. if (!has_explicit_init_) { if (std::optional opt = Function::GetGlobal("ffi._RegisterFFIInit")) { (*opt)(type_index_); } } } } /*! * \brief Define a readonly field. * * \tparam Class The class type. * \tparam T The field type. * \tparam Extra The extra arguments. * * \param name The name of the field. * \param field_ptr The pointer to the field. * \param extra The extra arguments that can be docstring or default value. * * \return The reflection definition. */ template TVM_FFI_INLINE ObjectDef& def_ro(const char* name, T BaseClass::* field_ptr, Extra&&... extra) { RegisterField(name, field_ptr, false, std::forward(extra)...); return *this; } /*! * \brief Define a read-write field. * * \tparam Class The class type. * \tparam T The field type. * \tparam Extra The extra arguments. * * \param name The name of the field. * \param field_ptr The pointer to the field. * \param extra The extra arguments that can be docstring or default value. * * \return The reflection definition. */ template TVM_FFI_INLINE ObjectDef& def_rw(const char* name, T BaseClass::* field_ptr, Extra&&... extra) { static_assert(Class::_type_mutable, "Only mutable classes are supported for writable fields"); RegisterField(name, field_ptr, true, std::forward(extra)...); return *this; } /*! * \brief Define a method. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the method. * \param func The function to be registered. * \param extra The extra arguments that can be docstring. * * \return The reflection definition. */ template TVM_FFI_INLINE ObjectDef& def(const char* name, Func&& func, Extra&&... extra) { RegisterMethod(name, false, std::forward(func), std::forward(extra)...); return *this; } /*! * \brief Define a static method. * * \tparam Func The function type. * \tparam Extra The extra arguments. * * \param name The name of the method. * \param func The function to be registered. * \param extra The extra arguments that can be docstring. * * \return The reflection definition. */ template TVM_FFI_INLINE ObjectDef& def_static(const char* name, Func&& func, Extra&&... extra) { RegisterMethod(name, true, std::forward(func), std::forward(extra)...); return *this; } /*! * \brief Register a constructor for this object type. * * This method registers a static `__init__` method that constructs an instance * of the object with the specified argument types. The constructor can be invoked * from Python or other FFI bindings. * * \tparam Args The argument types for the constructor. * \tparam Extra Additional arguments (e.g., docstring). * * \param init_func An instance of `init` specifying constructor signature. * \param extra Optional additional metadata such as docstring. * * \return Reference to this `ObjectDef` for method chaining. * * Example: * * \code{.cpp} * refl::ObjectDef() * .def(refl::init(), "Constructor docstring"); * \endcode */ template TVM_FFI_INLINE ObjectDef& def([[maybe_unused]] init init_func, Extra&&... extra) { has_explicit_init_ = true; // Register as TypeMethod (preserves type_schema metadata for Python stub generation). RegisterMethod(kInitMethodName, true, &init::template execute, std::forward(extra)...); // Also mirror into __ffi_init__ TypeAttrColumn for runtime dispatch. const TVMFFITypeInfo* tinfo = TVMFFIGetTypeInfo(type_index_); constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit); for (int32_t i = 0; i < tinfo->num_methods; ++i) { if (tinfo->methods[i].name.size == attr_name.size && std::strncmp(tinfo->methods[i].name.data, attr_name.data, attr_name.size) == 0) { TVM_FFI_CHECK_SAFE_CALL( TVMFFITypeRegisterAttr(type_index_, &attr_name, &tinfo->methods[i].method)); break; } } return *this; } private: template friend class OverloadObjectDef; /*! \brief If \p value is init(false), suppress auto-generated __ffi_init__. */ template void MaybeSuppressAutoInit(const T& value) { if constexpr (std::is_same_v, init<>>) { if (!value.include_) { has_explicit_init_ = true; } } } template void RegisterExtraInfo(ExtraArgs&&... extra_args) { TVMFFITypeMetadata info; info.total_size = sizeof(Class); info.structural_eq_hash_kind = Class::_type_s_eq_hash_kind; info.creator = nullptr; info.doc = TVMFFIByteArray{nullptr, 0}; if constexpr (std::is_default_constructible_v) { info.creator = ObjectCreatorDefault; } else if constexpr (std::is_constructible_v) { info.creator = ObjectCreatorUnsafeInit; } // apply extra info traits ((ApplyExtraInfoTrait(&info, std::forward(extra_args)), ...)); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMetadata(type_index_, &info)); } template void RegisterField(const char* name, T BaseClass::* field_ptr, bool writable, ExtraArgs&&... extra_args) { static_assert(std::is_base_of_v, "BaseClass must be a base class of Class"); FieldInfoBuilder info; info.name = TVMFFIByteArray{name, std::char_traits::length(name)}; info.field_static_type_index = TypeToFieldStaticTypeIndex::value; // store byte offset and setter, getter // so the same setter can be reused for all the same type info.offset = GetFieldByteOffsetToObject(field_ptr); info.size = sizeof(T); info.alignment = alignof(T); info.flags = 0; if (writable) { info.flags |= kTVMFFIFieldFlagBitMaskWritable; } info.getter = FieldGetter; info.setter = reinterpret_cast(FieldSetter); // initialize default value to nullptr info.default_value_or_factory = AnyView(nullptr).CopyToTVMFFIAny(); info.doc = TVMFFIByteArray{nullptr, 0}; info.metadata_.emplace_back("type_schema", ::tvm::ffi::details::TypeSchema::v()); // apply field info traits ((ApplyFieldInfoTrait(&info, std::forward(extra_args)), ...)); // call register std::string metadata_str = Metadata::ToJSON(info.metadata_); info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterField(type_index_, &info)); } // register a method template void RegisterMethod(const char* name, bool is_static, Func&& func, Extra&&... extra) { using FuncInfo = ::tvm::ffi::details::FunctionInfo>; MethodInfoBuilder info; info.name = TVMFFIByteArray{name, std::char_traits::length(name)}; info.doc = TVMFFIByteArray{nullptr, 0}; info.flags = 0; if (is_static) { info.flags |= kTVMFFIFieldFlagBitMaskIsStaticMethod; } // obtain the method function Function method = GetMethod(std::string(type_key_) + "." + name, std::forward(func)); info.method = AnyView(method).CopyToTVMFFIAny(); info.metadata_.emplace_back("type_schema", FuncInfo::TypeSchema()); // apply method info traits ((ApplyMethodInfoTrait(&info, std::forward(extra)), ...)); std::string metadata_str = Metadata::ToJSON(info.metadata_); info.metadata = TVMFFIByteArray{metadata_str.c_str(), metadata_str.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index_, &info)); } int32_t type_index_; const char* type_key_; bool has_explicit_init_{false}; static constexpr const char* kInitMethodName = type_attr::kInit; }; /*! * \brief Helper to register type attribute. * \tparam Class The class type. * \tparam ExtraArgs The extra arguments. * * \code{.cpp} * namespace refl = tvm::ffi::reflection; * refl::TypeAttrDef().def("func_attr", MyFunc); * \endcode */ template >> class TypeAttrDef : public ReflectionDefBase { public: /*! * \brief Constructor * \tparam ExtraArgs The extra arguments. * \param extra_args The extra arguments. */ template explicit TypeAttrDef(ExtraArgs&&... extra_args) : type_index_(Class::RuntimeTypeIndex()), type_key_(Class::_type_key) {} /*! * \brief Define a function-valued type attribute. * * \tparam Func The function type. * * \param name The name of the function. * \param func The function to be registered. * * \return The TypeAttrDef object. */ template TypeAttrDef& def(const char* name, Func&& func) { TVMFFIByteArray name_array = {name, std::char_traits::length(name)}; ffi::Function ffi_func = GetMethod(std::string(type_key_) + "." + name, std::forward(func)); TVMFFIAny value_any = AnyView(ffi_func).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any)); return *this; } /*! * \brief Define a constant-valued type attribute. * * \tparam T The type of the value. * * \param name The name of the attribute. * \param value The value of the attribute. * * \return The TypeAttrDef object. */ template TypeAttrDef& attr(const char* name, T value) { TVMFFIByteArray name_array = {name, std::char_traits::length(name)}; TVMFFIAny value_any = AnyView(value).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any)); return *this; } private: int32_t type_index_; const char* type_key_; }; /*! * \brief Ensure the type attribute column is presented in the system. * * \param name The name of the type attribute. */ inline void EnsureTypeAttrColumn(std::string_view name) { TVMFFIByteArray name_array = {name.data(), name.size()}; AnyView any_view(nullptr); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_array, reinterpret_cast(&any_view))); } /// \cond Doxygen_Suppress namespace details { /*! * \brief Implementation struct for overload_cast. * * Provides operator() overloads for each callable kind (free function, * non-const member, const member), in two flavors: full match where Args... * is the entire parameter list, and prefix match where Args... is a leading * prefix and the trailing parameters Rest... are deduced from the picked * overload's signature. */ template struct OverloadCastImpl { // The first triplet handles the case where Args... is the complete // parameter list of the picked overload. The second triplet handles // the prefix-match case where the picked overload has additional // trailing parameters Rest... beyond Args...; partial ordering picks // the first triplet when both apply, which lets the caller // disambiguate against shared-prefix overload sets by spelling the // full parameter list. template constexpr auto operator()(Ret (*fn)(Args...)) const noexcept { return fn; } template constexpr auto operator()(Ret (Cls::*pmf)(Args...), std::false_type = {}) const noexcept { return pmf; } template constexpr auto operator()(Ret (Cls::*pmf)(Args...) const, std::true_type) const noexcept { return pmf; } template constexpr auto operator()(Ret (*fn)(Args..., Rest...)) const noexcept { return fn; } template constexpr auto operator()(Ret (Cls::*pmf)(Args..., Rest...), std::false_type = {}) const noexcept { return pmf; } template constexpr auto operator()(Ret (Cls::*pmf)(Args..., Rest...) const, std::true_type) const noexcept { return pmf; } }; } // namespace details /// \endcond /*! * \brief Cast an overloaded callable to a specific overload, picked by * spelling out a parameter-type prefix that uniquely identifies it. * * `Args...` is matched against the leading parameters of each candidate * overload; the trailing parameter types (if any) are deduced from the * picked overload's signature. The returned value is a constexpr function * pointer (member or free) and can be used wherever a typed function * pointer is required, including as a non-type template argument. * * If the prefix matches multiple overloads (e.g. two overloads share the * same leading parameters), the call is ambiguous and the caller must * spell more parameters until exactly one overload matches. * * \note When picking a const-qualified member function, `refl::const_` must * be passed as the second argument even when it is the only overload * of its name. Without the tag the call does not compile. * * \note This helper can be more permissive than some `overload_cast` variants * in existing packages that require the full parameter list to be * spelled out: here a parameter-type prefix is accepted and the * trailing types are deduced from the picked overload. * * \code{.cpp} * class Pet { * public: * void Set(int); * void Set(const std::string&); * int Feed(const Cat*, int amount); * int Feed(const Dog*, int amount); * int Get(int); * int Get(int) const; * }; * * namespace refl = tvm::ffi::reflection; * * // Spell only the disambiguating first arg; the trailing `int amount` * // is deduced from the picked overload's signature. * auto p_feed_cat = refl::overload_cast(&Pet::Feed); * // decltype(p_feed_cat) == int (Pet::*)(const Cat*, int) * * // Spell the full parameter list when overloads share a prefix. * auto p_set_int = refl::overload_cast(&Pet::Set); * * // Const-qualified member — opt in via the const_ tag: * auto p_get_const = refl::overload_cast(&Pet::Get, refl::const_); * * // Use directly as a non-type template argument: * template struct UseAsTemplateArg { ... }; * using U = UseAsTemplateArg(&Pet::Feed)>; * \endcode */ template inline constexpr details::OverloadCastImpl overload_cast = {}; /// \cond Doxygen_Suppress // `const_`'s trailing underscore triggers RST hyperlink-reference syntax in // the exhale-generated per-variable page; suppress doc emission for it. // The symbol is still referenced (and rendered as inline literal) from the // overload_cast docstring above. inline constexpr auto const_ = std::true_type{}; /// \endcond } // namespace reflection } // namespace ffi } // namespace tvm #endif // TVM_FFI_REFLECTION_REGISTRY_H_ tvm-ffi-0.1.12/include/tvm/ffi/rvalue_ref.h000066400000000000000000000140711521067262500204540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/rvalue_ref.h * \brief Helper class to define rvalue reference type. */ #ifndef TVM_FFI_RVALUE_REF_H_ #define TVM_FFI_RVALUE_REF_H_ #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Helper class to define rvalue reference type. * * By default, FFI pass all values by lvalue reference. * * However, we do allow users to intentionally mark a function parameter * as RValueRef. In such cases, the caller can choose to pass parameter * wrapped by RValueRef to the function. In which case the parameter * can be directly moved by the callee. The caller can also choose to pass * a normal lvalue to the function, in such case a copy will be triggered. * * To keep FFI checking overhead minimal, we do not handle case when rvalue * is passed, but the callee did not declare the parameter as RValueRef. * * This design allows us to still leverage move semantics for parameters that * need copy on write scenarios (and requires an unique copy). * * \code{.cpp} * void Example() { * auto append = Function::FromTyped([](RValueRef> ref, int val) -> Array { * Array arr = *std::move(ref); * assert(arr.unique()); * arr.push_back(val); * return arr; * }); * Array a = Array({1, 2}); * // as we use rvalue ref to move a into append * // we keep a single copy of the Array without creating new copies during copy-on-write * a = append(RvalueRef(std::move(a)), 3); * assert(a.size() == 3); * } * \endcode */ template >> class RValueRef { public: /*! \brief the container type of the rvalue ref */ using ContainerType = typename TObjRef::ContainerType; /*! \brief only allow move constructor from rvalue of T */ explicit RValueRef(TObjRef&& data) : data_(details::ObjectUnsafe::ObjectPtrFromObjectRef(std::move(data))) {} /*! \brief return the data as rvalue */ TObjRef operator*() && { return TObjRef(std::move(data_)); } private: mutable ObjectPtr data_; template friend struct TypeTraits; }; template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public TypeTraitsBase { static constexpr bool storage_enabled = false; TVM_FFI_INLINE static void CopyToAnyView(const RValueRef& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIObjectRValueRef; result->zero_padding = 0; // store the address of the ObjectPtr, which allows us to move the value // and set the original ObjectPtr to nullptr result->v_ptr = &(src.data_); } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIObjectRValueRef) { ObjectPtr* rvalue_ref = reinterpret_cast*>(src->v_ptr); // object type does not match up, we need to try to convert the object // in this case we do not move the original rvalue ref since conversion creates a copy TVMFFIAny tmp_any; tmp_any.type_index = rvalue_ref->get()->type_index(); tmp_any.zero_padding = 0; tmp_any.v_obj = reinterpret_cast(rvalue_ref->get()); return "RValueRef<" + TypeTraits::GetMismatchTypeInfo(&tmp_any) + ">"; } else { return TypeTraits::GetMismatchTypeInfo(src); } } TVM_FFI_INLINE static std::optional> TryCastFromAnyView(const TVMFFIAny* src) { // first try rvalue conversion if (src->type_index == TypeIndex::kTVMFFIObjectRValueRef) { ObjectPtr* rvalue_ref = reinterpret_cast*>(src->v_ptr); TVMFFIAny tmp_any; tmp_any.type_index = rvalue_ref->get()->type_index(); tmp_any.zero_padding = 0; tmp_any.v_obj = reinterpret_cast(rvalue_ref->get()); // fast path, storage type matches, direct move the rvalue ref if (TypeTraits::CheckAnyStrict(&tmp_any)) { return RValueRef( details::ObjectUnsafe::ObjectRefFromObjectPtr(std::move(*rvalue_ref))); } if (std::optional opt = TypeTraits::TryCastFromAnyView(&tmp_any)) { // object type does not match up, we need to try to convert the object // in this case we do not move the original rvalue ref since conversion creates a copy return RValueRef(*std::move(opt)); } return std::nullopt; } // try lvalue conversion if (std::optional opt = TypeTraits::TryCastFromAnyView(src)) { return RValueRef(*std::move(opt)); } else { return std::nullopt; } } TVM_FFI_INLINE static std::string TypeStr() { return "RValueRef<" + TypeTraits::TypeStr() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { std::ostringstream oss; oss << R"({"type":")" << StaticTypeKey::kTVMFFIObjectRValueRef << R"(","args":[)"; oss << TypeTraits::TypeSchema(); oss << "]}"; return oss.str(); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_RVALUE_REF_H_ tvm-ffi-0.1.12/include/tvm/ffi/string.h000066400000000000000000001271701521067262500176350ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/string.h * \brief Runtime Bytes and String types. */ #ifndef TVM_FFI_STRING_H_ #define TVM_FFI_STRING_H_ #include #include #include #include #include #include #include #include #include #include #include #include #include #include // Note: We place string in tvm/ffi instead of tvm/ffi/container // because string itself needs special handling and is an inherent // core component for return string handling. // The following dependency relation holds // any -> string -> object /// \cond Doxygen_Suppress #ifdef _MSC_VER #define TVM_FFI_SNPRINTF _snprintf_s #pragma warning(push) #pragma warning(disable : 4244) #pragma warning(disable : 4127) #pragma warning(disable : 4702) #else #define TVM_FFI_SNPRINTF snprintf #endif /// \endcond namespace tvm { namespace ffi { namespace details { /*! * \brief Base class for bytes and string objects. */ class BytesObjBase : public Object, public TVMFFIByteArray {}; /*! * \brief An object representing bytes. * \note We use a separate object for bytes to follow Python convention * and indicate passing of raw bytes. * Bytes can be converted from/to string. */ class BytesObj : public BytesObjBase { public: static constexpr const uint32_t _type_index = TypeIndex::kTVMFFIBytes; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIBytes, BytesObj, Object); }; /*! \brief An object representing string. This is a POD type. */ class StringObj : public BytesObjBase { public: static constexpr const uint32_t _type_index = TypeIndex::kTVMFFIStr; static const constexpr bool _type_final = true; TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIStr, StringObj, Object); }; // String moved from std::string // without having to trigger a copy template class BytesObjStdImpl : public Base { public: explicit BytesObjStdImpl(std::string other) : data_{std::move(other)} { this->data = data_.data(); this->size = data_.size(); } private: std::string data_; }; /*! * \brief Helper cell class that can be used to back small string * \note Do not use directly, use String or Bytes instead */ class BytesBaseCell { public: BytesBaseCell() { // initialize to none data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } explicit BytesBaseCell(std::nullopt_t) { data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } BytesBaseCell(const BytesBaseCell& other) : data_(other.data_) { // NOLINT(*) if (data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::IncRefObjectHandle(data_.v_obj); } } BytesBaseCell(BytesBaseCell&& other) : data_(other.data_) { // NOLINT(*) other.data_.type_index = TypeIndex::kTVMFFINone; } BytesBaseCell& operator=(const BytesBaseCell& other) { BytesBaseCell(other).swap(*this); // NOLINT(*) return *this; } BytesBaseCell& operator=(BytesBaseCell&& other) noexcept { BytesBaseCell(std::move(other)).swap(*this); // NOLINT(*) return *this; } ~BytesBaseCell() { if (data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::DecRefObjectHandle(data_.v_obj); } } /*! * \brief Check if the cell is null * \return true if the cell is null, false otherwise */ bool operator==(std::nullopt_t) const { return data_.type_index == TypeIndex::kTVMFFINone; } /*! * \brief Check if the cell is not null * \return true if the cell is not null, false otherwise */ bool operator!=(std::nullopt_t) const { return data_.type_index != TypeIndex::kTVMFFINone; } /*! * \brief Swap this String with another string * \param other The other string */ void swap(BytesBaseCell& other) { // NOLINT(*) std::swap(data_, other.data_); } const char* data() const noexcept { if (data_.type_index < TypeIndex::kTVMFFIStaticObjectBegin) { return data_.v_bytes; } else { // NOLINTNEXTLINE(clang-analyzer-security.ArrayBound) return TVMFFIBytesGetByteArrayPtr(data_.v_obj)->data; } } size_t size() const noexcept { if (data_.type_index < TypeIndex::kTVMFFIStaticObjectBegin) { return data_.small_str_len; } else { // NOLINTNEXTLINE(clang-analyzer-security.ArrayBound) return TVMFFIBytesGetByteArrayPtr(data_.v_obj)->size; } } template void InitFromStd(std::string&& other, int32_t large_type_index) { // needs to be reset to none first for exception safety data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(&data_); ObjectPtr ptr = make_object>(std::move(other)); data_.v_obj = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(ptr)); data_.type_index = large_type_index; } /*! * \brief Create a new empty space for a string * \param size The size of the string * \param small_type_index The type index for the small string * \param large_type_index The type index for the large string * \note always reserve one byte for \0 compactibility * \return A pointer to the empty space */ template char* InitSpaceForSize(size_t size, int32_t small_type_index, int32_t large_type_index) { size_t kMaxSmallBytesLen = sizeof(int64_t) - 1; // first zero the content, this is important for exception safety data_.type_index = small_type_index; data_.zero_padding = 0; if (size <= kMaxSmallBytesLen) { // set up the size accordingly data_.small_str_len = static_cast(size); return data_.v_bytes; } else { // allocate from heap ObjectPtr ptr = make_inplace_array_object(size + 1); char* dest_data = reinterpret_cast(ptr.get()) + sizeof(LargeObj); ptr->data = dest_data; ptr->size = size; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(&data_); data_.v_obj = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(ptr)); // now reset the type index to str data_.type_index = large_type_index; return dest_data; } } void InitTypeIndex(int32_t type_index) { data_.type_index = type_index; } void MoveToAny(TVMFFIAny* result) { *result = data_; data_.type_index = TypeIndex::kTVMFFINone; data_.zero_padding = 0; data_.v_int64 = 0; } TVMFFIAny CopyToTVMFFIAny() const { return data_; } static BytesBaseCell CopyFromAnyView(const TVMFFIAny* src) { BytesBaseCell result(*src); if (result.data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { details::ObjectUnsafe::IncRefObjectHandle(result.data_.v_obj); } return result; } static BytesBaseCell MoveFromAny(TVMFFIAny* src) { BytesBaseCell result(*src); src->type_index = TypeIndex::kTVMFFINone; src->zero_padding = 0; src->v_int64 = 0; return result; } private: explicit BytesBaseCell(TVMFFIAny data) : data_(data) {} /*! \brief internal backing data */ TVMFFIAny data_; }; } // namespace details /*! * \brief Managed reference of byte array. */ class Bytes { public: /*! \brief default constructor */ Bytes() { data_.InitTypeIndex(TypeIndex::kTVMFFISmallBytes); } /*! * \brief constructor from size * * \param data The data pointer. * \param size The size of the char array. */ Bytes(const char* data, size_t size) { this->InitData(data, size); } /*! * \brief constructor from TVMFFIByteArray * * \param bytes a char array. */ Bytes(TVMFFIByteArray bytes) { // NOLINT(*) this->InitData(bytes.data, bytes.size); } /*! * \brief constructor from std::string * * \param other a char array. */ Bytes(const std::string& other) { // NOLINT(*) this->InitData(other.data(), other.size()); } /*! * \brief constructor from std::string * * \param other a char array. */ Bytes(std::string&& other) { // NOLINT(*) data_.InitFromStd(std::move(other), TypeIndex::kTVMFFIBytes); } /*! * \brief Swap this String with another string * \param other The other string */ void swap(Bytes& other) { // NOLINT(*) std::swap(data_, other.data_); } template Bytes& operator=(T&& other) { // copy-and-swap idiom Bytes(std::forward(other)).swap(*this); // NOLINT(*) return *this; } /*! * \brief Return the length of the string * * \return size_t string length */ size_t size() const { return data_.size(); } /*! * \brief Return the data pointer * * \return const char* data pointer */ const char* data() const { return data_.data(); } /*! * \brief Convert String to an std::string object * * \return std::string */ operator std::string() const { // NOLINT(google-explicit-constructor) return std::string{data(), size()}; } /*! * \brief Compare two char sequence * * \param lhs Pointers to the char array to compare * \param rhs Pointers to the char array to compare * \param lhs_count Length of the char array to compare * \param rhs_count Length of the char array to compare * \return int zero if both char sequences compare equal. negative if this * appear before other, positive otherwise. */ static int memncmp(const char* lhs, const char* rhs, size_t lhs_count, size_t rhs_count) { if (lhs == rhs && lhs_count == rhs_count) return 0; for (size_t i = 0; i < lhs_count && i < rhs_count; ++i) { if (lhs[i] < rhs[i]) return -1; if (lhs[i] > rhs[i]) return 1; } if (lhs_count < rhs_count) { return -1; } else if (lhs_count > rhs_count) { return 1; } else { return 0; } } /*! * \brief Compare two char sequence for equality * * \param lhs Pointers to the char array to compare * \param rhs Pointers to the char array to compare * \param lhs_count Length of the char array to compare * \param rhs_count Length of the char array to compare * * \return true if the two char sequences are equal, false otherwise. */ static bool memequal(const void* lhs, const void* rhs, size_t lhs_count, size_t rhs_count) { return lhs_count == rhs_count && (lhs == rhs || std::memcmp(lhs, rhs, lhs_count) == 0); } private: template friend struct TypeTraits; template friend class Optional; // internal backing cell details::BytesBaseCell data_; // create a new String from TVMFFIAny, must keep private explicit Bytes(details::BytesBaseCell data) : data_(std::move(data)) {} char* InitSpaceForSize(size_t size) { return data_.InitSpaceForSize(size, TypeIndex::kTVMFFISmallBytes, TypeIndex::kTVMFFIBytes); } void InitData(const char* data, size_t size) { char* dest_data = InitSpaceForSize(size); if (size > 0) { std::memcpy(dest_data, data, size); } // mainly to be compat with string dest_data[size] = '\0'; } }; /*! * \brief String container class. */ class String { public: /*! * \brief avoid misuse of nullptr */ String(std::nullptr_t) = delete; // NOLINT(*) /*! * \brief constructor */ String() { data_.InitTypeIndex(TypeIndex::kTVMFFISmallStr); } // constructors from Any /*! * \brief Copy constructor * \param other The other string */ String(const String& other) = default; // NOLINT(*) /*! * \brief Move constructor * \param other The other string */ String(String&& other) = default; // NOLINT(*) /*! * \brief Copy assignment operator * \param other The other string */ String& operator=(const String& other) = default; // NOLINT(*) /*! * \brief Move assignment operator * \param other The other string */ String& operator=(String&& other) = default; // NOLINT(*) /*! * \brief Swap this String with another string * \param other The other string */ void swap(String& other) noexcept { // NOLINT(*) std::swap(data_, other.data_); } /*! * \brief Copy assignment operator * \param other The other string */ String& operator=(const std::string& other) { String(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief Move assignment operator * \param other The other string */ String& operator=(std::string&& other) { String(std::move(other)).swap(*this); // NOLINT(*) return *this; } /*! * \brief Copy assignment operator * \param other The other string */ String& operator=(const char* other) { String(other).swap(*this); // NOLINT(*) return *this; } /*! * \brief constructor from raw string * * \param data The data pointer. * \param size The size of the char array. */ String(const char* data, size_t size) { this->InitData(data, size); } /*! * \brief constructor from raw string * * \param other a char array. * \note This constructor is marked as explicit to avoid implicit conversion * of nullptr value here to string, which then was used in comparison */ String(const char* other) { // NOLINT(*) this->InitData(other, std::char_traits::length(other)); } /*! * \brief Construct a new string object * \param other The std::string object to be copied */ String(const std::string& other) { // NOLINT(*) this->InitData(other.data(), other.size()); } /*! * \brief Construct a new string object * \param other The std::string object to be moved */ String(std::string&& other) { // NOLINT(*) // exception safety, first set to none so if exception is thrown // destructor works correctly data_.InitFromStd(std::move(other), TypeIndex::kTVMFFIStr); } /*! * \brief constructor from TVMFFIByteArray * * \param other a TVMFFIByteArray. */ explicit String(TVMFFIByteArray other) { this->InitData(other.data, other.size); } /*! * \brief Return the data pointer * * \return const char* data pointer */ const char* data() const noexcept { return data_.data(); } /*! * \brief Returns a pointer to the char array in the string. * * \return const char* */ const char* c_str() const noexcept { return data(); } /*! * \brief Return the length of the string * * \return size_t string length */ size_t size() const noexcept { return data_.size(); } /*! * \brief Compares this String object to other * * \param other The String to compare with. * * \return zero if both char sequences compare equal. negative if this appear * before other, positive otherwise. */ int compare(const String& other) const { return Bytes::memncmp(data(), other.data(), size(), other.size()); } /*! * \brief Compares this String object to other * * \param other The string to compare with. * * \return zero if both char sequences compare equal. negative if this appear * before other, positive otherwise. */ int compare(const std::string& other) const { return Bytes::memncmp(data(), other.data(), size(), other.size()); } /*! * \brief Compares this to other * * \param other The character array to compare with. * * \return zero if both char sequences compare equal. negative if this appear * before other, positive otherwise. */ int compare(const char* other) const { const char* this_data = data(); size_t this_size = size(); for (size_t i = 0; i < this_size; ++i) { // other is shorter than this if (other[i] == '\0') return 1; if (this_data[i] < other[i]) return -1; if (this_data[i] > other[i]) return 1; } // other equals this if (other[this_size] == '\0') return 0; // other longer than this return -1; } /*! * \brief Compares this to other * * \param other The TVMFFIByteArray to compare with. * * \return zero if both char sequences compare equal. negative if this appear * before other, positive otherwise. */ int compare(const TVMFFIByteArray& other) const { return Bytes::memncmp(data(), other.data, size(), other.size); } /*! * \brief Return the length of the string * * \return size_t string length */ size_t length() const { return size(); } /*! * \brief Retun if the string is empty * * \return true if empty, false otherwise. */ bool empty() const { return size() == 0; } /*! * \brief Read an element. * \param pos The position at which to read the character. * * \return The char at position */ char at(size_t pos) const { if (pos < size()) { return data()[pos]; } else { throw std::out_of_range("tvm::String index out of bounds"); } } /*! \brief Value returned by find() when no match is found */ static constexpr size_t npos = static_cast(-1); /*! * \brief Find the first occurrence of a substring * \param str The substring to search for * \param pos The position at which to start the search * \return The position of the first character of the first match, or npos if not found */ size_t find(const String& str, size_t pos = 0) const { return find(str.data(), pos, str.size()); } /*! * \brief Find the first occurrence of a substring * \param str The substring to search for * \param pos The position at which to start the search * \return The position of the first character of the first match, or npos if not found */ size_t find(const char* str, size_t pos = 0) const { return find(str, pos, std::strlen(str)); } /*! * \brief Find the first occurrence of a substring * \param str The substring to search for * \param pos The position at which to start the search * \param count The length of the substring * \return The position of the first character of the first match, or npos if not found */ size_t find(const char* str, size_t pos, size_t count) const { return std::string_view(data(), size()).find(std::string_view(str, count), pos); } /*! * \brief Returns a substring [pos, pos+count) * \param pos The position of the first character to include * \param count The length of the substring (default: until end of string) * \return A string containing the substring */ String substr(size_t pos = 0, size_t count = npos) const { if (pos > size()) { throw std::out_of_range("tvm::String substr index out of bounds"); } size_t rcount = std::min(count, size() - pos); return String(data() + pos, rcount); } /*! * \brief Check if the string starts with a prefix * \param prefix The prefix to check for * \return true if the string starts with prefix, false otherwise */ bool starts_with(const String& prefix) const { return starts_with(prefix.data(), prefix.size()); } /*! * \brief Check if the string starts with a prefix * \param prefix The prefix to check for * \return true if the string starts with prefix, false otherwise */ bool starts_with(std::string_view prefix) const { return starts_with(prefix.data(), prefix.size()); } /*! * \brief Check if the string starts with a prefix * \param prefix The prefix to check for * \return true if the string starts with prefix, false otherwise */ bool starts_with(const char* prefix) const { return starts_with(prefix, std::strlen(prefix)); } /*! * \brief Check if the string starts with a prefix * \param prefix The prefix to check for * \param count The length of the prefix * \return true if the string starts with prefix, false otherwise */ bool starts_with(const char* prefix, size_t count) const { if (count > size()) { return false; } return std::memcmp(data(), prefix, count) == 0; } /*! * \brief Check if the string ends with a suffix * \param suffix The suffix to check for * \return true if the string ends with suffix, false otherwise */ bool ends_with(const String& suffix) const { return ends_with(suffix.data(), suffix.size()); } /*! * \brief Check if the string ends with a suffix * \param suffix The suffix to check for * \return true if the string ends with suffix, false otherwise */ bool ends_with(std::string_view suffix) const { return ends_with(suffix.data(), suffix.size()); } /*! * \brief Check if the string ends with a suffix * \param suffix The suffix to check for * \return true if the string ends with suffix, false otherwise */ bool ends_with(const char* suffix) const { return ends_with(suffix, std::strlen(suffix)); } /*! * \brief Check if the string ends with a suffix * \param suffix The suffix to check for * \param count The length of the suffix * \return true if the string ends with suffix, false otherwise */ bool ends_with(const char* suffix, size_t count) const { if (count > size()) { return false; } return std::memcmp(data() + size() - count, suffix, count) == 0; } /*! * \brief Convert String to an std::string object * * \return std::string */ operator std::string() const { // NOLINT(google-explicit-constructor) return std::string{data(), size()}; } /*! * \brief Split the string by a delimiter character. * \param delim The delimiter character. * \return A vector of string_views pointing into this string's data. * \note The returned string_views are only valid while this String is alive. */ std::vector Split(char delim) const { std::vector ret; const char* start = data(); const char* end = start + size(); for (const char* p = start; p < end; ++p) { if (*p == delim) { ret.emplace_back(start, static_cast(p - start)); start = p + 1; } } ret.emplace_back(start, static_cast(end - start)); return ret; } private: template friend struct TypeTraits; template friend class Optional; // internal backing cell details::BytesBaseCell data_; // create a new String from TVMFFIAny, must keep private explicit String(details::BytesBaseCell data) : data_(std::move(data)) {} /*! * \brief Create a new empty space for a string * \param size The size of the string * \return A pointer to the empty space */ char* InitSpaceForSize(size_t size) { return data_.InitSpaceForSize(size, TypeIndex::kTVMFFISmallStr, TypeIndex::kTVMFFIStr); } void InitData(const char* data, size_t size) { char* dest_data = InitSpaceForSize(size); if (size > 0) { std::memcpy(dest_data, data, size); } dest_data[size] = '\0'; } /*! * \brief Concatenate two char sequences * * \param lhs Pointers to the lhs char array * \param lhs_size The size of the lhs char array * \param rhs Pointers to the rhs char array * \param rhs_size The size of the rhs char array * * \return The concatenated char sequence */ static String Concat(const char* lhs, size_t lhs_size, const char* rhs, size_t rhs_size) { String ret; // disable stringop-overflow and restrict warnings // gcc may produce false positive when we enable dest_data returned from small string path // Because compiler is not able to detect the condition that the path is only triggered via // size < kMaxSmallStrLen and can report it as a overflow case. #if (__GNUC__) && !(__clang__) #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wstringop-overflow" #pragma GCC diagnostic ignored "-Warray-bounds" #pragma GCC diagnostic ignored "-Wrestrict" #endif char* dest_data = ret.InitSpaceForSize(lhs_size + rhs_size); std::memcpy(dest_data, lhs, lhs_size); std::memcpy(dest_data + lhs_size, rhs, rhs_size); // NOLINTNEXTLINE(clang-analyzer-security.ArrayBound) dest_data[lhs_size + rhs_size] = '\0'; #if (__GNUC__) && !(__clang__) #pragma GCC diagnostic pop #endif return ret; } // Overload + operator friend String operator+(const String& lhs, const String& rhs); friend String operator+(const String& lhs, const std::string& rhs); friend String operator+(const std::string& lhs, const String& rhs); friend String operator+(const String& lhs, const char* rhs); friend String operator+(const char* lhs, const String& rhs); }; /*! * \brief Return a JSON-escaped version of the string (RFC 8259). * * Uses ``\\uXXXX`` for control characters, escapes ``\\/``, ``\\b``, ``\\f`` per the JSON spec. * Non-ASCII bytes are passed through as-is (valid UTF-8 is preserved). * * \param value The input string * \return The escaped string, quoted with double quotes */ inline String EscapeStringJSON(const String& value) { std::ostringstream oss; oss << '"'; const char* data = value.data(); const size_t size = value.size(); for (size_t i = 0; i < size; ++i) { switch (data[i]) { /// \cond Doxygen_Suppress #define TVM_FFI_ESCAPE_CHAR(pattern, val) \ case pattern: \ oss << (val); \ break TVM_FFI_ESCAPE_CHAR('\"', "\\\""); TVM_FFI_ESCAPE_CHAR('\\', "\\\\"); TVM_FFI_ESCAPE_CHAR('/', "\\/"); TVM_FFI_ESCAPE_CHAR('\b', "\\b"); TVM_FFI_ESCAPE_CHAR('\f', "\\f"); TVM_FFI_ESCAPE_CHAR('\n', "\\n"); TVM_FFI_ESCAPE_CHAR('\r', "\\r"); TVM_FFI_ESCAPE_CHAR('\t', "\\t"); #undef TVM_FFI_ESCAPE_CHAR /// \endcond default: { uint8_t u8_val = static_cast(data[i]); // this is a control character, print as \uXXXX if (u8_val < 0x20 || u8_val == 0x7f) { char buffer[8]; int size = TVM_FFI_SNPRINTF(buffer, sizeof(buffer), "\\u%04x", static_cast(data[i]) & 0xff); oss.write(buffer, size); } else { oss << data[i]; } break; } } } oss << '"'; return String(oss.str()); } /*! * \brief Escape a string for JSON output. * \deprecated Use EscapeStringJSON instead. * \param value The input string * \return The escaped string, quoted with double quotes */ [[deprecated("Use EscapeStringJSON instead")]] inline String EscapeString(const String& value) { return EscapeStringJSON(value); } /*! * \brief Return a Python-style escaped string representation. * * Handles ANSI escape sequences, UTF-8 multibyte characters, and standard * C escape sequences (\\n, \\t, \\r, \\\\, \\"). Uses \\xNN for control * characters and \\uXXXX / \\UXXXXXXXX for non-ASCII codepoints. * * \param value The input string to escape. * \return The escaped string, quoted with double quotes. */ inline String EscapedStringPy(const String& value) { const char* data = value.data(); const size_t length = value.size(); std::ostringstream oss; oss << '"'; for (size_t i = 0; i < length;) { unsigned char c = static_cast(data[i]); unsigned char d = (i + 1 < length) ? static_cast(data[i + 1]) : 0; // Detect ANSI escape sequences if (c == '\x1b' && d == '[') { size_t j = i + 2; while (j < length && (std::isdigit(static_cast(data[j])) || data[j] == ';')) { ++j; } if (j < length && (data[j] == 'm' || data[j] == 'K')) { oss << "\\x1b["; for (i += 2; i <= j; ++i) { oss << data[i]; } continue; } } // Handle ASCII C escape sequences switch (c) { case '\n': oss << "\\n"; ++i; continue; case '\t': oss << "\\t"; ++i; continue; case '\r': oss << "\\r"; ++i; continue; case '\\': oss << "\\\\"; ++i; continue; case '\"': oss << "\\\""; ++i; continue; default: break; } // Handle ASCII if ((c & 0x80) == 0) { if (c < 0x20 || c == 0x7f) { // Escape control characters as \xNN char buf[5]; TVM_FFI_SNPRINTF(buf, sizeof(buf), "\\x%02x", static_cast(c)); oss << buf; } else { oss << static_cast(c); } ++i; continue; } if ((c & 0xE0) == 0xC0 && i + 1 < length && (d & 0xC0) == 0x80) { int32_t codepoint = ((c & 0x1F) << 6) | (d & 0x3F); oss << "\\u" << std::hex << std::setw(4) << std::setfill('0') << codepoint; i += 2; } else if ((c & 0xF0) == 0xE0 && i + 2 < length) { unsigned char e = static_cast(data[i + 2]); if ((d & 0xC0) == 0x80 && (e & 0xC0) == 0x80) { int32_t codepoint = ((c & 0x0F) << 12) | ((d & 0x3F) << 6) | (e & 0x3F); oss << "\\u" << std::hex << std::setw(4) << std::setfill('0') << codepoint; i += 3; } else { oss << "\\x" << std::hex << std::setw(2) << std::setfill('0') << static_cast(c); ++i; } } else if ((c & 0xF8) == 0xF0 && i + 3 < length) { unsigned char e = static_cast(data[i + 2]); unsigned char f = static_cast(data[i + 3]); if ((d & 0xC0) == 0x80 && (e & 0xC0) == 0x80 && (f & 0xC0) == 0x80) { int32_t codepoint = ((c & 0x07) << 18) | ((d & 0x3F) << 12) | ((e & 0x3F) << 6) | (f & 0x3F); oss << "\\U" << std::hex << std::setw(8) << std::setfill('0') << codepoint; i += 4; } else { oss << "\\x" << std::hex << std::setw(2) << std::setfill('0') << static_cast(c); ++i; } } else { oss << "\\x" << std::hex << std::setw(2) << std::setfill('0') << static_cast(c); ++i; } oss.unsetf(std::ios::adjustfield | std::ios::basefield | std::ios::floatfield); oss.fill(' '); } oss << '"'; return String(oss.str()); } /*! \brief Convert TVMFFIByteArray to std::string_view */ TVM_FFI_INLINE std::string_view ToStringView(TVMFFIByteArray str) { return std::string_view(str.data, str.size); } /// \cond Doxygen_Suppress template <> inline constexpr bool use_default_type_traits_v = false; // specialize to enable implicit conversion from TVMFFIByteArray* template <> struct TypeTraits : public TypeTraitsBase { // bytes can be union type of small bytes and object, so keep it as any static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIAny; TVM_FFI_INLINE static void CopyToAnyView(const Bytes& src, TVMFFIAny* result) { *result = src.data_.CopyToTVMFFIAny(); } TVM_FFI_INLINE static void MoveToAny(Bytes src, TVMFFIAny* result) { src.data_.MoveToAny(result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFISmallBytes || src->type_index == TypeIndex::kTVMFFIBytes; } TVM_FFI_INLINE static Bytes CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { return Bytes(details::BytesBaseCell::CopyFromAnyView(src)); } TVM_FFI_INLINE static Bytes MoveFromAnyAfterCheck(TVMFFIAny* src) { return Bytes(details::BytesBaseCell::MoveFromAny(src)); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIByteArrayPtr) { return Bytes(*static_cast(src->v_ptr)); } if (src->type_index == TypeIndex::kTVMFFISmallBytes || src->type_index == TypeIndex::kTVMFFIBytes) { return Bytes(details::BytesBaseCell::CopyFromAnyView(src)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return "bytes"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIBytes) + R"("})"; } }; template <> inline constexpr bool use_default_type_traits_v = false; // specialize to enable implicit conversion from const char* template <> struct TypeTraits : public TypeTraitsBase { // string can be union type of small string and object, so keep it as any static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIAny; TVM_FFI_INLINE static void CopyToAnyView(const String& src, TVMFFIAny* result) { *result = src.data_.CopyToTVMFFIAny(); } TVM_FFI_INLINE static void MoveToAny(String src, TVMFFIAny* result) { src.data_.MoveToAny(result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFISmallStr || src->type_index == TypeIndex::kTVMFFIStr; } TVM_FFI_INLINE static String CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { return String(details::BytesBaseCell::CopyFromAnyView(src)); } TVM_FFI_INLINE static String MoveFromAnyAfterCheck(TVMFFIAny* src) { return String(details::BytesBaseCell::MoveFromAny(src)); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIRawStr) { return String(src->v_c_str); } if (src->type_index == TypeIndex::kTVMFFISmallStr || src->type_index == TypeIndex::kTVMFFIStr) { return String(details::BytesBaseCell::CopyFromAnyView(src)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return "str"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIStr) + R"("})"; } }; // const char*, requirement: not nullable, do not retain ownership template struct TypeTraits : public TypeTraitsBase { // NOTE: only enable implicit conversion into AnyView static constexpr bool storage_enabled = false; TVM_FFI_INLINE static void CopyToAnyView(const char src[N], TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIRawStr; result->zero_padding = 0; result->v_c_str = src; } TVM_FFI_INLINE static void MoveToAny(const char src[N], TVMFFIAny* result) { // when we need to move to any, convert to owned object first TypeTraits::MoveToAny(String(src), result); } }; template <> struct TypeTraits : public TypeTraitsBase { static constexpr bool storage_enabled = false; TVM_FFI_INLINE static void CopyToAnyView(const char* src, TVMFFIAny* result) { TVM_FFI_ICHECK_NOTNULL(src); result->type_index = TypeIndex::kTVMFFIRawStr; result->zero_padding = 0; result->v_c_str = src; } TVM_FFI_INLINE static void MoveToAny(const char* src, TVMFFIAny* result) { // when we need to move to any, convert to owned object first TypeTraits::MoveToAny(String(src), result); } // Do not allow const char* in a container, so we do not need CheckAnyStrict TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIRawStr) { return static_cast(src->v_c_str); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return "const char*"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"const char*"})"; } }; // TVMFFIByteArray, requirement: not nullable, do not retain ownership template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIByteArrayPtr; static constexpr bool storage_enabled = false; TVM_FFI_INLINE static void CopyToAnyView(TVMFFIByteArray* src, TVMFFIAny* result) { TVM_FFI_ICHECK_NOTNULL(src); result->type_index = TypeIndex::kTVMFFIByteArrayPtr; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_ptr = src; } TVM_FFI_INLINE static void MoveToAny(TVMFFIByteArray* src, TVMFFIAny* result) { TypeTraits::MoveToAny(Bytes(*src), result); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIByteArrayPtr) { return static_cast(src->v_ptr); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIByteArrayPtr; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIByteArrayPtr) + R"("})"; } }; template <> inline constexpr bool use_default_type_traits_v = false; template <> struct TypeTraits : public FallbackOnlyTraitsBase { TVM_FFI_INLINE static void CopyToAnyView(const std::string& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIRawStr; result->zero_padding = 0; result->v_c_str = src.c_str(); } TVM_FFI_INLINE static void MoveToAny(std::string src, TVMFFIAny* result) { // when we need to move to any, convert to owned object first TypeTraits::MoveToAny(String(std::move(src)), result); } TVM_FFI_INLINE static std::string TypeStr() { return "std::string"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"std::string"})"; } TVM_FFI_INLINE static std::string ConvertFallbackValue(const char* src) { return std::string(src); } TVM_FFI_INLINE static std::string ConvertFallbackValue(TVMFFIByteArray* src) { return std::string(src->data, src->size); } // NOLINTNEXTLINE(performance-unnecessary-value-param) TVM_FFI_INLINE static std::string ConvertFallbackValue(Bytes src) { return src.operator std::string(); } // NOLINTNEXTLINE(performance-unnecessary-value-param) TVM_FFI_INLINE static std::string ConvertFallbackValue(String src) { return src.operator std::string(); } }; inline String operator+(const String& lhs, const String& rhs) { size_t lhs_size = lhs.size(); size_t rhs_size = rhs.size(); return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size); } inline String operator+(const String& lhs, const std::string& rhs) { size_t lhs_size = lhs.size(); size_t rhs_size = rhs.size(); return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size); } inline String operator+(const std::string& lhs, const String& rhs) { size_t lhs_size = lhs.size(); size_t rhs_size = rhs.size(); return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size); } inline String operator+(const char* lhs, const String& rhs) { size_t lhs_size = std::strlen(lhs); size_t rhs_size = rhs.size(); return String::Concat(lhs, lhs_size, rhs.data(), rhs_size); } inline String operator+(const String& lhs, const char* rhs) { size_t lhs_size = lhs.size(); size_t rhs_size = std::strlen(rhs); return String::Concat(lhs.data(), lhs_size, rhs, rhs_size); } // Overload < operator inline bool operator<(std::nullptr_t, const String& rhs) = delete; inline bool operator<(const String& lhs, std::nullptr_t) = delete; inline bool operator<(const String& lhs, const std::string& rhs) { return lhs.compare(rhs) < 0; } inline bool operator<(const std::string& lhs, const String& rhs) { return rhs.compare(lhs) > 0; } inline bool operator<(const String& lhs, const String& rhs) { return lhs.compare(rhs) < 0; } inline bool operator<(const String& lhs, const char* rhs) { return lhs.compare(rhs) < 0; } inline bool operator<(const char* lhs, const String& rhs) { return rhs.compare(lhs) > 0; } // Overload > operator inline bool operator>(std::nullptr_t, const String& rhs) = delete; inline bool operator>(const String& lhs, std::nullptr_t) = delete; inline bool operator>(const String& lhs, const std::string& rhs) { return lhs.compare(rhs) > 0; } inline bool operator>(const std::string& lhs, const String& rhs) { return rhs.compare(lhs) < 0; } inline bool operator>(const String& lhs, const String& rhs) { return lhs.compare(rhs) > 0; } inline bool operator>(const String& lhs, const char* rhs) { return lhs.compare(rhs) > 0; } inline bool operator>(const char* lhs, const String& rhs) { return rhs.compare(lhs) < 0; } // Overload <= operator inline bool operator<=(std::nullptr_t, const String& rhs) = delete; inline bool operator<=(const String& lhs, std::nullptr_t) = delete; inline bool operator<=(const String& lhs, const std::string& rhs) { return lhs.compare(rhs) <= 0; } inline bool operator<=(const std::string& lhs, const String& rhs) { return rhs.compare(lhs) >= 0; } inline bool operator<=(const String& lhs, const String& rhs) { return lhs.compare(rhs) <= 0; } inline bool operator<=(const String& lhs, const char* rhs) { return lhs.compare(rhs) <= 0; } inline bool operator<=(const char* lhs, const String& rhs) { return rhs.compare(lhs) >= 0; } // Overload >= operator inline bool operator>=(std::nullptr_t, const String& rhs) = delete; inline bool operator>=(const String& lhs, std::nullptr_t) = delete; inline bool operator>=(const String& lhs, const std::string& rhs) { return lhs.compare(rhs) >= 0; } inline bool operator>=(const std::string& lhs, const String& rhs) { return rhs.compare(lhs) <= 0; } inline bool operator>=(const String& lhs, const String& rhs) { return lhs.compare(rhs) >= 0; } inline bool operator>=(const String& lhs, const char* rhs) { return lhs.compare(rhs) >= 0; } inline bool operator>=(const char* lhs, const String& rhs) { return rhs.compare(lhs) <= 0; } // delete Overload == operator for nullptr inline bool operator==(const String& lhs, std::nullptr_t) = delete; inline bool operator==(std::nullptr_t, const String& rhs) = delete; inline bool operator==(const String& lhs, const std::string& rhs) { return Bytes::memequal(lhs.data(), rhs.data(), lhs.size(), rhs.size()); } inline bool operator==(const std::string& lhs, const String& rhs) { return Bytes::memequal(lhs.data(), rhs.data(), lhs.size(), rhs.size()); } inline bool operator==(const String& lhs, const String& rhs) { return Bytes::memequal(lhs.data(), rhs.data(), lhs.size(), rhs.size()); } inline bool operator==(const String& lhs, const char* rhs) { return lhs.compare(rhs) == 0; } inline bool operator==(const char* lhs, const String& rhs) { return rhs.compare(lhs) == 0; } // Overload != operator inline bool operator!=(const String& lhs, std::nullptr_t) = delete; inline bool operator!=(std::nullptr_t, const String& rhs) = delete; inline bool operator!=(const String& lhs, const std::string& rhs) { return lhs.compare(rhs) != 0; } inline bool operator!=(const std::string& lhs, const String& rhs) { return rhs.compare(lhs) != 0; } inline bool operator!=(const String& lhs, const String& rhs) { return lhs.compare(rhs) != 0; } inline bool operator!=(const String& lhs, const char* rhs) { return lhs.compare(rhs) != 0; } inline bool operator!=(const char* lhs, const String& rhs) { return rhs.compare(lhs) != 0; } inline std::ostream& operator<<(std::ostream& out, const String& input) { out.write(input.data(), static_cast(input.size())); return out; } /// \endcond } // namespace ffi } // namespace tvm /// \cond Doxygen_Suppress namespace std { template <> struct hash<::tvm::ffi::Bytes> { std::size_t operator()(const ::tvm::ffi::Bytes& bytes) const { return std::hash()(std::string_view(bytes.data(), bytes.size())); } }; template <> struct hash<::tvm::ffi::String> { std::size_t operator()(const ::tvm::ffi::String& str) const { return std::hash()(std::string_view(str.data(), str.size())); } }; } // namespace std /// \endcond #endif // TVM_FFI_STRING_H_ tvm-ffi-0.1.12/include/tvm/ffi/tvm_ffi.h000066400000000000000000000041171521067262500177540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/tvm_ffi.h * \brief Umbrella header for the core TVM-FFI C++ APIs (excluding tvm/ffi/extra). * * \code{.cpp} * #include * \endcode */ #ifndef TVM_FFI_TVM_FFI_H_ #define TVM_FFI_TVM_FFI_H_ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #endif // TVM_FFI_TVM_FFI_H_ tvm-ffi-0.1.12/include/tvm/ffi/type_traits.h000066400000000000000000000760371521067262500207030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file tvm/ffi/object.h * \brief A managed object in the TVM FFI. */ #ifndef TVM_FFI_TYPE_TRAITS_H_ #define TVM_FFI_TYPE_TRAITS_H_ #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { class Any; /*! * \brief TypeTraits that specifies the conversion behavior from/to FFI Any. * * The function specifications of TypeTraits * * - CopyToAnyView: Convert a value T to AnyView * - MoveToAny: Move a value to Any * - CheckAnyStrict: Check if a Any stores a result of CopyToAnyView of current T. * - CopyFromAnyViewAfterCheck: Copy a value T from Any view after we pass CheckAnyStrict. * - MoveFromAnyAfterCheck: Move a value T from Any storage after we pass CheckAnyStrict. * - TryCastFromAnyView: Convert a AnyView to a T, we may apply type conversion. * - GetMismatchTypeInfo: Get the type key of a type when TryCastFromAnyView fails. * - TypeStr: Get the type key of a type * * It is possible that CheckAnyStrict is false but TryCastFromAnyView still works. * * For example, when Any x stores int, TypeTraits::CheckAnyStrict(x) will be false, * but TypeTraits::TryCastFromAnyView(x) will return a corresponding float value * via type conversion. * * CheckAnyStrict is mainly used in recursive container such as Array to * decide if a new Array needed to be created via recursive conversion, * or we can use the current container as is when converting to Array. * * A container array: Array satisfies the following invariant: * - `all(TypeTraits::CheckAnyStrict(x) for x in the array)`. */ template struct TypeTraits { /*! \brief Whether the type is enabled in FFI. */ static constexpr bool convert_enabled = false; /*! \brief Whether the type can appear as a storage type in Container */ static constexpr bool storage_enabled = false; }; /*! * \brief TypeTraits that removes const and reference keywords. * \tparam T the original type */ template using TypeTraitsNoCR = TypeTraits>>; template inline constexpr bool use_default_type_traits_v = true; struct TypeTraitsBase { static constexpr bool convert_enabled = true; static constexpr bool storage_enabled = true; static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIAny; // get mismatched type when result mismatches the trait. // this function is called after TryCastFromAnyView fails // to get more detailed type information in runtime // especially when the error involves nested container type TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* source) { return TypeIndexToTypeKey(source->type_index); } }; /*! * \brief Trait that maps a type to its field static type index * \tparam T the type * \return the field static type index */ template struct TypeToFieldStaticTypeIndex { /*! \brief The field static type index of the type */ static constexpr int32_t value = TypeIndex::kTVMFFIAny; }; template struct TypeToFieldStaticTypeIndex::convert_enabled>> { static constexpr int32_t value = TypeTraits::field_static_type_index; }; /*! * \brief Trait that maps a type to its runtime type index * \tparam T the type * \return the runtime type index */ template struct TypeToRuntimeTypeIndex { /*! * \brief Get the runtime type index of the type * \return the runtime type index */ static int32_t v() { return TypeToFieldStaticTypeIndex::value; } }; template struct TypeToRuntimeTypeIndex>> { static int32_t v() { return T::ContainerType::RuntimeTypeIndex(); } }; // None template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFINone; TVM_FFI_INLINE static void CopyToAnyView(const std::nullptr_t&, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFINone; result->zero_padding = 0; // invariant: the pointer field also equals nullptr // this will simplify same_as comparisons and hash result->v_int64 = 0; } TVM_FFI_INLINE static void MoveToAny(std::nullptr_t, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFINone; result->zero_padding = 0; // invariant: the pointer field also equals nullptr // this will simplify same_as comparisons and hash result->v_int64 = 0; } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFINone; } TVM_FFI_INLINE static std::nullptr_t CopyFromAnyViewAfterCheck(const TVMFFIAny*) { return nullptr; } TVM_FFI_INLINE static std::nullptr_t MoveFromAnyAfterCheck(TVMFFIAny*) { return nullptr; } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) { return nullptr; } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFINone; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFINone) + R"("})"; } }; /** * \brief A type that forbids implicit conversion from int to bool * * This type is used to prevent implicit conversion from int to bool. */ class StrictBool { public: /*! * \brief Constructor * \param value The value of the strict bool. */ StrictBool(bool value) : value_(value) {} // NOLINT(google-explicit-constructor) /*! *\brief Convert the strict bool to bool. * \return The value of the strict bool. */ operator bool() const { return value_; } // NOLINT(google-explicit-constructor) private: bool value_; }; template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIBool; TVM_FFI_INLINE static void CopyToAnyView(const StrictBool& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIBool; result->zero_padding = 0; result->v_int64 = static_cast(src); } TVM_FFI_INLINE static void MoveToAny(StrictBool src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIBool; } TVM_FFI_INLINE static StrictBool CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIBool); return static_cast(src->v_int64); } TVM_FFI_INLINE static StrictBool MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIBool) { return StrictBool(static_cast(src->v_int64)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIBool; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIBool) + R"("})"; } }; // Bool type, allow implicit casting from int template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIBool; TVM_FFI_INLINE static void CopyToAnyView(const bool& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIBool; result->zero_padding = 0; result->v_int64 = static_cast(src); } TVM_FFI_INLINE static void MoveToAny(bool src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIBool; } TVM_FFI_INLINE static bool CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIBool); return static_cast(src->v_int64); } TVM_FFI_INLINE static bool MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIInt || src->type_index == TypeIndex::kTVMFFIBool) { return static_cast(src->v_int64); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIBool; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIBool) + R"("})"; } }; // Integer POD values template struct TypeTraits>> : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIInt; TVM_FFI_INLINE static void CopyToAnyView(const Int& src, TVMFFIAny* result) { if constexpr (std::is_unsigned_v && sizeof(Int) >= sizeof(int64_t)) { if (src > static_cast(std::numeric_limits::max())) { TVM_FFI_THROW(OverflowError) << "Integer value " << src << " is too large to fit in int64_t. " << "Consider explicitly casting to int64_t first if this is intentional."; } } result->type_index = TypeIndex::kTVMFFIInt; result->zero_padding = 0; result->v_int64 = static_cast(src); // NOLINT(bugprone-signed-char-misuse) } TVM_FFI_INLINE static void MoveToAny(Int src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { // NOTE: CheckAnyStrict is always strict and should be consistent with MoveToAny return src->type_index == TypeIndex::kTVMFFIInt; } TVM_FFI_INLINE static Int CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIInt); return static_cast(src->v_int64); } TVM_FFI_INLINE static Int MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIInt || src->type_index == TypeIndex::kTVMFFIBool) { return Int(src->v_int64); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIInt; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIInt) + R"("})"; } }; /// \cond Doxygen_Suppress // trait to check if a type is an integeral enum // note that we need this trait so we can confirm underlying_type_t is an integral type // to avoid potential undefined behavior template > constexpr bool is_integeral_enum_v = false; template constexpr bool is_integeral_enum_v = std::is_integral_v>; /// \endcond // Enum Integer POD values template struct TypeTraits>> : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIInt; TVM_FFI_INLINE static void CopyToAnyView(const IntEnum& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIInt; result->zero_padding = 0; result->v_int64 = static_cast(src); } TVM_FFI_INLINE static void MoveToAny(IntEnum src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { // NOTE: CheckAnyStrict is always strict and should be consistent with MoveToAny return src->type_index == TypeIndex::kTVMFFIInt; } TVM_FFI_INLINE static IntEnum CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIInt); return static_cast(src->v_int64); } TVM_FFI_INLINE static IntEnum MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIInt || src->type_index == TypeIndex::kTVMFFIBool) { return static_cast(src->v_int64); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIInt; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIInt) + R"("})"; } }; // Float POD values template struct TypeTraits>> : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIFloat; TVM_FFI_INLINE static void CopyToAnyView(const Float& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIFloat; result->zero_padding = 0; result->v_float64 = static_cast(src); } TVM_FFI_INLINE static void MoveToAny(Float src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { // NOTE: CheckAnyStrict is always strict and should be consistent with MoveToAny return src->type_index == TypeIndex::kTVMFFIFloat; } TVM_FFI_INLINE static Float CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIFloat); return static_cast(src->v_float64); } TVM_FFI_INLINE static Float MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIFloat) { return Float(src->v_float64); } else if (src->type_index == TypeIndex::kTVMFFIInt || src->type_index == TypeIndex::kTVMFFIBool) { return Float(src->v_int64); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIFloat; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIFloat) + R"("})"; } }; // void* template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIOpaquePtr; TVM_FFI_INLINE static void CopyToAnyView(void* src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIOpaquePtr; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_ptr = src; } TVM_FFI_INLINE static void MoveToAny(void* src, TVMFFIAny* result) { CopyToAnyView(src, result); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { // NOTE: CheckAnyStrict is always strict and should be consistent with MoveToAny return src->type_index == TypeIndex::kTVMFFIOpaquePtr; } TVM_FFI_INLINE static void* CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIOpaquePtr); return src->v_ptr; } TVM_FFI_INLINE static void* MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIOpaquePtr) { return static_cast(src->v_ptr); } if (src->type_index == TypeIndex::kTVMFFINone) { return static_cast(nullptr); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIOpaquePtr; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIOpaquePtr) + R"("})"; } }; // Device template <> struct TypeTraits : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDevice; TVM_FFI_INLINE static void CopyToAnyView(const DLDevice& src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIDevice; result->zero_padding = 0; result->v_device = src; } TVM_FFI_INLINE static void MoveToAny(DLDevice src, TVMFFIAny* result) { result->type_index = TypeIndex::kTVMFFIDevice; result->zero_padding = 0; result->v_device = src; } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIDevice; } TVM_FFI_INLINE static DLDevice CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIDevice); return src->v_device; } TVM_FFI_INLINE static DLDevice MoveFromAnyAfterCheck(TVMFFIAny* src) { // POD type, we can just copy the value return CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIDevice) { return src->v_device; } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return StaticTypeKey::kTVMFFIDevice; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(StaticTypeKey::kTVMFFIDevice) + R"("})"; } }; // DLTensor*, requirement: not nullable, do not retain ownership template <> struct TypeTraits : public TypeTraitsBase { static constexpr bool storage_enabled = false; static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDLTensorPtr; TVM_FFI_INLINE static void CopyToAnyView(DLTensor* src, TVMFFIAny* result) { TVM_FFI_ICHECK_NOTNULL(src); result->type_index = TypeIndex::kTVMFFIDLTensorPtr; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_ptr = src; } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index == TypeIndex::kTVMFFIDLTensorPtr; } TVM_FFI_INLINE static DLTensor* CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { TVM_FFI_UNSAFE_ASSUME(src->type_index == TypeIndex::kTVMFFIDLTensorPtr); return static_cast(src->v_ptr); } TVM_FFI_INLINE static void MoveToAny(DLTensor*, TVMFFIAny*) { TVM_FFI_THROW(RuntimeError) << "DLTensor* cannot be held in Any as it does not retain ownership, use Tensor instead"; } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFIDLTensorPtr) { return static_cast(src->v_ptr); } else if (src->type_index == TypeIndex::kTVMFFITensor) { // Conversion from Tensor pointer to DLTensor // based on the assumption that Tensor always follows the TVMFFIObject header static_assert(sizeof(TVMFFIObject) == 24); return reinterpret_cast(reinterpret_cast(src->v_obj) + sizeof(TVMFFIObject)); } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return "DLTensor*"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"DLTensor*"})"; } }; // Traits for ObjectRef, None to ObjectRef will always fail. // use std::optional instead for nullable references. template struct ObjectRefTypeTraitsBase : public TypeTraitsBase { static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIObject; using ContainerType = typename TObjRef::ContainerType; TVM_FFI_INLINE static void CopyToAnyView(const TObjRef& src, TVMFFIAny* result) { if constexpr (TObjRef::_type_is_nullable) { if (!src.defined()) { TypeTraits::CopyToAnyView(nullptr, result); return; } } TVMFFIObject* obj_ptr = details::ObjectUnsafe::TVMFFIObjectPtrFromObjectRef(src); result->type_index = obj_ptr->type_index; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_obj = obj_ptr; } TVM_FFI_INLINE static void MoveToAny(TObjRef src, TVMFFIAny* result) { if constexpr (TObjRef::_type_is_nullable) { if (!src.defined()) { TypeTraits::CopyToAnyView(nullptr, result); return; } } TVMFFIObject* obj_ptr = details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(src)); result->type_index = obj_ptr->type_index; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_obj = obj_ptr; } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if constexpr (TObjRef::_type_is_nullable) { if (src->type_index == TypeIndex::kTVMFFINone) return true; } return (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && details::IsObjectInstance(src->type_index)); } TVM_FFI_INLINE static TObjRef CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { if constexpr (TObjRef::_type_is_nullable) { if (src->type_index == TypeIndex::kTVMFFINone) { return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); } } return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); } TVM_FFI_INLINE static TObjRef MoveFromAnyAfterCheck(TVMFFIAny* src) { if constexpr (TObjRef::_type_is_nullable) { if (src->type_index == TypeIndex::kTVMFFINone) { return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); } } // move out the object pointer ObjectPtr obj_ptr = details::ObjectUnsafe::ObjectPtrFromOwned(src->v_obj); // reset the src to nullptr TypeTraits::MoveToAny(nullptr, src); return details::ObjectUnsafe::ObjectRefFromObjectPtr(std::move(obj_ptr)); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if constexpr (TObjRef::_type_is_nullable) { if (src->type_index == TypeIndex::kTVMFFINone) { return details::ObjectUnsafe::ObjectRefFromObjectPtr(nullptr); } } if (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { if (details::IsObjectInstance(src->type_index)) { return details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); } } return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return ContainerType::_type_key; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(ContainerType::_type_key) + R"("})"; } }; template struct TypeTraits && use_default_type_traits_v>> : public ObjectRefTypeTraitsBase {}; /*! * \brief Helper class that convert to T only via the FallbackTypes * * The conversion will go through the FallbackTypes in the order * specified in the template parameter. * \tparam T The type of the target value. * \tparam FallbackTypes The type of the fallback value. * \note TypeTraits must be derived from this class and define * ConvertFallbackValue(FallbackType)->T for each FallbackType */ template struct FallbackOnlyTraitsBase : public TypeTraitsBase { // disable container for FallbackOnlyTraitsBase /// \cond Doxygen_Suppress static constexpr bool storage_enabled = false; TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { return TryFallbackTypes(src); } template TVM_FFI_INLINE static std::optional TryFallbackTypes(const TVMFFIAny* src) { static_assert(!std::is_same_v, "Using bool as FallbackType can cause bug because int will be detected as bool, " "use tvm::ffi::StrictBool instead"); if (auto opt_fallback = TypeTraits::TryCastFromAnyView(src)) { return TypeTraits::ConvertFallbackValue(*std::move(opt_fallback)); } if constexpr (sizeof...(Rest) > 0) { return TryFallbackTypes(src); } return std::nullopt; } /// \endcond }; /*! * \brief Helper class to define ObjectRef that can be auto-converted from a * fallback type, the Traits must be derived from it * and define a static methods named ConvertFallbackValue for each * FallbackType * * The conversion will go through the FallbackTypes in the order * specified in the template parameter. * \tparam ObjectRefType The type of the ObjectRef. * \tparam FallbackTypes The type of the fallback value. */ template struct ObjectRefWithFallbackTraitsBase : public ObjectRefTypeTraitsBase { /// \cond Doxygen_Suppress TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if (auto opt_obj = ObjectRefTypeTraitsBase::TryCastFromAnyView(src)) { return opt_obj; } // apply fallback types in TryCastFromAnyView return TryFallbackTypes(src); } template TVM_FFI_INLINE static std::optional TryFallbackTypes(const TVMFFIAny* src) { static_assert(!std::is_same_v, "Using bool as FallbackType can cause bug because int will be detected as bool, " "use tvm::ffi::StrictBool instead"); if (auto opt_fallback = TypeTraits::TryCastFromAnyView(src)) { return TypeTraits::ConvertFallbackValue(*std::move(opt_fallback)); } if constexpr (sizeof...(Rest) > 0) { return TryFallbackTypes(src); } return std::nullopt; } /// \endcond }; // Traits for weak pointer of object // NOTE: we require the weak pointer cast from template struct TypeTraits>> : public TypeTraitsBase { TVM_FFI_INLINE static void CopyToAnyView(TObject* src, TVMFFIAny* result) { TVMFFIObject* obj_ptr = details::ObjectUnsafe::GetHeader(src); result->type_index = obj_ptr->type_index; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_obj = obj_ptr; } TVM_FFI_INLINE static void MoveToAny(TObject* src, TVMFFIAny* result) { TVMFFIObject* obj_ptr = details::ObjectUnsafe::GetHeader(src); result->type_index = obj_ptr->type_index; result->zero_padding = 0; TVM_FFI_CLEAR_PTR_PADDING_IN_FFI_ANY(result); result->v_obj = obj_ptr; // needs to increase ref because original weak ptr do not own the code details::ObjectUnsafe::IncRefObjectHandle(result->v_obj); } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { return src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && details::IsObjectInstance(src->type_index); } TVM_FFI_INLINE static TObject* CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { if constexpr (!std::is_const_v) { static_assert(TObject::_type_mutable, "TObject must be mutable to enable cast from Any"); } return details::ObjectUnsafe::RawObjectPtrFromUnowned(src->v_obj); } TVM_FFI_INLINE static std::optional TryCastFromAnyView(const TVMFFIAny* src) { if constexpr (!std::is_const_v) { static_assert(TObject::_type_mutable, "TObject must be mutable to enable cast from Any"); } if (CheckAnyStrict(src)) return CopyFromAnyViewAfterCheck(src); return std::nullopt; } TVM_FFI_INLINE static std::string TypeStr() { return TObject::_type_key; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":")" + std::string(TObject::_type_key) + R"("})"; } }; template inline constexpr bool use_default_type_traits_v> = false; template struct TypeTraits> : public TypeTraitsBase { TVM_FFI_INLINE static void CopyToAnyView(const Optional& src, TVMFFIAny* result) { if (src.has_value()) { TypeTraits::CopyToAnyView(*src, result); } else { TypeTraits::CopyToAnyView(nullptr, result); } } TVM_FFI_INLINE static void MoveToAny(Optional src, TVMFFIAny* result) { if (src.has_value()) { TypeTraits::MoveToAny(*std::move(src), result); } else { TypeTraits::CopyToAnyView(nullptr, result); } } TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return true; return TypeTraits::CheckAnyStrict(src); } TVM_FFI_INLINE static Optional CopyFromAnyViewAfterCheck(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) { return Optional(std::nullopt); } return TypeTraits::CopyFromAnyViewAfterCheck(src); } TVM_FFI_INLINE static Optional MoveFromAnyAfterCheck(TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) { return Optional(std::nullopt); } return TypeTraits::MoveFromAnyAfterCheck(src); } TVM_FFI_INLINE static std::optional> TryCastFromAnyView(const TVMFFIAny* src) { if (src->type_index == TypeIndex::kTVMFFINone) return Optional(std::nullopt); if (std::optional opt = TypeTraits::TryCastFromAnyView(src)) { return Optional(*std::move(opt)); } else { // important to be explicit here // because nullopt can convert to std::optional(nullopt) which indicate success // return std::optional>(std::nullopt) to indicate failure return std::optional>(std::nullopt); } } TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) { return TypeTraits::GetMismatchTypeInfo(src); } TVM_FFI_INLINE static std::string TypeStr() { return "Optional<" + TypeTraits::TypeStr() + ">"; } TVM_FFI_INLINE static std::string TypeSchema() { return R"({"type":"Optional","args":[)" + details::TypeSchema::v() + "]}"; } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_TYPE_TRAITS_H_ tvm-ffi-0.1.12/licenses/000077500000000000000000000000001521067262500147565ustar00rootroot00000000000000tvm-ffi-0.1.12/licenses/LICENSE.dlpack.txt000066400000000000000000000261211521067262500200400ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "{}" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright 2017 by Contributors Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. tvm-ffi-0.1.12/licenses/LICENSE.libbacktrace.txt000066400000000000000000000026771521067262500212220ustar00rootroot00000000000000# Copyright (C) 2012-2016 Free Software Foundation, Inc. # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are # met: # (1) Redistributions of source code must retain the above copyright # notice, this list of conditions and the following disclaimer. # (2) Redistributions in binary form must reproduce the above copyright # notice, this list of conditions and the following disclaimer in # the documentation and/or other materials provided with the # distribution. # (3) The name of the author may not be used to # endorse or promote products derived from this software without # specific prior written permission. # THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR # IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE # DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, # INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) # HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, # STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING # IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. tvm-ffi-0.1.12/licenses/LICENSE.pytorch.txt000066400000000000000000000066101521067262500202730ustar00rootroot00000000000000From PyTorch: Copyright (c) 2016- Facebook, Inc (Adam Paszke) Copyright (c) 2014- Facebook, Inc (Soumith Chintala) Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) Copyright (c) 2011-2013 NYU (Clement Farabet) Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) Copyright (c) 2006 Idiap Research Institute (Samy Bengio) Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) From Caffe2: Copyright (c) 2016-present, Facebook Inc. All rights reserved. All contributions by Facebook: Copyright (c) 2016 Facebook Inc. All contributions by Google: Copyright (c) 2015 Google Inc. All rights reserved. All contributions by Yangqing Jia: Copyright (c) 2015 Yangqing Jia All rights reserved. All contributions by Kakao Brain: Copyright 2019-2020 Kakao Brain All contributions by Cruise LLC: Copyright (c) 2022 Cruise LLC. All rights reserved. All contributions by Tri Dao: Copyright (c) 2024 Tri Dao. All rights reserved. All contributions by Arm: Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates All contributions from Caffe: Copyright(c) 2013, 2014, 2015, the respective contributors All rights reserved. All other contributions: Copyright(c) 2015, 2016 the respective contributors All rights reserved. Caffe2 uses a copyright model similar to Caffe: each contributor holds copyright over their contributions to Caffe2. The project versioning records all such contribution and copyright details. If a contributor wants to further mark their specific copyright on a particular contribution, they should indicate their copyright solely in the commit message of the change when it is committed. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America and IDIAP Research Institute nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. tvm-ffi-0.1.12/licenses/NOTICE.pytorch.txt000066400000000000000000000561201521067262500201730ustar00rootroot00000000000000======================================================================= Software under third_party ======================================================================= Software libraries under third_party are provided as github submodule links, and their content is not part of the Caffe2 codebase. Their licences can be found under the respective software repositories. ======================================================================= Earlier BSD License ======================================================================= Early development of Caffe2 in 2015 and early 2016 is licensed under the BSD license. The license is attached below: All contributions by Facebook: Copyright (c) 2016 Facebook Inc. All contributions by Google: Copyright (c) 2015 Google Inc. All rights reserved. All contributions by Yangqing Jia: Copyright (c) 2015 Yangqing Jia All rights reserved. All contributions by Kakao Brain: Copyright 2019-2020 Kakao Brain All other contributions: Copyright(c) 2015, 2016 the respective contributors All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ======================================================================= Caffe's BSD License ======================================================================= Some parts of the caffe2 code is derived from the original Caffe code, which is created by Yangqing Jia and is now a BSD-licensed open-source project. The Caffe license is as follows: COPYRIGHT All contributions by the University of California: Copyright (c) 2014, The Regents of the University of California (Regents) All rights reserved. All other contributions: Copyright (c) 2014, the respective contributors All rights reserved. Caffe uses a shared copyright model: each contributor holds copyright over their contributions to Caffe. The project versioning records all such contribution and copyright details. If a contributor wants to further mark their specific copyright on a particular contribution, they should indicate their copyright solely in the commit message of the change when it is committed. LICENSE Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. CONTRIBUTION AGREEMENT By contributing to the BVLC/caffe repository through pull-request, comment, or otherwise, the contributor releases their content to the license and copyright terms herein. ======================================================================= Caffe2's Apache License ======================================================================= This repo contains Caffe2 code, which was previously licensed under Apache License Version 2.0: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. ======================================================================= Cephes's 3-Clause BSD License ======================================================================= Code derived from implementations in the Cephes Math Library should mention its derivation and reference the following license: 3-Clause BSD License for the Cephes Math Library Copyright (c) 2018, Steven Moshier All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. * Neither the name of the nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL Steven Moshier BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ======================================================================= SciPy's 3-Clause BSD License ======================================================================= Code derived from implementations in SciPy should mention its derivation and reference the following license: Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ======================================================================= Boost's 1.0 Software License ======================================================================= Code derived from implementations in Boost 1.0 should mention its derivation and reference the following license: Boost Software License - Version 1.0 - August 17th, 2003 Permission is hereby granted, free of charge, to any person or organization obtaining a copy of the software and accompanying documentation covered by this license (the "Software") to use, reproduce, display, distribute, execute, and transmit the Software, and to prepare derivative works of the Software, and to permit third-parties to whom the Software is furnished to do so, all subject to the following: The copyright notices in the Software and this entire statement, including the above license grant, this restriction and the following disclaimer, must be included in all copies of the Software, in whole or in part, and all derivative works of the Software, unless such copies or derivative works are solely in the form of machine-executable object code generated by a source language processor. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ======================================================================= PILLOW-SIMD Software License ======================================================================= Code derived from implementations in PILLOW-SIMD should mention its derivation and reference the following license: The Python Imaging Library (PIL) is Copyright © 1997-2011 by Secret Labs AB Copyright © 1995-2011 by Fredrik Lundh Pillow is the friendly PIL fork. It is Copyright © 2010-2022 by Alex Clark and contributors Like PIL, Pillow is licensed under the open source HPND License: By obtaining, using, and/or copying this software and/or its associated documentation, you agree that you have read, understood, and will comply with the following terms and conditions: Permission to use, copy, modify, and distribute this software and its associated documentation for any purpose and without fee is hereby granted, provided that the above copyright notice appears in all copies, and that both that copyright notice and this permission notice appear in supporting documentation, and that the name of Secret Labs AB or the author not be used in advertising or publicity pertaining to distribution of the software without specific, written prior permission. SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. tvm-ffi-0.1.12/pyproject.toml000066400000000000000000000170001521067262500160630ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [project] name = "apache-tvm-ffi" dynamic = ["version"] description = "tvm ffi" authors = [{ name = "TVM FFI team" }] readme = "README.md" license = { text = "Apache 2.0" } classifiers = [ "License :: OSI Approved :: Apache Software License", "Development Status :: 4 - Beta", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", ] keywords = ["machine learning", "inference"] requires-python = ">=3.8" dependencies = ["typing-extensions>=4.5"] [project.urls] Homepage = "https://github.com/apache/tvm-ffi" GitHub = "https://github.com/apache/tvm-ffi" [project.optional-dependencies] cpp = ["ninja"] [dependency-groups] torch = [ "torch; python_version < '3.14'", # pytorch does not yet ship with 3.14t "setuptools", # setuptools is needed by torch jit for best perf "ninja", "numpy", "ml_dtypes", ] test = [{ include-group = "torch" }, "pytest"] dev = [ { include-group = "test" }, "pre-commit", "ruff", "ty==0.0.15", "clang-format", "clang-tidy", "ipdb", "ipython", "cython>=3.0", "cmake", "scikit-build-core", "tomli", "setuptools-scm", ] docs = [ "autodocsumm", "breathe", "exhale", "linkify-it-py", "matplotlib", "myst-parser", "nbconvert", "nbsphinx", "nbstripout", # TODO: unpin when sphinx-toolbox fixes `from sphinx.ext.autodoc import logger` (removed in Sphinx 9) # Blocked by: sphinx-toolbox/utils.py imports logger from sphinx.ext.autodoc (broken as of v4.1.2) "sphinx<9", "sphinx-autobuild", "sphinx-book-theme", "sphinx-copybutton", "sphinx-design", "sphinx-reredirects", "sphinx-tabs", "sphinx-toolbox", "sphinx-autodoc-typehints", "sphinxcontrib-mermaid", "sphinxcontrib-napoleon", "sphinxcontrib_httpdomain", "setuptools", "setuptools-scm", "tomli", "urllib3", ] [project.scripts] tvm-ffi-config = "tvm_ffi.config:__main__" tvm-ffi-stubgen = "tvm_ffi.stub.cli:__main__" [build-system] requires = ["scikit-build-core>=0.10.0", "cython>=3.0", "setuptools-scm"] build-backend = "scikit_build_core.build" [tool.scikit-build] metadata.version.provider = "scikit_build_core.metadata.setuptools_scm" wheel.py-api = "cp312" minimum-version = "build-system.requires" ninja.version = ">=1.11" ninja.make-fallback = false # Build configuration build-dir = "build" build.verbose = true # Editable install configuration editable.rebuild = false editable.verbose = true # CMake configuration cmake.version = "CMakeLists.txt" cmake.build-type = "Release" cmake.args = [ "-DTVM_FFI_ATTACH_DEBUG_SYMBOLS=ON", "-DTVM_FFI_BUILD_TESTS=OFF", "-DTVM_FFI_BUILD_PYTHON_MODULE=ON", ] # Logging logging.level = "INFO" # Wheel configuration wheel.packages = ["python/tvm_ffi"] wheel.install-dir = "tvm_ffi" # Source distribution configuration sdist.include = [ # Build files "/CMakeLists.txt", "/pyproject.toml", "/cmake/**/*", # Source code "/src/**/*.cc", "/include/**/*", # python and cython "/python/tvm_ffi/**/*.py", "/python/tvm_ffi/**/*.pyx", "/python/tvm_ffi/**/*.pyi", "/python/tvm_ffi/py.typed", "/python/tvm_ffi/_version.py", # Third party files "/3rdparty/libbacktrace/**/*", "/3rdparty/dlpack/include/*/*", # Documentation and metadata "/docs/**/*", "/LICENSE", "/README.md", "/NOTICE", # Tests "/tests/**/*", ] sdist.exclude = [ "**/.git", "**/.github", "**/__pycache__", "**/*.pyc", "build", "dist", ] [tool.pytest.ini_options] testpaths = ["tests"] [tool.ruff] include = ["python/**/*.py", "tests/**/*.py"] line-length = 100 indent-width = 4 target-version = "py38" [tool.ruff.lint] select = [ "UP", # pyupgrade, https://docs.astral.sh/ruff/rules/#pyupgrade-up "PL", # pylint, https://docs.astral.sh/ruff/rules/#pylint-pl "I", # isort, https://docs.astral.sh/ruff/rules/#isort-i "RUF", # ruff, https://docs.astral.sh/ruff/rules/#ruff-specific-rules-ruf "NPY", # numpy, https://docs.astral.sh/ruff/rules/#numpy-specific-rules-npy "F", # pyflakes, https://docs.astral.sh/ruff/rules/#pyflakes-f "FA", # flake8-future-annotations, https://docs.astral.sh/ruff/rules/#flake8-future-annotations-fa "ANN", # flake8-annotations, https://docs.astral.sh/ruff/rules/#flake8-annotations-ann "PTH", # flake8-use-pathlib, https://docs.astral.sh/ruff/rules/#flake8-use-pathlib-pth "D", # pydocstyle, https://docs.astral.sh/ruff/rules/#pydocstyle-d ] ignore = [ "PLR2004", # pylint: magic-value-comparison "ANN401", # flake8-annotations: any-type "D105", # pydocstyle: undocumented-magic-method "D107", # pydocstyle: undocumented-public-init "D203", # pydocstyle: incorrect-blank-line-before-class "D213", # pydocstyle: multi-line-summary-second-line ] fixable = ["ALL"] unfixable = [] [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401"] # pyflakes: unused-import "tests/*" = [ "E741", # pycodestyle: ambiguous-variable-name "D100", # pydocstyle: undocumented-public-module "D101", # pydocstyle: undocumented-public-class "D103", # pydocstyle: undocumented-public-function "D107", # pydocstyle: undocumented-public-init "D205", # pydocstyle: missing-blank-line-after-summary ] [tool.ruff.lint.pylint] max-args = 10 [tool.ruff.format] quote-style = "double" indent-style = "space" skip-magic-trailing-comma = false line-ending = "auto" docstring-code-format = true docstring-code-line-length = 80 [tool.cibuildwheel] build-verbosity = 1 # only build up to cp312, cp312 # will be abi3 and can be used in future versions # ship 314t threaded nogil version build = ["cp38-*", "cp39-*", "cp310-*", "cp311-*", "cp312-*", "cp314t-*"] skip = ["*musllinux*"] # we only need to test on cp312 test-skip = ["cp38-*", "cp39-*", "cp310-*", "cp311-*"] # focus on testing abi3 wheel build-frontend = "build[uv]" test-command = "pytest {package}/tests/python -vvs" test-groups = ["test"] [tool.cibuildwheel.linux] archs = ["x86_64", "aarch64"] [tool.cibuildwheel.macos] archs = ["x86_64", "arm64"] environment = { MACOSX_DEPLOYMENT_TARGET = "10.14" } [tool.cibuildwheel.windows] archs = ["AMD64"] [tool.ty.environment] python-version = "3.9" extra-paths = ["python", "examples", "tests/python"] [tool.ty.src] include = ["python/tvm_ffi/**", "examples/**", "tests/python/**"] exclude = [".venv/**", "build/**", "dist/**"] [tool.ty.rules] unused-ignore-comment = "ignore" [tool.ty.analysis] allowed-unresolved-imports = [ "torch", "torch.*", "torch.utils.*", "my_ffi_extension", "my_ffi_extension.*", "_pytest.*", "cupy", "cupy.*", "paddle", "paddle.*", "triton", "triton.*", "cuda.bindings", "cuda.bindings.*", "torch_c_dlpack_ext", ] [tool.uv] exclude-newer = "14 days" [tool.uv.dependency-groups] docs = { requires-python = ">=3.13" } [tool.setuptools_scm] version_file = "python/tvm_ffi/_version.py" write_to = "python/tvm_ffi/_version.py" tvm-ffi-0.1.12/python/000077500000000000000000000000001521067262500144725ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/000077500000000000000000000000001521067262500161245ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/.gitignore000066400000000000000000000000271521067262500201130ustar00rootroot00000000000000core.cpp core.cpython* tvm-ffi-0.1.12/python/tvm_ffi/__init__.py000066400000000000000000000116601521067262500202410ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """TVM FFI Python package.""" # order matters here so we need to skip isort here # isort: skip_file import sys from typing import TYPE_CHECKING def _is_config_mode() -> bool: """Check user is invoking the config CLI entry.""" if sys.argv[0].endswith("tvm-ffi-config"): return True # sys.orig_argv is available only after python 3.10 if hasattr(sys, "orig_argv"): # Use orig_argv because Python strips the `tvm_ffi.config` from sys.argv when using -m. argv = sys.orig_argv for i, arg in enumerate(argv): if arg == "-m" and i + 1 < len(argv) and argv[i + 1] == "tvm_ffi.config": return True return False if TYPE_CHECKING or not _is_config_mode(): # Skip eager imports in CLI mode to avoid import # overhead in tvm-ffi-config command # HACK: try importing torch first, to avoid a potential # symbol conflict when both torch and tvm_ffi are imported. # This conflict can be reproduced in a very narrow scenario: # 1. GitHub action on Windows X64 # 2. Python 3.12 # 3. torch 2.9.0 try: import torch except ImportError: pass # Always load base libtvm_ffi before any other imports from . import libinfo LIB = libinfo.load_lib_ctypes("apache-tvm-ffi", "tvm_ffi", "RTLD_GLOBAL") # Enable package initialization from .registry import ( register_object, register_global_func, get_global_func, get_global_func_metadata, remove_global_func, init_ffi_api, ) from ._dtype import dtype from .core import Object, ObjectConvertible, Function, CAny, CContainerBase from ._convert import convert, convert_func from .error import register_error from ._tensor import Device, device, DLDeviceType from ._tensor import from_dlpack, Tensor, Shape from .container import Array, Dict, List, Map from .dataclasses.py_class import method from .module import Module, system_lib, load_module from .stream import StreamContext, get_raw_stream, use_raw_stream, use_torch_stream from .structural import ( StructuralKey, get_first_structural_mismatch, structural_equal, structural_hash, ) from . import serialization from . import access_path from . import dataclasses from . import structural from . import cpp # optional module to speedup dlpack conversion from . import _optional_torch_c_dlpack # import the dtype literals from ._dtype import ( bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float64, float32, float16, bfloat16, float8_e4m3fn, float8_e4m3fnuz, float8_e5m2, float8_e5m2fnuz, float8_e8m0fnu, float4_e2m1fnx2, ) elif sys.platform.startswith("win32"): # On Windows, load the library even in config CLI mode so the DLL search path # is set correctly (needed in some cases when test still loads cython extensions). from . import libinfo LIB = libinfo.load_lib_ctypes("apache-tvm-ffi", "tvm_ffi", "RTLD_GLOBAL") # normal version imports try: from ._version import __version__, __version_tuple__ except ImportError: __version__ = "0.0.0.dev0" __version_tuple__ = (0, 0, 0, "dev0", "7d34eb8ab.d20250913") __all__ = [ "LIB", "Array", "DLDeviceType", "Device", "Dict", "Function", "List", "Map", "Module", "Object", "ObjectConvertible", "Shape", "StreamContext", "StructuralKey", "Tensor", "__version__", "__version_tuple__", "access_path", "convert", "convert_func", "cpp", "dataclasses", "device", "dtype", "from_dlpack", "get_first_structural_mismatch", "get_global_func", "get_global_func_metadata", "get_raw_stream", "init_ffi_api", "load_module", "method", "register_error", "register_global_func", "register_object", "remove_global_func", "serialization", "structural", "structural_equal", "structural_hash", "system_lib", "use_raw_stream", "use_torch_stream", ] tvm-ffi-0.1.12/python/tvm_ffi/_convert.py000066400000000000000000000157131521067262500203240ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Conversion utilities to convert Python objects into TVM FFI values.""" from __future__ import annotations import ctypes from numbers import Number from types import ModuleType from typing import Any, Callable from . import _dtype, container, core try: import torch except ImportError: torch = None # ty: ignore[invalid-assignment] numpy: ModuleType | None = None try: import numpy except ImportError: pass def convert(value: Any) -> Any: # noqa: PLR0911,PLR0912 """Convert a Python object into TVM FFI values. This helper mirrors the automatic argument conversion that happens when calling FFI functions. It is primarily useful in tests or places where an explicit conversion is desired. Parameters ---------- value The Python object to be converted. Returns ------- ffi_obj The converted TVM FFI object. Examples -------- .. code-block:: python import tvm_ffi # Lists and tuples become tvm_ffi.Array a = tvm_ffi.convert([1, 2, 3]) assert isinstance(a, tvm_ffi.Array) # Dicts become tvm_ffi.Map m = tvm_ffi.convert({"a": 1, "b": 2}) assert isinstance(m, tvm_ffi.Map) # Strings and bytes become zero-copy FFI-aware types s = tvm_ffi.convert("hello") b = tvm_ffi.convert(b"bytes") assert isinstance(s, tvm_ffi.core.String) assert isinstance(b, tvm_ffi.core.Bytes) # Callables are wrapped as tvm_ffi.Function f = tvm_ffi.convert(lambda x: x + 1) assert isinstance(f, tvm_ffi.Function) # Array libraries that support DLPack export can be converted to Tensor import numpy as np x = tvm_ffi.convert(np.arange(4, dtype="int32")) assert isinstance(x, tvm_ffi.Tensor) Note ---- Function arguments to ffi function calls are automatically converted. So this function is mainly only used in internal or testing scenarios. """ if isinstance( value, (core.Object, core.PyNativeObject, bool, Number, ctypes.c_void_p, _dtype.dtype) ): return value elif isinstance(value, (tuple, list)): return container.Array(value) elif isinstance(value, dict): return container.Map(value) elif isinstance(value, str): return core.String(value) elif isinstance(value, (bytes, bytearray)): return core.Bytes(value) elif isinstance(value, core.ObjectConvertible): return value.asobject() elif callable(value): return core._convert_to_ffi_func(value) elif value is None: return None elif hasattr(value, "__dlpack__"): return core.from_dlpack(value) elif torch is not None and isinstance(value, torch.dtype): return core._convert_torch_dtype_to_ffi_dtype(value) elif numpy is not None and isinstance(value, numpy.dtype): return core._convert_numpy_dtype_to_ffi_dtype(value) elif hasattr(value, "__dlpack_data_type__"): cdtype = core._create_cdtype_from_tuple(core.DataType, *value.__dlpack_data_type__()) dtype = str.__new__(_dtype.dtype, str(cdtype)) dtype._tvm_ffi_dtype = cdtype return dtype elif isinstance(value, Exception): return core._convert_to_ffi_error(value) elif hasattr(value, "__tvm_ffi_object__"): return value.__tvm_ffi_object__() # keep rest protocol values as it is as they can be handled by ffi function elif hasattr(value, "__cuda_stream__"): return value elif hasattr(value, "__tvm_ffi_opaque_ptr__"): return value elif hasattr(value, "__dlpack_device__"): return value elif hasattr(value, "__tvm_ffi_int__"): return value elif hasattr(value, "__tvm_ffi_float__"): return value else: # in this case, it is an opaque python object return core._convert_to_opaque_object(value) def convert_func( pyfunc: Callable[..., Any], tensor_cls: type | None = None, ) -> Any: """Convert a Python callable to an FFI :py:class:`~tvm_ffi.Function`. This is the callable-specific sibling of :py:func:`tvm_ffi.convert`. It accepts one extra argument, ``tensor_cls``, that lets the caller specify how tensor arguments should be delivered to the Python callable when the resulting :py:class:`Function` is invoked from C++. :py:func:`tvm_ffi.convert` has no such knob — it always produces a :py:class:`Function` whose callback receives ``tvm_ffi.Tensor`` for tensor args. Parameters ---------- pyfunc : Callable The Python callable to wrap. tensor_cls : type, optional The class whose instances the callback should receive for tensor args. The class must expose a ``__dlpack_c_exchange_api__`` :py:class:`PyCapsule`; its capsule is threaded into the callback closure so tensor args are converted at the C level (via the DLPack exchange API) before the Python callback body runs — this is significantly faster than calling ``torch.from_dlpack(x)`` (or equivalent) inside the callback. Raises :py:class:`TypeError` if ``tensor_cls`` does not expose the attribute. When ``tensor_cls`` is ``None``, ``convert_func`` behaves like the callable branch of :py:func:`tvm_ffi.convert`. Returns ------- Function The wrapped FFI function. Examples -------- .. code-block:: python import torch import tvm_ffi # Without tensor_cls: same as tvm_ffi.convert(pyfunc) — the callback # receives tvm_ffi.Tensor for tensor args. f = tvm_ffi.convert_func(lambda x: x + 1) assert isinstance(f, tvm_ffi.Function) # With tensor_cls=torch.Tensor: the callback receives torch.Tensor # directly; the DLPack conversion happens in C before the body runs. def callback(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return a + b g = tvm_ffi.convert_func(callback, tensor_cls=torch.Tensor) See Also -------- :py:func:`tvm_ffi.convert` : Generic value-to-FFI conversion. Use this when you don't need to specify ``tensor_cls``. """ return core._convert_to_ffi_func(pyfunc, tensor_cls=tensor_cls) tvm-ffi-0.1.12/python/tvm_ffi/_dtype.py000066400000000000000000000242001521067262500177600ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Lightweight dtype wrapper for TVM FFI.""" # pylint: disable=invalid-name from __future__ import annotations from enum import IntEnum from typing import Any, ClassVar from . import core class DataTypeCode(IntEnum): """DLDataTypeCode code in DLTensor.""" INT = 0 UINT = 1 FLOAT = 2 HANDLE = 3 BFLOAT = 4 BOOL = 6 Float8E3M4 = 7 Float8E4M3 = 8 Float8E4M3B11FNUZ = 9 Float8E4M3FN = 10 Float8E4M3FNUZ = 11 Float8E5M2 = 12 Float8E5M2FNUZ = 13 Float8E8M0FNU = 14 Float6E2M3FN = 15 Float6E3M2FN = 16 Float4E2M1FN = 17 class dtype(str): """Lightweight data type in TVM FFI. It behaves like a Python :class:`str` but also carries an internal FFI representation. You can construct it from strings, NumPy/ML dtypes, or via :py:meth:`from_dlpack_data_type`. Parameters ---------- dtype_str The string representation of the dtype. Examples -------- .. code-block:: python import tvm_ffi # Create from string f32 = tvm_ffi.dtype("float32") assert f32.bits == 32 assert f32.itemsize == 4 # Adjust lanes to create vector types v4f32 = f32.with_lanes(4) assert v4f32 == "float32x4" # Round-trip from a DLPack (code, bits, lanes) triple f16 = tvm_ffi.dtype.from_dlpack_data_type((2, 16, 1)) assert f16 == "float16" Note ---- This class subclasses str so it can be directly passed into other array api's dtype arguments. """ __slots__ = ["_tvm_ffi_dtype"] _tvm_ffi_dtype: core.DataType _NUMPY_DTYPE_TO_STR: ClassVar[dict[Any, str]] = {} def __new__(cls, content: Any) -> dtype: content = str(content) val = str.__new__(cls, content) val._tvm_ffi_dtype = core.DataType(content) return val @staticmethod def from_dlpack_data_type(dltype_data_type: tuple[int, int, int]) -> dtype: """Create a dtype from a DLPack data type tuple. Parameters ---------- dltype_data_type The DLPack data type tuple ``(type_code, bits, lanes)``. Returns ------- dtype The created dtype. Examples -------- .. code-block:: python import tvm_ffi # Create float16 and int8 directly from DLPack triples f16 = tvm_ffi.dtype.from_dlpack_data_type((2, 16, 1)) i8 = tvm_ffi.dtype.from_dlpack_data_type((0, 8, 1)) assert f16 == "float16" assert i8 == "int8" See Also -------- :py:class:`tvm_ffi.dtype` User-facing dtype wrapper. :py:meth:`tvm_ffi.dtype.with_lanes` Create vector dtypes from a scalar base. """ cdtype = core._create_cdtype_from_tuple( core.DataType, dltype_data_type[0], dltype_data_type[1], dltype_data_type[2], ) val = str.__new__(dtype, str(cdtype)) val._tvm_ffi_dtype = cdtype return val def __repr__(self) -> str: return f"dtype('{self}')" def with_lanes(self, lanes: int) -> dtype: """Create a new dtype with the given number of lanes. Parameters ---------- lanes The number of lanes for the resulting vector type. Returns ------- dtype The new dtype with the given number of lanes. Examples -------- .. code-block:: python import tvm_ffi f32 = tvm_ffi.dtype("float32") v4f32 = f32.with_lanes(4) assert v4f32 == "float32x4" assert v4f32.bits == f32.bits and v4f32.lanes == 4 See Also -------- :py:meth:`tvm_ffi.dtype.from_dlpack_data_type` Construct from a DLPack ``(code, bits, lanes)`` triple. """ cdtype = core._create_cdtype_from_tuple( core.DataType, self._tvm_ffi_dtype.type_code, self._tvm_ffi_dtype.bits, lanes, ) val = str.__new__(dtype, str(cdtype)) val._tvm_ffi_dtype = cdtype return val @property def itemsize(self) -> int: """Size of one element in bytes. The size is computed as ``bits * lanes // 8``. When the number of lanes is greater than 1, the ``itemsize`` represents the byte size of the vector element. Examples -------- .. code-block:: python import tvm_ffi assert tvm_ffi.dtype("float32").itemsize == 4 assert tvm_ffi.dtype("float32").with_lanes(4).itemsize == 16 See Also -------- :py:attr:`tvm_ffi.dtype.bits` Bit width of the scalar base type. :py:attr:`tvm_ffi.dtype.lanes` Number of lanes for vector types. :py:meth:`tvm_ffi.dtype.with_lanes` Create a vector dtype from a scalar base. """ return self._tvm_ffi_dtype.itemsize @property def type_code(self) -> int: """Integer DLDataTypeCode of the scalar base type. Examples -------- .. code-block:: python import tvm_ffi f32 = tvm_ffi.dtype("float32") # The type code is an integer following DLPack conventions assert isinstance(f32.type_code, int) # Consistent with constructing from an explicit (code, bits, lanes) assert ( f32.type_code == tvm_ffi.dtype.from_dlpack_data_type((2, 32, 1)).type_code ) See Also -------- :py:meth:`tvm_ffi.dtype.from_dlpack_data_type` Construct a dtype from a DLPack ``(code, bits, lanes)`` triple. """ return self._tvm_ffi_dtype.type_code @property def bits(self) -> int: """Number of bits of the scalar base type. Examples -------- .. code-block:: python import tvm_ffi assert tvm_ffi.dtype("int8").bits == 8 v4f32 = tvm_ffi.dtype("float32").with_lanes(4) assert v4f32.bits == 32 # per-lane bit width See Also -------- :py:attr:`tvm_ffi.dtype.itemsize` Byte size accounting for lanes. :py:attr:`tvm_ffi.dtype.lanes` Number of lanes for vector types. """ return self._tvm_ffi_dtype.bits @property def lanes(self) -> int: """Number of lanes (for vector types). Returns ``1`` for scalar dtypes and the lane count for vector dtypes created via :py:meth:`tvm_ffi.dtype.with_lanes`. Examples -------- .. code-block:: python import tvm_ffi assert tvm_ffi.dtype("float32").lanes == 1 assert tvm_ffi.dtype("float32").with_lanes(4).lanes == 4 See Also -------- :py:meth:`tvm_ffi.dtype.with_lanes` Create a vector dtype from a scalar base. :py:attr:`tvm_ffi.dtype.itemsize` Byte size accounting for lanes. """ return self._tvm_ffi_dtype.lanes try: # this helps to make numpy as optional # although almost in all cases we want numpy import numpy as np dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.bool_)] = "bool" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int8)] = "int8" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int16)] = "int16" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int32)] = "int32" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int64)] = "int64" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint8)] = "uint8" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint16)] = "uint16" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint32)] = "uint32" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint64)] = "uint64" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float16)] = "float16" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float32)] = "float32" dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float64)] = "float64" if hasattr(np, "float_"): dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float64)] = "float64" except ImportError: pass try: import ml_dtypes dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.bfloat16)] = "bfloat16" dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e4m3fn)] = "float8_e4m3fn" dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e5m2)] = "float8_e5m2" if hasattr(ml_dtypes, "float4_e2m1fn"): # ml_dtypes >= 0.5.0 dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float4_e2m1fn)] = "float4_e2m1fn" except ImportError: pass core._set_class_dtype(dtype) # list of common dtype literals in machine learning systems apps # note that we can always cover more dtypes via explicit construction # from dlpack data type tuple # align with choice of numpy 2.0, which moved away from bool_ to bool bool = dtype("bool") int8 = dtype("int8") int16 = dtype("int16") int32 = dtype("int32") int64 = dtype("int64") uint8 = dtype("uint8") uint16 = dtype("uint16") uint32 = dtype("uint32") uint64 = dtype("uint64") float64 = dtype("float64") float32 = dtype("float32") float16 = dtype("float16") bfloat16 = dtype("bfloat16") # float8 dtypes float8_e4m3fn = dtype("float8_e4m3fn") float8_e4m3fnuz = dtype("float8_e4m3fnuz") float8_e5m2 = dtype("float8_e5m2") float8_e5m2fnuz = dtype("float8_e5m2fnuz") float8_e8m0fnu = dtype("float8_e8m0fnu") # float4x2 dtypes float4_e2m1fnx2 = dtype("float4_e2m1fnx2") # alias for torch naming pattern for f4x2 float4_e2m1fn_x2 = float4_e2m1fnx2 tvm-ffi-0.1.12/python/tvm_ffi/_dunder.py000066400000000000000000000413071521067262500201230ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Dunder method installation for ``@c_class`` and ``@py_class``.""" from __future__ import annotations import inspect from typing import TYPE_CHECKING, Any, Callable from . import core from .core import TypeInfo, object_repr if TYPE_CHECKING: from .core import Function def _make_init( type_cls: type, type_info: TypeInfo, ffi_init: Function, py_class_mode: bool = False, ) -> Callable[..., None]: """Build ``__init__`` that delegates to ``__ffi_init__``. Both ``@c_class`` and ``@py_class`` use the same constructor-call path (``self.__init_handle_by_constructor__(ffi_init, *args)``). The only difference is how ``super().__init__()`` from a subclass is handled: * **c_class** — raises ``TypeError`` (subclass must define its own init). * **py_class** — silently skips when the C++ handle is already set, so ``super().__init__()`` is a harmless no-op. Parameters ---------- type_cls The class to build an __init__ for. type_info The TypeInfo for *type_cls*. ffi_init The C++ ``__ffi_init__`` resolved at install time. py_class_mode If True, use a ``chandle`` guard instead of a ``TypeError`` guard. """ sig = _make_init_signature(type_info) kwargs_obj = core.KWARGS missing = core.MISSING has_post_init = hasattr(type_cls, "__post_init__") if py_class_mode: def __init__(self: Any, *args: Any, **kwargs: Any) -> None: if self.__chandle__() != 0: return ffi_args: list[Any] = list(args) if kwargs: ffi_args.append(kwargs_obj) for key, val in kwargs.items(): if val is not missing: ffi_args.append(key) ffi_args.append(val) self.__init_handle_by_constructor__(ffi_init, *ffi_args) if has_post_init: self.__post_init__() else: type_name = type_cls.__name__ def __init__(self: Any, *args: Any, **kwargs: Any) -> None: if type_info is not type(self).__tvm_ffi_type_info__: raise TypeError( f"Calling `super().__init__()` is not supported for @c_class " f"`{type_name}`. Use @py_class for inheritance with manual " f"__init__, or define `{type(self).__name__}` with init=True." ) ffi_args: list[Any] = list(args) if kwargs: ffi_args.append(kwargs_obj) for key, val in kwargs.items(): if val is not missing: ffi_args.append(key) ffi_args.append(val) self.__init_handle_by_constructor__(ffi_init, *ffi_args) if has_post_init: self.__post_init__() __init__.__signature__ = sig # ty: ignore[invalid-assignment] __init__.__qualname__ = f"{type_cls.__qualname__}.__init__" __init__.__module__ = type_cls.__module__ return __init__ def _make_init_signature(type_info: TypeInfo) -> inspect.Signature: """Build an ``inspect.Signature`` from reflection field metadata. Walks the parent chain (parent-first) to collect all ``init=True`` fields, reorders required-before-optional within each group, and returns a Signature for introspection. """ positional: list[tuple[str, bool]] = [] # (name, has_default) kw_only: list[tuple[str, bool]] = [] # (name, has_default) # Walk the parent chain to collect all fields (parent-first order). all_fields: list[Any] = [] ti: TypeInfo | None = type_info chain: list[TypeInfo] = [] while ti is not None: chain.append(ti) ti = ti.parent_type_info for ancestor_info in reversed(chain): all_fields.extend(ancestor_info.fields) for field in all_fields: if not field.c_init: continue if field.c_kw_only: kw_only.append((field.name, field.c_has_default)) else: positional.append((field.name, field.c_has_default)) # Required params must come before optional ones within each group. pos_required = [(n, d) for n, d in positional if not d] pos_default = [(n, d) for n, d in positional if d] kw_required = [(n, d) for n, d in kw_only if not d] kw_default = [(n, d) for n, d in kw_only if d] params: list[inspect.Parameter] = [] params.append(inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)) for name, _has_default in pos_required: params.append(inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD)) for name, _has_default in pos_default: params.append( inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD, default=core.MISSING) ) for name, _has_default in kw_required: params.append(inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY)) for name, _has_default in kw_default: params.append(inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY, default=core.MISSING)) return inspect.Signature(params) # --------------------------------------------------------------------------- # Copy / deepcopy / replace factories # --------------------------------------------------------------------------- def _make_copy(type_info: TypeInfo, shallow_copy_fn: Function | None) -> Callable[..., Any]: """Build ``__copy__`` with ``__ffi_shallow_copy__`` resolved at install time.""" if shallow_copy_fn is not None: def __copy__(self: Any) -> Any: return shallow_copy_fn(self) else: def __copy__(self: Any) -> Any: raise TypeError( f"Type `{type(self).__name__}` does not support copy. " f"The underlying C++ type is not copy-constructible." ) return __copy__ def _make_deepcopy(_type_info: TypeInfo) -> Callable[..., Any]: """Build ``__deepcopy__`` with ``DeepCopy`` resolved at install time.""" from . import _ffi_api # noqa: PLC0415 deep_copy = _ffi_api.DeepCopy def __deepcopy__(self: Any, memo: Any = None) -> Any: return deep_copy(self) return __deepcopy__ def _make_replace(_type_info: TypeInfo) -> Callable[..., Any]: """Build ``__replace__`` with ``copy.copy`` resolved at install time.""" import copy # noqa: PLC0415 copy_copy = copy.copy def __replace__(self: Any, **kwargs: Any) -> Any: obj = copy_copy(self) cls = type(obj) for key, value in kwargs.items(): getattr(cls, key).set(obj, value) return obj return __replace__ # --------------------------------------------------------------------------- # __init__ installation # --------------------------------------------------------------------------- def _set_match_args(cls: type, type_info: TypeInfo) -> None: """Set ``cls.__match_args__`` from reflected fields. Mirrors stdlib :func:`dataclasses.dataclass` semantics: the tuple contains the names of positional ``__init__`` fields (``init=True`` and ``kw_only=False``), walking the parent chain in parent-first order. If ``cls`` already defines ``__match_args__`` in its own ``__dict__``, it is left untouched. """ if "__match_args__" in cls.__dict__: return chain: list[TypeInfo] = [] ti: TypeInfo | None = type_info while ti is not None: chain.append(ti) ti = ti.parent_type_info names: list[str] = [] for ancestor in reversed(chain): for tf in ancestor.fields or (): df = tf.dataclass_field if df is None: continue if df.init and not df.kw_only: names.append(tf.name) setattr(cls, "__match_args__", tuple(names)) def _install_dataclass_dunders( # noqa: PLR0912, PLR0915 cls: type, *, init: bool, repr: bool, eq: bool, order: bool, unsafe_hash: bool, match_args: bool = True, py_class_mode: bool = False, ) -> None: """Install structural dunder methods on *cls*. Each dunder delegates to the corresponding C++ recursive structural operation (``RecursiveEq``, ``RecursiveHash``, ``RecursiveLt``, etc.). If the user already defined a dunder in the class body (i.e. it exists in ``cls.__dict__``), it is left untouched. Parameters ---------- cls The class to install dunders on. Must have been processed by :func:`register_object` first (so ``__tvm_ffi_type_info__`` exists). init If True, install ``__init__`` from C++ reflection metadata. repr If True, install :func:`~tvm_ffi.core.object_repr` as ``__repr__``. eq If True, install ``__eq__`` and ``__ne__`` using ``RecursiveEq``. Returns ``NotImplemented`` for unrelated types so Python can fall back to identity comparison. order If True, install ``__lt__``, ``__le__``, ``__gt__``, ``__ge__`` using ``RecursiveLt``/``Le``/``Gt``/``Ge``. Returns ``NotImplemented`` for unrelated types. unsafe_hash If True, install ``__hash__`` using ``RecursiveHash``. match_args If True (default), set ``cls.__match_args__`` to the tuple of positional ``__init__`` field names for use with ``match`` statements. Skipped when the class already defines ``__match_args__`` in its body. py_class_mode If True, use a ``chandle`` guard for ``__init__`` so that ``super().__init__()`` is a no-op, and wrap user-defined ``__init__`` to allocate via ``__ffi_init__`` before user code. """ from . import _ffi_api # noqa: PLC0415 type_info: TypeInfo = cls.__tvm_ffi_type_info__ # type: ignore[assignment] type_index: int = type_info.type_index # Look up __ffi_init__ from TypeMethod (preferred) or TypeAttrColumn (fallback). ffi_init: Function | None = None for method in type_info.methods: if method.name == "__ffi_init__": ffi_init = method.func break if ffi_init is None: ffi_init = core._lookup_type_attr(type_index, "__ffi_init__") ffi_new: Function | None = core._lookup_type_attr(type_index, "__ffi_new__") ffi_shallow_copy: Function | None = core._lookup_type_attr(type_index, "__ffi_shallow_copy__") # __init__ # ┌─────────────────────┬──────────────────────┬──────────────────────┐ # │ │ user __init__ │ no user __init__ │ # ├─────────────────────┼──────────────────────┼──────────────────────┤ # │ c_class, init=True │ keep as-is │ _make_init │ # │ c_class, init=False │ keep as-is │ TypeError guard │ # │ py_class, init=True │ wrap chandle guard │ _make_init(py_class) │ # │ py_class, init=False│ wrap chandle guard │ TypeError guard │ # └─────────────────────┴──────────────────────┴──────────────────────┘ if "__init__" not in cls.__dict__: if init and ffi_init is not None: cls.__init__ = _make_init( # type: ignore[attr-defined] cls, type_info, ffi_init=ffi_init, py_class_mode=py_class_mode, ) elif not init: # init=False, no user __init__: TypeError guard msg = ( f"`{cls.__name__}` cannot be constructed directly. " f"Define a custom __init__ or use a factory method." ) def __init___(self: Any, *args: Any, **kwargs: Any) -> None: raise TypeError(msg) __init___.__qualname__ = f"{cls.__qualname__}.__init__" __init___.__module__ = cls.__module__ cls.__init__ = __init___ # type: ignore[attr-defined] elif py_class_mode and ffi_new is not None: # User-defined __init__: wrap with chandle guard so the C++ object # is allocated (via __ffi_new__) before the user's code runs. # We use __ffi_new__ (zero-arg allocator) rather than __ffi_init__ # because the user's __init__ signature may not match the field # layout. super().__init__() from within the user's code becomes # a no-op because chandle is already set. import functools # noqa: PLC0415 user_init = cls.__dict__["__init__"] @functools.wraps(user_init) def __init__(self: Any, *args: Any, **kwargs: Any) -> None: if self.__chandle__() == 0: self.__init_handle_by_constructor__(ffi_new) user_init(self, *args, **kwargs) cls.__init__ = __init__ # type: ignore[attr-defined] # __repr__ if repr and "__repr__" not in cls.__dict__: cls.__repr__ = object_repr # type: ignore[attr-defined] def _is_comparable(self: Any, other: Any) -> bool: """Return True if *self* and *other* share a type hierarchy.""" return isinstance(other, type(self)) or isinstance(self, type(other)) # Semantic families (__eq__/__ne__, __lt__/__le__/__gt__/__ge__) are # treated as a unit: if the user defines any member, the whole family # is skipped so generated and user-defined methods don't disagree. if eq and not ({"__eq__", "__ne__"} & set(cls.__dict__)): recursive_eq = _ffi_api.RecursiveEq def __eq__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return recursive_eq(self, other) def __ne__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return not recursive_eq(self, other) cls.__eq__ = __eq__ # type: ignore[attr-defined] cls.__ne__ = __ne__ # type: ignore[attr-defined] if order and not ({"__lt__", "__le__", "__gt__", "__ge__"} & set(cls.__dict__)): recursive_lt = _ffi_api.RecursiveLt recursive_le = _ffi_api.RecursiveLe recursive_gt = _ffi_api.RecursiveGt recursive_ge = _ffi_api.RecursiveGe def __lt__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return recursive_lt(self, other) def __le__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return recursive_le(self, other) def __gt__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return recursive_gt(self, other) def __ge__(self: Any, other: Any) -> bool: if not _is_comparable(self, other): return NotImplemented return recursive_ge(self, other) cls.__lt__ = __lt__ # type: ignore[attr-defined] cls.__le__ = __le__ # type: ignore[attr-defined] cls.__gt__ = __gt__ # type: ignore[attr-defined] cls.__ge__ = __ge__ # type: ignore[attr-defined] if unsafe_hash and "__hash__" not in cls.__dict__: recursive_hash = _ffi_api.RecursiveHash def __hash__(self: Any) -> int: return recursive_hash(self) cls.__hash__ = __hash__ # type: ignore[attr-defined] # __copy__ / __deepcopy__ / __replace__ if "__copy__" not in cls.__dict__: cls.__copy__ = _make_copy(type_info, ffi_shallow_copy) # type: ignore[attr-defined] if "__deepcopy__" not in cls.__dict__: cls.__deepcopy__ = _make_deepcopy(type_info) # type: ignore[attr-defined] if "__replace__" not in cls.__dict__: cls.__replace__ = _make_replace(type_info) # type: ignore[attr-defined] if match_args: _set_match_args(cls, type_info) tvm-ffi-0.1.12/python/tvm_ffi/_ffi_api.py000066400000000000000000000202541521067262500202350ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """FFI API.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from .registry import init_ffi_api as _FFI_INIT_FUNC from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence from ctypes import c_void_p from tvm_ffi import Device, Module, Object, StructuralKey as _StructuralKey from tvm_ffi.access_path import AccessPath from typing import Any, Callable # isort: on # fmt: on # tvm-ffi-stubgen(end) # tvm-ffi-stubgen(begin): global/ffi@.registry # fmt: off _FFI_INIT_FUNC("ffi", __name__) if TYPE_CHECKING: def Array(*args: Any) -> Any: ... def ArrayContains(_0: Sequence[Any], _1: Any, /) -> bool: ... def ArrayGetItem(_0: Sequence[Any], _1: int, /) -> Any: ... def ArraySize(_0: Sequence[Any], /) -> int: ... def Bytes(_0: bytes, /) -> bytes: ... def ContainerFindFirstNonCPUDevice(_0: Any, /) -> Device: ... def DeepCopy(_0: Any, /) -> Any: ... def Dict(*args: Any) -> Any: ... def DictClear(_0: MutableMapping[Any, Any], /) -> None: ... def DictCount(_0: MutableMapping[Any, Any], _1: Any, /) -> int: ... def DictErase(_0: MutableMapping[Any, Any], _1: Any, /) -> None: ... def DictForwardIterFunctor(_0: MutableMapping[Any, Any], /) -> Callable[..., Any]: ... def DictGetItem(_0: MutableMapping[Any, Any], _1: Any, /) -> Any: ... def DictGetItemOrMissing(_0: MutableMapping[Any, Any], _1: Any, /) -> Any: ... def DictSetItem(_0: MutableMapping[Any, Any], _1: Any, _2: Any, /) -> None: ... def DictSize(_0: MutableMapping[Any, Any], /) -> int: ... def FromJSONGraph(_0: Any, /) -> Any: ... def FromJSONGraphString(_0: str, /) -> Any: ... def FunctionFromExternC(_0: c_void_p, _1: c_void_p, _2: c_void_p, /) -> Callable[..., Any]: ... def FunctionListGlobalNamesFunctor() -> Callable[..., Any]: ... def FunctionRemoveGlobal(_0: str, /) -> bool: ... def GetFirstStructuralMismatch(_0: Any, _1: Any, _2: bool, _3: bool, /) -> tuple[AccessPath, AccessPath] | None: ... def GetGlobalFuncMetadata(_0: str, /) -> str: ... def GetInvalidObject() -> Object: ... def GetKwargsObject() -> Object: ... def GetRegisteredTypeKeys() -> Sequence[str]: ... def List(*args: Any) -> Any: ... def ListAppend(_0: MutableSequence[Any], _1: Any, /) -> None: ... def ListClear(_0: MutableSequence[Any], /) -> None: ... def ListContains(_0: MutableSequence[Any], _1: Any, /) -> bool: ... def ListErase(_0: MutableSequence[Any], _1: int, /) -> None: ... def ListEraseRange(_0: MutableSequence[Any], _1: int, _2: int, /) -> None: ... def ListGetItem(_0: MutableSequence[Any], _1: int, /) -> Any: ... def ListInsert(_0: MutableSequence[Any], _1: int, _2: Any, /) -> None: ... def ListPop(_0: MutableSequence[Any], _1: int, /) -> Any: ... def ListReplaceSlice(_0: MutableSequence[Any], _1: int, _2: int, _3: MutableSequence[Any], /) -> None: ... def ListReverse(_0: MutableSequence[Any], /) -> None: ... def ListSetItem(_0: MutableSequence[Any], _1: int, _2: Any, /) -> None: ... def ListSize(_0: MutableSequence[Any], /) -> int: ... def MakeFieldGetter(_0: int, /) -> int: ... def MakeFieldSetter(_0: int, _1: int, _2: int, /) -> Callable[..., Any]: ... def MakeObjectFromPackedArgs(*args: Any) -> Any: ... def Map(*args: Any) -> Any: ... def MapCount(_0: Mapping[Any, Any], _1: Any, /) -> int: ... def MapForwardIterFunctor(_0: Mapping[Any, Any], /) -> Callable[..., Any]: ... def MapGetItem(_0: Mapping[Any, Any], _1: Any, /) -> Any: ... def MapGetItemOrMissing(_0: Mapping[Any, Any], _1: Any, /) -> Any: ... def MapSize(_0: Mapping[Any, Any], /) -> int: ... def ModuleClearImports(_0: Module, /) -> None: ... def ModuleGetFunction(_0: Module, _1: str, _2: bool, /) -> Callable[..., Any] | None: ... def ModuleGetFunctionDoc(_0: Module, _1: str, _2: bool, /) -> str | None: ... def ModuleGetFunctionMetadata(_0: Module, _1: str, _2: bool, /) -> str | None: ... def ModuleGetKind(_0: Module, /) -> str: ... def ModuleGetPropertyMask(_0: Module, /) -> int: ... def ModuleGetWriteFormats(_0: Module, /) -> Sequence[str]: ... def ModuleGlobalsAdd(_0: Module, /) -> None: ... def ModuleGlobalsRemove(_0: Module, /) -> None: ... def ModuleImplementsFunction(_0: Module, _1: str, _2: bool, /) -> bool: ... def ModuleImportModule(_0: Module, _1: Module, /) -> None: ... def ModuleInspectSource(_0: Module, _1: str, /) -> str: ... def ModuleLoadFromFile(_0: str, /) -> Module: ... def ModuleWriteToFile(_0: Module, _1: str, _2: str, /) -> None: ... def RecursiveEq(_0: Any, _1: Any, /) -> bool: ... def RecursiveGe(_0: Any, _1: Any, /) -> bool: ... def RecursiveGt(_0: Any, _1: Any, /) -> bool: ... def RecursiveHash(_0: Any, /) -> int: ... def RecursiveLe(_0: Any, _1: Any, /) -> bool: ... def RecursiveLt(_0: Any, _1: Any, /) -> bool: ... def ReprPrint(_0: Any, /) -> str: ... def Shape(*args: Any) -> Any: ... def String(_0: str, /) -> str: ... def StructuralEqual(_0: Any, _1: Any, _2: bool, _3: bool, /) -> bool: ... def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ... def StructuralKey(_0: Any, /) -> _StructuralKey: ... def StructuralKeyEqual(_0: Any, _1: Any, /) -> bool: ... def SystemLib(*args: Any) -> Any: ... def ToJSONGraph(_0: Any, _1: Any, /) -> Any: ... def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ... def _PyClassRegisterTypeAttrColumns(_0: int, _1: int, /) -> None: ... def _RegisterFFIInit(_0: int, /) -> None: ... # fmt: on # tvm-ffi-stubgen(end) __all__ = [ # tvm-ffi-stubgen(begin): __all__ "Array", "ArrayContains", "ArrayGetItem", "ArraySize", "Bytes", "ContainerFindFirstNonCPUDevice", "DeepCopy", "Dict", "DictClear", "DictCount", "DictErase", "DictForwardIterFunctor", "DictGetItem", "DictGetItemOrMissing", "DictSetItem", "DictSize", "FromJSONGraph", "FromJSONGraphString", "FunctionFromExternC", "FunctionListGlobalNamesFunctor", "FunctionRemoveGlobal", "GetFirstStructuralMismatch", "GetGlobalFuncMetadata", "GetInvalidObject", "GetKwargsObject", "GetRegisteredTypeKeys", "List", "ListAppend", "ListClear", "ListContains", "ListErase", "ListEraseRange", "ListGetItem", "ListInsert", "ListPop", "ListReplaceSlice", "ListReverse", "ListSetItem", "ListSize", "MakeFieldGetter", "MakeFieldSetter", "MakeObjectFromPackedArgs", "Map", "MapCount", "MapForwardIterFunctor", "MapGetItem", "MapGetItemOrMissing", "MapSize", "ModuleClearImports", "ModuleGetFunction", "ModuleGetFunctionDoc", "ModuleGetFunctionMetadata", "ModuleGetKind", "ModuleGetPropertyMask", "ModuleGetWriteFormats", "ModuleGlobalsAdd", "ModuleGlobalsRemove", "ModuleImplementsFunction", "ModuleImportModule", "ModuleInspectSource", "ModuleLoadFromFile", "ModuleWriteToFile", "RecursiveEq", "RecursiveGe", "RecursiveGt", "RecursiveHash", "RecursiveLe", "RecursiveLt", "ReprPrint", "Shape", "String", "StructuralEqual", "StructuralHash", "StructuralKey", "StructuralKeyEqual", "SystemLib", "ToJSONGraph", "ToJSONGraphString", "_PyClassRegisterTypeAttrColumns", "_RegisterFFIInit", # tvm-ffi-stubgen(end) ] tvm-ffi-0.1.12/python/tvm_ffi/_optional_torch_c_dlpack.py000066400000000000000000000170211521067262500235020ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Optional module to support faster DLPack conversion. This is an optional module to support faster DLPack conversion for torch. Some of the changes are merged but not yet released, so it is used as a stop gap to support faster DLPack conversion. This file contains source code from PyTorch: License: licenses/LICENSE.pytorch.txt This module only serves as temp measure and will likely be phased away and deleted after changes landed and released in pytorch. This module will load slowly at first time due to JITing, subsequent calls will be much faster. """ from __future__ import annotations import ctypes import logging import os import subprocess import sys import warnings from pathlib import Path from typing import Any logger = logging.getLogger(__name__) def _torch_extension_device(torch_module: Any) -> str: """Return the torch backend name used in the optional extension library name.""" if torch_module.cuda.is_available(): if getattr(torch_module.version, "cuda", None) is not None: return "cuda" if getattr(torch_module.version, "hip", None) is not None: return "rocm" return "cuda" return "cpu" def _create_dlpack_exchange_api_capsule(ptr_as_int: int) -> Any: """Create a PyCapsule wrapping the DLPack exchange API pointer. Parameters ---------- ptr_as_int : int The pointer to the DLPack exchange API as an integer. Returns ------- capsule : PyCapsule A PyCapsule object wrapping the pointer with name "dlpack_exchange_api". """ capsule_name = b"dlpack_exchange_api" pythonapi = ctypes.pythonapi pythonapi.PyCapsule_New.restype = ctypes.py_object pythonapi.PyCapsule_New.argtypes = [ ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ] capsule = pythonapi.PyCapsule_New(ctypes.c_void_p(ptr_as_int), capsule_name, None) return capsule def _check_and_update_dlpack_c_exchange_api(tensor_cls: object) -> bool: """Check if the DLPack exchange API is available and update the __dlpack_c_exchange_api__ attribute.""" if hasattr(tensor_cls, "__dlpack_c_exchange_api__"): return True # legacy path compactibility handling if hasattr(tensor_cls, "__c_dlpack_exchange_api__"): c_dlpack_attribute = tensor_cls.__c_dlpack_exchange_api__ if isinstance(c_dlpack_attribute, int): setattr( tensor_cls, "__dlpack_c_exchange_api__", _create_dlpack_exchange_api_capsule(c_dlpack_attribute), ) else: setattr(tensor_cls, "__dlpack_c_exchange_api__", c_dlpack_attribute) return True return False def load_torch_c_dlpack_extension() -> Any: # noqa: PLR0912, PLR0915 try: import torch # noqa: PLC0415 import torch.version # noqa: PLC0415 if _check_and_update_dlpack_c_exchange_api(torch.Tensor): # skip loading the extension if the __dlpack_c_exchange_api__ # attribute is already set so we don't have to do it in # newer version of PyTorch return None except ImportError: return None """Load the torch c dlpack extension.""" try: import torch_c_dlpack_ext # noqa: PLC0415, F401 if _check_and_update_dlpack_c_exchange_api(torch.Tensor): return None except ImportError: pass except AttributeError: # When torch_c_dlpack_ext and torch have different ABI # `ctypes.CDLL` will raise an `AttributeError`. # Keep trying JIT pass try: # check whether a JIT shared library is built in cache cache_dir = Path(os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")).expanduser() addon_output_dir = cache_dir major, minor = torch.__version__.split(".")[:2] device = _torch_extension_device(torch) suffix = ".dll" if sys.platform.startswith("win") else ".so" libname = f"libtorch_c_dlpack_addon_torch{major}{minor}-{device}{suffix}" lib_path = addon_output_dir / libname if not lib_path.exists(): logger.debug("JIT-compiling torch-c-dlpack-ext to cache...") build_script_path = ( Path(__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py" ) args = [ sys.executable, str(build_script_path), "--output-dir", str(cache_dir), "--libname", libname, ] if device == "cuda": args.append("--build-with-cuda") elif device == "rocm": args.append("--build-with-rocm") # use capture_output to reduce noise when building the torch c dlpack addon result = subprocess.run(args, check=False, capture_output=True) if result.returncode != 0: msg = [f"Build failed with status {result.returncode}"] if result.stdout: msg.append(f"stdout:\n{result.stdout.decode('utf-8')}") if result.stderr: msg.append(f"stderr:\n{result.stderr.decode('utf-8')}") raise RuntimeError("\n".join(msg)) if not lib_path.exists(): raise RuntimeError("Failed to build torch c dlpack addon.") lib = ctypes.CDLL(str(lib_path)) func = lib.TorchDLPackExchangeAPIPtr func.restype = ctypes.c_uint64 func.argtypes = [] # Create a PyCapsule from the pointer capsule = _create_dlpack_exchange_api_capsule(func()) # Set the DLPackExchangeAPI pointer on the class setattr(torch.Tensor, "__dlpack_c_exchange_api__", capsule) return lib except ImportError: pass except Exception: warnings.warn( "Failed to JIT torch c dlpack extension, EnvTensorAllocator will not be enabled.\n" "We recommend installing via `pip install torch-c-dlpack-ext`" ) return None def patch_torch_cuda_stream_protocol() -> None: """Load the torch cuda stream protocol for older versions of torch.""" try: import torch # noqa: PLC0415 if not torch.cuda.is_available(): return if not hasattr(torch.cuda.Stream, "__cuda_stream__"): def __torch_cuda_stream__(self: torch.cuda.Stream) -> tuple[int, int]: """Return the version number and the cuda stream.""" return (0, self.cuda_stream) setattr(torch.cuda.Stream, "__cuda_stream__", __torch_cuda_stream__) except ImportError: pass if os.environ.get("TVM_FFI_DISABLE_TORCH_C_DLPACK", "0") == "0": _LIB = load_torch_c_dlpack_extension() # keep a reference to the loaded shared library patch_torch_cuda_stream_protocol() tvm-ffi-0.1.12/python/tvm_ffi/_tensor.py000066400000000000000000000065411521067262500201550ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tensor related objects and functions.""" from __future__ import annotations # we name it as _tensor.py to avoid potential future case # if we also want to expose a tensor function in the root namespace from numbers import Integral from typing import Any from . import _ffi_api, core, registry from .core import ( Device, DLDeviceType, PyNativeObject, Tensor, _shape_obj_get_py_tuple, from_dlpack, ) @registry.register_object("ffi.Shape") class Shape(tuple, PyNativeObject): """Shape tuple that represents :cpp:class:`tvm::ffi::Shape` returned by an FFI call. Notes ----- This class subclasses :class:`tuple` so it can be used in most places where :class:`tuple` is used in Python array APIs. Examples -------- .. code-block:: python import numpy as np import tvm_ffi x = tvm_ffi.from_dlpack(np.arange(6, dtype="int32").reshape(2, 3)) assert x.shape == (2, 3) """ _tvm_ffi_cached_object: Any # tvm-ffi-stubgen(begin): object/ffi.Shape # fmt: off # fmt: on # tvm-ffi-stubgen(end) def __new__(cls, content: tuple[int, ...]) -> Shape: if any(not isinstance(x, Integral) for x in content): raise ValueError("Shape must be a tuple of integers") val: Shape = tuple.__new__(cls, content) val.__init_cached_object_by_constructor__(_ffi_api.Shape, *content) return val # pylint: disable=no-self-argument def __from_tvm_ffi_object__(cls, obj: Any) -> Shape: """Construct from a given tvm object.""" content = _shape_obj_get_py_tuple(obj) val: Shape = tuple.__new__(cls, content) # ty: ignore[invalid-argument-type] val._tvm_ffi_cached_object = obj return val def device(device_type: str | int | DLDeviceType, index: int | None = None) -> Device: """Construct a TVM FFI device with given device type and index. Parameters ---------- device_type: str or int The device type or name. index: int, optional The device index. Returns ------- device: tvm_ffi.Device Examples -------- Device can be used to create reflection of device by string representation of the device type. .. code-block:: python import tvm_ffi assert tvm_ffi.device("cuda:0") == tvm_ffi.device("cuda", 0) assert tvm_ffi.device("cpu:0") == tvm_ffi.device("cpu", 0) """ # must refer to core._CLASS_DEVICE so we pick up override here return core._CLASS_DEVICE(device_type, index) __all__ = ["DLDeviceType", "Device", "Tensor", "device", "from_dlpack"] tvm-ffi-0.1.12/python/tvm_ffi/access_path.py000066400000000000000000000161501521067262500207560ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name """Access path classes.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Sequence from tvm_ffi import Object from typing import Any # isort: on # fmt: on # tvm-ffi-stubgen(end) from enum import IntEnum from .core import Object from .registry import register_object class AccessKind(IntEnum): """Kinds of access steps in an access path.""" ATTR = 0 ARRAY_ITEM = 1 MAP_ITEM = 2 ATTR_MISSING = 3 ARRAY_ITEM_MISSING = 4 MAP_ITEM_MISSING = 5 @register_object("ffi.reflection.AccessStep") class AccessStep(Object): """Access step container.""" # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessStep # fmt: off kind: int key: Any # fmt: on # tvm-ffi-stubgen(end) @register_object("ffi.reflection.AccessPath") class AccessPath(Object): """Access path container. It describes how to reach a nested attribute or item inside a complex FFI object by recording a sequence of steps (attribute, array index, or map key). It is primarily used by diagnostics to pinpoint structural mismatches. Examples -------- .. code-block:: python from tvm_ffi.access_path import AccessPath root = AccessPath.root() # Build a path equivalent to obj.layer.weight[2] p = root.attr("layer").attr("weight").array_item(2) assert isinstance(p, AccessPath) """ # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessPath -> ffi.access_path.AccessPath # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessPath # fmt: off parent: Object | None step: AccessStep | None depth: int if TYPE_CHECKING: @staticmethod def _root() -> AccessPath: ... def _extend(self, _1: AccessStep, /) -> AccessPath: ... def _attr(self, _1: str, /) -> AccessPath: ... def _array_item(self, _1: int, /) -> AccessPath: ... def _map_item(self, _1: Any, /) -> AccessPath: ... def _attr_missing(self, _1: str, /) -> AccessPath: ... def _array_item_missing(self, _1: int, /) -> AccessPath: ... def _map_item_missing(self, _1: Any, /) -> AccessPath: ... def _is_prefix_of(self, _1: AccessPath, /) -> bool: ... def _to_steps(self, /) -> Sequence[AccessStep]: ... def _path_equal(self, _1: AccessPath, /) -> bool: ... # fmt: on # tvm-ffi-stubgen(end) def __init__(self) -> None: """Disallow direct construction; use `AccessPath.root()` instead.""" raise ValueError( "AccessPath can't be initialized directly. " "Use AccessPath.root() to create a path to the root object" ) @staticmethod def root() -> AccessPath: """Create a root access path. Returns ------- AccessPath A path representing the root of an object graph. """ return AccessPath._root() def __eq__(self, other: Any) -> bool: """Return whether two access paths are equal.""" if not isinstance(other, AccessPath): return False return self._path_equal(other) def __ne__(self, other: Any) -> bool: """Return whether two access paths are not equal.""" if not isinstance(other, AccessPath): return True return not self._path_equal(other) def is_prefix_of(self, other: AccessPath) -> bool: """Check if this access path is a prefix of another access path. Parameters ---------- other The access path to check if it is a prefix of this access path Returns ------- bool True if this access path is a prefix of the other access path, False otherwise """ return self._is_prefix_of(other) def attr(self, attr_key: str) -> AccessPath: """Create an access path to the attribute of the current object. Parameters ---------- attr_key The key of the attribute to access Returns ------- AccessPath The extended access path """ return self._attr(attr_key) def attr_missing(self, attr_key: str) -> AccessPath: """Create an access path that indicate an attribute is missing. Parameters ---------- attr_key The key of the attribute to access Returns ------- AccessPath The extended access path """ return self._attr_missing(attr_key) def array_item(self, index: int) -> AccessPath: """Create an access path to the item of the current array. Parameters ---------- index The index of the item to access Returns ------- AccessPath The extended access path """ return self._array_item(index) def array_item_missing(self, index: int) -> AccessPath: """Create an access path that indicate an array item is missing. Parameters ---------- index The index of the item to access Returns ------- AccessPath The extended access path """ return self._array_item_missing(index) def map_item(self, key: Any) -> AccessPath: """Create an access path to the item of the current map. Parameters ---------- key The key of the item to access Returns ------- AccessPath The extended access path """ return self._map_item(key) def map_item_missing(self, key: Any) -> AccessPath: """Create an access path that indicate a map item is missing. Parameters ---------- key The key of the item to access Returns ------- AccessPath The extended access path """ return self._map_item_missing(key) def to_steps(self) -> Sequence[AccessStep]: """Convert the access path to a list of access steps. Returns ------- access_steps The list of access steps """ return self._to_steps() __hash__ = Object.__hash__ tvm-ffi-0.1.12/python/tvm_ffi/config.py000066400000000000000000000073171521067262500177530ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Config utilities for finding paths to lib and headers.""" import argparse import sys from pathlib import Path from . import libinfo def _find_libdir() -> str: """Find the library directory for tvm-ffi.""" libtvm_ffi = Path(libinfo.find_libtvm_ffi()) return str(libtvm_ffi.parent) def __main__() -> None: # noqa: PLR0912 """Parse CLI args and print build and include configuration paths.""" parser = argparse.ArgumentParser( description="Get various configuration information needed to compile with tvm-ffi" ) parser.add_argument("--includedir", action="store_true", help="Print include directory") parser.add_argument( "--dlpack-includedir", action="store_true", help="Print dlpack include directory", ) parser.add_argument("--cmakedir", action="store_true", help="Print library directory") parser.add_argument("--sourcedir", action="store_true", help="Print source directory") parser.add_argument("--libfiles", action="store_true", help="Fully qualified library filenames") parser.add_argument("--libdir", action="store_true", help="Print library directory") parser.add_argument("--libs", action="store_true", help="Libraries to be linked") parser.add_argument("--cython-lib-path", action="store_true", help="Print cython path") parser.add_argument("--cxxflags", action="store_true", help="Print cxx flags") parser.add_argument("--cflags", action="store_true", help="Print c flags") parser.add_argument("--ldflags", action="store_true", help="Print ld flags") args = parser.parse_args() # print help when no arguments are provided if len(sys.argv) == 1: parser.print_help() return if args.includedir: print(libinfo.find_include_path()) if args.dlpack_includedir: print(libinfo.find_dlpack_include_path()) if args.cmakedir: print(libinfo.find_cmake_path()) if args.libdir: print(_find_libdir()) if args.libfiles: if sys.platform.startswith("win32"): print(libinfo.find_windows_implib()) else: print(libinfo.find_libtvm_ffi()) if args.sourcedir: print(libinfo.find_source_path()) if args.cython_lib_path: print(libinfo.find_cython_lib()) if args.cxxflags: include_dir = libinfo.find_include_path() dlpack_include_dir = libinfo.find_dlpack_include_path() print(f"-I{include_dir} -I{dlpack_include_dir} -std=c++17") if args.cflags: include_dir = libinfo.find_include_path() dlpack_include_dir = libinfo.find_dlpack_include_path() print(f"-I{include_dir} -I{dlpack_include_dir}") if args.libs: if sys.platform.startswith("win32"): print(libinfo.find_windows_implib()) else: print("-ltvm_ffi") if args.ldflags: if not sys.platform.startswith("win32"): print(f"-L{_find_libdir()}") if __name__ == "__main__": __main__() tvm-ffi-0.1.12/python/tvm_ffi/container.py000066400000000000000000000573111521067262500204670ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Container classes.""" from __future__ import annotations import itertools import operator import sys from typing import ( Any, Callable, SupportsIndex, TypeVar, cast, overload, ) from . import _ffi_api, core from .registry import register_object if sys.version_info >= (3, 9): # PEP 585 generics from collections.abc import ( ItemsView as ItemsViewBase, ) from collections.abc import ( Iterable, Iterator, Mapping, MutableMapping, MutableSequence, Sequence, ) from collections.abc import ( KeysView as KeysViewBase, ) from collections.abc import ( ValuesView as ValuesViewBase, ) else: # Python 3.8 # workarounds for python 3.8 # typing-module generics (subscriptable on 3.8) from typing import ( ItemsView as ItemsViewBase, ) from typing import ( Iterable, Iterator, Mapping, MutableMapping, MutableSequence, Sequence, ) from typing import ( KeysView as KeysViewBase, ) from typing import ( ValuesView as ValuesViewBase, ) __all__ = ["Array", "Dict", "List", "Map"] T = TypeVar("T") K = TypeVar("K") V = TypeVar("V") def _sequence_compare_other(this: object, other: object) -> object: """Normalize plain Python sequences for structural container equality.""" if isinstance(other, (str, bytes, Mapping)): return NotImplemented if isinstance(other, Sequence): try: return type(this)(other) except (TypeError, ValueError): return NotImplemented return NotImplemented def _sequence_contains( this: Sequence[Any], value: object, ffi_contains: Callable[[Any, object], bool], ) -> bool: """Containment with a Python-level structural fallback for nested sequences.""" if ffi_contains(this, value): return True if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): return False try: search_value = type(this)(value) # ty: ignore[too-many-positional-arguments] except (TypeError, ValueError): return False for item in this: if item == search_value: return True return False _DefaultT = TypeVar("_DefaultT") from .core import MISSING def getitem_helper( obj: Any, elem_getter: Callable[[Any, int], T], length: int, idx: SupportsIndex | slice, ) -> T | list[T]: """Implement a pythonic __getitem__ helper. Parameters ---------- obj The original object elem_getter A simple function that takes index and return a single element. length The size of the array idx The argument passed to getitem Returns ------- result The element for integer indices or a :class:`list` for slices. """ if isinstance(idx, slice): start, stop, step = idx.indices(length) return [elem_getter(obj, i) for i in range(start, stop, step)] index = normalize_index(length, idx) return elem_getter(obj, index) def normalize_index(length: int, idx: SupportsIndex) -> int: """Normalize and bounds-check a Python index.""" try: index = operator.index(idx) except TypeError as exc: # pragma: no cover - defensive, matches list behaviour raise TypeError(f"indices must be integers or slices, not {type(idx).__name__}") from exc if index < -length or index >= length: raise IndexError(f"Index out of range. size: {length}, got index {index}") if index < 0: index += length return index @register_object("ffi.Array") class Array(core.CContainerBase, core.Object, Sequence[T]): """Array container that represents a sequence of values in the FFI. :py:func:`tvm_ffi.convert` will map python list/tuple to this class. Parameters ---------- input_list The list of values to be stored in the array. Examples -------- .. code-block:: python import tvm_ffi a = tvm_ffi.Array([1, 2, 3]) assert tuple(a) == (1, 2, 3) Notes ----- For structural equality and hashing, use ``structural_equal`` and ``structural_hash`` APIs. See Also -------- :py:func:`tvm_ffi.convert` """ # tvm-ffi-stubgen(begin): object/ffi.Array # fmt: off # fmt: on # tvm-ffi-stubgen(end) def __deepcopy__(self, memo: Any = None) -> Any: return _ffi_api.DeepCopy(self) def __init__(self, input_list: Iterable[T]) -> None: """Construct an Array from a Python sequence.""" self.__init_handle_by_constructor__(_ffi_api.Array, *input_list) @overload def __getitem__(self, idx: SupportsIndex, /) -> T: ... @overload def __getitem__(self, idx: slice, /) -> list[T]: ... def __getitem__(self, idx: SupportsIndex | slice, /) -> T | list[T]: # ty: ignore[invalid-method-override] """Return one element or a list for a slice.""" length = len(self) result = getitem_helper(self, _ffi_api.ArrayGetItem, length, idx) return result def __len__(self) -> int: """Return the number of elements in the array.""" return _ffi_api.ArraySize(self) def __iter__(self) -> Iterator[T]: """Iterate over the elements in the array.""" length = len(self) for i in range(length): yield self[i] def __contains__(self, value: object) -> bool: """Check if the array contains a value.""" return _sequence_contains(self, value, _ffi_api.ArrayContains) def __eq__(self, other: object) -> bool: """Structural equality.""" if isinstance(other, type(self)) or isinstance(self, type(other)): return _ffi_api.RecursiveEq(self, other) other = _sequence_compare_other(self, other) if other is NotImplemented: return NotImplemented return _ffi_api.RecursiveEq(self, other) def __ne__(self, other: object) -> bool: """Structural inequality.""" result = self.__eq__(other) if result is NotImplemented: return NotImplemented return not result def __hash__(self) -> int: """Structural hash.""" return _ffi_api.RecursiveHash(self) def __bool__(self) -> bool: """Return True if the array is non-empty.""" return len(self) > 0 def __add__(self, other: Iterable[T]) -> Array[T]: """Concatenate two arrays.""" return type(self)(itertools.chain(self, other)) def __radd__(self, other: Iterable[T]) -> Array[T]: """Concatenate two arrays.""" return type(self)(itertools.chain(other, self)) @register_object("ffi.List") class List(core.CContainerBase, core.Object, MutableSequence[T]): """Mutable list container that represents a mutable sequence in the FFI.""" # tvm-ffi-stubgen(begin): object/ffi.List # fmt: off # fmt: on # tvm-ffi-stubgen(end) def __deepcopy__(self, memo: Any = None) -> Any: return _ffi_api.DeepCopy(self) def __init__(self, input_list: Iterable[T] = ()) -> None: """Construct a List from a Python sequence.""" self.__init_handle_by_constructor__(_ffi_api.List, *input_list) @overload def __getitem__(self, idx: SupportsIndex, /) -> T: ... @overload def __getitem__(self, idx: slice, /) -> list[T]: ... def __getitem__(self, idx: SupportsIndex | slice, /) -> T | list[T]: # ty: ignore[invalid-method-override] """Return one element or a list for a slice.""" length = len(self) return getitem_helper(self, _ffi_api.ListGetItem, length, idx) @overload def __setitem__(self, index: SupportsIndex, value: T) -> None: ... @overload def __setitem__(self, index: slice[int | None], value: Iterable[T]) -> None: ... def __setitem__(self, index: SupportsIndex | slice[int | None], value: T | Iterable[T]) -> None: """Set one element or assign a slice.""" if isinstance(index, slice): replacement = list(cast(Iterable[T], value)) length = len(self) start, stop, step = index.indices(length) if step != 1: target_indices = list(range(start, stop, step)) if len(replacement) != len(target_indices): raise ValueError( "attempt to assign sequence of size " f"{len(replacement)} to extended slice of size {len(target_indices)}" ) for i, item in zip(target_indices, replacement): _ffi_api.ListSetItem(self, i, item) return stop = max(stop, start) _ffi_api.ListReplaceSlice(self, start, stop, type(self)(replacement)) return normalized_index = normalize_index(len(self), index) _ffi_api.ListSetItem(self, normalized_index, cast(T, value)) @overload def __delitem__(self, index: SupportsIndex) -> None: ... @overload def __delitem__(self, index: slice[int | None]) -> None: ... def __delitem__(self, index: SupportsIndex | slice[int | None]) -> None: """Delete one element or a slice.""" if isinstance(index, slice): length = len(self) start, stop, step = index.indices(length) if step == 1: stop = max(stop, start) _ffi_api.ListEraseRange(self, start, stop) else: # Delete indices from high to low so that earlier deletions # do not shift the positions of later ones. indices = ( reversed(range(start, stop, step)) if step > 0 else range(start, stop, step) ) for i in indices: _ffi_api.ListErase(self, i) return normalized_index = normalize_index(len(self), index) _ffi_api.ListErase(self, normalized_index) def insert(self, index: int, value: T) -> None: """Insert value before index.""" length = len(self) if index < 0: index = max(0, index + length) else: index = min(index, length) _ffi_api.ListInsert(self, index, value) def append(self, value: T) -> None: """Append one value to the tail.""" _ffi_api.ListAppend(self, value) def clear(self) -> None: """Remove all elements from the list.""" _ffi_api.ListClear(self) def reverse(self) -> None: """Reverse the list in-place.""" _ffi_api.ListReverse(self) def pop(self, index: int = -1) -> T: """Remove and return item at index (default last).""" length = len(self) if length == 0: raise IndexError("pop from empty list") normalized_index = normalize_index(length, index) return cast(T, _ffi_api.ListPop(self, normalized_index)) def extend(self, values: Iterable[T]) -> None: """Append elements from an iterable.""" end = len(self) self[end:end] = values def __len__(self) -> int: """Return the number of elements in the list.""" return _ffi_api.ListSize(self) def __iter__(self) -> Iterator[T]: """Iterate over the elements in the list.""" length = len(self) for i in range(length): yield cast(T, _ffi_api.ListGetItem(self, i)) def __contains__(self, value: object) -> bool: """Check if the list contains a value.""" return _sequence_contains(self, value, _ffi_api.ListContains) def __eq__(self, other: object) -> bool: """Structural equality.""" if isinstance(other, type(self)) or isinstance(self, type(other)): return _ffi_api.RecursiveEq(self, other) other = _sequence_compare_other(self, other) if other is NotImplemented: return NotImplemented return _ffi_api.RecursiveEq(self, other) def __ne__(self, other: object) -> bool: """Structural inequality.""" result = self.__eq__(other) if result is NotImplemented: return NotImplemented return not result def __hash__(self) -> int: """Structural hash.""" return _ffi_api.RecursiveHash(self) def __bool__(self) -> bool: """Return True if the list is non-empty.""" return len(self) > 0 def __add__(self, other: Iterable[T]) -> List[T]: """Concatenate two lists.""" return type(self)(itertools.chain(self, other)) def __radd__(self, other: Iterable[T]) -> List[T]: """Concatenate two lists.""" return type(self)(itertools.chain(other, self)) class KeysView(KeysViewBase[K]): """Helper class to return keys view.""" def __init__( self, backend_map: Map[K, V] | Dict[K, V], iter_functor_getter: Callable[..., Callable[[int], Any]] | None = None, ) -> None: self._backend_map = backend_map self._iter_functor_getter = iter_functor_getter or _ffi_api.MapForwardIterFunctor def __len__(self) -> int: return len(self._backend_map) def __iter__(self) -> Iterator[K]: size = len(self._backend_map) functor: Callable[[int], Any] = self._iter_functor_getter(self._backend_map) for _ in range(size): key = cast(K, functor(0)) yield key if not functor(2): break def __contains__(self, k: object) -> bool: # ty: ignore[invalid-method-override] return k in self._backend_map class ValuesView(ValuesViewBase[V]): """Helper class to return values view.""" def __init__( self, backend_map: Map[K, V] | Dict[K, V], iter_functor_getter: Callable[..., Callable[[int], Any]] | None = None, ) -> None: self._backend_map = backend_map self._iter_functor_getter = iter_functor_getter or _ffi_api.MapForwardIterFunctor def __len__(self) -> int: return len(self._backend_map) def __iter__(self) -> Iterator[V]: size = len(self._backend_map) functor: Callable[[int], Any] = self._iter_functor_getter(self._backend_map) for _ in range(size): value = cast(V, functor(1)) yield value if not functor(2): break class ItemsView(ItemsViewBase[K, V]): """Helper class to return items view.""" def __init__( self, backend_map: Map[K, V] | Dict[K, V], iter_functor_getter: Callable[..., Callable[[int], Any]] | None = None, ) -> None: self._backend_map = backend_map self._iter_functor_getter = iter_functor_getter or _ffi_api.MapForwardIterFunctor def __len__(self) -> int: return len(self._backend_map) def __iter__(self) -> Iterator[tuple[K, V]]: size = len(self._backend_map) functor: Callable[[int], Any] = self._iter_functor_getter(self._backend_map) for _ in range(size): key = cast(K, functor(0)) value = cast(V, functor(1)) yield (key, value) if not functor(2): break def __contains__(self, item: object) -> bool: if not isinstance(item, tuple) or len(item) != 2: return False key, value = item actual_value = self._backend_map.get(key, MISSING) # ty: ignore[invalid-argument-type] if actual_value is MISSING: return False # TODO(@junrus): Is `__eq__` the right method to use here? return actual_value == value @register_object("ffi.Map") class Map(core.CContainerBase, core.Object, Mapping[K, V]): """Map container. :py:func:`tvm_ffi.convert` will map python dict to this class. Parameters ---------- input_dict The dictionary of values to be stored in the map. Examples -------- .. code-block:: python import tvm_ffi amap = tvm_ffi.Map({"a": 1, "b": 2}) assert len(amap) == 2 assert amap["a"] == 1 assert amap["b"] == 2 Notes ----- For structural equality and hashing, use ``structural_equal`` and ``structural_hash`` APIs. See Also -------- :py:func:`tvm_ffi.convert` """ # tvm-ffi-stubgen(begin): object/ffi.Map # fmt: off # fmt: on # tvm-ffi-stubgen(end) def __deepcopy__(self, memo: Any = None) -> Any: return _ffi_api.DeepCopy(self) def __init__(self, input_dict: Mapping[K, V]) -> None: """Construct a Map from a Python mapping.""" list_kvs: list[Any] = [] for k, v in input_dict.items(): list_kvs.append(k) list_kvs.append(v) self.__init_handle_by_constructor__(_ffi_api.Map, *list_kvs) def __getitem__(self, k: K) -> V: """Return the value for key `k` or raise KeyError.""" return cast(V, _ffi_api.MapGetItem(self, k)) def __contains__(self, k: object) -> bool: """Return True if the map contains key `k`.""" return _ffi_api.MapCount(self, k) != 0 def __eq__(self, other: object) -> bool: """Structural equality.""" if not (isinstance(other, type(self)) or isinstance(self, type(other))): return NotImplemented return _ffi_api.RecursiveEq(self, other) def __ne__(self, other: object) -> bool: """Structural inequality.""" result = self.__eq__(other) if result is NotImplemented: return NotImplemented return not result def __hash__(self) -> int: """Structural hash.""" return _ffi_api.RecursiveHash(self) def keys(self) -> KeysView[K]: """Return a dynamic view of the map's keys.""" return KeysView(self) def values(self) -> ValuesView[V]: """Return a dynamic view of the map's values.""" return ValuesView(self) def items(self) -> ItemsView[K, V]: """Get the items from the map.""" return ItemsView(self) def __len__(self) -> int: """Return the number of items in the map.""" return _ffi_api.MapSize(self) def __bool__(self) -> bool: """Return True if the map is non-empty.""" return len(self) > 0 def __iter__(self) -> Iterator[K]: """Iterate over the map's keys.""" return iter(self.keys()) @overload def get(self, key: K) -> V | None: ... @overload def get(self, key: K, default: V | _DefaultT) -> V | _DefaultT: ... def get(self, key: K, default: V | _DefaultT | None = None) -> V | _DefaultT | None: """Get an element with a default value. Parameters ---------- key The attribute key. default The default object. Returns ------- value The result value. """ ret = _ffi_api.MapGetItemOrMissing(self, key) if MISSING.same_as(ret): return default return ret @register_object("ffi.Dict") class Dict(core.CContainerBase, core.Object, MutableMapping[K, V]): """Mutable dictionary container with shared reference semantics. Unlike :class:`Map`, ``Dict`` does NOT implement copy-on-write. Mutations happen directly on the underlying shared object. All Python references sharing the same ``Dict`` see mutations immediately. Parameters ---------- input_dict The dictionary of values to be stored. Examples -------- .. code-block:: python import tvm_ffi d = tvm_ffi.Dict({"a": 1, "b": 2}) d["c"] = 3 assert len(d) == 3 """ def __deepcopy__(self, memo: Any = None) -> Any: return _ffi_api.DeepCopy(self) def __init__(self, input_dict: Mapping[K, V] | None = None) -> None: """Construct a Dict from a Python mapping.""" list_kvs: list[Any] = [] if input_dict is not None: for k, v in input_dict.items(): list_kvs.append(k) list_kvs.append(v) self.__init_handle_by_constructor__(_ffi_api.Dict, *list_kvs) def __getitem__(self, k: K) -> V: """Return the value for key `k` or raise KeyError.""" return cast(V, _ffi_api.DictGetItem(self, k)) def __setitem__(self, k: K, v: V) -> None: """Set the value for key `k`.""" _ffi_api.DictSetItem(self, k, v) def __delitem__(self, k: K) -> None: """Delete the entry for key `k`.""" if _ffi_api.DictCount(self, k) == 0: raise KeyError(k) _ffi_api.DictErase(self, k) def __contains__(self, k: object) -> bool: """Return True if the dict contains key `k`.""" return _ffi_api.DictCount(self, k) != 0 def __eq__(self, other: object) -> bool: """Structural equality.""" if not (isinstance(other, type(self)) or isinstance(self, type(other))): return NotImplemented return _ffi_api.RecursiveEq(self, other) def __ne__(self, other: object) -> bool: """Structural inequality.""" result = self.__eq__(other) if result is NotImplemented: return NotImplemented return not result def __hash__(self) -> int: """Structural hash.""" return _ffi_api.RecursiveHash(self) def __len__(self) -> int: """Return the number of items in the dict.""" return _ffi_api.DictSize(self) def __bool__(self) -> bool: """Return True if the dict is non-empty.""" return len(self) > 0 def __iter__(self) -> Iterator[K]: """Iterate over the dict's keys.""" return iter(self.keys()) def keys(self) -> KeysView[K]: """Return a dynamic view of the dict's keys.""" return KeysView(self, _ffi_api.DictForwardIterFunctor) def values(self) -> ValuesView[V]: """Return a dynamic view of the dict's values.""" return ValuesView(self, _ffi_api.DictForwardIterFunctor) def items(self) -> ItemsView[K, V]: """Get the items from the dict.""" return ItemsView(self, _ffi_api.DictForwardIterFunctor) @overload def get(self, key: K) -> V | None: ... @overload def get(self, key: K, default: V | _DefaultT) -> V | _DefaultT: ... def get(self, key: K, default: V | _DefaultT | None = None) -> V | _DefaultT | None: """Get an element with a default value.""" ret = _ffi_api.DictGetItemOrMissing(self, key) if MISSING.same_as(ret): return default return ret def pop(self, key: K, *args: V | _DefaultT) -> V | _DefaultT: """Remove and return value for key, or default if not present.""" if len(args) > 1: raise TypeError(f"pop expected at most 2 arguments, got {1 + len(args)}") ret = _ffi_api.DictGetItemOrMissing(self, key) if MISSING.same_as(ret): if args: return args[0] raise KeyError(key) _ffi_api.DictErase(self, key) return cast(V, ret) def clear(self) -> None: """Remove all elements from the dict.""" _ffi_api.DictClear(self) def update(self, other: Mapping[K, V]) -> None: # type: ignore[override] """Update the dict from a mapping.""" for k, v in other.items(): self[k] = v tvm-ffi-0.1.12/python/tvm_ffi/core.pyi000066400000000000000000000246301521067262500176040ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Typestubs for Cython.""" from __future__ import annotations import types from ctypes import c_void_p from enum import IntEnum from typing import Any, Callable # Public module-level variables referenced by Python code MISSING: Object KWARGS: Object ERROR_NAME_TO_TYPE: dict[str, type] ERROR_TYPE_TO_NAME: dict[type, str] _WITH_APPEND_BACKTRACE: Callable[[BaseException, str], BaseException] | None _TRACEBACK_TO_BACKTRACE_STR: Callable[[types.TracebackType | None], str] | None # DLPack protocol version (defined in tensor.pxi) __dlpack_version__: tuple[int, int] class CObject: def __ctypes_handle__(self) -> Any: ... def __chandle__(self) -> int: ... @property def id_(self) -> int: ... def __reduce__(self) -> Any: ... def __getstate__(self) -> dict[str, Any]: ... def __setstate__(self, state: dict[str, Any]) -> None: ... def __repr__(self) -> str: ... def __eq__(self, other: Any) -> bool: ... def __ne__(self, other: Any) -> bool: ... def __hash__(self) -> int: ... def __init_handle_by_constructor__(self, fconstructor: Any, *args: Any) -> None: ... def same_as(self, other: Any) -> bool: ... def is_(self, other: Any) -> bool: ... def _move(self) -> ObjectRValueRef: ... def __move_handle_from__(self, other: CObject) -> None: ... class CContainerBase(CObject): ... class Object(CObject): ... def object_repr(obj: CObject) -> str: ... class ObjectConvertible: def asobject(self) -> Object: ... class ObjectRValueRef: obj: Object def __init__(self, obj: Object) -> None: ... class OpaquePyObject(Object): def pyobject(self) -> Any: ... class PyNativeObject: __slots__: list[str] def __init_cached_object_by_constructor__(self, fconstructor: Any, *args: Any) -> None: ... def _register_object_by_index(type_index: int, type_cls: type) -> TypeInfo: ... def _object_type_key_to_index(type_key: str) -> int | None: ... def _set_type_cls(type_info: TypeInfo, type_cls: type) -> None: ... def _lookup_or_register_type_info_from_type_key(type_key: str) -> TypeInfo: ... def _lookup_type_attr(type_index: int, attr_key: str) -> Any: ... def _register_type_attr(type_index: int, attr_key: str, value: Any) -> None: ... def _type_cls_to_type_info(type_cls: type) -> TypeInfo | None: ... class Error(Object): def __init__(self, kind: str, message: str, backtrace: str) -> None: ... def update_backtrace(self, backtrace: str) -> None: ... def py_error(self) -> BaseException: ... @property def kind(self) -> str: ... @property def message(self) -> str: ... @property def backtrace(self) -> str: ... def _convert_to_ffi_error(error: BaseException) -> Error: ... def _env_set_current_stream( device_type: int, device_id: int, stream: int | c_void_p ) -> int | c_void_p: ... def _env_get_current_stream(device_type: int, device_id: int) -> int: ... class DataType: def __init__(self, dtype_str: str) -> None: ... def __reduce__(self) -> Any: ... def __eq__(self, other: Any) -> bool: ... def __ne__(self, other: Any) -> bool: ... @property def type_code(self) -> int: ... @property def bits(self) -> int: ... @property def lanes(self) -> int: ... @property def itemsize(self) -> int: ... def __str__(self) -> str: ... def _set_class_dtype(cls: type) -> None: ... def _convert_torch_dtype_to_ffi_dtype(torch_dtype: Any) -> DataType: ... def _convert_numpy_dtype_to_ffi_dtype(numpy_dtype: Any) -> DataType: ... def _create_cdtype_from_tuple( cls: type[DataType], code: int, bits: int, lanes: int ) -> DataType: ... class DLDeviceType(IntEnum): kDLCPU = 1 kDLCUDA = 2 kDLCUDAHost = 3 kDLOpenCL = 4 kDLVulkan = 7 kDLMetal = 8 kDLVPI = 9 kDLROCM = 10 kDLROCMHost = 11 kDLExtDev = 12 kDLCUDAManaged = 13 kDLOneAPI = 14 kDLWebGPU = 15 kDLHexagon = 16 kDLTrn = 17 class Device: def __init__(self, device_type: str | int, index: int | None = None) -> None: ... def __reduce__(self) -> Any: ... def __eq__(self, other: Any) -> bool: ... def __ne__(self, other: Any) -> bool: ... def __str__(self) -> str: ... def __repr__(self) -> str: ... def __hash__(self) -> int: ... def __device_type_name__(self) -> str: ... @property def type(self) -> str: ... @property def index(self) -> int: ... def dlpack_device_type(self) -> int: ... def _set_class_device(cls: type) -> None: ... _CLASS_DEVICE: type[Device] def _shape_obj_get_py_tuple(obj: Any) -> tuple[int, ...]: ... class Tensor(Object): @property def shape(self) -> tuple[int, ...]: ... @property def ndim(self) -> int: ... def numel(self) -> int: ... def size(self, idx: int) -> int: ... def is_contiguous(self) -> bool: ... @property def strides(self) -> tuple[int, ...]: ... @property def dtype(self) -> Any: ... @property def device(self) -> Device: ... def _to_dlpack(self) -> Any: ... def _to_dlpack_versioned(self) -> Any: ... def __dlpack_device__(self) -> tuple[int, int]: ... def __dlpack__( self, *, stream: Any | None = None, max_version: tuple[int, int] | None = None, dl_device: tuple[int, int] | None = None, copy: bool | None = None, ) -> Any: ... _CLASS_TENSOR: type[Tensor] = Tensor def _set_class_tensor(cls: type[Tensor]) -> None: ... def from_dlpack( ext_tensor: Any, *, require_alignment: int = ..., require_contiguous: bool = ... ) -> Tensor: ... class DLTensorTestWrapper: __dlpack_c_exchange_api__: int def __init__(self, tensor: Tensor) -> None: ... def __tvm_ffi_env_stream__(self) -> int: ... def __dlpack_device__(self) -> tuple[int, int]: ... def __dlpack__(self, **kwargs: Any) -> Any: ... def _dltensor_test_wrapper_c_dlpack_from_pyobject_as_intptr() -> int: ... class Function(Object): @property def release_gil(self) -> bool: ... @release_gil.setter def release_gil(self, value: bool) -> None: ... def __call__(self, *args: Any) -> Any: ... @staticmethod def __from_extern_c__(c_symbol: int, *, keep_alive_object: Any | None = None) -> Function: ... @staticmethod def __from_mlir_packed_safe_call__( mlir_packed_symbol: int, *, keep_alive_object: Any | None = None ) -> Function: ... def _register_global_func( name: str, pyfunc: Callable[..., Any] | Function, override: bool ) -> Function: ... def _get_global_func(name: str, allow_missing: bool) -> Function | None: ... def _convert_to_ffi_func(pyfunc: Callable[..., Any], tensor_cls: type | None = ...) -> Function: ... def _convert_to_opaque_object(pyobject: Any) -> OpaquePyObject: ... def _print_debug_info() -> None: ... def _register_py_class(parent_type_info: TypeInfo, type_key: str, type_cls: type) -> TypeInfo: ... def _register_fields(type_info: TypeInfo, fields: list[Any]) -> list[TypeField]: ... class String(str, PyNativeObject): __slots__ = ["_tvm_ffi_cached_object"] _tvm_ffi_cached_object: Object | None def __new__(cls, value: str) -> String: ... # pylint: disable=no-self-argument def __from_tvm_ffi_object__(cls, obj: Any) -> String: ... class Bytes(bytes, PyNativeObject): __slots__ = ["_tvm_ffi_cached_object"] _tvm_ffi_cached_object: Object | None def __new__(cls, value: bytes) -> Bytes: ... # pylint: disable=no-self-argument def __from_tvm_ffi_object__(cls, obj: Any) -> Bytes: ... # --------------------------------------------------------------------------- # Owned FFI value container (from cython/object.pxi) # --------------------------------------------------------------------------- class CAny: def __init__(self, value: Any = None) -> None: ... @property def type_index(self) -> int: ... def _to_py_class_value(value: CAny) -> Any: ... # --------------------------------------------------------------------------- # Type reflection metadata (from cython/type_info.pxi) # --------------------------------------------------------------------------- class TypeSchema: origin: str args: tuple[TypeSchema, ...] = () origin_type_index: int def __init__( self, origin: str, args: tuple[TypeSchema, ...] = (), origin_type_index: int = ..., ) -> None: ... @staticmethod def from_json_obj(obj: dict[str, Any]) -> TypeSchema: ... @staticmethod def from_json_str(s: str) -> TypeSchema: ... @staticmethod def from_type_index(type_index: int, args: tuple[TypeSchema, ...] = ()) -> TypeSchema: ... @staticmethod def from_annotation(annotation: object) -> TypeSchema: ... def repr(self, ty_map: Callable[[str], str] | None = None) -> str: ... def check_value(self, value: object) -> None: ... def convert(self, value: object) -> CAny: ... def to_json(self) -> dict[str, Any]: ... class TypeField: name: str doc: str | None size: int offset: int frozen: bool metadata: dict[str, Any] getter: Any setter: Any ty: TypeSchema | None c_init: bool c_kw_only: bool c_has_default: bool c_default: Any c_default_factory: Any c_repr: bool c_compare: bool c_hash: bool c_structural_eq: str | None dataclass_field: Any | None def as_property(self, cls: type) -> property: ... class TypeMethod: name: str doc: str | None func: Any is_static: bool metadata: dict[str, Any] def as_callable(self, cls: type) -> Callable[..., Any]: ... class TypeInfo: type_cls: type | None type_index: int type_key: str type_ancestors: list[int] fields: list[TypeField] methods: list[TypeMethod] parent_type_info: TypeInfo | None def _register_fields(self, fields: list[Any]) -> None: ... def prototype_py(self) -> str: ... tvm-ffi-0.1.12/python/tvm_ffi/cpp/000077500000000000000000000000001521067262500167065ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/cpp/__init__.py000066400000000000000000000020321521067262500210140ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """C++ integration helpers for building and loading inline modules.""" from .dtype import to_cpp_dtype from .extension import build, build_inline, load, load_inline __all__ = [ "build", "build_inline", "load", "load_inline", "to_cpp_dtype", ] tvm-ffi-0.1.12/python/tvm_ffi/cpp/dtype.py000066400000000000000000000062211521067262500204060ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities for C++ dtype conversion.""" from __future__ import annotations import functools from typing import Any, Literal CPU_DTYPE_MAP = { "int8": "int8_t", "int16": "int16_t", "int32": "int32_t", "int64": "int64_t", "uint8": "uint8_t", "uint16": "uint16_t", "uint32": "uint32_t", "uint64": "uint64_t", "float32": "float", "float64": "double", "bool": "bool", } CUDA_DTYPE_MAP = { "float16": "__half", "bfloat16": "__nv_bfloat16", "float8_e4m3fn": "__nv_fp8_e4m3", # "float8_e4m3fnuz": "__nv_fp8_e4m3", "float8_e5m2": "__nv_fp8_e5m2", # "float8_e5m2fnuz": "__nv_fp8_e5m2", "float8_e8m0fnu": "__nv_fp8_e8m0", "float4_e2m1": "__nv_fp4_e2m1", "float4_e2m1fn_x2": "__nv_fp4x2_e2m1", } ROCM_DTYPE_MAP = { "float16": "__half", "bfloat16": "__hip_bfloat16", "float8_e4m3fn": "__hip_fp8_e4m3", "float8_e4m3fnuz": "__hip_fp8_e4m3_fnuz", "float8_e5m2": "__hip_fp8_e5m2", "float8_e5m2fnuz": "__hip_fp8_e5m2_fnuz", "float4_e2m1": "__hip_fp4_e2m1", "float4_e2m1fn_x2": "__hip_fp4x2_e2m1", } @functools.lru_cache(maxsize=None) def _determine_backend_once() -> Literal["cpu", "cuda", "rocm"]: try: import torch # noqa: PLC0415 import torch.version # noqa: PLC0415 if torch.cuda.is_available(): if torch.version.cuda is not None: return "cuda" elif torch.version.hip is not None: return "rocm" except ImportError: pass return "cpu" def to_cpp_dtype(dtype_str: str | Any) -> str: """Convert a dtype to its corresponding C++ dtype string. Parameters ---------- dtype_str : `str` or `torch.dtype` The dtype string or object to convert. Returns ------- str The corresponding C++ dtype string. """ if not isinstance(dtype_str, str): dtype_str = str(dtype_str) if dtype_str.startswith("torch."): dtype_str = dtype_str[6:] cpp_str = CPU_DTYPE_MAP.get(dtype_str) if cpp_str is not None: return cpp_str backend = _determine_backend_once() if backend in ("cuda", "rocm"): dtype_map = CUDA_DTYPE_MAP if backend == "cuda" else ROCM_DTYPE_MAP cpp_str = dtype_map.get(dtype_str) if cpp_str is not None: return cpp_str raise ValueError(f"Unsupported dtype string: {dtype_str} for {backend = }") tvm-ffi-0.1.12/python/tvm_ffi/cpp/extension.py000066400000000000000000001650611521067262500213050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Build and load C++/CUDA sources into a tvm_ffi Module using Ninja.""" from __future__ import annotations import functools import hashlib import logging import os import shutil import subprocess import sys from collections.abc import Mapping, Sequence from contextlib import nullcontext from pathlib import Path from typing import Any, Literal from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path, find_libtvm_ffi from tvm_ffi.module import Module, load_module from tvm_ffi.utils import FileLock IS_WINDOWS = sys.platform == "win32" BACKEND_STR = Literal["cuda", "hip"] logger = logging.getLogger(__name__) @functools.lru_cache def _detect_gpu_backend() -> BACKEND_STR: """Auto-detect whether to use CUDA or HIP (ROCm). Returns 'hip' if ROCm/HIP is available, 'cuda' otherwise. """ # Check environment variable override first backend = os.environ.get("TVM_FFI_GPU_BACKEND", "").lower() if backend in ("cuda", "hip"): return backend # type: ignore[return-value] try: _find_rocm_home() return "hip" except RuntimeError: return "cuda" def _resolve_gpu_backend(backend: str | None) -> BACKEND_STR: if backend is not None: if backend in ("cuda", "hip"): return backend # type: ignore[return-value] raise ValueError(f"Invalid backend: {backend}. Supported backends are 'cuda' and 'hip'.") return _detect_gpu_backend() def _hash_sources( cpp_source: str | None, cuda_source: str | None, cpp_files: Sequence[str] | None, cuda_files: Sequence[str] | None, functions: Sequence[str] | Mapping[str, str], extra_cflags: Sequence[str], extra_cuda_cflags: Sequence[str], extra_ldflags: Sequence[str], extra_include_paths: Sequence[str], embed_cubin: Mapping[str, bytes] | None = None, ) -> str: """Generate a unique hash for the given sources and functions.""" m = hashlib.sha256() def _hash(obj: Any) -> None: if obj is None: m.update(b"None") elif isinstance(obj, str): m.update(b"str") m.update(obj.encode("utf-8")) elif isinstance(obj, bytes): m.update(b"bytes") m.update(obj) elif isinstance(obj, Mapping): m.update(b"Mapping") for key in sorted(obj.keys()): item = obj[key] _hash(key) _hash(item) elif isinstance(obj, Sequence): m.update(b"Sequence") for item in obj: _hash(item) else: raise ValueError(f"Unsupported type: {type(obj)}") _hash( ( cpp_source, cuda_source, sorted(cpp_files) if cpp_files is not None else None, sorted(cuda_files) if cuda_files is not None else None, functions, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, embed_cubin, ) ) return m.hexdigest()[:16] def _maybe_write(path: str, content: str) -> None: """Write content to path if it does not already exist with the same content.""" p = Path(path) if p.exists(): with p.open() as f: existing_content = f.read() if existing_content == content: return with p.open("w") as f: f.write(content) @functools.lru_cache def _find_cuda_home() -> str: """Find the CUDA install path.""" # Guess #1 cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") if cuda_home is None: # Guess #2 nvcc_path = shutil.which("nvcc") if nvcc_path is not None: cuda_home = str(Path(nvcc_path).parent.parent) else: # Guess #3 if IS_WINDOWS: cuda_root = Path("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA") cuda_homes = list(cuda_root.glob("v*.*")) if len(cuda_homes) == 0: raise RuntimeError( "Could not find CUDA installation. Please set CUDA_HOME environment variable." ) cuda_home = str(cuda_homes[0]) else: cuda_home = "/usr/local/cuda" if not Path(cuda_home).exists(): raise RuntimeError( "Could not find CUDA installation. Please set CUDA_HOME environment variable." ) return cuda_home def _get_cuda_target() -> str: """Get the CUDA target architecture flag.""" if "TVM_FFI_CUDA_ARCH_LIST" in os.environ: arch_list = os.environ["TVM_FFI_CUDA_ARCH_LIST"].split() # e.g., "8.9 9.0a" flags = [] for arch in arch_list: if len(arch.split(".")) != 2: raise ValueError(f"Invalid CUDA architecture: {arch}") major, minor = arch.split(".") flags.append(f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}") return " ".join(flags) else: try: status = subprocess.run( args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], capture_output=True, check=True, ) compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] major, minor = compute_cap.split(".") return f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" except Exception: try: # For old drivers, there is no compute_cap, but we can use the GPU name to determine the architecture. ampere_arch_map = { "A100": ("8", "0"), "A10": ("8", "6"), } status = subprocess.run( args=["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, check=True, text=True, ) gpu_name = status.stdout.strip().split("\n")[0] for gpu_key, (major, minor) in ampere_arch_map.items(): if gpu_key in gpu_name: return f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" except (subprocess.CalledProcessError, FileNotFoundError): pass raise RuntimeError( "Could not detect CUDA compute_cap automatically. Please set TVM_FFI_CUDA_ARCH_LIST environment variable." ) @functools.lru_cache def _find_rocm_home() -> str: """Find the ROCm install path.""" # Guess #1: check environment variables rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") if rocm_home is None: hipcc_path = shutil.which("hipcc") # Guess #2: find hipcc in PATH and resolve ROCm home from it if hipcc_path is not None: rocm_home = str(Path(hipcc_path).resolve().parent.parent) if Path(rocm_home).name == "hip": rocm_home = str(Path(rocm_home).parent) else: # Guess #3: use default installation path rocm_home = "/opt/rocm" if not Path(rocm_home).exists(): raise RuntimeError( "Could not find ROCm installation. Please set ROCM_HOME environment variable." ) return rocm_home def _get_rocm_target() -> list[str]: """Get the ROCm target architecture flags (--offload-arch=gfxXXXX).""" if "TVM_FFI_ROCM_ARCH_LIST" in os.environ: arch_list = os.environ["TVM_FFI_ROCM_ARCH_LIST"].split() # e.g., "gfx90a gfx942" return [f"--offload-arch={arch}" for arch in arch_list] # Try rocm_agent_enumerator try: agent_enum = str(Path(_find_rocm_home()) / "bin" / "rocm_agent_enumerator") if not Path(agent_enum).exists(): agent_enum = "rocm_agent_enumerator" status = subprocess.run(args=[agent_enum], capture_output=True, check=True, text=True) archs = list( dict.fromkeys( line.strip() for line in status.stdout.strip().split("\n") if line.strip() and line.strip() != "gfx000" ) ) if archs: return [f"--offload-arch={arch}" for arch in archs] except (subprocess.CalledProcessError, FileNotFoundError): pass # Try rocminfo try: status = subprocess.run(args=["rocminfo"], capture_output=True, check=True, text=True) archs = list( dict.fromkeys( line.split(":")[-1].strip() for line in status.stdout.split("\n") if "Name:" in line and "gfx" in line.lower() and line.split(":")[-1].strip() != "gfx000" ) ) if archs: return [f"--offload-arch={arch}" for arch in archs] except (subprocess.CalledProcessError, FileNotFoundError): pass raise RuntimeError( "Could not detect ROCm GPU architecture automatically. " "Please set TVM_FFI_ROCM_ARCH_LIST environment variable (e.g. 'gfx90a gfx942')." ) def _run_command_in_dev_prompt( args: list[str], cwd: str | os.PathLike[str], capture_output: bool, ) -> subprocess.CompletedProcess: """Locates the Developer Command Prompt and runs a command within its environment.""" try: # Path to vswhere.exe vswhere_path = str( Path(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)")) / "Microsoft Visual Studio" / "Installer" / "vswhere.exe" ) if not Path(vswhere_path).exists(): raise FileNotFoundError("vswhere.exe not found.") # Find the Visual Studio installation path vs_install_path = subprocess.run( [ vswhere_path, "-latest", "-prerelease", "-products", "*", "-property", "installationPath", ], capture_output=True, text=True, check=True, ).stdout.strip() if not vs_install_path: raise FileNotFoundError("No Visual Studio installation found.") # Construct the path to the VsDevCmd.bat file vsdevcmd_path = str(Path(vs_install_path) / "Common7" / "Tools" / "VsDevCmd.bat") if not Path(vsdevcmd_path).exists(): raise FileNotFoundError(f"VsDevCmd.bat not found at: {vsdevcmd_path}") # Use cmd.exe to run the batch file and then your command. # The /k flag keeps the command prompt open after the batch file runs. # The "&" symbol chains the commands. cmd_command = '"{vsdevcmd_path}" -arch=x64 & {command}'.format( vsdevcmd_path=vsdevcmd_path, command=" ".join(args) ) # Execute the command in a new shell return subprocess.run( cmd_command, check=False, cwd=cwd, capture_output=capture_output, shell=True ) except (FileNotFoundError, subprocess.CalledProcessError) as e: raise RuntimeError( "Failed to run the following command in MSVC developer environment: {}".format( " ".join(args) ) ) from e def _generate_ninja_build( # noqa: PLR0915, PLR0912 name: str, extra_cflags: Sequence[str], extra_cuda_cflags: Sequence[str], extra_ldflags: Sequence[str], extra_include_paths: Sequence[str], sources: Sequence[str], embed_cubin: Mapping[str, bytes] | None = None, backend: str | None = None, output: str | None = None, ) -> str: """Generate the content of build.ninja for building the module.""" # Determine output format from extension if output is not None: out_ext = Path(output).suffix.lower() object_mode = out_ext in (".o", ".obj") output_name = output else: object_mode = False output_name = f"{name}{'.dll' if IS_WINDOWS else '.so'}" has_cuda_sources = any(Path(s).suffix.lower() == ".cu" for s in sources) with_hip = backend == "hip" with_cuda = backend == "cuda" with_backend = with_hip or with_cuda or has_cuda_sources if has_cuda_sources and not (with_hip or with_cuda): # Auto-detect backend from available GPU detected = _resolve_gpu_backend(None) with_hip = detected == "hip" with_cuda = detected == "cuda" default_include_paths = [find_include_path(), find_dlpack_include_path()] tvm_ffi_lib = Path(find_libtvm_ffi()) tvm_ffi_lib_path = str(tvm_ffi_lib.parent) tvm_ffi_lib_name = tvm_ffi_lib.stem if IS_WINDOWS: default_cflags = ["/O2", "/MD"] default_cxxflags = ["/std:c++17", "/MD", "/EHsc"] _win_warnings = [ "/wd4819", "/wd4251", "/wd4244", "/wd4267", "/wd4275", "/wd4018", "/wd4190", "/wd4624", "/wd4067", "/wd4068", ] default_cflags += _win_warnings default_cxxflags += _win_warnings default_cuda_cflags = ["-Xcompiler", "/std:c++17", "/O2"] default_ldflags = [ "/DLL", f"/LIBPATH:{tvm_ffi_lib_path}", f"{tvm_ffi_lib_name}.lib", ] else: default_cflags = ["-fPIC", "-O2"] default_cxxflags = ["-std=c++17", "-fPIC", "-O2"] default_cuda_cflags = ["-std=c++17", "-O2"] default_ldflags = ["-shared", f"-L{tvm_ffi_lib_path}", "-ltvm_ffi"] if with_hip: rocm_home = _find_rocm_home() default_cuda_cflags += ["-fPIC", "-D__HIP_PLATFORM_AMD__=1", "-fno-gpu-rdc"] default_cuda_cflags += _get_rocm_target() default_include_paths.append(str(Path(rocm_home) / "include")) default_ldflags += [ f"-L{Path(rocm_home) / 'lib'!s}", "-lamdhip64", ] if with_cuda: default_cuda_cflags = ["-Xcompiler", "-fPIC", *default_cuda_cflags] default_cuda_cflags += [_get_cuda_target()] default_ldflags += [ "-L{}".format(str(Path(_find_cuda_home()) / "lib64")), "-lcudart", # cuda runtime library ] extra_cflags_list = [flag.strip() for flag in extra_cflags] cflags = default_cflags + extra_cflags_list cxxflags = default_cxxflags + extra_cflags_list cuda_cflags = default_cuda_cflags + [flag.strip() for flag in extra_cuda_cflags] ldflags = default_ldflags + [flag.strip() for flag in extra_ldflags] include_paths = default_include_paths + [ str(Path(path).resolve()) for path in extra_include_paths ] # append include paths for path in include_paths: inc = "-I{}".format(path.replace(":", "$:")) cflags.append(inc) cxxflags.append(inc) cuda_cflags.append(inc) # Classify sources by extension to determine which rules are needed with_c = any(Path(s).suffix.lower() == ".c" for s in sources) ninja: list[str] = [] ninja.append("ninja_required_version = 1.3") ninja.append("cxx = {}".format(os.environ.get("CXX", "cl" if IS_WINDOWS else "c++"))) ninja.append("cxxflags = {}".format(" ".join(cxxflags))) if with_c: ninja.append("cc = {}".format(os.environ.get("CC", "cl" if IS_WINDOWS else "cc"))) ninja.append("cflags = {}".format(" ".join(cflags))) if with_backend: if with_hip: ninja.append("nvcc = {}".format(str(Path(_find_rocm_home()) / "bin" / "hipcc"))) if with_cuda: ninja.append("nvcc = {}".format(str(Path(_find_cuda_home()) / "bin" / "nvcc"))) ninja.append("cuda_cflags = {}".format(" ".join(cuda_cflags))) ninja.append("ldflags = {}".format(" ".join(ldflags))) # rules ninja.append("") ninja.append("rule compile") if IS_WINDOWS: ninja.append(" command = $cxx /showIncludes $cxxflags -c $in /Fo$out") ninja.append(" deps = msvc") else: ninja.append(" depfile = $out.d") ninja.append(" deps = gcc") ninja.append(" command = $cxx -MMD -MF $out.d $cxxflags -c $in -o $out") ninja.append("") if with_c: ninja.append("rule c_compile") if IS_WINDOWS: ninja.append(" command = $cc /showIncludes $cflags -c $in /Fo$out") ninja.append(" deps = msvc") else: ninja.append(" depfile = $out.d") ninja.append(" deps = gcc") ninja.append(" command = $cc -MMD -MF $out.d $cflags -c $in -o $out") ninja.append("") if with_backend: ninja.append("rule compile_cuda") ninja.append(" depfile = $out.d") ninja.append(" deps = gcc") if with_hip: ninja.append(" command = $nvcc $cuda_cflags -c $in -o $out") else: ninja.append( " command = $nvcc --generate-dependencies-with-compile --dependency-output $out.d $cuda_cflags -c $in -o $out" ) ninja.append("") # Add rules for object merging and cubin embedding (Unix only) if not IS_WINDOWS: ninja.append("rule merge_objects") ninja.append(" command = ld -r -o $out $in") ninja.append("") if embed_cubin: ninja.append("rule embed_cubin") ninja.append( f" command = {sys.executable} -m tvm_ffi.utils.embed_cubin --output-obj $out --input-obj $in --cubin $cubin --name $name" ) ninja.append("") if not object_mode: ninja.append("rule link") if IS_WINDOWS: ninja.append(" command = $cxx $in /link $ldflags /out:$out") else: ninja.append(" command = $cxx $in $ldflags -o $out") ninja.append("") # build targets — dispatch by file extension obj_files: list[str] = [] c_idx = cpp_idx = cuda_idx = 0 for src in sorted(sources): ext = Path(src).suffix.lower() escaped = src.replace(":", "$:") if ext in (".o", ".obj"): # Pre-compiled object file: pass directly to linker obj_files.append(escaped) elif ext == ".c": obj_name = f"c_{c_idx}.o" ninja.append(f"build {obj_name}: c_compile {escaped}") obj_files.append(obj_name) c_idx += 1 elif ext == ".cu": obj_name = f"cuda_{cuda_idx}.o" ninja.append(f"build {obj_name}: compile_cuda {escaped}") obj_files.append(obj_name) cuda_idx += 1 else: # .cc, .cpp, .cxx — default to C++ compilation obj_name = f"cpp_{cpp_idx}.o" ninja.append(f"build {obj_name}: compile {escaped}") obj_files.append(obj_name) cpp_idx += 1 if object_mode: # Object-only output: merge all object files into the target. if not IS_WINDOWS: ninja.append(f"build {output_name}: merge_objects {' '.join(obj_files)}") ninja.append("") ninja.append(f"default {output_name}") else: # Windows: no ld -r available; default to the first intermediate object ninja.append(f"default {obj_files[0]}") ninja.append("") return "\n".join(ninja) # For Unix systems with embed_cubin, use a 3-step process: # 1. Merge all object files into a unified object file # 2. Embed each cubin into the unified object file (chain them) # 3. Link the final object file into a shared library if not IS_WINDOWS and embed_cubin: # Step 1: Merge object files into unified.o unified_obj = "unified.o" obj_files_str = " ".join(obj_files) ninja.append(f"build {unified_obj}: merge_objects {obj_files_str}") ninja.append("") # Step 2: Chain embed_cubin operations for each cubin current_obj = unified_obj for cubin_name in sorted(embed_cubin.keys()): # Create next object file name next_obj = f"unified_with_{cubin_name}.o" cubin_file = f"{cubin_name}.cubin" # Add ninja build rule ninja.append(f"build {next_obj}: embed_cubin {current_obj}") ninja.append(f" cubin = {cubin_file}") ninja.append(f" name = {cubin_name}") ninja.append("") current_obj = next_obj # Step 3: Link the final object file ninja.append(f"build {output_name}: link {current_obj}") ninja.append("") else: # Directly link object files (for Windows or no cubin embedding) link_files_str = " ".join(obj_files) ninja.append(f"build {output_name}: link {link_files_str}") ninja.append("") # default target ninja.append(f"default {output_name}") ninja.append("") return "\n".join(ninja) def build_ninja(build_dir: str) -> None: """Build the module in the given build directory using ninja.""" command = ["ninja", "-v"] num_workers = os.environ.get("MAX_JOBS", None) if num_workers is not None: command += ["-j", num_workers] if IS_WINDOWS: status = _run_command_in_dev_prompt(args=command, cwd=build_dir, capture_output=True) else: status = subprocess.run(check=False, args=command, cwd=build_dir, capture_output=True) encoding = "oem" if IS_WINDOWS else "utf-8" if status.returncode != 0: msg = [f"ninja exited with status {status.returncode}"] if status.stdout: msg.append(f"stdout:\n{status.stdout.decode(encoding)}") if status.stderr: msg.append(f"stderr:\n{status.stderr.decode(encoding)}") raise RuntimeError("\n".join(msg)) LOG_BUILD = os.environ.get("TVM_FFI_CPP_EXTENSION_LOG_BUILD", "0") if LOG_BUILD in ("1", "stdout"): logger.info("ninja build stdout:\n%s", status.stdout.decode(encoding)) if LOG_BUILD in ("1", "stderr"): logger.info("ninja build stderr:\n%s", status.stderr.decode(encoding)) # Translation table for escaping C++ string literals _CPP_ESCAPE_TABLE = str.maketrans( { "\\": "\\\\", '"': '\\"', "\n": "\\n", "\r": "\\r", "\t": "\\t", } ) def _escape_cpp_string_literal(s: str) -> str: """Escape special characters for C++ string literals.""" return s.translate(_CPP_ESCAPE_TABLE) def _decorate_with_tvm_ffi(source: str, functions: Mapping[str, str]) -> str: """Decorate the given source code with TVM FFI export macros.""" sources = [ "#include ", "#include ", "#include ", "#include ", "#include ", "", source, ] for func_name, func_doc in functions.items(): sources.append(f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({func_name}, {func_name});") if func_doc: # Escape the docstring for C++ string literal escaped_doc = _escape_cpp_string_literal(func_doc) sources.append(f'TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC({func_name}, "{escaped_doc}");') sources.append("") return "\n".join(sources) def _str_seq2list(seq: Sequence[str] | str | None) -> list[str]: if seq is None: return [] elif isinstance(seq, str): return [seq] else: return list(seq) def _build_impl( # noqa: PLR0913 name: str, sources: Sequence[str] | str | None, extra_cflags: Sequence[str] | None, extra_cuda_cflags: Sequence[str] | None, extra_ldflags: Sequence[str] | None, extra_include_paths: Sequence[str] | None, build_directory: str | None, need_lock: bool = True, embed_cubin: Mapping[str, bytes] | None = None, backend: str | None = None, output: str | None = None, ) -> str: """Real implementation of build function.""" # need to resolve the path to make it unique source_path_list = [str(Path(p).resolve()) for p in _str_seq2list(sources)] assert source_path_list, "sources must be provided." has_cuda = any(Path(p).suffix.lower() == ".cu" for p in source_path_list) resolved_backend = _resolve_gpu_backend(backend) if has_cuda else None extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else [] extra_cflags_list = list(extra_cflags) if extra_cflags is not None else [] extra_cuda_cflags_list = list(extra_cuda_cflags) if extra_cuda_cflags is not None else [] extra_include_paths_list = list(extra_include_paths) if extra_include_paths is not None else [] build_dir: Path if build_directory is None: cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", str(Path("~/.cache/tvm-ffi").expanduser())) source_hash: str = _hash_sources( None, None, source_path_list, None, {}, extra_cflags_list, extra_cuda_cflags_list, extra_ldflags_list, extra_include_paths_list, embed_cubin, ) build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}" else: build_dir = Path(build_directory).resolve() build_dir.mkdir(parents=True, exist_ok=True) # CUBIN embedding is only supported on Unix systems if embed_cubin and IS_WINDOWS: raise NotImplementedError("CUBIN embedding is not yet supported on Windows") # Write CUBIN files to build directory if needed (for Unix systems) # These will be embedded using the embed_cubin utility during ninja build if embed_cubin: for cubin_name, cubin_bytes in embed_cubin.items(): cubin_path = build_dir / f"{cubin_name}.cubin" cubin_path.write_bytes(cubin_bytes) # generate build.ninja ninja_source = _generate_ninja_build( name=name, extra_cflags=extra_cflags_list, extra_cuda_cflags=extra_cuda_cflags_list, extra_ldflags=extra_ldflags_list, extra_include_paths=extra_include_paths_list, sources=source_path_list, embed_cubin=embed_cubin, backend=resolved_backend, output=output, ) # may not hold lock when build_directory is specified, prevent deadlock with FileLock(str(build_dir / "lock")) if need_lock else nullcontext(): # write build.ninja if it does not already exist _maybe_write(str(build_dir / "build.ninja"), ninja_source) # build the module build_ninja(str(build_dir)) # Determine the output filename (mirrors _generate_ninja_build logic) if output is not None: out_ext = Path(output).suffix.lower() object_mode = out_ext in (".o", ".obj") output_name = Path(output).name else: object_mode = False output_name = f"{name}{'.dll' if IS_WINDOWS else '.so'}" if object_mode and IS_WINDOWS: # Windows has no ld -r; the actual target is the first intermediate object. # The name must match _generate_ninja_build: c_0.o / cpp_0.o / cuda_0.o. first_ext = Path(sorted(source_path_list)[0]).suffix.lower() if source_path_list else "" if first_ext == ".c": obj_name = "c_0.o" elif first_ext == ".cu": obj_name = "cuda_0.o" else: obj_name = "cpp_0.o" return str((build_dir / obj_name).resolve()) return str((build_dir / output_name).resolve()) def build_inline( # noqa: PLR0913 name: str, *, cpp_sources: Sequence[str] | str | None = None, cuda_sources: Sequence[str] | str | None = None, functions: Mapping[str, str] | Sequence[str] | str | None = None, extra_cflags: Sequence[str] | None = None, extra_cuda_cflags: Sequence[str] | None = None, extra_ldflags: Sequence[str] | None = None, extra_include_paths: Sequence[str] | None = None, build_directory: str | None = None, embed_cubin: Mapping[str, bytes] | None = None, backend: str | None = None, output: str | None = None, ) -> str: """Compile and build a C++/CUDA module from inline source code. This function compiles the given C++ and/or CUDA source code into a shared library or object file. Both ``cpp_sources`` and ``cuda_sources`` are compiled to an object file. When ``output`` is ``None`` (the default) or has a shared-library extension (``.so``, ``.dll``), object files are linked into a shared library. When ``output`` has an object-file extension (``.o``, ``.obj``), linking is skipped and the path to the object file is returned directly. The ``functions`` parameter is used to specify which functions in the source code should be exported to the tvm ffi module. It can be a mapping, a sequence, or a single string. When a mapping is given, the keys are the names of the exported functions, and the values are docstrings for the functions. When a sequence of string is given, they are the function names needed to be exported, and the docstrings are set to empty strings. A single function name can also be given as a string, indicating that only one function is to be exported. Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags`` parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional flags for your specific use case. The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can include any header from tvm ffi in your source code. You can also provide additional include paths via the ``extra_include_paths`` parameter and include custom headers in your source code. The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used. The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified, the default cache directory is ``~/.cache/tvm-ffi``. Parameters ---------- name The name of the tvm ffi module. cpp_sources The C++ source code. It can be a list of sources or a single source. cuda_sources The CUDA source code. It can be a list of sources or a single source. functions The functions in cpp_sources or cuda_source that will be exported to the tvm ffi module. When a mapping is given, the keys are the names of the exported functions, and the values are docstrings for the functions (use an empty string to skip documentation for specific functions). When a sequence or a single string is given, they are the functions needed to be exported, and the docstrings are set to empty strings. A single function name can also be given as a string. When cpp_sources is given, the functions must be declared (not necessarily defined) in the cpp_sources. When cpp_sources is not given, the functions must be defined in the cuda_sources. If not specified, no function will be exported. extra_cflags The extra compiler flags for C++ compilation. The default flags are: - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2'] - On Windows: ['/std:c++17', '/O2'] extra_cuda_cflags The extra compiler flags for CUDA compilation. extra_ldflags The extra linker flags. The default flags are: - On Linux/macOS: ['-shared'] - On Windows: ['/DLL'] extra_include_paths The extra include paths. build_directory The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to specify the cache directory. embed_cubin: Mapping[str, bytes], optional A mapping from CUBIN module names to CUBIN binary data. TVM-FFI provides a macro `TVM_FFI_EMBED_CUBIN(name)` to embed CUBIN data into the compiled shared library. The keys should match the names used in `TVM_FFI_EMBED_CUBIN(name)` calls in the C++ source code. The values are the CUBIN binary data bytes. The embedded CUBIN kernels can be accessed by the macro `TVM_FFI_EMBED_CUBIN_GET_KERNEL(name, kernel_name)` defined in the `tvm/ffi/extra/cuda/cubin_launcher.h` header. See the `examples/cubin_launcher` directory for examples how to use cubin launcher to launch CUBIN kernels in TVM-FFI. backend The GPU backend to use. It can be "cuda" or "hip". If not specified, the backend will be automatically determined based on the available GPU and the provided source code. output Output filename that determines the build type from its extension. When ``None`` (the default), builds a shared library (``.so`` on Unix, ``.dll`` on Windows). Use an object-file extension (e.g., ``"hello.o"``) to skip linking and produce a relocatable object file. The file is placed in the build directory. Returns ------- path: str The path to the built shared library or object file. Example ------- .. code-block:: python import torch from tvm_ffi import Module import tvm_ffi.cpp # define the cpp source code cpp_source = ''' void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } ''' # compile the cpp source code and load the module lib_path: str = tvm_ffi.cpp.build_inline( name="hello", cpp_sources=cpp_source, functions="add_one_cpu", ) # load the module mod: Module = tvm_ffi.load_module(lib_path) # use the function from the loaded module to perform x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) """ cpp_source_list = _str_seq2list(cpp_sources) cpp_source = "\n".join(cpp_source_list) with_cpp = bool(cpp_source_list) del cpp_source_list cuda_source_list = _str_seq2list(cuda_sources) cuda_source = "\n".join(cuda_source_list) with_backend = bool(cuda_source_list) del cuda_source_list extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else [] extra_cflags_list = list(extra_cflags) if extra_cflags is not None else [] extra_cuda_cflags_list = list(extra_cuda_cflags) if extra_cuda_cflags is not None else [] extra_include_paths_list = list(extra_include_paths) if extra_include_paths is not None else [] # add function registration code to sources if functions is None: function_map: dict[str, str] = {} elif isinstance(functions, str): function_map = {functions: ""} elif isinstance(functions, Mapping): function_map = dict(functions) else: function_map = {name: "" for name in functions} if with_cpp: cpp_source = _decorate_with_tvm_ffi(cpp_source, function_map) cuda_source = _decorate_with_tvm_ffi(cuda_source, {}) else: cpp_source = _decorate_with_tvm_ffi(cpp_source, {}) cuda_source = _decorate_with_tvm_ffi(cuda_source, function_map) # determine the cache dir for the built module build_dir: Path if build_directory is None: cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", str(Path("~/.cache/tvm-ffi").expanduser())) source_hash: str = _hash_sources( cpp_source, cuda_source, None, None, function_map, extra_cflags_list, extra_cuda_cflags_list, extra_ldflags_list, extra_include_paths_list, embed_cubin, ) build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}" else: build_dir = Path(build_directory).resolve() build_dir.mkdir(parents=True, exist_ok=True) cpp_file = str((build_dir / "main.cpp").resolve()) cuda_file = str((build_dir / "cuda.cu").resolve()) with FileLock(str(build_dir / "lock")): # write source files if they do not already exist _maybe_write(cpp_file, cpp_source) if with_backend: _maybe_write(cuda_file, cuda_source) src_files = [] if with_cpp: src_files.append(cpp_file) if with_backend: src_files.append(cuda_file) return _build_impl( name=name, sources=src_files, extra_cflags=extra_cflags_list, extra_cuda_cflags=extra_cuda_cflags_list, extra_ldflags=extra_ldflags_list, extra_include_paths=extra_include_paths_list, build_directory=str(build_dir), need_lock=False, # already hold the lock embed_cubin=embed_cubin, backend=backend, output=output, ) def load_inline( # noqa: PLR0913 name: str, *, cpp_sources: Sequence[str] | str | None = None, cuda_sources: Sequence[str] | str | None = None, functions: Mapping[str, str] | Sequence[str] | str | None = None, extra_cflags: Sequence[str] | None = None, extra_cuda_cflags: Sequence[str] | None = None, extra_ldflags: Sequence[str] | None = None, extra_include_paths: Sequence[str] | None = None, build_directory: str | None = None, embed_cubin: Mapping[str, bytes] | None = None, keep_module_alive: bool = True, backend: str | None = None, ) -> Module: """Compile, build and load a C++/CUDA module from inline source code. This function compiles the given C++ and/or CUDA source code into a shared library. Both ``cpp_sources`` and ``cuda_sources`` are compiled to an object file, and then linked together into a shared library. It's possible to only provide cpp_sources or cuda_sources. The ``functions`` parameter is used to specify which functions in the source code should be exported to the tvm ffi module. It can be a mapping, a sequence, or a single string. When a mapping is given, the keys are the names of the exported functions, and the values are docstrings for the functions. When a sequence of string is given, they are the function names needed to be exported, and the docstrings are set to empty strings. A single function name can also be given as a string, indicating that only one function is to be exported. Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags`` parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional flags for your specific use case. The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can include any header from tvm ffi in your source code. You can also provide additional include paths via the ``extra_include_paths`` parameter and include custom headers in your source code. The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used. The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified, the default cache directory is ``~/.cache/tvm-ffi``. Parameters ---------- name The name of the tvm ffi module. cpp_sources The C++ source code. It can be a list of sources or a single source. cuda_sources The CUDA source code. It can be a list of sources or a single source. functions The functions in cpp_sources or cuda_source that will be exported to the tvm ffi module. When a mapping is given, the keys are the names of the exported functions, and the values are docstrings for the functions (use an empty string to skip documentation for specific functions). When a sequence or a single string is given, they are the functions needed to be exported, and the docstrings are set to empty strings. A single function name can also be given as a string. When cpp_sources is given, the functions must be declared (not necessarily defined) in the cpp_sources. When cpp_sources is not given, the functions must be defined in the cuda_sources. If not specified, no function will be exported. extra_cflags The extra compiler flags for C++ compilation. The default flags are: - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2'] - On Windows: ['/std:c++17', '/O2'] extra_cuda_cflags The extra compiler flags for CUDA compilation. extra_ldflags The extra linker flags. The default flags are: - On Linux/macOS: ['-shared'] - On Windows: ['/DLL'] extra_include_paths The extra include paths. build_directory The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to specify the cache directory. embed_cubin A mapping from CUBIN module names to CUBIN binary data. When provided, the CUBIN data will be embedded into the compiled shared library using objcopy, making it accessible via the TVM_FFI_EMBED_CUBIN macro. The keys should match the names used in TVM_FFI_EMBED_CUBIN calls in the C++ source code. keep_module_alive Whether to keep the module alive. If True, the module will be kept alive for the duration of the program until libtvm_ffi.so is unloaded. backend The GPU backend to use. It can be "cuda" or "hip". If not specified, the backend will be automatically determined based on the available GPU and the provided source code. Returns ------- mod: Module The loaded tvm ffi module. See Also -------- :py:func:`tvm_ffi.load_module` Example ------- .. code-block:: python import torch from tvm_ffi import Module import tvm_ffi.cpp # define the cpp source code cpp_source = ''' void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } ''' # compile the cpp source code and load the module mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=cpp_source, functions="add_one_cpu", ) # use the function from the loaded module to perform x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) """ return load_module( build_inline( name=name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, functions=functions, extra_cflags=extra_cflags, extra_cuda_cflags=extra_cuda_cflags, extra_ldflags=extra_ldflags, extra_include_paths=extra_include_paths, build_directory=build_directory, embed_cubin=embed_cubin, backend=backend, ), keep_module_alive=keep_module_alive, ) def build( # noqa: PLR0913 name: str, *, sources: Sequence[str] | str | None = None, cpp_files: Sequence[str] | str | None = None, cuda_files: Sequence[str] | str | None = None, extra_cflags: Sequence[str] | None = None, extra_cuda_cflags: Sequence[str] | None = None, extra_ldflags: Sequence[str] | None = None, extra_include_paths: Sequence[str] | None = None, build_directory: str | None = None, backend: str | None = None, output: str | None = None, ) -> str: """Compile and build a C/C++/CUDA module from source files. This function compiles the given C, C++, and/or CUDA source files into a shared library or object file. The compiler is selected automatically based on file extension: - ``.c`` — compiled with the C compiler (``$CC``) - ``.cc``, ``.cpp``, ``.cxx`` — compiled with the C++ compiler (``$CXX``) - ``.o``, ``.obj`` — pre-compiled objects, passed directly to the linker When ``output`` is ``None`` (the default) or has a shared-library extension, object files are linked into a shared library. When ``output`` has an object-file extension (``.o``, ``.obj``), linking is skipped and the path to the object file is returned. Note that this function does not automatically export functions to the tvm ffi module. You need to manually use the TVM FFI export macros (e.g., ``TVM_FFI_DLL_EXPORT_TYPED_FUNC``) in your source files to export functions. This gives you more control over which functions are exported and how they are exported. Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags`` parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional flags for your specific use case. The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can include any header from tvm ffi in your source files. You can also provide additional include paths via the ``extra_include_paths`` parameter and include custom headers in your source code. The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used. The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified, the default cache directory is ``~/.cache/tvm-ffi``. The C compiler is controlled by the ``$CC`` environment variable (default: ``cc`` on Unix, ``cl`` on Windows). The C++ compiler is controlled by the ``$CXX`` environment variable (default: ``c++`` on Unix, ``cl`` on Windows). Parameters ---------- name The name of the tvm ffi module. sources Source files to compile. The compiler is auto-detected from the file extension: - ``.c`` → C compiler (``$CC``) - ``.cc``, ``.cpp``, ``.cxx`` → C++ compiler (``$CXX``) - ``.cu`` → CUDA/HIP compiler (``nvcc`` or ``hipcc``) - ``.o``, ``.obj`` → pre-compiled objects, passed directly to the linker It can be a list of file paths or a single file path. cpp_files Alias for ``sources``, kept for backward compatibility. cuda_files Alias for ``sources``, kept for backward compatibility. extra_cflags Extra compiler flags applied to both C and C++ compilation. The C++ default flags are: - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2'] - On Windows: ['/std:c++17', '/MD', '/O2'] The C default flags omit ``-std=c++17`` and ``/EHsc``. extra_cuda_cflags The extra compiler flags for CUDA compilation. The default flags are: - ['-Xcompiler', '-fPIC', '-std=c++17', '-O2'] (Linux/macOS) - ['-Xcompiler', '/std:c++17', '/O2'] (Windows) extra_ldflags The extra linker flags. The default flags are: - On Linux/macOS: ['-shared', '-L', '-ltvm_ffi'] - On Windows: ['/DLL', '/LIBPATH:', '.lib'] extra_include_paths The extra include paths for header files. Both absolute and relative paths are supported. build_directory The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to specify the cache directory. backend The GPU backend to use. It can be "cuda" or "hip". If not specified, the backend will be automatically determined based on the available GPU and the provided source code. output Output filename that determines the build type from its extension. When ``None`` (the default), builds a shared library (``.so`` on Unix, ``.dll`` on Windows). Use an object-file extension (e.g., ``"my_ops.o"``) to skip linking and produce a relocatable object file. The file is placed in the build directory. Returns ------- path: str The path to the built shared library or object file. Example ------- .. code-block:: python import torch from tvm_ffi import Module import tvm_ffi.cpp # Assume we have a C++ source file "my_ops.cpp" with the following content: # ```cpp # #include # #include # #include # #include # #include # # void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { # TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; # DLDataType f32_dtype{kDLFloat, 32, 1}; # TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; # TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; # TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; # TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; # for (int i = 0; i < x.size(0); ++i) { # static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; # } # } # # TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu); # ``` # compile the cpp source file and get the library path lib_path: str = tvm_ffi.cpp.build( name="my_ops", sources="my_ops.cpp", ) # load the module mod: Module = tvm_ffi.load_module(lib_path) # use the function from the loaded module x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) """ # Merge sources, cpp_files, and cuda_files (backward compat aliases) merged = _str_seq2list(sources) + _str_seq2list(cpp_files) + _str_seq2list(cuda_files) return _build_impl( name=name, sources=merged or None, extra_cflags=extra_cflags, extra_cuda_cflags=extra_cuda_cflags, extra_ldflags=extra_ldflags, extra_include_paths=extra_include_paths, build_directory=build_directory, need_lock=True, backend=backend, output=output, ) def load( # noqa: PLR0913 name: str, *, sources: Sequence[str] | str | None = None, cpp_files: Sequence[str] | str | None = None, cuda_files: Sequence[str] | str | None = None, extra_cflags: Sequence[str] | None = None, extra_cuda_cflags: Sequence[str] | None = None, extra_ldflags: Sequence[str] | None = None, extra_include_paths: Sequence[str] | None = None, build_directory: str | None = None, keep_module_alive: bool = True, backend: str | None = None, ) -> Module: """Compile, build and load a C/C++/CUDA module from source files. This function compiles the given source files into a shared library and loads it as a tvm ffi module. The compiler is selected automatically based on file extension. Note that this function does not automatically export functions to the tvm ffi module. You need to manually use the TVM FFI export macros (e.g., :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC`) in your source files to export functions. This gives you more control over which functions are exported and how they are exported. Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags`` parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional flags for your specific use case. The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can include any header from tvm ffi in your source files. You can also provide additional include paths via the ``extra_include_paths`` parameter and include custom headers in your source code. The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used. The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified, the default cache directory is ``~/.cache/tvm-ffi``. Parameters ---------- name The name of the tvm ffi module. sources Source files to compile. The compiler is auto-detected from the file extension: ``.c`` → C, ``.cc``/``.cpp``/``.cxx`` → C++, ``.cu`` → CUDA/HIP, ``.o``/``.obj`` → linker passthrough. It can be a list of file paths or a single file path. cpp_files Alias for ``sources``, kept for backward compatibility. cuda_files Alias for ``sources``, kept for backward compatibility. extra_cflags The extra compiler flags for C++ compilation. The default flags are: - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2'] - On Windows: ['/std:c++17', '/MD', '/O2'] extra_cuda_cflags The extra compiler flags for CUDA compilation. The default flags are: - ['-Xcompiler', '-fPIC', '-std=c++17', '-O2'] (Linux/macOS) - ['-Xcompiler', '/std:c++17', '/O2'] (Windows) extra_ldflags The extra linker flags. The default flags are: - On Linux/macOS: ['-shared', '-L', '-ltvm_ffi'] - On Windows: ['/DLL', '/LIBPATH:', '.lib'] extra_include_paths The extra include paths for header files. Both absolute and relative paths are supported. build_directory The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to specify the cache directory. keep_module_alive Whether to keep the module alive. If True, the module will be kept alive for the duration of the program until libtvm_ffi.so is unloaded. backend The GPU backend to use. It can be "cuda" or "hip". If not specified, the backend will be automatically determined based on the available GPU and the provided source code. Returns ------- mod: Module The loaded tvm ffi module. See Also -------- :py:func:`tvm_ffi.load_module` Example ------- .. code-block:: python import torch from tvm_ffi import Module import tvm_ffi.cpp # Assume we have a C++ source file "my_ops.cpp" with the following content: # ```cpp # #include # #include # #include # #include # #include # # void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { # TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; # DLDataType f32_dtype{kDLFloat, 32, 1}; # TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; # TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; # TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; # TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; # for (int i = 0; i < x.size(0); ++i) { # static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; # } # } # # TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu); # ``` # compile the cpp source file and load the module mod: Module = tvm_ffi.cpp.load( name="my_ops", sources="my_ops.cpp", ) # use the function from the loaded module x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32) y = torch.empty_like(x) mod.add_one_cpu(x, y) torch.testing.assert_close(x + 1, y) """ return load_module( build( name=name, sources=sources, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=extra_cflags, extra_cuda_cflags=extra_cuda_cflags, extra_ldflags=extra_ldflags, extra_include_paths=extra_include_paths, build_directory=build_directory, backend=backend, ), keep_module_alive=keep_module_alive, ) tvm-ffi-0.1.12/python/tvm_ffi/cpp/nvrtc.py000066400000000000000000000140271521067262500204200ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """NVRTC (NVIDIA Runtime Compilation) utilities for compiling CUDA source to CUBIN.""" from __future__ import annotations from typing import Sequence def nvrtc_compile( # noqa: PLR0912, PLR0915 source: str, *, name: str = "kernel.cu", arch: str | None = None, extra_opts: Sequence[str] | None = None, ) -> bytes: """Compile CUDA source code to CUBIN using NVRTC. This function uses the NVIDIA Runtime Compilation (NVRTC) library to compile CUDA C++ source code into a CUBIN binary that can be loaded and executed using the CUDA Driver API. Parameters ---------- source : str The CUDA C++ source code to compile. name : str, optional The name to use for the source file (for error messages). Default: "kernel.cu" arch : str, optional The target GPU architecture (e.g., "sm_75", "sm_80", "sm_89"). If not specified, attempts to auto-detect from the current GPU. extra_opts : Sequence[str], optional Additional compilation options to pass to NVRTC (e.g., ["-I/path/to/include", "-DDEFINE=1"]). Returns ------- bytes The compiled CUBIN binary data. Raises ------ RuntimeError If NVRTC compilation fails or CUDA bindings are not available. Example ------- .. code-block:: python from tvm_ffi.cpp import nvrtc cuda_source = ''' extern "C" __global__ void add_one(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } ''' cubin_bytes = nvrtc.nvrtc_compile(cuda_source) # Use cubin_bytes with tvm_ffi.cpp.load_inline and embed_cubin parameter """ try: from cuda.bindings import driver, nvrtc # noqa: PLC0415 except ImportError as e: raise RuntimeError( "CUDA bindings not available. Install with: pip install cuda-python" ) from e # Auto-detect architecture if not specified if arch is None: try: # Initialize CUDA driver API (result,) = driver.cuInit(0) if result != driver.CUresult.CUDA_SUCCESS: raise RuntimeError(f"Failed to initialize CUDA driver: {result}") # Get current device result, device = driver.cuCtxGetDevice() if result != driver.CUresult.CUDA_SUCCESS: # Try to get device 0 if no context exists device = 0 # Get compute capability result, major = driver.cuDeviceGetAttribute( driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device ) if result != driver.CUresult.CUDA_SUCCESS: raise RuntimeError(f"Failed to get compute capability major: {result}") result, minor = driver.cuDeviceGetAttribute( driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device ) if result != driver.CUresult.CUDA_SUCCESS: raise RuntimeError(f"Failed to get compute capability minor: {result}") arch = f"sm_{major}{minor}" except Exception as e: # Fallback to a reasonable default raise RuntimeError( f"Failed to auto-detect GPU architecture: {e}. " "Please specify 'arch' parameter explicitly." ) from e # Create program result, prog = nvrtc.nvrtcCreateProgram(str.encode(source), str.encode(name), 0, None, None) if result != nvrtc.nvrtcResult.NVRTC_SUCCESS: raise RuntimeError(f"Failed to create NVRTC program: {result}") # Compile options opts = [ b"--gpu-architecture=" + arch.encode(), b"-default-device", ] # Add extra options if provided if extra_opts: opts.extend([opt.encode() if isinstance(opt, str) else opt for opt in extra_opts]) # Compile (result,) = nvrtc.nvrtcCompileProgram(prog, len(opts), opts) if result != nvrtc.nvrtcResult.NVRTC_SUCCESS: # Get compilation log result_log, log_size = nvrtc.nvrtcGetProgramLogSize(prog) if result_log == nvrtc.nvrtcResult.NVRTC_SUCCESS and log_size > 0: log_buf = b" " * log_size (result_log,) = nvrtc.nvrtcGetProgramLog(prog, log_buf) if result_log == nvrtc.nvrtcResult.NVRTC_SUCCESS: error_msg = f"NVRTC compilation failed:\n{log_buf.decode('utf-8')}" else: error_msg = f"NVRTC compilation failed (couldn't get log): {result}" else: error_msg = f"NVRTC compilation failed: {result}" nvrtc.nvrtcDestroyProgram(prog) raise RuntimeError(error_msg) # Get CUBIN result, cubin_size = nvrtc.nvrtcGetCUBINSize(prog) if result != nvrtc.nvrtcResult.NVRTC_SUCCESS: nvrtc.nvrtcDestroyProgram(prog) raise RuntimeError(f"Failed to get CUBIN size from NVRTC: {result}") cubin_buf = b" " * cubin_size (result,) = nvrtc.nvrtcGetCUBIN(prog, cubin_buf) if result != nvrtc.nvrtcResult.NVRTC_SUCCESS: nvrtc.nvrtcDestroyProgram(prog) raise RuntimeError(f"Failed to get CUBIN from NVRTC: {result}") # Clean up nvrtc.nvrtcDestroyProgram(prog) return cubin_buf tvm-ffi-0.1.12/python/tvm_ffi/cython/000077500000000000000000000000001521067262500174305ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/cython/base.pxi000066400000000000000000000504011521067262500210640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import ctypes from libc.stdint cimport int32_t, int64_t, uint64_t, uint32_t, uint8_t, int16_t from libc.string cimport memcpy from libcpp.vector cimport vector from cpython.bytes cimport PyBytes_AsStringAndSize, PyBytes_FromStringAndSize, PyBytes_AsString from cpython cimport Py_INCREF, Py_DECREF, Py_REFCNT from cpython cimport PyErr_CheckSignals, PyGILState_Ensure, PyGILState_Release, PyObject from cpython cimport pycapsule, PyCapsule_Destructor from cpython cimport PyErr_SetNone cdef extern from "dlpack/dlpack.h": int DLPACK_MAJOR_VERSION int DLPACK_MINOR_VERSION cdef enum: kDLCPU = 1, kDLCUDA = 2, kDLCUDAHost = 3, kDLOpenCL = 4, kDLVulkan = 7, kDLMetal = 8, kDLVPI = 9, kDLROCM = 10, kDLROCMHost = 11, kDLExtDev = 12, kDLCUDAManaged = 13, kDLOneAPI = 14, kDLWebGPU = 15, kDLHexagon = 16, kDLMAIA = 17 kDLTrn = 18 ctypedef struct DLDataType: uint8_t code uint8_t bits int16_t lanes ctypedef struct DLDevice: int device_type int device_id ctypedef struct DLTensor: void* data DLDevice device int ndim DLDataType dtype int64_t* shape int64_t* strides uint64_t byte_offset ctypedef struct DLPackVersion: uint32_t major uint32_t minor ctypedef struct DLManagedTensor: DLTensor dl_tensor void* manager_ctx void (*deleter)(DLManagedTensor* self) ctypedef struct DLManagedTensorVersioned: DLPackVersion version DLTensor dl_tensor void* manager_ctx void (*deleter)(DLManagedTensorVersioned* self) uint64_t flags # DLPack Exchange API function pointer types ctypedef int (*DLPackManagedTensorAllocator)( DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx, void (*SetError)(void* error_ctx, const char* kind, const char* message) ) noexcept ctypedef int (*DLPackManagedTensorFromPyObjectNoSync)( void* py_object, DLManagedTensorVersioned** out ) except -1 ctypedef int (*DLPackManagedTensorToPyObjectNoSync)( DLManagedTensorVersioned* tensor, void** out_py_object ) except -1 ctypedef int (*DLPackCurrentWorkStream)( int device_type, int32_t device_id, void** out_current_stream ) except -1 ctypedef int (*DLPackDLTensorFromPyObjectNoSync)( void* py_object, DLTensor* out ) except -1 ctypedef struct DLPackExchangeAPIHeader: DLPackVersion version DLPackExchangeAPIHeader* prev_api ctypedef struct DLPackExchangeAPI: DLPackExchangeAPIHeader header DLPackManagedTensorAllocator managed_tensor_allocator DLPackManagedTensorFromPyObjectNoSync managed_tensor_from_py_object_no_sync DLPackManagedTensorToPyObjectNoSync managed_tensor_to_py_object_no_sync DLPackDLTensorFromPyObjectNoSync dltensor_from_py_object_no_sync DLPackCurrentWorkStream current_work_stream # Cython binding for TVM FFI C API cdef extern from "tvm/ffi/c_api.h": cdef enum TVMFFITypeIndex: kTVMFFIAny = -1 kTVMFFINone = 0 kTVMFFIInt = 1 kTVMFFIBool = 2 kTVMFFIFloat = 3 kTVMFFIOpaquePtr = 4 kTVMFFIDataType = 5 kTVMFFIDevice = 6 kTVMFFIDLTensorPtr = 7 kTVMFFIRawStr = 8 kTVMFFIByteArrayPtr = 9 kTVMFFIObjectRValueRef = 10 kTVMFFISmallStr = 11 kTVMFFISmallBytes = 12 kTVMFFIStaticObjectBegin = 64 kTVMFFIObject = 64 kTVMFFIStr = 65 kTVMFFIBytes = 66 kTVMFFIError = 67 kTVMFFIFunction = 68 kTVMFFIShape = 69 kTVMFFITensor = 70 kTVMFFIArray = 71 kTVMFFIMap = 72 kTVMFFIModule = 73 kTVMFFIOpaquePyObject = 74 kTVMFFIList = 75 kTVMFFIDict = 76 ctypedef void* TVMFFIObjectHandle ctypedef struct TVMFFIObject: uint64_t combined_ref_count int32_t type_index uint32_t __padding void (*deleter)(void* self, int flags) ctypedef struct TVMFFIAny: int32_t type_index int32_t zero_padding int64_t v_int64 double v_float64 void* v_ptr TVMFFIObject* v_obj const char* v_c_str DLDataType v_dtype DLDevice v_device ctypedef struct TVMFFIByteArray: const char* data size_t size ctypedef struct TVMFFIOpaqueObjectCell: void* handle ctypedef struct TVMFFIShapeCell: const int64_t* data size_t size ctypedef enum TVMFFIBacktraceUpdateMode: kTVMFFIBacktraceUpdateModeReplace = 0 kTVMFFIBacktraceUpdateModeAppend = 1 ctypedef struct TVMFFIErrorCell: TVMFFIByteArray kind TVMFFIByteArray message TVMFFIByteArray backtrace void (*update_backtrace)( TVMFFIObjectHandle self, const TVMFFIByteArray* backtrace, int32_t update_mode ) TVMFFIObjectHandle cause_chain TVMFFIObjectHandle extra_context ctypedef int (*TVMFFISafeCallType)( void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) noexcept cdef enum TVMFFIFieldFlagBitMask: kTVMFFIFieldFlagBitMaskWritable = 1 << 0 kTVMFFIFieldFlagBitMaskHasDefault = 1 << 1 kTVMFFIFieldFlagBitMaskIsStaticMethod = 1 << 2 kTVMFFIFieldFlagBitMaskSEqHashIgnore = 1 << 3 kTVMFFIFieldFlagBitMaskSEqHashDefRecursive = 1 << 4 kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5 kTVMFFIFieldFlagBitMaskReprOff = 1 << 6 kTVMFFIFieldFlagBitMaskCompareOff = 1 << 7 kTVMFFIFieldFlagBitMaskHashOff = 1 << 8 kTVMFFIFieldFlagBitMaskInitOff = 1 << 9 kTVMFFIFieldFlagBitMaskKwOnly = 1 << 10 kTVMFFIFieldFlagBitSetterIsFunctionObj = 1 << 11 kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive = 1 << 12 ctypedef int (*TVMFFIFieldGetter)(void* field, TVMFFIAny* result) noexcept ctypedef int (*TVMFFIFieldSetter)(void* field, const TVMFFIAny* value) noexcept ctypedef int (*TVMFFIObjectCreator)(TVMFFIObjectHandle* result) noexcept ctypedef struct TVMFFIFieldInfo: TVMFFIByteArray name TVMFFIByteArray doc TVMFFIByteArray metadata int64_t flags int64_t size int64_t alignment int64_t offset TVMFFIFieldGetter getter void* setter TVMFFIAny default_value_or_factory int32_t field_static_type_index ctypedef struct TVMFFIMethodInfo: TVMFFIByteArray name TVMFFIByteArray doc TVMFFIByteArray metadata int64_t flags TVMFFIAny method cdef enum TVMFFISEqHashKind: kTVMFFISEqHashKindUnsupported = 0 kTVMFFISEqHashKindTreeNode = 1 kTVMFFISEqHashKindFreeVar = 2 kTVMFFISEqHashKindDAGNode = 3 kTVMFFISEqHashKindConstTreeNode = 4 kTVMFFISEqHashKindUniqueInstance = 5 cdef enum TVMFFIDefRegionKind: kTVMFFIDefRegionKindNone = 0 kTVMFFIDefRegionKindRecursive = 1 kTVMFFIDefRegionKindNonRecursive = 2 ctypedef struct TVMFFITypeMetadata: TVMFFIByteArray doc TVMFFIObjectCreator creator int32_t total_size TVMFFISEqHashKind structural_eq_hash_kind ctypedef struct TVMFFITypeInfo: int32_t type_index int32_t type_depth TVMFFIByteArray type_key const TVMFFITypeInfo** type_ancestors uint64_t type_key_hash int32_t num_fields int32_t num_methods const TVMFFIFieldInfo* fields const TVMFFIMethodInfo* methods const TVMFFITypeMetadata* metadata ctypedef struct TVMFFITypeAttrColumn: const TVMFFIAny* data int32_t size int32_t begin_index int TVMFFIObjectDecRef(TVMFFIObjectHandle obj) nogil int TVMFFIObjectIncRef(TVMFFIObjectHandle obj) nogil int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index, void (*deleter)(void*), TVMFFIObjectHandle* out) nogil int TVMFFIObjectGetTypeIndex(TVMFFIObjectHandle obj) nogil int TVMFFIFunctionCall(TVMFFIObjectHandle func, TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) nogil int TVMFFIFunctionCreate(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void*), TVMFFIObjectHandle* out) nogil int TVMFFIAnyViewToOwnedAny(const TVMFFIAny* any_view, TVMFFIAny* out) nogil int TVMFFIFunctionSetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle f, int override) nogil int TVMFFIFunctionGetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle* out) nogil void TVMFFIErrorMoveFromRaised(TVMFFIObjectHandle* result) nogil void TVMFFIErrorSetRaised(TVMFFIObjectHandle error) nogil int TVMFFIErrorCreate(TVMFFIByteArray* kind, TVMFFIByteArray* message, TVMFFIByteArray* backtrace, TVMFFIObjectHandle* out) nogil int TVMFFITypeKeyToIndex(TVMFFIByteArray* type_key, int32_t* out_tindex) nogil int TVMFFIStringFromByteArray(TVMFFIByteArray* input_, TVMFFIAny* out) nogil int TVMFFIBytesFromByteArray(TVMFFIByteArray* input_, TVMFFIAny* out) nogil int TVMFFIDataTypeFromString(TVMFFIByteArray* str, DLDataType* out) nogil int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIAny* out) nogil const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const char* func, int cross_ffi_boundary) nogil int TVMFFITensorFromDLPack(DLManagedTensor* src, int32_t require_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out) nogil int TVMFFITensorFromDLPackVersioned(DLManagedTensorVersioned* src, int32_t require_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out) nogil int TVMFFITensorToDLPack(TVMFFIObjectHandle src, DLManagedTensor** out) nogil int TVMFFITensorToDLPackVersioned(TVMFFIObjectHandle src, DLManagedTensorVersioned** out) nogil const TVMFFITypeInfo* TVMFFIGetTypeInfo(int32_t type_index) nogil TVMFFIByteArray TVMFFISmallBytesGetContentByteArray(const TVMFFIAny* value) nogil TVMFFIByteArray* TVMFFIBytesGetByteArrayPtr(TVMFFIObjectHandle obj) nogil TVMFFIErrorCell* TVMFFIErrorGetCellPtr(TVMFFIObjectHandle obj) nogil TVMFFIOpaqueObjectCell* TVMFFIOpaqueObjectGetCellPtr(TVMFFIObjectHandle obj) nogil TVMFFIShapeCell* TVMFFIShapeGetCellPtr(TVMFFIObjectHandle obj) nogil DLTensor* TVMFFITensorGetDLTensorPtr(TVMFFIObjectHandle obj) nogil DLDevice TVMFFIDLDeviceFromIntPair(int32_t device_type, int32_t device_id) nogil const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray* attr_name) nogil int32_t TVMFFITypeGetOrAllocIndex( const TVMFFIByteArray* type_key, int32_t static_type_index, int32_t type_depth, int32_t num_child_slots, int32_t child_slots_can_overflow, int32_t parent_type_index ) nogil int TVMFFITypeRegisterField(int32_t type_index, const TVMFFIFieldInfo* info) nogil int TVMFFITypeRegisterMethod(int32_t type_index, const TVMFFIMethodInfo* info) nogil int TVMFFITypeRegisterMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata) nogil int TVMFFITypeRegisterAttr(int32_t type_index, const TVMFFIByteArray* attr_name, const TVMFFIAny* attr_value) nogil void TVMFFIErrorSetRaisedFromCStr(const char* kind, const char* message) nogil cdef extern from "tvm/ffi/extra/c_env_api.h": ctypedef void* TVMFFIStreamHandle int TVMFFIEnvRegisterCAPI(const char* name, void* ptr) nogil void* TVMFFIEnvGetStream(int32_t device_type, int32_t device_id) nogil int TVMFFIEnvSetStream(int32_t device_type, int32_t device_id, TVMFFIStreamHandle stream, TVMFFIStreamHandle* opt_out_original_stream) nogil def _env_set_current_stream(int device_type, int device_id, uint64_t stream): cdef TVMFFIStreamHandle prev_stream = NULL CHECK_CALL(TVMFFIEnvSetStream( device_type, device_id, stream, &prev_stream)) return prev_stream def _env_get_current_stream(int device_type, int device_id): cdef void* current_stream current_stream = TVMFFIEnvGetStream(device_type, device_id) return current_stream cdef extern from "tvm_ffi_python_helpers.h": # no need to expose fields of the call context setter data structure ctypedef struct TVMFFIPyCallContext: int device_type int device_id TVMFFIStreamHandle stream const DLPackExchangeAPI* dlpack_c_exchange_api ctypedef struct TVMFFIPyArgSetter: int (*func)(TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out) except -1 const DLPackExchangeAPI* dlpack_c_exchange_api # The main call function int TVMFFIPyFuncCall( void* chandle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, int release_gil, const DLPackExchangeAPI** out_ctx_dlpack_api ) except -1 int TVMFFIPyConstructorCall( void* chandle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, TVMFFIPyCallContext* parent_ctx ) except -1 int TVMFFIPyCallFieldSetter( void* field_setter, int64_t field_flags, void* field_ptr, PyObject* py_arg, int* c_api_ret_code ) except -1 int TVMFFIPyPyObjectToFFIAny( PyObject* py_arg, TVMFFIAny* out, int* c_api_ret_code ) except -1 int TVMFFIPySetArgumentGenericDispatcher( TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1 size_t TVMFFIPyGetArgDispatchMapSize() noexcept void TVMFFIPyPushTempFFIObject(TVMFFIPyCallContext* ctx, TVMFFIObjectHandle arg) noexcept void TVMFFIPyPushTempPyObject(TVMFFIPyCallContext* ctx, PyObject* arg) noexcept void TVMFFIPyPushExtraTempPyObject(TVMFFIPyCallContext* ctx, PyObject* arg) # the predefined setters for common POD types int TVMFFIPyArgSetterFloat_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1 int TVMFFIPyArgSetterInt_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1 int TVMFFIPyArgSetterBool_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1 int TVMFFIPyArgSetterNone_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1 # Callback arg setter types — view-based AnyView -> PyObject conversion # used by the C++ -> Python callback path (PyCallback). ctypedef int (*TVMFFIPyCallbackArgSetterCallback)(TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out) except -1 ctypedef struct TVMFFIPyCallbackArgSetter: TVMFFIPyCallbackArgSetterCallback func # Built-in C callback arg setters for POD types. int TVMFFIPyCallbackArgSetterNone_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny*, PyObject** out) except -1 int TVMFFIPyCallbackArgSetterBool_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) except -1 int TVMFFIPyCallbackArgSetterInt_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) except -1 int TVMFFIPyCallbackArgSetterFloat_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) except -1 int TVMFFIPyCallbackArgSetterSmallStr_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) except -1 int TVMFFIPyCallbackArgSetterSmallBytes_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) except -1 int TVMFFIPyCallback(void* context, const TVMFFIAny* packed_args, int32_t num_args, TVMFFIAny* result) noexcept # Closure + convert helper for the PyCallback path. Returns the raw FFI rc; # callers use CHECK_CALL to translate the TLS FFI error into a Python exception. int TVMFFIPyConvertPyCallback(PyObject* callable, const DLPackExchangeAPI* dlpack_api, TVMFFIObjectHandle* out_handle) noexcept # MLIRPackedSafeCall void* TVMFFIPyMLIRPackedSafeCallCreate(void (*mlir_packed_safe_call)(void**) noexcept, PyObject* keep_alive_object) int TVMFFIPyMLIRPackedSafeCallInvoke(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* rv) void TVMFFIPyMLIRPackedSafeCallDeleter(void* self) # deleter for python objects void TVMFFIPyObjectDeleter(void* py_obj) noexcept nogil # dummy target for testing int TVMFFITestingDummyTarget() nogil cdef class ByteArrayArg: __slots__ = () cdef TVMFFIByteArray cdata cdef object py_data def __cinit__(self, py_data): if isinstance(py_data, bytearray): py_data = bytes(py_data) cdef char* data cdef Py_ssize_t size self.py_data = py_data PyBytes_AsStringAndSize(py_data, &data, &size) self.cdata.data = data self.cdata.size = size cdef inline TVMFFIByteArray* cptr(self): return &self.cdata cdef inline py_str(const char* x): """Convert a c_char_p to a python string Parameters ---------- x : c_char_p A char pointer that can be passed to C API """ return x.decode("utf-8") cdef inline str bytearray_to_str(const TVMFFIByteArray* x): return PyBytes_FromStringAndSize(x.data, x.size).decode("utf-8") cdef inline bytes bytearray_to_bytes(const TVMFFIByteArray* x): return PyBytes_FromStringAndSize(x.data, x.size) cdef inline c_str(pystr): """Create ctypes char * from a python string Parameters ---------- string : string type python string Returns ------- str : c_char_p A char pointer that can be passed to C API """ return pystr.encode("utf-8") cdef inline object ctypes_handle(void* chandle): """Cast C handle to ctypes handle.""" return ctypes.cast(chandle, ctypes.c_void_p) cdef inline void* c_handle(object handle): """Cast C types handle to c handle.""" cdef unsigned long long v_ptr cdef object value = handle.value v_ptr = 0 if value is None else value return (v_ptr) cdef _init_env_api(): # Initialize env api for signal handling # Also registers the gil state release and ensure as PyErr_CheckSignals # function is called with gil released and we need to regrab the gil CHECK_CALL(TVMFFIEnvRegisterCAPI(c_str("PyErr_CheckSignals"), PyErr_CheckSignals)) CHECK_CALL(TVMFFIEnvRegisterCAPI(c_str("PyGILState_Ensure"), PyGILState_Ensure)) CHECK_CALL(TVMFFIEnvRegisterCAPI(c_str("PyGILState_Release"), PyGILState_Release)) _init_env_api() # ensure testing is linked and we can run testcases TVMFFITestingDummyTarget() tvm-ffi-0.1.12/python/tvm_ffi/cython/core.pyx000066400000000000000000000034051521067262500211240ustar00rootroot00000000000000# cython: freethreading_compatible = True # cython: language_level=3 # cython: annotation_typing=False # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # N.B. Make sure `_register_object_by_index` is called in inheritance order, # where the base class has to be registered before the derived class. # Otherwise, `TypeInfo.parent_type_info` may not be properly propagated to the derived class. include "./base.pxi" include "./type_info.pxi" include "./object.pxi" _register_object_by_index(kTVMFFIObject, Object) include "./error.pxi" _register_object_by_index(kTVMFFIError, Error) include "./dtype.pxi" _register_object_by_index(kTVMFFIDataType, DataType) include "./device.pxi" _register_object_by_index(kTVMFFIDevice, Device) include "./string.pxi" _register_object_by_index(kTVMFFIStr, String) _register_object_by_index(kTVMFFIBytes, Bytes) include "./tensor.pxi" _register_object_by_index(kTVMFFITensor, Tensor) include "./function.pxi" _register_object_by_index(kTVMFFIFunction, Function) include "./pycallback.pxi" include "./pyclass_type_converter.pxi" tvm-ffi-0.1.12/python/tvm_ffi/cython/device.pxi000066400000000000000000000160261521067262500214160ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from enum import IntEnum from numbers import Integral from typing import Any, Optional _CLASS_DEVICE = None def _set_class_device(cls): global _CLASS_DEVICE _CLASS_DEVICE = cls def _create_device_from_tuple(cls, device_type, device_id): cdef DLDevice cdevice = TVMFFIDLDeviceFromIntPair(device_type, device_id) ret = cls.__new__(cls) (ret).cdevice = cdevice return ret class DLDeviceType(IntEnum): """Enumeration mirroring DLPack's `DLDeviceType `_ Values can be compared against :py:meth:`Device.dlpack_device_type`. Examples -------- .. code-block:: python import tvm_ffi dev = tvm_ffi.device("cuda", 0) assert dev.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCUDA """ kDLCPU = 1 kDLCUDA = 2 kDLCUDAHost = 3 kDLOpenCL = 4 kDLVulkan = 7 kDLMetal = 8 kDLVPI = 9 kDLROCM = 10 kDLROCMHost = 11 kDLExtDev = 12 kDLCUDAManaged = 13 kDLOneAPI = 14 kDLWebGPU = 15 kDLHexagon = 16 kDLMAIA = 17 kDLTrn = 18 cdef class Device: """A device descriptor used by TVM FFI and DLPack. A :class:`Device` identifies a placement (e.g. CPU, CUDA GPU) and a device index within that placement. Most users construct devices using :func:`tvm_ffi.device`. Parameters ---------- device_type : Union[str, int] A device type name (e.g. ``"cpu"``, ``"cuda"``) or a DLPack device type code. index : int, optional Zero-based device index (defaults to ``0`` when omitted). Examples -------- .. code-block:: python import tvm_ffi dev = tvm_ffi.device("cuda:0") assert dev.type == "cuda" assert dev.index == 0 assert str(dev) == "cuda:0" """ __slots__ = () cdef DLDevice cdevice _DEVICE_TYPE_TO_NAME = { DLDeviceType.kDLCPU: "cpu", DLDeviceType.kDLCUDA: "cuda", DLDeviceType.kDLCUDAHost: "cuda_host", DLDeviceType.kDLCUDAManaged: "cuda_managed", DLDeviceType.kDLOpenCL: "opencl", DLDeviceType.kDLVulkan: "vulkan", DLDeviceType.kDLMetal: "metal", DLDeviceType.kDLVPI: "vpi", DLDeviceType.kDLROCM: "rocm", DLDeviceType.kDLROCMHost: "rocm_host", DLDeviceType.kDLExtDev: "ext_dev", DLDeviceType.kDLOneAPI: "oneapi", DLDeviceType.kDLWebGPU: "webgpu", DLDeviceType.kDLHexagon: "hexagon", DLDeviceType.kDLMAIA: "maia", DLDeviceType.kDLTrn: "trn", } _DEVICE_NAME_TO_TYPE = { "llvm": DLDeviceType.kDLCPU, "cpu": DLDeviceType.kDLCPU, "c": DLDeviceType.kDLCPU, "test": DLDeviceType.kDLCPU, "cuda": DLDeviceType.kDLCUDA, "nvptx": DLDeviceType.kDLCUDA, "cl": DLDeviceType.kDLOpenCL, "opencl": DLDeviceType.kDLOpenCL, "vulkan": DLDeviceType.kDLVulkan, "metal": DLDeviceType.kDLMetal, "vpi": DLDeviceType.kDLVPI, "rocm": DLDeviceType.kDLROCM, "ext_dev": DLDeviceType.kDLExtDev, "hexagon": DLDeviceType.kDLHexagon, "webgpu": DLDeviceType.kDLWebGPU, "maia": DLDeviceType.kDLMAIA, "trn": DLDeviceType.kDLTrn, } def __init__(self, device_type: str | int, index: Optional[Integral] = None) -> None: device_type_or_name = device_type index = index if index is not None else 0 if isinstance(device_type_or_name, str): # skip suffix annotations device_type_or_name = device_type_or_name.split(" ")[0] parts = device_type_or_name.split(":") if len(parts) < 1 or len(parts) > 2: raise ValueError(f"Invalid device: {device_type_or_name}") if parts[0] not in self._DEVICE_NAME_TO_TYPE: raise ValueError(f"Unknown device: {parts[0]}") device_type = self._DEVICE_NAME_TO_TYPE[parts[0]] if len(parts) == 2: try: index = int(parts[1]) except ValueError: raise ValueError(f"Invalid device index: {parts[1]}") else: device_type = device_type_or_name if not isinstance(index, Integral): if hasattr(index, "item") and callable(index.item): index = index.item() if not isinstance(index, Integral): raise TypeError(f"Invalid device index: {index}") self.cdevice = TVMFFIDLDeviceFromIntPair(device_type, index) def __reduce__(self) -> Any: cls = type(self) return (_create_device_from_tuple, (cls, self.cdevice.device_type, self.cdevice.device_id)) def __eq__(self, other: object) -> bool: if not isinstance(other, Device): return False return ( self.cdevice.device_type == (other).cdevice.device_type and self.cdevice.device_id == (other).cdevice.device_id ) def __ne__(self, other: object) -> bool: return not self.__eq__(other) def __str__(self) -> str: cdef int dev_type = self.cdevice.device_type name = self.__device_type_name__() index = self.cdevice.device_id return f"{name}:{index}" def __repr__(self) -> str: cdef int dev_type = self.cdevice.device_type name = self.__device_type_name__() index = self.cdevice.device_id return f"device(type='{name}', index={index})" def __hash__(self) -> int: return hash((self.cdevice.device_type, self.cdevice.device_id)) def __device_type_name__(self) -> str: """Return the canonical device type name (e.g. ``"cuda"``).""" return self._DEVICE_TYPE_TO_NAME[self.cdevice.device_type] @property def type(self) -> str: """Device type name such as ``"cpu"`` or ``"cuda"``.""" return self.__device_type_name__() @property def index(self) -> int: """Zero-based device index.""" return self.cdevice.device_id def dlpack_device_type(self) -> int: """Return the corresponding :class:`DLDeviceType` enum value.""" return self.cdevice.device_type cdef inline object make_ret_device(TVMFFIAny result): ret = _CLASS_DEVICE.__new__(_CLASS_DEVICE) (ret).cdevice = result.v_device return ret _set_class_device(Device) tvm-ffi-0.1.12/python/tvm_ffi/cython/dtype.pxi000066400000000000000000000217761521067262500213140ustar00rootroot00000000000000 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import os from typing import Any if os.environ.get("TVM_FFI_BUILD_DOCS", "0") == "0": try: # optionally import torch and setup torch related utils import torch except ImportError: torch = None try: # optionally import numpy import numpy except ImportError: numpy = None try: # optionally import ml_dtypes import ml_dtypes except ImportError: ml_dtypes = None else: torch = None numpy = None ml_dtypes = None _CLASS_DTYPE = None def _set_class_dtype(cls): global _CLASS_DTYPE _CLASS_DTYPE = cls def _create_cdtype_from_tuple(cls, code, bits, lanes): cdef DLDataType cdtype cdtype.code = code cdtype.bits = bits cdtype.lanes = lanes ret = cls.__new__(cls) (ret).cdtype = cdtype return ret cdef class DataType: """Internal wrapper around ``DLDataType``. This is a low-level representation used by the FFI layer. It is not intended as a user-facing API. For user code, prefer :class:`tvm_ffi.dtype`, which behaves like a Python ``str`` and integrates with array libraries. Examples -------- .. code-block:: python # Prefer the user-facing helper d = tvm_ffi.dtype("int32") assert d.bits == 32 assert str(d) == "int32" """ __slots__ = () cdef DLDataType cdtype def __init__(self, dtype_str: str) -> None: cdef ByteArrayArg dtype_str_arg = ByteArrayArg(c_str(dtype_str)) CHECK_CALL(TVMFFIDataTypeFromString(dtype_str_arg.cptr(), &(self.cdtype))) def __reduce__(self) -> Any: cls = type(self) return (_create_cdtype_from_tuple, (cls, self.cdtype.code, self.cdtype.bits, self.cdtype.lanes)) def __eq__(self, other: object) -> bool: if not isinstance(other, DataType): return False return ( self.cdtype.code == other.cdtype.code and self.cdtype.bits == other.cdtype.bits and self.cdtype.lanes == other.cdtype.lanes ) def __ne__(self, other: object) -> bool: return not self.__eq__(other) def __hash__(self) -> int: return hash((self.cdtype.code, self.cdtype.bits, self.cdtype.lanes)) @property def type_code(self) -> int: """Integer DLDataTypeCode of the scalar base type.""" return self.cdtype.code @property def bits(self) -> int: """Number of bits of the scalar base type.""" return self.cdtype.bits @property def lanes(self) -> int: """Number of lanes (for vector types).""" return self.cdtype.lanes @property def itemsize(self) -> int: """Size of one element in bytes (``bits * lanes // 8``). When the number of lanes is greater than 1, the ``itemsize`` is the size of the vector type. Returns ------- int The number of bytes of a single element of this data type. """ lanes_as_int = self.cdtype.lanes if lanes_as_int < 0: raise ValueError("Cannot determine itemsize for scalable vector types") return (self.cdtype.bits * self.cdtype.lanes + 7) // 8 def __str__(self) -> str: cdef TVMFFIAny temp_any cdef TVMFFIByteArray* bytes_ptr cdef TVMFFIByteArray bytes CHECK_CALL(TVMFFIDataTypeToString(&(self.cdtype), &temp_any)) if temp_any.type_index == kTVMFFISmallStr: bytes = TVMFFISmallBytesGetContentByteArray(&temp_any) res = bytearray_to_str(&bytes) return res bytes_ptr = TVMFFIBytesGetByteArrayPtr(temp_any.v_obj) res = bytearray_to_str(bytes_ptr) CHECK_CALL(TVMFFIObjectDecRef(temp_any.v_obj)) return res cdef inline object make_dtype_from_dl_data_type(DLDataType dl_data_type): cdtype = DataType.__new__(DataType) (cdtype).cdtype = dl_data_type val = str.__new__(_CLASS_DTYPE, cdtype.__str__()) val._tvm_ffi_dtype = cdtype return val cdef inline object make_ret_dtype(TVMFFIAny result): return make_dtype_from_dl_data_type(result.v_dtype) cdef TORCH_DTYPE_TO_DL_DATA_TYPE = {} cdef NUMPY_DTYPE_TO_DL_DATA_TYPE = {} cdef MLDTYPES_DTYPE_TO_DL_DATA_TYPE = {} if torch is not None: TORCH_DTYPE_TO_DL_DATA_TYPE = { torch.int8: DLDataType(0, 8, 1), torch.short: DLDataType(0, 16, 1), torch.int16: DLDataType(0, 16, 1), torch.int32: DLDataType(0, 32, 1), torch.int: DLDataType(0, 32, 1), torch.int64: DLDataType(0, 64, 1), torch.long: DLDataType(0, 64, 1), torch.uint8: DLDataType(1, 8, 1), torch.float16: DLDataType(2, 16, 1), torch.half: DLDataType(2, 16, 1), torch.float32: DLDataType(2, 32, 1), torch.float: DLDataType(2, 32, 1), torch.float64: DLDataType(2, 64, 1), torch.double: DLDataType(2, 64, 1), torch.bfloat16: DLDataType(4, 16, 1), torch.bool: DLDataType(6, 8, 1), } extra_types = [ ("uint16", DLDataType(1, 16, 1)), ("uint32", DLDataType(1, 32, 1)), ("uint64", DLDataType(1, 64, 1)), ("float8_e4m3fn", DLDataType(10, 8, 1)), ("float8_e4m3fnuz", DLDataType(11, 8, 1)), ("float8_e5m2", DLDataType(12, 8, 1)), ("float8_e5m2fnuz", DLDataType(13, 8, 1)), ("float8_e8m0fnu", DLDataType(14, 8, 1)), ("float4_e2m1fn_x2", DLDataType(17, 4, 2)), ] for attr_name, dl_dtype in extra_types: if hasattr(torch, attr_name): TORCH_DTYPE_TO_DL_DATA_TYPE[getattr(torch, attr_name)] = dl_dtype def _convert_torch_dtype_to_ffi_dtype(torch_dtype): cdef DLDataType dl_data_type = TORCH_DTYPE_TO_DL_DATA_TYPE[torch_dtype] return make_dtype_from_dl_data_type(dl_data_type) else: def _convert_torch_dtype_to_ffi_dtype(torch_dtype): raise ValueError("torch not found") if ml_dtypes is not None: MLDTYPES_DTYPE_TO_DL_DATA_TYPE = { numpy.dtype(ml_dtypes.int4): DLDataType(0, 4, 1), numpy.dtype(ml_dtypes.uint4): DLDataType(1, 4, 1), numpy.dtype(ml_dtypes.bfloat16): DLDataType(4, 16, 1), numpy.dtype(ml_dtypes.float8_e4m3b11fnuz): DLDataType(9, 8, 1), numpy.dtype(ml_dtypes.float8_e4m3fn): DLDataType(10, 8, 1), numpy.dtype(ml_dtypes.float8_e4m3fnuz): DLDataType(11, 8, 1), numpy.dtype(ml_dtypes.float8_e5m2): DLDataType(12, 8, 1), numpy.dtype(ml_dtypes.float8_e5m2fnuz): DLDataType(13, 8, 1), } if hasattr(ml_dtypes, "int2"): # ml_dtypes >= 0.5.0 MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.int2)] = DLDataType(0, 2, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.uint2)] = DLDataType(1, 2, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float8_e3m4)] = DLDataType(7, 8, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float8_e4m3)] = DLDataType(8, 8, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float8_e8m0fnu)] = DLDataType(14, 8, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float6_e2m3fn)] = DLDataType(15, 6, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float6_e3m2fn)] = DLDataType(16, 6, 1) MLDTYPES_DTYPE_TO_DL_DATA_TYPE[numpy.dtype(ml_dtypes.float4_e2m1fn)] = DLDataType(17, 4, 1) if numpy is not None: NUMPY_DTYPE_TO_DL_DATA_TYPE = { numpy.dtype(numpy.int8): DLDataType(0, 8, 1), numpy.dtype(numpy.int16): DLDataType(0, 16, 1), numpy.dtype(numpy.int32): DLDataType(0, 32, 1), numpy.dtype(numpy.int64): DLDataType(0, 64, 1), numpy.dtype(numpy.uint8): DLDataType(1, 8, 1), numpy.dtype(numpy.uint16): DLDataType(1, 16, 1), numpy.dtype(numpy.uint32): DLDataType(1, 32, 1), numpy.dtype(numpy.uint64): DLDataType(1, 64, 1), numpy.dtype(numpy.float16): DLDataType(2, 16, 1), numpy.dtype(numpy.float32): DLDataType(2, 32, 1), numpy.dtype(numpy.float64): DLDataType(2, 64, 1), **MLDTYPES_DTYPE_TO_DL_DATA_TYPE, } def _convert_numpy_dtype_to_ffi_dtype(numpy_dtype): cdef DLDataType cdtype = NUMPY_DTYPE_TO_DL_DATA_TYPE[numpy_dtype] return make_dtype_from_dl_data_type(cdtype) else: def _convert_numpy_dtype_to_ffi_dtype(numpy_dtype): raise ValueError("numpy not found") tvm-ffi-0.1.12/python/tvm_ffi/cython/error.pxi000066400000000000000000000157051521067262500213130ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # error handling for FFI import types import re from typing import Callable, Optional ERROR_NAME_TO_TYPE: dict[str, type] = {} ERROR_TYPE_TO_NAME: dict[type, str] = {} _WITH_APPEND_BACKTRACE: Optional[Callable[[BaseException, str], BaseException]] = None _TRACEBACK_TO_BACKTRACE_STR: Optional[Callable[[types.TracebackType | None], str]] = None cdef class Error(CObject): """Base class for FFI errors. An :class:`Error` is a lightweight wrapper around a concrete Python exception raised by FFI calls. It stores the error ``kind`` (e.g. ``"ValueError"``), the message, and a serialized FFI backtrace that can be re-attached to produce a Python traceback. Users normally interact with specific error subclasses that are registered via :func:`tvm_ffi.error.register_error`. Notes ----- Do not directly raise this object. Instead, use :py:meth:`py_error` to convert it to a Python exception and raise that. """ __slots__ = () def __init__(self, kind: str, message: str, backtrace: str): """Construct an error wrapper. Parameters ---------- kind : str Name of the Python exception type (e.g. ``"ValueError"``). message : str The error message from the FFI side. backtrace : str Serialized backtrace encoded by the runtime. """ cdef ByteArrayArg kind_arg = ByteArrayArg(c_str(kind)) cdef ByteArrayArg message_arg = ByteArrayArg(c_str(message)) cdef ByteArrayArg backtrace_arg = ByteArrayArg(c_str(backtrace)) cdef TVMFFIObjectHandle out cdef int ret = TVMFFIErrorCreate( kind_arg.cptr(), message_arg.cptr(), backtrace_arg.cptr(), &out ) if ret != 0: raise MemoryError("Failed to create error object") (self).chandle = out def update_backtrace(self, backtrace: str) -> None: """Replace the stored backtrace string with ``backtrace``. Parameters ---------- backtrace : str The backtrace to store. The internal storage is reverse of Python's traceback order to simplify appending during propagation; it is reversed again when rendered. """ cdef ByteArrayArg backtrace_arg = ByteArrayArg(c_str(backtrace)) TVMFFIErrorGetCellPtr(self.chandle).update_backtrace( self.chandle, backtrace_arg.cptr(), kTVMFFIBacktraceUpdateModeReplace ) def py_error(self) -> BaseException: """Return a Python :class:`BaseException` instance for this error.""" error_cls = ERROR_NAME_TO_TYPE.get(self.kind, RuntimeError) py_error = error_cls(self.message) py_error = _WITH_APPEND_BACKTRACE(py_error, self.backtrace) py_error.__tvm_ffi_error__ = self return py_error @property def kind(self): return bytearray_to_str(&(TVMFFIErrorGetCellPtr(self.chandle).kind)) @property def message(self): return bytearray_to_str(&(TVMFFIErrorGetCellPtr(self.chandle).message)) @property def backtrace(self): return bytearray_to_str(&(TVMFFIErrorGetCellPtr(self.chandle).backtrace)) @property def extra_context(self): """Optional structured payload attached to this error. Returns ``None`` if nothing is attached. May be inspected via the appropriate type-specific helpers. """ cdef TVMFFIObjectHandle ctx_handle = TVMFFIErrorGetCellPtr(self.chandle).extra_context if ctx_handle == NULL: return None # Build an owned Any from the unowned handle by incrementing the refcount. cdef TVMFFIAny any_val any_val.type_index = TVMFFIObjectGetTypeIndex(ctx_handle) any_val.v_obj = ctx_handle TVMFFIObjectIncRef(ctx_handle) return make_ret_object(any_val) cdef inline Error move_from_last_error(): # raise last error error = Error.__new__(Error) TVMFFIErrorMoveFromRaised(&(error).chandle) return error cdef inline int raise_existing_error() except -2: return -2 cdef inline int set_last_ffi_error(error) except -1: """Set the last FFI error""" cdef Error ffi_error kind = ERROR_TYPE_TO_NAME.get(type(error), "RuntimeError") message = error.__str__() # NOTE: backtrace storage convention is reverse of python traceback py_backtrace = _TRACEBACK_TO_BACKTRACE_STR(error.__traceback__) c_backtrace = bytearray_to_str(TVMFFIBacktrace(NULL, 0, NULL, 0)) # error comes from an exception thrown from C++ side if hasattr(error, "__tvm_ffi_error__"): # already have stack trace ffi_error = error.__tvm_ffi_error__ # attach the python backtrace together with the C++ backtrace to get full trace ffi_error.update_backtrace(py_backtrace + c_backtrace) TVMFFIErrorSetRaised(ffi_error.chandle) else: ffi_error = Error(kind, message, py_backtrace + c_backtrace) TVMFFIErrorSetRaised(ffi_error.chandle) def _convert_to_ffi_error(error: BaseException) -> Error: """Convert the python error to the FFI error""" py_backtrace = _TRACEBACK_TO_BACKTRACE_STR(error.__traceback__) if hasattr(error, "__tvm_ffi_error__"): error.__tvm_ffi_error__.update_backtrace(py_backtrace) return error.__tvm_ffi_error__ else: kind = ERROR_TYPE_TO_NAME.get(type(error), "RuntimeError") message = error.__str__() return Error(kind, message, py_backtrace) cdef public int TVMFFICyErrorSetRaisedFromPyError(PyObject* py_err) noexcept: """Set the last FFI error from a Python exception. Parameters ---------- py_err : PyObject* The Python exception to set as the last FFI error. """ set_last_ffi_error(py_err) return -1 cdef inline int CHECK_CALL(int ret) except -2: """Check the return code of the C API function call""" if ret == 0: return 0 # backward compact with error already set case # TODO(tqchen): remove after we move beyond a few versions. if ret == -2: raise raise_existing_error() error = move_from_last_error() if error.kind == "EnvErrorAlreadySet": raise raise_existing_error() raise error.py_error() tvm-ffi-0.1.12/python/tvm_ffi/cython/function.pxi000066400000000000000000001242571521067262500220120ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import ctypes import threading import os from numbers import Integral, Real from typing import Any, Callable if os.environ.get("TVM_FFI_BUILD_DOCS", "0") == "0": try: # optionally import torch and setup torch related utils import torch except ImportError: torch = None try: # optionally import numpy import numpy except ImportError: numpy = None try: from cuda.bindings import driver as cuda_driver except ImportError: cuda_driver = None else: torch = None numpy = None cuda_driver = None cdef int _RELEASE_GIL_BY_DEFAULT = int( os.environ.get("TVM_FFI_RELEASE_GIL_BY_DEFAULT", "1") ) cdef inline object make_ret_small_str(TVMFFIAny result): """convert small string to return value.""" cdef TVMFFIByteArray bytes bytes = TVMFFISmallBytesGetContentByteArray(&result) return bytearray_to_str(&bytes) cdef inline object make_ret_small_bytes(TVMFFIAny result): """convert small bytes to return value.""" cdef TVMFFIByteArray bytes bytes = TVMFFISmallBytesGetContentByteArray(&result) return bytearray_to_bytes(&bytes) cdef inline object make_ret(TVMFFIAny result, const DLPackExchangeAPI* c_ctx_dlpack_api = NULL): """convert result to return value.""" cdef int32_t type_index type_index = result.type_index if type_index == kTVMFFITensor: # specially handle Tensor as it needs a special dltensor field return make_tensor_from_any(result, c_ctx_dlpack_api) elif type_index == kTVMFFIOpaquePyObject: return make_ret_opaque_object(result) elif type_index >= kTVMFFIStaticObjectBegin: obj = make_ret_object(result) if c_ctx_dlpack_api != NULL and isinstance(obj, CContainerBase): (obj)._dlpack_exchange_api = c_ctx_dlpack_api return obj # the following code should be optimized to switch case if type_index == kTVMFFINone: return None elif type_index == kTVMFFIBool: return bool(result.v_int64) elif type_index == kTVMFFIInt: return result.v_int64 elif type_index == kTVMFFIFloat: return result.v_float64 elif type_index == kTVMFFISmallStr: return make_ret_small_str(result) elif type_index == kTVMFFISmallBytes: return make_ret_small_bytes(result) elif type_index == kTVMFFIOpaquePtr: return ctypes_handle(result.v_ptr) elif type_index == kTVMFFIDataType: return make_ret_dtype(result) elif type_index == kTVMFFIDevice: return make_ret_device(result) elif type_index == kTVMFFIDLTensorPtr: return make_ret_dltensor(result) elif type_index == kTVMFFIObjectRValueRef: raise ValueError("Return value cannot be ObjectRValueRef") elif type_index == kTVMFFIByteArrayPtr: raise ValueError("Return value cannot be ByteArrayPtr") elif type_index == kTVMFFIRawStr: raise ValueError("Return value cannot be RawStr") raise ValueError("Unhandled type index %d" % type_index) # ---------------------------------------------------------------------------- # Helper to simplify calling constructor # ---------------------------------------------------------------------------- cdef inline int ConstructorCall(void* constructor_handle, PyObject* py_arg_tuple, void** handle, TVMFFIPyCallContext* parent_ctx) except -1: """Call contructor of a handle function""" cdef TVMFFIAny result cdef int c_api_ret_code # IMPORTANT: caller need to initialize result->type_index to kTVMFFINone result.type_index = kTVMFFINone result.v_int64 = 0 TVMFFIPyConstructorCall( constructor_handle, py_arg_tuple, &result, &c_api_ret_code, parent_ctx ) CHECK_CALL(c_api_ret_code) handle[0] = result.v_ptr return 0 # ---------------------------------------------------------------------------- # Implementation of setters using same naming style as TVMFFIPyArgSetterXXX_ # ---------------------------------------------------------------------------- cdef int TVMFFIPyArgSetterTensor_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* arg, TVMFFIAny* out ) except -1: if (arg).chandle != NULL: out.type_index = kTVMFFITensor out.v_ptr = (arg).chandle else: out.type_index = kTVMFFIDLTensorPtr out.v_ptr = (arg).cdltensor return 0 cdef int TVMFFIPyArgSetterObject_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* arg, TVMFFIAny* out ) except -1: out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle return 0 cdef int TVMFFIPyArgSetterContainerObject_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* arg, TVMFFIAny* out ) except -1: """Setter for container objects (Array, List, Map, Dict). Propagates DLPack exchange API tag and scans for stream context. """ cdef TVMFFIAny scan_args[1] cdef TVMFFIAny scan_result cdef void* stream = NULL out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle cdef const DLPackExchangeAPI* api = (arg)._dlpack_exchange_api if api != NULL: if ctx.dlpack_c_exchange_api == NULL: ctx.dlpack_c_exchange_api = api if ctx.device_type == -1 and api.current_work_stream != NULL: # Call C++ to find the first non-CPU tensor device in one shot. scan_args[0].type_index = out.type_index scan_args[0].v_obj = (arg).chandle scan_result.type_index = kTVMFFINone scan_result.v_int64 = 0 CHECK_CALL(TVMFFIFunctionCall( (_FFI_CONTAINER_FIND_FIRST_NON_CPU_DEVICE).chandle, scan_args, 1, &scan_result)) if scan_result.type_index == kTVMFFIDevice and scan_result.v_device.device_type != kDLCPU: ctx.device_type = scan_result.v_device.device_type ctx.device_id = scan_result.v_device.device_id api.current_work_stream( scan_result.v_device.device_type, scan_result.v_device.device_id, &stream) ctx.stream = stream return 0 cdef int TVMFFIPyArgSetterDLPackExchangeAPI_( TVMFFIPyArgSetter* this, TVMFFIPyCallContext* ctx, PyObject* arg, TVMFFIAny* out ) except -1: cdef DLManagedTensorVersioned* temp_managed_tensor cdef TVMFFIObjectHandle temp_chandle cdef void* current_stream = NULL cdef const DLPackExchangeAPI* exchange_api = this.dlpack_c_exchange_api # Set the exchange API in context ctx.dlpack_c_exchange_api = exchange_api # Convert PyObject to DLPack using the struct's function pointer if exchange_api.managed_tensor_from_py_object_no_sync(arg, &temp_managed_tensor) != 0: return -1 # Query current stream from producer if device is not CPU if temp_managed_tensor.dl_tensor.device.device_type != kDLCPU: if ctx.device_type == -1 and exchange_api.current_work_stream != NULL: # First time seeing a device, query the stream if exchange_api.current_work_stream( temp_managed_tensor.dl_tensor.device.device_type, temp_managed_tensor.dl_tensor.device.device_id, ¤t_stream ) == 0: ctx.stream = current_stream ctx.device_type = temp_managed_tensor.dl_tensor.device.device_type ctx.device_id = temp_managed_tensor.dl_tensor.device.device_id # Convert to TVM Tensor if TVMFFITensorFromDLPackVersioned(temp_managed_tensor, 0, 0, &temp_chandle) != 0: # recycle the managed tensor to avoid leak if temp_managed_tensor.deleter != NULL: temp_managed_tensor.deleter(temp_managed_tensor) raise BufferError("Failed to convert DLManagedTensorVersioned to ffi.Tensor") out.type_index = kTVMFFITensor out.v_ptr = temp_chandle TVMFFIPyPushTempFFIObject(ctx, temp_chandle) return 0 cdef int TorchManagedTensorToPyObjectNoSyncFallback_( DLManagedTensorVersioned* dltensor, void** py_obj_out ) except -1: # a bit convoluted but ok as a fallback cdef TVMFFIObjectHandle temp_chandle if TVMFFITensorFromDLPackVersioned(dltensor, 0, 0, &temp_chandle) != 0: return -1 tensor = make_tensor_from_chandle(temp_chandle) torch_tensor = torch.from_dlpack(tensor) Py_INCREF(torch_tensor) py_obj_out[0] = (torch_tensor) return 0 cdef inline const DLPackExchangeAPI* GetTorchFallbackExchangeAPI() noexcept: global _torch_fallback_exchange_api _torch_fallback_exchange_api.header.version.major = DLPACK_MAJOR_VERSION _torch_fallback_exchange_api.header.version.minor = DLPACK_MINOR_VERSION _torch_fallback_exchange_api.header.prev_api = NULL _torch_fallback_exchange_api.managed_tensor_allocator = NULL _torch_fallback_exchange_api.managed_tensor_from_py_object_no_sync = NULL _torch_fallback_exchange_api.managed_tensor_to_py_object_no_sync = ( TorchManagedTensorToPyObjectNoSyncFallback_ ) _torch_fallback_exchange_api.dltensor_from_py_object_no_sync = NULL _torch_fallback_exchange_api.current_work_stream = NULL return &_torch_fallback_exchange_api # Static storage for the fallback exchange API cdef DLPackExchangeAPI _torch_fallback_exchange_api cdef int TVMFFIPyArgSetterTorchFallback_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Current setter for torch.Tensor, go through python and not as fast as c exporter""" # TODO(tqchen): remove this once torch always support fast DLPack importer cdef object arg = py_arg cdef long long temp_ptr is_cuda = arg.is_cuda arg = from_dlpack(torch.utils.dlpack.to_dlpack(arg)) out.type_index = kTVMFFITensor out.v_ptr = (arg).chandle temp_dltensor = TVMFFITensorGetDLTensorPtr((arg).chandle) ctx.dlpack_c_exchange_api = GetTorchFallbackExchangeAPI() # record the stream and device for torch context if is_cuda and ctx.device_type == -1: ctx.device_type = temp_dltensor.device.device_type ctx.device_id = temp_dltensor.device.device_id # This is an API that dynamo and other uses to get the raw stream from torch temp_ptr = torch._C._cuda_getCurrentRawStream(temp_dltensor.device.device_id) ctx.stream = temp_ptr # push to temp and clear the handle TVMFFIPyPushTempPyObject(ctx, arg) return 0 cdef int TVMFFIPyArgSetterDLPack_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for __dlpack__ mechanism through python, not as fast as c exporter""" cdef TVMFFIObjectHandle temp_chandle cdef object arg = py_arg _from_dlpack_universal(arg, 0, 0, &temp_chandle) out.type_index = kTVMFFITensor out.v_ptr = temp_chandle # record the stream from the source framework context when possible temp_dltensor = TVMFFITensorGetDLTensorPtr(temp_chandle) if (temp_dltensor.device.device_type != kDLCPU and ctx.device_type != -1): # __tvm_ffi_env_stream__ returns the expected stream that should be set # through TVMFFIEnvSetStream when calling a TVM FFI function if hasattr(arg, "__tvm_ffi_env_stream__"): # Ideally projects should directly setup their stream context API # write through by also calling TVMFFIEnvSetStream # so we do not need this protocol to do exchange ctx.device_type = temp_dltensor.device.device_type ctx.device_id = temp_dltensor.device.device_id temp_ptr= arg.__tvm_ffi_env_stream__() ctx.stream = temp_ptr TVMFFIPyPushTempFFIObject(ctx, temp_chandle) return 0 cdef int TVMFFIPyArgSetterIntegral_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Integral""" cdef object arg = py_arg out.type_index = kTVMFFIInt # keep it in cython so it will also check for fallback cases # where the arg is not exactly the int class out.v_int64 = arg return 0 cdef int TVMFFIPyArgSetterReal_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Real""" cdef object arg = py_arg out.type_index = kTVMFFIFloat # keep it in cython so it will also check for fallback cases # where the arg is not exactly the float class out.v_float64 = arg return 0 cdef int TVMFFIPyArgSetterFFIObjectProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for objects that implement the `__tvm_ffi_object__` protocol.""" cdef object arg = py_arg cdef TVMFFIObjectHandle temp_chandle cdef CObject obj = arg.__tvm_ffi_object__() cdef long ref_count = Py_REFCNT(obj) temp_chandle = obj.chandle out.type_index = TVMFFIObjectGetTypeIndex(temp_chandle) out.v_ptr = temp_chandle if ref_count == 1: # keep alive the tensor, since the tensor is temporary # and will be freed after we exit here TVMFFIObjectIncRef(temp_chandle) TVMFFIPyPushTempFFIObject(ctx, temp_chandle) return 0 cdef int TVMFFIPyArgSetterCUDAStreamProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for cuda stream protocol""" cdef object arg = py_arg # cuda stream is a subclass of str, so this check occur before str cdef tuple cu_stream_tuple = arg.__cuda_stream__() cdef long long long_ptr = cu_stream_tuple[1] out.type_index = kTVMFFIOpaquePtr out.v_ptr = long_ptr return 0 cdef int TVMFFIPyArgSetterCUDADriverStreamFallback_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for cuda.bindings.driver.CUstream as a fallback without __cuda_stream__ protocol""" cdef object arg = py_arg # call driver stream cdef long long long_ptr = int(arg) out.type_index = kTVMFFIOpaquePtr out.v_ptr = long_ptr return 0 cdef int TVMFFIPyArgSetterDType_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for dtype""" cdef object arg = py_arg # dtype is a subclass of str, so this check occur before str arg = arg._tvm_ffi_dtype out.type_index = kTVMFFIDataType out.v_dtype = (arg).cdtype return 0 cdef int TVMFFIPyArgSetterDevice_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for device""" cdef object arg = py_arg out.type_index = kTVMFFIDevice out.v_device = (arg).cdevice return 0 cdef int TVMFFIPyArgSetterDLPackDeviceProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for dlpack device protocol""" cdef object arg = py_arg cdef tuple dlpack_device = arg.__dlpack_device__() out.type_index = kTVMFFIDevice out.v_device = TVMFFIDLDeviceFromIntPair( dlpack_device[0], dlpack_device[1] ) return 0 cdef int TVMFFIPyArgSetterStr_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for str""" cdef object arg = py_arg cdef bytes tstr = arg.encode("utf-8") cdef char* data cdef Py_ssize_t size cdef TVMFFIByteArray cdata PyBytes_AsStringAndSize(tstr, &data, &size) cdata.data = data cdata.size = size CHECK_CALL(TVMFFIStringFromByteArray(&cdata, out)) if out.type_index >= kTVMFFIStaticObjectBegin: TVMFFIPyPushTempFFIObject(ctx, out.v_ptr) return 0 cdef int TVMFFIPyArgSetterPyNativeObjectStr_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Specially handle String as its _tvm_ffi_cached_object may be empty""" cdef object arg = py_arg # need to check if the arg is a large string returned from ffi if arg._tvm_ffi_cached_object is not None: arg = arg._tvm_ffi_cached_object out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle return 0 return TVMFFIPyArgSetterStr_(handle, ctx, py_arg, out) cdef int TVMFFIPyArgSetterBytes_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for bytes""" cdef object arg = py_arg if isinstance(arg, bytearray): arg = bytes(arg) cdef char* data cdef Py_ssize_t size cdef TVMFFIByteArray cdata PyBytes_AsStringAndSize(arg, &data, &size) cdata.data = data cdata.size = size CHECK_CALL(TVMFFIBytesFromByteArray(&cdata, out)) if out.type_index >= kTVMFFIStaticObjectBegin: TVMFFIPyPushTempFFIObject(ctx, out.v_ptr) return 0 cdef int TVMFFIPyArgSetterPyNativeObjectBytes_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Specially handle Bytes as its _tvm_ffi_cached_object may be empty""" cdef object arg = py_arg # need to check if the arg is a large bytes returned from ffi if arg._tvm_ffi_cached_object is not None: arg = arg._tvm_ffi_cached_object out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle return 0 return TVMFFIPyArgSetterBytes_(handle, ctx, py_arg, out) cdef int TVMFFIPyArgSetterPyNativeObjectGeneral_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Specially handle Object as its _tvm_ffi_cached_object may be empty""" cdef object arg = py_arg if arg._tvm_ffi_cached_object is None: raise ValueError(f"_tvm_ffi_cached_object is None for {type(arg)}") assert arg._tvm_ffi_cached_object is not None arg = arg._tvm_ffi_cached_object out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle return 0 cdef int TVMFFIPyArgSetterCtypesVoidPtr_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for ctypes.c_void_p""" out.type_index = kTVMFFIOpaquePtr out.v_ptr = c_handle(py_arg) return 0 cdef int TVMFFIPyArgSetterFFIOpaquePtrCompatible_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for objects that implement the `__tvm_ffi_opaque_ptr__` protocol.""" cdef object arg = py_arg cdef long long long_ptr = arg.__tvm_ffi_opaque_ptr__() out.type_index = kTVMFFIOpaquePtr out.v_ptr = long_ptr return 0 cdef int TVMFFIPyArgSetterObjectRValueRef_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for ObjectRValueRef""" cdef object arg = py_arg out.type_index = kTVMFFIObjectRValueRef out.v_ptr = &(((arg.obj)).chandle) return 0 cdef int TVMFFIPyArgSetterCallable_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Callable""" cdef TVMFFIObjectHandle chandle CHECK_CALL(TVMFFIPyConvertPyCallback(py_arg, NULL, &chandle)) out.type_index = TVMFFIObjectGetTypeIndex(chandle) out.v_ptr = chandle TVMFFIPyPushTempFFIObject(ctx, chandle) return 0 cdef int TVMFFIPyArgSetterException_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Exception""" cdef object arg = py_arg arg = _convert_to_ffi_error(arg) out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle TVMFFIPyPushTempPyObject(ctx, arg) return 0 cdef int TVMFFIPyArgSetterTuple_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Tuple""" # recursively construct a new tuple cdef TVMFFIObjectHandle chandle ConstructorCall(_CONSTRUCTOR_ARRAY.chandle, py_arg, &chandle, ctx) out.type_index = TVMFFIObjectGetTypeIndex(chandle) out.v_ptr = chandle TVMFFIPyPushTempFFIObject(ctx, chandle) return 0 cdef int TVMFFIPyArgSetterTupleLike_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for TupleLike""" # recursively construct a new tuple cdef tuple tuple_arg = tuple(py_arg) cdef TVMFFIObjectHandle chandle ConstructorCall(_CONSTRUCTOR_ARRAY.chandle, tuple_arg, &chandle, ctx) out.type_index = TVMFFIObjectGetTypeIndex(chandle) out.v_ptr = chandle TVMFFIPyPushTempFFIObject(ctx, chandle) return 0 cdef int TVMFFIPyArgSetterMap_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for Map""" # recursively construct a new map cdef dict dict_arg = py_arg cdef list list_kvs = [] for k, v in dict_arg.items(): list_kvs.append(k) list_kvs.append(v) cdef tuple_arg_kvs = tuple(list_kvs) cdef TVMFFIObjectHandle chandle ConstructorCall(_CONSTRUCTOR_MAP.chandle, tuple_arg_kvs, &chandle, ctx) out.type_index = TVMFFIObjectGetTypeIndex(chandle) out.v_ptr = chandle TVMFFIPyPushTempFFIObject(ctx, chandle) return 0 cdef int TVMFFIPyArgSetterObjectConvertible_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for ObjectConvertible""" # recursively construct a new map cdef object arg = py_arg arg = arg.asobject() out.type_index = TVMFFIObjectGetTypeIndex((arg).chandle) out.v_ptr = (arg).chandle TVMFFIPyPushTempPyObject(ctx, arg) cdef int TVMFFIPyArgSetterFallback_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Fallback setter for all other types""" cdef object arg = py_arg cdef TVMFFIObjectHandle chandle _convert_to_opaque_object_handle(arg, &chandle) out.type_index = kTVMFFIOpaquePyObject out.v_ptr = chandle TVMFFIPyPushTempFFIObject(ctx, chandle) cdef int TVMFFIPyArgSetterDTypeFromTorch_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for torch dtype""" cdef py_obj = py_arg if py_obj not in TORCH_DTYPE_TO_DL_DATA_TYPE: raise ValueError("Unsupported torch dtype: ", py_obj) out.type_index = kTVMFFIDataType out.v_dtype = TORCH_DTYPE_TO_DL_DATA_TYPE[py_obj] return 0 cdef int TVMFFIPyArgSetterDTypeFromNumpy_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for torch dtype""" cdef py_obj = py_arg if py_obj not in NUMPY_DTYPE_TO_DL_DATA_TYPE: raise ValueError("Unsupported numpy or ml_dtypes dtype: ", py_obj) out.type_index = kTVMFFIDataType out.v_dtype = NUMPY_DTYPE_TO_DL_DATA_TYPE[py_obj] return 0 cdef int TVMFFIPyArgSetterDLPackDataTypeProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for dtype protocol""" cdef object arg = py_arg cdef tuple dltype_data_type = arg.__dlpack_data_type__() out.type_index = kTVMFFIDataType out.v_dtype.code = dltype_data_type[0] out.v_dtype.bits = dltype_data_type[1] out.v_dtype.lanes = dltype_data_type[2] return 0 cdef int TVMFFIPyArgSetterIntProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for class with __tvm_ffi_int__() method""" cdef object arg = py_arg out.type_index = kTVMFFIInt out.v_int64 = (arg.__tvm_ffi_int__()) return 0 cdef int TVMFFIPyArgSetterFloatProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for class with __tvm_ffi_float__() method""" cdef object arg = py_arg out.type_index = kTVMFFIFloat out.v_float64 = (arg.__tvm_ffi_float__()) return 0 cdef int TVMFFIPyArgSetterFFIValueProtocol_( TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out ) except -1: """Setter for class with __tvm_ffi_value__() method""" cdef object arg = py_arg cdef object ffi_value_py_obj = arg.__tvm_ffi_value__() cdef PyObject* ffi_value_py_obj_ptr = ffi_value_py_obj # keep alive the python object since this is a temporary object # we must push to extra temp py objects stack to avoid overflow the temp py objects stack TVMFFIPyPushExtraTempPyObject(ctx, ffi_value_py_obj_ptr) return TVMFFIPySetArgumentGenericDispatcher(ctx, ffi_value_py_obj_ptr, out) cdef _DISPATCH_TYPE_KEEP_ALIVE = set() cdef _DISPATCH_TYPE_KEEP_ALIVE_LOCK = threading.Lock() cdef public int TVMFFICyArgSetterFactory(PyObject* value, TVMFFIPyArgSetter* out) except -1: """ Factory function that creates an argument setter for a given Python argument type. """ # NOTE: the order of checks matter here # becase each argument may satisfy multiple checks # priortize native types over external types cdef object arg = value cdef long long temp_ptr # The C++ dispatcher dispatches the argument passing by TYPE(obj) pointer which # is non-owning. This means that there is the following edge case: # - type A is registered through dispatcher # - type A gets garbage collected (because it is a local type) # - type B is created and uses the same memory address as type A # # Then when we pass in type B, it will mistakenly use the dispatch function for type A # # To prevent this, we keep alive the types that are registered through dispatcher # by adding them to _DISPATCH_TYPE_KEEP_ALIVE # # NOTE that the total number of types that are registered through dispatcher is expected # to be limited in practice so we can afford to keep them alive # Lock is used to ensure thread-safety for future thread-free python case with _DISPATCH_TYPE_KEEP_ALIVE_LOCK: _DISPATCH_TYPE_KEEP_ALIVE.add(type(arg)) if arg is None: out.func = TVMFFIPyArgSetterNone_ return 0 if isinstance(arg, Tensor): out.func = TVMFFIPyArgSetterTensor_ return 0 if isinstance(arg, CContainerBase): out.func = TVMFFIPyArgSetterContainerObject_ return 0 if isinstance(arg, CObject): out.func = TVMFFIPyArgSetterObject_ return 0 if isinstance(arg, ObjectRValueRef): out.func = TVMFFIPyArgSetterObjectRValueRef_ return 0 arg_class = type(arg) if hasattr(arg_class, "__tvm_ffi_object__"): # can directly map to tvm ffi object # usually used for solutions that takes subclass of ffi.Object # as a member variable out.func = TVMFFIPyArgSetterFFIObjectProtocol_ return 0 if os.environ.get("TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API", "0") != "1": # Check for DLPackExchangeAPI struct (new approach) # This is checked on the CLASS, not the instance if hasattr(arg_class, "__dlpack_c_exchange_api__"): out.func = TVMFFIPyArgSetterDLPackExchangeAPI_ _get_dlpack_exchange_api(arg_class.__dlpack_c_exchange_api__, &(out.dlpack_c_exchange_api)) return 0 if hasattr(arg_class, "__cuda_stream__"): # cuda stream protocol out.func = TVMFFIPyArgSetterCUDAStreamProtocol_ return 0 if cuda_driver is not None and isinstance(arg, cuda_driver.CUstream): # TODO(tqchen): remove this once cuda-python supports __cuda_stream__ protocol out.func = TVMFFIPyArgSetterCUDADriverStreamFallback_ return 0 if torch is not None and isinstance(arg, torch.Tensor): out.func = TVMFFIPyArgSetterTorchFallback_ return 0 if hasattr(arg_class, "__dlpack__"): out.func = TVMFFIPyArgSetterDLPack_ return 0 if isinstance(arg, bool): # A python `bool` is a subclass of `int`, so this check # must occur before `Integral`. out.func = TVMFFIPyArgSetterBool_ return 0 if isinstance(arg, Integral): # must occur before Real check # cannot simply use TVMFFIPyArgSetterInt # because Integral may not be exactly the int class out.func = TVMFFIPyArgSetterIntegral_ return 0 if isinstance(arg, Real): # cannot simply use TVMFFIPyArgSetterFloat # because Real may not be exactly the float class out.func = TVMFFIPyArgSetterReal_ return 0 # dtype is a subclass of str, so this check must occur before str if isinstance(arg, _CLASS_DTYPE): out.func = TVMFFIPyArgSetterDType_ return 0 if isinstance(arg, _CLASS_DEVICE): out.func = TVMFFIPyArgSetterDevice_ return 0 if isinstance(arg, PyNativeObject): # check for PyNativeObject # this check must happen before str/bytes/tuple if isinstance(arg, str): out.func = TVMFFIPyArgSetterPyNativeObjectStr_ return 0 if isinstance(arg, bytes): out.func = TVMFFIPyArgSetterPyNativeObjectBytes_ return 0 out.func = TVMFFIPyArgSetterPyNativeObjectGeneral_ return 0 if isinstance(arg, str): out.func = TVMFFIPyArgSetterStr_ return 0 if isinstance(arg, (bytes, bytearray)): out.func = TVMFFIPyArgSetterBytes_ return 0 if isinstance(arg, tuple): out.func = TVMFFIPyArgSetterTuple_ return 0 if isinstance(arg, list): out.func = TVMFFIPyArgSetterTupleLike_ return 0 if isinstance(arg, dict): out.func = TVMFFIPyArgSetterMap_ return 0 if isinstance(arg, ctypes.c_void_p): out.func = TVMFFIPyArgSetterCtypesVoidPtr_ return 0 if hasattr(arg_class, "__tvm_ffi_opaque_ptr__"): out.func = TVMFFIPyArgSetterFFIOpaquePtrCompatible_ return 0 if callable(arg): out.func = TVMFFIPyArgSetterCallable_ return 0 if torch is not None and isinstance(arg, torch.dtype): out.func = TVMFFIPyArgSetterDTypeFromTorch_ return 0 if numpy is not None and isinstance(arg, numpy.dtype): out.func = TVMFFIPyArgSetterDTypeFromNumpy_ return 0 if hasattr(arg_class, "__dlpack_data_type__"): # prefer dlpack as it covers all DLDataType struct out.func = TVMFFIPyArgSetterDLPackDataTypeProtocol_ return 0 if hasattr(arg_class, "__dlpack_device__") and not hasattr(arg_class, "__dlpack__"): # if a class have __dlpack_device__ but not __dlpack__ # then it is a DLPack device protocol out.func = TVMFFIPyArgSetterDLPackDeviceProtocol_ return 0 if hasattr(arg_class, "__tvm_ffi_int__"): out.func = TVMFFIPyArgSetterIntProtocol_ return 0 if hasattr(arg_class, "__tvm_ffi_float__"): out.func = TVMFFIPyArgSetterFloatProtocol_ return 0 if hasattr(arg_class, "__tvm_ffi_value__"): out.func = TVMFFIPyArgSetterFFIValueProtocol_ return 0 if isinstance(arg, Exception): out.func = TVMFFIPyArgSetterException_ return 0 if isinstance(arg, ObjectConvertible): out.func = TVMFFIPyArgSetterObjectConvertible_ return 0 # default to opaque object out.func = TVMFFIPyArgSetterFallback_ return 0 # --------------------------------------------------------------------------------------------- # Implementation of function calling # --------------------------------------------------------------------------------------------- cdef class Function(CObject): """Callable wrapper around a TVM FFI function. Instances are obtained by converting Python callables with :func:`tvm_ffi.convert`, or by looking up globally-registered FFI functions using :func:`tvm_ffi.get_global_func`. Examples -------- .. code-block:: python @tvm_ffi.register_global_func("my.add") def add(a, b): return a + b f = tvm_ffi.get_global_func("my.add") assert isinstance(f, tvm_ffi.Function) assert f(1, 2) == 3 See Also -------- :py:func:`tvm_ffi.register_global_func` Register a Python callable as a global FFI function. :py:func:`tvm_ffi.get_global_func` Look up a previously registered global FFI function by name. """ cdef int c_release_gil cdef dict __dict__ def __cinit__(self) -> None: self.c_release_gil = _RELEASE_GIL_BY_DEFAULT property release_gil: """Whether calls release the Python GIL while executing.""" def __get__(self) -> bool: return self.c_release_gil != 0 def __set__(self, value: bool) -> None: self.c_release_gil = value def __call__(self, *args: Any) -> Any: """Invoke the wrapped FFI function with ``args``.""" cdef TVMFFIAny result cdef int c_api_ret_code cdef const DLPackExchangeAPI* c_ctx_dlpack_api = NULL # IMPORTANT: caller need to initialize result->type_index to kTVMFFINone result.type_index = kTVMFFINone result.v_int64 = 0 TVMFFIPyFuncCall( (self).chandle, args, &result, &c_api_ret_code, self.release_gil, &c_ctx_dlpack_api ) # NOTE: logic is same as check_call # directly inline here to simplify the resulting trace if c_api_ret_code == 0: return make_ret(result, c_ctx_dlpack_api) # backward compact with error already set case # TODO(tqchen): remove after we move beyond a few versions. if c_api_ret_code == -2: raise raise_existing_error() # epecial handle env error already set error = move_from_last_error() if error.kind == "EnvErrorAlreadySet": raise raise_existing_error() raise error.py_error() @staticmethod def __from_extern_c__( c_symbol: int, *, keep_alive_object: object | None = None ) -> Function: """Convert a function from extern C address. Parameters ---------- c_symbol : int Function pointer to the safe call function. The function pointer must ignore the first argument, which is the function handle. keep_alive_object : object Optional object to be captured and kept alive. Usually this can be the execution engine that JIT-compiled the function to ensure we keep the execution environment alive as long as the function is alive. Returns ------- Function The function object. """ cdef TVMFFIObjectHandle chandle # must first convert to int64_t cdef int64_t c_symbol_as_long_long = c_symbol cdef void* safe_call_addr_ptr = c_symbol_as_long_long cdef PyObject* closure_py_obj = keep_alive_object cdef int ret_code if keep_alive_object is None: ret_code = TVMFFIFunctionCreate( NULL, safe_call_addr_ptr, NULL, &chandle ) else: # otherwise, we use Python object Py_INCREF(keep_alive_object) ret_code = TVMFFIFunctionCreate( closure_py_obj, safe_call_addr_ptr, TVMFFIPyObjectDeleter, &chandle ) if ret_code != 0: # cleanup during error handling Py_DECREF(keep_alive_object) CHECK_CALL(ret_code) func = Function.__new__(Function) (func).chandle = chandle return func @staticmethod def __from_mlir_packed_safe_call__( mlir_packed_symbol: int, *, keep_alive_object: object | None = None ) -> Function: """Convert a function from MLIR packed safe call function pointer. Parameters ---------- mlir_packed_symbol : int Function pointer to the MLIR packed call function that represents a safe call function. keep_alive_object : object Optional object to be captured and kept alive. Usually this can be the execution engine that JIT-compiled the function to ensure we keep the execution environment alive as long as the function is alive. Returns ------- Function The function object. """ cdef TVMFFIObjectHandle chandle # must first convert to int64_t cdef int64_t c_symbol_as_long_long = mlir_packed_symbol cdef void* packed_call_addr_ptr = c_symbol_as_long_long cdef PyObject* keepalive_py_obj if keep_alive_object is None: keepalive_py_obj = NULL else: keepalive_py_obj = keep_alive_object cdef void* mlir_packed_safe_call = TVMFFIPyMLIRPackedSafeCallCreate( packed_call_addr_ptr, keepalive_py_obj ) cdef int ret_code ret_code = TVMFFIFunctionCreate( mlir_packed_safe_call, TVMFFIPyMLIRPackedSafeCallInvoke, TVMFFIPyMLIRPackedSafeCallDeleter, &chandle ) if ret_code != 0: # cleanup during error handling TVMFFIPyMLIRPackedSafeCallDeleter(mlir_packed_safe_call) CHECK_CALL(ret_code) func = Function.__new__(Function) (func).chandle = chandle return func def _register_global_func(name: str, pyfunc: Callable[..., Any] | Function, override: bool) -> Function: cdef TVMFFIObjectHandle chandle cdef int c_api_ret_code cdef int ioverride = override cdef ByteArrayArg name_arg = ByteArrayArg(c_str(name)) if not isinstance(pyfunc, Function): pyfunc = _convert_to_ffi_func(pyfunc) CHECK_CALL(TVMFFIFunctionSetGlobal(name_arg.cptr(), (pyfunc).chandle, ioverride)) return pyfunc def _get_global_func(name: str, allow_missing: bool): cdef TVMFFIObjectHandle chandle cdef ByteArrayArg name_arg = ByteArrayArg(c_str(name)) CHECK_CALL(TVMFFIFunctionGetGlobal(name_arg.cptr(), &chandle)) if chandle != NULL: ret = Function.__new__(Function) (ret).chandle = chandle return ret if allow_missing: return None raise ValueError("Cannot find global function %s" % name) cdef inline int _convert_to_opaque_object_handle( object pyobject, TVMFFIObjectHandle* out_handle ) except -1: """Convert a python object to TVM FFI opaque object handle""" Py_INCREF(pyobject) CHECK_CALL(TVMFFIObjectCreateOpaque( (pyobject), kTVMFFIOpaquePyObject, TVMFFIPyObjectDeleter, out_handle)) return 0 def _convert_to_opaque_object(object pyobject: Any) -> OpaquePyObject: """Convert a python object to TVM FFI opaque object""" cdef TVMFFIObjectHandle chandle _convert_to_opaque_object_handle(pyobject, &chandle) ret = OpaquePyObject.__new__(OpaquePyObject) (ret).chandle = chandle return ret cdef extern from *: """ static void TVMFFITestingCallDeleterWithoutThreadState(void* py_obj) { PyThreadState* thread_state = PyEval_SaveThread(); TVMFFIPyObjectDeleter(py_obj); PyEval_RestoreThread(thread_state); } """ void TVMFFITestingCallDeleterWithoutThreadState(void* py_obj) def _testing_drop_last_ref_without_thread_state() -> None: """Drop the last Python ref from a detached-thread-state region.""" cdef object pyobject = {} cdef PyObject* py_obj = pyobject Py_INCREF(pyobject) pyobject = None TVMFFITestingCallDeleterWithoutThreadState(py_obj) def _print_debug_info() -> None: """Get the size of the arg dispatch map""" cdef size_t size = TVMFFIPyGetArgDispatchMapSize() print(f"TVMFFIPyGetArgDispatchMapSize: {size}") cdef Function _OBJECT_FROM_JSON_GRAPH_STR = _get_global_func("ffi.FromJSONGraphString", True) cdef Function _OBJECT_TO_JSON_GRAPH_STR = _get_global_func("ffi.ToJSONGraphString", True) cdef Function _CONSTRUCTOR_ARRAY = _get_global_func("ffi.Array", True) cdef Function _CONSTRUCTOR_MAP = _get_global_func("ffi.Map", True) cdef Function _FFI_CONTAINER_FIND_FIRST_NON_CPU_DEVICE = _get_global_func("ffi.ContainerFindFirstNonCPUDevice", True) cdef Function _MAKE_FILED_GETTER = _get_global_func("ffi.MakeFieldGetter", True) cdef Function _MAKE_FIELD_SETTER = _get_global_func("ffi.MakeFieldSetter", True) cdef Function _PYCLS_REGISTER = _get_global_func("ffi._PyClassRegisterTypeAttrColumns", True) MISSING = _get_global_func("ffi.GetInvalidObject", False)() KWARGS = _get_global_func("ffi.GetKwargsObject", False)() tvm-ffi-0.1.12/python/tvm_ffi/cython/object.pxi000066400000000000000000000751621521067262500214330ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import json from abc import ABCMeta from typing import Any _CLASS_OBJECT = None def _set_class_object(cls): global _CLASS_OBJECT _CLASS_OBJECT = cls def object_repr(obj: "CObject") -> str: """Return a human-readable repr of *obj* via ``ffi.ReprPrint``. Falls back to ``TypeName(handle)`` if ``ReprPrint`` is unavailable. """ if (obj).chandle == NULL: return type(obj).__name__ + "(chandle=None)" try: from tvm_ffi._ffi_api import ReprPrint return str(ReprPrint(obj)) except Exception: # noqa: BLE001 # Silently fall back: repr must never raise. return type(obj).__name__ + "(" + str(obj.__chandle__()) + ")" def _new_object(cls): """Helper function for pickle""" return cls.__new__(cls) class ObjectConvertible: """Base class for Python classes convertible to :class:`Object`. Subclasses implement :py:meth:`asobject` to produce an :class:`Object` instance used by the FFI runtime. """ def asobject(self) -> "Object": """Return an :class:`Object` view of this value. This method is used by the conversion helpers (e.g. :func:`tvm_ffi.convert`) when a Python value needs to be passed into FFI calls. Returns ------- tvm_ffi.core.Object """ raise NotImplementedError() class ObjectRValueRef: """Rvalue reference wrapper used to express move semantics. Instances are created from :py:meth:`Object._move` and signal to the FFI layer that ownership of the underlying handle can be transferred. Parameters ---------- obj : tvm_ffi.core.Object The source object from which to move the underlying handle. """ __slots__ = ["obj"] def __init__(self, obj): self.obj = obj cdef class CObject: """Cython base class for TVM FFI objects. This extension type owns the low-level handle. Prefer subclassing :class:`Object` in Python to enforce slots policy. """ __slots__ = () cdef void* chandle def __cinit__(self): # initialize chandle to NULL to avoid leak in # case of error before chandle is set self.chandle = NULL def __dealloc__(self): if self.chandle != NULL: CHECK_CALL(TVMFFIObjectDecRef(self.chandle)) self.chandle = NULL def __ctypes_handle__(self) -> object: return ctypes_handle(self.chandle) def __chandle__(self) -> int: cdef uint64_t chandle = self.chandle return chandle @property def id_(self) -> int: """The integer address of the underlying FFI handle. Alias for :py:meth:`__chandle__`. Returns ``0`` when the handle is NULL. """ cdef uint64_t chandle = self.chandle return chandle def __reduce__(self): cls = type(self) return (_new_object, (cls,), self.__getstate__()) def __getstate__(self) -> dict[str, Any]: if _OBJECT_TO_JSON_GRAPH_STR is None: raise RuntimeError("ffi.ToJSONGraphString is not registered, make sure build project with extra API") if not self.__chandle__() == 0: # need to explicit convert to str in case String # returned and triggered another infinite recursion in get state return {"handle": str(_OBJECT_TO_JSON_GRAPH_STR(self, None))} return {"handle": None} def __setstate__(self, state: dict[str, Any]) -> None: # pylint: disable=assigning-non-slot, assignment-from-no-return if _OBJECT_FROM_JSON_GRAPH_STR is None: raise RuntimeError("ffi.FromJSONGraphString is not registered, make sure build project with extra API") handle = state["handle"] if handle is not None: self.__init_handle_by_constructor__(_OBJECT_FROM_JSON_GRAPH_STR, handle) else: self.chandle = NULL def __repr__(self) -> str: return object_repr(self) def __eq__(self, other: object) -> bool: return self.same_as(other) def __ne__(self, other: object) -> bool: return not self.__eq__(other) def __hash__(self) -> int: cdef uint64_t hash_value = self.chandle return hash_value def same_as(self, other: object) -> bool: return isinstance(other, CObject) and self.chandle == (other).chandle def is_(self, other: object) -> bool: """Return ``True`` if both references point to the same FFI handle. Alias for :py:meth:`same_as`. Checks identity of the underlying handle rather than performing a structural, value-based comparison. """ return isinstance(other, CObject) and self.chandle == (other).chandle def __move_handle_from__(self, other: CObject) -> None: self.chandle = (other).chandle (other).chandle = NULL def __init_handle_by_constructor__(self, fconstructor: Any, *args: Any) -> None: # avoid error raised during construction. self.chandle = NULL cdef void* chandle ConstructorCall( (fconstructor).chandle, args, &chandle, NULL) self.chandle = chandle cdef class CContainerBase(CObject): """Cython base for container types that support lazy DLPack conversion. Stores a ``DLPackExchangeAPI*`` tag so that element access on a returned container can automatically convert ``ffi.Tensor`` to the framework tensor type (e.g. ``torch.Tensor``). """ # Raw pointer to the DLPack exchange API struct. Not ref-counted. # # Lifetime safety: the two sources of this pointer are both # effectively process-lifetime: # # 1. __dlpack_c_exchange_api__ (e.g. torch.Tensor) — points to a # static struct in the framework's C++ runtime. The source # type is kept alive by _DISPATCH_TYPE_KEEP_ALIVE (set in # TVMFFICyArgSetterFactory), which prevents module unloading. # # 2. GetTorchFallbackExchangeAPI() — returns the address of a # module-level Cython static; lives for the entire process. # # The DLPack spec also mandates that DLPackExchangeAPI* must stay # alive throughout the lifetime of the process (dlpack.h line 600). cdef const DLPackExchangeAPI* _dlpack_exchange_api def __cinit__(self): self._dlpack_exchange_api = NULL class _ObjectSlotsMeta(ABCMeta): def __new__(mcls, name: str, bases: tuple[type, ...], ns: dict[str, Any], **kwargs: Any): if "__slots__" not in ns: ns["__slots__"] = () return super().__new__(mcls, name, bases, ns, **kwargs) def __init__(cls, name: str, bases: tuple[type, ...], ns: dict[str, Any], **kwargs: Any): super().__init__(name, bases, ns, **kwargs) class Object(CObject, metaclass=_ObjectSlotsMeta): """Base class of all TVM FFI objects. This is the root Python type for objects backed by the TVM FFI runtime. Each instance references a handle to a C++ runtime object. Python subclasses typically correspond to C++ runtime types and are registered via :py:meth:`tvm_ffi.register_object`. Notes ----- - Equality of two :py:class:`Object` instances uses underlying handle identity unless an overridden implementation is provided on the concrete type. Use :py:meth:`same_as` to check whether two references point to the same underlying object. - Subclasses that omit ``__slots__`` get ``__slots__ = ()`` injected automatically by the metaclass. To allow a per-instance ``__dict__``, declare ``__slots__ = ("__dict__",)`` explicitly in the class body. - Most users interact with subclasses (e.g. :class:`Tensor`, :class:`Function`) rather than :py:class:`Object` directly. Examples -------- Constructing objects is typically performed by Python wrappers that call into registered constructors on the FFI side. .. code-block:: python import tvm_ffi.testing # Acquire a testing object constructed through FFI obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=12) assert isinstance(obj, tvm_ffi.Object) assert obj.same_as(obj) Subclasses can declare explicit slots when needed. .. code-block:: python @tvm_ffi.register_object("my.MyObject") class MyObject(tvm_ffi.Object): __slots__ = () Subclasses that need a per-instance ``__dict__`` (e.g. for attribute caching) can opt in explicitly. .. code-block:: python @tvm_ffi.register_object("my.MyDynObject") class MyDynObject(tvm_ffi.Object): __slots__ = ("__dict__",) """ __slots__ = () def same_as(self, other: object) -> bool: """Return ``True`` if both references point to the same object. This checks identity of the underlying FFI handle rather than performing a structural, value-based comparison. Parameters ---------- other : object The object to compare against. Returns ------- bool Examples -------- .. code-block:: python import tvm_ffi.testing x = tvm_ffi.testing.create_object("testing.TestObjectBase") y = x z = tvm_ffi.testing.create_object("testing.TestObjectBase") assert x.same_as(y) assert not x.same_as(z) """ return CObject.same_as(self, other) def _move(self) -> ObjectRValueRef: """Create an rvalue reference that transfers ownership. The returned :class:`ObjectRValueRef` indicates move semantics to the FFI layer, and is intended for performance-sensitive paths that wish to avoid an additional retain/release pair. Notes ----- After a successful move, the original object should be treated as invalid on the FFI side. Do not rely on the handle after transferring. Returns ------- ObjectRValueRef The rvalue reference wrapper. """ return ObjectRValueRef(self) def __move_handle_from__(self, other: CObject) -> None: """Steal the FFI handle from ``other``. Internal helper used by the runtime to implement move semantics. Users should prefer :py:meth:`_move`. """ CObject.__move_handle_from__(self, other) def __init_handle_by_constructor__(self, fconstructor: Any, *args: Any) -> None: """Initialize the handle by calling constructor function. Parameters ---------- fconstructor : Function Constructor function. args: list of objects The arguments to the constructor Notes ----- We have a special calling convention to call constructor functions. So the return handle is directly set into the Node object instead of creating a new Node. """ CObject.__init_handle_by_constructor__(self, fconstructor, *args) cdef class OpaquePyObject(CObject): """Wrapper that carries an arbitrary Python object across the FFI. The contained object is held with correct reference counting, and can be recovered on the Python side using :py:meth:`pyobject`. Notes ----- ``OpaquePyObject`` is useful when a Python value must traverse the FFI boundary without conversion into a native FFI type. """ __slots__ = () def pyobject(self) -> object: """Return the original Python object held by this wrapper.""" cdef object obj cdef PyObject* py_handle py_handle = (TVMFFIOpaqueObjectGetCellPtr(self.chandle).handle) obj = py_handle return obj class PyNativeObject: """Base class for TVM objects that also inherit Python builtins. This mixin is used by Python-native proxy types such as :class:`String` and :class:`Bytes`, which subclass :class:`str` and :class:`bytes` respectively while also carrying an attached FFI object for zero-copy exchange with the runtime when beneficial. """ __slots__ = [] def __init_cached_object_by_constructor__(self, fconstructor: Any, *args: Any) -> None: """Initialize the internal _tvm_ffi_cached_object by calling constructor function. Parameters ---------- fconstructor : Function Constructor function. args: list of objects The arguments to the constructor Note ---- We have a special calling convention to call constructor functions. So the return object is directly set into the object """ obj = _CLASS_OBJECT.__new__(_CLASS_OBJECT) obj.__init_handle_by_constructor__(fconstructor, *args) self._tvm_ffi_cached_object = obj def _object_type_key_to_index(str type_key): """get the type index of object class""" cdef int32_t tidx type_key_arg = ByteArrayArg(c_str(type_key)) if TVMFFITypeKeyToIndex(type_key_arg.cptr(), &tidx) == 0: return tidx return None cdef inline str _type_index_to_key(int32_t tindex): """get the type key of object class""" cdef const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(tindex) cdef const TVMFFIByteArray* type_key if info == NULL: return "" type_key = &(info.type_key) return bytearray_to_str(type_key) cdef inline object make_ret_opaque_object(TVMFFIAny result): obj = OpaquePyObject.__new__(OpaquePyObject) (obj).chandle = result.v_obj return obj.pyobject() cdef inline object make_fallback_cls_for_type_index(int32_t type_index): cdef str type_key = _type_index_to_key(type_index) cdef object type_info = _lookup_or_register_type_info_from_type_key(type_key) cdef object parent_type_info = type_info.parent_type_info assert type_info.type_cls is None # Ensure parent classes are created first assert parent_type_info is not None if parent_type_info.type_cls is None: # recursively create parent class first make_fallback_cls_for_type_index(parent_type_info.type_index) assert parent_type_info.type_cls is not None # Create `type_info.type_cls` now class cls(parent_type_info.type_cls): __slots__ = () cls.__tvm_ffi_type_info__ = type_info cls.__name__ = type_key.split(".")[-1] cls.__qualname__ = type_key cls.__module__ = ".".join(type_key.split(".")[:-1]) cls.__doc__ = ( f"Auto-generated fallback class for {type_key}.\n" "This class is generated because the class is not registered.\n" "Please do not use this class directly, instead register the class\n" "using `register_object` decorator." ) for field in type_info.fields: setattr(cls, field.name, field.as_property(cls)) for method in type_info.methods: setattr(cls, method.name, method.as_callable(cls)) # Update the registry type_info.type_cls = cls _update_registry(type_index, type_key, type_info, cls) return cls cdef inline object make_ret_object(TVMFFIAny result): cdef int32_t type_index cdef object cls, obj type_index = result.type_index if type_index < len(TYPE_INDEX_TO_CLS) and (cls := TYPE_INDEX_TO_CLS[type_index]) is not None: if issubclass(cls, PyNativeObject): obj = Object.__new__(Object) (obj).chandle = result.v_obj return cls.__from_tvm_ffi_object__(cls, obj) else: # Slow path: object is not found in registered entry # In this case create a dummy stub class for future usage. # For every unregistered class, this slow path will be triggered only once. cls = make_fallback_cls_for_type_index(type_index) obj = cls.__new__(cls) (obj).chandle = result.v_obj return obj cdef _get_method_from_method_info(const TVMFFIMethodInfo* method): cdef TVMFFIAny result CHECK_CALL(TVMFFIAnyViewToOwnedAny(&(method.method), &result)) return make_ret(result) cdef _type_info_create_from_type_key(object type_cls, str type_key): cdef const TVMFFIFieldInfo* field cdef const TVMFFIMethodInfo* method cdef const TVMFFITypeInfo* info cdef int32_t type_index cdef list ancestors = [] cdef int ancestor cdef dict metadata_obj cdef object fields = [] cdef object methods = [] cdef str type_schema_json cdef FieldGetter getter cdef FieldSetter setter cdef bint has_default cdef bint default_from_factory cdef TVMFFIAny owned_default cdef object c_default cdef object c_default_factory cdef object c_structural_eq cdef ByteArrayArg type_key_arg = ByteArrayArg(c_str(type_key)) # NOTE: `type_key_arg` must be kept alive until after the call to `TVMFFITypeKeyToIndex`, # because Cython doesn't defer the destruction of `type_key_arg` until after the call. if TVMFFITypeKeyToIndex(type_key_arg.cptr(), &type_index) != 0: raise ValueError(f"Cannot find type key: {type_key}") info = TVMFFIGetTypeInfo(type_index) for i in range(info.num_fields): field = &(info.fields[i]) getter = FieldGetter.__new__(FieldGetter) (getter).getter = field.getter (getter).offset = field.offset setter = FieldSetter.__new__(FieldSetter) (setter).setter = field.setter (setter).offset = field.offset (setter).flags = field.flags metadata_obj = json.loads(bytearray_to_str(&field.metadata)) if field.metadata.size != 0 else {} # Decode the static default value or factory (if any) registered by C++. has_default = (field.flags & kTVMFFIFieldFlagBitMaskHasDefault) != 0 default_from_factory = (field.flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory) != 0 c_default = MISSING c_default_factory = MISSING if has_default: CHECK_CALL(TVMFFIAnyViewToOwnedAny(&field.default_value_or_factory, &owned_default)) if default_from_factory: c_default_factory = make_ret(owned_default) else: c_default = make_ret(owned_default) # Decode SEqHashIgnore / SEqHashDef* into the Field.structural_eq vocabulary. if (field.flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) != 0: c_structural_eq = "ignore" elif (field.flags & kTVMFFIFieldFlagBitMaskSEqHashDefRecursive) != 0: c_structural_eq = "def-recursive" elif (field.flags & kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive) != 0: c_structural_eq = "def-non-recursive" else: c_structural_eq = None fields.append( TypeField( name=bytearray_to_str(&field.name), doc=bytearray_to_str(&field.doc) if field.doc.size != 0 else None, size=field.size, offset=field.offset, frozen=(field.flags & kTVMFFIFieldFlagBitMaskWritable) == 0, metadata=metadata_obj, getter=getter, setter=setter, c_init=(field.flags & kTVMFFIFieldFlagBitMaskInitOff) == 0, c_kw_only=(field.flags & kTVMFFIFieldFlagBitMaskKwOnly) != 0, c_has_default=has_default, c_default=c_default, c_default_factory=c_default_factory, c_repr=(field.flags & kTVMFFIFieldFlagBitMaskReprOff) == 0, c_compare=(field.flags & kTVMFFIFieldFlagBitMaskCompareOff) == 0, c_hash=(field.flags & kTVMFFIFieldFlagBitMaskHashOff) == 0, c_structural_eq=c_structural_eq, ) ) for i in range(info.num_methods): method = &(info.methods[i]) metadata_obj = json.loads(bytearray_to_str(&method.metadata)) if method.metadata.size != 0 else {} methods.append( TypeMethod( name=bytearray_to_str(&method.name), doc=bytearray_to_str(&method.doc) if method.doc.size != 0 else None, func=_get_method_from_method_info(method), is_static=(method.flags & kTVMFFIFieldFlagBitMaskIsStaticMethod) != 0, metadata=metadata_obj, ) ) for i in range(info.type_depth): ancestor = info.type_ancestors[i].type_index ancestors.append(ancestor) return TypeInfo( type_cls=type_cls, type_index=type_index, type_key=bytearray_to_str(&info.type_key), type_ancestors=ancestors, fields=fields, methods=methods, parent_type_info=None, ) cdef _update_registry(int type_index, object type_key, object type_info, object type_cls): cdef int extra = type_index + 1 - len(TYPE_INDEX_TO_INFO) assert len(TYPE_INDEX_TO_INFO) == len(TYPE_INDEX_TO_CLS) if extra > 0: TYPE_INDEX_TO_INFO.extend([None] * extra) TYPE_INDEX_TO_CLS.extend([None] * extra) TYPE_INDEX_TO_CLS[type_index] = type_cls TYPE_INDEX_TO_INFO[type_index] = type_info TYPE_KEY_TO_INFO[type_key] = type_info if type_cls is not None: TYPE_CLS_TO_INFO[type_cls] = type_info def _register_object_by_index(int type_index, object type_cls): global TYPE_INDEX_TO_INFO, TYPE_KEY_TO_INFO, TYPE_INDEX_TO_CLS cdef str type_key = _type_index_to_key(type_index) cdef object info = _type_info_create_from_type_key(type_cls, type_key) _update_registry(type_index, type_key, info, type_cls) return info def _set_type_cls(object type_info, object type_cls): global TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS, TYPE_CLS_TO_INFO assert type_info.type_cls is None, f"Type already registered for {type_info.type_key}" assert TYPE_INDEX_TO_INFO[type_info.type_index] is type_info assert TYPE_KEY_TO_INFO[type_info.type_key] is type_info type_info.type_cls = type_cls TYPE_INDEX_TO_CLS[type_info.type_index] = type_cls TYPE_CLS_TO_INFO[type_cls] = type_info def _lookup_or_register_type_info_from_type_key(type_key: str) -> TypeInfo: if info := TYPE_KEY_TO_INFO.get(type_key, None): return info info = _type_info_create_from_type_key(None, type_key) _update_registry(info.type_index, type_key, info, None) return info def _register_py_class(parent_type_info, str type_key, object type_cls): """Register a new Python-defined TVM-FFI type. Allocates a dynamic type index for *type_key* as a child of *parent_type_info* and registers it in the global type tables. Parameters ---------- parent_type_info : TypeInfo The parent type's TypeInfo (e.g., Object's TypeInfo). type_key : str The unique type key string for the new type. type_cls : type The Python class to associate with this type. Returns ------- TypeInfo The newly created TypeInfo with ``fields=None`` (pending registration). Raises ------ ValueError If *type_key* is already registered. """ # Reject duplicate type keys if type_key in TYPE_KEY_TO_INFO: raise ValueError( f"Type key '{type_key}' is already registered" ) cdef int32_t parent_type_index = parent_type_info.type_index cdef int32_t parent_type_depth = len(parent_type_info.type_ancestors) cdef int32_t type_depth = parent_type_depth + 1 cdef ByteArrayArg type_key_arg = ByteArrayArg(c_str(type_key)) cdef int32_t type_index # Allocate a new type index # static_type_index=-1 means dynamic allocation # num_child_slots=0, child_slots_can_overflow=1 type_index = TVMFFITypeGetOrAllocIndex( type_key_arg.cptr(), -1, # static_type_index (dynamic) type_depth, 0, # num_child_slots 1, # child_slots_can_overflow parent_type_index, ) # Build ancestors list cdef list ancestors = list(parent_type_info.type_ancestors) ancestors.append(parent_type_index) # Create TypeInfo with fields=None (pending _register_fields call) cdef object info = TypeInfo( type_cls=type_cls, type_index=type_index, type_key=type_key, type_ancestors=ancestors, fields=None, methods=[], parent_type_info=parent_type_info, ) _update_registry(type_index, type_key, info, type_cls) return info def _rollback_py_class(object type_info): """Roll back a ``_register_py_class`` call from the Python-level registry. Called by ``@py_class`` when phase-2 (field validation) fails, so the type key can be reused after the user fixes the error. The C-level type index is permanently consumed (cannot be reclaimed), but the Python dicts are cleaned up so that a retry does not hit "already registered". """ cdef int32_t idx = type_info.type_index cdef str key = type_info.type_key cdef object cls = type_info.type_cls TYPE_KEY_TO_INFO.pop(key, None) if cls is not None: TYPE_CLS_TO_INFO.pop(cls, None) if 0 <= idx < len(TYPE_INDEX_TO_INFO): TYPE_INDEX_TO_INFO[idx] = None TYPE_INDEX_TO_CLS[idx] = None def _lookup_type_attr(type_index: int32_t, attr_key: str) -> Any: cdef ByteArrayArg attr_key_bytes = ByteArrayArg(c_str(attr_key)) cdef const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&attr_key_bytes.cdata) cdef TVMFFIAny data cdef int32_t offset if column == NULL: return None offset = type_index - column.begin_index if offset < 0 or offset >= column.size: return None CHECK_CALL(TVMFFIAnyViewToOwnedAny(&(column.data[offset]), &data)) return make_ret(data) def _register_type_attr(type_index: int32_t, attr_key: str, value: object) -> None: """Register a value for the ``(type_index, attr_key)`` slot. Wraps :c:func:`TVMFFITypeRegisterAttr`, which raises :class:`RuntimeError` if a value is already registered for the slot. To update the stored value, register a mutable container (e.g. ``Dict``/``List``) once and mutate it in place on subsequent calls. ``TVMFFIPyPyObjectToFFIAny`` produces a non-owning :c:type:`TVMFFIAny` view of *value*; ``TVMFFITypeRegisterAttr`` incref's the underlying object when it stores the slot, so no explicit refcount management is needed here. """ cdef ByteArrayArg attr_key_bytes = ByteArrayArg(c_str(attr_key)) cdef TVMFFIAny temp cdef int c_api_ret_code temp.type_index = kTVMFFINone temp.v_int64 = 0 TVMFFIPyPyObjectToFFIAny( value, &temp, &c_api_ret_code, ) CHECK_CALL(c_api_ret_code) CHECK_CALL(TVMFFITypeRegisterAttr(type_index, &attr_key_bytes.cdata, &temp)) def _type_cls_to_type_info(type_cls: type) -> TypeInfo | None: return TYPE_CLS_TO_INFO.get(type_cls, None) cdef list TYPE_INDEX_TO_CLS = [] cdef list TYPE_INDEX_TO_INFO = [] cdef dict TYPE_CLS_TO_INFO = {} cdef dict TYPE_KEY_TO_INFO = {} _set_class_object(Object) # --------------------------------------------------------------------------- # CAny: Owned TVMFFIAny value container # --------------------------------------------------------------------------- cdef class CAny: """Owned :c:type:`TVMFFIAny` value container. Holds sole ownership of the underlying value. For object types (``type_index >= kTVMFFIStaticObjectBegin``), the reference is properly ref-counted and released in ``__dealloc__``. Use :meth:`to_py` to recover the Python object. """ cdef TVMFFIAny cdata def __cinit__(self): """Initialize the contained value to ``None``.""" self.cdata.type_index = kTVMFFINone self.cdata.v_int64 = 0 def __init__(self, value=None): """Pack a Python value into an owned :c:type:`TVMFFIAny`. Uses ``TVMFFIPyPyObjectToFFIAny`` to produce a non-owning AnyView, then ``TVMFFIAnyViewToOwnedAny`` to convert to an owned Any. Parameters ---------- value : object, optional The Python value to pack. When ``None`` (the default), the container stays in the ``kTVMFFINone`` state set by ``__cinit__``. """ if value is None: return cdef TVMFFIAny temp cdef int c_api_ret_code temp.type_index = kTVMFFINone temp.v_int64 = 0 TVMFFIPyPyObjectToFFIAny( value, &temp, &c_api_ret_code ) CHECK_CALL(c_api_ret_code) CHECK_CALL(TVMFFIAnyViewToOwnedAny(&temp, &self.cdata)) def __dealloc__(self): """Release owned object reference, if any.""" if self.cdata.type_index >= kTVMFFIStaticObjectBegin: if self.cdata.v_obj != NULL: CHECK_CALL(TVMFFIObjectDecRef(self.cdata.v_obj)) self.cdata.v_obj = NULL @property def type_index(self) -> int: """The TVM FFI type index of the contained value.""" return self.cdata.type_index def __repr__(self) -> str: """Return a developer-friendly representation.""" cdef int32_t ti = self.cdata.type_index if ti == kTVMFFINone: return "CAny(None)" elif ti == kTVMFFIInt: return f"CAny(int={self.cdata.v_int64})" elif ti == kTVMFFIFloat: return f"CAny(float={self.cdata.v_float64})" elif ti == kTVMFFIBool: return f"CAny(bool={bool(self.cdata.v_int64)})" elif ti >= kTVMFFIStaticObjectBegin: return f"CAny(object, type_index={ti})" else: return f"CAny(type_index={ti})" cpdef object _to_py_class_value(CAny self): """Convert a CAny to a Python object (module-level cdef for direct C dispatch).""" cdef TVMFFIAny copy = self.cdata if copy.type_index >= kTVMFFIStaticObjectBegin: if copy.v_obj != NULL: TVMFFIObjectIncRef(copy.v_obj) cdef object result = make_ret(copy) # Promote inline SmallStr/SmallBytes to their FFI wrapper types # so that convert().to_py() always yields tvm_ffi.String / tvm_ffi.Bytes. if copy.type_index == kTVMFFISmallStr: return String(result) if copy.type_index == kTVMFFISmallBytes: return Bytes(result) return result tvm-ffi-0.1.12/python/tvm_ffi/cython/pycallback.pxi000066400000000000000000000272011521067262500222610ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ---------------------------------------------------------------------------- # Cython-side implementation of the C++ -> Python callback path. # # The C++ side (tvm_ffi_python_helpers.h + TVMFFIPyCallManager::PyCallback) # calls into the functions defined here via function pointers registered in # callback_arg_dispatch_table_. Each per-type setter takes a borrowed # AnyView and produces a new-reference PyObject*. # # The caller will grab the thread state before caling into each individual setter. # # This file also hosts `_convert_to_ffi_func`, the Cython entry point that # wraps a Python callable as a FFI Function backed by a TVMFFIPyCallback # closure (see TVMFFIPyConvertPyCallback in the header). # ---------------------------------------------------------------------------- cdef int TVMFFIPyCallbackArgSetterTensor_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFITensor -> ffi.Tensor or torch.Tensor (via DLPack). The DLPack branch is inlined rather than delegated to ``make_tensor_from_chandle`` so we can pass a borrowed chandle: ``TensorObj::ToDLPackVersioned`` incs internally, so the inc/dec pair that ``make_tensor_from_chandle`` requires on an owned chandle is pure waste here. The non-DLPack branch upgrades to owned and reuses ``make_tensor_from_chandle`` for consistency with the RValueRef path. """ cdef TVMFFIObjectHandle chandle = arg.v_ptr cdef DLManagedTensorVersioned* dlpack cdef void* py_obj if api != NULL and api.managed_tensor_to_py_object_no_sync != NULL: if TVMFFITensorToDLPackVersioned(chandle, &dlpack) == 0: try: api.managed_tensor_to_py_object_no_sync(dlpack, &py_obj) except Exception: dlpack.deleter(dlpack) raise # py_obj already holds +1 from the DLPack import; transfer to caller. out[0] = py_obj return 0 # Non-DLPack path: upgrade borrowed -> owned, wrap via make_tensor_from_chandle. TVMFFIObjectIncRef(chandle) obj = make_tensor_from_chandle(chandle, NULL) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterObject_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for generic static object types (type_index >= kTVMFFIStaticObjectBegin).""" cdef TVMFFIObjectHandle chandle = arg.v_ptr TVMFFIObjectIncRef(chandle) try: obj = make_ret_object(arg[0]) if api != NULL and isinstance(obj, CContainerBase): (obj)._dlpack_exchange_api = api except BaseException: TVMFFIObjectDecRef(chandle) raise Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterOpaquePyObject_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIOpaquePyObject -> underlying Python object. Inlined equivalent of `make_ret_opaque_object`: reads the cell's Python handle directly, skipping the throwaway OpaquePyObject wrapper that would otherwise be created just to extract the handle. `arg` is borrowed, but the cell stays alive for the callback's duration, so the handle is safe to read without inc'ing the chandle. """ cdef PyObject* py_handle = TVMFFIOpaqueObjectGetCellPtr(arg.v_ptr).handle cdef object obj = py_handle Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterOpaquePtr_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIOpaquePtr -> ctypes.c_void_p.""" obj = ctypes_handle(arg.v_ptr) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterDataType_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIDataType -> DataType.""" obj = make_ret_dtype(arg[0]) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterDevice_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIDevice -> Device.""" obj = make_ret_device(arg[0]) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterDLTensorPtr_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIDLTensorPtr -> ffi.Tensor (via DLTensor pointer).""" obj = make_ret_dltensor(arg[0]) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterRawStr_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIRawStr -> Python str (UTF-8 decode). ``arg.v_c_str`` is a non-owning ``const char*`` pointer into C-side storage that remains valid for the duration of the callback invocation. We copy the contents into a Python ``str`` immediately, so there is no dangling-pointer concern after the setter returns. """ obj = arg.v_c_str.decode("utf-8") Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterByteArrayPtr_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIByteArrayPtr -> Python bytes. ``arg.v_ptr`` is a non-owning ``TVMFFIByteArray*`` pointer into C-side storage valid for the callback's lifetime. ``bytearray_to_bytes`` copies the raw bytes into a new Python ``bytes`` object immediately. """ obj = bytearray_to_bytes(arg.v_ptr) Py_INCREF(obj) out[0] = obj return 0 cdef int TVMFFIPyCallbackArgSetterRValueRef_( TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* api, const TVMFFIAny* arg, PyObject** out ) except -1: """Callback arg setter for kTVMFFIObjectRValueRef. For RValueRef, ``arg.v_ptr`` is an ``Object**`` (address of the caller's mutable slot holding the moved chandle), not the chandle itself. We read the slot, null it out to prevent a double-move, and wrap WITHOUT inc'ing (the move already gave us the +1). """ cdef TVMFFIObjectHandle* slot_ptr = arg.v_ptr cdef TVMFFIObjectHandle chandle = slot_ptr[0] if chandle == NULL: raise ValueError("RValueRef already moved") # mark as moved before constructing wrappers (so error paths don't double-move) slot_ptr[0] = NULL cdef int32_t actual_type_index = TVMFFIObjectGetTypeIndex(chandle) cdef TVMFFIAny synthesized synthesized.type_index = actual_type_index synthesized.zero_padding = 0 synthesized.v_int64 = 0 synthesized.v_ptr = chandle try: if actual_type_index == kTVMFFITensor: obj = make_tensor_from_chandle(chandle, api) else: obj = make_ret_object(synthesized) if api != NULL and isinstance(obj, CContainerBase): (obj)._dlpack_exchange_api = api except BaseException: # Caller's moved +1 needs to be released on error. TVMFFIObjectDecRef(chandle) raise Py_INCREF(obj) out[0] = obj return 0 cdef public int TVMFFICyCallbackArgSetterFactory(int32_t type_index, TVMFFIPyCallbackArgSetter* out) except -1: """Factory that creates callback arg setters for a given type index. POD setters live in tvm_ffi_python_helpers.h (header-inline); object-bearing setters are the Cython functions above. """ if type_index >= kTVMFFIStaticObjectBegin: if type_index == kTVMFFITensor: out.func = TVMFFIPyCallbackArgSetterTensor_ elif type_index == kTVMFFIOpaquePyObject: out.func = TVMFFIPyCallbackArgSetterOpaquePyObject_ else: out.func = TVMFFIPyCallbackArgSetterObject_ return 0 if type_index == kTVMFFINone: out.func = TVMFFIPyCallbackArgSetterNone_ elif type_index == kTVMFFIBool: out.func = TVMFFIPyCallbackArgSetterBool_ elif type_index == kTVMFFIInt: out.func = TVMFFIPyCallbackArgSetterInt_ elif type_index == kTVMFFIFloat: out.func = TVMFFIPyCallbackArgSetterFloat_ elif type_index == kTVMFFISmallStr: out.func = TVMFFIPyCallbackArgSetterSmallStr_ elif type_index == kTVMFFISmallBytes: out.func = TVMFFIPyCallbackArgSetterSmallBytes_ elif type_index == kTVMFFIOpaquePtr: out.func = TVMFFIPyCallbackArgSetterOpaquePtr_ elif type_index == kTVMFFIDataType: out.func = TVMFFIPyCallbackArgSetterDataType_ elif type_index == kTVMFFIDevice: out.func = TVMFFIPyCallbackArgSetterDevice_ elif type_index == kTVMFFIDLTensorPtr: out.func = TVMFFIPyCallbackArgSetterDLTensorPtr_ elif type_index == kTVMFFIObjectRValueRef: out.func = TVMFFIPyCallbackArgSetterRValueRef_ elif type_index == kTVMFFIByteArrayPtr: out.func = TVMFFIPyCallbackArgSetterByteArrayPtr_ elif type_index == kTVMFFIRawStr: out.func = TVMFFIPyCallbackArgSetterRawStr_ else: raise ValueError("Unhandled type index %d" % type_index) return 0 def _convert_to_ffi_func( object pyfunc: Callable[..., Any], object tensor_cls: object = None, ) -> Function: """Convert a python function to a TVM FFI Function. Parameters ---------- pyfunc : Callable The Python callable to wrap. Incref'd into a TVMFFIPyCallbackClosure. tensor_cls : type, optional If given, its ``__dlpack_c_exchange_api__`` capsule is threaded into the closure and used when constructing tensor return values inside the callback. Returns ------- Function The wrapped FFI function. """ cdef TVMFFIObjectHandle chandle cdef const DLPackExchangeAPI* dlpack_api = NULL if tensor_cls is not None: if not hasattr(tensor_cls, "__dlpack_c_exchange_api__"): raise TypeError( f"tensor_cls {tensor_cls!r} must expose __dlpack_c_exchange_api__" ) _get_dlpack_exchange_api( tensor_cls.__dlpack_c_exchange_api__, &dlpack_api ) CHECK_CALL(TVMFFIPyConvertPyCallback(pyfunc, dlpack_api, &chandle)) ret = Function.__new__(Function) (ret).chandle = chandle return ret tvm-ffi-0.1.12/python/tvm_ffi/cython/pyclass_type_converter.pxi000066400000000000000000001052361521067262500247670ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Type converter implementation for TypeSchema. Provides the ``_type_convert_impl`` function used by ``TypeSchema.convert`` and ``TypeSchema.check_value``. Each ``_TypeConverter`` dispatches directly to a Cython cdef function that returns a fully materialized :class:`CAny`. Container converters recurse into their child schemas, rebuild the target FFI container shape, and then wrap the final result in ``CAny``. """ import ctypes import os from numbers import Integral, Real from collections.abc import Mapping cdef object _INT64_MIN = -(1 << 63) cdef object _INT64_MAX = (1 << 63) - 1 cdef int _VALUE_PROTOCOL_MAX_DEPTH = 64 cdef str _TYPE_ATTR_FFI_CONVERT = "__ffi_convert__" cdef str _TYPE_ATTR_ENUM_VALUE_ENTRIES = "__ffi_enum_value_entries__" cdef class _TypeConverter ctypedef CAny (*_dispatch_fn_t)(_TypeConverter, object, bint*) except * # --------------------------------------------------------------------------- # cdef class _TypeConverter — holds dispatch state as C-level struct fields # --------------------------------------------------------------------------- cdef class _TypeConverter: """Pre-built converter holding a C function pointer and sub-converters.""" cdef _dispatch_fn_t dispatch cdef int32_t type_index cdef tuple subs cdef str err_hint cdef Function _fn_convert cdef bint _fn_convert_ready @property def fn_convert(self): cdef object attr assert self.type_index >= kTVMFFIStaticObjectBegin if not self._fn_convert_ready: attr = _lookup_type_attr(self.type_index, _TYPE_ATTR_FFI_CONVERT) if attr is not None: assert isinstance(attr, Function) self._fn_convert = attr else: self._fn_convert = None self._fn_convert_ready = True return self._fn_convert class _ConvertError(Exception): """Internal exception used to signal conversion failure.""" __slots__ = () def __init__(self, message): super().__init__(message) @property def message(self): return self.args[0] # --------------------------------------------------------------------------- # Converters (1/N): Simple value converters # --------------------------------------------------------------------------- cdef CAny _tc_convert_any(_TypeConverter _conv, object value, bint* changed) except *: """Any: accept any marshalable FFI value.""" cdef CAny packed assert _CLASS_DEVICE is not None assert _CLASS_DTYPE is not None packed = CAnyChecked(value, "Any", value) if not isinstance( value, ( type(None), bool, int, float, ctypes.c_void_p, String, Bytes, Tensor, DataType, CObject, _CLASS_DEVICE, _CLASS_DTYPE, ), ): changed[0] = True return packed cdef CAny _tc_convert_none(_TypeConverter _conv, object value, bint* changed) except *: """None accepts: None only.""" if value is None: return CAny(None) raise _ConvertError(f"expected None, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_int(_TypeConverter _conv, object value, bint* changed) except *: """int accepts: int, bool, Integral, __tvm_ffi_int__ protocol.""" cdef object ivalue if isinstance(value, bool): changed[0] = True return CAny(int(value)) if isinstance(value, int): if not (_INT64_MIN <= value <= _INT64_MAX): raise _ConvertError( f"integer {value} out of int64 range [{_INT64_MIN}, {_INT64_MAX}]" ) return CAny(value) if isinstance(value, Integral): try: ivalue = int(value) except Exception as err: raise _ConvertError(f"int() failed for {type(value).__qualname__}: {err}") from None if not (_INT64_MIN <= ivalue <= _INT64_MAX): raise _ConvertError( f"integer {ivalue} out of int64 range [{_INT64_MIN}, {_INT64_MAX}]" ) changed[0] = True return CAny(ivalue) if hasattr(type(value), "__tvm_ffi_int__"): changed[0] = True return CAnyChecked(value, "int", value) raise _ConvertError(f"expected int, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_float(_TypeConverter _conv, object value, bint* changed) except *: """float accepts: float, int, bool, Real, __tvm_ffi_float__ protocol.""" cdef object fvalue if isinstance(value, float): return CAny(value) if isinstance(value, (int, bool)): changed[0] = True return CAny(float(value)) if isinstance(value, (Integral, Real)): try: fvalue = float(value) except Exception as err: raise _ConvertError(f"float() failed for {type(value).__qualname__}: {err}") from None changed[0] = True return CAny(fvalue) if hasattr(type(value), "__tvm_ffi_float__"): changed[0] = True return CAnyChecked(value, "float", value) raise _ConvertError(f"expected float, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_bool(_TypeConverter _conv, object value, bint* changed) except *: """bool accepts: bool, int, Integral.""" cdef object bvalue if isinstance(value, bool): return CAny(value) if isinstance(value, Integral): # TODO: do we coerce Integral to bool? try: bvalue = bool(value) except Exception as err: raise _ConvertError(f"bool() failed for {type(value).__qualname__}: {err}") from None changed[0] = True return CAny(bvalue) raise _ConvertError(f"expected bool, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_str(_TypeConverter _conv, object value, bint* changed) except *: """str accepts: str only. Short strings use SmallStr (inline in Any).""" if isinstance(value, String): return CAny(value) if isinstance(value, str): changed[0] = True return CAny(value) raise _ConvertError(f"expected str, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_bytes(_TypeConverter _conv, object value, bint* changed) except *: """bytes accepts: bytes, bytearray. Short bytes use SmallBytes (inline in Any).""" if isinstance(value, Bytes): return CAny(value) if isinstance(value, bytes): changed[0] = True return CAny(value) if isinstance(value, bytearray): changed[0] = True return CAny(bytes(value)) raise _ConvertError(f"expected bytes, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_device(_TypeConverter _conv, object value, bint* changed) except *: """Device accepts: Device and __dlpack_device__ without __dlpack__.""" cdef object vtype = type(value) assert _CLASS_DEVICE is not None if isinstance(value, _CLASS_DEVICE): return CAny(value) if hasattr(vtype, "__dlpack_device__") and not hasattr(vtype, "__dlpack__"): changed[0] = True return CAnyChecked(value, "Device", value) raise _ConvertError(f"expected Device, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_dtype(_TypeConverter _conv, object value, bint* changed) except *: """dtype accepts: DataType, dtype wrapper, str and dtype protocols.""" cdef object dtype_value assert _CLASS_DTYPE is not None if isinstance(value, (DataType, _CLASS_DTYPE)): return CAny(value) if isinstance(value, str): try: dtype_value = DataType(value) except Exception: raise _ConvertError(f"expected dtype, got invalid dtype string {value!r}") from None changed[0] = True return CAny(dtype_value) if ( (torch is not None and isinstance(value, torch.dtype)) or (numpy is not None and isinstance(value, numpy.dtype)) or hasattr(value, "__dlpack_data_type__") ): changed[0] = True return CAnyChecked(value, "dtype", value) raise _ConvertError(f"expected dtype, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_opaque_ptr(_TypeConverter _conv, object value, bint* changed) except *: """ctypes.c_void_p accepts ctypes.c_void_p, None and opaque pointer protocols.""" cdef object vtype = type(value) if isinstance(value, ctypes.c_void_p): return CAny(value) # TODO: noticed that `OpaquePtr(nullptr) != None` - need to double check if this is correct if value is None: changed[0] = True return CAny(ctypes.c_void_p(None)) if hasattr(vtype, "__tvm_ffi_opaque_ptr__") or hasattr(vtype, "__cuda_stream__"): changed[0] = True return CAnyChecked(value, "ctypes.c_void_p", value) raise _ConvertError(f"expected ctypes.c_void_p, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_tensor(_TypeConverter _conv, object value, bint* changed) except *: """Tensor accepts Tensor, Tensor subtypes and DLPack exporters.""" cdef object vtype = type(value) if isinstance(value, Tensor): return CAny(value) if hasattr(vtype, "__dlpack__"): changed[0] = True return CAnyChecked(value, "Tensor", value) if os.environ.get("TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API", "0") != "1": if hasattr(vtype, "__dlpack_c_exchange_api__"): changed[0] = True return CAnyChecked(value, "Tensor", value) raise _ConvertError(f"expected Tensor, got {_tc_describe_value_type(value)}") cdef CAny _tc_convert_callable(_TypeConverter _conv, object value, bint* changed) except *: """Callable accepts Function and any Python callable.""" cdef Function func if isinstance(value, Function): return CAny(value) if callable(value): if isinstance(value, CObject): raise _ConvertError(f"expected Callable, got {_tc_describe_value_type(value)}") changed[0] = True return CAnyChecked(value, "Callable", value) raise _ConvertError(f"expected Callable, got {_tc_describe_value_type(value)}") # --------------------------------------------------------------------------- # Converters (2/N): Sequence/Mapping Containers # --------------------------------------------------------------------------- cdef CAny _tc_convert_array(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for Array[T]. Accepts Array or List CObjects (cross-type).""" from tvm_ffi.container import Array return _tc_convert_seq(conv, value, changed, Array) cdef CAny _tc_convert_list(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for List[T]. Accepts List or Array CObjects (cross-type).""" from tvm_ffi.container import List return _tc_convert_seq(conv, value, changed, List) cdef CAny _tc_convert_map(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for Map[K, V]. Accepts Map or Dict CObjects (cross-type).""" from tvm_ffi import _ffi_api from tvm_ffi.container import Map return _tc_convert_mapping(conv, value, changed, Map, _ffi_api.Map) cdef CAny _tc_convert_dict(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for Dict[K, V]. Accepts Dict or Map CObjects (cross-type).""" from tvm_ffi import _ffi_api from tvm_ffi.container import Dict return _tc_convert_mapping(conv, value, changed, Dict, _ffi_api.Dict) cdef CAny _tc_convert_seq(_TypeConverter conv, object value, bint* changed, object seq_type) except *: from tvm_ffi.container import Array, List cdef _TypeConverter elem_conv = conv.subs[0] if conv.subs else None if not isinstance(value, (Array, List, list, tuple)): raise _ConvertError(f"expected {seq_type.__name__}, got {_tc_describe_value_type(value)}") if elem_conv is None and isinstance(value, seq_type): return CAny(value) cdef list converted = [] cdef bint item_changed cdef object raw_item cdef int i = 0 cdef CAny item for raw_item in value: if elem_conv is not None: item_changed = False try: item = _type_convert_dispatch_with_fallback(elem_conv, raw_item, &item_changed) except _ConvertError as err: raise _ConvertError(f"element [{i}]: {err.message}") from None if item_changed: changed[0] = True converted.append(_to_py_class_value(item)) else: converted.append(raw_item) else: converted.append(raw_item) i += 1 if isinstance(value, seq_type) and not changed[0]: return CAny(value) changed[0] = True return CAnyChecked(seq_type(converted), seq_type.__name__, value) cdef CAny _tc_convert_tuple(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for tuple[T1, T2, ...] or bare tuple.""" cdef int i cdef int n cdef list converted = [] cdef CAny item cdef bint item_changed cdef object raw_item from tvm_ffi.container import Array, List if not isinstance(value, (Array, List, list, tuple)): raise _ConvertError(f"expected tuple, got {_tc_describe_value_type(value)}") if conv.subs is None: if isinstance(value, Array): return CAny(value) changed[0] = True return CAnyChecked(Array(value), "tuple", value) n = len(conv.subs) if len(value) != n: raise _ConvertError( f"expected tuple of length {n}, got {type(value).__name__} of length {len(value)}" ) for i in range(n): raw_item = value[i] item_changed = False try: item = _type_convert_dispatch_with_fallback( <_TypeConverter>(conv.subs[i]), raw_item, &item_changed, ) except _ConvertError as err: raise _ConvertError(f"element [{i}]: {err.message}") from None if item_changed: changed[0] = True converted.append(_to_py_class_value(item)) else: converted.append(raw_item) if isinstance(value, Array) and not changed[0]: return CAny(value) changed[0] = True return CAnyChecked(Array(converted), "tuple", value) cdef CAny _tc_convert_mapping( _TypeConverter conv, object value, bint* changed, object mapping_type, object constructor, ) except *: cdef _TypeConverter key_conv = conv.subs[0] if conv.subs else None cdef _TypeConverter val_conv = conv.subs[1] if conv.subs else None cdef list list_kvs = [] cdef CAny item cdef bint item_changed cdef object raw_key cdef object raw_val cdef object new_key cdef object new_val cdef object mapping cdef str expected = mapping_type.__name__ if not isinstance(value, Mapping): raise _ConvertError(f"expected {expected}, got {_tc_describe_value_type(value)}") if key_conv is None and val_conv is None and isinstance(value, mapping_type): return CAny(value) for raw_key, raw_val in value.items(): new_key = raw_key if key_conv is not None: item_changed = False try: item = _type_convert_dispatch_with_fallback(key_conv, raw_key, &item_changed) except _ConvertError as err: raise _ConvertError(f"key {raw_key!r}: {err.message}") from None if item_changed: changed[0] = True new_key = _to_py_class_value(item) new_val = raw_val if val_conv is not None: item_changed = False try: item = _type_convert_dispatch_with_fallback(val_conv, raw_val, &item_changed) except _ConvertError as err: raise _ConvertError(f"value for key {raw_key!r}: {err.message}") from None if item_changed: changed[0] = True new_val = _to_py_class_value(item) list_kvs.append(new_key) list_kvs.append(new_val) if isinstance(value, mapping_type) and not changed[0]: return CAny(value) changed[0] = True try: mapping = constructor(*list_kvs) except _ConvertError: raise except Exception as err: raise _ConvertError( f"expected {expected}, got {_tc_describe_value_type(value)}: {err}" ) from None return CAnyChecked(mapping, expected, value) # --------------------------------------------------------------------------- # Converters (3/N): Optional and Union # --------------------------------------------------------------------------- cdef CAny _tc_convert_optional(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for Optional[T]: None passthrough or inner dispatch.""" if value is None: return CAny(None) return _type_convert_dispatch_with_fallback(<_TypeConverter>(conv.subs[0]), value, changed) cdef CAny _tc_convert_union(_TypeConverter conv, object value, bint* changed) except *: """Dispatch for Union[T1, T2, ...].""" cdef _TypeConverter alt cdef bint alt_changed cdef CAny result for alt_obj in conv.subs: alt = <_TypeConverter>alt_obj try: alt_changed = False result = alt.dispatch(alt, value, &alt_changed) changed[0] = alt_changed return result except _ConvertError: pass raise _ConvertError(f"expected {conv.err_hint}, got {_tc_describe_value_type(value)}") # --------------------------------------------------------------------------- # Converters (4/N): Object Types # --------------------------------------------------------------------------- cdef inline object _tc_get_registered_cls(int32_t type_index): if 0 <= type_index < len(TYPE_INDEX_TO_CLS): return TYPE_INDEX_TO_CLS[type_index] return None cdef inline bint _tc_registered_cls_has_base( int32_t type_index, str module_name, str base_name, ) except *: cdef object cls = _tc_get_registered_cls(type_index) cdef object base if cls is None: return False for base in cls.__mro__: if ( getattr(base, "__module__", None) == module_name and getattr(base, "__name__", None) == base_name ): return True return False cdef object _tc_find_payload_enum_variant( int32_t type_index, object enum_cls, object payload ) except *: """Resolve *payload* to its singleton variant (``None`` if no match). Primary path: O(1) lookup in the cross-language value-entries column (``__ffi_enum_value_entries__``). Falls back to an O(n) linear scan over ``enum_cls.all_entries()`` when the column has no entry for *payload* — needed so correctness is preserved for variants whose creators haven't populated the column. """ cdef object value_entries cdef object variant value_entries = _lookup_type_attr(type_index, _TYPE_ATTR_ENUM_VALUE_ENTRIES) if value_entries is not None: variant = value_entries.get(payload) if variant is not None: return variant for variant in enum_cls.all_entries(): if variant.value == payload: return variant return None cdef object _tc_normalize_int_enum_payload(object value, bint* matched) except *: cdef object ivalue matched[0] = False if isinstance(value, bool): matched[0] = True return int(value) if isinstance(value, int): if not (_INT64_MIN <= value <= _INT64_MAX): raise _ConvertError( f"integer {value} out of int64 range [{_INT64_MIN}, {_INT64_MAX}]" ) matched[0] = True return value if isinstance(value, Integral): try: ivalue = int(value) except Exception as err: raise _ConvertError(f"int() failed for {type(value).__qualname__}: {err}") from None if not (_INT64_MIN <= ivalue <= _INT64_MAX): raise _ConvertError( f"integer {ivalue} out of int64 range [{_INT64_MIN}, {_INT64_MAX}]" ) matched[0] = True return ivalue return None cdef CAny _tc_convert_object_marshaled(_TypeConverter conv, object value) except *: cdef int32_t actual_type_index = kTVMFFINone cdef CAny packed cdef CAny converted cdef Function fn_convert cdef object err = None packed = CAnyChecked(value, conv.err_hint, value) fn_convert = conv.fn_convert try: if fn_convert is not None: converted = CAny.__new__(CAny) CHECK_CALL(TVMFFIFunctionCall(fn_convert.chandle, &packed.cdata, 1, &converted.cdata)) else: converted = packed except Exception as err_: err = err_ else: actual_type_index = converted.type_index if actual_type_index >= kTVMFFIStaticObjectBegin: if _tc_type_index_is_instance(actual_type_index, conv.type_index): return converted raise _ConvertError(f"expected {conv.err_hint}, got {_tc_describe_value_type(value)}") from err cdef CAny _tc_convert_object(_TypeConverter conv, object value, bint* changed) except *: """Convert *value* to an object compatible with ``conv.type_index``.""" # TODO: SmallStr and SmallBytes => ObjectRef conversion is not supported yet cdef int32_t actual_type_index = kTVMFFINone # Step 1: existing FFI objects that already satisfy the target schema are passthrough. assert conv.type_index >= kTVMFFIStaticObjectBegin if isinstance(value, CObject): actual_type_index = TVMFFIObjectGetTypeIndex((value).chandle) if _tc_type_index_is_instance(actual_type_index, conv.type_index): return CAny(value) changed[0] = True # Step 2: pack, and convert to the target type. return _tc_convert_object_marshaled(conv, value) cdef CAny _tc_convert_int_enum(_TypeConverter conv, object value, bint* changed) except *: """Convert *value* to an IntEnum-compatible object.""" cdef int32_t actual_type_index = kTVMFFINone cdef object target_cls cdef object ivalue cdef object variant cdef bint is_int_like = False assert conv.type_index >= kTVMFFIStaticObjectBegin if isinstance(value, CObject): actual_type_index = TVMFFIObjectGetTypeIndex((value).chandle) if _tc_type_index_is_instance(actual_type_index, conv.type_index): return CAny(value) target_cls = _tc_get_registered_cls(conv.type_index) if target_cls is not None: ivalue = _tc_normalize_int_enum_payload(value, &is_int_like) if is_int_like: changed[0] = True variant = _tc_find_payload_enum_variant(conv.type_index, target_cls, ivalue) if variant is not None: return CAny(variant) changed[0] = True return _tc_convert_object_marshaled(conv, value) cdef CAny _tc_convert_str_enum(_TypeConverter conv, object value, bint* changed) except *: """Convert *value* to a StrEnum-compatible object.""" cdef int32_t actual_type_index = kTVMFFINone cdef object target_cls cdef object variant assert conv.type_index >= kTVMFFIStaticObjectBegin if isinstance(value, CObject): actual_type_index = TVMFFIObjectGetTypeIndex((value).chandle) if _tc_type_index_is_instance(actual_type_index, conv.type_index): return CAny(value) target_cls = _tc_get_registered_cls(conv.type_index) if target_cls is not None and isinstance(value, str): changed[0] = True variant = _tc_find_payload_enum_variant(conv.type_index, target_cls, value) if variant is not None: return CAny(variant) changed[0] = True return _tc_convert_object_marshaled(conv, value) cdef inline bint _tc_type_index_is_instance(int32_t actual_tindex, int32_t target_tindex) noexcept: """Check if *actual_tindex* is *target_tindex* or a subclass thereof.""" # TODO: this can be optimized by looking up `TYPE_INDEX_TO_INFO` if actual_tindex == target_tindex: return True cdef const TVMFFITypeInfo* actual_info = TVMFFIGetTypeInfo(actual_tindex) if actual_info == NULL: return False cdef const TVMFFITypeInfo* target_info = TVMFFIGetTypeInfo(target_tindex) if target_info == NULL: return False cdef int32_t target_depth = target_info.type_depth if actual_info.type_depth <= target_depth: return False return actual_info.type_ancestors[target_depth].type_index == target_tindex # --------------------------------------------------------------------------- # Helper: describe the Python type of a value for error messages # --------------------------------------------------------------------------- cdef str _tc_describe_value_type(object value): """Return a human-readable type description for *value*.""" cdef object type_info if value is None: return "None" if isinstance(value, bool): return "bool" if isinstance(value, int): return "int" if isinstance(value, float): return "float" if isinstance(value, str): return "str" if isinstance(value, (bytes, bytearray)): return "bytes" if isinstance(value, CObject): type_info = getattr(type(value), "__tvm_ffi_type_info__", None) if type_info is not None: return type_info.type_key return _type_index_to_key(TVMFFIObjectGetTypeIndex((value).chandle)) return type(value).__qualname__ cdef CAny CAnyChecked(object value, str expected, object original_value) except *: """Pack *value* into CAny and normalize packing failures to _ConvertError.""" try: return CAny(value) except _ConvertError: raise except Exception as err: raise _ConvertError( f"expected {expected}, got {_tc_describe_value_type(original_value)}: {err}" ) from None # --------------------------------------------------------------------------- # Builder (runs once per TypeSchema at construction time) # --------------------------------------------------------------------------- def _build_converter(schema): """Build a ``_TypeConverter`` for *schema*.""" cdef _TypeConverter conv = _TypeConverter.__new__(_TypeConverter) cdef _TypeConverter sub_conv cdef _TypeConverter key_conv cdef _TypeConverter val_conv origin = schema.origin args = schema.args origin_tindex = schema.origin_type_index conv.err_hint = origin def _to_type_converter_or_none(object schema_arg): if schema_arg.origin == "Any": return None return <_TypeConverter>(schema_arg._converter) if origin_tindex == kTVMFFIAny or origin == "Any": conv.dispatch = _tc_convert_any conv.err_hint = "Any" return conv if origin == "Optional": conv.dispatch = _tc_convert_optional conv.subs = (<_TypeConverter>(args[0]._converter),) return conv if origin == "Union": conv.dispatch = _tc_convert_union conv.subs = tuple(<_TypeConverter>(a._converter) for a in args) conv.err_hint = " | ".join(repr(a) for a in args) return conv if origin == "int": conv.dispatch = _tc_convert_int return conv if origin == "float": conv.dispatch = _tc_convert_float return conv if origin == "bool": conv.dispatch = _tc_convert_bool return conv if origin == "None": conv.dispatch = _tc_convert_none return conv if origin == "str": conv.dispatch = _tc_convert_str return conv if origin == "bytes": conv.dispatch = _tc_convert_bytes return conv if origin == "Device": conv.dispatch = _tc_convert_device return conv if origin == "dtype": conv.dispatch = _tc_convert_dtype return conv if origin == "ctypes.c_void_p": conv.dispatch = _tc_convert_opaque_ptr return conv if origin == "Tensor": conv.dispatch = _tc_convert_tensor return conv if origin == "Callable": conv.dispatch = _tc_convert_callable return conv if origin == "Array": conv.dispatch = _tc_convert_array if len(args) > 0: sub_conv = _to_type_converter_or_none(args[0]) if sub_conv is not None: conv.subs = (sub_conv,) return conv if origin in ("List", "list"): conv.dispatch = _tc_convert_list conv.err_hint = "List" if len(args) > 0: sub_conv = _to_type_converter_or_none(args[0]) if sub_conv is not None: conv.subs = (sub_conv,) return conv if origin == "Map": conv.dispatch = _tc_convert_map if len(args) == 2: key_conv = _to_type_converter_or_none(args[0]) val_conv = _to_type_converter_or_none(args[1]) if key_conv is not None or val_conv is not None: conv.subs = (key_conv, val_conv) return conv if origin in ("Dict", "dict"): conv.dispatch = _tc_convert_dict conv.err_hint = "Dict" if len(args) == 2: key_conv = _to_type_converter_or_none(args[0]) val_conv = _to_type_converter_or_none(args[1]) if key_conv is not None or val_conv is not None: conv.subs = (key_conv, val_conv) return conv if origin == "tuple": conv.dispatch = _tc_convert_tuple if args is not None: conv.subs = tuple(<_TypeConverter>(a._converter) for a in args) return conv if origin == "Object": conv.dispatch = _tc_convert_object conv.type_index = kTVMFFIObject conv.err_hint = "Object" return conv if origin_tindex >= kTVMFFIStaticObjectBegin: if _tc_registered_cls_has_base(origin_tindex, "tvm_ffi.dataclasses.enum", "IntEnum"): conv.dispatch = _tc_convert_int_enum elif _tc_registered_cls_has_base(origin_tindex, "tvm_ffi.dataclasses.enum", "StrEnum"): conv.dispatch = _tc_convert_str_enum else: conv.dispatch = _tc_convert_object conv.type_index = origin_tindex conv.err_hint = origin return conv tindex = _object_type_key_to_index(origin) if tindex is not None: if _tc_registered_cls_has_base(tindex, "tvm_ffi.dataclasses.enum", "IntEnum"): conv.dispatch = _tc_convert_int_enum elif _tc_registered_cls_has_base(tindex, "tvm_ffi.dataclasses.enum", "StrEnum"): conv.dispatch = _tc_convert_str_enum else: conv.dispatch = _tc_convert_object conv.type_index = tindex conv.err_hint = origin return conv raise TypeError(f"unknown TypeSchema origin: {origin!r}") # --------------------------------------------------------------------------- # Eager protocol normalization and dispatch # --------------------------------------------------------------------------- cdef void _tc_raise_eager_value_protocol_error(_TypeConverter conv, object value) except *: if conv.dispatch == _tc_convert_optional: _tc_raise_eager_value_protocol_error(<_TypeConverter>(conv.subs[0]), value) if conv.err_hint == "Any": raise _ConvertError(f"failed to convert Any from {_tc_describe_value_type(value)}") raise _ConvertError(f"expected {conv.err_hint}, got {_tc_describe_value_type(value)}") cdef object _tc_eager_protocol_step(object value, bint* stalled_value_protocol) except *: cdef object vtype cdef object inner if isinstance(value, (Tensor, CObject, ObjectRValueRef, PyNativeObject)): return value vtype = type(value) if hasattr(vtype, "__tvm_ffi_object__"): try: return value.__tvm_ffi_object__() except Exception: raise _ConvertError( f"__tvm_ffi_object__() failed for {_tc_describe_value_type(value)}" ) from None if hasattr(vtype, "__tvm_ffi_value__"): try: inner = value.__tvm_ffi_value__() except Exception: # Report the schema mismatch instead of leaking the raw # __tvm_ffi_value__ implementation error. stalled_value_protocol[0] = True return value if inner is value: stalled_value_protocol[0] = True return inner if isinstance(value, ObjectConvertible): # Normalize ObjectConvertible eagerly so nested Union/container dispatch # sees the inner FFI object instead of the Python wrapper. try: inner = value.asobject() except Exception: raise _ConvertError(f"asobject() failed for {_tc_describe_value_type(value)}") from None if not isinstance(inner, CObject): raise _ConvertError( f"asobject() returned {_tc_describe_value_type(inner)} " f"for {_tc_describe_value_type(value)}" ) return inner return value cdef CAny _type_convert_dispatch_with_fallback(_TypeConverter conv, object value, bint* changed) except *: """Dispatch after eager protocol normalization with cycle protection.""" cdef int depth = 0 cdef object inner cdef bint stalled_value_protocol cdef bint used_value_protocol = False cdef CAny result while True: stalled_value_protocol = False inner = _tc_eager_protocol_step(value, &stalled_value_protocol) if stalled_value_protocol: _tc_raise_eager_value_protocol_error(conv, value) if inner is value: break depth += 1 if depth > _VALUE_PROTOCOL_MAX_DEPTH: raise _ConvertError("infinite __tvm_ffi_value__ cycle detected") from None used_value_protocol = True value = inner changed[0] = False result = conv.dispatch(conv, value, changed) if used_value_protocol: changed[0] = True return result # --------------------------------------------------------------------------- # Main dispatcher (thin entry point from Python-level TypeSchema methods) # --------------------------------------------------------------------------- cdef CAny _type_convert_impl(_TypeConverter converter, object value) except *: """Dispatch to the C-level converter.""" cdef bint changed return _type_convert_dispatch_with_fallback(converter, value, &changed) tvm-ffi-0.1.12/python/tvm_ffi/cython/string.pxi000066400000000000000000000061141521067262500214620ustar00rootroot00000000000000 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # helper class for string/bytes handling cdef inline str _string_obj_get_py_str(obj): cdef TVMFFIByteArray* bytes = TVMFFIBytesGetByteArrayPtr((obj).chandle) return bytearray_to_str(bytes) cdef inline bytes _bytes_obj_get_py_bytes(obj): cdef TVMFFIByteArray* bytes = TVMFFIBytesGetByteArrayPtr((obj).chandle) return bytearray_to_bytes(bytes) from typing import Any class String(str, PyNativeObject): __slots__ = ["_tvm_ffi_cached_object"] """UTF-8 string that interoperates with FFI while behaving like ``str``. ``String`` is a :class:`str` subclass that can travel across the FFI boundary without copying for large payloads. For most Python APIs, using a plain ``str`` works seamlessly; the runtime converts to and from ``String`` as needed. Examples -------- .. code-block:: python fecho = tvm_ffi.get_global_func("testing.echo") s = tvm_ffi.core.String("hello") assert fecho(s) == "hello" assert fecho("world") == "world" """ _tvm_ffi_cached_object: Object | None def __new__(cls, value: str) -> "String": val = str.__new__(cls, value) val._tvm_ffi_cached_object = None return val # pylint: disable=no-self-argument def __from_tvm_ffi_object__(cls, obj: Any) -> "String": """Construct a ``String`` from an FFI object (internal).""" content = _string_obj_get_py_str(obj) val = str.__new__(cls, content) val._tvm_ffi_cached_object = obj return val class Bytes(bytes, PyNativeObject): """Byte buffer that interoperates with FFI while behaving like ``bytes``. Like :class:`String`, this class enables zero-copy exchange for large data. Most Python code can use ``bytes`` directly; the FFI layer constructs :class:`Bytes` as needed. """ _tvm_ffi_cached_object: Object | None def __new__(cls, value: bytes) -> "Bytes": val = bytes.__new__(cls, value) val._tvm_ffi_cached_object = None return val # pylint: disable=no-self-argument def __from_tvm_ffi_object__(cls, obj: Any) -> "Bytes": """Construct ``Bytes`` from an FFI object (internal).""" content = _bytes_obj_get_py_bytes(obj) val = bytes.__new__(cls, content) val._tvm_ffi_cached_object = obj return val tvm-ffi-0.1.12/python/tvm_ffi/cython/tensor.pxi000066400000000000000000000510771521067262500214760ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from typing import Any __dlpack_version__: tuple[int, int] = (DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION) _CLASS_TENSOR = None def _set_class_tensor(cls): global _CLASS_TENSOR _CLASS_TENSOR = cls cdef const char* _c_str_dltensor = "dltensor" cdef const char* _c_str_used_dltensor = "used_dltensor" cdef const char* _c_str_dltensor_versioned = "dltensor_versioned" cdef const char* _c_str_used_dltensor_versioned = "used_dltensor_versioned" cdef const char* _c_str_dlpack_exchange_api = "dlpack_exchange_api" cdef int _get_dlpack_exchange_api( object dlpack_exchange_api_obj, const DLPackExchangeAPI** out_ptr ) except -1: if isinstance(dlpack_exchange_api_obj, int): out_ptr[0] = (dlpack_exchange_api_obj) return 0 if pycapsule.PyCapsule_IsValid(dlpack_exchange_api_obj, _c_str_dlpack_exchange_api): out_ptr[0] = pycapsule.PyCapsule_GetPointer( dlpack_exchange_api_obj, _c_str_dlpack_exchange_api ) return 0 raise ValueError("Expect a dlpack_exchange_api field") cdef void _c_dlpack_deleter(object pycaps): cdef DLManagedTensor* dltensor if pycapsule.PyCapsule_IsValid(pycaps, _c_str_dltensor): dltensor = pycapsule.PyCapsule_GetPointer(pycaps, _c_str_dltensor) dltensor.deleter(dltensor) cdef void _c_dlpack_versioned_deleter(object pycaps): cdef DLManagedTensorVersioned* dltensor if pycapsule.PyCapsule_IsValid(pycaps, _c_str_dltensor_versioned): dltensor = pycapsule.PyCapsule_GetPointer( pycaps, _c_str_dltensor_versioned) dltensor.deleter(dltensor) cdef inline int _from_dlpack( object dltensor, int require_alignment, int require_contiguous, TVMFFIObjectHandle* out ) except -1: cdef DLManagedTensor* ptr cdef int c_api_ret_code cdef int c_req_alignment = require_alignment cdef int c_req_contiguous = require_contiguous if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor): ptr = pycapsule.PyCapsule_GetPointer(dltensor, _c_str_dltensor) c_api_ret_code = TVMFFITensorFromDLPack( ptr, c_req_alignment, c_req_contiguous, out) CHECK_CALL(c_api_ret_code) # set name and destructor to be empty pycapsule.PyCapsule_SetDestructor(dltensor, NULL) pycapsule.PyCapsule_SetName(dltensor, _c_str_used_dltensor) return 0 raise ValueError("Expect a dltensor field, PyCapsule can only be consumed once") cdef inline int _from_dlpack_versioned( object dltensor, int require_alignment, int require_contiguous, TVMFFIObjectHandle* out ) except -1: cdef DLManagedTensorVersioned* ptr cdef int c_api_ret_code cdef int c_req_alignment = require_alignment cdef int c_req_contiguous = require_contiguous if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor_versioned): ptr = pycapsule.PyCapsule_GetPointer( dltensor, _c_str_dltensor_versioned) c_api_ret_code = TVMFFITensorFromDLPackVersioned( ptr, c_req_alignment, c_req_contiguous, out) CHECK_CALL(c_api_ret_code) # set name and destructor to be empty pycapsule.PyCapsule_SetDestructor(dltensor, NULL) pycapsule.PyCapsule_SetName(dltensor, _c_str_used_dltensor_versioned) return 0 raise ValueError("Expect a dltensor_versioned field, PyCapsule can only be consumed once") cdef inline int _from_dlpack_exchange_api( object ext_tensor, const DLPackExchangeAPI* exchange_api, int require_alignment, int require_contiguous, TVMFFIObjectHandle* out ) except -1: cdef DLManagedTensorVersioned* temp_managed_tensor cdef PyObject* ext_tensor_pyobj = ext_tensor if exchange_api.managed_tensor_from_py_object_no_sync(ext_tensor_pyobj, &temp_managed_tensor) != 0: return -1 # Convert to TVM Tensor if TVMFFITensorFromDLPackVersioned( temp_managed_tensor, require_alignment, require_contiguous, out ) != 0: # recycle the managed tensor to avoid leak if temp_managed_tensor.deleter != NULL: temp_managed_tensor.deleter(temp_managed_tensor) raise BufferError("Failed to convert DLManagedTensorVersioned to ffi.Tensor") return 0 cdef inline int _from_dlpack_universal( object ext_tensor, int require_alignment, int require_contiguous, TVMFFIObjectHandle* out ) except -1: # as of most frameworks do not yet support v1.1 # move to false as most frameworks get upgraded. cdef int favor_legacy_dlpack = True cdef const DLPackExchangeAPI* exchange_api = NULL if hasattr(ext_tensor, "__dlpack_c_exchange_api__"): try: _get_dlpack_exchange_api(ext_tensor.__dlpack_c_exchange_api__, &exchange_api) return _from_dlpack_exchange_api( ext_tensor, exchange_api, require_alignment, require_contiguous, out ) except BufferError: pass if hasattr(ext_tensor, "__dlpack__"): if favor_legacy_dlpack: return _from_dlpack( ext_tensor.__dlpack__(), require_alignment, require_contiguous, out ) else: try: return _from_dlpack_versioned( ext_tensor.__dlpack__(max_version=__dlpack_version__), require_alignment, require_contiguous, out ) except TypeError: return _from_dlpack( ext_tensor.__dlpack__(), require_alignment, require_contiguous, out ) else: if pycapsule.PyCapsule_IsValid(ext_tensor, _c_str_dltensor_versioned): return _from_dlpack_versioned( ext_tensor, require_alignment, require_contiguous, out ) elif pycapsule.PyCapsule_IsValid(ext_tensor, _c_str_dltensor): return _from_dlpack( ext_tensor, require_alignment, require_contiguous, out ) else: raise TypeError("Expect from_dlpack to take either a compatible tensor or PyCapsule") def from_dlpack( ext_tensor: Any, *, require_alignment: int = 0, require_contiguous: bool = False ) -> Tensor: """Import a foreign array that implements the DLPack producer protocol. Parameters ---------- ext_tensor An object supporting :py:meth:`__dlpack__ ` and :py:meth:`__dlpack_device__ `. require_alignment If greater than zero, require the underlying data pointer to be aligned to this many bytes. Misaligned inputs raise :class:`ValueError`. require_contiguous : bool, optional When True, require the layout to be contiguous. Non-contiguous inputs raise :class:`ValueError`. Returns ------- Tensor A TVM FFI :class:`Tensor` that references the same memory. Examples -------- .. code-block:: python import numpy as np import tvm_ffi x_np = np.arange(8, dtype="int32") x = tvm_ffi.from_dlpack(x_np) y_np = np.from_dlpack(x) assert np.shares_memory(x_np, y_np) """ # noqa: E501 cdef TVMFFIObjectHandle chandle _from_dlpack_universal(ext_tensor, require_alignment, require_contiguous, &chandle) return make_tensor_from_chandle(chandle) # helper class for shape handling def _shape_obj_get_py_tuple(obj: "CObject") -> tuple[int, ...]: cdef TVMFFIShapeCell* shape = TVMFFIShapeGetCellPtr((obj).chandle) return tuple(shape.data[i] for i in range(shape.size)) def _make_strides_from_shape(tuple shape: tuple[int, ...]) -> tuple[int, ...]: cdef int64_t expected_stride = 1 cdef list strides = [] cdef int64_t ndim = len(shape) cdef int64_t reverse_index for i in range(ndim): reverse_index = ndim - i - 1 strides.append(expected_stride) expected_stride *= shape[reverse_index] return tuple(reversed(strides)) cdef class Tensor(CObject): """Managed n-dimensional array compatible with DLPack. It provides zero-copy interoperability with array libraries through the DLPack protocol. Instances are typically created with :func:`from_dlpack` or returned from FFI functions. Examples -------- .. code-block:: python import numpy as np import tvm_ffi x = tvm_ffi.from_dlpack(np.arange(6, dtype="int32")) assert x.shape == (6,) assert x.dtype == tvm_ffi.dtype("int32") # Round-trip through NumPy using DLPack np.testing.assert_equal(np.from_dlpack(x), np.arange(6, dtype="int32")) """ __slots__ = () cdef DLTensor* cdltensor @property def shape(self) -> tuple[int, ...]: """Tensor shape as a tuple of integers.""" return tuple(self.cdltensor.shape[i] for i in range(self.cdltensor.ndim)) @property def ndim(self) -> int: """Number of dimensions of the tensor.""" return self.cdltensor.ndim def numel(self) -> int: """Total number of elements in the tensor.""" cdef int64_t count = 1 cdef int i for i in range(self.cdltensor.ndim): count *= self.cdltensor.shape[i] return count def size(self, idx: int) -> int: """Get the size of the ``idx``-th dimension. Negative ``idx`` counts from the last dimension.""" cdef int ndim = self.cdltensor.ndim if idx < -ndim or idx >= ndim: raise IndexError( f"Dimension {idx} out of range for tensor with {ndim} dimensions" ) if idx < 0: idx += ndim return self.cdltensor.shape[idx] def is_contiguous(self) -> bool: """True if the Tensor is C-contiguous (row-major), False otherwise.""" if self.cdltensor.strides == NULL: return True # An empty tensor (numel == 0) is trivially contiguous regardless of strides, # matching NumPy/PyTorch semantics. cdef int i cdef int k for i in range(self.cdltensor.ndim): if self.cdltensor.shape[i] == 0: return True cdef int64_t expected_stride = 1 for i in range(self.cdltensor.ndim, 0, -1): k = i - 1 if self.cdltensor.shape[k] == 1: continue if self.cdltensor.strides[k] != expected_stride: return False expected_stride *= self.cdltensor.shape[k] return True @property def strides(self) -> tuple[int, ...]: """Tensor strides as a tuple of integers.""" if self.cdltensor.strides == NULL: return _make_strides_from_shape(self.shape) return tuple(self.cdltensor.strides[i] for i in range(self.cdltensor.ndim)) @property def dtype(self) -> Any: """Data type as :class:`tvm_ffi.dtype` (``str`` subclass).""" cdef TVMFFIAny dtype_any dtype_any.v_dtype = self.cdltensor.dtype return make_ret_dtype(dtype_any) @property def device(self) -> Device: """The :class:`Device` on which the tensor is placed.""" cdef TVMFFIAny device_any device_any.v_device = self.cdltensor.device return make_ret_device(device_any) def _to_dlpack(self) -> object: """Return a DLPack capsule representing this tensor (internal).""" cdef DLManagedTensor* dltensor cdef int c_api_ret_code c_api_ret_code = TVMFFITensorToDLPack(self.chandle, &dltensor) CHECK_CALL(c_api_ret_code) return pycapsule.PyCapsule_New(dltensor, _c_str_dltensor, _c_dlpack_deleter) def _to_dlpack_versioned(self) -> object: """Return a versioned DLPack capsule (internal).""" cdef DLManagedTensorVersioned* dltensor cdef int c_api_ret_code c_api_ret_code = TVMFFITensorToDLPackVersioned(self.chandle, &dltensor) CHECK_CALL(c_api_ret_code) return pycapsule.PyCapsule_New( dltensor, _c_str_dltensor_versioned, _c_dlpack_versioned_deleter) def __dlpack_device__(self) -> tuple[int, int]: """Implement the standard :py:meth:`__dlpack_device__ ` protocol.""" cdef int device_type = self.cdltensor.device.device_type cdef int device_id = self.cdltensor.device.device_id return (device_type, device_id) def __dlpack__( self, *, stream: Any | None = None, max_version: tuple[int, int] | None = None, dl_device: tuple[int, int] | None = None, copy: bool | None = None, ) -> object: """Implement the standard :py:meth:`__dlpack__ ` protocol. Parameters ---------- stream Framework-specific stream/context object. max_version Upper bound on the supported DLPack version of the consumer. When ``None``, use the built-in protocol version. dl_device Override the device reported by :py:meth:`__dlpack_device__`. copy If ``True``, produce a copy rather than exporting in-place. Raises ------ BufferError If the requested behavior cannot be satisfied. """ # noqa: E501 if max_version is None: # Keep and use the DLPack 0.X implementation # Note: from March 2025 onwards (but ideally as late as # possible), it's okay to raise BufferError here return self._to_dlpack() else: # We get to produce `DLManagedTensorVersioned` now. Note that # our_own_dlpack_version is the max version that the *producer* # supports and fills in the `DLManagedTensorVersioned::version` # field if max_version[0] >= __dlpack_version__[0]: if dl_device is not None and dl_device != self.__dlpack_device__(): raise BufferError("dl_device of different type not supported") if copy is not None and copy: raise BufferError("copy not yet supported") return self._to_dlpack_versioned() elif max_version[0] < 1: return self.__ctypes_handle__to_dlpack() else: raise BufferError(f"Unsupported max_version {max_version}") _set_class_tensor(Tensor) cdef int _dltensor_test_wrapper_from_pyobject( void* obj, DLManagedTensorVersioned** out ) except -1: """DLPackExchangeAPI: managed_tensor_from_py_object_no_sync""" cdef PyObject* py_obj = obj cdef DLTensorTestWrapper wrapper = py_obj return TVMFFITensorToDLPackVersioned(wrapper.tensor.chandle, out) cdef int _dltensor_test_wrapper_to_pyobject( DLManagedTensorVersioned* tensor, void** out_py_object ) except -1: """DLPackExchangeAPI: managed_tensor_to_py_object_no_sync""" cdef TVMFFIObjectHandle temp_chandle if TVMFFITensorFromDLPackVersioned(tensor, 0, 0, &temp_chandle) != 0: return -1 py_tensor = make_tensor_from_chandle(temp_chandle) Py_INCREF(py_tensor) out_py_object[0] = (py_tensor) return 0 cdef int _dltensor_test_wrapper_current_work_stream( int device_type, int32_t device_id, void** out_stream ) except -1: """DLPackExchangeAPI: current_work_stream""" if device_type != kDLCPU: out_stream[0] = TVMFFIEnvGetStream(device_type, device_id) return 0 # Module-level static DLPackExchangeAPI for DLTensorTestWrapper cdef DLPackExchangeAPI _dltensor_test_wrapper_static_api cdef DLPackExchangeAPI* _dltensor_test_wrapper_get_exchange_api() noexcept: """Get the static DLPackExchangeAPI instance for DLTensorTestWrapper.""" global _dltensor_test_wrapper_static_api # Initialize header using macros from dlpack.h _dltensor_test_wrapper_static_api.header.version.major = DLPACK_MAJOR_VERSION _dltensor_test_wrapper_static_api.header.version.minor = DLPACK_MINOR_VERSION _dltensor_test_wrapper_static_api.header.prev_api = NULL # Initialize function pointers _dltensor_test_wrapper_static_api.managed_tensor_allocator = NULL _dltensor_test_wrapper_static_api.managed_tensor_from_py_object_no_sync = ( _dltensor_test_wrapper_from_pyobject ) _dltensor_test_wrapper_static_api.managed_tensor_to_py_object_no_sync = ( _dltensor_test_wrapper_to_pyobject ) _dltensor_test_wrapper_static_api.dltensor_from_py_object_no_sync = NULL _dltensor_test_wrapper_static_api.current_work_stream = ( _dltensor_test_wrapper_current_work_stream ) return &_dltensor_test_wrapper_static_api cdef class DLTensorTestWrapper: """Wrapper of a Tensor that exposes DLPack protocol, only for testing purpose. """ __slots__ = () __dlpack_c_exchange_api__ = pycapsule.PyCapsule_New( _dltensor_test_wrapper_get_exchange_api(), b"dlpack_exchange_api", NULL ) cdef Tensor tensor cdef dict __dict__ def __init__(self, tensor: Tensor) -> None: self.tensor = tensor def __tvm_ffi_env_stream__(self) -> int: cdef TVMFFIStreamHandle stream cdef long long stream_as_int cdef int c_api_ret_code stream = TVMFFIEnvGetStream( self.tensor.cdltensor.device.device_type, self.tensor.cdltensor.device.device_id) stream_as_int = stream return stream_as_int def __dlpack_device__(self) -> tuple[int, int]: return self.tensor.__dlpack_device__() def __dlpack__(self, *, **kwargs: Any) -> object: return self.tensor.__dlpack__(**kwargs) cdef inline object make_ret_dltensor(TVMFFIAny result): cdef DLTensor* dltensor dltensor = result.v_ptr tensor = _CLASS_TENSOR.__new__(_CLASS_TENSOR) (tensor).chandle = NULL (tensor).cdltensor = dltensor return tensor cdef inline object make_tensor_from_chandle( TVMFFIObjectHandle chandle, const DLPackExchangeAPI* c_ctx_dlpack_api = NULL ): cdef object tensor cdef void* py_obj cdef DLManagedTensorVersioned* dlpack if c_ctx_dlpack_api != NULL and c_ctx_dlpack_api.managed_tensor_to_py_object_no_sync != NULL: # try convert and import into the environment array if possible if TVMFFITensorToDLPackVersioned(chandle, &dlpack) == 0: try: # note that py_obj already holds an extra reference to the tensor # so we need to decref it after the conversion c_ctx_dlpack_api.managed_tensor_to_py_object_no_sync(dlpack, &py_obj) tensor = (py_obj) Py_DECREF(tensor) # decref original handle to prevent leak. # note that DLManagedTensor also hold a reference to the tensor # so we need to decref the original handle if the conversion is successful TVMFFIObjectDecRef(chandle) return tensor except Exception: # call the deleter to free the memory since we will continue to use the chandle dlpack.deleter(dlpack) pass # default return the tensor tensor = _CLASS_TENSOR.__new__(_CLASS_TENSOR) (tensor).chandle = chandle (tensor).cdltensor = TVMFFITensorGetDLTensorPtr(chandle) return tensor cdef inline object make_tensor_from_any(TVMFFIAny any, const DLPackExchangeAPI* c_ctx_dlpack_api): return make_tensor_from_chandle(any.v_ptr, c_ctx_dlpack_api) tvm-ffi-0.1.12/python/tvm_ffi/cython/tvm_ffi_python_helpers.h000066400000000000000000001365351521067262500243730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file tvm_ffi_python_helpers.h * \brief C++ based helpers for the Python FFI call to optimize performance. */ #ifndef TVM_FFI_PYTHON_HELPERS_H_ #define TVM_FFI_PYTHON_HELPERS_H_ #include #include #include // Define here to avoid dependencies on non-c headers for now #ifndef TVM_FFI_INLINE #if defined(_MSC_VER) #define TVM_FFI_INLINE [[msvc::forceinline]] inline #else #define TVM_FFI_INLINE [[gnu::always_inline]] inline #endif #endif // Local mirror of TVM_FFI_COLD_CODE / TVM_FFI_PREDICT_* from // . The Cython helper deliberately avoids that header // (keeps the include surface c-headers-only), so we duplicate the macro // definitions here. Keep these in sync with base_details.h: same expansion on // GCC/Clang, no-op on MSVC. #ifndef TVM_FFI_COLD_CODE #if defined(__GNUC__) || defined(__clang__) #define TVM_FFI_COLD_CODE [[gnu::cold]] #else #define TVM_FFI_COLD_CODE #endif #endif #ifndef TVM_FFI_PREDICT_FALSE #if defined(__GNUC__) || defined(__clang__) #define TVM_FFI_PREDICT_FALSE(cond) (__builtin_expect(static_cast(cond), 0)) #define TVM_FFI_PREDICT_TRUE(cond) (__builtin_expect(static_cast(cond), 1)) #else #define TVM_FFI_PREDICT_FALSE(cond) (cond) #define TVM_FFI_PREDICT_TRUE(cond) (cond) #endif #endif #include #include #include #include #include #include ///-------------------------------------------------------------------------------- /// We deliberately designed the data structure and function to be C-style // prefixed with TVMFFIPy so they can be easily invoked through Cython. ///-------------------------------------------------------------------------------- //------------------------------------------------------------------------------------ // Helpers for Python thread-state attachment //------------------------------------------------------------------------------------ // // On classic builds, PyGILState_Ensure attaches the current thread and acquires the GIL. // On free-threaded builds, there is no process-wide GIL to acquire, but CPython still // requires an attached thread state before manipulating Python refcounts. class TVMFFIPyWithAttachedThreadState { public: TVMFFIPyWithAttachedThreadState() noexcept { gstate_ = PyGILState_Ensure(); } ~TVMFFIPyWithAttachedThreadState() { PyGILState_Release(gstate_); } private: PyGILState_STATE gstate_; }; /*! * \brief Closure state carried as the resource handle for an FFI function that * wraps a Python callable and optional exchange api for tensor handling. * * Created by TVMFFIPyConvertPyCallback and released by * TVMFFIPyCallbackClosure::Deleter when the FFI function is destroyed. */ struct TVMFFIPyCallbackClosure { /*! \brief Strong reference to the Python callable. */ PyObject* callable; /*! \brief Optional DLPack exchange API used when constructing Tensor returns. */ const DLPackExchangeAPI* dlpack_exchange_api; /*! * \brief Deleter registered with TVMFFIFunctionCreate. Runs on FFI function destroy. * * Releases the closure's strong Python reference and frees the closure. */ static void Deleter(void* context) noexcept { TVMFFIPyWithAttachedThreadState thread_state; auto* closure = static_cast(context); Py_DecRef(closure->callable); delete closure; } }; /*! * \brief Thread-local call stack used by TVMFFIPyCallContext. */ class TVMFFIPyCallStack { public: /*! \brief The stack of arguments */ std::vector args_stack; /*! \brief The top of the argument call stack currently */ int64_t args_stack_top = 0; /*! * \brief The stack of extra temporary Python objects that may not fit into * one temp per argument budget, mainly used by value protocol. */ std::vector extra_temp_py_objects_stack; /*! \brief Constructor to initialize the call stack */ TVMFFIPyCallStack() { // keep it 4K as default stack size so it is page aligned constexpr size_t kDefaultStackSize = 4096; // fit everything roughly 4K stack args_stack.resize(kDefaultStackSize / sizeof(TVMFFIAny)); extra_temp_py_objects_stack.reserve(16); } }; //--------------------------------------------------------------------------------------------- // Support for Python -> FFI function calls. //--------------------------------------------------------------------------------------------- /*! * \brief Context for each ffi call to track the stream, device and temporary arguments. */ class TVMFFIPyCallContext { public: /*! \brief The workspace for the packed arguments */ TVMFFIAny* packed_args = nullptr; /*! \brief Detected device type, if any */ int device_type = -1; /*! \brief Detected device id, if any */ int device_id = 0; /*! \brief Detected stream, if any */ void* stream = nullptr; /*! \brief the DLPack exchange API, if any */ const DLPackExchangeAPI* dlpack_c_exchange_api{nullptr}; /*! \brief pointer to the call stack space */ TVMFFIPyCallStack* call_stack = nullptr; /*! \brief the temporary arguments to be recycled */ void** temp_ffi_objects = nullptr; /*! \brief the temporary arguments to be recycled */ void** temp_py_objects = nullptr; /*! \brief the number of temporary arguments */ int num_temp_ffi_objects = 0; /*! \brief the number of temporary arguments */ int num_temp_py_objects = 0; /*! \brief RAII guard constructor to create a TVMFFIPyCallContext */ TVMFFIPyCallContext(TVMFFIPyCallStack* call_stack, int64_t num_args) : call_stack(call_stack) { // In most cases, it will try to allocate from temp_stack, // then allocate from heap if the request goes beyond the stack size. static_assert(sizeof(TVMFFIAny) >= (sizeof(void*) * 2)); static_assert(alignof(TVMFFIAny) % alignof(void*) == 0); old_args_stack_top_ = call_stack->args_stack_top; int64_t requested_count = num_args * 2; TVMFFIAny* stack_head = call_stack->args_stack.data() + call_stack->args_stack_top; if (call_stack->args_stack_top + requested_count > static_cast(call_stack->args_stack.size())) { // allocate from heap heap_ptr_ = new TVMFFIAny[requested_count]; stack_head = heap_ptr_; } else { call_stack->args_stack_top += requested_count; } this->packed_args = stack_head; // by default we co-locate the temporary arguments with packed arguments // for better cache locality with one temp per argument budget. this->temp_ffi_objects = reinterpret_cast(stack_head + num_args); this->temp_py_objects = this->temp_ffi_objects + num_args; this->old_extra_temp_py_objects_stack_top_ = call_stack->extra_temp_py_objects_stack.size(); } ~TVMFFIPyCallContext() { TVMFFIPyWithAttachedThreadState thread_state; try { // recycle the temporary arguments if any for (int i = 0; i < this->num_temp_ffi_objects; ++i) { TVMFFIObjectDecRef(this->temp_ffi_objects[i]); } for (int i = 0; i < this->num_temp_py_objects; ++i) { Py_DecRef(static_cast(this->temp_py_objects[i])); } for (size_t i = old_extra_temp_py_objects_stack_top_; i < call_stack->extra_temp_py_objects_stack.size(); ++i) { Py_DecRef(static_cast(call_stack->extra_temp_py_objects_stack[i])); } call_stack->extra_temp_py_objects_stack.resize(old_extra_temp_py_objects_stack_top_); } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); } // now recycle the memory of the call stack if (heap_ptr_ == nullptr) { call_stack->args_stack_top = old_args_stack_top_; } else { delete[] heap_ptr_; } } private: /*! \brief the heap pointer */ TVMFFIAny* heap_ptr_ = nullptr; /*! \brief the old stack top */ size_t old_args_stack_top_; /*! \brief the begin index of the temporary Python objects stack */ size_t old_extra_temp_py_objects_stack_top_; }; /*! \brief Argument setter for a given python argument. */ struct TVMFFIPyArgSetter { /*! * \brief Function pointer to invoke the setter. * \param self Pointer to this, this should be TVMFFIPyArgSetter* * \param call_ctx The call context. * \param arg The python argument to be set * \param out The output argument. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ int (*func)(TVMFFIPyArgSetter* self, TVMFFIPyCallContext* call_ctx, PyObject* arg, TVMFFIAny* out); /*! * \brief Optional DLPackExchangeAPI struct pointer. * This is the new struct-based approach that bundles all DLPack exchange functions. */ const DLPackExchangeAPI* dlpack_c_exchange_api{nullptr}; /*! * \brief Invoke the setter. * \param call_ctx The call context. * \param arg The python argument to be set * \param out The output argument. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ TVM_FFI_INLINE int operator()(TVMFFIPyCallContext* call_ctx, PyObject* arg, TVMFFIAny* out) const { return (*func)(const_cast(this), call_ctx, arg, out); } }; //--------------------------------------------------------------------------------------------- // The following section contains predefined setters for common POD types // They ar not meant to be used directly, but instead being registered to TVMFFIPyCallManager //--------------------------------------------------------------------------------------------- int TVMFFIPyArgSetterFloat_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) noexcept { out->type_index = kTVMFFIFloat; // this function getsdispatched when type is already float, so no need to worry about error out->v_float64 = PyFloat_AsDouble(arg); return 0; } int TVMFFIPyArgSetterInt_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) noexcept { int overflow = 0; out->type_index = kTVMFFIInt; out->v_int64 = PyLong_AsLongLongAndOverflow(arg, &overflow); if (TVM_FFI_PREDICT_FALSE(overflow != 0)) { PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to int64_t"); return -1; } return 0; } int TVMFFIPyArgSetterBool_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) noexcept { out->type_index = kTVMFFIBool; // this function getsdispatched when type is already bool, so no need to worry about error out->v_int64 = PyLong_AsLong(arg); return 0; } int TVMFFIPyArgSetterNone_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) noexcept { out->type_index = kTVMFFINone; out->v_int64 = 0; return 0; } //--------------------------------------------------------------------------------------------- // Support for PyCallback function calls. //--------------------------------------------------------------------------------------------- /*! * \brief Context for a C -> Python callback call. * * Owns a temporary PyObject* array that holds arguments converted from the * packed FFI call. Space is first taken from the thread-local args_stack on * TVMFFIPyCallStack; if insufficient, we fall back to the heap. * * Unlike TVMFFIPyCallContext::~TVMFFIPyCallContext, this destructor does NOT * attach a thread state — callers are expected to already hold one * (e.g. via TVMFFIPyWithAttachedThreadState at the top of the callback). * * The destructor also decrefs every PyObject* pushed into py_args[0 .. * num_active_py_args-1], tracking the pushed count via `num_active_py_args`. */ class TVMFFIPyCallbackContext { public: /*! \brief The temporary PyObject* slots for Python call arguments. */ PyObject** py_args = nullptr; /*! \brief How many slots have a live reference and need decref on cleanup. */ int32_t num_active_py_args = 0; /*! \brief Number of total argument slots allocated. */ int32_t num_args = 0; TVMFFIPyCallbackContext(TVMFFIPyCallStack* call_stack, int32_t num_args) : num_args(num_args), call_stack_(call_stack) { static_assert(sizeof(TVMFFIAny) % sizeof(PyObject*) == 0); // slots needed in the unit of TVMFFIAny int64_t slots_needed = (static_cast(num_args) * sizeof(PyObject*) + sizeof(TVMFFIAny) - 1) / sizeof(TVMFFIAny); old_args_stack_top_ = call_stack->args_stack_top; if (call_stack->args_stack_top + slots_needed <= static_cast(call_stack->args_stack.size())) { py_args = reinterpret_cast(call_stack->args_stack.data() + call_stack->args_stack_top); call_stack->args_stack_top += slots_needed; } else { heap_ptr_ = new PyObject*[num_args]; py_args = heap_ptr_; } } ~TVMFFIPyCallbackContext() { // caller must already hold an attached thread state; do NOT re-attach. // we ensure that all the pyargs are not null for (int32_t i = 0; i < num_active_py_args; ++i) { Py_DecRef(py_args[i]); } if (heap_ptr_ == nullptr) { call_stack_->args_stack_top = old_args_stack_top_; } else { delete[] heap_ptr_; } } private: TVMFFIPyCallStack* call_stack_ = nullptr; int64_t old_args_stack_top_ = 0; PyObject** heap_ptr_ = nullptr; }; /*! * \brief A callback arg setter entry registered to handle efficient callback argument conversion. */ struct TVMFFIPyCallbackArgSetter { /*! * \brief Callback type that converts a borrowed TVMFFIAny (AnyView) to a new-reference PyObject*. * \param handle Pointer to the TVMFFIPyCallbackArgSetter (for per-type state). * \param dlpack_exchange_api The DLPack exchange API (may be NULL). * \param arg The TVMFFIAny value to convert (borrowed; setter must inc if it transfers * ownership). * \param out Output: a new-reference PyObject*. * \return 0 on success, -1 on failure (PyErr set). */ int (*func)(TVMFFIPyCallbackArgSetter* handle, const DLPackExchangeAPI* dlpack_exchange_api, const TVMFFIAny* arg, PyObject** out); }; // common callback arg setters that can be quikcly implemented in C++ and used by cython factory // note that PyErr is propagated back so we just need to return -1 on failure. int TVMFFIPyCallbackArgSetterNone_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny*, PyObject** out) noexcept { Py_IncRef(Py_None); *out = Py_None; return 0; } int TVMFFIPyCallbackArgSetterBool_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) noexcept { *out = PyBool_FromLong(static_cast(arg->v_int64)); return (*out != nullptr) ? 0 : -1; } int TVMFFIPyCallbackArgSetterInt_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) noexcept { *out = PyLong_FromLongLong(arg->v_int64); return (*out != nullptr) ? 0 : -1; } int TVMFFIPyCallbackArgSetterFloat_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) noexcept { *out = PyFloat_FromDouble(arg->v_float64); return (*out != nullptr) ? 0 : -1; } int TVMFFIPyCallbackArgSetterSmallStr_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) noexcept { TVMFFIByteArray ba = TVMFFISmallBytesGetContentByteArray(arg); *out = PyUnicode_DecodeUTF8(ba.data, static_cast(ba.size), nullptr); return (*out != nullptr) ? 0 : -1; } int TVMFFIPyCallbackArgSetterSmallBytes_(TVMFFIPyCallbackArgSetter*, const DLPackExchangeAPI*, const TVMFFIAny* arg, PyObject** out) noexcept { TVMFFIByteArray ba = TVMFFISmallBytesGetContentByteArray(arg); *out = PyBytes_FromStringAndSize(ba.data, static_cast(ba.size)); return (*out != nullptr) ? 0 : -1; } ///-------------------------------------------------------------------------------- /// Declaring functions defined in Cython to be invoked by the C++ implementation. /// in all cases PyErr is propagated back so we just need to return -1 on failure. ///-------------------------------------------------------------------------------- /* * \brief Set the error raised from Python to the FFI side. * \param py_err The Python error to be set. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ __PYX_EXTERN_C int TVMFFICyErrorSetRaisedFromPyError(PyObject* py_err); /* * \brief Create an argument setter for a given Python argument type. * \param arg The Python argument to be set. * \param out The output argument setter. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ __PYX_EXTERN_C int TVMFFICyArgSetterFactory(PyObject* arg, TVMFFIPyArgSetter* out); /* * \brief Create a callback arg setter for a given type index. * \param type_index The type index of the argument. * \param out The output callback arg setter. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ __PYX_EXTERN_C int TVMFFICyCallbackArgSetterFactory(int32_t type_index, TVMFFIPyCallbackArgSetter* out); //--------------------------------------------------------------------------------------------- // The function call manager section //--------------------------------------------------------------------------------------------- /*! * \brief A manager class that handles python ffi calls. */ class TVMFFIPyCallManager { public: /*! * \brief Get the thread local call manager. * \return The thread local call manager. */ static TVMFFIPyCallManager* ThreadLocal() { static thread_local TVMFFIPyCallManager inst; return &inst; } /*! * \brief Call a function with a variable number of arguments * \param func_handle The handle of the function to call * \param py_arg_tuple The arguments to the function * \param result The result of the function * \param c_api_ret_code The return code of the C-call * \param release_gil Whether to release the GIL * \param optional_out_ctx_dlpack_api The DLPack exchange API to be used for the result * \return 0 on when there is no python error, -1 on python error * \note When an error happens on FFI side, we should return 0 and set c_api_ret_code */ TVM_FFI_INLINE int FuncCall(void* func_handle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, bool release_gil, const DLPackExchangeAPI** optional_out_ctx_dlpack_api) { int64_t num_args = PyTuple_Size(py_arg_tuple); if (TVM_FFI_PREDICT_FALSE(num_args == -1)) return -1; try { // allocate a call stack TVMFFIPyCallContext ctx(&call_stack_, num_args); // Iterate over the arguments and set them for (int64_t i = 0; i < num_args; ++i) { PyObject* py_arg = PyTuple_GetItem(py_arg_tuple, i); TVMFFIAny* c_arg = ctx.packed_args + i; if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1; } TVMFFIStreamHandle prev_stream = nullptr; DLPackManagedTensorAllocator prev_tensor_allocator = nullptr; // setup stream context if needed if (ctx.device_type != -1) { c_api_ret_code[0] = TVMFFIEnvSetStream(ctx.device_type, ctx.device_id, ctx.stream, &prev_stream); // setting failed, directly return if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0; } if (ctx.dlpack_c_exchange_api != nullptr && ctx.dlpack_c_exchange_api->managed_tensor_allocator != nullptr) { c_api_ret_code[0] = TVMFFIEnvSetDLPackManagedTensorAllocator( ctx.dlpack_c_exchange_api->managed_tensor_allocator, 0, &prev_tensor_allocator); if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0; } // call the function if (release_gil) { // release the GIL Py_BEGIN_ALLOW_THREADS; c_api_ret_code[0] = TVMFFIFunctionCall(func_handle, ctx.packed_args, num_args, result); Py_END_ALLOW_THREADS; } else { c_api_ret_code[0] = TVMFFIFunctionCall(func_handle, ctx.packed_args, num_args, result); } // restore the original stream if (ctx.device_type != -1 && prev_stream != ctx.stream) { // always try recover first, even if error happens if (TVM_FFI_PREDICT_FALSE( TVMFFIEnvSetStream(ctx.device_type, ctx.device_id, prev_stream, nullptr) != 0)) { // recover failed, set python error PyErr_SetString(PyExc_RuntimeError, "Failed to recover stream"); return -1; } } if (ctx.dlpack_c_exchange_api != nullptr && prev_tensor_allocator != ctx.dlpack_c_exchange_api->managed_tensor_allocator) { // note: we cannot set the error value to c_api_ret_code[0] here because it // will be overwritten by the error value from the function call if (TVM_FFI_PREDICT_FALSE( TVMFFIEnvSetDLPackManagedTensorAllocator(prev_tensor_allocator, 0, nullptr) != 0)) { PyErr_SetString(PyExc_RuntimeError, "Failed to recover DLPack managed tensor allocator"); return -1; } // return error after if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0; } if (optional_out_ctx_dlpack_api != nullptr && ctx.dlpack_c_exchange_api != nullptr) { *optional_out_ctx_dlpack_api = ctx.dlpack_c_exchange_api; } return 0; } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); return -1; } } /* * \brief Call a constructor with a variable number of arguments * * This function is similar to FuncCall, but it will not set the * stream and tensor allocator, instead, it will synchronize the TVMFFIPyCallContext * with the parent context. This behavior is needed for nested conversion of arguments * where detected argument setting needs to be synchronized with final call. * * This function will also not release the GIL since constructor call is usually cheap. * * \param func_handle The handle of the constructor to call * \param py_arg_tuple The arguments to the constructor * \param result The result of the constructor * \param c_api_ret_code The return code of the constructor * \param parent_ctx The parent call context to * \return 0 on success, -1 on failure */ TVM_FFI_INLINE int ConstructorCall(void* func_handle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, TVMFFIPyCallContext* parent_ctx) { int64_t num_args = PyTuple_Size(py_arg_tuple); if (TVM_FFI_PREDICT_FALSE(num_args == -1)) return -1; try { // allocate a call stack TVMFFIPyCallContext ctx(&call_stack_, num_args); // Iterate over the arguments and set them for (int64_t i = 0; i < num_args; ++i) { PyObject* py_arg = PyTuple_GetItem(py_arg_tuple, i); TVMFFIAny* c_arg = ctx.packed_args + i; if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1; } c_api_ret_code[0] = TVMFFIFunctionCall(func_handle, ctx.packed_args, num_args, result); // propagate the call context to the parent context if (parent_ctx != nullptr) { // stream and current device information if (parent_ctx->device_type == -1) { parent_ctx->device_type = ctx.device_type; parent_ctx->device_id = ctx.device_id; parent_ctx->stream = ctx.stream; } // DLPack exchange API if (parent_ctx->dlpack_c_exchange_api == nullptr) { parent_ctx->dlpack_c_exchange_api = ctx.dlpack_c_exchange_api; } } return 0; } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); return -1; } } TVM_FFI_INLINE int SetField(void* field_setter, int64_t field_flags, void* field_ptr, PyObject* py_arg, int* c_api_ret_code) { try { TVMFFIPyCallContext ctx(&call_stack_, 1); TVMFFIAny* c_arg = ctx.packed_args; if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1; if (!(field_flags & kTVMFFIFieldFlagBitSetterIsFunctionObj)) { auto setter = reinterpret_cast(field_setter); c_api_ret_code[0] = (*setter)(field_ptr, c_arg); } else { TVMFFIAny args[2]{}; args[0].type_index = kTVMFFIOpaquePtr; args[0].v_ptr = field_ptr; args[1] = *c_arg; TVMFFIAny result{}; result.type_index = kTVMFFINone; c_api_ret_code[0] = TVMFFIFunctionCall(static_cast(field_setter), args, 2, &result); } return 0; } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); return -1; } } TVM_FFI_INLINE int PyObjectToFFIAny(PyObject* py_arg, TVMFFIAny* out, int* c_api_ret_code) { try { TVMFFIPyCallContext ctx(&call_stack_, 1); TVMFFIAny* c_arg = ctx.packed_args; if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1; c_api_ret_code[0] = TVMFFIAnyViewToOwnedAny(c_arg, out); return 0; } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); return -1; } } /*! * \brief Set an py_arg to out. * \param ctx The call context * \param py_arg The python argument to be set * \param out The output argument * \return 0 on success, -1 on failure */ TVM_FFI_INLINE int SetArgument(TVMFFIPyCallContext* ctx, PyObject* py_arg, TVMFFIAny* out) { PyTypeObject* py_type = Py_TYPE(py_arg); // pre-zero the output argument, modulo the type index out->type_index = kTVMFFINone; out->zero_padding = 0; out->v_int64 = 0; // find the pre-cached setter // This class is thread-local, so we don't need to worry about race condition auto it = arg_dispatch_map_.find(py_type); if (TVM_FFI_PREDICT_TRUE(it != arg_dispatch_map_.end())) { TVMFFIPyArgSetter setter = it->second; // if error happens, propagate it back if (TVM_FFI_PREDICT_FALSE(setter(ctx, py_arg, out) != 0)) return -1; } else { // no dispatch found, query and create a new one. TVMFFIPyArgSetter setter; // propagate python error back if (TVM_FFI_PREDICT_FALSE(TVMFFICyArgSetterFactory(py_arg, &setter) != 0)) { return -1; } // update dispatch table arg_dispatch_map_.emplace(py_type, setter); if (TVM_FFI_PREDICT_FALSE(setter(ctx, py_arg, out) != 0)) return -1; } return 0; } /*! * \brief Get the size of the arg dispatch map * \return The size of the arg dispatch map */ size_t GetArgDispatchMapSize() const { return arg_dispatch_map_.size(); } /*! * \brief Convert a borrowed TVMFFIAny (AnyView) into a new-reference PyObject*. * \param dlpack_exchange_api The DLPack exchange API (may be NULL). * \param arg The borrowed TVMFFIAny to convert. * \param py_arg The output PyObject*. * \return 0 on success, -1 on failure. PyError should be set if -1 is returned. */ TVM_FFI_INLINE int SetPyCallbackArg(const DLPackExchangeAPI* dlpack_exchange_api, const TVMFFIAny* arg, PyObject** out) { size_t type_index = static_cast(arg->type_index); // Mirror of SetArgument for the C++ -> Python callback path: each per-type // callback arg setter is responsible for its own refcount. // hot path: cached hit if (type_index < callback_arg_dispatch_table_.size() && callback_arg_dispatch_table_[type_index].func != nullptr) { TVMFFIPyCallbackArgSetter* setter = &callback_arg_dispatch_table_[type_index]; return setter->func(setter, dlpack_exchange_api, arg, out); } // cold path: grow and populate via factory if (type_index >= callback_arg_dispatch_table_.size()) { // initialize empty entries with nullptr callback_arg_dispatch_table_.resize(type_index + 1, TVMFFIPyCallbackArgSetter{nullptr}); } TVMFFIPyCallbackArgSetter* setter = &callback_arg_dispatch_table_[type_index]; if (TVMFFICyCallbackArgSetterFactory(static_cast(type_index), setter) != 0) { return -1; } return setter->func(setter, dlpack_exchange_api, arg, out); } /*! * \brief Python Callback function entry * * \param context The TVMFFIPyCallbackClosure* holding the Python callable * and optional DLPack exchange API. * \param packed_args The packed FFI arguments. * \param num_args Number of arguments. * \param result Output FFI result. * \return 0 on success, -1 on error. */ TVM_FFI_INLINE int PyCallback(void* context, const TVMFFIAny* packed_args, int32_t num_args, TVMFFIAny* result) noexcept { TVMFFIPyWithAttachedThreadState thread_state; auto* closure = static_cast(context); // Wrap the body in try/catch so any C++ exception raised by the stack // allocator (TVMFFIPyCallbackContext / TVMFFIPyCallContext), dispatch // table resize in SetPyCallbackArg, or unordered_map::emplace in // SetArgument is converted into a PyErr + FFI error instead of // triggering std::terminate via the noexcept contract. try { TVMFFIPyCallbackContext cb_ctx(&call_stack_, num_args); // Step 1: Convert each packed arg (borrowed AnyView) to a PyObject* for (int32_t i = 0; i < num_args; ++i) { if (TVM_FFI_PREDICT_FALSE(SetPyCallbackArg(closure->dlpack_exchange_api, &packed_args[i], &cb_ctx.py_args[i]) != 0)) { ForwardPyErrorToFFI(); return -1; } // must set active arguments count to ensure correct recycling cb_ctx.num_active_py_args = i + 1; } // Step 2: Call the Python function via vectorcall. Wrap py_result in // a RAII guard so its +1 is released on every exit path, including // the C++ exception path (e.g., bad_alloc from ret_ctx construction // or SetArgument's emplace). #if PY_VERSION_HEX >= 0x03090000 PyObject* py_result_raw = PyObject_Vectorcall(closure->callable, cb_ctx.py_args, static_cast(num_args), nullptr); #else // backward compatibility for Python 3.8 PyObject* py_result_raw = _PyObject_Vectorcall(closure->callable, cb_ctx.py_args, static_cast(num_args), nullptr); #endif struct PyResultGuard { PyObject* p; ~PyResultGuard() { if (p != nullptr) Py_DecRef(p); } } py_result{py_result_raw}; if (py_result.p == Py_None) { // fast path for Py_None result->type_index = kTVMFFINone; result->zero_padding = 0; result->v_int64 = 0; return 0; } else if (py_result.p != nullptr) { // normal return // Use SetArgument on a temporary view slot, then promote to owned. // Note: SetArgument only borrows py_result's chandle into `view`; it // does NOT inc. py_result must stay alive until AFTER // TVMFFIAnyViewToOwnedAny has promoted the view to an owned ref, // otherwise dec'ing py_result first could free the underlying object // (e.g. if py_result owns the last ref to a freshly created tensor). // The guard's destructor runs AFTER the return value is computed. TVMFFIPyCallContext ret_ctx(&call_stack_, 1); TVMFFIAny* view = ret_ctx.packed_args; if (TVM_FFI_PREDICT_FALSE(SetArgument(&ret_ctx, py_result.p, view) != 0)) { ForwardPyErrorToFFI(); return -1; } // TLS FFI error set on failure. return TVMFFIAnyViewToOwnedAny(view, result); } else { // vectorcall failed ForwardPyErrorToFFI(); return -1; } } catch (const std::exception& ex) { // very rare, catch c++ exception and set python error PyErr_SetString(PyExc_RuntimeError, ex.what()); ForwardPyErrorToFFI(); return -1; } } /*! * \brief Fetch the current Python exception and forward it to * TVMFFICyErrorSetRaisedFromPyError, then clear the Python error indicator. * * This helper correctly extracts the exception *value* (not just the type * returned by PyErr_Occurred()) so that set_last_ffi_error can access the * message and traceback. */ TVM_FFI_COLD_CODE static void ForwardPyErrorToFFI() noexcept { #if PY_VERSION_HEX >= 0x030C0000 // Python 3.12+: PyErr_Fetch / PyErr_NormalizeException are deprecated. // PyErr_GetRaisedException returns an already-normalized exception // instance and clears the indicator. Traceback is attached as usual. PyObject* pvalue = PyErr_GetRaisedException(); if (pvalue != nullptr) { TVMFFICyErrorSetRaisedFromPyError(pvalue); Py_DecRef(pvalue); } #else // Python 3.9 - 3.11. PyObject *ptype, *pvalue, *ptraceback; PyErr_Fetch(&ptype, &pvalue, &ptraceback); PyErr_NormalizeException(&ptype, &pvalue, &ptraceback); if (ptraceback != nullptr) { PyException_SetTraceback(pvalue, ptraceback); } TVMFFICyErrorSetRaisedFromPyError(pvalue); Py_DecRef(ptype); Py_DecRef(pvalue); Py_DecRef(ptraceback); #endif } private: TVMFFIPyCallManager() { static constexpr size_t kDefaultDispatchCapacity = 32; arg_dispatch_map_.reserve(kDefaultDispatchCapacity); // Pre-allocate callback arg dispatch table for static type indices static constexpr size_t kDefaultCallbackArgDispatchCapacity = 128; callback_arg_dispatch_table_.resize(kDefaultCallbackArgDispatchCapacity); } // internal arg dispatch map: type -> argument setter std::unordered_map arg_dispatch_map_; // call stack TVMFFIPyCallStack call_stack_; // callback arg setter dispatch table indexed by type_index (view-based path // used by PyCallback; see TVMFFIPyCallbackArgSetter docs above) std::vector callback_arg_dispatch_table_; }; /*! * \brief Call a function with a variable number of arguments * \param func_handle The handle of the function to call * \param py_arg_tuple The arguments to the function * \param result The result of the function * \param c_api_ret_code The return code of the function * \param release_gil Whether to release the GIL * \param out_ctx_dlpack_api The DLPack exchange API to be used for the result * \return 0 on success, nonzero on failure */ TVM_FFI_INLINE int TVMFFIPyFuncCall(void* func_handle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, bool release_gil = true, const DLPackExchangeAPI** out_ctx_dlpack_api = nullptr) { return TVMFFIPyCallManager::ThreadLocal()->FuncCall( func_handle, py_arg_tuple, result, c_api_ret_code, release_gil, out_ctx_dlpack_api); } /*! * \brief Call a constructor function with a variable number of arguments * * This function is similar to TVMFFIPyFuncCall, but it will not set the * stream and tensor allocator. Instead, it will synchronize the TVMFFIPyCallContext * with the parent context. This behavior is needed for nested conversion of arguments * where detected argument settings need to be synchronized with the final call. * * This function will also not release the GIL since constructor call is usually cheap. * * \param func_handle The handle of the function to call * \param py_arg_tuple The arguments to the constructor * \param result The result of the constructor * \param c_api_ret_code The return code of the constructor * \param parent_ctx The parent call context * \param release_gil Whether to release the GIL * \param out_dlpack_exporter The DLPack exporter to be used for the result * \return 0 on success, nonzero on failure */ TVM_FFI_INLINE int TVMFFIPyConstructorCall(void* func_handle, PyObject* py_arg_tuple, TVMFFIAny* result, int* c_api_ret_code, TVMFFIPyCallContext* parent_ctx) { return TVMFFIPyCallManager::ThreadLocal()->ConstructorCall(func_handle, py_arg_tuple, result, c_api_ret_code, parent_ctx); } /*! * \brief Set a field of a FFI object * \param field_setter The field setter (function pointer or FunctionObj handle) * \param field_flags The field flags (to dispatch between function pointer and FunctionObj) * \param field_ptr The pointer to the field * \param py_arg The python argument to be set * \param c_api_ret_code The return code of the function * \return 0 on success, nonzero on failure */ TVM_FFI_INLINE int TVMFFIPyCallFieldSetter(void* field_setter, int64_t field_flags, void* field_ptr, PyObject* py_arg, int* c_api_ret_code) { return TVMFFIPyCallManager::ThreadLocal()->SetField(field_setter, field_flags, field_ptr, py_arg, c_api_ret_code); } /*! * \brief Set an python argument to a FFI Any using the generic dispatcher in call manager * \param ctx The call context * \param py_arg_tvm_ffi_value The python argument to be set using the __tvm_ffi_value__ protocol * \param out The output argument * \return 0 on success, nonzero on failure */ TVM_FFI_INLINE int TVMFFIPySetArgumentGenericDispatcher(TVMFFIPyCallContext* ctx, PyObject* py_arg_tvm_ffi_value, TVMFFIAny* out) { return TVMFFIPyCallManager::ThreadLocal()->SetArgument(ctx, py_arg_tvm_ffi_value, out); } /*! * \brief Convert a Python object to a FFI Any * \param py_arg The python argument to be set * \param out The output argument * \param c_api_ret_code The return code of the function * \return 0 on success, nonzero on failure */ TVM_FFI_INLINE int TVMFFIPyPyObjectToFFIAny(PyObject* py_arg, TVMFFIAny* out, int* c_api_ret_code) { return TVMFFIPyCallManager::ThreadLocal()->PyObjectToFFIAny(py_arg, out, c_api_ret_code); } /*! * \brief Get the size of the arg dispatch map * \return The size of the arg dispatch map */ TVM_FFI_INLINE size_t TVMFFIPyGetArgDispatchMapSize() { return TVMFFIPyCallManager::ThreadLocal()->GetArgDispatchMapSize(); } //--------------------------------------------------------------------------------------------- // Free function wrapper for the Python callback path. // Mirrors the pattern of TVMFFIPyFuncCall / TVMFFIPyConstructorCall: a top-level // TVM_FFI_INLINE free function that forwards to the thread-local manager. //--------------------------------------------------------------------------------------------- /*! * \brief C-callable Python callback entry point (TVMFFISafeCallType shape). * * Forwards to TVMFFIPyCallManager::ThreadLocal()->PyCallback. Designed to be * installed as the safe_call pointer for FFI functions that wrap a Python * callable. * * \note The `context` argument is interpreted as a TVMFFIPyCallbackClosure* * by the manager (see TVMFFIPyConvertPyCallback). */ TVM_FFI_INLINE int TVMFFIPyCallback(void* context, const TVMFFIAny* packed_args, int32_t num_args, TVMFFIAny* result) noexcept { return TVMFFIPyCallManager::ThreadLocal()->PyCallback(context, packed_args, num_args, result); } /*! * \brief Create an FFI function handle from a Python callable + optional DLPack exchange API. * * Allocates a TVMFFIPyCallbackClosure on the heap, IncRefs the callable, and * registers it with the FFI function-creation API using TVMFFIPyCallback as the * safe-call entry point and TVMFFIPyCallbackClosure::Deleter as the deleter. * * Returns the raw FFI return code (TLS FFI error set on failure). The Cython * caller uses CHECK_CALL to translate it into a Python exception. * * \param callable The Python callable to wrap. Must be non-NULL. * \param dlpack_api Optional DLPack exchange API. May be NULL. * \param out_handle Destination for the new FFI function handle. * \return The return code from TVMFFIFunctionCreate (0 on success). */ TVM_FFI_INLINE int TVMFFIPyConvertPyCallback(PyObject* callable, const DLPackExchangeAPI* dlpack_api, TVMFFIObjectHandle* out_handle) noexcept { // Use nothrow new: plain `new` can throw std::bad_alloc, which in this // noexcept function would trigger std::terminate. On allocation failure, // set PyErr and return -1 so the Cython caller's CHECK_CALL surfaces it. auto* raw = new (std::nothrow) TVMFFIPyCallbackClosure{callable, dlpack_api}; if (raw == nullptr) { PyErr_NoMemory(); return -1; } // The callable's +1 is owned by the closure; TVMFFIPyCallbackClosure::Deleter // is responsible for Py_DecRef on destruction. By wiring the same Deleter as // the unique_ptr deleter, the failure path below (unique_ptr unwind) runs // the same cleanup as the success path (invoked by the FFI runtime). Py_IncRef(callable); std::unique_ptr closure( raw, &TVMFFIPyCallbackClosure::Deleter); int rc = TVMFFIFunctionCreate(closure.get(), &TVMFFIPyCallback, &TVMFFIPyCallbackClosure::Deleter, out_handle); // On success, transfer ownership to the FFI function; on failure, let // unique_ptr unwind via Deleter (decrefs the callable, frees the closure). if (rc == 0) closure.release(); return rc; } /*! * \brief Push a temporary FFI object to the call context that will be recycled after the call * \param ctx The call context * \param arg The FFI object to push */ TVM_FFI_INLINE void TVMFFIPyPushTempFFIObject(TVMFFIPyCallContext* ctx, TVMFFIObjectHandle arg) noexcept { // invariance: each ArgSetter can have at most one temporary Python object // so it ensures that we won't overflow the temporary Python object stack ctx->temp_ffi_objects[ctx->num_temp_ffi_objects++] = arg; } /*! * \brief Push a temporary Python object to the call context that will be recycled after the call * \param ctx The call context * \param arg The Python object to push */ TVM_FFI_INLINE void TVMFFIPyPushTempPyObject(TVMFFIPyCallContext* ctx, PyObject* arg) noexcept { // invariance: each ArgSetter can have at most one temporary Python object // so it ensures that we won't overflow the temporary Python object stack Py_IncRef(arg); ctx->temp_py_objects[ctx->num_temp_py_objects++] = arg; } /*! * \brief Push Extra temporary Python object to the call context that may go beyond one temp per * argument budget, mainly used by value protocol. * \param ctx The call context * \param arg The Python object to push */ TVM_FFI_INLINE void TVMFFIPyPushExtraTempPyObject(TVMFFIPyCallContext* ctx, PyObject* arg) { Py_IncRef(arg); ctx->call_stack->extra_temp_py_objects_stack.emplace_back(arg); } //---------------------------------------------------------- // Helpers for MLIR redirection //---------------------------------------------------------- /*! * \brief Function specialization that leverages MLIR packed safe call definitions. * * The MLIR execution engine generates functions that correspond to the packed signature. * As of now, it is hard to access the raw extern C function pointer of SafeCall * directly when we declare the signature in LLVM dialect. * * Note that in theory, the MLIR execution engine should be able to support * some form of "extern C" feature that directly exposes the function pointers * of C-compatible functions with an attribute tag. So we keep this feature * in the Python helper layer for now in case the MLIR execution engine supports it in the future. * * This helper enables us to create ffi::Function from the MLIR packed * safe call function pointer instead of following the redirection pattern * in `TVMFFIPyMLIRPackedSafeCall::Invoke`. * * \sa TVMFFIPyMLIRPackedSafeCall::Invoke */ class TVMFFIPyMLIRPackedSafeCall { public: TVMFFIPyMLIRPackedSafeCall(void (*mlir_packed_safe_call)(void**), PyObject* keep_alive_object) : mlir_packed_safe_call_(mlir_packed_safe_call), keep_alive_object_(keep_alive_object) { if (keep_alive_object_) { Py_IncRef(keep_alive_object_); } } ~TVMFFIPyMLIRPackedSafeCall() { TVMFFIPyWithAttachedThreadState thread_state; if (keep_alive_object_) { Py_DecRef(keep_alive_object_); } } static int Invoke(void* func, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* rv) { TVMFFIPyMLIRPackedSafeCall* self = reinterpret_cast(func); int ret_code = 0; void* handle = nullptr; void* mlir_args[] = {&handle, const_cast(&args), &num_args, &rv, &ret_code}; (*self->mlir_packed_safe_call_)(mlir_args); return ret_code; } static void Deleter(void* self) { delete static_cast(self); } private: void (*mlir_packed_safe_call_)(void**); PyObject* keep_alive_object_; }; /*! * \brief Create a TVMFFIPyMLIRPackedSafeCall handle * \param mlir_packed_safe_call The MLIR packed safe call function * \param keep_alive_object The keep alive object * \return The TVMFFIPyMLIRPackedSafeCall object */ void* TVMFFIPyMLIRPackedSafeCallCreate(void (*mlir_packed_safe_call)(void**), PyObject* keep_alive_object) { return new TVMFFIPyMLIRPackedSafeCall(mlir_packed_safe_call, keep_alive_object); } /*! * \brief Call the MLIR packed safe call function * \param self The TVMFFIPyMLIRPackedSafeCall object * \param args The arguments * \param num_args The number of arguments * \param rv The result * \return The return code */ int TVMFFIPyMLIRPackedSafeCallInvoke(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* rv) { return TVMFFIPyMLIRPackedSafeCall::Invoke(self, args, num_args, rv); } /*! * \brief Delete the TVMFFIPyMLIRPackedSafeCall object * \param self The TVMFFIPyMLIRPackedSafeCall object */ void TVMFFIPyMLIRPackedSafeCallDeleter(void* self) { return TVMFFIPyMLIRPackedSafeCall::Deleter(self); } /*! * \brief Deleter for Python objects * \param py_obj The Python object to delete */ extern "C" void TVMFFIPyObjectDeleter(void* py_obj) noexcept { TVMFFIPyWithAttachedThreadState thread_state; Py_DecRef(static_cast(py_obj)); } /* * \brief Dummy target to ensure testing is linked and we can run testcases */ extern "C" TVM_FFI_DLL int TVMFFITestingDummyTarget(); #endif // TVM_FFI_PYTHON_HELPERS_H_ tvm-ffi-0.1.12/python/tvm_ffi/cython/type_info.pxi000066400000000000000000001274251521067262500221610ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import dataclasses import json import typing import collections.abc from functools import cached_property from typing import Optional, Any from io import StringIO try: from types import UnionType as _UnionType except ImportError: _UnionType = None cdef class FieldGetter: cdef dict __dict__ cdef TVMFFIFieldGetter getter cdef int64_t offset def __call__(self, CObject obj): cdef TVMFFIAny result cdef int c_api_ret_code cdef void* field_ptr = ((obj).chandle) + self.offset result.type_index = kTVMFFINone result.v_int64 = 0 c_api_ret_code = self.getter(field_ptr, &result) CHECK_CALL(c_api_ret_code) return make_ret(result) cdef class FieldSetter: cdef dict __dict__ cdef void* setter cdef int64_t offset cdef int64_t flags def __call__(self, CObject obj, value): cdef int c_api_ret_code cdef void* field_ptr = ((obj).chandle) + self.offset TVMFFIPyCallFieldSetter( self.setter, self.flags, field_ptr, value, &c_api_ret_code ) # NOTE: logic is same as check_call # directly inline here to simplify backtrace if c_api_ret_code == 0: return # backward compact with error already set case # TODO(tqchen): remove after we move beyond a few versions. if c_api_ret_code == -2: raise raise_existing_error() # epecial handle env error already set error = move_from_last_error() if error.kind == "EnvErrorAlreadySet": raise raise_existing_error() raise error.py_error() _TYPE_SCHEMA_ORIGIN_CONVERTER = { # A few Python-native types "Variant": "Union", "Optional": "Optional", "Tuple": "tuple", "ffi.Function": "Callable", "ffi.Array": "Array", "ffi.List": "List", "ffi.Map": "Map", "ffi.Dict": "Dict", # OpaquePyObject accepts any Python value at the FFI boundary (the C++ # side wraps it opaquely), so mapping to "Any" is semantically correct. "ffi.OpaquePyObject": "Any", "ffi.Object": "Object", "ffi.Tensor": "Tensor", "DLTensor*": "Tensor", # ctype types "void*": "ctypes.c_void_p", # bytes "TVMFFIByteArray*": "bytes", "ffi.SmallBytes": "bytes", "ffi.Bytes": "bytes", # strings "std::string": "str", "const char*": "str", "ffi.SmallStr": "str", "ffi.String": "str", "DataType": "dtype", # C++ STL types (emitted by TypeTraits in include/tvm/ffi/extra/stl.h) "std::vector": "Array", "std::optional": "Optional", "std::variant": "Union", "std::tuple": "tuple", "std::map": "Map", "std::unordered_map": "Map", "std::function": "Callable", # Rvalue reference (C++ move semantics). Python has no move semantics, # so the checker treats it as a plain Object reference. "ObjectRValueRef": "Object", } # Sentinel for structural types (Optional, Union) that have no single type_index _ORIGIN_TYPE_INDEX_STRUCTURAL = -2 # Sentinel for unknown/unresolved origins _ORIGIN_TYPE_INDEX_UNKNOWN = -3 # Map origin string -> type_index for known types _ORIGIN_TO_TYPE_INDEX = { "None": kTVMFFINone, "int": kTVMFFIInt, "bool": kTVMFFIBool, "float": kTVMFFIFloat, "str": kTVMFFIStr, "bytes": kTVMFFIBytes, "Device": kTVMFFIDevice, "dtype": kTVMFFIDataType, "ctypes.c_void_p": kTVMFFIOpaquePtr, "Tensor": kTVMFFITensor, "Object": kTVMFFIObject, "Callable": kTVMFFIFunction, "Array": kTVMFFIArray, "List": kTVMFFIList, "Map": kTVMFFIMap, "Dict": kTVMFFIDict, "Any": kTVMFFIAny, } # Reverse map: type_index -> origin string _TYPE_INDEX_TO_ORIGIN = {v: k for k, v in _ORIGIN_TO_TYPE_INDEX.items()} # Low-level type indices that alias canonical origins _TYPE_INDEX_TO_ORIGIN[kTVMFFIDLTensorPtr] = "Tensor" _TYPE_INDEX_TO_ORIGIN[kTVMFFIRawStr] = "str" _TYPE_INDEX_TO_ORIGIN[kTVMFFIByteArrayPtr] = "bytes" _TYPE_INDEX_TO_ORIGIN[kTVMFFISmallStr] = "str" _TYPE_INDEX_TO_ORIGIN[kTVMFFISmallBytes] = "bytes" _TYPE_INDEX_TO_ORIGIN[kTVMFFIObjectRValueRef] = "Object" @dataclasses.dataclass(repr=False) class TypeSchema: """Type schema that describes a TVM FFI type. The schema is expressed using a compact JSON-compatible structure and can be rendered as a Python typing string with :py:meth:`repr`. """ origin: str args: tuple["TypeSchema", ...] | None = None origin_type_index: int = dataclasses.field(default=_ORIGIN_TYPE_INDEX_UNKNOWN, repr=False) def __post_init__(self): origin = self.origin args = self.args if args is not None and not isinstance(args, tuple): args = tuple(args) self.args = args if origin != "tuple" and args is None: args = () self.args = args if origin == "Union": if len(args) < 2: raise ValueError("Union must have at least two arguments") elif origin == "Optional": if len(args) != 1: raise ValueError("Optional must have exactly one argument") elif origin in ("list", "Array", "List"): if len(args) not in (0, 1): raise ValueError(f"{origin} must have 0 or 1 argument") if args == (): self.args = (TypeSchema("Any"),) elif origin in ("dict", "Map", "Dict"): if len(args) not in (0, 2): raise ValueError(f"{origin} must have 0 or 2 arguments") if args == (): self.args = (TypeSchema("Any"), TypeSchema("Any")) elif origin == "tuple": pass # tuple can have arbitrary number of arguments # Compute origin_type_index if not already set if self.origin_type_index == _ORIGIN_TYPE_INDEX_UNKNOWN: if origin in ("Optional", "Union"): self.origin_type_index = _ORIGIN_TYPE_INDEX_STRUCTURAL elif origin in _ORIGIN_TO_TYPE_INDEX: self.origin_type_index = _ORIGIN_TO_TYPE_INDEX[origin] else: # Try to resolve as a registered object type key tindex = _object_type_key_to_index(origin) if tindex is not None: self.origin_type_index = tindex @cached_property def _converter(self): """Lazily build the type converter on first use. Deferred construction ensures all object types are registered by the time the converter is built. Raises TypeError for unresolvable origins. """ return _build_converter(self) def __repr__(self) -> str: return self.repr(ty_map=None) @staticmethod def from_json_obj(obj: dict[str, Any]) -> "TypeSchema": """Construct a :class:`TypeSchema` from a parsed JSON object. Non-dict elements in the ``"args"`` list (e.g., numeric lengths emitted by ``std::array`` TypeTraits) are silently skipped. """ if not isinstance(obj, dict) or "type" not in obj: raise TypeError( f"expected schema dict with 'type' key, got {type(obj).__name__}" ) origin = obj["type"] origin = _TYPE_SCHEMA_ORIGIN_CONVERTER.get(origin, origin) if "args" not in obj: return TypeSchema(origin) raw_args = obj["args"] if not isinstance(raw_args, (list, tuple)): raw_args = () args = tuple( TypeSchema.from_json_obj(a) for a in raw_args if isinstance(a, dict) ) return TypeSchema(origin, args) @staticmethod def from_json_str(s: str) -> "TypeSchema": """Construct a :class:`TypeSchema` from a JSON string.""" return TypeSchema.from_json_obj(json.loads(s)) @staticmethod def from_type_index(type_index: int, args: "tuple[TypeSchema, ...]" = ()) -> "TypeSchema": """Construct a :class:`TypeSchema` from a type_index and optional args. Parameters ---------- type_index : int A valid TVM FFI type index (e.g., ``kTVMFFIInt``, ``kTVMFFIArray``, or an object type index from ``_object_type_key_to_index``). Passing an unregistered index triggers a fatal C++ assertion; callers must ensure the index was obtained from the type registry. args : tuple[TypeSchema, ...], optional Type arguments for parameterized types (e.g., element type for Array). Returns ------- TypeSchema A new schema with the origin resolved from the type index. """ origin = _TYPE_INDEX_TO_ORIGIN.get(type_index, None) if origin is None: origin = _type_index_to_key(type_index) return TypeSchema(origin, args, origin_type_index=type_index) @staticmethod def from_annotation(annotation: object) -> "TypeSchema": """Construct a :class:`TypeSchema` from a Python type annotation. Parameters ---------- annotation : object A Python type annotation such as ``int``, ``list[int]``, ``Optional[str]``, ``Union[int, str]``, ``Callable[[int], str]``, or a registered :class:`CObject` subclass. Returns ------- TypeSchema The corresponding schema. Raises ------ TypeError If the annotation cannot be mapped to a TypeSchema. Examples -------- >>> TypeSchema.from_annotation(int) int >>> TypeSchema.from_annotation(list[int]) List[int] >>> TypeSchema.from_annotation(tuple[int, ...]) Array[int] """ # --- Singletons --- if annotation is type(None) or annotation is None: return TypeSchema("None") if annotation is typing.Any: return TypeSchema("Any") # --- Bare builtin scalar types --- if annotation is bool: return TypeSchema("bool") if annotation is int: return TypeSchema("int") if annotation is float: return TypeSchema("float") if annotation is str: return TypeSchema("str") if annotation is bytes: return TypeSchema("bytes") # --- Bare container types (unparameterised) --- if annotation is list: return TypeSchema("List") if annotation is dict: return TypeSchema("Dict") if annotation is tuple: return TypeSchema("tuple") if annotation is collections.abc.Callable: return TypeSchema("Callable") # --- Python 3.10+ union syntax (X | Y) --- if _UnionType is not None and isinstance(annotation, _UnionType): return _annotation_union(typing.get_args(annotation)) # --- Generic aliases (list[int], Optional[T], etc.) --- origin = typing.get_origin(annotation) targs = typing.get_args(annotation) if origin is typing.Union: return _annotation_union(targs) if origin is list: if len(targs) > 1: raise TypeError( f"list takes at most 1 type argument, got {len(targs)}" ) if targs: return TypeSchema("List", (TypeSchema.from_annotation(targs[0]),)) return TypeSchema("List") if origin is dict: if len(targs) == 1 or len(targs) > 2: raise TypeError( f"dict requires 0 or 2 type arguments, got {len(targs)}" ) if len(targs) == 2: return TypeSchema("Dict", ( TypeSchema.from_annotation(targs[0]), TypeSchema.from_annotation(targs[1]), )) return TypeSchema("Dict") if origin is tuple: if len(targs) == 2 and targs[1] is Ellipsis: # tuple[T, ...] → homogeneous variable-length → Array return TypeSchema("Array", (TypeSchema.from_annotation(targs[0]),)) if targs: return TypeSchema( "tuple", tuple(TypeSchema.from_annotation(a) for a in targs), ) if annotation is not tuple: return TypeSchema("tuple", ()) return TypeSchema("tuple") if origin is collections.abc.Callable: if len(targs) == 2: params, ret = targs ret_schema = TypeSchema.from_annotation(ret) if isinstance(params, list): # Callable[[P1, P2], R] → (R, P1, P2) param_schemas = tuple( TypeSchema.from_annotation(p) for p in params ) return TypeSchema("Callable", (ret_schema,) + param_schemas) # Callable[..., R] return TypeSchema("Callable", (ret_schema,)) return TypeSchema("Callable") # --- Parameterised CObject subclasses (Array[int], Dict[str, V], …) --- if isinstance(origin, type) and issubclass(origin, CObject): return _annotation_cobject(origin, targs) # --- Bare (unparameterised) CObject subclasses --- if isinstance(annotation, type) and issubclass(annotation, CObject): return _annotation_cobject(annotation, ()) # --- PyNativeObject subclasses (String, Bytes) --- if isinstance(annotation, type) and issubclass(annotation, PyNativeObject): if issubclass(annotation, str): return TypeSchema("str") if issubclass(annotation, bytes): return TypeSchema("bytes") # --- Non-CObject cdef classes with known origins --- if annotation is DataType or (_CLASS_DTYPE is not None and annotation is _CLASS_DTYPE): return TypeSchema("dtype") if annotation is Device or (_CLASS_DEVICE is not None and annotation is _CLASS_DEVICE): return TypeSchema("Device") # --- ctypes.c_void_p --- import ctypes as _ctypes if annotation is _ctypes.c_void_p: return TypeSchema("ctypes.c_void_p") # --- Types with __dlpack__ protocol (e.g. torch.Tensor) → Tensor --- if isinstance(annotation, type) and hasattr(annotation, "__dlpack__"): return TypeSchema("Tensor") raise TypeError( f"Cannot convert {annotation!r} to TypeSchema" ) def check_value(self, value: object) -> None: """Validate that *value* is compatible with this type schema. Parameters ---------- value : object The Python value to check. Raises ------ TypeError If the value is not compatible with the schema, with a human-readable error message describing the mismatch. """ try: _type_convert_impl(self._converter, value) except RecursionError: raise TypeError( f"type check failed for {self!r}: " f"infinite __tvm_ffi_value__ cycle detected" ) from None except _ConvertError as err: raise TypeError(f"type check failed for {self!r}: {err.message}") from None def convert(self, value: object) -> "CAny": """Convert *value* according to this type schema, returning a :class:`CAny`. Applies the same implicit conversions as the C++ FFI ``TypeTraits::TryCastFromAnyView`` rules. The result is always a :class:`CAny` instance that owns the converted value. Use ``_to_py_class_value(result)`` to recover the Python object. Parameters ---------- value : object The Python value to convert. Returns ------- CAny The converted value wrapped in a CAny. Raises ------ TypeError If the value cannot be converted to this schema's type. """ try: return _type_convert_impl(self._converter, value) except RecursionError: raise TypeError( f"type conversion failed for {self!r}: " f"infinite __tvm_ffi_value__ cycle detected" ) from None except _ConvertError as err: raise TypeError(f"type conversion failed for {self!r}: {err.message}") from None def repr(self, ty_map: "Optional[Callable[[str], str]]" = None) -> str: """Render a human-readable representation of this schema. Parameters ---------- ty_map : Callable[[str], str], optional A mapping function applied to the schema origin name before rendering (e.g. map ``"Array" -> "Array"`` and ``"Map" -> "Map"``). If ``None``, the raw origin is used. Returns ------- str A readable string using Python typing syntax. Formats include: - Unions as ``"T1 | T2"`` - Optional as ``"T | None"`` - Callables as ``"Callable[[arg1, ...], ret]"`` - Containers as ``"origin[arg1, ...]"`` Examples -------- .. code-block:: python # From JSON emitted by the runtime s = TypeSchema.from_json_str('{"type":"Optional","args":[{"type":"int"}]}') assert s.repr() == "int | None" # Callable where the first arg is return type, remaining are parameters s = TypeSchema("Callable", (TypeSchema("int"), TypeSchema("str"))) assert s.repr() == "Callable[[str], int]" # Container types from C++ FFI schemas s = TypeSchema.from_json_str('{"type":"ffi.Map","args":[{"type":"str"},{"type":"int"}]}') assert s.repr() == "Map[str, int]" s = TypeSchema.from_json_str('{"type":"ffi.Array","args":[{"type":"int"}]}') assert s.repr() == "Array[int]" """ if ty_map is None: origin = self.origin else: origin = ty_map(self.origin) schema_args = self.args args = [i.repr(ty_map) for i in (() if schema_args is None else schema_args)] if origin == "Union": return " | ".join(args) elif origin == "Optional": return args[0] + " | None" elif origin == "Callable": if not args: return "Callable[..., Any]" else: ret = args[0] args = ", ".join(args[1:]) return f"Callable[[{args}], {ret}]" elif origin == "tuple" and schema_args == (): return "tuple[()]" elif not args: return origin else: args = ", ".join(args) return f"{origin}[{args}]" def to_json(self) -> dict[str, Any]: """Convert a TypeSchema to a JSON-compatible dict.""" if self.args is not None and (self.args or self.origin == "tuple"): return { "type": self.origin, "args": [a.to_json() for a in self.args], } return {"type": self.origin} def _annotation_union(args): """Convert Union type args to a TypeSchema (Optional or Union).""" non_none = tuple(a for a in args if a is not type(None)) has_none = len(non_none) < len(args) converted = tuple(TypeSchema.from_annotation(a) for a in non_none) if has_none: if len(non_none) == 1: return TypeSchema("Optional", converted) return TypeSchema("Optional", (TypeSchema("Union", converted),)) return TypeSchema("Union", converted) def _annotation_cobject(cls, targs): """Handle a CObject subclass (bare or parameterised) in from_annotation.""" info = TYPE_CLS_TO_INFO.get(cls) if info is None: raise TypeError( f"CObject subclass {cls!r} is not registered " f"in TYPE_CLS_TO_INFO; use @register_object to register it" ) # Prefer canonical short origin from _TYPE_INDEX_TO_ORIGIN (e.g. "Array") # over the registered type_key (e.g. "ffi.Array") when available. origin = _TYPE_INDEX_TO_ORIGIN.get(info.type_index, info.type_key) n = len(targs) if n > 0: if origin in ("Array", "List"): if n != 1: raise TypeError( f"{origin} requires 1 type argument, got {n}" ) elif origin in ("Map", "Dict"): if n != 2: raise TypeError( f"{origin} requires 2 type arguments, got {n}" ) arg_schemas = tuple(TypeSchema.from_annotation(a) for a in targs) return TypeSchema(origin, arg_schemas, origin_type_index=info.type_index) return TypeSchema(origin, origin_type_index=info.type_index) class FFIProperty(property): """Property descriptor for FFI-backed fields. When *frozen* is True the public setter (``fset``) is suppressed so that normal attribute assignment raises ``AttributeError``. The real setter is stashed in :attr:`_fset` and exposed via the :meth:`set` escape-hatch. """ def __init__(self, fget, fset, frozen, fdel=None, doc=None): super().__init__(fget, None if frozen else fset, fdel, doc) self._fset = fset def set(self, obj, value): """Force-set the field value, bypassing the frozen guard.""" self._fset(obj, value) @dataclasses.dataclass(eq=False) class TypeField: """Description of a single reflected field on an FFI-backed type.""" name: str doc: Optional[str] size: int offset: int frozen: bool metadata: dict[str, Any] getter: FieldGetter setter: FieldSetter ty: Optional[TypeSchema] = None c_init: bool = True c_kw_only: bool = False c_has_default: bool = False # ``c_default`` / ``c_default_factory`` are populated from the C++ # reflection layer (``TVMFFIFieldInfo.default_value_or_factory``) for # ``@c_class`` types, and from the ``Field`` descriptors for # ``@py_class`` types. Both default to :data:`MISSING` when no # default / factory has been registered. c_default: Any = dataclasses.field(default_factory=lambda: MISSING) c_default_factory: Any = dataclasses.field(default_factory=lambda: MISSING) # Presentation / structural flags decoded from the reflection layer. # ``c_repr`` / ``c_compare`` / ``c_hash`` default to True and flip off # when ``refl::repr(false)`` / ``refl::compare(false)`` / ``refl::hash(false)`` # (or the corresponding ``@py_class`` ``field(...)`` kwargs) are set. # ``c_structural_eq`` is ``None`` / ``"ignore"`` / ``"def"`` matching the # ``Field.structural_eq`` vocabulary. c_repr: bool = True c_compare: bool = True c_hash: bool = True c_structural_eq: Optional[str] = None dataclass_field: Any = None def __post_init__(self): assert self.setter is not None assert self.getter is not None def as_property(self, object cls): """Create an :class:`FFIProperty` descriptor for this field on ``cls``.""" cdef str name = self.name cdef FieldGetter fget = self.getter cdef FieldSetter fset = self.setter cdef object ret fget.__name__ = fset.__name__ = name fget.__module__ = fset.__module__ = cls.__module__ fget.__qualname__ = fset.__qualname__ = f"{cls.__qualname__}.{name}" ret = FFIProperty( fget=fget, fset=fset, frozen=self.frozen, ) if self.doc: ret.__doc__ = self.doc fget.__doc__ = self.doc fset.__doc__ = self.doc return ret @dataclasses.dataclass(eq=False) class TypeMethod: """Description of a single reflected method on an FFI-backed type.""" name: str doc: Optional[str] func: object metadata: dict[str, Any] is_static: bool def __post_init__(self): assert callable(self.func) def as_callable(self, object cls): """Create a Python method attribute for this method on ``cls``.""" cdef str name = self.name cdef object func = self.func if not self.is_static: func = _member_method_wrapper(func) func.__module__ = cls.__module__ func.__name__ = name func.__qualname__ = f"{cls.__qualname__}.{name}" if self.doc: func.__doc__ = self.doc if self.is_static: func = staticmethod(func) return func @dataclasses.dataclass(eq=False) class TypeInfo: """Aggregated type information required to build a proxy class.""" type_cls: Optional[type] type_index: int type_key: str type_ancestors: list[int] fields: Optional[list[TypeField]] methods: list[TypeMethod] parent_type_info: Optional[TypeInfo] def __post_init__(self): cdef int parent_type_index cdef str parent_type_key # Assert no duplicate field names within this type's own fields. if self.fields is not None: seen = set() for f in self.fields: assert f.name not in seen, ( f"duplicate field name {f.name!r} in TypeInfo for {self.type_key!r}; " f"TypeInfo.fields must only contain the type's own fields" ) seen.add(f.name) if not self.type_ancestors: return parent_type_index = self.type_ancestors[-1] parent_type_key = _type_index_to_key(parent_type_index) # ensure parent is registered self.parent_type_info = _lookup_or_register_type_info_from_type_key(parent_type_key) # Warn if own fields shadow any ancestor field. if self.fields and self.parent_type_info is not None: parent_names = set() ti = self.parent_type_info while ti is not None: if ti.fields: for f in ti.fields: parent_names.add(f.name) ti = ti.parent_type_info for f in self.fields: if f.name in parent_names: import warnings warnings.warn( f"Field {f.name!r} in {self.type_key!r} duplicates " f"an ancestor field. Child types should not " f"re-register inherited fields.", stacklevel=2, ) @cached_property def total_size(self) -> int: """Total object size in bytes (header + all fields). For native C++ types with metadata, returns metadata.total_size. For Python-defined types, computes from field layout. """ cdef const TVMFFITypeInfo* c_info = TVMFFIGetTypeInfo(self.type_index) if c_info != NULL and c_info.metadata != NULL: return c_info.metadata.total_size if self.parent_type_info is None: raise ValueError(f"Cannot find parent type of {type(self)}") cdef int64_t end = self.parent_type_info.total_size assert end >= sizeof(TVMFFIObject) for f in self.fields: end = max(end, f.offset + f.size) return (end + 7) & ~7 # align to 8 bytes def _register_fields(self, fields, structure_kind=None): """Register Field descriptors and set up __ffi_new__/__ffi_init__. Delegates to the module-level _register_fields function, stores the resulting list[TypeField] on self.fields, then reads back methods registered by C++ via _read_back_methods. Can only be called once (fields must be None beforehand). Parameters ---------- fields : list[Field] The Field descriptors to register. structure_kind : int | None The structural equality/hashing kind (``TVMFFISEqHashKind`` integer). ``None`` or ``0`` means unsupported (no metadata registered). """ assert self.fields is None, ( f"_register_fields already called for {self.type_key!r}" ) self.fields = _register_fields(self, fields, structure_kind) self._read_back_methods() def _register_py_methods(self, py_methods=None, type_attr_names=frozenset()): """Register user-defined dunder hooks and re-read the method table. Each entry whose name is in *type_attr_names* is registered as a TypeAttrColumn entry (for C++ dispatch); the value need not be callable (e.g. ``__ffi_ir_traits__``). All other entries are registered as TypeMethod (for reflection introspection). Regardless, the full method list is always re-read from the C type table so that system-generated methods (``__ffi_init__``, ``__ffi_shallow_copy__``) are picked up. Parameters ---------- py_methods : list[tuple[str, Any, bool]] | None Each entry is ``(name, value, is_static)``. type_attr_names : frozenset[str] Names to register as TypeAttrColumn instead of TypeMethod. """ if py_methods: _register_py_methods(self.type_index, py_methods, type_attr_names) self._read_back_methods() def _read_back_methods(self): """Read methods from the C type table into self.methods. Called after C++ registers __ffi_init__, __ffi_shallow_copy__, etc. """ cdef const TVMFFITypeInfo* c_info = TVMFFIGetTypeInfo(self.type_index) cdef const TVMFFIMethodInfo* mi self.methods = [] for i in range(c_info.num_methods): mi = &(c_info.methods[i]) self.methods.append(TypeMethod( name=bytearray_to_str(&mi.name), doc=bytearray_to_str(&mi.doc) if mi.doc.size != 0 else None, func=_get_method_from_method_info(mi), is_static=(mi.flags & kTVMFFIFieldFlagBitMaskIsStaticMethod) != 0, metadata=json.loads(bytearray_to_str(&mi.metadata)) if mi.metadata.size != 0 else {}, )) # --------------------------------------------------------------------------- # Python-defined type field registration helpers # --------------------------------------------------------------------------- # Native layout for each TypeSchema origin: (size, alignment, field_static_type_index) _ORIGIN_NATIVE_LAYOUT = { "int": (8, 8, kTVMFFIInt), "float": (8, 8, kTVMFFIFloat), "bool": (1, 1, kTVMFFIBool), "ctypes.c_void_p": (8, 8, kTVMFFIOpaquePtr), "dtype": (4, 2, kTVMFFIDataType), "Device": (8, 4, kTVMFFIDevice), "Any": (16, 8, -1), # kTVMFFIAny = -1 # str/bytes can be SmallStr/SmallBytes (inline, not ObjectRef), # so store as Any (16 bytes) to handle both inline and heap variants. "str": (16, 8, -1), "bytes": (16, 8, -1), # Optional/Union can hold any type including inline scalars "Optional": (16, 8, -1), "Union": (16, 8, -1), } cdef _register_one_field( int32_t type_index, object py_field, int64_t offset, int64_t size, int64_t alignment, int32_t field_type_index, TVMFFIFieldGetter getter, CObject setter_fn, ): """Build a TVMFFIFieldInfo and register it for the given type.""" cdef TVMFFIFieldInfo info cdef int c_api_ret_code # --- name --- name_bytes = c_str(py_field.name) cdef ByteArrayArg name_arg = ByteArrayArg(name_bytes) info.name = name_arg.cdata # --- doc --- cdef ByteArrayArg doc_arg if py_field.doc is not None: doc_bytes = c_str(py_field.doc) doc_arg = ByteArrayArg(doc_bytes) info.doc = doc_arg.cdata else: info.doc.data = NULL info.doc.size = 0 # --- metadata (JSON with type_schema) --- metadata_str = json.dumps({"type_schema": py_field._ty_schema.to_json()}) metadata_bytes = c_str(metadata_str) cdef ByteArrayArg metadata_arg = ByteArrayArg(metadata_bytes) info.metadata = metadata_arg.cdata # --- flags --- cdef int64_t flags = kTVMFFIFieldFlagBitMaskWritable | kTVMFFIFieldFlagBitSetterIsFunctionObj if py_field.default is not MISSING or py_field.default_factory is not MISSING: flags |= kTVMFFIFieldFlagBitMaskHasDefault if py_field.default_factory is not MISSING: flags |= kTVMFFIFieldFlagBitMaskDefaultFromFactory if not py_field.init: flags |= kTVMFFIFieldFlagBitMaskInitOff if not py_field.repr: flags |= kTVMFFIFieldFlagBitMaskReprOff if not py_field.hash: flags |= kTVMFFIFieldFlagBitMaskHashOff if not py_field.compare: flags |= kTVMFFIFieldFlagBitMaskCompareOff if py_field.kw_only: flags |= kTVMFFIFieldFlagBitMaskKwOnly # Structural equality/hashing field annotations cdef object field_structure = getattr(py_field, "structural_eq", None) if field_structure == "ignore": flags |= kTVMFFIFieldFlagBitMaskSEqHashIgnore elif field_structure == "def" or field_structure == "def-recursive": # ``"def"`` is the legacy short form, kept as a Python-side synonym for # ``"def-recursive"`` since the C-level rename of the underlying flag # (``kTVMFFIFieldFlagBitMaskSEqHashDef`` -> ``...SEqHashDefRecursive``) # only changed the constant name, not the recursive semantics. flags |= kTVMFFIFieldFlagBitMaskSEqHashDefRecursive elif field_structure == "def-non-recursive": flags |= kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive info.flags = flags # --- native layout --- info.size = size info.alignment = alignment info.offset = offset # --- getter / setter --- info.getter = getter info.setter = setter_fn.chandle # --- default value --- cdef TVMFFIAny default_any default_any.type_index = kTVMFFINone default_any.v_int64 = 0 # Determine which Python object (if any) to store as the default. # No memory leak: TVMFFIAny is a POD struct; TVMFFITypeRegisterField # copies the bytes into the type table, which owns the reference. cdef object default_obj = MISSING if py_field.default is not MISSING: default_obj = py_field.default elif py_field.default_factory is not MISSING: default_obj = py_field.default_factory if default_obj is not MISSING: TVMFFIPyPyObjectToFFIAny( default_obj, &default_any, &c_api_ret_code ) CHECK_CALL(c_api_ret_code) info.default_value_or_factory = default_any # --- field_static_type_index --- info.field_static_type_index = field_type_index CHECK_CALL(TVMFFITypeRegisterField(type_index, &info)) cdef int _f_type_convert(void* type_converter, const TVMFFIAny* value, TVMFFIAny* result) noexcept with gil: """C callback for type conversion, called from C++ MakeFieldSetter. Parameters ---------- type_converter : void* A PyObject* pointing to a _TypeConverter instance (borrowed reference). value : const TVMFFIAny* The packed value to convert (borrowed from the caller). result : TVMFFIAny* Output: the converted value (caller takes ownership). Returns 0 on success, -1 on error (error stored in TLS via set_last_ffi_error). """ cdef TVMFFIAny temp cdef _TypeConverter conv cdef CAny cany try: # Unpack the packed AnyView to a Python object. # We must IncRef if it's an object, because make_ret takes ownership. temp = value[0] if temp.type_index >= kTVMFFIStaticObjectBegin: if temp.v_obj != NULL: TVMFFIObjectIncRef(temp.v_obj) py_value = make_ret(temp) # Dispatch directly through the C-level converter conv = <_TypeConverter>type_converter cany = _type_convert_impl(conv, py_value) # Transfer ownership from CAny to result (zero cany to prevent double-free) result[0] = cany.cdata cany.cdata.type_index = kTVMFFINone cany.cdata.v_int64 = 0 return 0 except Exception as err: set_last_ffi_error(err) return -1 def _register_fields(type_info, fields, structure_kind=None): """Register Field descriptors for a Python-defined type and set up __ffi_new__/__ffi_init__. For each Field: 1. Computes native layout (size, alignment, offset) 2. Obtains a C getter function pointer 3. Creates a FunctionObj setter with type conversion 4. Registers via TVMFFITypeRegisterField After all fields, registers __ffi_new__ (object allocator), __ffi_init__ (auto-generated constructor), and optionally type metadata (structural_eq_hash_kind). Parameters ---------- type_info : TypeInfo The TypeInfo of the type being defined. fields : list[Field] The Field descriptors to register. structure_kind : int | None The structural equality/hashing kind (``TVMFFISEqHashKind`` integer). ``None`` or ``0`` means unsupported (no metadata registered). Returns ------- list[TypeField] The registered field descriptors. """ cdef int32_t type_index = type_info.type_index # Start field offsets AFTER all parent fields (not at fixed offset 24). # This is critical for inheritance: child fields must not overlap parent memory. cdef int64_t current_offset = type_info.parent_type_info.total_size cdef int64_t size, alignment cdef int32_t field_type_index cdef TVMFFIFieldGetter getter cdef FieldGetter fgetter cdef FieldSetter fsetter cdef list type_fields = [] for py_field in fields: # 1. Get layout layout = _ORIGIN_NATIVE_LAYOUT.get(py_field._ty_schema.origin, (8, 8, kTVMFFIObject)) size = layout[0] alignment = layout[1] field_type_index = layout[2] # 2. Compute offset (align up) current_offset = (current_offset + alignment - 1) & ~(alignment - 1) field_offset = current_offset current_offset += size # 3. Get getter (C function pointer) and setter (FunctionObj). # Pointers are transported as int64_t through the FFI boundary. getter = _MAKE_FILED_GETTER(field_type_index) setter_fn = _MAKE_FIELD_SETTER( field_type_index, py_field._ty_schema._converter, &_f_type_convert, ) # 4. Register field in the C type table _register_one_field( type_index, py_field, field_offset, size, alignment, field_type_index, getter, setter_fn, ) # 5. Build the Python-side TypeField descriptor fgetter = FieldGetter.__new__(FieldGetter) fgetter.getter = getter fgetter.offset = field_offset fsetter = FieldSetter.__new__(FieldSetter) fsetter.setter = setter_fn.chandle fsetter.offset = field_offset fsetter.flags = (kTVMFFIFieldFlagBitMaskWritable | kTVMFFIFieldFlagBitSetterIsFunctionObj) type_fields.append( TypeField( name=py_field.name, doc=py_field.doc, size=size, offset=field_offset, frozen=py_field.frozen, metadata={"type_schema": py_field._ty_schema.to_json()}, getter=fgetter, setter=fsetter, ty=py_field._ty_schema, c_init=py_field.init, c_kw_only=py_field.kw_only, c_has_default=(py_field.default is not MISSING or py_field.default_factory is not MISSING), c_default=py_field.default, c_default_factory=py_field.default_factory, c_repr=py_field.repr, c_compare=py_field.compare, c_hash=bool(py_field.hash), c_structural_eq=py_field.structural_eq, ) ) # Align total size to 8 bytes cdef int64_t total_size = (current_offset + 7) & ~7 if total_size < sizeof(TVMFFIObject): total_size = sizeof(TVMFFIObject) # 7. Register __ffi_new__, __ffi_shallow_copy__, __ffi_init__ TypeAttrColumns _PYCLS_REGISTER(type_index, total_size) # 8. Register type metadata (structural_eq_hash_kind) if specified. if structure_kind is not None and structure_kind != 0: _register_type_metadata(type_index, total_size, structure_kind) return type_fields cdef _register_type_metadata(int32_t type_index, int32_t total_size, int structure_kind): """Register TVMFFITypeMetadata for the given type with structural eq/hash kind.""" cdef TVMFFITypeMetadata metadata metadata.doc.data = NULL metadata.doc.size = 0 metadata.creator = NULL metadata.total_size = total_size metadata.structural_eq_hash_kind = structure_kind CHECK_CALL(TVMFFITypeRegisterMetadata(type_index, &metadata)) cdef _register_py_methods(int32_t type_index, list py_methods, frozenset type_attr_names): """Register user-defined methods and type attrs as TypeAttrColumn or TypeMethod. For each entry in *py_methods*: 1. Convert the Python object to a ``TVMFFIAny``. 2. If the name is in *type_attr_names*, register as TypeAttrColumn (for C++ dispatch via ``TypeAttrColumn``). The value need not be callable (e.g. ``__ffi_ir_traits__`` is an Object instance). 3. Otherwise, register as TypeMethod (for reflection introspection via ``TypeInfo.methods``). Parameters ---------- type_index : int The runtime type index of the type. py_methods : list[tuple[str, Any, bool]] Each entry is ``(name, value, is_static)``. type_attr_names : frozenset[str] Names to register as TypeAttrColumn instead of TypeMethod. """ cdef TVMFFIMethodInfo method_info cdef TVMFFIAny func_any cdef TVMFFIAny sentinel_any cdef int c_api_ret_code cdef ByteArrayArg name_arg sentinel_any.type_index = kTVMFFINone sentinel_any.v_int64 = 0 for name, func, is_static in py_methods: func_any.type_index = kTVMFFINone func_any.v_int64 = 0 try: name_bytes = c_str(name) name_arg = ByteArrayArg(name_bytes) # Convert Python object -> TVMFFIAny TVMFFIPyPyObjectToFFIAny( func, &func_any, &c_api_ret_code, ) CHECK_CALL(c_api_ret_code) if name in type_attr_names: # Register as TypeAttrColumn (for C++ dispatch) CHECK_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_arg.cdata, &sentinel_any)) CHECK_CALL(TVMFFITypeRegisterAttr(type_index, &name_arg.cdata, &func_any)) else: # Register as TypeMethod (for reflection introspection) method_info.name = name_arg.cdata method_info.doc.data = NULL method_info.doc.size = 0 method_info.flags = kTVMFFIFieldFlagBitMaskIsStaticMethod if is_static else 0 method_info.method = func_any method_info.metadata.data = NULL method_info.metadata.size = 0 CHECK_CALL(TVMFFITypeRegisterMethod(type_index, &method_info)) finally: if func_any.type_index >= kTVMFFIStaticObjectBegin and func_any.v_obj != NULL: TVMFFIObjectDecRef(func_any.v_obj) def _member_method_wrapper(method_func: Callable[..., Any]) -> Callable[..., Any]: def wrapper(self: Any, *args: Any) -> Any: return method_func(self, *args) return wrapper tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/000077500000000000000000000000001521067262500204135ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/__init__.py000066400000000000000000000026501521067262500225270ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """FFI dataclass decorators: ``c_class`` for C++-backed types, ``py_class`` for Python-defined types.""" from tvm_ffi.core import MISSING, Object from .c_class import c_class from .common import asdict, astuple, fields, is_dataclass, replace from .enum import Enum, EnumAttrMap, IntEnum, StrEnum, auto, entry from .field import KW_ONLY, Field, field from .py_class import py_class __all__ = [ "KW_ONLY", "MISSING", "Enum", "EnumAttrMap", "Field", "IntEnum", "Object", "StrEnum", "asdict", "astuple", "auto", "c_class", "entry", "field", "fields", "is_dataclass", "py_class", "replace", ] tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/c_class.py000066400000000000000000000137361521067262500224060ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """The ``c_class`` decorator: register_object + structural dunders.""" from __future__ import annotations import typing from collections.abc import Callable from typing import Any, TypeVar from typing_extensions import dataclass_transform from .field import Field _T = TypeVar("_T", bound=type) def _attach_field_objects(cls: type, type_info: Any) -> None: """Populate ``TypeField.dataclass_field`` for every own reflected field. ``@c_class`` fields originate from C++ reflection, so there is no user-supplied :class:`Field`. We synthesize one per ``TypeField`` and stash it on ``TypeField.dataclass_field`` so :func:`~tvm_ffi.dataclasses.fields` can return it. """ try: hints = typing.get_type_hints(cls) except Exception: hints = {} for tf in type_info.fields: f = Field( name=tf.name, _ty_schema=tf.ty, default=tf.c_default, default_factory=tf.c_default_factory, frozen=tf.frozen, init=tf.c_init, repr=tf.c_repr, hash=tf.c_hash, compare=tf.c_compare, kw_only=tf.c_kw_only, structural_eq=tf.c_structural_eq, doc=tf.doc, ) f.type = hints.get(tf.name) tf.dataclass_field = f @dataclass_transform(eq_default=False, order_default=False) def c_class( type_key: str, *, init: bool = True, repr: bool = True, eq: bool = False, order: bool = False, unsafe_hash: bool = False, match_args: bool = True, ) -> Callable[[_T], _T]: """Register a C++ FFI class and install structural dunder methods. Combines :func:`~tvm_ffi.register_object` with structural comparison, hashing, and ordering derived from the C++ reflection metadata. User-defined dunders in the class body are never overwritten. Parameters ---------- type_key The reflection key that identifies the C++ type in the FFI registry. Must match a key already registered on the C++ side via ``TVM_FFI_DECLARE_OBJECT_INFO``. init If True (default), install ``__init__`` from C++ reflection metadata. The generated ``__init__`` respects ``Init()``, ``KwOnly()``, and ``Default()`` traits declared on each C++ field. If the class body already defines ``__init__``, it is kept. repr If True (default), install ``__repr__`` using :func:`~tvm_ffi.core.object_repr`, which formats the object via the C++ ``ReprPrint`` visitor. Skipped if the class body already defines ``__repr__``. eq If True, install ``__eq__`` and ``__ne__`` using the C++ recursive structural comparison (``RecursiveEq``). Returns ``NotImplemented`` for unrelated types. Defaults to False. order If True, install ``__lt__``, ``__le__``, ``__gt__``, ``__ge__`` using the C++ recursive comparators. Returns ``NotImplemented`` for unrelated types. Defaults to False. unsafe_hash If True, install ``__hash__`` using ``RecursiveHash``. Called *unsafe* because mutable fields contribute to the hash, so mutating an object while it is in a set or dict key will break invariants. Defaults to False. match_args If True (default), set ``__match_args__`` to a tuple of the positional ``__init__`` field names (``init=True`` and not ``kw_only``), enabling ``match`` statements. Ignored when the class body already defines ``__match_args__``. Returns ------- Callable[[type], type] A class decorator. Examples -------- Basic usage with default settings (``init`` and ``repr`` enabled): .. code-block:: python @c_class("my.Point") class Point(Object): x: float y: float Enable structural equality, hashing, and ordering: .. code-block:: python @c_class("my.Point", eq=True, unsafe_hash=True, order=True) class Point(Object): x: float y: float See Also -------- :func:`tvm_ffi.register_object` Lower-level decorator that only registers the type without installing structural dunders. """ from .._dunder import _install_dataclass_dunders # noqa: PLC0415 from ..registry import ( # noqa: PLC0415 _warn_missing_field_annotations, register_object, ) def decorator(cls: _T) -> _T: cls = register_object(type_key, init=False)(cls) type_info = getattr(cls, "__tvm_ffi_type_info__", None) assert type_info is not None _warn_missing_field_annotations(cls, type_info, stacklevel=2) _attach_field_objects(cls, type_info) _install_dataclass_dunders( cls, init=init, repr=repr, eq=eq, order=order, unsafe_hash=unsafe_hash, match_args=match_args, ) # Marker: distinguishes @c_class / @py_class types from FFI containers # (Array, List, Map, Dict) that also have __tvm_ffi_type_info__ but are # not dataclasses. Used by is_dataclass() in common.py. setattr(cls, "__tvm_ffi_is_dataclass__", True) return cls return decorator tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/common.py000066400000000000000000000227731521067262500222700ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """``dataclasses``-style helpers unified over stdlib, ``@c_class``, and ``@py_class``.""" from __future__ import annotations import copy from collections.abc import Callable from typing import Any from ..container import Array, Dict, List, Map from .field import Field __all__ = ["asdict", "astuple", "fields", "is_dataclass", "replace"] # Exact-type fast path for atomic (immutable, non-recursive) values. Mirrors # :data:`dataclasses._ATOMIC_TYPES` from the standard library. _ATOMIC_TYPES: frozenset[type] = frozenset( { bool, bytes, complex, float, int, str, type(None), type(Ellipsis), type(NotImplemented), } ) def is_dataclass(obj: Any) -> bool: """Return True if ``obj`` is a ``@c_class`` / ``@py_class`` type or instance. Returns False for FFI container types (:class:`~tvm_ffi.Array`, :class:`~tvm_ffi.List`, :class:`~tvm_ffi.Map`, :class:`~tvm_ffi.Dict`) even though they also carry ``__tvm_ffi_type_info__``; those are reflected through :func:`~tvm_ffi.register_object` directly, not through the dataclass decorators. """ cls = obj if isinstance(obj, type) else type(obj) return getattr(cls, "__tvm_ffi_is_dataclass__", False) is True def fields(obj_or_cls: Any) -> tuple[Field, ...]: """Return the :class:`~tvm_ffi.dataclasses.Field` descriptors for a type. Accepts a ``@c_class`` / ``@py_class`` type or instance and walks the parent chain so inherited fields appear parent-first, matching the order of the auto-generated ``__init__``. Raises ------ TypeError If ``obj_or_cls`` is not a ``@c_class`` / ``@py_class`` type or instance. """ if not is_dataclass(obj_or_cls): raise TypeError( f"fields() argument must be a c_class or py_class type or instance, " f"got {type(obj_or_cls).__name__}" ) cls = obj_or_cls if isinstance(obj_or_cls, type) else type(obj_or_cls) ti = getattr(cls, "__tvm_ffi_type_info__", None) chain = [] while ti is not None: chain.append(ti) ti = ti.parent_type_info out: list[Field] = [] for ti in reversed(chain): for tf in ti.fields or (): if tf.dataclass_field is not None: out.append(tf.dataclass_field) return tuple(out) def replace(obj: Any, /, **changes: Any) -> Any: """Return a copy of ``obj`` with selected fields replaced. Drop-in for :func:`dataclasses.replace` for FFI-backed instances: the call is forwarded to ``obj.__replace__`` (installed by the decorator), which uses the ``FFIProperty.set()`` escape hatch so frozen fields are still replaceable. """ return obj.__replace__(**changes) def _is_ffi_dataclass_instance(obj: Any) -> bool: """Return True when *obj* is a ``@c_class`` / ``@py_class`` **instance** (not a type).""" if isinstance(obj, type): return False return is_dataclass(obj) def _asdict_inner( # noqa: PLR0911, PLR0912 obj: Any, dict_factory: Callable[..., Any] ) -> Any: obj_type = type(obj) if obj_type in _ATOMIC_TYPES: return obj # FFI containers are treated as their stdlib analogues so the result is # plain Python data — handy for JSON serialisation, the main use case. if isinstance(obj, (Array, List)): return [_asdict_inner(v, dict_factory) for v in obj] if isinstance(obj, (Map, Dict)): return dict_factory( [ (_asdict_inner(k, dict_factory), _asdict_inner(v, dict_factory)) for k, v in obj.items() ] ) if _is_ffi_dataclass_instance(obj): fs = fields(obj) if dict_factory is dict: return {f.name: _asdict_inner(getattr(obj, f.name), dict) for f in fs} # ty: ignore[invalid-argument-type] return dict_factory( [(f.name, _asdict_inner(getattr(obj, f.name), dict_factory)) for f in fs] # ty: ignore[invalid-argument-type] ) if obj_type is list: return [_asdict_inner(v, dict_factory) for v in obj] if obj_type is dict: return { _asdict_inner(k, dict_factory): _asdict_inner(v, dict_factory) for k, v in obj.items() } if obj_type is tuple: return tuple(_asdict_inner(v, dict_factory) for v in obj) if issubclass(obj_type, tuple): if hasattr(obj, "_fields"): # namedtuple return obj_type(*[_asdict_inner(v, dict_factory) for v in obj]) return obj_type(_asdict_inner(v, dict_factory) for v in obj) if issubclass(obj_type, dict): if hasattr(obj_type, "default_factory"): result = obj_type(obj.default_factory) for k, v in obj.items(): result[_asdict_inner(k, dict_factory)] = _asdict_inner(v, dict_factory) return result return obj_type( (_asdict_inner(k, dict_factory), _asdict_inner(v, dict_factory)) for k, v in obj.items() ) if issubclass(obj_type, list): return obj_type(_asdict_inner(v, dict_factory) for v in obj) return copy.deepcopy(obj) def _astuple_inner(obj: Any, tuple_factory: Callable[..., Any]) -> Any: # noqa: PLR0911 obj_type = type(obj) if obj_type in _ATOMIC_TYPES: return obj if isinstance(obj, (Array, List)): return [_astuple_inner(v, tuple_factory) for v in obj] if isinstance(obj, (Map, Dict)): return { _astuple_inner(k, tuple_factory): _astuple_inner(v, tuple_factory) for k, v in obj.items() } if _is_ffi_dataclass_instance(obj): return tuple_factory( [_astuple_inner(getattr(obj, f.name), tuple_factory) for f in fields(obj)] # ty: ignore[invalid-argument-type] ) if isinstance(obj, tuple) and hasattr(obj, "_fields"): # namedtuple return obj_type(*[_astuple_inner(v, tuple_factory) for v in obj]) if isinstance(obj, (list, tuple)): return obj_type(_astuple_inner(v, tuple_factory) for v in obj) if isinstance(obj, dict): if hasattr(obj_type, "default_factory"): result = obj_type(obj.default_factory) for k, v in obj.items(): result[_astuple_inner(k, tuple_factory)] = _astuple_inner(v, tuple_factory) return result return obj_type( (_astuple_inner(k, tuple_factory), _astuple_inner(v, tuple_factory)) for k, v in obj.items() ) return copy.deepcopy(obj) def asdict(obj: Any, *, dict_factory: Callable[..., Any] = dict) -> Any: r"""Return the fields of a ``@c_class`` / ``@py_class`` instance as a new dict. Mirrors :func:`dataclasses.asdict` from the standard library. The function recurses into nested FFI dataclass instances, FFI containers (:class:`~tvm_ffi.Array`, :class:`~tvm_ffi.List`, :class:`~tvm_ffi.Map`, :class:`~tvm_ffi.Dict`), and the built-in ``list`` / ``tuple`` / ``dict``. FFI sequence containers are converted to Python ``list``\ s and FFI mapping containers to Python ``dict``\ s so the result is plain Python data, e.g. for JSON serialisation. Any other value is copied with :func:`copy.deepcopy`. Parameters ---------- obj A ``@c_class`` / ``@py_class`` instance. Passing a type raises :class:`TypeError`. dict_factory Callable used to construct the outer mapping and any nested mapping recursed from an FFI dataclass. Defaults to :class:`dict`. Raises ------ TypeError If ``obj`` is not a ``@c_class`` / ``@py_class`` instance. """ if not _is_ffi_dataclass_instance(obj): raise TypeError("asdict() should be called on c_class / py_class instances") return _asdict_inner(obj, dict_factory) def astuple(obj: Any, *, tuple_factory: Callable[..., Any] = tuple) -> Any: """Return the fields of a ``@c_class`` / ``@py_class`` instance as a new tuple. Mirrors :func:`dataclasses.astuple` from the standard library. The function recurses into nested FFI dataclass instances, FFI containers (:class:`~tvm_ffi.Array`, :class:`~tvm_ffi.List`, :class:`~tvm_ffi.Map`, :class:`~tvm_ffi.Dict`), and the built-in ``list`` / ``tuple`` / ``dict``. Any other value is copied with :func:`copy.deepcopy`. Parameters ---------- obj A ``@c_class`` / ``@py_class`` instance. Passing a type raises :class:`TypeError`. tuple_factory Callable used to construct the outer tuple and any nested tuple recursed from an FFI dataclass. Defaults to :class:`tuple`. Raises ------ TypeError If ``obj`` is not a ``@c_class`` / ``@py_class`` instance. """ if not _is_ffi_dataclass_instance(obj): raise TypeError("astuple() should be called on c_class / py_class instances") return _astuple_inner(obj, tuple_factory) tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/enum.py000066400000000000000000001240201521067262500217300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Cross-language enum types: named, frozen, ordinal-indexed singletons. An ``Enum`` subclass has one of two usage modes, distinguished by whether its ``type_key`` is already registered in the FFI type system: * **Closed Python enum** — fresh ``type_key``, variants declared once in the class body. Behavior matches ``enum.Enum``. * **Cross-language registry** — ``type_key`` also registered in C++ (or another Python module). Python and C++ both contribute variants to the same per-class registry, and consumers attach *extensible attributes* to variants from any module at any time. See :class:`Enum` for declaration forms and :meth:`Enum.def_attr` for extensible attributes. Storage layout (mirrors ``include/tvm/ffi/enum.h``): * ``__ffi_enum_entries__`` — ``Dict[str, Enum]``, name → variant. * ``__ffi_enum_attrs__`` — ``Dict[str, List[Any]]``, extensible-attr name → column indexed by each variant's ordinal. """ from __future__ import annotations import sys import typing from collections.abc import Callable, Iterator from typing import Any, ClassVar from typing_extensions import dataclass_transform from .. import core from ..container import Dict, List from ..core import Object from .c_class import c_class from .field import Field, field from .py_class import py_class __all__ = [ "ENUM_ATTRS_ATTR", "ENUM_ENTRIES_ATTR", "ENUM_VALUE_ENTRIES_ATTR", "Enum", "EnumAttrMap", "IntEnum", "StrEnum", "auto", "entry", ] #: TypeAttr column storing ``Dict[str, Enum]`` (instance name → singleton). ENUM_ENTRIES_ATTR = "__ffi_enum_entries__" #: TypeAttr column storing ``Dict[str, List[Any]]`` of per-variant attrs. ENUM_ATTRS_ATTR = "__ffi_enum_attrs__" #: TypeAttr column storing ``Dict[Any, Enum]`` (payload → singleton) — #: parallel to :data:`ENUM_ENTRIES_ATTR` but keyed by the user-visible #: payload (``IntEnum.value`` / ``StrEnum.value``). Populated by every #: creator of a payload-carrying variant; consumed by FFI converters for #: O(1) payload → variant resolution. ENUM_VALUE_ENTRIES_ATTR = "__ffi_enum_value_entries__" # --------------------------------------------------------------------------- # entry() sentinel # --------------------------------------------------------------------------- class _EnumEntry: """Sentinel produced by :func:`entry`; consumed by ``Enum.__init_subclass__``. Holds the positional and keyword arguments forwarded to the subclass's ``__init__`` when the variant is materialized. ``_value`` and ``_name`` are auto-assigned (dense ordinal and class-body name) and must not appear in the captured arguments. """ __slots__ = ("args", "kwargs") def __init__(self, *args: Any, **kwargs: Any) -> None: self.args: tuple[Any, ...] = args self.kwargs: dict[str, Any] = kwargs def __repr__(self) -> str: parts = [repr(a) for a in self.args] parts.extend(f"{k}={v!r}" for k, v in self.kwargs.items()) return f"entry({', '.join(parts)})" def entry(*args: Any, **kwargs: Any) -> Any: """Declare a new enum variant with values for its declared fields. ``entry(...)`` is a class-body sentinel; it never produces a real instance. At class creation, :meth:`Enum.__init_subclass__` scans for these sentinels and, for each one, constructs a singleton variant by forwarding the captured positional and keyword arguments to the subclass's ``__init__``, together with an auto-assigned :attr:`~Enum._value` (dense ordinal) and :attr:`~Enum._name` (class-body name). Prefer :func:`auto` when a variant has no declared fields beyond the auto-assigned ordinal and name — it expresses intent without the empty-arg-list noise. When the enum's ``type_key`` is C++-backed (registered via ``refl::ObjectDef``), only keyword arguments are supported — field values are assigned via reflected setters keyed by name. Passing positional arguments in that case raises :class:`TypeError`. ``entry(_value=...)`` and ``entry(_name=...)`` always raise :class:`TypeError` because those fields are auto-assigned. Examples -------- Variant with declared fields: .. code-block:: python from typing import ClassVar class Activation(Enum, type_key="my.Activation"): output_zero: bool is_monotonic: bool relu: ClassVar[Activation] = entry(output_zero=True, is_monotonic=True) gelu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=False) Returns ------- object An opaque sentinel. The declared return type is ``Any`` so that ``ClassVar[Cls] = entry(...)`` type-checks even though the sentinel is not a real ``Cls``. """ return _EnumEntry(*args, **kwargs) def auto() -> Any: """Declare a new enum variant with no declared fields. Semantically equivalent to :func:`entry` called with no arguments but reads more clearly for the common case where a variant differs from its siblings only by name and ordinal. The resulting singleton has only the auto-assigned :attr:`~Enum._value` and :attr:`~Enum._name`. ``auto()`` registers a *new* Python-side variant; it is not the right tool for binding to a pre-existing C++-registered entry (use a bare ``ClassVar[Cls]`` annotation for that — see :class:`Enum`). Examples -------- .. code-block:: python class Status(Enum, type_key="my.Status"): ok = auto() err = auto() retry = auto() assert Status.ok._value == 0 assert Status.err._name == "err" Returns ------- object An opaque sentinel, the same kind returned by :func:`entry`. The declared return type is ``Any`` so that both ``name = auto()`` and ``name: ClassVar[Cls] = auto()`` type-check. """ return _EnumEntry() # --------------------------------------------------------------------------- # Class-level helpers # --------------------------------------------------------------------------- class _ClassProperty: """Read-only descriptor whose getter receives the owning class. Used for ``attr_dict`` so it works as a class-level attribute access (e.g., ``Op.attr_dict["has_side_effects"]``) without needing a metaclass. """ __slots__ = ("_fget",) def __init__(self, fget: Callable[[type], Any]) -> None: self._fget = fget def __get__(self, instance: Any, owner: type | None = None) -> Any: cls = owner if owner is not None else type(instance) return self._fget(cls) def _install_payload_enum_behaviors(cls: type, *, user_defined_repr: bool) -> None: """Install stdlib-like instance behavior for ``IntEnum`` / ``StrEnum``. Payload enums have two identities that plain ``Enum`` does not: the FFI singleton identity (``Priority.low.same_as(...)``) and a user-visible raw payload (``Priority.low.value == 10``). These methods make the Python surface follow the payload where users expect stdlib-like behavior, while keeping the object identity available through ``same_as``. This runs after ``py_class`` / ``c_class`` because those decorators may install generated object dunders. Installing here lets payload enums override generated repr/equality/hash defaults, but still preserves methods explicitly written in the enum subclass body. """ def _payload_enum_eq(self: Enum, other: object) -> Any: """Compare payload enum variants by payload and accept raw payloads. Two variants of the same enum class compare like their ``value`` fields, and a variant can compare directly with its raw payload type (``Priority.low == 10``). Unrelated objects return ``NotImplemented`` so Python can try reflected comparison or fall back normally. """ if isinstance(other, type(self)): return self.value == other.value # type: ignore[attr-defined] payload_type = getattr(type(self), "__ffi_enum_payload_value_type__", None) if payload_type is not None and isinstance(other, payload_type): return self.value == other # type: ignore[attr-defined] return NotImplemented def _payload_enum_ne(self: Enum, other: object) -> Any: """Negate payload equality while preserving ``NotImplemented`` semantics.""" if (eq := _payload_enum_eq(self, other)) is not NotImplemented: return not eq return NotImplemented def _payload_enum_str(self: Enum) -> str: """Render as the raw payload, matching ``int`` / ``str`` enum expectations.""" return str(self.value) # type: ignore[attr-defined] def _payload_enum_hash(self: Enum) -> int: """Hash like the raw payload so equality and hashing stay consistent.""" return hash(self.value) # type: ignore[attr-defined] _PAYLOAD_ENUM_DUNDERS = { "__eq__": _payload_enum_eq, "__ne__": _payload_enum_ne, "__str__": _payload_enum_str, "__hash__": _payload_enum_hash, } def _payload_enum_name(self: Enum) -> str: """Expose the declaration name without making plain ``Enum`` grow ``name``.""" return str(self._name) def _payload_enum_repr(self: Enum) -> str: """Render as ``Class.member``, matching Python payload enum conventions.""" return f"{type(self).__name__}.{self.name}" # ty: ignore[unresolved-attribute] if "name" not in cls.__dict__: cls.name = property(_payload_enum_name) # ty: ignore[unresolved-attribute] if not user_defined_repr: cls.__repr__ = _payload_enum_repr # type: ignore[attr-defined] for name, method in _PAYLOAD_ENUM_DUNDERS.items(): if name not in cls.__dict__: setattr(cls, name, method) # --------------------------------------------------------------------------- # Enum base + EnumAttrMap # --------------------------------------------------------------------------- class _EnumMeta(type(Object)): """Metaclass for :class:`Enum` that wires class-level iteration and length. Kept as a dedicated subclass of ``type(Object)`` so ``iter(MyEnum)`` and ``len(MyEnum)`` only resolve on :class:`Enum` subclasses, rather than mutating the shared ``Object`` metaclass globally. """ def __iter__(cls) -> Iterator[Any]: return iter(_ordered_entries(cls)) # ty: ignore[unresolved-attribute] def __len__(cls) -> int: return len(_ordered_entries(cls)) # ty: ignore[unresolved-attribute] def __call__(cls, *args: Any, **kwargs: Any) -> Any: """Construct enum variants from existing variants, payloads, or names. Normal object construction is still delegated to ``Object`` / dataclass initialization. The single-argument path is intercepted to match stdlib enum lookup behavior: * ``Priority(Priority.low)`` returns the same variant. * ``Priority(10)`` returns the member whose payload value is ``10``. * ``Priority("low")`` returns the member declared with that name. Payload lookup is only enabled for classes with ``__ffi_enum_payload_value_type__`` so plain enums do not accidentally interpret arbitrary class-body values as public payloads. """ if not kwargs and len(args) == 1: value = args[0] if isinstance(value, cls): return value payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None) if payload_value_type is not None and isinstance(value, payload_value_type): value_entries = _value_entries_dict(cls) if value_entries is not None and value in value_entries: return value_entries[value] if isinstance(value, str): entries = _entries_dict(cls) if entries is not None and value in entries: return entries[value] return super().__call__(*args, **kwargs) @dataclass_transform( eq_default=False, order_default=False, frozen_default=True, field_specifiers=(Field, field, entry, auto), ) @c_class("ffi.Enum", init=False) class Enum(Object, metaclass=_EnumMeta): """A named-singleton registry with cross-language identity. Subclasses declare variants: frozen, named, ordinal-indexed singletons — the familiar enum pattern. Unlike ``enum.Enum``, an ``Enum`` subclass bound to an FFI-registered ``type_key`` has an **open variant set**: C++ translation units and other Python modules binding the same ``type_key`` can contribute variants to the shared registry. Per-variant metadata can also be attached post-hoc via :meth:`def_attr` as an *extensible attribute*, outside the class definition. For **closed, Python-only enums**, use a fresh ``type_key`` with :func:`auto` / :func:`entry` — behavior matches ``enum.Enum``. Attributes ---------- _value : int Dense ordinal assigned at registration (0-indexed per class). _name : str The variant's string name key (e.g., ``"Add"`` for ``Op.Add``). Closed Python enum ------------------ Pick a fresh ``type_key`` and list variants with :func:`auto` or :func:`entry`. The variant set is fixed at class-definition time. .. code-block:: python class Priority(Enum, type_key="my.Priority"): low = auto() medium = auto() high = auto() # Variants with declared fields — values supplied via entry(...). class Activation(Enum, type_key="my.Activation"): output_zero: bool is_monotonic: bool relu: ClassVar[Activation] = entry(output_zero=True, is_monotonic=True) gelu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=False) Cross-language registry ----------------------- When ``type_key`` is already registered (typically by C++), the Python class *binds* to the existing type rather than creating a new one. Bare ``ClassVar[Cls]`` annotations bind to variants already registered on the C++ side; :func:`entry` / :func:`auto` still register fresh Python variants whose ordinals extend past the C++ ones. All variants — regardless of origin — land in the same per-class registry and are visible to every binder of the same ``type_key``. .. code-block:: python # Registered in C++ via refl::EnumDef("Alpha")... . class Variant(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Variant] # binds to C++-registered "Alpha" Beta: ClassVar[Variant] # binds to C++-registered "Beta" Declaration forms ----------------- Four shapes are supported in the class body: 1. ``name = auto()`` — new variant with no declared fields. 2. ``name: ClassVar[Cls] = entry(**kwargs)`` — new variant; ``kwargs`` populate declared fields. 3. ``name = entry(**kwargs)`` — same as (2), without the ``ClassVar`` annotation. 4. ``name: ClassVar[Cls]`` — in cross-language mode, binds to an existing C++-registered variant (error if unknown); otherwise registers a new Python variant with only the auto-assigned :attr:`_value` and :attr:`_name` (equivalent to ``name = auto()``). Integer literals (``ok = 0``) are rejected: :attr:`_value` is auto-assigned, so a user-supplied ordinal would either silently duplicate or conflict. ``entry(_value=...)`` and ``entry(_name=...)`` raise :class:`TypeError` at class-body time. Differences from ``enum.Enum`` ------------------------------ * **Same**: hidden identity fields :attr:`_name`, :attr:`_value`, iteration, identity comparison; closed-set behavior when ``type_key`` is fresh. * **Extended**: ``entry(**kwargs)`` replaces the tuple-RHS idiom; ``dataclass_transform`` gives native type-checker support; open registry when ``type_key`` is shared across languages. * **Different**: :attr:`_value` is always the ordinal; :meth:`def_attr` adds extensible attributes outside the class schema. Use :class:`IntEnum` / :class:`StrEnum` when you need a user-visible ``value`` field. * **Not provided**: ``Flag`` / ``IntFlag``, member aliasing, ``_missing_`` hook. Subclasses inherit :meth:`get`, :meth:`all_entries`, :meth:`def_attr`, and the ``attr_dict`` class-level view. """ __slots__ = () _value: int _name: str __ffi_enum_payload_value_type__: object = None def __init_subclass__( cls, *, type_key: str | None = None, frozen: bool = True, init: bool = True, repr: bool = True, **kwargs: Any, ) -> None: super().__init_subclass__(**kwargs) if type_key is None: return user_defined_repr = "__repr__" in cls.__dict__ binders, python_entries = _collect_entry_declarations(cls) cxx_backed = core._object_type_key_to_index(type_key) is not None if cxx_backed: c_class(type_key, init=init, repr=repr)(cls) else: py_class(type_key, frozen=frozen)(cls) if getattr(cls, "__ffi_enum_payload_value_type__", None) is not None: _install_payload_enum_behaviors(cls, user_defined_repr=user_defined_repr) _resolve_entries(cls, binders, python_entries, type_key=type_key, cxx_backed=cxx_backed) @classmethod def get(cls, name: str) -> Enum: """Return the variant named *name*, or raise :class:`KeyError`.""" entries = _entries_dict(cls) if entries is not None and name in entries: return entries[name] raise KeyError(f"{cls.__name__} has no variant named {name!r}") @classmethod def all_entries(cls) -> Iterator[Enum]: """Iterate over all variants, in ordinal (``_value``) order.""" return iter(_ordered_entries(cls)) @_ClassProperty def attr_dict(cls: type) -> Any: """Live ``Dict[str, List[Any]]`` backing every extensible attribute. The outer dict is keyed by extensible-attribute name; each value is a list indexed by variant ordinal. Prefer :meth:`def_attr` for normal per-variant reads and writes; this property is for bulk inspection and for reading values written by C++ ``EnumDef::set_attr``. """ return _attrs_dict(cls) or Dict({}) @classmethod def def_attr( cls, name: str, *, default: Any = core.MISSING, ) -> EnumAttrMap: """Declare an *extensible attribute* column on this enum. Extensible attributes let any consumer associate per-variant data outside the enum's class-body schema — a lowering function attached to an operator by a code generator, a cost model registered only on some targets, a documentation string added after the fact. Writes are last-write-wins for the same ``(variant, name)`` pair and visible to every consumer that calls :meth:`def_attr` with the same name on the same enum, including C++ code writing via ``EnumDef::set_attr``. Extensible attributes differ from **declared fields**: ==================== ======================== ========================== Concept Lives on Added by ==================== ======================== ========================== Declared field The variant object Enum author, in class body Extensible attribute ``__ffi_enum_attrs__`` Any consumer, any time ==================== ======================== ========================== Rule of thumb: if the data is part of *what a variant is*, declare a field in the class body; if it's part of *what a consumer wants to know*, attach it with :meth:`def_attr`. Parameters ---------- name The extensible-attribute name (e.g., ``"has_side_effects"``). Writes go to ``attr_dict[name]`` as a list indexed by each variant's ordinal. default Value returned by ``attr[variant]`` when nothing was registered for that variant. Left as ``MISSING`` to raise :class:`KeyError` on unset variants instead. The default is a property of *this* :class:`EnumAttrMap` view, not of the underlying column: calling :meth:`def_attr` twice with the same ``name`` but different defaults creates two views that share every explicit write but may disagree on unset variants — e.g., ``Op.def_attr("cost", default=0)`` and ``Op.def_attr("cost", default=-1)`` return ``0`` and ``-1`` respectively for a variant that was never written to. Returns ------- EnumAttrMap Mutable view over the column. Use ``variant in attr`` to distinguish an explicit write from a default-hit. ``None`` is reserved as the "unset" sentinel (matching C++ ``EnumDef::set_attr`` padding), so ``attr[variant] = None`` raises :class:`TypeError` — store a typed wrapper (e.g. a ``0``/``False`` flag) when you need a falsy-but-present value. Notes ----- :meth:`def_attr` is not a way to add fields to the enum's schema, subclass frozen variants, or bypass the frozen-instance invariant via ``setattr`` — for that, declare a field in the class body instead. """ return EnumAttrMap(cls, name, default=default) class EnumAttrMap: """Mutable per-variant view over an extensible-attribute column. Returned by :meth:`Enum.def_attr`. Writes go to a ``List[Any]`` column keyed by extensible-attribute name inside the per-class ``__ffi_enum_attrs__`` dict. The list is indexed by each variant's ordinal (``variant._value``) and padded with ``None`` as new variants are registered. The column is shared across every consumer — including C++ code writing via ``EnumDef::set_attr`` — and the data is not a field on the variant object. See :meth:`Enum.def_attr` for full semantics. ``None`` is reserved as the column's "unset" sentinel (matching the C++ ``Any(nullptr)`` padding used by ``EnumDef::set_attr``), so :meth:`__setitem__` rejects ``None`` with :class:`TypeError` — an explicit ``attr[variant] = None`` would otherwise be indistinguishable from never-written and surprise ``variant in attr`` / :meth:`__getitem__` readers. To "clear" a previously written value, register a mutable container once and mutate it in place. """ __slots__ = ("_default", "_enum_cls", "_name") def __init__(self, enum_cls: type, name: str, *, default: Any = core.MISSING) -> None: self._enum_cls = enum_cls self._name = name self._default = default def _ordinal_of(self, variant: object) -> int: if not isinstance(variant, self._enum_cls): raise TypeError( f"{self._enum_cls.__name__}.def_attr({self._name!r}) expects a " f"{self._enum_cls.__name__} variant, got {type(variant).__name__}" ) return int(variant._value) # type: ignore[attr-defined] def _column(self, *, create: bool) -> Any | None: """Return the ``List[Any]`` column for this attribute; create if missing. Returns ``None`` when ``create`` is false and the column doesn't exist. """ attrs = _attrs_dict(self._enum_cls) if not create else _ensure_attrs_dict(self._enum_cls) if attrs is None: return None if self._name in attrs: return attrs[self._name] if not create: return None col = List([]) attrs[self._name] = col return col def __setitem__(self, variant: object, value: Any) -> None: if value is None: raise TypeError( f"{self._enum_cls.__name__}.def_attr({self._name!r}): " f"None is reserved as the 'unset' sentinel for extensible " f"attributes and cannot be written explicitly." ) ordinal = self._ordinal_of(variant) col = self._column(create=True) assert col is not None # create=True always materialises the column. while len(col) <= ordinal: col.append(None) col[ordinal] = value def __getitem__(self, variant: object) -> Any: ordinal = self._ordinal_of(variant) col = self._column(create=False) if col is not None and ordinal < len(col): v = col[ordinal] if v is not None: return v if self._default is core.MISSING: raise KeyError( f"{self._enum_cls.__name__}.{variant._name} has no " # type: ignore[attr-defined] f"extensible attribute {self._name!r} set" ) return self._default def __contains__(self, variant: object) -> bool: if not isinstance(variant, self._enum_cls): return False try: ordinal = self._ordinal_of(variant) except TypeError: return False col = self._column(create=False) return col is not None and ordinal < len(col) and col[ordinal] is not None def get(self, variant: object, default: Any = None) -> Any: """Return the value for *variant*, or *default* if unset or foreign.""" if not isinstance(variant, self._enum_cls): return default try: return self[variant] except KeyError: return default @property def name(self) -> str: """The extensible-attribute name passed to :meth:`Enum.def_attr`.""" return self._name class IntEnum(Enum): """Enum variant base with a user-visible ``value: int`` field.""" __slots__ = () value: int def __init_subclass__( cls, *, type_key: str | None = None, frozen: bool = True, init: bool = True, repr: bool = True, **kwargs: Any, ) -> None: _prepare_payload_enum_subclass(cls, value_type=int, base_name="IntEnum") super().__init_subclass__( type_key=type_key, frozen=frozen, init=init, repr=repr, **kwargs, ) class StrEnum(Enum): """Enum variant base with a user-visible ``value: str`` field.""" __slots__ = () value: str def __init_subclass__( cls, *, type_key: str | None = None, frozen: bool = True, init: bool = True, repr: bool = True, **kwargs: Any, ) -> None: _prepare_payload_enum_subclass(cls, value_type=str, base_name="StrEnum") super().__init_subclass__( type_key=type_key, frozen=frozen, init=init, repr=repr, **kwargs, ) # --------------------------------------------------------------------------- # TypeAttr accessors # --------------------------------------------------------------------------- def _entries_dict(cls: type) -> Any: type_info = getattr(cls, "__tvm_ffi_type_info__", None) if type_info is None: return None return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR) def _attrs_dict(cls: type) -> Any: type_info = getattr(cls, "__tvm_ffi_type_info__", None) if type_info is None: return None return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR) def _value_entries_dict(cls: type) -> Any: """Return the live ``Dict[Any, Enum]`` payload-to-variant index, or None.""" type_info = getattr(cls, "__tvm_ffi_type_info__", None) if type_info is None: return None return core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) def _ordered_entries(cls: type) -> list[Any]: """Return all variants ordered by ordinal.""" entries = _entries_dict(cls) if entries is None: return [] ordered: list[Any] = [None] * len(entries) for inst in entries.values(): ordered[int(inst._value)] = inst return ordered def _ensure_entries_dict(cls: type) -> Any: """Return the live ``__ffi_enum_entries__`` dict, registering it if absent.""" type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] entries = core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR) if entries is not None: return entries entries = Dict({}) core._register_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR, entries) # Re-read so mutations go through the ref owned by the registry. return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR) def _ensure_attrs_dict(cls: type) -> Any: """Return the live ``__ffi_enum_attrs__`` dict, registering it if absent.""" type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] attrs = core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR) if attrs is not None: return attrs attrs = Dict({}) core._register_type_attr(type_info.type_index, ENUM_ATTRS_ATTR, attrs) return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR) def _ensure_value_entries_dict(cls: type) -> Any: """Return the live ``__ffi_enum_value_entries__`` dict, registering it if absent.""" type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] entries = core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) if entries is not None: return entries entries = Dict({}) core._register_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR, entries) return core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) # --------------------------------------------------------------------------- # Class-body scanning + entry materialisation # --------------------------------------------------------------------------- def _collect_entry_declarations( cls: type, ) -> tuple[list[str], dict[str, _EnumEntry]]: """Scan *cls.__dict__* for variant declarations. Returns ``(binders, python_entries)`` in declaration order: * *binders* — names annotated as ``ClassVar[Cls]`` with no assigned value. Each either binds to an existing C++-registered entry with the same name or registers a new blank Python entry. * *python_entries* — names assigned an ``entry(...)`` sentinel (with or without a ``ClassVar`` annotation). Each registers a new Python entry using the captured args/kwargs. Matched assignments are removed from ``cls.__dict__`` so that ``@c_class`` / ``@py_class`` don't misinterpret them as field defaults or class constants. """ annotations = _own_annotations(cls) dict_keys = set(cls.__dict__.keys()) binders: list[str] = [] ordinary_fields: set[str] = set() payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None) for name, ann in annotations.items(): if name.startswith("_"): continue if _is_class_var(ann): if name not in dict_keys: binders.append(name) else: ordinary_fields.add(name) python_entries: dict[str, _EnumEntry] = {} for name, value in list(cls.__dict__.items()): if name.startswith("_"): continue if isinstance(value, _EnumEntry): python_entries[name] = value try: delattr(cls, name) except AttributeError: pass elif payload_value_type is not None and name not in ordinary_fields: if isinstance(value, (staticmethod, classmethod, property)) or callable(value): continue python_entries[name] = _EnumEntry(value=value) try: delattr(cls, name) except AttributeError: pass return binders, python_entries def _resolve_entries( cls: type, binders: list[str], python_entries: dict[str, _EnumEntry], *, type_key: str, cxx_backed: bool, ) -> None: """Materialise *binders* and *python_entries* into class-attribute singletons. Processing order matches declaration order: ``binders`` first (because their annotations appear before any class-body assignments), then ``python_entries`` in their class-body order. Each newly registered entry gets a dense ordinal equal to the current entries-dict size, so ordinals stay compact and stable across registrations. A cxx-backed enum (``type_key`` was already registered in the FFI type system before this Python subclass was created) supports mixing C++ and Python entries: bare ``ClassVar[Cls]`` binders must name an existing C++-registered entry, but ``entry(...)``/``auto()`` sentinels may add fresh Python-side entries whose ordinals extend past the C++ entries. """ entries = _ensure_entries_dict(cls) payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None) for name in binders: if name in entries: # Already materialised — either C++-registered or previously bound. instance = entries[name] setattr(cls, name, instance) _index_payload(cls, instance, payload_value_type) continue if cxx_backed: raise _cxx_backed_unknown_binder_error(cls, name, type_key, entries) ordinal = len(entries) instance = _instantiate(cls, args=(), kwargs={}, ordinal=ordinal, name=name) entries[name] = instance setattr(cls, name, instance) _index_payload(cls, instance, payload_value_type) for name, e in python_entries.items(): if name in entries: raise RuntimeError( f"Duplicate enum entry {name!r} for {cls.__name__}: already " f"registered as ordinal {int(entries[name]._value)}." ) if "_value" in e.kwargs or "_name" in e.kwargs: raise TypeError( f"{cls.__name__}.{name}: `_value` and `_name` are auto-assigned " f"and must not appear in entry(...) arguments." ) ordinal = len(entries) instance = _instantiate_entry( cls, entry=e, ordinal=ordinal, name=name, cxx_backed=cxx_backed ) entries[name] = instance setattr(cls, name, instance) _index_payload(cls, instance, payload_value_type) def _index_payload(cls: type, instance: Any, payload_value_type: type | None) -> None: """Record ``(instance.value → instance)`` in the value-entries column. No-op for non-payload enums. For payload enums, the first writer of a given payload wins — matches the "first-match" semantics of the linear scan that FFI converters fall back to when this column is incomplete. """ if payload_value_type is None: return payload = getattr(instance, "value", None) if payload is None: return value_entries = _ensure_value_entries_dict(cls) if payload not in value_entries: value_entries[payload] = instance def _cxx_backed_unknown_binder_error( cls: type, name: str, type_key: str, entries: Any, ) -> RuntimeError: """Build a descriptive error for an unbindable bare ``ClassVar`` binder. A bare ``ClassVar[Cls]`` annotation on a cxx-backed enum means "bind to an existing C++ entry with this name" — if the C++ registry has no such entry, the declaration is almost always a typo. For adding a *new* Python-side variant on a cxx-backed enum, use ``entry(...)`` or ``auto()`` instead. """ known = list(entries.keys()) if entries is not None else [] known_str = ", ".join(repr(k) for k in known) if known else "" return RuntimeError( f"Cannot bind enum variant {name!r} on {cls.__name__}: the FFI " f"type {type_key!r} is already registered in C++ with entries " f"[{known_str}], but has no C++ entry named {name!r}. " f"Bare ``ClassVar[{cls.__name__}]`` binders on a C++-backed enum " f"must name an entry already registered in C++; they cannot " f"introduce new variants from Python. " f"If this was a typo, double-check the spelling against the known " f"entries above (`{name}: ClassVar[{cls.__name__}]`); if you meant " f"to add a new Python-side variant, use `{name} = auto()` or " f"`{name} = entry(...)` instead." ) def _instantiate( cls: type, *, args: tuple[Any, ...], kwargs: dict[str, Any], ordinal: int, name: str, ) -> Any: """Construct a subclass instance with auto-assigned ``_value``/``_name``.""" merged = dict(kwargs) merged["_value"] = ordinal merged["_name"] = name return cls(*args, **merged) def _instantiate_entry( cls: type, *, entry: _EnumEntry, ordinal: int, name: str, cxx_backed: bool, ) -> Any: """Instantiate an enum entry and normalize construction errors.""" try: if cxx_backed: return _instantiate_cxx_backed( cls, args=entry.args, kwargs=entry.kwargs, ordinal=ordinal, name=name ) return _instantiate(cls, args=entry.args, kwargs=entry.kwargs, ordinal=ordinal, name=name) except Exception as err: raise TypeError(f"{cls.__name__}.{name}: invalid enum entry: {err}") from None def _instantiate_cxx_backed( cls: type, *, args: tuple[Any, ...], kwargs: dict[str, Any], ordinal: int, name: str, ) -> Any: """Construct a new variant of a cxx-backed enum without going through ``__init__``. C++-backed enums whose underlying type is registered with ``refl::init(false)`` (e.g. any subclass of ``EnumObj`` in C++) have no ``__ffi_init__``, so the usual ``cls(_value=..., _name=...)`` path is not available. Mirror ``reflection::EnumDef`` by allocating a blank instance via ``__ffi_new__`` and populating fields through the frozen-setter escape hatch exposed on the reflected property descriptors. """ if args: raise TypeError( f"{cls.__name__}.{name}: positional `entry(...)` args are not " f"supported when extending a C++-backed enum; use keyword " f"arguments naming reflected fields." ) type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] ffi_new = core._lookup_type_attr(type_info.type_index, "__ffi_new__") if ffi_new is None: raise RuntimeError( f"Cannot add Python enum variant {name!r} on {cls.__name__}: " f"its C++ type has no ``__ffi_new__`` allocator registered, so " f"blank instances cannot be created from Python." ) instance = ffi_new() for key in ("_value", "_name", *kwargs.keys()): descriptor = getattr(cls, key, None) if descriptor is None or not hasattr(descriptor, "set"): raise TypeError( f"{cls.__name__}.{name}: cannot set field {key!r} on a " f"C++-backed enum — no reflected setter is available." ) getattr(cls, "_value").set(instance, ordinal) getattr(cls, "_name").set(instance, name) for k, v in kwargs.items(): getattr(cls, k).set(instance, v) return instance # --------------------------------------------------------------------------- # Annotation helpers # --------------------------------------------------------------------------- def _own_annotations(cls: type) -> dict[str, Any]: """Return *cls*'s own annotations dict (not inherited).""" if sys.version_info >= (3, 14): return dict(getattr(cls, "__annotations__", {}) or {}) return dict(cls.__dict__.get("__annotations__", {})) def _is_class_var(annotation: Any) -> bool: """Return True if *annotation* is ``ClassVar`` or ``ClassVar[...]``.""" if annotation is ClassVar: return True if typing.get_origin(annotation) is ClassVar: return True if isinstance(annotation, str): stripped = annotation.replace(" ", "") return stripped.startswith("ClassVar") or stripped.startswith("typing.ClassVar") return False def _annotation_matches_expected(annotation: Any, expected_type: type) -> bool: """Return True if *annotation* matches the required payload-field type.""" if annotation is expected_type: return True if isinstance(annotation, str): stripped = annotation.replace(" ", "") expected = expected_type.__name__ return stripped in {expected, f"builtins.{expected}"} return False def _prepare_payload_enum_subclass( cls: type[Enum], *, value_type: type[Any], base_name: str ) -> None: """Inject and validate the user-visible ``value`` field for payload enums.""" if "value" in cls.__dict__: raise TypeError( f"{base_name} reserves `value` as a declared field; use " "`entry(value=...)` or ` = ` to assign each " "variant's payload." ) annotations = _own_annotations(cls) existing = annotations.get("value") if existing is not None: if _is_class_var(existing): raise TypeError(f"{base_name} reserves `value` as a declared field, not a ClassVar.") if not _annotation_matches_expected(existing, value_type): raise TypeError( f"{base_name} requires `value: {value_type.__name__}` on subclasses, " f"got {existing!r}." ) return updated = _own_annotations(cls) updated["value"] = value_type cls.__annotations__ = updated cls.__ffi_enum_payload_value_type__ = value_type tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/field.py000066400000000000000000000246101521067262500220530ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Field descriptor and ``field()`` helper for Python-defined TVM-FFI types.""" from __future__ import annotations import sys from collections.abc import Callable from typing import Any, ClassVar from ..core import MISSING, TypeSchema # Re-export the stdlib KW_ONLY sentinel so type checkers recognise # ``_: KW_ONLY`` as a keyword-only boundary rather than a real field. # dataclasses.KW_ONLY was added in Python 3.10; on older runtimes we # define a class sentinel (a class, not an instance, so that ``_: KW_ONLY`` # is a valid type annotation for static analysers targeting 3.9). if sys.version_info >= (3, 10): from dataclasses import KW_ONLY else: class KW_ONLY: """Sentinel type: annotations after ``_: KW_ONLY`` are keyword-only.""" class Field: """Descriptor for a single field in a Python-defined TVM-FFI type. When constructed directly (low-level API), *name* and *_ty_schema* should be provided. When returned by :func:`field` (``@py_class`` workflow), both are ``None`` and filled in by the decorator. Parameters ---------- name : str | None The field name. ``None`` when created via :func:`field`; filled in by the ``@py_class`` decorator. _ty_schema : TypeSchema | None Private: the internal :class:`TypeSchema` used by the reflection layer. ``None`` when created via :func:`field`; filled in by the ``@py_class`` decorator. Consumers should use :attr:`type` instead. type : Any The resolved Python annotation (e.g. ``int``, ``list[str]``, ``Optional[X]``). Filled in by the ``@py_class`` / ``@c_class`` decorator via :func:`typing.get_type_hints`; ``None`` until then or when the annotation cannot be resolved. default : object Default value for the field. Mutually exclusive with *default_factory*. ``MISSING`` when not set. default_factory : Callable[[], object] | None A zero-argument callable that produces the default value. Mutually exclusive with *default*. ``None`` when not set. frozen : bool Whether this field is read-only after ``__init__``. init : bool Whether this field appears in the auto-generated ``__init__``. repr : bool Whether this field appears in ``__repr__`` output. hash : bool | None Whether this field participates in recursive hashing. ``None`` means "follow *compare*" (the native dataclass default). compare : bool Whether this field participates in recursive comparison. kw_only : bool | None Whether this field is keyword-only in ``__init__``. ``None`` means "inherit from the decorator-level *kw_only* flag". structural_eq : str | None Structural equality/hashing annotation for this field. Valid values are: - ``None`` (default): the field participates normally in structural comparison and hashing. - ``"ignore"``: the field is excluded from structural equality and hashing entirely (e.g. source spans, caches). - ``"def-recursive"`` (alias: ``"def"``): the field is a **recursive definition region** that introduces new variable bindings. Free variables encountered anywhere in this field's subtree (including inside the var's own sub-fields) are mapped by position. One example is function parameter lists, where the value var and any shape parameters in its type are co-introduced at the same site. - ``"def-non-recursive"``: the field is a **non-recursive definition region**. Only the immediate free var(s) at this field's value bind; free vars inside their sub-fields must resolve against an outer binding (use semantics). One example is a normal binding whose value type contains shape parameters that reference outer-scope vars. doc : str | None Optional docstring for the field. """ __slots__ = ( "_ty_schema", "compare", "default", "default_factory", "doc", "frozen", "hash", "init", "kw_only", "name", "repr", "structural_eq", "type", ) name: str | None _ty_schema: TypeSchema | None type: Any default: object default_factory: Callable[[], object] | None frozen: bool init: bool repr: bool hash: bool | None compare: bool kw_only: bool | None structural_eq: str | None doc: str | None #: Valid values for the *structural_eq* parameter. #: #: ``"def"`` is kept as a Python-side alias for ``"def-recursive"`` to #: preserve back-compat with code written against the old single-flag #: ``SEqHashDef`` API. _VALID_STRUCTURAL_EQ_VALUES: ClassVar[frozenset[str | None]] = frozenset( {None, "ignore", "def", "def-recursive", "def-non-recursive"} ) def __init__( # noqa: PLR0913 self, name: str | None = None, _ty_schema: TypeSchema | None = None, *, default: object = MISSING, default_factory: Callable[[], object] | None = MISSING, # type: ignore[assignment] frozen: bool = False, init: bool = True, repr: bool = True, hash: bool | None = True, compare: bool = False, kw_only: bool | None = False, structural_eq: str | None = None, doc: str | None = None, ) -> None: # MISSING means "parameter not provided". # An explicit None from the user fails the callable() check, # matching stdlib dataclasses semantics. if default_factory is not MISSING: if default is not MISSING: raise ValueError("cannot specify both default and default_factory") if not callable(default_factory): raise TypeError( f"default_factory must be a callable, got {type(default_factory).__name__}" ) if structural_eq not in Field._VALID_STRUCTURAL_EQ_VALUES: raise ValueError( f"structural_eq must be one of " f"{sorted(Field._VALID_STRUCTURAL_EQ_VALUES, key=str)}, " f"got {structural_eq!r}" ) self.name = name self._ty_schema = _ty_schema self.type = None self.default = default self.default_factory = default_factory self.frozen = frozen self.init = init self.repr = repr self.hash = hash self.compare = compare self.kw_only = kw_only self.structural_eq = structural_eq self.doc = doc def field( *, default: object = MISSING, default_factory: Callable[[], object] | None = MISSING, # type: ignore[assignment] frozen: bool = False, init: bool = True, repr: bool = True, hash: bool | None = None, compare: bool = True, kw_only: bool | None = None, structural_eq: str | None = None, doc: str | None = None, ) -> Any: """Customize a field in a ``@py_class``-decorated class. Returns a :class:`Field` sentinel whose *name* and *_ty_schema* are ``None``. The ``@py_class`` decorator fills them in later from the class annotations. The return type is ``Any`` because ``dataclass_transform`` field specifiers must be assignable to any annotated type (e.g. ``x: int = field(default=0)``). Parameters ---------- default Default value for the field. Mutually exclusive with *default_factory*. default_factory A zero-argument callable that produces the default value. Mutually exclusive with *default*. frozen Whether this field is read-only after ``__init__``. When True, the Python property descriptor has no setter; use the ``type(obj).field_name.set(obj, value)`` escape hatch when mutation is necessary. init Whether this field appears in the auto-generated ``__init__``. repr Whether this field appears in ``__repr__`` output. hash Whether this field participates in recursive hashing. ``None`` (default) means "follow *compare*". compare Whether this field participates in recursive comparison. kw_only Whether this field is keyword-only in ``__init__``. ``None`` means "inherit from the decorator-level ``kw_only`` flag". structural_eq Structural equality/hashing annotation. ``None`` (default) means the field participates normally. ``"ignore"`` excludes the field from structural comparison and hashing. ``"def-recursive"`` (alias ``"def"``) marks the field as a recursive definition region: free vars in the field's whole subtree bind. ``"def-non-recursive"`` marks it as a non-recursive definition region: only immediate free vars bind; nested free vars must resolve against an outer binding. doc Optional docstring for the field. Returns ------- Any A :class:`Field` sentinel recognised by ``@py_class``. Examples -------- .. code-block:: python @py_class class Point(Object): x: float y: float = field(default=0.0, repr=False) @py_class(structural_eq="tree") class MyFunc(Object): params: Array = field(structural_eq="def") body: Expr span: Object = field(structural_eq="ignore") """ return Field( default=default, default_factory=default_factory, frozen=frozen, init=init, repr=repr, hash=hash, compare=compare, kw_only=kw_only, structural_eq=structural_eq, doc=doc, ) tvm-ffi-0.1.12/python/tvm_ffi/dataclasses/py_class.py000066400000000000000000001005611521067262500226050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """The ``py_class`` decorator: Python-defined FFI classes with dataclass semantics.""" from __future__ import annotations import sys import typing from collections.abc import Callable from copy import copy from dataclasses import dataclass from typing import Any, ClassVar, TypeVar from typing_extensions import dataclass_transform from .. import core from .._dunder import _install_dataclass_dunders from ..core import MISSING, TypeSchema from ..registry import _add_class_attrs from .field import KW_ONLY, Field, field _T = TypeVar("_T", bound=type) # --------------------------------------------------------------------------- # Module-level state # --------------------------------------------------------------------------- # # Registration happens in two phases: # # Phase 1 (_register_type_without_fields) # Allocates a C-level type index and inserts the class into the # global type registry. This must happen early so that self- # referential and mutually-referential annotations can resolve # the class via ``TypeSchema.from_annotation()``. Phase 1 always # succeeds (or raises immediately for non-Object parents). # # Phase 2 (_register_fields_into_type) # Resolves string annotations via ``typing.get_type_hints``, # converts them to ``TypeSchema`` / ``Field`` objects, validates # field ordering, registers fields with the Cython layer, and # installs ``__init__``, ``__repr__``, ``__eq__``, etc. # # If ``get_type_hints`` raises ``NameError`` (forward reference # not yet defined), the class is added to ``_PENDING_CLASSES`` # and retried after each successful phase-2. If phase-2 fails # for any other reason, ``_rollback_registration`` undoes phase-1 # so the type key can be reused. # --------------------------------------------------------------------------- @dataclass class _PendingClass: """Bookkeeping for a class whose annotations couldn't be resolved yet.""" cls: type type_info: Any # core.TypeInfo globalns: dict[str, Any] params: dict[str, Any] #: Classes whose phase-2 (field registration) was deferred because #: ``typing.get_type_hints`` raised ``NameError`` on an unresolved #: forward reference. Retried after each successful phase-2 via #: :func:`_flush_pending`. _PENDING_CLASSES: list[_PendingClass] = [] #: Per-module mapping of ``class.__name__ → class`` for every #: ``@py_class``-decorated type. Used as *localns* when resolving #: annotations so that mutual references between classes in the same #: module work even before the second class is assigned to the module #: variable by Python. _PY_CLASS_BY_MODULE: dict[str, dict[str, type]] = {} # --------------------------------------------------------------------------- # Phase 1: type registration # --------------------------------------------------------------------------- def _register_type_without_fields(cls: type, type_key: str | None) -> Any: """Phase 1: allocate type index and register the type (always succeeds).""" parent_info: core.TypeInfo | None = None for base in cls.__bases__: parent_info = core._type_cls_to_type_info(base) if parent_info is None: parent_info = getattr(base, "__tvm_ffi_type_info__", None) if parent_info is not None: break if parent_info is None: raise TypeError( f"{cls.__name__} must inherit from a registered FFI Object type (e.g. tvm_ffi.Object)" ) if type_key is None: type_key = f"{cls.__module__}.{cls.__qualname__}" info = core._register_py_class(parent_info, type_key, cls) setattr(cls, "__tvm_ffi_type_info__", info) # Register in resolution namespace so sibling classes can find us _PY_CLASS_BY_MODULE.setdefault(cls.__module__, {})[cls.__name__] = cls return info def _rollback_registration(cls: type, type_info: Any) -> None: """Undo :func:`_register_type_without_fields` after a phase-2 failure. The C-level type index is permanently consumed (cannot be reclaimed), but the Python-level registry dicts are cleaned up so a retry with the same type key does not hit "already registered". """ # Remove from the Cython-level registry dicts (TYPE_KEY_TO_INFO, # TYPE_CLS_TO_INFO, TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS). core._rollback_py_class(type_info) # ty: ignore[unresolved-attribute] # Remove from our own module-level resolution namespace. _PY_CLASS_BY_MODULE.get(cls.__module__, {}).pop(cls.__name__, None) for attr in ("__tvm_ffi_type_info__", "__tvm_ffi_is_dataclass__"): try: delattr(cls, attr) except AttributeError: pass # --------------------------------------------------------------------------- # Phase 2: annotation resolution, field registration, dunder installation # --------------------------------------------------------------------------- def _own_annotations(cls: type) -> dict[str, Any]: """Return annotations declared directly on *cls*.""" # Python 3.14+ (PEP 749): annotations are lazily evaluated via # __annotate__ and no longer stored directly in __dict__. getattr() # triggers evaluation and returns per-class annotations correctly. # On Python < 3.14, getattr() follows MRO and returns *parent* # annotations when the child has none — use __dict__ to avoid that. if sys.version_info >= (3, 14): return getattr(cls, "__annotations__", {}) return cls.__dict__.get("__annotations__", {}) def _field_owner_classes(cls: type) -> list[type]: """Classes whose annotations become this type's own fields.""" registered_parent = next( (b for b in cls.__mro__[1:] if "__tvm_ffi_type_info__" in b.__dict__), object ) represented = set(registered_parent.__mro__) return [ b for b in reversed(cls.__mro__) if b is not object and b not in represented and _own_annotations(b) ] def _collect_own_fields( # noqa: PLR0912 cls: type, owner: type, hints: dict[str, Any], kw_only_active: bool, decorator_frozen: bool, ) -> tuple[list[Field], bool]: """Parse own annotations into :class:`Field` objects. - Skips ``ClassVar`` annotations. - Handles ``KW_ONLY`` sentinel. - Extracts ``Field`` metadata from class attributes (set via :func:`field`). - Handles bare defaults (non-``Field`` values). - Converts resolved types to ``TypeSchema``. - Resolves ``hash=None`` to follow ``compare``. """ fields: list[Field] = [] own_annotations = _own_annotations(owner) for name in own_annotations: resolved_type = hints.get(name) # Skip ClassVar if ( resolved_type is None or resolved_type is ClassVar or typing.get_origin(resolved_type) is ClassVar ): continue # KW_ONLY sentinel if resolved_type is KW_ONLY: kw_only_active = True if owner is cls and name in cls.__dict__: try: delattr(cls, name) except AttributeError: pass continue # Extract Field from class dict (inline of _pop_field_from_class) class_val = owner.__dict__.get(name, MISSING) if isinstance(class_val, Field): f = class_val if owner is cls else copy(class_val) elif class_val is not MISSING: f = field(default=class_val) else: f = field() if owner is cls and class_val is not MISSING: try: delattr(cls, name) except AttributeError: pass # Fill in name, schema, and resolved type (set by the decorator, not the user) f.name = name f._ty_schema = TypeSchema.from_annotation(resolved_type) f.type = resolved_type # Resolve kw_only: None means "inherit from decorator" if f.kw_only is None: f.kw_only = kw_only_active # Apply class-level frozen when the field doesn't explicitly set it if decorator_frozen and not f.frozen: f.frozen = True # Resolve hash=None → follow compare (native dataclass semantics) if f.hash is None: f.hash = f.compare fields.append(f) return fields, kw_only_active def method(fn: Any) -> Any: """Mark a ``@py_class`` method for FFI TypeMethod registration. Decorate any staticmethod or plain instance method on a ``@py_class`` body to have it land in the C-level ``TVMFFITypeInfo.methods[]`` table. Once registered, the method is resolvable by name from any FFI consumer — Python-side reflection via ``TypeInfo.methods``, C++, Rust — through the same path already used by C++-defined methods declared via ``refl::ObjectDef().def(...)``. Example:: from tvm_ffi import Object, method from tvm_ffi.dataclasses import py_class @py_class("example.Node") class Node(Object): x: int @method def label(self) -> str: return f"N({self.x})" # The method is now in ``TypeInfo.methods`` and FFI-callable: info = Node.__tvm_ffi_type_info__ fn = next(m.func for m in info.methods if m.name == "label") fn(Node(x=7)) # -> "N(7)" ``staticmethod`` is supported: the marker is written onto the underlying function and unwrapped at registration time. Plain functions are also accepted — the marker lives on the function object directly. ``classmethod`` is rejected at decoration time because its ``cls``-first dispatch does not match the packed-call convention. """ if isinstance(fn, staticmethod): fn.__func__.__ffi_method__ = True return fn if isinstance(fn, classmethod): raise TypeError( "@tvm_ffi.method: @classmethod is not supported for FFI " "TypeMethod registration — the classmethod's ``cls`` first " "arg does not match the packed-call convention. Use " "@staticmethod or a plain instance method instead.", ) if not callable(fn): raise TypeError( f"@tvm_ffi.method: expected a callable, got {type(fn).__name__}.", ) fn.__ffi_method__ = True return fn def _is_method_marked(value: Any) -> bool: """Return True when ``value`` is a callable marked by :func:`method`.""" if isinstance(value, (staticmethod, classmethod)): return getattr(value.__func__, "__ffi_method__", False) is True if callable(value): return getattr(value, "__ffi_method__", False) is True return False def _validate_method_name(cls: type, name: str) -> None: """Reject ``@method``-marked names that collide with reserved namespaces. Names in :data:`_FFI_TYPE_ATTR_NAMES` and Python-protocol dunders are not allowed for ``@method`` — they are routed through the TypeAttrColumn / Python-protocol paths instead. """ if name in _FFI_TYPE_ATTR_NAMES: raise NameError( f"@py_class({cls.__name__!r}): {name!r} is a TypeAttrColumn " "name — define it directly on the class body without " "``@method``; the FFI system routes it to ``TVMFFITypeRegisterAttr``.", ) if name.startswith("__ffi_"): raise NameError( f"@py_class({cls.__name__!r}): {name!r} starts with the " "reserved ``__ffi_`` prefix. Pick a different name for your " "``@method``-decorated method.", ) if name.startswith("__") and name.endswith("__"): raise NameError( f"@py_class({cls.__name__!r}): {name!r} is a Python protocol " "dunder — these are reserved for Python semantics and cannot " "be registered as FFI TypeMethods.", ) def _collect_py_methods(cls: type) -> list[tuple[str, Any, bool]] | None: """Extract FFI-registered entries from a ``@py_class`` body. Two sources are collected: 1. **TypeAttrColumn dunders** — names in :data:`_FFI_RECOGNIZED_METHODS` that appear in ``cls.__dict__``. Both callables (e.g. ``__ffi_repr__``) and non-callable values flow here; the Cython layer routes them to ``TVMFFITypeRegisterAttr`` based on name. 2. **User TypeMethods** — every callable in ``cls.__dict__`` marked with :func:`method`. Registered via ``TVMFFITypeRegisterMethod`` so the method is resolvable by name from any FFI consumer (introspection through ``TypeInfo.methods``, name-based lookup from C++ / Rust, etc.). The decorator pattern keeps the per-class declaration co-located with the method body; no separate allowlist. Validation runs at registration time — reserved ``__ffi_*`` names and Python protocol dunders cannot be ``@method``-decorated; those are reserved by the TypeAttrColumn and Python semantics respectively. Returns the ``(name, value, is_static)`` list, or :data:`None` when no entries were found. """ methods: list[tuple[str, Any, bool]] = [] for name, value in cls.__dict__.items(): marked = _is_method_marked(value) if name not in _FFI_RECOGNIZED_METHODS and not marked: continue if marked: _validate_method_name(cls, name) # In every case, registering a classmethod as a TypeMethod is # wrong: the packed-call convention places ``self`` (an instance) # in slot 0, but classmethod's descriptor binds slot 0 to the # class. if isinstance(value, classmethod): raise TypeError( f"@py_class({cls.__name__!r}): {name!r} is wrapped by " "@classmethod, which is incompatible with FFI " "registration — the cls-first arg breaks the packed-call " "convention. Use @staticmethod or a plain instance " "method. If you wrote ``@classmethod @method``, swap to " "``@staticmethod @method`` (or drop @classmethod).", ) is_static = isinstance(value, staticmethod) func = value.__func__ if is_static else value methods.append((name, func, is_static)) return methods if methods else None def _build_localns(cls: type, *, cross_module: bool = False) -> dict[str, Any]: """Build the localns dict for resolving ``cls``'s annotations. By default, includes only classes from ``cls.__module__``, preserving standard Python name resolution semantics. When ``cross_module=True``, also includes classes from all other registered modules as a fallback — this is needed when ``cls`` has a forward reference to a class in another module that can't appear in ``cls.__module__``'s globals due to a circular import (e.g. the target is imported only under ``if TYPE_CHECKING:``). Cross-module entries are added with ``setdefault`` so same-module classes and the class itself always take precedence over foreign classes with the same ``__name__``. """ localns = dict(_PY_CLASS_BY_MODULE.get(cls.__module__, {})) localns[cls.__name__] = cls if cross_module: for mod_name, mod_classes in list(_PY_CLASS_BY_MODULE.items()): if mod_name == cls.__module__: continue for name, klass in mod_classes.items(): localns.setdefault(name, klass) return localns def _register_fields_into_type( cls: type, type_info: Any, globalns: dict[str, Any], params: dict[str, Any], ) -> bool: """Phase 2: resolve annotations, register fields, install dunders. Returns True on success, False if forward references are unresolved. """ # Resolve string annotations to types; return False (defer) on NameError. # # First try with module-scoped localns (standard Python name resolution). # On NameError, retry with a cross-module localns that includes classes # from every registered module — this handles circular imports where the # target of a forward reference is imported only under TYPE_CHECKING and # therefore never enters the declaring module's globals. kwargs: dict[str, Any] = {"globalns": globalns, "localns": _build_localns(cls)} if sys.version_info >= (3, 11): kwargs["include_extras"] = True try: hints = typing.get_type_hints(cls, **kwargs) except (NameError, AttributeError): kwargs["localns"] = _build_localns(cls, cross_module=True) try: hints = typing.get_type_hints(cls, **kwargs) except (NameError, AttributeError): return False fields_map: dict[str, Field] = {} kw_only_active = params["kw_only"] for owner in _field_owner_classes(cls): owner_fields, kw_only_active = _collect_own_fields( cls, owner, hints, kw_only_active, params["frozen"], ) for f in owner_fields: assert f.name is not None fields_map[f.name] = f own_fields = list(fields_map.values()) py_methods = _collect_py_methods(cls) # Register fields and type-level structural eq/hash kind with the C layer. structure_kind = _STRUCTURE_KIND_MAP.get(params.get("structural_eq")) type_info._register_fields(own_fields, structure_kind) # Attach the user's Field sentinel to each TypeField so the # ``tvm_ffi.dataclasses.fields()`` compat layer can recover defaults # and default_factory values. _register_fields preserves order, so # own_fields and type_info.fields line up 1:1. for py_field, type_field in zip(own_fields, type_info.fields): type_field.dataclass_field = py_field # Register user-defined dunder methods and read back system-generated ones. # Non-callable entries whose names are in _FFI_TYPE_ATTR_NAMES are routed # to TVMFFITypeRegisterAttr by the Cython layer. type_info._register_py_methods(py_methods, type_attr_names=_FFI_TYPE_ATTR_NAMES) _add_class_attrs(cls, type_info) # Remove deferred __init__ and restore user-defined __init__ if saved if "_py_class_deferred_init" in cls.__dict__: # Always remove the deferred wrapper if "__init__" in cls.__dict__: delattr(cls, "__init__") try: delattr(cls, "_py_class_deferred_init") except AttributeError: pass # Restore user-defined __init__ if it was saved user_init = cls.__dict__.get("_py_class_user_init") if user_init is not None: cls.__init__ = user_init delattr(cls, "_py_class_user_init") _install_dataclass_dunders( cls, init=params["init"], repr=params["repr"], eq=params["eq"], order=params["order"], unsafe_hash=params["unsafe_hash"], match_args=params["match_args"], py_class_mode=True, ) return True # --------------------------------------------------------------------------- # Deferred resolution (when phase 2 cannot run at decoration time) # --------------------------------------------------------------------------- def _flush_pending() -> None: """Retry all pending classes. Called after each successful phase 2.""" changed = True while changed: changed = False remaining: list[_PendingClass] = [] for entry in _PENDING_CLASSES: if _register_fields_into_type(entry.cls, entry.type_info, entry.globalns, entry.params): changed = True else: remaining.append(entry) _PENDING_CLASSES[:] = remaining def _raise_unresolved_forward_reference(cls: type, globalns: dict[str, Any]) -> None: """Raise :class:`TypeError` listing the annotations that cannot be resolved.""" localns = _build_localns(cls, cross_module=True) owners = _field_owner_classes(cls) localns.update({owner.__name__: owner for owner in owners}) unresolved: list[str] = [] for owner in owners: for name, ann_str in _own_annotations(owner).items(): if isinstance(ann_str, str): try: eval(ann_str, globalns, localns) except NameError: unresolved.append(f"{name}: {ann_str}") raise TypeError( f"Cannot instantiate {cls.__name__}: unresolved forward references: {unresolved}" ) def _make_temporary_init( cls: type, type_info: Any, globalns: dict[str, Any], params: dict[str, Any] ) -> Callable[[...], None]: def __init__(self: Any, *args: Any, **kwargs: Any) -> None: if type_info.fields is None: try: if not _register_fields_into_type(cls, type_info, globalns, params): _raise_unresolved_forward_reference(cls, globalns) # cls stays in _PENDING_CLASSES after phase-2 succeeds; drop it # before _flush_pending so the loop doesn't hit the Cython-level # "_register_fields already called" assertion on a second pass. _PENDING_CLASSES[:] = [p for p in _PENDING_CLASSES if p.cls is not cls] _flush_pending() except Exception: # Remove from pending list and roll back so the type key can be reused. _PENDING_CLASSES[:] = [p for p in _PENDING_CLASSES if p.cls is not cls] _rollback_registration(cls, type_info) raise # cls.__init__ has been replaced by the real init (or restored user init) cls.__init__(self, *args, **kwargs) __init__.__qualname__ = f"{cls.__qualname__}.__init__" __init__.__module__ = cls.__module__ return __init__ def _install_deferred_init( cls: type, type_info: Any, globalns: dict[str, Any], params: dict[str, Any], ) -> None: """Install a temporary ``__init__`` that completes registration on first call. Preserves a user-defined ``__init__`` if present in *cls.__dict__*; it is restored by :func:`_register_fields_into_type` after registration completes so that ``_install_dataclass_dunders`` sees it and skips auto-generation. """ # Save user-defined __init__ before overwriting user_init = cls.__dict__.get("__init__") if user_init is not None: cls._py_class_user_init = user_init # type: ignore[attr-defined] cls.__init__ = _make_temporary_init(cls, type_info, globalns, params) # type: ignore[assignment] cls._py_class_deferred_init = True # type: ignore[attr-defined] # --------------------------------------------------------------------------- # Main decorator # --------------------------------------------------------------------------- #: Mapping from Python string names to C-level ``TVMFFISEqHashKind`` enum values. _STRUCTURE_KIND_MAP: dict[str | None, int] = { None: 0, # kTVMFFISEqHashKindUnsupported (default; no metadata registered) "tree": 1, # kTVMFFISEqHashKindTreeNode "var": 2, # kTVMFFISEqHashKindFreeVar "dag": 3, # kTVMFFISEqHashKindDAGNode "const-tree": 4, # kTVMFFISEqHashKindConstTreeNode "singleton": 5, # kTVMFFISEqHashKindUniqueInstance } #: Names that should be registered as TypeAttrColumn entries (for C++ #: dispatch via ``TypeAttrColumn``), NOT as TypeMethod. #: #: See ``reflection::type_attr`` in ``accessor.h`` for the C++ constants. _FFI_TYPE_ATTR_NAMES: frozenset[str] = frozenset( { "__ffi_repr__", "__ffi_hash__", "__ffi_eq__", "__ffi_compare__", "__ffi_convert__", "__any_hash__", "__any_equal__", "__s_equal__", "__s_hash__", "__data_to_json__", "__data_from_json__", } ) #: Allowlist of dunder names that ``_collect_py_methods`` collects from #: the class body. Names in ``_FFI_TYPE_ATTR_NAMES`` are registered as #: TypeAttrColumn entries; all other names are registered as TypeMethod. #: #: System-managed names (``__ffi_new__``, ``__ffi_init__``, #: ``__ffi_shallow_copy__``) are intentionally #: absent because the C++ runtime generates them. _FFI_RECOGNIZED_METHODS: frozenset[str] = _FFI_TYPE_ATTR_NAMES @dataclass_transform( eq_default=False, order_default=False, field_specifiers=(field, Field), ) def py_class( # noqa: PLR0913 cls_or_type_key: type | str | None = None, /, *, type_key: str | None = None, frozen: bool = False, init: bool = True, repr: bool = True, eq: bool = False, order: bool = False, unsafe_hash: bool = False, match_args: bool = True, kw_only: bool = False, structural_eq: str | None = None, slots: bool = True, ) -> Callable[[_T], _T] | _T: """Register a Python-defined FFI class with dataclass-style semantics. Can be used as: .. code-block:: python @py_class # bare decorator class Point(Object): x: float y: float @py_class("my.Point") # with explicit type_key class Point(Object): ... @py_class(eq=True, order=True) # with options class Point(Object): ... @py_class("my.Point", eq=True) # both class Point(Object): ... @py_class(structural_eq="tree") # structural eq/hash kind class MyNode(Object): value: int span: Object = field(structural_eq="ignore") Parameters ---------- cls_or_type_key When a string, used as the FFI type key. When a type (bare decorator usage), the class to decorate. type_key Explicit FFI type key. Auto-generated from ``{module}.{qualname}`` when omitted. frozen If True, all fields are read-only after ``__init__`` by default. Individual fields can still be marked ``field(frozen=True)`` on a non-frozen class. Use ``type(obj).field_name.set(obj, value)`` as an escape hatch when mutation is necessary. init If True (default), generate ``__init__`` from field annotations. repr If True (default), generate ``__repr__``. eq If True, generate ``__eq__`` and ``__ne__`` using recursive field-wise content equality. Default False, in which case the class inherits the pointer-based ``__eq__`` from ``Object`` (``a == b`` is equivalent to ``a.same_as(b)``). If the class body defines ``__eq__`` or ``__ne__``, the generator is skipped and the user definition is preserved. order If True, generate ``__lt__``, ``__le__``, ``__gt__``, ``__ge__``. Requires ``eq=True``. unsafe_hash If True, generate ``__hash__`` using recursive field-wise content hashing (unsafe for mutable objects). Default False, in which case the class inherits the handle-address ``__hash__`` from ``Object``. A user-defined ``__hash__`` in the class body is preserved. match_args If True (default), set ``__match_args__`` to a tuple of the positional ``__init__`` field names (``init=True`` and not ``kw_only``), enabling ``match`` statements. Ignored when the class body already defines ``__match_args__``. kw_only If True, all fields are keyword-only in ``__init__`` by default. structural_eq Structural equality/hashing kind for this type. Controls how instances participate in ``structural_equal`` and ``structural_hash``. Valid values are: - ``None`` (default): structural comparison is not supported. - ``"tree"``: content-based comparison, the safe default for most IR nodes. - ``"var"``: compared by binding position, for variable types. - ``"dag"``: content + sharing-aware comparison, for dataflow graph nodes. - ``"const-tree"``: like ``"tree"`` with a pointer-equality fast path (only safe for types with no transitive ``"var"`` children). - ``"singleton"``: pointer equality only, for singleton types. This parameter is **independent** of ``eq`` / ``unsafe_hash``: it only configures how ``structural_equal`` / ``structural_hash`` walk the object in C++ and never installs or alters Python-level ``__eq__`` / ``__hash__``. See Notes below. slots Accepted for ``dataclass_transform`` compatibility. Object subclasses always use ``__slots__ = ()`` via the metaclass. Returns ------- Callable | type A class decorator, or the decorated class (bare usage). Notes ----- Three orthogonal equality/hashing mechanisms coexist on a ``@py_class`` type, each controlled by an independent knob: - ``a == b`` / ``hash(a)`` — selected by ``eq`` / ``unsafe_hash`` params (or user-defined ``__eq__`` / ``__hash__`` in the class body). Default: pointer-based ``same_as`` and handle-address hash, inherited from ``Object``. - ``structural_equal(a, b)`` / ``structural_hash(a)`` — selected by the ``structural_eq`` param. Default (``None``): structural comparison is unsupported for this type. - ``a.same_as(b)`` — always available; always pointer comparison. The typical pattern is to leave ``eq`` / ``unsafe_hash`` at their defaults so ``==`` and ``hash()`` stay cheap and pointer-based (ideal for pass-internal bookkeeping such as visited-set tracking), and call ``structural_equal`` / ``structural_hash`` explicitly at the points that require the heavy semantic check. Combining ``eq=True`` (or ``unsafe_hash=True``) with a ``structural_eq`` kind is legal but gives the type two different recursive equalities — a Python-level one for ``==`` and a C++ structural one for ``structural_equal`` — which rarely coincide. Prefer setting only one. """ if order and not eq: raise ValueError("order=True requires eq=True") if structural_eq not in _STRUCTURE_KIND_MAP: raise ValueError( f"structural_eq must be one of " f"{sorted(k for k in _STRUCTURE_KIND_MAP if k is not None)}" f" or None, got {structural_eq!r}" ) effective_type_key = type_key params: dict[str, Any] = { "frozen": frozen, "init": init, "repr": repr, "eq": eq, "order": order, "unsafe_hash": unsafe_hash, "match_args": match_args, "kw_only": kw_only, "structural_eq": structural_eq, } def decorator(cls: _T) -> _T: nonlocal effective_type_key globalns = getattr(sys.modules.get(cls.__module__, None), "__dict__", {}) info = _register_type_without_fields(cls, effective_type_key) try: if _register_fields_into_type(cls, info, globalns, params): _flush_pending() else: _PENDING_CLASSES.append(_PendingClass(cls, info, globalns, params)) _install_deferred_init(cls, info, globalns, params) except Exception: # Phase-2 failed (bad annotation, field ordering, etc.). # Roll back phase-1 so the type key can be reused after # the user fixes the error. _rollback_registration(cls, info) raise # Marker: distinguishes @c_class / @py_class types from FFI containers # (Array, List, Map, Dict) that also have __tvm_ffi_type_info__ but are # not dataclasses. Used by is_dataclass() in common.py. setattr(cls, "__tvm_ffi_is_dataclass__", True) return cls # Handle different calling conventions: # @py_class → cls_or_type_key is the class # @py_class("key") → cls_or_type_key is a string # @py_class() → cls_or_type_key is None # @py_class(eq=True) → cls_or_type_key is None if cls_or_type_key is None: return decorator if isinstance(cls_or_type_key, str): effective_type_key = cls_or_type_key return decorator if isinstance(cls_or_type_key, type): return decorator(cls_or_type_key) raise TypeError(f"py_class: expected str or type, got {type(cls_or_type_key)}") tvm-ffi-0.1.12/python/tvm_ffi/error.py000066400000000000000000000226721521067262500176400ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name """Error handling.""" from __future__ import annotations import ast import re import sys import types from typing import Any from . import core def _parse_backtrace(backtrace: str) -> list[tuple[str, int, str]]: """Parse the backtrace string into a list of (filename, lineno, func). Parameters ---------- backtrace The backtrace string. Returns ------- result The list of (filename, lineno, func) """ pattern = r'File "(.+?)", line (\d+), in (.+)' result = [] for line in backtrace.split("\n"): match = re.match(pattern, line.strip()) if match: try: filename = match.group(1) lineno = int(match.group(2)) func = match.group(3) result.append((filename, lineno, func)) except ValueError: pass return result class TracebackManager: """Helper to manage traceback generation.""" def __init__(self) -> None: """Initialize the traceback manager and its cache.""" self._code_cache: dict[tuple[str, int, str], types.CodeType] = {} def _get_cached_code_object(self, filename: str, lineno: int, func: str) -> types.CodeType: # Hack to create a code object that points to the correct # line number and function name key = (filename, lineno, func) # cache the code object to avoid re-creating it if key in self._code_cache: return self._code_cache[key] # Parse to AST and zero out column info # since column info are not accurate in original trace tree = ast.parse("_getframe()", filename=filename, mode="eval") for node in ast.walk(tree): if hasattr(node, "col_offset"): node.col_offset = 0 # ty: ignore[invalid-assignment] if hasattr(node, "end_col_offset"): node.end_col_offset = 0 # ty: ignore[invalid-assignment] # call into get frame, bt changes the context code_object = compile(tree, filename, "eval") # replace the function name and line number code_object = code_object.replace(co_name=func, co_firstlineno=lineno) self._code_cache[key] = code_object return code_object def _create_frame(self, filename: str, lineno: int, func: str) -> types.FrameType: """Create a frame object from the filename, lineno, and func.""" code_object = self._get_cached_code_object(filename, lineno, func) # call into get frame, but changes the context so the code # points to the correct frame context = {"_getframe": sys._getframe} # pylint: disable=eval-used return eval(code_object, context, context) def append_traceback( self, tb: types.TracebackType | None, filename: str, lineno: int, func: str, ) -> types.TracebackType: """Append a traceback to the given traceback. Parameters ---------- tb The traceback to append to. filename The filename of the traceback lineno The line number of the traceback func The function name of the traceback Returns ------- new_tb The new traceback with the appended frame. """ # This approach avoids binding the created frame object to a local variable # in `append_traceback`, which would create a reference cycle. By using a # nested function, the frame object is a temporary that is not held by # the locals of `append_traceback`. See the diagram in `_with_append_backtrace` # and PR #327 for more details. def create( tb: types.TracebackType | None, frame: types.FrameType, lineno: int ) -> types.TracebackType: return types.TracebackType(tb, frame, frame.f_lasti, lineno) return create(tb, self._create_frame(filename, lineno, func), lineno) _TRACEBACK_MANAGER = TracebackManager() def _with_append_backtrace(py_error: BaseException, backtrace: str) -> BaseException: """Append the backtrace to the py_error and return it.""" # We manually delete py_error and tb to avoid reference cycle, making it faster to gc the locals inside the frame # please see pull request #327 for more details # # Memory Cycle Diagram: # # [Stack Frames] [Heap Objects] # +-------------------+ # | outside functions | -----------------------> [ Tensor ] # +-------------------+ (Held by cycle, slow to free) # ^ # | f_back # +-------------------+ locals py_error # | py_error (this) | -----+--------------> [ BaseException ] # +-------------------+ | | # ^ | | (with_traceback) # | f_back | v # +-------------------+ +--------------> [ Traceback Obj ] # | append_traceback | tb | # +-------------------+ | # ^ | # | f_back | # +-------------------+ | # | _create_frame | | # +-------------------+ | # ^ | # | f_back | # +-------------------+ | # | _get_frame | <----------------------------+ # +-------------------+ (Cycle closes here) tb = py_error.__traceback__ try: for filename, lineno, func in _parse_backtrace(backtrace): tb = _TRACEBACK_MANAGER.append_traceback(tb, filename, lineno, func) return py_error.with_traceback(tb) finally: # We explicitly break the reference cycle here. The `finally` block is # executed just before the function returns, after the `return` expression # in the `try` block has been evaluated. Deleting `py_error` and `tb` # here ensures they are not held by this function's frame's locals, # which resolves the cycle. del py_error, tb def _traceback_to_backtrace_str(tb: types.TracebackType | None) -> str: """Convert the traceback to a string.""" lines = [] while tb is not None: frame = tb.tb_frame lineno = tb.tb_lineno filename = frame.f_code.co_filename funcname = frame.f_code.co_name lines.append(f' File "{filename}", line {lineno}, in {funcname}\n') tb = tb.tb_next # needs to reverse the order of the lines so backtrace stores in # the reverse order of python traceback return "".join(reversed(lines)) core._WITH_APPEND_BACKTRACE = _with_append_backtrace core._TRACEBACK_TO_BACKTRACE_STR = _traceback_to_backtrace_str def register_error( name_or_cls: str | type | None = None, cls: type | None = None, ) -> Any: """Register an error class so it can be recognized by the ffi error handler. Parameters ---------- name_or_cls The name of the error class. cls The class to register. Returns ------- fregister Register function if f is not specified. Examples -------- .. code-block:: python import tvm_ffi # Register a custom Python exception so tvm_ffi.Error maps to it @tvm_ffi.error.register_error class MyError(RuntimeError): pass # Convert a Python exception to an FFI Error and back ffi_err = tvm_ffi.convert(MyError("boom")) py_err = ffi_err.py_error() assert isinstance(py_err, MyError) """ if isinstance(name_or_cls, type): cls = name_or_cls name_or_cls = cls.__name__ def register(mycls: type) -> type: """Register the error class name with the FFI core.""" err_name = name_or_cls if isinstance(name_or_cls, str) else mycls.__name__ core.ERROR_NAME_TO_TYPE[err_name] = mycls core.ERROR_TYPE_TO_NAME[mycls] = err_name return mycls if cls is None: return register return register(cls) register_error("RuntimeError", RuntimeError) register_error("ValueError", ValueError) register_error("TypeError", TypeError) register_error("AttributeError", AttributeError) register_error("KeyError", KeyError) register_error("IndexError", IndexError) register_error("AssertionError", AssertionError) register_error("MemoryError", MemoryError) tvm-ffi-0.1.12/python/tvm_ffi/libinfo.py000066400000000000000000000324231521067262500201240ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities to locate tvm_ffi libraries, headers, and helper include paths. This module also provides helpers to locate and load platform-specific shared libraries by a target_name (e.g., ``tvm_ffi`` -> ``libtvm_ffi.so`` on Linux). """ from __future__ import annotations import ctypes import importlib.metadata as im import os import sys from pathlib import Path from typing import TYPE_CHECKING, Callable if TYPE_CHECKING: from .module import Module def find_libtvm_ffi() -> str: """Find libtvm_ffi. Returns ------- path The full path to the located library. """ candidate = _find_library_by_basename("apache-tvm-ffi", "tvm_ffi") if ret := _resolve_and_validate([candidate], cond=lambda _: True): return ret raise RuntimeError("Cannot find libtvm_ffi") def find_windows_implib() -> str: """Find and return the Windows import library path for tvm_ffi.lib.""" # implib = _find_library_by_basename("apache-tvm-ffi", "tvm_ffi").parent / "tvm_ffi.lib" # ret = _resolve_to_str(implib) candidate = _find_library_by_basename("apache-tvm-ffi", "tvm_ffi").parent / "tvm_ffi.lib" if ret := _resolve_and_validate([candidate], cond=lambda _: True): return ret raise RuntimeError("Cannot find implib tvm_ffi.lib") def find_source_path() -> str: """Find packaged source home path.""" if ret := _resolve_and_validate( paths=[ _rel_top_directory(), _dev_top_directory(), ], cond=lambda p: (p / "src").is_dir(), ): return ret raise RuntimeError("Cannot find home path.") def find_cmake_path() -> str: """Find the preferred cmake path.""" if ret := _resolve_and_validate( paths=[ _rel_top_directory() / "share" / "cmake" / "tvm_ffi", # Standard install _dev_top_directory() / "cmake", # Development mode ], cond=lambda p: p.is_dir(), ): return ret raise RuntimeError("Cannot find cmake path.") def find_include_path() -> str: """Find header files for C compilation.""" if ret := _resolve_and_validate( paths=[ _rel_top_directory() / "include", _dev_top_directory() / "include", ], cond=lambda p: p.is_dir(), ): return ret raise RuntimeError("Cannot find include path.") def find_dlpack_include_path() -> str: """Find dlpack header files for C compilation.""" if ret := _resolve_and_validate( paths=[ _rel_top_directory() / "include", _dev_top_directory() / "3rdparty" / "dlpack" / "include", ], cond=lambda p: (p / "dlpack").is_dir(), ): return ret raise RuntimeError("Cannot find dlpack include path.") def find_cython_lib() -> str: """Find the path to tvm cython.""" from tvm_ffi import core # noqa: PLC0415 try: return str(Path(core.__file__).resolve()) except OSError: pass raise RuntimeError("Cannot find tvm cython path.") def find_python_helper_include_path() -> str: """Find header files for C compilation.""" if ret := _resolve_and_validate( paths=[ _rel_top_directory() / "include", _dev_top_directory() / "python" / "tvm_ffi" / "cython", ], cond=lambda p: (p / "tvm_ffi_python_helpers.h").is_file(), ): return ret raise RuntimeError("Cannot find python helper include path.") def include_paths() -> list[str]: """Find all include paths needed for FFI related compilation.""" return sorted( { find_include_path(), find_dlpack_include_path(), find_python_helper_include_path(), } ) def load_lib_ctypes( package: str, target_name: str, mode: str, extra_lib_paths: list[Path] | None = None, ) -> ctypes.CDLL: """Load the tvm_ffi shared library by searching likely paths. Parameters ---------- package The package name where the library is expected to be found. For example, ``"apache-tvm-ffi"`` is the package name of `tvm-ffi`. target_name Name of the CMake target, e.g., ``"tvm_ffi"``. It is used to derive the platform-specific shared library name, e.g., ``"libtvm_ffi.so"`` on Linux, ``"tvm_ffi.dll"`` on Windows. mode The mode to load the shared library. See `ctypes.${MODE}` for details. Usually it is either ``"RTLD_LOCAL"`` or ``"RTLD_GLOBAL"``. extra_lib_paths Optional list of additional directories to search for the shared library before falling back to the built-in dev candidates and PATH-derived dirs. Useful when the caller's package does not ship a wheel-style ``RECORD`` but knows where its build artifacts land (e.g. ``/build/lib``). Every element must be a :class:`pathlib.Path`. Returns ------- The loaded shared library. """ lib_path: Path = _find_library_by_basename(package, target_name, extra_lib_paths) # The dll search path need to be added explicitly in windows if sys.platform.startswith("win32"): os.add_dll_directory(str(lib_path.parent)) return ctypes.CDLL(str(lib_path), getattr(ctypes, mode)) def load_lib_module( package: str, target_name: str, keep_module_alive: bool = True, extra_lib_paths: list[Path] | None = None, ) -> Module: """Load the tvm_ffi shared library by searching likely paths. Parameters ---------- package The package name where the library is expected to be found. For example, ``"apache-tvm-ffi"`` is the package name of `tvm-ffi`. target_name Name of the CMake target, e.g., ``"tvm_ffi"``. It is used to derive the platform-specific shared library name, e.g., ``"libtvm_ffi.so"`` on Linux, ``"tvm_ffi.dll"`` on Windows. keep_module_alive Whether to keep the loaded module alive to prevent it from being unloaded. extra_lib_paths Optional list of additional directories to search for the shared library before falling back to the built-in dev candidates and PATH-derived dirs. Useful when the caller's package does not ship a wheel-style ``RECORD`` but knows where its build artifacts land (e.g. ``/build/lib``). Every element must be a :class:`pathlib.Path`. Returns ------- The loaded shared library. """ from .module import load_module # noqa: PLC0415 lib_path: Path = _find_library_by_basename(package, target_name, extra_lib_paths) # The dll search path need to be added explicitly in windows if sys.platform.startswith("win32"): os.add_dll_directory(str(lib_path.parent)) return load_module(lib_path, keep_module_alive=keep_module_alive) def _find_library_by_basename( # noqa: PLR0912 package: str, target_name: str, extra_lib_paths: list[Path] | None = None, ) -> Path: """Find a shared library by target_name name across known directories. Parameters ---------- package The package name where the library is expected to be found. target_name Base name (e.g., ``"tvm_ffi"`` or ``"tvm_ffi_testing"``). extra_lib_paths Optional list of additional directories to search **before** the built-in self-anchored fallback dirs and the ``PATH``/ ``LD_LIBRARY_PATH``/``DYLD_LIBRARY_PATH``-derived dirs. Caller-supplied directories take precedence over the built-ins, so a foreign caller (e.g. ``package="tvm"``) can point at its own build tree without relying on ``tvm_ffi``'s file location. Every element must be a :class:`pathlib.Path`. Returns ------- path The full path to the located library. Raises ------ TypeError If any element of ``extra_lib_paths`` is not a :class:`pathlib.Path`. RuntimeError If the library cannot be found in any of the candidate directories. The error message lists every directory searched. """ if extra_lib_paths is not None: for i, p in enumerate(extra_lib_paths): if not isinstance(p, Path): raise TypeError( f"extra_lib_paths[{i}] must be a pathlib.Path, got {type(p).__name__}: {p!r}" ) if sys.platform.startswith("win32"): lib_dll_names = (f"{target_name}.dll",) elif sys.platform.startswith("darwin"): # Prefer dylib, also allow .so for some toolchains lib_dll_names = (f"lib{target_name}.dylib", f"lib{target_name}.so") else: # Linux, FreeBSD, etc lib_dll_names = (f"lib{target_name}.so",) # Use `importlib.metadata` is the most reliable way to find package data files try: dist: im.PathDistribution = im.distribution(package) # ty: ignore[invalid-assignment] record = dist.read_text("RECORD") or "" for line in record.splitlines(): partial_path, *_ = line.split(",") if partial_path.endswith(lib_dll_names): try: path = (dist._path.parent / partial_path).resolve() except OSError: continue if path.name in lib_dll_names: return path except (im.PackageNotFoundError, OSError): # Distribution metadata may be missing when the caller's package is on # PYTHONPATH without a wheel install. Fall through to the dev fallback. pass # **Fallback**. it's possible that the library is not built as part of Python ecosystem, # e.g. Use PYTHONPATH to point to dev package, and CMake + Makefiles to build the shared library. dll_paths: list[Path] = [] # Case 1. Caller-supplied directories take precedence — these let a foreign # caller (e.g. ``package="tvm"``) point at its own build tree without # relying on ``tvm_ffi``'s file location. if extra_lib_paths is not None: dll_paths.extend(extra_lib_paths) # Case 2. Built-in self-anchored fallback (back-compat for the self-call). dll_paths.extend( [ _rel_top_directory() / "build" / "lib", _rel_top_directory() / "lib", _dev_top_directory() / "build" / "lib", _dev_top_directory() / "lib", ] ) # Case 3. PATH-related environment variables. if sys.platform.startswith("win32"): dll_paths.extend(Path(p) for p in _split_env_var("PATH", ";")) elif sys.platform.startswith("darwin"): dll_paths.extend(Path(p) for p in _split_env_var("DYLD_LIBRARY_PATH", ":")) dll_paths.extend(Path(p) for p in _split_env_var("PATH", ":")) else: dll_paths.extend(Path(p) for p in _split_env_var("LD_LIBRARY_PATH", ":")) dll_paths.extend(Path(p) for p in _split_env_var("PATH", ":")) # Search for the library in candidate directories for dll_dir in dll_paths: for lib_dll_name in lib_dll_names: try: path = (dll_dir / lib_dll_name).resolve() if path.is_file(): return path except OSError: continue raise RuntimeError( f"Cannot find library {', '.join(lib_dll_names)}; " f"searched directories:\n " + "\n ".join(str(p) for p in dll_paths) ) def _split_env_var(env_var: str, split: str) -> list[str]: """Split an environment variable string. Parameters ---------- env_var Name of environment variable. split String to split env_var on. Returns ------- splits If env_var exists, split env_var. Otherwise, empty list. """ if os.environ.get(env_var, None): return [p.strip() for p in os.environ[env_var].split(split)] return [] def _rel_top_directory() -> Path: """Get the current directory of this file.""" return Path(__file__).parent def _dev_top_directory() -> Path: """Get the top-level development directory.""" return _rel_top_directory() / ".." / ".." def _resolve_and_validate( paths: list[Path], cond: Callable[[Path], bool | Path], ) -> str | None: """For all paths that resolve properly, find the 1st one that meets the specified condition. M. B. This code path gracefully handles broken paths, symlinks, or permission issues, and is required for robust library discovery in all public APIs in this file. """ for path in paths: try: resolved = path.resolve() ret = cond(resolved) except (OSError, AssertionError): continue if isinstance(ret, Path): return str(ret) elif ret is True: return str(resolved) return None tvm-ffi-0.1.12/python/tvm_ffi/module.py000066400000000000000000000344501521067262500177710ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Module related objects and functions.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Sequence from typing import Any # isort: on # fmt: on # tvm-ffi-stubgen(end) import json from collections.abc import Sequence from enum import IntEnum from os import PathLike, fspath from typing import Any, ClassVar, cast from . import _ffi_api, core from .registry import register_object __all__ = ["Module", "ModulePropertyMask", "load_module", "system_lib"] class ModulePropertyMask(IntEnum): """Runtime Module Property Mask.""" BINARY_SERIALIZABLE = 0b001 RUNNABLE = 0b010 COMPILATION_EXPORTABLE = 0b100 @register_object("ffi.Module") class Module(core.Object): """Module container for dynamically loaded Module. Examples -------- .. code-block:: python import tvm_ffi # Load the module from a shared library mod = tvm_ffi.load_module("path/to/library.so") # Call exported function mod.func_name(*args) # Query function metadata (type signature) metadata = mod.get_function_metadata("func_name") # Query function documentation (if available) doc = mod.get_function_doc("func_name") See Also -------- :py:func:`tvm_ffi.load_module` :py:meth:`get_function_metadata` :py:meth:`get_function_doc` Notes ----- If you load a module within a local scope, be careful when any called function creates and returns an object. The memory deallocation routines are part of the library's code. If the module is unloaded before the object is destroyed, the deleter may call an invalid address. Keep the module loaded until all returned objects are deleted. You can safely use returned objects inside a nested function that finishes before the module goes out of scope. When possible, consider keeping the module alive in a long-lived/global scope (for example, in a global state) to avoid premature unloading. .. code-block:: python def bad_pattern(x): # Bad: unload order of `tensor` and `mod` is not guaranteed mod = tvm_ffi.load_module("path/to/library.so") # ... do something with the tensor tensor = mod.func_create_and_return_tensor(x) def good_pattern(x): # Good: `tensor` is freed before `mod` goes out of scope mod = tvm_ffi.load_module("path/to/library.so") def run_some_tests(): tensor = mod.func_create_and_return_tensor(x) # ... do something with the tensor run_some_tests() """ # tvm-ffi-stubgen(begin): object/ffi.Module # fmt: off imports_: Sequence[Any] # fmt: on # tvm-ffi-stubgen(end) entry_name: ClassVar[str] = "main" # constant for entry function name __slots__ = ("__dict__",) @property def kind(self) -> str: """Get type key of the module.""" return _ffi_api.ModuleGetKind(self) @property def imports(self) -> list[Module]: """Get imported modules. Returns ------- modules The module """ return self.imports_ # ty: ignore[invalid-return-type] def implements_function(self, name: str, query_imports: bool = False) -> bool: """Return True if the module defines a global function. Notes ----- ``implements_function(name)`` does not guarantee that ``get_function(name)`` will return a callable, because some module kinds (e.g. a source-only module) may not provide a packed function implementation until further compilation occurs. However, a non-null result from ``get_function(name)`` should imply the module implements the function. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- True if module (or one of its imports) has a definition for name. """ return _ffi_api.ModuleImplementsFunction(self, name, query_imports) def __getattr__(self, name: str) -> core.Function: """Accessor to allow getting functions as attributes.""" try: func = self.get_function(name) except AttributeError as exc: raise AttributeError(f"Module has no function '{name}'") from exc setattr(self, name, func) return func def get_function(self, name: str, query_imports: bool = False) -> core.Function: """Get function from the module. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- The result function. """ func = _ffi_api.ModuleGetFunction(self, name, query_imports) func = cast(core.Function, func) if func is None: raise AttributeError(f"Module has no function '{name}'") return func def get_function_metadata( self, name: str, query_imports: bool = False ) -> dict[str, Any] | None: """Get metadata for a function exported from the module. This retrieves metadata for functions exported via c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` and when c:macro:`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is on, which includes type schema information. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- metadata A dictionary containing function metadata. The ``type_schema`` field encodes the callable signature. Examples -------- .. code-block:: python import tvm_ffi from tvm_ffi.core import TypeSchema import json mod = tvm_ffi.load_module("add_one_cpu.so") metadata = mod.get_function_metadata("add_one_cpu") schema = TypeSchema.from_json_str(metadata["type_schema"]) print(schema) # Shows function signature See Also -------- :py:func:`tvm_ffi.get_global_func_metadata` Get metadata for global registry functions. """ metadata_str = _ffi_api.ModuleGetFunctionMetadata(self, name, query_imports) if metadata_str is None: return None return json.loads(metadata_str) def get_function_doc(self, name: str, query_imports: bool = False) -> str | None: """Get documentation string for a function exported from the module. This retrieves documentation for functions exported via c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- doc : str or None The documentation string if available, None otherwise. Examples -------- .. code-block:: python import tvm_ffi mod = tvm_ffi.load_module("mylib.so") doc = mod.get_function_doc("process_batch") if doc: print(doc) See Also -------- :py:meth:`get_function_metadata` Get metadata including type schema. """ doc_str = _ffi_api.ModuleGetFunctionDoc(self, name, query_imports) return doc_str if doc_str else None def import_module(self, module: Module) -> None: """Add module to the import list of current one. Parameters ---------- module The other module. """ _ffi_api.ModuleImportModule(self, module) def __getitem__(self, name: str) -> core.Function: """Return function by name using item access (module["func"]).""" if not isinstance(name, str): raise ValueError("Can only take string as function name") return self.get_function(name) def __call__(self, *args: Any) -> Any: """Call the module's entry function (`main`).""" # pylint: disable=not-callable return self.main(*args) def inspect_source(self, fmt: str = "") -> str: """Get source code from module, if available. Parameters ---------- fmt The specified format. Returns ------- The result source code. """ return _ffi_api.ModuleInspectSource(self, fmt) def get_write_formats(self) -> Sequence[str]: """Get the format of the module.""" return _ffi_api.ModuleGetWriteFormats(self) def get_property_mask(self) -> int: """Get the runtime module property mask. The mapping is stated in ModulePropertyMask. Returns ------- Bitmask of runtime module property """ return _ffi_api.ModuleGetPropertyMask(self) def is_binary_serializable(self) -> bool: """Return whether the module is binary serializable (supports save_to_bytes). Returns ------- True if the module is binary serializable. """ return (self.get_property_mask() & ModulePropertyMask.BINARY_SERIALIZABLE) != 0 def is_runnable(self) -> bool: """Return whether the module is runnable (supports get_function). Returns ------- True if the module is runnable. """ return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0 def is_compilation_exportable(self) -> bool: """Return whether the module is compilation exportable. write_to_file is supported for object or source. Returns ------- True if the module is compilation exportable. """ return (self.get_property_mask() & ModulePropertyMask.COMPILATION_EXPORTABLE) != 0 def clear_imports(self) -> None: """Remove all imports of the module.""" _ffi_api.ModuleClearImports(self) def write_to_file(self, file_name: str, fmt: str = "") -> None: """Write the current module to file. Parameters ---------- file_name The name of the file. fmt The format of the file. See Also -------- tvm.runtime.Module.export_library : export the module to shared library. """ _ffi_api.ModuleWriteToFile(self, file_name, fmt) def system_lib(symbol_prefix: str = "") -> Module: """Get system-wide library module singleton with functions prefixed by ``__tvm_ffi_{symbol_prefix}``. The library module contains symbols that are registered via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`. .. note:: The system lib is intended to be statically linked and loaded during the entire lifecycle of the program. If you want dynamic loading features, use DSO modules instead. Parameters ---------- symbol_prefix Optional symbol prefix that can be used for search. When we lookup a symbol symbol_prefix + name will first be searched, then the name without symbol_prefix. Returns ------- The system-wide library module. Examples -------- Register the function ``test_symbol_add_one`` in C++ with the name ``__tvm_ffi_test_symbol_add_one`` via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`. .. code-block:: cpp // A function to be registered in the system lib static int test_symbol_add_one(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* ret) { TVM_FFI_SAFE_CALL_BEGIN(); TVM_FFI_CHECK(num_args == 1, "Expected 1 argument, but got: " + std::to_string(num_args)); int64_t x = reinterpret_cast(args)[0].cast(); reinterpret_cast(ret)[0] = x + 1; TVM_FFI_SAFE_CALL_END(); } // Register the function with name `test_symbol_add_one` prefixed by `__tvm_ffi_` int _ = TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one", reinterpret_cast(test_symbol_add_one)); Look up and call the function from Python: .. code-block:: python import tvm_ffi mod: tvm_ffi.Module = tvm_ffi.system_lib( "testing." ) # symbols prefixed with `__tvm_ffi_testing.` func: tvm_ffi.Function = mod["add_one"] # looks up `__tvm_ffi_testing.add_one` assert func(10) == 11 """ return _ffi_api.SystemLib(symbol_prefix) def load_module(path: str | PathLike, keep_module_alive: bool = True) -> Module: """Load module from file. Parameters ---------- path The path to the module file. keep_module_alive Whether to keep the module alive. If True, the module will be kept alive for the duration of the program until libtvm_ffi.so is unloaded. Returns ------- The loaded module Examples -------- .. code-block:: python import tvm_ffi from pathlib import Path # Works with string paths mod = tvm_ffi.load_module("path/to/module.so") # Also works with pathlib.Path objects mod = tvm_ffi.load_module(Path("path/to/module.so")) mod.func_name(*args) See Also -------- :py:class:`tvm_ffi.Module` """ path = fspath(path) mod = _ffi_api.ModuleLoadFromFile(path) if keep_module_alive: _ffi_api.ModuleGlobalsAdd(mod) return mod tvm-ffi-0.1.12/python/tvm_ffi/py.typed000066400000000000000000000000611521067262500176200ustar00rootroot00000000000000# Marker file for PEP 561. Keep this file empty. tvm-ffi-0.1.12/python/tvm_ffi/registry.py000066400000000000000000000357061521067262500203610ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """FFI registry to register function and objects.""" from __future__ import annotations import json import sys import warnings from typing import Any, Callable, Literal, Sequence, TypeVar, overload from . import core from .core import Function, TypeInfo # whether we simplify skip unknown objects regtistration _SKIP_UNKNOWN_OBJECTS = False _T = TypeVar("_T", bound=type) def register_object( type_key: str | None = None, *, init: bool = True, ) -> Callable[[_T], _T]: """Register object type. Parameters ---------- type_key The type key of the node. It requires ``type_key`` to be registered already on the C++ side. If not specified, the class name will be used. init If True (default), install ``__init__`` from the C++ ``__ffi_init__`` TypeAttrColumn when available, or a TypeError guard for ``Object`` subclasses that lack one. Set to False when a subsequent decorator (e.g. ``@c_class``) will handle ``__init__`` installation. Notes ----- All :class:`Object` subclasses get ``__slots__ = ()`` by default via the metaclass, preventing per-instance ``__dict__``. To opt out and allow arbitrary instance attributes, declare ``__slots__ = ("__dict__",)`` explicitly in the class body:: @tvm_ffi.register_object("test.MyObject") class MyObject(Object): __slots__ = ("__dict__",) Examples -------- The following code registers MyObject using type key "test.MyObject", if the type key is already registered on the C++ side. .. code-block:: python @tvm_ffi.register_object("test.MyObject") class MyObject(Object): pass """ def _register(cls: _T, object_name: str) -> _T: """Register the object type with the FFI core.""" type_index = core._object_type_key_to_index(object_name) if type_index is None: if _SKIP_UNKNOWN_OBJECTS: return cls raise ValueError(f"Cannot find object type index for {object_name}") info = core._register_object_by_index(type_index, cls) _add_class_attrs(type_cls=cls, type_info=info) setattr(cls, "__tvm_ffi_type_info__", info) if init: _install_init(cls, info) return cls if isinstance(type_key, str): def _decorator_with_name(cls: _T) -> _T: return _register(cls, type_key) return _decorator_with_name def _decorator_default(cls: _T) -> _T: return _register(cls, cls.__name__) if type_key is None: return _decorator_default if isinstance(type_key, type): return _decorator_default(type_key) raise TypeError("type_key must be a string, type, or None") def register_global_func( func_name: str | Callable[..., Any], f: Callable[..., Any] | None = None, override: bool = False, ) -> Any: """Register global function. Parameters ---------- func_name The function name f The function to be registered. override Whether override existing entry. Returns ------- fregister Register function if f is not specified. Examples -------- .. code-block:: python import tvm_ffi # we can use decorator to register a function @tvm_ffi.register_global_func("mytest.echo") def echo(x): return x # After registering, we can get the function by its name f = tvm_ffi.get_global_func("mytest.echo") assert f(1) == 1 # we can also directly register a function tvm_ffi.register_global_func("mytest.add_one", lambda x: x + 1) f = tvm_ffi.get_global_func("mytest.add_one") assert f(1) == 2 See Also -------- :py:func:`tvm_ffi.get_global_func` :py:func:`tvm_ffi.remove_global_func` """ if not isinstance(func_name, str): f = func_name func_name = f.__name__ # ty: ignore[unresolved-attribute] if not isinstance(func_name, str): raise ValueError("expect string function name") def register(myf: Callable[..., Any]) -> Any: """Register the global function with the FFI core.""" return core._register_global_func(func_name, myf, override) if f is not None: return register(f) return register @overload def get_global_func(name: str, allow_missing: Literal[True]) -> core.Function | None: ... @overload def get_global_func(name: str, allow_missing: Literal[False] = False) -> core.Function: ... def get_global_func(name: str, allow_missing: bool = False) -> core.Function | None: """Get a global function by name. Parameters ---------- name The name of the global function allow_missing Whether allow missing function or raise an error. Returns ------- func The function to be returned, ``None`` if function is missing. Examples -------- .. code-block:: python import tvm_ffi @tvm_ffi.register_global_func("demo.echo") def echo(x): return x f = tvm_ffi.get_global_func("demo.echo") assert f(123) == 123 See Also -------- :py:func:`tvm_ffi.register_global_func` """ return core._get_global_func(name, allow_missing) def list_global_func_names() -> list[str]: """Get list of global functions registered. Returns ------- names List of global functions names. """ name_functor = get_global_func("ffi.FunctionListGlobalNamesFunctor")() num_names = name_functor(-1) return [name_functor(i) for i in range(num_names)] def remove_global_func(name: str) -> None: """Remove a global function by name. Parameters ---------- name The name of the global function. Examples -------- .. code-block:: python import tvm_ffi @tvm_ffi.register_global_func("my.temp") def temp(): return 42 assert tvm_ffi.get_global_func("my.temp", allow_missing=True) is not None tvm_ffi.remove_global_func("my.temp") assert tvm_ffi.get_global_func("my.temp", allow_missing=True) is None See Also -------- :py:func:`tvm_ffi.register_global_func` :py:func:`tvm_ffi.get_global_func` """ get_global_func("ffi.FunctionRemoveGlobal")(name) def get_global_func_metadata(name: str) -> dict[str, Any]: """Get metadata (including type schema) for a global function. Parameters ---------- name The name of the global function. Returns ------- metadata A dictionary containing function metadata. The ``type_schema`` field encodes the callable signature. Examples -------- .. code-block:: python import tvm_ffi meta = tvm_ffi.get_global_func_metadata("testing.add_one") print(meta) See Also -------- :py:func:`tvm_ffi.get_global_func` Retrieve a callable for an existing global function. :py:func:`tvm_ffi.register_global_func` Register a Python callable as a global FFI function. """ metadata_json = get_global_func("ffi.GetGlobalFuncMetadata")(name) return json.loads(metadata_json) if metadata_json else {} def init_ffi_api(namespace: str, target_module_name: str | None = None) -> None: """Initialize register ffi api functions into a given module. Parameters ---------- namespace The namespace of the source registry target_module_name The target module name if different from namespace Examples -------- A typical usage pattern is to create a _ffi_api.py file to register the functions under a given module. The following code populates all registered global functions prefixed with ``mypackage.`` into the current module, then we can call the function through ``_ffi_api.func_name(*args)`` which will call into the registered global function "mypackage.func_name". .. code-block:: python # _ffi_api.py import tvm_ffi tvm_ffi.init_ffi_api("mypackage", __name__) """ target_module_name = target_module_name if target_module_name else namespace if namespace.startswith("tvm."): prefix = namespace[4:] else: prefix = namespace target_module = sys.modules[target_module_name] for name in list_global_func_names(): if not name.startswith(prefix): continue fname = name[len(prefix) + 1 :] if fname.find(".") != -1: continue f = get_global_func(name) setattr(f, "__name__", fname) setattr(target_module, fname, f) def _install_init(cls: type, type_info: TypeInfo) -> None: """Install ``__init__`` from ``__ffi_init__`` TypeMethod or TypeAttrColumn. Skipped if the class body already defines ``__init__``. This ensures that ``register_object`` alone provides a working constructor, maintaining the invariant that ``c_class`` is a full alias of ``register_object`` + dunder installation. When no ``__ffi_init__`` is available and the class is an ``Object`` subclass, a TypeError guard is installed to prevent segfaults from uninitialised handles. """ if "__init__" in cls.__dict__: return # Look up __ffi_init__ from TypeMethod (preferred) or TypeAttrColumn (fallback). ffi_init = None for method in type_info.methods: if method.name == "__ffi_init__": ffi_init = method.func break if ffi_init is None: ffi_init = core._lookup_type_attr(type_info.type_index, "__ffi_init__") if ffi_init is not None: from ._dunder import _make_init # noqa: PLC0415 cls.__init__ = _make_init( # type: ignore[attr-defined] cls, type_info, ffi_init=ffi_init, ) elif issubclass(cls, core.Object): type_name = cls.__name__ def __init__(self: Any, *args: Any, **kwargs: Any) -> None: raise TypeError( f"`{type_name}` cannot be constructed directly. " f"Define a custom __init__ or use a factory method." ) __init__.__qualname__ = f"{cls.__qualname__}.__init__" __init__.__module__ = cls.__module__ cls.__init__ = __init__ # type: ignore[attr-defined] def _add_class_attrs(type_cls: type, type_info: TypeInfo) -> type: for field in type_info.fields: name = field.name if name not in type_cls.__dict__: # skip attributes defined directly on this class setattr(type_cls, name, field.as_property(type_cls)) has_ffi_init = False for method in type_info.methods: name = method.name if name == "__ffi_init__": _install_ffi_init_attr(type_cls, type_info, method.func) has_ffi_init = True continue if not hasattr(type_cls, name): setattr(type_cls, name, method.as_callable(type_cls)) # Also check TypeAttrColumn for auto-generated __ffi_init__. if not has_ffi_init: ffi_init = core._lookup_type_attr(type_info.type_index, "__ffi_init__") if ffi_init is not None: _install_ffi_init_attr(type_cls, type_info, ffi_init) return type_cls def _install_ffi_init_attr(cls: type, type_info: TypeInfo, ffi_init: Function) -> None: """Install ``__ffi_init__`` as a method that delegates to ``__init_handle_by_constructor__``. Custom ``__init__`` methods call ``self.__ffi_init__(*args, **kwargs)`` to construct the underlying C++ object. This installs a wrapper that translates that call into ``self.__init_handle_by_constructor__(ffi_init, *ffi_args)`` with kwargs packed using the FFI KWARGS protocol. The wrapper includes a type-owner guard (same as ``_make_init``) to prevent subclasses from accidentally using a parent's ``__ffi_init__``. """ kwargs_obj = core.KWARGS missing = core.MISSING type_name = cls.__name__ def __ffi_init__(self: Any, *args: Any, **kwargs: Any) -> None: if type_info is not type(self).__tvm_ffi_type_info__: raise TypeError( f"Calling `{type_name}.__ffi_init__()` on a `{type(self).__name__}` " f"instance is not supported. Define `{type(self).__name__}` with init=True." ) ffi_args: list[Any] = list(args) if kwargs: ffi_args.append(kwargs_obj) for key, val in kwargs.items(): if val is not missing: ffi_args.append(key) ffi_args.append(val) self.__init_handle_by_constructor__(ffi_init, *ffi_args) __ffi_init__.__qualname__ = f"{cls.__qualname__}.__ffi_init__" __ffi_init__.__module__ = cls.__module__ cls.__ffi_init__ = __ffi_init__ # type: ignore[attr-defined] def _warn_missing_field_annotations(cls: type, type_info: TypeInfo, *, stacklevel: int) -> None: """Emit a warning if any C++ reflected fields lack Python annotations on *cls*. Only checks fields owned by *type_info* (not inherited from parents). Only checks annotations defined directly on *cls* (``cls.__dict__``), so parent annotations do not suppress warnings for child-level fields. """ reflected_names = {field.name for field in type_info.fields} if not reflected_names: return own_annotations = cls.__dict__.get("__annotations__", {}) missing = sorted(reflected_names - set(own_annotations)) if missing: missing_str = ", ".join(missing) warnings.warn( f"@c_class({type_info.type_key!r}): class `{cls.__qualname__}` does not " f"annotate the following reflected field(s): {missing_str}. " f"Add type annotations (e.g. `field_name: type`) to the class body " f"for IDE support and documentation.", UserWarning, stacklevel=stacklevel, ) def get_registered_type_keys() -> Sequence[str]: """Get the list of valid type keys registered to TVM-FFI. Returns ------- type_keys List of valid type keys. """ return get_global_func("ffi.GetRegisteredTypeKeys")() __all__ = [ "get_global_func", "get_global_func_metadata", "get_registered_type_keys", "init_ffi_api", "list_global_func_names", "register_global_func", "register_object", "remove_global_func", ] tvm-ffi-0.1.12/python/tvm_ffi/serialization.py000066400000000000000000000046361521067262500213640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities for serializing and deserializing FFI object graphs. These helpers produce a stable JSON graph representation that preserves object identity and references. It is useful for debugging and for lightweight persistence when pickling is not available. """ from __future__ import annotations from typing import Any from . import _ffi_api def to_json_graph_str(obj: Any, metadata: dict[str, Any] | None = None) -> str: """Dump an object to a JSON graph string. The JSON graph is a textual representation of the object graph that preserves shared references. It can be used for debugging or simple persistence. Parameters ---------- obj The object to save. metadata Extra metadata to save into the json graph string. Returns ------- json_str The JSON graph string. Examples -------- .. code-block:: python import tvm_ffi a = tvm_ffi.convert([1, 2, 3]) s = tvm_ffi.serialization.to_json_graph_str(a) b = tvm_ffi.serialization.from_json_graph_str(s) assert list(b) == [1, 2, 3] """ return _ffi_api.ToJSONGraphString(obj, metadata) def from_json_graph_str(json_str: str) -> Any: """Load an object from a JSON graph string. The JSON graph string is produced by :py:func:`to_json_graph_str` and can be converted back into the corresponding FFI-backed objects. Parameters ---------- json_str The JSON graph string to load. Returns ------- obj The loaded object. """ return _ffi_api.FromJSONGraphString(json_str) __all__ = ["from_json_graph_str", "to_json_graph_str"] tvm-ffi-0.1.12/python/tvm_ffi/stream.py000066400000000000000000000146571521067262500200060ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name """Stream context.""" from __future__ import annotations from ctypes import c_void_p from typing import Any from . import core from ._tensor import device class StreamContext: """Represent a stream context in the FFI system. StreamContext helps setup ffi environment stream by python `with` statement. When entering `with` scope, it caches the current environment stream and setup the given new stream. When exiting `with` scope, it recovers the stream to the cached environment stream. Parameters ---------- device The device to which the stream belongs. stream The stream handle. See Also -------- :py:func:`tvm_ffi.use_raw_stream`, :py:func:`tvm_ffi.use_torch_stream` """ def __init__(self, device: core.Device, stream: int | c_void_p) -> None: """Initialize a stream context with a device and stream handle.""" self.device_type = device.dlpack_device_type() self.device_id = device.index self.stream = stream def __enter__(self) -> StreamContext: """Enter the context and set the current stream.""" self.prev_stream = core._env_set_current_stream( self.device_type, self.device_id, self.stream ) return self def __exit__(self, *args: Any) -> None: """Exit the context and restore the previous stream.""" self.prev_stream = core._env_set_current_stream( self.device_type, self.device_id, self.prev_stream ) try: import torch class TorchStreamContext: """Context manager that syncs Torch and FFI stream contexts.""" def __init__(self, context: Any) -> None: """Initialize with an optional Torch stream/graph context wrapper.""" self.torch_context = context def __enter__(self) -> TorchStreamContext: """Enter both Torch and FFI stream contexts.""" if self.torch_context: self.torch_context.__enter__() current_stream = torch.cuda.current_stream() self.ffi_context = StreamContext( device(str(current_stream.device)), current_stream.cuda_stream ) self.ffi_context.__enter__() return self def __exit__(self, *args: Any) -> None: """Exit both Torch and FFI stream contexts.""" if self.torch_context: self.torch_context.__exit__(*args) self.ffi_context.__exit__(*args) def use_torch_stream(context: Any = None) -> TorchStreamContext: """Create an FFI stream context with a Torch stream or graph. cuda graph or current stream if `None` provided. Parameters ---------- context The wrapped torch stream or cuda graph. Returns ------- context The ffi stream context wrapping torch stream context. Examples -------- .. code-block:: python s = torch.cuda.Stream() with tvm_ffi.use_torch_stream(torch.cuda.stream(s)): ... g = torch.cuda.CUDAGraph() with tvm_ffi.use_torch_stream(torch.cuda.graph(g)): ... Note ---- When working with a raw ``cudaStream_t`` handle, use :py:func:`tvm_ffi.use_raw_stream` instead. """ return TorchStreamContext(context) except ImportError: def use_torch_stream(context: Any = None) -> TorchStreamContext: """Raise an informative error when Torch is unavailable.""" raise ImportError("Cannot import torch") def use_raw_stream(device: core.Device, stream: int | c_void_p) -> StreamContext: """Create an FFI stream context with the given device and stream handle. Parameters ---------- device The device to which the stream belongs. stream The stream handle (for example, a CUDA ``cudaStream_t`` as an integer, or ``0``). Returns ------- context The FFI stream context. Examples -------- The example below uses a CPU device and a dummy stream handle. On CUDA, pass a real ``cudaStream_t`` integer. .. code-block:: python import tvm_ffi dev = tvm_ffi.device("cpu:0") with tvm_ffi.use_raw_stream(dev, 0): # Within the context, the current stream for this device is set assert tvm_ffi.get_raw_stream(dev) == 0 See Also -------- :py:func:`tvm_ffi.use_torch_stream` Use a Torch stream or CUDA graph as the source of truth. :py:func:`tvm_ffi.get_raw_stream` Query the current FFI stream for a device. """ if not isinstance(stream, (int, c_void_p)): raise ValueError( "use_raw_stream only accepts int or c_void_p as stream input, " "try use_torch_stream when using torch.cuda.Stream or torch.cuda.graph" ) return StreamContext(device, stream) def get_raw_stream(device: core.Device) -> int: """Get the current FFI stream of a given device. Parameters ---------- device The device to which the stream belongs. Returns ------- stream The current FFI stream as an integer handle. Examples -------- .. code-block:: python import tvm_ffi dev = tvm_ffi.device("cpu:0") # Default stream is implementation-defined; set it explicitly with tvm_ffi.use_raw_stream(dev, 0): assert tvm_ffi.get_raw_stream(dev) == 0 See Also -------- :py:func:`tvm_ffi.use_raw_stream` Set the current stream for a device. """ return core._env_get_current_stream(device.dlpack_device_type(), device.index) tvm-ffi-0.1.12/python/tvm_ffi/structural.py000066400000000000000000000154451521067262500207170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name """Structural helper objects and functions.""" from __future__ import annotations from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .access_path import AccessPath from . import _ffi_api from .core import Object from .registry import register_object __all__ = [ "StructuralKey", "get_first_structural_mismatch", "structural_equal", "structural_hash", ] def structural_equal( lhs: Any, rhs: Any, map_free_vars: bool = False, skip_tensor_content: bool = False ) -> bool: """Check structural equality between two values. Structural equality compares the *shape/content structure* of two values instead of Python object identity. For container-like values, this means recursive comparison of elements/fields. For object types that provide structural equal hooks, those hooks are used. Parameters ---------- lhs Left-hand side value. rhs Right-hand side value. map_free_vars Whether free variables (variables without a definition site) can be mapped to each other during comparison. skip_tensor_content Whether to skip tensor data content when comparing tensors. Returns ------- result Whether the two values are structurally equal. Examples -------- .. code-block:: python import tvm_ffi assert tvm_ffi.structural_equal([1, 2, 3], [1, 2, 3]) assert not tvm_ffi.structural_equal([1, 2, 3], [1, 2, 4]) See Also -------- :py:func:`tvm_ffi.structural_hash` Hash function compatible with structural equality. :py:func:`tvm_ffi.get_first_structural_mismatch` Mismatch diagnostics with access paths. """ return _ffi_api.StructuralEqual(lhs, rhs, map_free_vars, skip_tensor_content) def structural_hash( value: Any, map_free_vars: bool = False, skip_tensor_content: bool = False ) -> int: """Compute structural hash of a value. This hash is designed to be consistent with :py:func:`structural_equal` under the same options. If two values are structurally equal, they should have the same structural hash. Parameters ---------- value Input value to hash. map_free_vars Whether free variables mapped to each other during hashing. skip_tensor_content Whether tensor data content is ignored when hashing tensors. Returns ------- hash_value Structural hash value. Examples -------- .. code-block:: python import tvm_ffi h0 = tvm_ffi.structural_hash([1, 2, 3]) h1 = tvm_ffi.structural_hash([1, 2, 3]) assert h0 == h1 Notes ----- Structural hash is intended for hash-table bucketing and fast pre-checks. Always use structural equality to confirm semantic equivalence. """ # need to mask the result so negative values are converted to u64 # this is because hash value were stored as int64_t in the C++ code return _ffi_api.StructuralHash(value, map_free_vars, skip_tensor_content) & 0xFFFFFFFFFFFFFFFF def get_first_structural_mismatch( lhs: Any, rhs: Any, map_free_vars: bool = False, skip_tensor_content: bool = False ) -> tuple[AccessPath, AccessPath] | None: """Like structural_equal(), but returns the AccessPath pair of the first detected mismatch. Parameters ---------- lhs The left operand. rhs The right operand. map_free_vars Whether free variables (i.e. variables without a definition site) should be mapped as equal to each other. skip_tensor_content Whether to skip the data content of tensor. Returns ------- mismatch: tuple[AccessPath, AccessPath] | None `None` if `lhs` and `rhs` are structurally equal. Otherwise, a tuple of two AccessPath objects that point to the first detected mismatch. """ return _ffi_api.GetFirstStructuralMismatch(lhs, rhs, map_free_vars, skip_tensor_content) @register_object("ffi.StructuralKey") class StructuralKey(Object): """Hash-cached structural key wrapper. This wrapper can be used to hint that a dict uses structural equality and hash for the key. Examples -------- Use ``StructuralKey`` with Python dictionaries when you want key lookup by structural semantics: .. code-block:: python import tvm_ffi k0 = tvm_ffi.StructuralKey([1, 2, 3]) k1 = tvm_ffi.StructuralKey([1, 2, 3]) k2 = tvm_ffi.StructuralKey([1, 2, 4]) d = {k0: "value-a", k2: "value-b"} assert d[k1] == "value-a" # k1 matches k0 structurally assert d[k2] == "value-b" It can also be used directly with :py:class:`tvm_ffi.Map`: .. code-block:: python m = tvm_ffi.Map({k0: 1, k1: 2}) assert len(m) == 1 assert m[k0] == 2 See Also -------- :py:func:`tvm_ffi.structural_equal` Structural equality comparison. :py:func:`tvm_ffi.structural_hash` Structural hash computation. """ # tvm-ffi-stubgen(begin): object/ffi.StructuralKey # fmt: off key: Any hash_i64: int if TYPE_CHECKING: def __init__(self, key: Any, hash_i64: int) -> None: ... def __ffi_init__(self, _0: Any, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) def __init__(self, key: Any) -> None: """Create a structural key from ``key``. Parameters ---------- key The underlying value used for structural hash/equality. """ self.__init_handle_by_constructor__(_ffi_api.StructuralKey, key) def __hash__(self) -> int: """Return cached structural hash.""" # need to mask the result so negative values are converted to u64 # this is because hash value were stored as int64_t in the C++ code return self.hash_i64 & 0xFFFFFFFFFFFFFFFF def __eq__(self, other: Any) -> bool: """Compare by structural equality.""" return isinstance(other, StructuralKey) and _ffi_api.StructuralKeyEqual(self, other) tvm-ffi-0.1.12/python/tvm_ffi/stub/000077500000000000000000000000001521067262500171015ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/stub/__init__.py000066400000000000000000000014711521067262500212150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """A series of stub generator tools.""" tvm-ffi-0.1.12/python/tvm_ffi/stub/cli.py000066400000000000000000000333431521067262500202300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """TVM-FFI Stub Generator (``tvm-ffi-stubgen``).""" from __future__ import annotations import argparse import ctypes import importlib import sys import traceback from pathlib import Path from . import codegen as G from . import consts as C from .file_utils import FileInfo, collect_files from .lib_state import ( collect_global_funcs, collect_type_keys, object_info_from_type_key, toposort_objects, ) from .utils import FuncInfo, ImportItem, InitConfig, Options def __main__() -> int: """Command line entry point for ``tvm-ffi-stubgen``. This generates in-place type stubs inside special ``tvm-ffi-stubgen`` blocks in the given files or directories. See the module docstring for an overview and examples of the block syntax. """ opt = _parse_args() for imp in opt.imports or []: importlib.import_module(imp) dlls = [ctypes.CDLL(lib) for lib in opt.dlls] files: list[FileInfo] = collect_files([Path(f) for f in opt.files]) global_funcs: dict[str, list[FuncInfo]] = collect_global_funcs() init_path: Path | None = None if opt.files: init_path = Path(opt.files[0]).resolve() if init_path.is_file(): init_path = init_path.parent # Stage 1: Collect information # - type maps: `tvm-ffi-stubgen(ty-map)` # - defined global functions: `tvm-ffi-stubgen(begin): global/...` # - defined object types: `tvm-ffi-stubgen(begin): object/...` ty_map: dict[str, str] = C.TY_MAP_DEFAULTS.copy() for file in files: try: _stage_1(file, ty_map) except Exception: print( f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' ) # Stage 2. Generate stubs if they are not defined on the file. if opt.init: assert init_path is not None, "init-path could not be determined" _stage_2( files, ty_map, init_cfg=opt.init, init_path=init_path, global_funcs=global_funcs, ) # Stage 3: Process # - `tvm-ffi-stubgen(begin): global/...` # - `tvm-ffi-stubgen(begin): object/...` for file in files: if opt.verbose: print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}") try: _stage_3(file, opt, ty_map, global_funcs) except Exception: print( f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' ) del dlls return 0 def _stage_1( file: FileInfo, ty_map: dict[str, str], ) -> None: for code in file.code_blocks: if code.kind == "ty-map": try: assert isinstance(code.param, str) lhs, rhs = code.param.split("->") except ValueError as e: raise ValueError( f"Invalid ty_map format at line {code.lineno_start}. Example: `A.B -> C.D`" ) from e ty_map[lhs.strip()] = rhs.strip() def _stage_2( files: list[FileInfo], ty_map: dict[str, str], init_cfg: InitConfig, init_path: Path, global_funcs: dict[str, list[FuncInfo]], ) -> None: def _find_or_insert_file(path: Path) -> FileInfo: ret: FileInfo | None if not path.exists(): ret = FileInfo(path=path, lines=(), code_blocks=[]) else: for file in files: if path.samefile(file.path): return file ret = FileInfo.from_file(file=path, include_empty=True) assert ret is not None, f"Failed to read file: {path}" files.append(ret) return ret # Step 0. Find out functions and classes already defined on files. defined_func_prefixes: set[str] = { code.param[0] for file in files for code in file.code_blocks if code.kind == "global" } defined_objs: set[str] = { # ty: ignore[invalid-assignment] code.param for file in files for code in file.code_blocks if code.kind == "object" } | C.BUILTIN_TYPE_KEYS # Step 0. Generate missing `_ffi_api.py` and `__init__.py` under each prefix. prefix_filter = init_cfg.prefix.strip() if prefix_filter and not prefix_filter.endswith("."): prefix_filter += "." root_prefix = prefix_filter.rstrip(".") prefixes: dict[str, list[str]] = collect_type_keys() for prefix in global_funcs: prefixes.setdefault(prefix, []) for prefix, obj_names in prefixes.items(): if not (prefix == root_prefix or prefix.startswith(prefix_filter)): continue funcs = sorted( [] if prefix in defined_func_prefixes else global_funcs.get(prefix, []), key=lambda f: f.schema.name, ) objs = sorted(set(obj_names) - defined_objs) object_infos = toposort_objects(objs) if not funcs and not object_infos: continue # Step 1. Create target directory if not exists directory = init_path / prefix.replace(".", "/") directory.mkdir(parents=True, exist_ok=True) # Step 2. Generate `_ffi_api.py` target_path = directory / "_ffi_api.py" target_file = _find_or_insert_file(target_path) with target_path.open("a", encoding="utf-8") as f: f.write( G.generate_ffi_api( target_file.code_blocks, ty_map, prefix, object_infos, init_cfg, is_root=prefix == root_prefix, ) ) target_file.reload() # Step 3. Generate `__init__.py` target_path = directory / "__init__.py" target_file = _find_or_insert_file(target_path) with target_path.open("a", encoding="utf-8") as f: f.write(G.generate_init(target_file.code_blocks, prefix, submodule="_ffi_api")) target_file.reload() def _stage_3( # noqa: PLR0912 file: FileInfo, opt: Options, ty_map: dict[str, str], global_funcs: dict[str, list[FuncInfo]], ) -> None: defined_funcs: set[str] = set() defined_types: set[str] = set() imports: list[ImportItem] = [] ffi_load_lib_imported = False # Stage 1. Collect `tvm-ffi-stubgen(import-object): ...` for code in file.code_blocks: if code.kind == "import-object": name, type_checking_only, alias = code.param imports.append( ImportItem( name, type_checking_only=( bool(type_checking_only) and isinstance(type_checking_only, str) and type_checking_only.lower() == "true" ), alias=alias if alias else None, ) ) if (alias and alias == "_FFI_LOAD_LIB") or name.endswith("libinfo.load_lib_module"): ffi_load_lib_imported = True # Stage 2. Process `tvm-ffi-stubgen(begin): global/...` for code in file.code_blocks: if code.kind == "global": funcs = global_funcs.get(code.param[0], []) for func in funcs: defined_funcs.add(func.schema.name) G.generate_global_funcs(code, funcs, ty_map, imports, opt) # Stage 3. Process `tvm-ffi-stubgen(begin): object/...` for code in file.code_blocks: if code.kind == "object": type_key = code.param assert isinstance(type_key, str) obj_info = object_info_from_type_key(type_key) type_key = ty_map.get(type_key, type_key) full_name = ImportItem(type_key).full_name defined_types.add(full_name) G.generate_object(code, ty_map, imports, opt, obj_info) # Stage 4. Add imports for used types. imports = [i for i in imports if i.full_name not in defined_types] for code in file.code_blocks: if code.kind == "import-section": G.generate_import_section(code, imports, opt) break # Only one import block per file is supported for now. # Stage 5. Add `__all__` for defined classes and functions. for code in file.code_blocks: if code.kind == "__all__": export_names = defined_funcs | defined_types if ffi_load_lib_imported: export_names = export_names | {"LIB"} G.generate_all(code, export_names, opt) break # Only one __all__ block per file is supported for now. # Stage 6. Process `tvm-ffi-stubgen(begin): export/...` for code in file.code_blocks: if code.kind == "export": G.generate_export(code) # Finalize: write back to file file.update(verbose=opt.verbose, dry_run=opt.dry_run) def _parse_args() -> Options: class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelpFormatter): pass def _split_list_arg(arg: str | None) -> list[str]: if not arg: return [] return [item.strip() for item in arg.split(";") if item.strip()] parser = argparse.ArgumentParser( prog="tvm-ffi-stubgen", description=( "Generate type stubs for TVM FFI extensions. It supports two modes\n" "- In `--init-*` mode, it generates missing `_ffi_api.py` and `__init__.py` files, " "based on the registered global functions and object types in the loaded libraries.\n" "- In normal mode, it processes the given files/directories in-place, generating " "type stubs inside special `tvm-ffi-stubgen` directive blocks.\n\n" f"Documentation: {C.TERM_CYAN}{C.DOC_URL}{C.TERM_RESET}." ), formatter_class=HelpFormatter, ) parser.add_argument( "--imports", type=str, default="", metavar="IMPORTS", help=( "Additional imports to load before generation, separated by ';' " "(e.g. 'pkgA;pkgB.submodule')." ), ) parser.add_argument( "--dlls", type=str, default="", metavar="LIBS", help=( "Shared libraries to preload before generation (e.g. TVM runtime or " "your extension), separated by ';'. This ensures global function and " "object metadata is available. Platform-specific suffixes like " ".so/.dylib/.dll are supported." ), ) parser.add_argument( "--init-pypkg", type=str, default="", help=( "Python package name to generate stubs for (e.g. apache-tvm-ffi). " "Required together with --init-lib and --init-prefix." ), ) parser.add_argument( "--init-lib", type=str, default="", help=( "CMake target that produces the shared library to load for stub generation " "(e.g. tvm_ffi_shared). Required together with --init-pypkg and " "--init-prefix." ), ) parser.add_argument( "--init-prefix", type=str, default="", help=( "Global function/object prefix to include when generating stubs " "(e.g. tvm_ffi.). Required together with --init-pypkg and --init-lib." ), ) parser.add_argument( "--indent", type=int, default=4, help=( "Extra spaces added inside each generated block, relative to the " f"indentation of the corresponding '{C.STUB_BEGIN}' line." ), ) parser.add_argument( "files", nargs="*", metavar="PATH", help=( "Files or directories to process. Directories are scanned recursively; " "only .py and .pyi files are modified. Use tvm-ffi-stubgen directives to " "select where stubs are generated." ), ) parser.add_argument( "--verbose", action="store_true", help=( "Print a unified diff of changes to each file. This is useful for " "debugging or previewing changes before applying them." ), ) parser.add_argument( "--dry-run", action="store_true", help=( "Don't write changes to files. This is useful for previewing changes " "without modifying any files." ), ) args = parser.parse_args() init_flags = [args.init_pypkg, args.init_lib, args.init_prefix] init_cfg: InitConfig | None = None if any(init_flags): if not all(init_flags): parser.error("--init-pypkg, --init-lib, and --init-prefix must be provided together") init_cfg = InitConfig( pkg=args.init_pypkg, shared_target=args.init_lib, prefix=args.init_prefix, ) if not args.files: parser.print_help() sys.exit(1) return Options( imports=_split_list_arg(args.imports), dlls=_split_list_arg(args.dlls), init=init_cfg, indent=args.indent, files=args.files, verbose=args.verbose, dry_run=args.dry_run, ) if __name__ == "__main__": sys.exit(__main__()) tvm-ffi-0.1.12/python/tvm_ffi/stub/codegen.py000066400000000000000000000237141521067262500210660ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Code generation logic for the `tvm-ffi-stubgen` tool.""" from __future__ import annotations from typing import Callable from . import consts as C from .file_utils import CodeBlock from .utils import FuncInfo, ImportItem, InitConfig, ObjectInfo, Options def _type_suffix_and_record( ty_map: dict[str, str], imports: list[ImportItem], func_names: set[str] | None = None, ) -> Callable[[str], str]: def _run(name: str) -> str: nonlocal ty_map, imports name = ty_map.get(name, name) suffix = name.rsplit(".", 1)[-1] if "." in name: alias = None if func_names and suffix in func_names: alias = f"_{suffix}" imports.append(ImportItem(name, type_checking_only=True, alias=alias)) if alias: return alias return suffix return _run def generate_global_funcs( code: CodeBlock, global_funcs: list[FuncInfo], ty_map: dict[str, str], imports: list[ImportItem], opt: Options, ) -> None: """Generate function signatures for global functions. It processes: global/${prefix}@${import_from="tvm_ffi") """ assert len(code.lines) >= 2 if not global_funcs: return assert isinstance(code.param, tuple) prefix, import_from = code.param if not import_from: import_from = "tvm_ffi" imports.extend( [ ImportItem( f"{import_from}.init_ffi_api", type_checking_only=False, alias="_FFI_INIT_FUNC", ), ImportItem( "typing.TYPE_CHECKING", type_checking_only=False, ), ] ) func_names = {f.schema.name.rsplit(".", 1)[-1] for f in global_funcs} fn_ty_map = _type_suffix_and_record(ty_map, imports, func_names=func_names) results: list[str] = [ "# fmt: off", f'_FFI_INIT_FUNC("{prefix}", __name__)', "if TYPE_CHECKING:", *[func.gen(fn_ty_map, indent=opt.indent) for func in global_funcs], "# fmt: on", ] indent = " " * code.indent code.lines = [ code.lines[0], *[indent + line for line in results], code.lines[-1], ] def generate_object( code: CodeBlock, ty_map: dict[str, str], imports: list[ImportItem], opt: Options, obj_info: ObjectInfo, ) -> None: """Generate a class definition for an object type. It processes: object/${type_key} """ assert len(code.lines) >= 2 info = obj_info method_names = {m.schema.name.rsplit(".", 1)[-1] for m in info.methods} fn_ty_map = _type_suffix_and_record(ty_map, imports, func_names=method_names) init_lines = info.gen_init(fn_ty_map, indent=opt.indent) ffi_init_lines = info.gen_ffi_init(fn_ty_map, indent=opt.indent) type_checking_lines = [ *init_lines, *ffi_init_lines, *info.gen_methods(fn_ty_map, indent=opt.indent), ] if type_checking_lines: imports.append( ImportItem( "typing.TYPE_CHECKING", type_checking_only=False, ) ) results = [ "# fmt: off", *info.gen_fields(fn_ty_map, indent=0), "if TYPE_CHECKING:", *type_checking_lines, "# fmt: on", ] else: results = [ "# fmt: off", *info.gen_fields(fn_ty_map, indent=0), "# fmt: on", ] indent = " " * code.indent code.lines = [ code.lines[0], *[indent + line for line in results], code.lines[-1], ] def generate_import_section( code: CodeBlock, imports: list[ImportItem], opt: Options, ) -> None: """Generate import statements for the types used in the stub. It processes: import-section """ imports_concrete: dict[str, list[ImportItem]] = {} imports_ty_check: dict[str, list[ImportItem]] = {} for item in imports: if item.type_checking_only: imports_ty_check.setdefault(item.mod, []).append(item) else: imports_concrete.setdefault(item.mod, []).append(item) if imports_ty_check: imports_concrete.setdefault("typing", []).append( ImportItem("typing.TYPE_CHECKING", type_checking_only=True) ) def _make_line(mod: str, items: list[ImportItem], indent: int) -> str: items.sort(key=lambda item: item.name) names = ", ".join(sorted(set(item.name_with_alias for item in items))) indent_str = " " * indent if mod: return f"{indent_str}from {mod} import {names}" else: return f"{indent_str}import {names}" results: list[str] = [] if imports_concrete: results.extend( _make_line(mod, imports_concrete[mod], indent=0) for mod in sorted(imports_concrete) ) if imports_ty_check: results.append("if TYPE_CHECKING:") results.extend( _make_line(mod, imports_ty_check[mod], opt.indent) for mod in sorted(imports_ty_check) ) if results: code.lines = [ code.lines[0], "# fmt: off", "# isort: off", "from __future__ import annotations", *results, "# isort: on", "# fmt: on", code.lines[-1], ] def generate_all(code: CodeBlock, names: set[str], opt: Options) -> None: """Generate an `__all__` variable for the given names.""" assert len(code.lines) >= 2 if not names: return indent = " " * code.indent names = {f.rsplit(".", 1)[-1] for f in names} def _sort_key(name: str) -> tuple[int, str]: if name.isupper(): return (0, name) if name and name[0].isupper() and not "_" in name: return (1, name) return (2, name) code.lines = [ code.lines[0], *[f'{indent}"{name}",' for name in sorted(names, key=_sort_key)], code.lines[-1], ] def generate_export(code: CodeBlock) -> None: """Generate an `__all__` variable for the given names.""" assert len(code.lines) >= 2 mod = code.param code.lines = [ code.lines[0], "# fmt: off", "# isort: off", f"from .{mod} import * # noqa: F403", f"from .{mod} import __all__ as {mod}__all__", 'if "__all__" not in globals():', " __all__ = []", f"__all__.extend({mod}__all__)", "# isort: on", "# fmt: on", code.lines[-1], ] def generate_ffi_api( code_blocks: list[CodeBlock], ty_map: dict[str, str], module_name: str, object_infos: list[ObjectInfo], init_cfg: InitConfig, is_root: bool, ) -> str: """Generate the initial FFI API stub code for a given module.""" # TODO(@junrus): New code is appended to the end of the file. # We should consider a more sophisticated approach. append = "" # Part 0. Imports if not code_blocks: append += f"""\"\"\"FFI API bindings for {module_name}.\"\"\"\n""" if not any(code.kind == "import-section" for code in code_blocks): append += C.PROMPT_IMPORT_SECTION # Part 1. Library loading if is_root: append += C._prompt_import_object("tvm_ffi.libinfo.load_lib_module", "_FFI_LOAD_LIB") append += f"""LIB = _FFI_LOAD_LIB("{init_cfg.pkg}", "{init_cfg.shared_target}")\n""" # Part 2. Global functions if not any(code.kind == "global" for code in code_blocks): append += C._prompt_globals(module_name) # Part 3. Object types if object_infos: append += C._prompt_import_object("tvm_ffi.register_object", "_FFI_REG_OBJ") defined_type_keys = {info.type_key for info in object_infos if info.type_key} for info in object_infos: type_key = info.type_key parent_type_key = info.parent_type_key if type_key is None: continue # Canonicalize type key names type_key = ty_map.get(type_key, type_key) type_name = type_key.rsplit(".", 1)[-1] parent_type_key = ( ty_map.get(parent_type_key, parent_type_key) if parent_type_key else parent_type_key ) parent_type_name = parent_type_key.rsplit(".", 1)[-1] if parent_type_key else "Object" # Import parent type keys if they are not defined in the current module if parent_type_key and parent_type_key not in defined_type_keys: parent_type_name = "_" + parent_type_key.replace(".", "_") append += C._prompt_import_object(parent_type_key, parent_type_name) # Generate class definition append += C._prompt_class_def( type_name, type_key, parent_type_name, ) # Part 4. __all__ if not any(code.kind == "__all__" for code in code_blocks): append += C.PROMPT_ALL_SECTION return append def generate_init( code_blocks: list[CodeBlock], module_name: str, submodule: str = "_ffi_api", ) -> str: """Generate the `__init__.py` file for the `tvm_ffi` package.""" code = f""" {C.STUB_BEGIN} export/{submodule} {C.STUB_END} """ if not code_blocks: return f"""\"\"\"Package {module_name}.\"\"\"\n""" + code if not any(code.kind == "export" for code in code_blocks): return code return "" tvm-ffi-0.1.12/python/tvm_ffi/stub/consts.py000066400000000000000000000057441521067262500207760ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Constants used in stub generation.""" from __future__ import annotations from typing import Literal from typing_extensions import TypeAlias STUB_PREFIX = "# tvm-ffi-stubgen(" STUB_BEGIN = f"{STUB_PREFIX}begin):" STUB_END = f"{STUB_PREFIX}end)" STUB_TY_MAP = f"{STUB_PREFIX}ty-map):" STUB_IMPORT_OBJECT = f"{STUB_PREFIX}import-object):" STUB_SKIP_FILE = f"{STUB_PREFIX}skip-file)" STUB_BLOCK_KINDS: TypeAlias = Literal[ "global", "object", "ty-map", "import-section", "import-object", "export", "__all__", None, ] TERM_RESET = "\033[0m" TERM_BOLD = "\033[1m" TERM_BLACK = "\033[30m" TERM_RED = "\033[31m" TERM_GREEN = "\033[32m" TERM_YELLOW = "\033[33m" TERM_BLUE = "\033[34m" TERM_MAGENTA = "\033[35m" TERM_CYAN = "\033[36m" TERM_WHITE = "\033[37m" DOC_URL = "https://tvm.apache.org/ffi/packaging/stubgen.html" DEFAULT_SOURCE_EXTS = {".py", ".pyi"} TY_MAP_DEFAULTS = { "Any": "typing.Any", "Callable": "typing.Callable", "Array": "collections.abc.Sequence", "List": "collections.abc.MutableSequence", "Map": "collections.abc.Mapping", "Dict": "collections.abc.MutableMapping", "Object": "ffi.Object", "Tensor": "ffi.Tensor", "dtype": "ffi.dtype", "Device": "ffi.Device", } # TODO(@junrushao): Make it configurable MOD_MAP = { "testing": "tvm_ffi.testing", "ffi": "tvm_ffi", } FN_NAME_MAP: dict[str, str] = {} BUILTIN_TYPE_KEYS = { "ffi.Bytes", "ffi.Error", "ffi.Function", "ffi.Object", "ffi.OpaquePyObject", "ffi.SmallBytes", "ffi.SmallStr", "ffi.String", "ffi.Tensor", } def _prompt_globals(mod: str) -> str: return f"""{STUB_BEGIN} global/{mod} {STUB_END} """ def _prompt_class_def(type_name: str, type_key: str, parent_type_name: str) -> str: return f'''@_FFI_REG_OBJ("{type_key}") class {type_name}({parent_type_name}): """FFI binding for `{type_key}`.""" {STUB_BEGIN} object/{type_key} {STUB_END}\n\n''' def _prompt_import_object(type_key: str, type_name: str) -> str: return f"""{STUB_IMPORT_OBJECT} {type_key};False;{type_name}\n""" PROMPT_IMPORT_SECTION = f""" {STUB_BEGIN} import-section {STUB_END} """ PROMPT_ALL_SECTION = f""" __all__ = [ {STUB_BEGIN} __all__ {STUB_END} ] """ tvm-ffi-0.1.12/python/tvm_ffi/stub/file_utils.py000066400000000000000000000254131521067262500216170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities for parsing and generating stub files for TVM FFI.""" from __future__ import annotations import dataclasses import difflib import os import traceback from pathlib import Path from typing import Callable, Generator, Iterable from . import consts as C @dataclasses.dataclass class CodeBlock: """A block of code to be generated in a stub file.""" kind: C.STUB_BLOCK_KINDS param: str | tuple[str, ...] lineno_start: int lineno_end: int | None lines: list[str] def __post_init__(self) -> None: """Validate the code block after initialization.""" assert self.kind in { "global", "object", "ty-map", "import-section", "import-object", "export", "__all__", None, } @property def indent(self) -> int: """Calculate the indentation level of the block based on the first line.""" if not self.lines: return 0 first_line = self.lines[0] return len(first_line) - len(first_line.lstrip(" ")) @staticmethod def from_begin_line(lineo: int, line: str) -> CodeBlock: """Parse a line to create a CodeBlock if it contains a stub begin marker.""" if line.startswith(C.STUB_TY_MAP): line = line[len(C.STUB_TY_MAP) :].strip() return CodeBlock( kind="ty-map", param=line, lineno_start=lineo, lineno_end=lineo, lines=[], ) elif line.startswith(C.STUB_IMPORT_OBJECT): line = line[len(C.STUB_IMPORT_OBJECT) :].strip() splits = [p.strip() for p in line.split(";")] if len(splits) < 3: splits += [""] * (3 - len(splits)) return CodeBlock( kind="import-object", param=tuple(splits), lineno_start=lineo, lineno_end=lineo, lines=[], ) assert line.startswith(C.STUB_BEGIN) param: str | tuple[str, ...] stub = line[len(C.STUB_BEGIN) :].strip() if stub.startswith("global/"): kind = "global" param = stub[len("global/") :].strip() if "@" in param: param = tuple(param.split("@")) else: param = (param, "") elif stub.startswith("object/"): kind = "object" param = stub[len("object/") :].strip() elif stub.startswith("ty-map/"): kind = "ty-map" param = stub[len("ty-map/") :].strip() elif stub == "import-section": kind = "import-section" param = "" elif stub.startswith("export/"): kind = "export" param = stub[len("export/") :].strip() elif stub == "__all__": kind = "__all__" param = "" else: raise ValueError(f"Unknown stub type `{stub}` at line {lineo}") return CodeBlock( kind=kind, param=param, lineno_start=lineo, lineno_end=None, lines=[], ) @dataclasses.dataclass class FileInfo: """Information about a file being processed.""" path: Path lines: tuple[str, ...] code_blocks: list[CodeBlock] def update(self, verbose: bool, dry_run: bool) -> bool: """Update the file's lines based on the current code blocks and optionally show a diff.""" new_lines = tuple(line for block in self.code_blocks for line in block.lines) if self.lines == new_lines: if verbose: print(f"{C.TERM_CYAN}-----> Unchanged{C.TERM_RESET}") return False if verbose: for line in difflib.unified_diff(self.lines, new_lines, lineterm=""): # Skip placeholder headers when fromfile/tofile are unspecified if line.startswith("---") or line.startswith("+++"): continue if line.startswith("-") and not line.startswith("---"): print(f"{C.TERM_RED}{line}{C.TERM_RESET}") # Red for removals elif line.startswith("+") and not line.startswith("+++"): print(f"{C.TERM_GREEN}{line}{C.TERM_RESET}") # Green for additions elif line.startswith("?"): print(f"{C.TERM_YELLOW}{line}{C.TERM_RESET}") # Yellow for hints else: print(line) self.lines = new_lines if not dry_run: self.path.write_text("\n".join(self.lines) + "\n", encoding="utf-8") return True @staticmethod def from_file(file: Path, include_empty: bool = False) -> FileInfo | None: # noqa: PLR0912 """Parse a file to extract code blocks based on stub markers.""" assert file.is_file(), f"Expected a file, but got: {file}" file = file.resolve() has_marker = False lines: list[str] = file.read_text(encoding="utf-8").splitlines() for _, line in enumerate(lines, start=1): if line.strip().startswith(C.STUB_SKIP_FILE): return None if line.strip().startswith(C.STUB_PREFIX): has_marker = True if not has_marker and not include_empty: return None del has_marker codes: list[CodeBlock] = [] code: CodeBlock | None = None for lineno, line in enumerate(lines, 1): clean_line = line.strip() if clean_line.startswith(C.STUB_BEGIN): # Process "# tvm-ffi-stubgen(begin)" if code is not None: raise ValueError(f"Nested stub not permitted at line {lineno}") code = CodeBlock.from_begin_line(lineno, clean_line) code.lineno_start = lineno code.lines.append(line) elif clean_line.startswith(C.STUB_END): # Process "# tvm-ffi-stubgen(end)" if code is None: raise ValueError(f"Unmatched `{C.STUB_END}` found at line {lineno}") code.lineno_end = lineno code.lines.append(line) codes.append(code) code = None elif clean_line.startswith(C.STUB_TY_MAP): # Process "# tvm-ffi-stubgen(ty_map)" ty_code = CodeBlock.from_begin_line(lineno, clean_line) ty_code.lineno_end = lineno ty_code.lines.append(line) codes.append(ty_code) del ty_code elif clean_line.startswith(C.STUB_IMPORT_OBJECT): # Process "# tvm-ffi-stubgen(import-object)" imp_code = CodeBlock.from_begin_line(lineno, clean_line) imp_code.lineno_end = lineno imp_code.lines.append(line) codes.append(imp_code) del imp_code elif clean_line.startswith(C.STUB_PREFIX): raise ValueError(f"Unknown stub type at line {lineno}: {clean_line}") elif code is None: # Process a plain line outside of any stub block codes.append( CodeBlock( kind=None, param="", lineno_start=lineno, lineno_end=lineno, lines=[line], ) ) else: # Process a line inside a stub block code.lines.append(line) if code is not None: raise ValueError("Unclosed stub block at end of file") return FileInfo(path=file, lines=tuple(lines), code_blocks=codes) def reload(self) -> None: """Reload the code blocks from disk while preserving original `lines`.""" source = FileInfo.from_file(self.path) assert source is not None, f"File no longer exists or valid: {self.path}" self.code_blocks = source.code_blocks def collect_files(paths: list[Path]) -> list[FileInfo]: """Collect all files from the given paths and parse them into FileInfo objects.""" def _on_error(e: Exception) -> None: print( f"{C.TERM_RED}[Error]\n{traceback.format_exc()}{C.TERM_RESET}", end="", flush=True, ) def _walk_recursive() -> Generator[Path, None, None]: for p in paths: if p.is_file(): yield p continue for root, _dirs, files in path_walk(p, follow_symlinks=False, on_error=_on_error): for file in files: f = Path(root) / file if f.suffix.lower() not in C.DEFAULT_SOURCE_EXTS: continue yield f filenames = list(_walk_recursive()) filenames = sorted(filenames, key=lambda f: str(f)) files = [] for file in filenames: try: content = FileInfo.from_file(file) except Exception as e: _on_error(e) else: if content is not None: files.append(content) return files def path_walk( p: Path, *, top_down: bool = True, on_error: Callable[[Exception], None] | None = None, follow_symlinks: bool = False, ) -> Iterable[tuple[Path, list[str], list[str]]]: """Compat wrapper for Path.walk (3.12+) with a fallback for < 3.12.""" if not p.exists(): return # Python 3.12+ - just delegate to `Path.walk` if hasattr(p, "walk"): yield from p.walk( # ty: ignore[call-non-callable] top_down=top_down, on_error=on_error, follow_symlinks=follow_symlinks, ) return # Python < 3.12 - use `os.walk`` for root_str, dirnames, filenames in os.walk( p, topdown=top_down, onerror=on_error, followlinks=follow_symlinks, ): root = Path(root_str) # dirnames and filenames are lists of *names*, not full paths, # just like Path.walk()'s documented behavior. yield root, dirnames, filenames tvm-ffi-0.1.12/python/tvm_ffi/stub/lib_state.py000066400000000000000000000110131521067262500214150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Stateful helpers for querying TVM FFI runtime metadata.""" from __future__ import annotations import functools import heapq from collections import defaultdict from tvm_ffi._ffi_api import GetRegisteredTypeKeys from tvm_ffi.core import TypeSchema, _lookup_or_register_type_info_from_type_key from tvm_ffi.registry import get_global_func_metadata, list_global_func_names from . import consts as C from .utils import FuncInfo, NamedTypeSchema, ObjectInfo @functools.lru_cache(maxsize=None) def object_info_from_type_key(type_key: str) -> ObjectInfo: """Construct an `ObjectInfo` from an object type key.""" type_info = _lookup_or_register_type_info_from_type_key(str(type_key)) assert type_info.type_key == type_key return ObjectInfo.from_type_info(type_info) def collect_global_funcs() -> dict[str, list[FuncInfo]]: """Collect global functions from TVM FFI's global registry.""" global_funcs: dict[str, list[FuncInfo]] = {} for name in list_global_func_names(): try: prefix, _ = name.rsplit(".", 1) except ValueError: print(f"{C.TERM_YELLOW}[Skipped] Invalid name in global function: {name}{C.TERM_RESET}") else: try: global_funcs.setdefault(prefix, []).append(_func_info_from_global_name(name)) except Exception: print(f"{C.TERM_YELLOW}[Skipped] Function has no type schema: {name}{C.TERM_RESET}") for k in list(global_funcs.keys()): global_funcs[k].sort(key=lambda x: x.schema.name) return global_funcs def collect_type_keys() -> dict[str, list[str]]: """Collect registered object type keys from TVM FFI's global registry.""" global_objects: dict[str, list[str]] = {} for type_key in GetRegisteredTypeKeys(): try: prefix, _ = type_key.rsplit(".", 1) except ValueError: pass else: global_objects.setdefault(prefix, []).append(type_key) for k in list(global_objects.keys()): global_objects[k].sort() return global_objects def toposort_objects(type_keys: list[str]) -> list[ObjectInfo]: """Collect ObjectInfo objects for type keys, topologically sorted by inheritance.""" # Remove duplicates while preserving order. unique_type_keys = list(dict.fromkeys(type_keys)) infos: dict[str, ObjectInfo] = { type_key: object_info_from_type_key(type_key) for type_key in unique_type_keys } child_types: dict[str, list[str]] = defaultdict(list) in_degree: dict[str, int] = defaultdict(int) for type_key, info in infos.items(): parent_type_key = info.parent_type_key if parent_type_key in infos: child_types[parent_type_key].append(type_key) in_degree[type_key] += 1 in_degree[parent_type_key] += 0 else: in_degree[type_key] += 0 for children in child_types.values(): children.sort() queue: list[str] = [ty for ty, deg in in_degree.items() if deg == 0] heapq.heapify(queue) sorted_keys: list[str] = [] while queue: type_key = heapq.heappop(queue) sorted_keys.append(type_key) for child_type_key in child_types[type_key]: in_degree[child_type_key] -= 1 if in_degree[child_type_key] == 0: heapq.heappush(queue, child_type_key) assert len(sorted_keys) == len(infos) return [infos[type_key] for type_key in sorted_keys] @functools.lru_cache(maxsize=None) def _func_info_from_global_name(name: str) -> FuncInfo: """Construct a `FuncInfo` from a global function name.""" return FuncInfo( schema=NamedTypeSchema( name=name, schema=TypeSchema.from_json_str(get_global_func_metadata(name)["type_schema"]), ), is_member=False, ) tvm-ffi-0.1.12/python/tvm_ffi/stub/utils.py000066400000000000000000000347621521067262500206270ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Common utilities for the `tvm-ffi-stubgen` tool.""" from __future__ import annotations import dataclasses from io import StringIO from typing import Any, Callable from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr from . import consts as C def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema: """Parse a type schema from either a JSON string or an already-parsed dict.""" if isinstance(raw, dict): return TypeSchema.from_json_obj(raw) return TypeSchema.from_json_str(raw) @dataclasses.dataclass class InitConfig: """Configuration for generating new stubs. Examples -------- If we are generating type stubs for Python package `my-ffi-extension`, and the CMake target that generates the shared library is `my_ffi_extension_shared`, then we can run the following command to generate the stubs: --init-pypkg my-ffi-extension --init-lib my_ffi_extension_shared --init-prefix my_ffi_extension. """ pkg: str """Name of the Python package to generate stubs for, e.g. apache-tvm-ffi (instead of tvm_ffi)""" shared_target: str """Name of CMake target that generates the shared library, e.g. tvm_ffi_shared This is used to determine the name of the shared library file. - macOS: lib{shared_target}.dylib or lib{shared_target}.so - Linux: lib{shared_target}.so - Windows: {shared_target}.dll """ prefix: str """Only generate stubs for global function and objects with the given prefix, e.g. `tvm_ffi.`""" @dataclasses.dataclass class Options: """Command line options for stub generation.""" imports: list[str] = dataclasses.field(default_factory=list) dlls: list[str] = dataclasses.field(default_factory=list) init: InitConfig | None = None indent: int = 4 files: list[str] = dataclasses.field(default_factory=list) verbose: bool = False dry_run: bool = False @dataclasses.dataclass(frozen=True, eq=True) class ImportItem: """An import statement item.""" mod: str name: str type_checking_only: bool = False alias: str | None = None def __init__( self, name: str, type_checking_only: bool = False, alias: str | None = None, ) -> None: """Initialize an `ImportItem` with the given module name and optional alias.""" if "." in name: mod, name = name.rsplit(".", 1) for mod_prefix, mod_replacement in C.MOD_MAP.items(): if mod.startswith(mod_prefix): mod = mod.replace(mod_prefix, mod_replacement, 1) break else: mod = "" object.__setattr__(self, "mod", mod) object.__setattr__(self, "name", name) object.__setattr__(self, "type_checking_only", type_checking_only) object.__setattr__(self, "alias", alias) @property def name_with_alias(self) -> str: """Generate a string of the form `name as alias` if an alias is set, otherwise just `name`.""" return f"{self.name} as {self.alias}" if self.alias else self.name @property def full_name(self) -> str: """Generate a string of the form `mod.name` or `name` if no module is set.""" return f"{self.mod}.{self.name}" if self.mod else self.name def __repr__(self) -> str: """Generate an import statement string for this item.""" return str(self) def __str__(self) -> str: """Generate an import statement string for this item.""" if self.mod: ret = f"from {self.mod} import {self.name_with_alias}" else: ret = f"import {self.name_with_alias}" return ret @dataclasses.dataclass(init=False) class NamedTypeSchema(TypeSchema): """A type schema with an associated name.""" name: str def __init__(self, name: str, schema: TypeSchema) -> None: """Initialize a `NamedTypeSchema` with the given name and type schema.""" super().__init__(origin=schema.origin, args=schema.args) self.name = name @dataclasses.dataclass class FuncInfo: """Information of a function.""" schema: NamedTypeSchema is_member: bool @staticmethod def from_schema(name: str, schema: TypeSchema, *, is_member: bool = False) -> FuncInfo: """Construct a `FuncInfo` from a name and its type schema.""" return FuncInfo(schema=NamedTypeSchema(name=name, schema=schema), is_member=is_member) def gen(self, ty_map: Callable[[str], str], indent: int) -> str: """Generate a function signature string for this function.""" try: _, func_name = self.schema.name.rsplit(".", 1) except ValueError: func_name = self.schema.name buf = StringIO() buf.write(" " * indent) buf.write(f"def {func_name}(") if self.schema.origin != "Callable": raise ValueError(f"Expected Callable type schema, but got: {self.schema}") if not self.schema.args: ty_map("Any") buf.write("*args: Any) -> Any: ...") return buf.getvalue() arg_ret = self.schema.args[0] arg_args = self.schema.args[1:] for i, arg in enumerate(arg_args): if self.is_member and i == 0: buf.write("self, ") else: buf.write(f"_{i}: ") buf.write(arg.repr(ty_map)) buf.write(", ") if arg_args: buf.write("/") buf.write(") -> ") buf.write(arg_ret.repr(ty_map)) buf.write(": ...") return buf.getvalue() @dataclasses.dataclass class InitFieldInfo: """A field that participates in the auto-generated ``__init__``.""" name: str schema: NamedTypeSchema kw_only: bool has_default: bool @dataclasses.dataclass class ObjectInfo: """Information of an object type, including its fields and methods.""" fields: list[NamedTypeSchema] methods: list[FuncInfo] type_key: str | None = None parent_type_key: str | None = None init_fields: list[InitFieldInfo] = dataclasses.field(default_factory=list) has_init: bool = False @staticmethod def from_type_info(type_info: TypeInfo) -> ObjectInfo: """Construct an `ObjectInfo` from a `TypeInfo` instance.""" parent_type_key: str | None = None if type_info.parent_type_info is not None: parent_type_key = type_info.parent_type_info.type_key # Detect __ffi_init__ from TypeMethod or TypeAttrColumn. has_init = any(m.name == "__ffi_init__" for m in type_info.methods) if not has_init: has_init = _lookup_type_attr(type_info.type_index, "__ffi_init__") is not None # Walk parent chain (parent-first) to collect all init-eligible fields. init_fields: list[InitFieldInfo] = [] if has_init: ti: TypeInfo | None = type_info chain: list[TypeInfo] = [] while ti is not None: chain.append(ti) ti = ti.parent_type_info for ancestor_info in reversed(chain): for field in ancestor_info.fields: if not field.c_init: continue init_fields.append( InitFieldInfo( name=field.name, schema=NamedTypeSchema( name=field.name, schema=_parse_type_schema(field.metadata["type_schema"]), ), kw_only=field.c_kw_only, has_default=field.c_has_default, ) ) return ObjectInfo( fields=[ NamedTypeSchema( name=field.name, schema=_parse_type_schema(field.metadata["type_schema"]), ) for field in type_info.fields ], methods=[ FuncInfo( schema=NamedTypeSchema( name=C.FN_NAME_MAP.get(method.name, method.name), schema=_parse_type_schema(method.metadata["type_schema"]), ), is_member=not method.is_static, ) for method in type_info.methods ], type_key=type_info.type_key, parent_type_key=parent_type_key, init_fields=init_fields, has_init=has_init, ) def gen_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate field definitions for this object.""" indent_str = " " * indent return [f"{indent_str}{field.name}: {field.repr(ty_map)}" for field in self.fields] def gen_methods(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate method definitions for this object.""" indent_str = " " * indent ret = [] for method in self.methods: func_name = method.schema.name.rsplit(".", 1)[-1] if func_name == "__ffi_init__": # __ffi_init__ is installed as an instance method (self, *args, **kwargs) -> None # by _install_ffi_init_attr, regardless of the C++ static registration. ret.append(self._gen_ffi_init_from_method(method, ty_map, indent)) continue if not method.is_member: ret.append(f"{indent_str}@staticmethod") ret.append(method.gen(ty_map, indent)) return ret @staticmethod def _gen_ffi_init_from_method( method: FuncInfo, ty_map: Callable[[str], str], indent: int ) -> str: """Render ``__ffi_init__`` TypeMethod as an instance method returning None.""" indent_str = " " * indent schema = method.schema # Subclass __ffi_init__ signatures legitimately differ from the parent # (different fields → different constructor params), so suppress LSP. ignore = " # ty: ignore[invalid-method-override]" if schema.origin != "Callable" or not schema.args: ty_map("Any") return f"{indent_str}def __ffi_init__(self, *args: Any) -> None: ...{ignore}" # schema.args[0] is return type, schema.args[1:] are param types. parts: list[str] = [] for i, arg in enumerate(schema.args[1:]): parts.append(f"_{i}: {arg.repr(ty_map)}") if parts: params = ", ".join(parts) return f"{indent_str}def __ffi_init__(self, {params}, /) -> None: ...{ignore}" return f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}" def gen_ffi_init(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate a ``__ffi_init__`` stub when it's not already in TypeMethod. For types whose ``__ffi_init__`` is auto-generated by ``RegisterFFIInit`` (TypeAttrColumn only), synthesize a static-method stub from field metadata. Types that already have ``__ffi_init__`` in TypeMethod (from explicit ``refl::init<>``) get it via ``gen_methods`` instead. """ if not self.has_init: return [] # If __ffi_init__ is already in methods (from TypeMethod), gen_methods handles it. if any(m.schema.name.rsplit(".", 1)[-1] == "__ffi_init__" for m in self.methods): return [] return self._gen_ffi_init_from_fields(ty_map, indent) def gen_init(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate an ``__init__`` stub from init-eligible field metadata.""" if not self.has_init: return [] return self._gen_init_from_fields(ty_map, indent) def _format_field_params(self, ty_map: Callable[[str], str]) -> str: """Format init-eligible fields as a parameter string with defaults and kw_only.""" positional = [f for f in self.init_fields if not f.kw_only] kw_only = [f for f in self.init_fields if f.kw_only] pos_required = [f for f in positional if not f.has_default] pos_default = [f for f in positional if f.has_default] kw_required = [f for f in kw_only if not f.has_default] kw_default = [f for f in kw_only if f.has_default] parts: list[str] = [] for f in pos_required: parts.append(f"{f.name}: {f.schema.repr(ty_map)}") for f in pos_default: parts.append(f"{f.name}: {f.schema.repr(ty_map)} = ...") if kw_required or kw_default: parts.append("*") for f in kw_required: parts.append(f"{f.name}: {f.schema.repr(ty_map)}") for f in kw_default: parts.append(f"{f.name}: {f.schema.repr(ty_map)} = ...") return ", ".join(parts) def _gen_init_from_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate ``__init__`` from init-eligible field metadata (auto-generated init).""" indent_str = " " * indent params = self._format_field_params(ty_map) if params: return [f"{indent_str}def __init__(self, {params}) -> None: ..."] return [f"{indent_str}def __init__(self) -> None: ..."] def _gen_ffi_init_from_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]: """Generate ``__ffi_init__`` stub from field metadata for auto-generated init.""" indent_str = " " * indent # Subclass __ffi_init__ signatures legitimately differ from the parent # (different fields → different constructor params), so suppress LSP. ignore = " # ty: ignore[invalid-method-override]" params = self._format_field_params(ty_map) if params: return [f"{indent_str}def __ffi_init__(self, {params}) -> None: ...{ignore}"] return [f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}"] tvm-ffi-0.1.12/python/tvm_ffi/testing/000077500000000000000000000000001521067262500176015ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/testing/__init__.py000066400000000000000000000026621521067262500217200ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Testing utilities.""" from ._ffi_api import * # noqa: F403 from .testing import ( TestCompare, TestCustomCompare, TestCustomHash, TestEqWithoutHash, TestHash, TestIntPair, TestNonCopyable, TestObjectBase, TestObjectDerived, _SchemaAllTypes, _TestCxxAutoInit, _TestCxxAutoInitAllInitOff, _TestCxxAutoInitChild, _TestCxxAutoInitKwOnlyDefaults, _TestCxxAutoInitParent, _TestCxxAutoInitSimple, _TestCxxClassBase, _TestCxxClassDerived, _TestCxxClassDerivedDerived, _TestCxxInitSubset, _TestCxxKwOnly, _TestCxxNoAutoInit, add_one, create_object, make_unregistered_object, ) tvm-ffi-0.1.12/python/tvm_ffi/testing/_ffi_api.py000066400000000000000000000152131521067262500217110ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """FFI API for namespace `testing`.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from ..registry import init_ffi_api as _FFI_INIT_FUNC from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence from tvm_ffi import Device, Object, Tensor, dtype from tvm_ffi.testing import TestIntPair from typing import Any, Callable # isort: on # fmt: on # tvm-ffi-stubgen(end) # tvm-ffi-stubgen(begin): global/testing@..registry # fmt: off _FFI_INIT_FUNC("testing", __name__) if TYPE_CHECKING: def TestIntPairSum(_0: TestIntPair, /) -> int: ... def add_one(_0: int, /) -> int: ... def apply(*args: Any) -> Any: ... def echo(*args: Any) -> Any: ... def get_add_one_c_symbol() -> int: ... def get_mlir_add_one_c_symbol() -> int: ... def make_array_with_mixed(_0: Tensor, _1: int, /) -> Sequence[Any]: ... def make_array_with_tensor(_0: Tensor, /) -> Sequence[Any]: ... def make_dict_with_tensor(_0: Tensor, /) -> MutableMapping[str, Any]: ... def make_empty_array_with_tensor_input(_0: Tensor, /) -> Sequence[Any]: ... def make_list_with_tensor(_0: Tensor, _1: int, /) -> MutableSequence[Any]: ... def make_map_with_tensor(_0: Tensor, /) -> Mapping[str, Any]: ... def make_nested_array_with_tensor(_0: Tensor, /) -> Sequence[Any]: ... def make_nested_map_with_tensor(_0: Tensor, _1: Tensor, /) -> Mapping[str, Any]: ... def make_unregistered_object() -> Object: ... def nop(*args: Any) -> Any: ... def object_use_count(_0: Object, /) -> int: ... def optional_tensor_view_has_value(_0: Tensor | None, /) -> bool: ... def run_check_signal(_0: int, /) -> None: ... def schema_arr_map_opt(_0: Sequence[int | None], _1: Mapping[str, Sequence[int]], _2: str | None, /) -> Mapping[str, Sequence[int]]: ... def schema_id_any(_0: Any, /) -> Any: ... def schema_id_arr(_0: Sequence[Any], /) -> Sequence[Any]: ... def schema_id_arr_int(_0: Sequence[int], /) -> Sequence[int]: ... def schema_id_arr_obj(_0: Sequence[Object], /) -> Sequence[Object]: ... def schema_id_arr_str(_0: Sequence[str], /) -> Sequence[str]: ... def schema_id_bool(_0: bool, /) -> bool: ... def schema_id_bytes(_0: bytes, /) -> bytes: ... def schema_id_device(_0: Device, /) -> Device: ... def schema_id_dict_str_int(_0: MutableMapping[str, int], /) -> MutableMapping[str, int]: ... def schema_id_dict_str_str(_0: MutableMapping[str, str], /) -> MutableMapping[str, str]: ... def schema_id_dltensor(_0: Tensor, /) -> Tensor: ... def schema_id_dtype(_0: dtype, /) -> dtype: ... def schema_id_float(_0: float, /) -> float: ... def schema_id_func(_0: Callable[..., Any], /) -> Callable[..., Any]: ... def schema_id_func_typed(_0: Callable[[int, float, Callable[..., Any]], None], /) -> Callable[[int, float, Callable[..., Any]], None]: ... def schema_id_int(_0: int, /) -> int: ... def schema_id_list_int(_0: MutableSequence[int], /) -> MutableSequence[int]: ... def schema_id_list_obj(_0: MutableSequence[Object], /) -> MutableSequence[Object]: ... def schema_id_list_str(_0: MutableSequence[str], /) -> MutableSequence[str]: ... def schema_id_map(_0: Mapping[Any, Any], /) -> Mapping[Any, Any]: ... def schema_id_map_str_int(_0: Mapping[str, int], /) -> Mapping[str, int]: ... def schema_id_map_str_obj(_0: Mapping[str, Object], /) -> Mapping[str, Object]: ... def schema_id_map_str_str(_0: Mapping[str, str], /) -> Mapping[str, str]: ... def schema_id_object(_0: Object, /) -> Object: ... def schema_id_opt_int(_0: int | None, /) -> int | None: ... def schema_id_opt_obj(_0: Object | None, /) -> Object | None: ... def schema_id_opt_str(_0: str | None, /) -> str | None: ... def schema_id_string(_0: str, /) -> str: ... def schema_id_tensor(_0: Tensor, /) -> Tensor: ... def schema_id_variant_int_str(_0: int | str, /) -> int | str: ... def schema_no_args() -> int: ... def schema_no_args_no_return() -> None: ... def schema_no_return(_0: int, /) -> None: ... def schema_packed(*args: Any) -> Any: ... def schema_tensor_view_input(_0: Tensor, /) -> None: ... def schema_variant_mix(_0: int | str | Sequence[int], /) -> int | str | Sequence[int]: ... def test_raise_error(_0: str, _1: str, /) -> None: ... # fmt: on # tvm-ffi-stubgen(end) __all__ = [ # tvm-ffi-stubgen(begin): __all__ "TestIntPairSum", "add_one", "apply", "echo", "get_add_one_c_symbol", "get_mlir_add_one_c_symbol", "make_array_with_mixed", "make_array_with_tensor", "make_dict_with_tensor", "make_empty_array_with_tensor_input", "make_list_with_tensor", "make_map_with_tensor", "make_nested_array_with_tensor", "make_nested_map_with_tensor", "make_unregistered_object", "nop", "object_use_count", "optional_tensor_view_has_value", "run_check_signal", "schema_arr_map_opt", "schema_id_any", "schema_id_arr", "schema_id_arr_int", "schema_id_arr_obj", "schema_id_arr_str", "schema_id_bool", "schema_id_bytes", "schema_id_device", "schema_id_dict_str_int", "schema_id_dict_str_str", "schema_id_dltensor", "schema_id_dtype", "schema_id_float", "schema_id_func", "schema_id_func_typed", "schema_id_int", "schema_id_list_int", "schema_id_list_obj", "schema_id_list_str", "schema_id_map", "schema_id_map_str_int", "schema_id_map_str_obj", "schema_id_map_str_str", "schema_id_object", "schema_id_opt_int", "schema_id_opt_obj", "schema_id_opt_str", "schema_id_string", "schema_id_tensor", "schema_id_variant_int_str", "schema_no_args", "schema_no_args_no_return", "schema_no_return", "schema_packed", "schema_tensor_view_input", "schema_variant_mix", "test_raise_error", # tvm-ffi-stubgen(end) ] tvm-ffi-0.1.12/python/tvm_ffi/testing/testing.py000066400000000000000000000376351521067262500216460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Testing utilities.""" # ruff: noqa: D102 # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Mapping, Sequence from tvm_ffi import Device, dtype from typing import Any # isort: on # fmt: on # tvm-ffi-stubgen(end) import sys from typing import Any, ClassVar import pytest from tvm_ffi import Object, get_global_func from tvm_ffi.dataclasses import c_class from .. import _ffi_api from .. import core as tvm_ffi_core requires_py39 = pytest.mark.skipif( sys.version_info < (3, 9), reason="requires Python 3.9+", ) requires_py310 = pytest.mark.skipif( sys.version_info < (3, 10), reason="requires Python 3.10+", ) requires_py312 = pytest.mark.skipif( sys.version_info < (3, 12), reason="requires Python 3.12+", ) requires_py313 = pytest.mark.skipif( sys.version_info < (3, 13), reason="requires Python 3.13+", ) @c_class("testing.TestObjectBase") class TestObjectBase(Object): """Test object base class.""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestObjectBase # fmt: off v_i64: int v_f64: float v_str: str if TYPE_CHECKING: def __init__(self, v_i64: int = ..., v_f64: float = ..., v_str: str = ...) -> None: ... def __ffi_init__(self, v_i64: int = ..., v_f64: float = ..., v_str: str = ...) -> None: ... # ty: ignore[invalid-method-override] def add_i64(self, _1: int, /) -> int: ... # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestIntPair") class TestIntPair(Object): """Test Int Pair.""" __test__: ClassVar[bool] = False # tvm-ffi-stubgen(begin): object/testing.TestIntPair # fmt: off a: int b: int if TYPE_CHECKING: def __init__(self, a: int, b: int) -> None: ... def __ffi_init__(self, _0: int, _1: int, /) -> None: ... # ty: ignore[invalid-method-override] def sum(self, /) -> int: ... # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestObjectDerived") class TestObjectDerived(TestObjectBase): """Test object derived class.""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestObjectDerived # fmt: off v_map: Mapping[Any, Any] v_array: Sequence[Any] if TYPE_CHECKING: def __init__(self, v_map: Mapping[Any, Any], v_array: Sequence[Any], v_i64: int = ..., v_f64: float = ..., v_str: str = ...) -> None: ... def __ffi_init__(self, v_map: Mapping[Any, Any], v_array: Sequence[Any], v_i64: int = ..., v_f64: float = ..., v_str: str = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestNonCopyable") class TestNonCopyable(Object): """Test object with deleted copy constructor.""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestNonCopyable # fmt: off value: int if TYPE_CHECKING: def __init__(self, value: int) -> None: ... def __ffi_init__(self, _0: int, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.SchemaAllTypes") class _SchemaAllTypes: # tvm-ffi-stubgen(ty-map): testing.SchemaAllTypes -> testing._SchemaAllTypes # tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes # fmt: off v_bool: bool v_int: int v_float: float v_device: Device v_dtype: dtype v_string: str v_bytes: bytes v_opt_int: int | None v_opt_str: str | None v_arr_int: Sequence[int] v_arr_str: Sequence[str] v_map_str_int: Mapping[str, int] v_map_str_arr_int: Mapping[str, Sequence[int]] v_variant: str | Sequence[int] | Mapping[str, int] v_opt_arr_variant: Sequence[int | str] | None if TYPE_CHECKING: def add_int(self, _1: int, /) -> int: ... def append_int(self, _1: Sequence[int], _2: int, /) -> Sequence[int]: ... def maybe_concat(self, _1: str | None, _2: str | None, /) -> str | None: ... def merge_map(self, _1: Mapping[str, Sequence[int]], _2: Mapping[str, Sequence[int]], /) -> Mapping[str, Sequence[int]]: ... @staticmethod def make_with(_0: int, _1: float, _2: str, /) -> _SchemaAllTypes: ... # fmt: on # tvm-ffi-stubgen(end) def create_object(type_key: str, **kwargs: Any) -> Object: """Make an object by reflection. Parameters ---------- type_key The type key of the object. kwargs The keyword arguments to the object. Returns ------- obj The created object. Note ---- This function is only used for testing purposes and should not be used in other cases. """ args = [type_key] for k, v in kwargs.items(): args.append(k) args.append(v) return _ffi_api.MakeObjectFromPackedArgs(*args) def make_unregistered_object() -> Object: """Return an object whose type is not registered on the Python side.""" return get_global_func("testing.make_unregistered_object")() def add_one(x: int) -> int: """Add one to the input integer.""" return get_global_func("testing.add_one")(x) @c_class("testing.TestCompare") class TestCompare(Object): """Test object with Compare(false) on ignored_field.""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestCompare # fmt: off key: int name: str ignored_field: int if TYPE_CHECKING: def __init__(self, key: int, name: str, ignored_field: int) -> None: ... def __ffi_init__(self, _0: int, _1: str, _2: int, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestHash") class TestHash(Object): """Test object with Hash(false) on hash_ignored.""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestHash # fmt: off key: int name: str hash_ignored: int if TYPE_CHECKING: def __init__(self, key: int, name: str, hash_ignored: int) -> None: ... def __ffi_init__(self, _0: int, _1: str, _2: int, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCustomHash") class TestCustomHash(Object): """Test object with custom __ffi_hash__ hook (hashes only key).""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestCustomHash # fmt: off key: int label: str if TYPE_CHECKING: def __init__(self, key: int, label: str) -> None: ... def __ffi_init__(self, _0: int, _1: str, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCustomCompare") class TestCustomCompare(Object): """Test object with custom __ffi_eq__/__ffi_compare__ hooks (compares only key).""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestCustomCompare # fmt: off key: int label: str if TYPE_CHECKING: def __init__(self, key: int, label: str) -> None: ... def __ffi_init__(self, _0: int, _1: str, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestEqWithoutHash") class TestEqWithoutHash(Object): """Test object with __ffi_eq__ but no __ffi_hash__ (exercises hash guard).""" __test__ = False # tvm-ffi-stubgen(begin): object/testing.TestEqWithoutHash # fmt: off key: int label: str if TYPE_CHECKING: def __init__(self, key: int, label: str) -> None: ... def __ffi_init__(self, _0: int, _1: str, /) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxClassBase") class _TestCxxClassBase(Object): # tvm-ffi-stubgen(ty-map): testing.TestCxxClassBase -> testing._TestCxxClassBase # tvm-ffi-stubgen(begin): object/testing.TestCxxClassBase # fmt: off v_i64: int v_i32: int if TYPE_CHECKING: def __init__(self, v_i64: int, v_i32: int) -> None: ... def __ffi_init__(self, v_i64: int, v_i32: int) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) not_field_1 = 1 not_field_2: ClassVar[int] = 2 def __init__(self, v_i64: int, v_i32: int) -> None: ti = getattr(type(self), "__tvm_ffi_type_info__") ffi_init = tvm_ffi_core._lookup_type_attr(ti.type_index, "__ffi_init__") self.__init_handle_by_constructor__(ffi_init, v_i64 + 1, v_i32 + 2) @c_class("testing.TestCxxClassDerived", eq=True, order=True, unsafe_hash=True) class _TestCxxClassDerived(_TestCxxClassBase): # tvm-ffi-stubgen(ty-map): testing.TestCxxClassDerived -> testing._TestCxxClassDerived # tvm-ffi-stubgen(begin): object/testing.TestCxxClassDerived # fmt: off v_f64: float v_f32: float if TYPE_CHECKING: def __init__(self, v_i64: int, v_i32: int, v_f64: float, v_f32: float = ...) -> None: ... def __ffi_init__(self, v_i64: int, v_i32: int, v_f64: float, v_f32: float = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxClassDerivedDerived") class _TestCxxClassDerivedDerived(_TestCxxClassDerived): # tvm-ffi-stubgen(ty-map): testing.TestCxxClassDerivedDerived -> testing._TestCxxClassDerivedDerived # tvm-ffi-stubgen(begin): object/testing.TestCxxClassDerivedDerived # fmt: off v_str: str v_bool: bool if TYPE_CHECKING: def __init__(self, v_i64: int, v_i32: int, v_f64: float, v_bool: bool, v_f32: float = ..., v_str: str = ...) -> None: ... def __ffi_init__(self, v_i64: int, v_i32: int, v_f64: float, v_bool: bool, v_f32: float = ..., v_str: str = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxInitSubset") class _TestCxxInitSubset(Object): # tvm-ffi-stubgen(ty-map): testing.TestCxxInitSubset -> testing._TestCxxInitSubset # tvm-ffi-stubgen(begin): object/testing.TestCxxInitSubset # fmt: off required_field: int optional_field: int note: str if TYPE_CHECKING: def __init__(self, required_field: int) -> None: ... def __ffi_init__(self, required_field: int) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxKwOnly") class _TestCxxKwOnly(Object): # tvm-ffi-stubgen(ty-map): testing.TestCxxKwOnly -> testing._TestCxxKwOnly # tvm-ffi-stubgen(begin): object/testing.TestCxxKwOnly # fmt: off x: int y: int z: int w: int if TYPE_CHECKING: def __init__(self, *, x: int, y: int, z: int, w: int = ...) -> None: ... def __ffi_init__(self, *, x: int, y: int, z: int, w: int = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInit") class _TestCxxAutoInit(Object): """Test object with init(false) on b and KwOnly(true) on c.""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInit -> testing._TestCxxAutoInit # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInit # fmt: off a: int b: int c: int d: int if TYPE_CHECKING: def __init__(self, a: int, d: int = ..., *, c: int) -> None: ... def __ffi_init__(self, a: int, d: int = ..., *, c: int) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInitSimple") class _TestCxxAutoInitSimple(Object): """Test object with all fields positional (no init/KwOnly traits).""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitSimple -> testing._TestCxxAutoInitSimple # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitSimple # fmt: off x: int y: int if TYPE_CHECKING: def __init__(self, x: int, y: int) -> None: ... def __ffi_init__(self, x: int, y: int) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInitAllInitOff") class _TestCxxAutoInitAllInitOff(Object): """Test object with all fields excluded from auto-init (init(false)).""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitAllInitOff -> testing._TestCxxAutoInitAllInitOff # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitAllInitOff # fmt: off x: int y: int z: int if TYPE_CHECKING: def __init__(self) -> None: ... def __ffi_init__(self) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInitKwOnlyDefaults") class _TestCxxAutoInitKwOnlyDefaults(Object): """Test object with mixed positional/kw-only/default/init=False fields.""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitKwOnlyDefaults -> testing._TestCxxAutoInitKwOnlyDefaults # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitKwOnlyDefaults # fmt: off p_required: int p_default: int k_required: int k_default: int hidden: int if TYPE_CHECKING: def __init__(self, p_required: int, p_default: int = ..., *, k_required: int, k_default: int = ...) -> None: ... def __ffi_init__(self, p_required: int, p_default: int = ..., *, k_required: int, k_default: int = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxNoAutoInit", init=False) class _TestCxxNoAutoInit(Object): """Test object with init(false) at class level — no __ffi_init__ generated.""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxNoAutoInit -> testing._TestCxxNoAutoInit # tvm-ffi-stubgen(begin): object/testing.TestCxxNoAutoInit # fmt: off x: int y: int # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInitParent") class _TestCxxAutoInitParent(Object): """Parent object for inheritance auto-init tests.""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitParent -> testing._TestCxxAutoInitParent # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitParent # fmt: off parent_required: int parent_default: int if TYPE_CHECKING: def __init__(self, parent_required: int, parent_default: int = ...) -> None: ... def __ffi_init__(self, parent_required: int, parent_default: int = ...) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) @c_class("testing.TestCxxAutoInitChild") class _TestCxxAutoInitChild(_TestCxxAutoInitParent): """Child object for inheritance auto-init tests.""" __test__ = False # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitChild -> testing._TestCxxAutoInitChild # tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitChild # fmt: off child_required: int child_kw_only: int if TYPE_CHECKING: def __init__(self, parent_required: int, child_required: int, parent_default: int = ..., *, child_kw_only: int) -> None: ... def __ffi_init__(self, parent_required: int, child_required: int, parent_default: int = ..., *, child_kw_only: int) -> None: ... # ty: ignore[invalid-method-override] # fmt: on # tvm-ffi-stubgen(end) tvm-ffi-0.1.12/python/tvm_ffi/utils/000077500000000000000000000000001521067262500172645ustar00rootroot00000000000000tvm-ffi-0.1.12/python/tvm_ffi/utils/__init__.py000066400000000000000000000016021521067262500213740ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities used by the tvm_ffi Python package.""" from . import kwargs_wrapper from .lockfile import FileLock tvm-ffi-0.1.12/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py000066400000000000000000000736201521067262500260300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Build Torch C DLPack Addon.""" from __future__ import annotations import argparse import os import shutil import subprocess import sys import sysconfig import tempfile from collections.abc import Sequence from pathlib import Path import torch import torch.torch_version import torch.utils.cpp_extension # Important: to avoid cyclic dependency, we avoid import tvm_ffi names at top level here. IS_WINDOWS = sys.platform == "win32" IS_DARWIN = sys.platform == "darwin" cpp_source = """ #include #include #include #include #ifdef BUILD_WITH_CUDA #include #endif #ifdef BUILD_WITH_ROCM #include #include #endif using namespace std; namespace at { namespace { DLDataType getDLDataTypeForDLPackv1(const Tensor& t) { DLDataType dtype; dtype.lanes = 1; dtype.bits = t.element_size() * 8; switch (t.scalar_type()) { case ScalarType::UInt1: case ScalarType::UInt2: case ScalarType::UInt3: case ScalarType::UInt4: case ScalarType::UInt5: case ScalarType::UInt6: case ScalarType::UInt7: case ScalarType::Byte: case ScalarType::UInt16: case ScalarType::UInt32: case ScalarType::UInt64: dtype.code = DLDataTypeCode::kDLUInt; break; #if (TORCH_VERSION_MAJOR > 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 6) case ScalarType::Int1: case ScalarType::Int2: case ScalarType::Int3: case ScalarType::Int4: case ScalarType::Int5: case ScalarType::Int6: case ScalarType::Int7: #endif case ScalarType::Char: case ScalarType::Short: case ScalarType::Int: case ScalarType::Long: dtype.code = DLDataTypeCode::kDLInt; break; case ScalarType::Half: case ScalarType::Float: case ScalarType::Double: dtype.code = DLDataTypeCode::kDLFloat; break; case ScalarType::Bool: dtype.code = DLDataTypeCode::kDLBool; break; case ScalarType::ComplexHalf: case ScalarType::ComplexFloat: case ScalarType::ComplexDouble: dtype.code = DLDataTypeCode::kDLComplex; break; case ScalarType::BFloat16: dtype.code = DLDataTypeCode::kDLBfloat; break; case ScalarType::Float8_e5m2: dtype.code = DLDataTypeCode::kDLFloat8_e5m2; break; case ScalarType::Float8_e5m2fnuz: dtype.code = DLDataTypeCode::kDLFloat8_e5m2fnuz; break; case ScalarType::Float8_e4m3fn: dtype.code = DLDataTypeCode::kDLFloat8_e4m3fn; break; case ScalarType::Float8_e4m3fnuz: dtype.code = DLDataTypeCode::kDLFloat8_e4m3fnuz; break; #if (TORCH_VERSION_MAJOR > 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 7) case ScalarType::Float8_e8m0fnu: dtype.code = DLDataTypeCode::kDLFloat8_e8m0fnu; break; #endif #if (TORCH_VERSION_MAJOR > 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 8) case ScalarType::Float4_e2m1fn_x2: dtype.code = DLDataTypeCode::kDLFloat4_e2m1fn; dtype.lanes = 2; dtype.bits = 4; break; #endif default: TORCH_CHECK(false, "Unsupported scalar type: "); } return dtype; } DLDevice torchDeviceToDLDeviceForDLPackv1(at::Device device) { DLDevice ctx; ctx.device_id = (device.is_cuda() || device.is_privateuseone()) ? static_cast(static_cast(device.index())) : 0; switch (device.type()) { case DeviceType::CPU: ctx.device_type = DLDeviceType::kDLCPU; break; case DeviceType::CUDA: #ifdef USE_ROCM ctx.device_type = DLDeviceType::kDLROCM; #else ctx.device_type = DLDeviceType::kDLCUDA; #endif break; case DeviceType::OPENCL: ctx.device_type = DLDeviceType::kDLOpenCL; break; case DeviceType::HIP: ctx.device_type = DLDeviceType::kDLROCM; break; case DeviceType::XPU: ctx.device_type = DLDeviceType::kDLOneAPI; ctx.device_id = at::detail::getXPUHooks().getGlobalIdxFromDevice(device); break; case DeviceType::MAIA: ctx.device_type = DLDeviceType::kDLMAIA; break; case DeviceType::PrivateUse1: ctx.device_type = DLDeviceType::kDLExtDev; break; case DeviceType::MPS: ctx.device_type = DLDeviceType::kDLMetal; break; default: TORCH_CHECK(false, "Cannot pack tensors on " + device.str()); } return ctx; } template struct ATenDLMTensor { Tensor handle; T tensor{}; }; template void deleter(T* arg) { delete static_cast*>(arg->manager_ctx); } // Adds version information for DLManagedTensorVersioned. // This is a no-op for the other types. template void fillVersion(T* tensor) {} template <> void fillVersion(DLManagedTensorVersioned* tensor) { tensor->flags = 0; tensor->version.major = DLPACK_MAJOR_VERSION; tensor->version.minor = DLPACK_MINOR_VERSION; } // This function returns a shared_ptr to memory managed DLpack tensor // constructed out of ATen tensor template T* toDLPackImpl(const Tensor& src) { ATenDLMTensor* atDLMTensor(new ATenDLMTensor); atDLMTensor->handle = src; atDLMTensor->tensor.manager_ctx = atDLMTensor; atDLMTensor->tensor.deleter = &deleter; atDLMTensor->tensor.dl_tensor.data = src.data_ptr(); atDLMTensor->tensor.dl_tensor.device = torchDeviceToDLDeviceForDLPackv1(src.device()); atDLMTensor->tensor.dl_tensor.ndim = static_cast(src.dim()); atDLMTensor->tensor.dl_tensor.dtype = getDLDataTypeForDLPackv1(src); atDLMTensor->tensor.dl_tensor.shape = const_cast(src.sizes().data()); atDLMTensor->tensor.dl_tensor.strides = const_cast(src.strides().data()); atDLMTensor->tensor.dl_tensor.byte_offset = 0; fillVersion(&atDLMTensor->tensor); return &(atDLMTensor->tensor); } static Device getATenDeviceForDLPackv1(DLDeviceType type, c10::DeviceIndex index, void* data = nullptr) { switch (type) { case DLDeviceType::kDLCPU: return at::Device(DeviceType::CPU); #ifndef USE_ROCM // if we are compiled under HIP, we cannot do cuda case DLDeviceType::kDLCUDA: return at::Device(DeviceType::CUDA, index); #endif case DLDeviceType::kDLOpenCL: return at::Device(DeviceType::OPENCL, index); case DLDeviceType::kDLROCM: #ifdef USE_ROCM // this looks funny, we need to return CUDA here to masquerade return at::Device(DeviceType::CUDA, index); #else return at::Device(DeviceType::HIP, index); #endif case DLDeviceType::kDLOneAPI: TORCH_CHECK(data != nullptr, "Can't get ATen device for XPU without XPU data."); return at::detail::getXPUHooks().getDeviceFromPtr(data); case DLDeviceType::kDLMAIA: return at::Device(DeviceType::MAIA, index); case DLDeviceType::kDLExtDev: return at::Device(DeviceType::PrivateUse1, index); case DLDeviceType::kDLMetal: return at::Device(DeviceType::MPS, index); default: TORCH_CHECK(false, "Unsupported device_type: ", std::to_string(type)); } } ScalarType toScalarTypeForDLPackv1(const DLDataType& dtype) { ScalarType stype = ScalarType::Undefined; #if TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 8 if (dtype.code != DLDataTypeCode::kDLFloat4_e2m1fn) { TORCH_CHECK(dtype.lanes == 1, "ATen does not support lanes != 1 for dtype code", std::to_string(dtype.code)); } #endif switch (dtype.code) { case DLDataTypeCode::kDLUInt: switch (dtype.bits) { case 8: stype = ScalarType::Byte; break; case 16: stype = ScalarType::UInt16; break; case 32: stype = ScalarType::UInt32; break; case 64: stype = ScalarType::UInt64; break; default: TORCH_CHECK(false, "Unsupported kUInt bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLInt: switch (dtype.bits) { case 8: stype = ScalarType::Char; break; case 16: stype = ScalarType::Short; break; case 32: stype = ScalarType::Int; break; case 64: stype = ScalarType::Long; break; default: TORCH_CHECK(false, "Unsupported kInt bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLFloat: switch (dtype.bits) { case 16: stype = ScalarType::Half; break; case 32: stype = ScalarType::Float; break; case 64: stype = ScalarType::Double; break; default: TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLBfloat: switch (dtype.bits) { case 16: stype = ScalarType::BFloat16; break; default: TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLComplex: switch (dtype.bits) { case 32: stype = ScalarType::ComplexHalf; break; case 64: stype = ScalarType::ComplexFloat; break; case 128: stype = ScalarType::ComplexDouble; break; default: TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLBool: switch (dtype.bits) { case 8: stype = ScalarType::Bool; break; default: TORCH_CHECK(false, "Unsupported kDLBool bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLFloat8_e5m2: switch (dtype.bits) { case 8: stype = ScalarType::Float8_e5m2; break; default: TORCH_CHECK(false, "Unsupported kDLFloat8_e5m2 bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLFloat8_e5m2fnuz: switch (dtype.bits) { case 8: stype = ScalarType::Float8_e5m2fnuz; break; default: TORCH_CHECK(false, "Unsupported kDLFloat8_e5m2fnuz bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLFloat8_e4m3fn: switch (dtype.bits) { case 8: stype = ScalarType::Float8_e4m3fn; break; default: TORCH_CHECK(false, "Unsupported kDLFloat8_e4m3fn bits ", std::to_string(dtype.bits)); } break; case DLDataTypeCode::kDLFloat8_e4m3fnuz: switch (dtype.bits) { case 8: stype = ScalarType::Float8_e4m3fnuz; break; default: TORCH_CHECK(false, "Unsupported kDLFloat8_e4m3fnuz bits ", std::to_string(dtype.bits)); } break; #if (TORCH_VERSION_MAJOR > 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 7) case DLDataTypeCode::kDLFloat8_e8m0fnu: switch (dtype.bits) { case 8: stype = ScalarType::Float8_e8m0fnu; break; default: TORCH_CHECK(false, "Unsupported kDLFloat8_e8m0fnu bits ", std::to_string(dtype.bits)); } break; #endif #if (TORCH_VERSION_MAJOR > 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 8) case DLDataTypeCode::kDLFloat4_e2m1fn: switch (dtype.bits) { case 4: switch (dtype.lanes) { case 2: stype = ScalarType::Float4_e2m1fn_x2; break; default: TORCH_CHECK(false, "Unsupported kDLFloat4_e2m1fn lanes ", std::to_string(dtype.lanes)); } break; default: TORCH_CHECK(false, "Unsupported kDLFloat4_e2m1fn bits ", std::to_string(dtype.bits)); } break; #endif default: TORCH_CHECK(false, "Unsupported code ", std::to_string(dtype.code)); } return stype; } // This function constructs a Tensor from a memory managed DLPack which // may be represented as either: DLManagedTensor and DLManagedTensorVersioned. template at::Tensor fromDLPackImpl(T* src, std::function deleter) { if (!deleter) { deleter = [src](void* self [[maybe_unused]]) { if (src->deleter) { src->deleter(src); } }; } DLTensor& dl_tensor = src->dl_tensor; Device device = getATenDeviceForDLPackv1(dl_tensor.device.device_type, dl_tensor.device.device_id, dl_tensor.data); ScalarType stype = toScalarTypeForDLPackv1(dl_tensor.dtype); if (!dl_tensor.strides) { return at::from_blob(dl_tensor.data, IntArrayRef(dl_tensor.shape, dl_tensor.ndim), std::move(deleter), at::device(device).dtype(stype), {device}); } return at::from_blob(dl_tensor.data, IntArrayRef(dl_tensor.shape, dl_tensor.ndim), IntArrayRef(dl_tensor.strides, dl_tensor.ndim), deleter, at::device(device).dtype(stype), {device}); } void toDLPackNonOwningImpl(const Tensor& tensor, DLTensor& out) { // Fill in the pre-allocated DLTensor struct with direct pointers // This is a non-owning conversion - the caller owns the tensor // and must keep it alive for the duration of DLTensor usage out.data = tensor.data_ptr(); out.device = torchDeviceToDLDeviceForDLPackv1(tensor.device()); out.ndim = static_cast(tensor.dim()); out.dtype = getDLDataTypeForDLPackv1(tensor); // sizes() and strides() return pointers to TensorImpl's stable storage // which remains valid as long as the tensor is alive out.shape = const_cast(tensor.sizes().data()); out.strides = const_cast(tensor.strides().data()); out.byte_offset = 0; } } // namespace } // namespace at struct TorchDLPackExchangeAPI : public DLPackExchangeAPI { TorchDLPackExchangeAPI() { header.version.major = DLPACK_MAJOR_VERSION; header.version.minor = DLPACK_MINOR_VERSION; header.prev_api = nullptr; managed_tensor_allocator = ManagedTensorAllocator; managed_tensor_from_py_object_no_sync = ManagedTensorFromPyObjectNoSync; managed_tensor_to_py_object_no_sync = ManagedTensorToPyObjectNoSync; dltensor_from_py_object_no_sync = DLTensorFromPyObjectNoSync; current_work_stream = CurrentWorkStream; } static const DLPackExchangeAPI* Global() { static TorchDLPackExchangeAPI inst; return &inst; } private: static int DLTensorFromPyObjectNoSync(void* py_obj, DLTensor* out) { try { // Use handle (non-owning) to avoid unnecessary refcount operations py::handle handle(static_cast(py_obj)); at::Tensor tensor = handle.cast(); at::toDLPackNonOwningImpl(tensor, *out); return 0; } catch (const std::exception& e) { PyErr_SetString(PyExc_RuntimeError, e.what()); return -1; } } static int ManagedTensorFromPyObjectNoSync(void* py_obj, DLManagedTensorVersioned** out) { try { py::handle handle(static_cast(py_obj)); at::Tensor tensor = handle.cast(); *out = at::toDLPackImpl(tensor); return 0; } catch (const std::exception& e) { PyErr_SetString(PyExc_RuntimeError, e.what()); return -1; } } // Get current CUDA/ROCm work stream static int CurrentWorkStream(DLDeviceType device_type, int32_t device_id, void** out_stream) { try { #ifdef BUILD_WITH_ROCM if (device_type == kDLROCM || device_type == kDLCUDA) { *out_stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(device_id).stream(); return 0; } #endif #ifdef BUILD_WITH_CUDA if (device_type == kDLCUDA) { *out_stream = at::cuda::getCurrentCUDAStream(device_id).stream(); return 0; } #endif // For CPU and other devices, return NULL (no stream concept) *out_stream = nullptr; return 0; } catch (const std::exception& e) { PyErr_SetString(PyExc_RuntimeError, e.what()); return -1; } } static int ManagedTensorToPyObjectNoSync(DLManagedTensorVersioned* src, void** py_obj_out) { try { at::Tensor tensor = at::fromDLPackImpl(src, nullptr); *py_obj_out = THPVariable_Wrap(tensor); return 0; } catch (const std::exception& e) { PyErr_SetString(PyExc_RuntimeError, e.what()); return -1; } } static int ManagedTensorAllocator(DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx, void (*SetError)(void* error_ctx, const char* kind, const char* message)) { try { at::IntArrayRef shape(prototype->shape, prototype->shape + prototype->ndim); at::TensorOptions options = at::TensorOptions() .dtype(at::toScalarType(prototype->dtype)) .device(at::getATenDeviceForDLPackv1(prototype->device.device_type, prototype->device.device_id)); at::Tensor tensor = at::empty(shape, options); *out = at::toDLPackImpl(tensor); return 0; } catch (const std::exception& e) { SetError(error_ctx, "MemoryError", e.what()); return -1; } } }; // defien a cross-platgorm macro to export the symbol #ifdef _WIN32 #define DLL_EXPORT __declspec(dllexport) #else #define DLL_EXPORT __attribute__((visibility("default"))) #endif extern "C" DLL_EXPORT int64_t TorchDLPackExchangeAPIPtr() { return reinterpret_cast(TorchDLPackExchangeAPI::Global()); } """ def parse_env_flags(env_var_name: str) -> list[str]: env_flags = os.environ.get(env_var_name) if env_flags: try: import shlex # noqa: PLC0415 return shlex.split(env_flags) except ValueError as e: print( f"Warning: Could not parse {env_var_name} with shlex: {e}. Falling back to simple split.", file=sys.stderr, ) return env_flags.split() return [] def _run_build_on_linux_like( build_dir: Path, libname: str, source_path: Path, extra_cflags: Sequence[str], extra_ldflags: Sequence[str], extra_include_paths: Sequence[str], ) -> None: """Build the module directly by invoking compiler commands (non-Windows only).""" from tvm_ffi.libinfo import find_dlpack_include_path # noqa: PLC0415 default_cflags = ["-std=c++17", "-fPIC", "-O3", "-fvisibility=hidden"] # Platform-specific linker flags if IS_DARWIN: # macOS uses @loader_path instead of $ORIGIN default_ldflags = ["-shared", "-Wl,-rpath,@loader_path"] else: # Linux uses $ORIGIN default_ldflags = ["-shared", "-Wl,-rpath,$ORIGIN"] cflags = default_cflags + [flag.strip() for flag in extra_cflags] ldflags = default_ldflags + [flag.strip() for flag in extra_ldflags] include_paths = [find_dlpack_include_path()] + [ str(Path(path).resolve()) for path in extra_include_paths ] # append include paths for path in include_paths: cflags.extend(["-I", str(path)]) # Get compiler and build paths cxx = os.environ.get("CXX", "c++") source_path_resolved = source_path.resolve() lib_path = build_dir / libname # Build command: compile and link in one step build_cmd = [cxx, *cflags, str(source_path_resolved), *ldflags, "-o", str(lib_path)] # Run build command status = subprocess.run(build_cmd, cwd=str(build_dir), capture_output=True, check=False) if status.returncode != 0: msg = [f"Build failed with status {status.returncode}"] if status.stdout: msg.append(f"stdout:\n{status.stdout.decode('utf-8')}") if status.stderr: msg.append(f"stderr:\n{status.stderr.decode('utf-8')}") raise RuntimeError("\n".join(msg)) def _generate_ninja_build_windows( build_dir: Path, libname: str, source_path: Path, extra_cflags: Sequence[str], extra_ldflags: Sequence[str], extra_include_paths: Sequence[str], ) -> None: """Generate the content of build.ninja for building the module on Windows.""" from tvm_ffi.libinfo import find_dlpack_include_path # noqa: PLC0415 default_cflags = [ "/std:c++17", "/MD", "/wd4819", "/wd4251", "/wd4244", "/wd4267", "/wd4275", "/wd4018", "/wd4190", "/wd4624", "/wd4067", "/wd4068", "/EHsc", ] default_ldflags = ["/DLL"] cflags = default_cflags + [flag.strip() for flag in extra_cflags] ldflags = default_ldflags + [flag.strip() for flag in extra_ldflags] include_paths = [find_dlpack_include_path()] + [ str(Path(path).resolve()) for path in extra_include_paths ] # append include paths for path in include_paths: path_str = str(path) if " " in path_str: path_str = f'"{path_str}"' path_str = path_str.replace(":", "$:") cflags.append(f"-I{path_str}") # flags ninja = [] ninja.append("ninja_required_version = 1.3") ninja.append("cxx = {}".format(os.environ.get("CXX", "cl"))) ninja.append("cflags = {}".format(" ".join(cflags))) ninja.append("ldflags = {}".format(" ".join(ldflags))) # rules ninja.append("") ninja.append("rule compile") ninja.append(" command = $cxx /showIncludes $cflags -c $in /Fo$out") ninja.append(" deps = msvc") ninja.append("") ninja.append("rule link") ninja.append(" command = $cxx $in /link $ldflags /out:$out") ninja.append("") # build targets obj_name = "main.obj" ninja.append( "build {}: compile {}".format(obj_name, str(source_path.resolve()).replace(":", "$:")) ) # Use appropriate extension based on platform ninja.append(f"build {libname}: link {obj_name}") ninja.append("") # default target ninja.append(f"default {libname}") ninja.append("") with open(build_dir / "build.ninja", "w") as f: # noqa: PTH123 f.write("\n".join(ninja)) def get_torch_include_paths(build_with_cuda: bool) -> Sequence[str]: """Get the include paths for building with torch.""" if torch.__version__ >= torch.torch_version.TorchVersion("2.6.0"): return torch.utils.cpp_extension.include_paths( device_type="cuda" if build_with_cuda else "cpu" ) else: return torch.utils.cpp_extension.include_paths(cuda=build_with_cuda) # ty: ignore[unknown-argument] def main() -> None: # noqa: PLR0912, PLR0915 """Build the torch c dlpack extension.""" # we need to set the following env to avoid tvm_ffi to build the torch c-dlpack addon during importing os.environ["TVM_FFI_DISABLE_TORCH_C_DLPACK"] = "1" from tvm_ffi.utils.lockfile import FileLock # noqa: PLC0415 parser = argparse.ArgumentParser( description="Build the torch c dlpack extension. After building, a shared library will be placed in the output directory.", ) parser.add_argument( "--build-dir", type=str, required=False, help="Directory to store the built extension library. If not provided, a temporary directory will be used.", ) parser.add_argument( "--output-dir", type=str, required=False, default=str(Path(os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")).expanduser()), help="Directory to store the built extension library. If not specified, the default cache directory of tvm-ffi will be used.", ) parser.add_argument( "--build-with-cuda", action="store_true", help="Build with CUDA support.", ) parser.add_argument( "--build-with-rocm", action="store_true", help="Build with ROCm support.", ) parser.add_argument( "--libname", type=str, default="auto", help="The name of the generated library. It can be a name 'auto' to auto-generate a name following 'libtorch_c_dlpack_addon_torch{version.major}{version.minor}-cpu/cuda.{extension}'.", ) args = parser.parse_args() if args.build_with_cuda and args.build_with_rocm: raise ValueError("Cannot enable both CUDA and ROCm at the same time.") # resolve build directory if args.build_dir is None: build_dir = Path(tempfile.mkdtemp(prefix="tvm-ffi-torch-c-dlpack-")) else: build_dir = Path(args.build_dir) build_dir = build_dir.resolve() if not build_dir.exists(): build_dir.mkdir(parents=True, exist_ok=True) # resolve library name if args.libname == "auto": major, minor = torch.__version__.split(".")[:2] if args.build_with_cuda: device = "cuda" elif args.build_with_rocm: device = "rocm" else: device = "cpu" suffix = ".dll" if IS_WINDOWS else ".so" libname = f"libtorch_c_dlpack_addon_torch{major}{minor}-{device}{suffix}" else: libname = args.libname tmp_libname = libname + ".tmp" # create output directory is not exists output_dir = Path(args.output_dir).expanduser() if not output_dir.exists(): output_dir.mkdir(parents=True, exist_ok=True) with FileLock(str(output_dir / (libname + ".lock"))): if (output_dir / libname).exists(): # already built return # write the source source_path = build_dir / "addon.cc" with open(source_path, "w") as f: # noqa: PTH123 f.write(cpp_source) # resolve configs include_paths = [] ldflags = [] cflags = [] include_paths.append(sysconfig.get_paths()["include"]) if args.build_with_cuda: cflags.append("-DBUILD_WITH_CUDA") elif args.build_with_rocm: cflags.extend(torch.utils.cpp_extension.COMMON_HIP_FLAGS) cflags.append("-DBUILD_WITH_ROCM") include_paths.extend(get_torch_include_paths(args.build_with_cuda or args.build_with_rocm)) # use CXX11 ABI if torch.compiled_with_cxx11_abi(): cflags.append("-D_GLIBCXX_USE_CXX11_ABI=1") else: cflags.append("-D_GLIBCXX_USE_CXX11_ABI=0") for lib_dir in torch.utils.cpp_extension.library_paths(): if IS_WINDOWS: ldflags.append(f"/LIBPATH:{lib_dir}") else: ldflags.extend(["-L", str(lib_dir)]) # Add all required PyTorch libraries if IS_WINDOWS: # On Windows, use .lib format for linking ldflags.extend(["c10.lib", "torch.lib", "torch_cpu.lib", "torch_python.lib"]) if args.build_with_cuda: ldflags.extend(["torch_cuda.lib", "c10_cuda.lib"]) else: # On Unix/macOS, use -l format for linking ldflags.extend(["-lc10", "-ltorch", "-ltorch_cpu", "-ltorch_python"]) if args.build_with_cuda: ldflags.extend(["-ltorch_cuda", "-lc10_cuda"]) # Add Python library linking if IS_WINDOWS: python_lib = f"python{sys.version_info.major}.lib" python_libdir_list = [ sysconfig.get_config_var("LIBDIR"), str(Path(sys.base_exec_prefix) / "libs"), ] if ( sysconfig.get_path("include") is not None and (Path(sysconfig.get_path("include")).parent / "libs").exists() ): python_libdir_list.append( str((Path(sysconfig.get_path("include")).parent / "libs").resolve()) ) for python_libdir in python_libdir_list: if python_libdir and (Path(python_libdir) / python_lib).exists(): ldflags.append(f"/LIBPATH:{python_libdir.replace(':', '$:')}") break if IS_DARWIN: python_libdir = sysconfig.get_config_var("LIBDIR") if python_libdir: py_version = f"python{sysconfig.get_python_version()}" ldflags.append(f"-L{python_libdir}") ldflags.append(f"-l{py_version}") env_ldflags = parse_env_flags("TVM_FFI_JIT_EXTRA_LDFLAGS") if env_ldflags: ldflags.extend(env_ldflags) env_cflags = parse_env_flags("TVM_FFI_JIT_EXTRA_CFLAGS") if env_cflags: cflags.extend(env_cflags) # build the shared library if IS_WINDOWS: # Use ninja on Windows _generate_ninja_build_windows( build_dir=build_dir, libname=tmp_libname, source_path=source_path, extra_cflags=cflags, extra_ldflags=ldflags, extra_include_paths=include_paths, ) from tvm_ffi.cpp.extension import build_ninja # noqa: PLC0415 build_ninja(build_dir=str(build_dir)) else: # Use direct command on non-Windows _run_build_on_linux_like( build_dir=build_dir, libname=tmp_libname, source_path=source_path, extra_cflags=cflags, extra_ldflags=ldflags, extra_include_paths=include_paths, ) # rename the tmp file to final libname shutil.move(str(build_dir / tmp_libname), str(output_dir / libname)) if __name__ == "__main__": main() tvm-ffi-0.1.12/python/tvm_ffi/utils/embed_cubin.py000066400000000000000000000345411521067262500221010ustar00rootroot00000000000000#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utility script to embed CUBIN data into existing object files. This script takes an object file containing C++ code that uses TVM_FFI_EMBED_CUBIN and embeds the CUBIN binary data directly into it, creating a new combined object file. Usage: python -m tvm_ffi.utils.embed_cubin \ --output-obj new_obj.o \ --input-obj old_obj.o \ --cubin my-cubin.cubin \ --name my_module The output object file will contain both: - The original C++ code from input-obj - The embedded CUBIN data with symbols: __tvm_ffi__cubin_ __tvm_ffi__cubin__end These symbols are accessed in C++ code using: TVM_FFI_EMBED_CUBIN(my_module); """ from __future__ import annotations import argparse import subprocess import sys import tempfile import traceback from pathlib import Path def embed_cubin( # noqa: PLR0912, PLR0915 cubin_path: Path, input_obj_path: Path, output_obj_path: Path, name: str, verbose: bool = False, ) -> None: """Embed a CUBIN file into an existing object file. This function takes a CUBIN binary file and merges it into an existing object file that contains C++ code using TVM_FFI_EMBED_CUBIN. The result is a new object file containing both the original code and the embedded CUBIN data. The process involves: 1. Creating an intermediate CUBIN object file using `ld -r -b binary` 2. Adding `.note.GNU-stack` section to CUBIN object for security 3. Using `objcopy` to rename symbols to the format expected by TVM-FFI: - _binary__start → __tvm_ffi__cubin_ - _binary__end → __tvm_ffi__cubin__end 4. Using `ld -r` (relocatable link) to merge the CUBIN object with the input object 5. Localizing the symbols to prevent conflicts when multiple object files use the same name Parameters ---------- cubin_path : Path Path to the input CUBIN file. input_obj_path : Path Path to the existing object file containing C++ code with TVM_FFI_EMBED_CUBIN. output_obj_path : Path Path to the output object file (will contain both code and CUBIN data). name : str Name used in the TVM_FFI_EMBED_CUBIN macro (e.g., "my_module"). This name must match the name used in the C++ code. verbose : bool, optional If True, print detailed command output. Default is False. Raises ------ FileNotFoundError If the input CUBIN file or input object file does not exist. RuntimeError If `ld` or `objcopy` commands fail. Examples -------- In Python, ```python embed_cubin( Path("kernel.cubin"), Path("old.o"), Path("new.o"), "my_kernels", ) ``` Then in C++ code, ```C++ TVM_FFI_EMBED_CUBIN(my_kernels); auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "kernel_name"); ``` (in the source that was compiled to old.o). """ if not cubin_path.exists(): raise FileNotFoundError(f"CUBIN file not found: {cubin_path}") if not input_obj_path.exists(): raise FileNotFoundError(f"Input object file not found: {input_obj_path}") if verbose: print(f"Input CUBIN: {cubin_path} ({cubin_path.stat().st_size} bytes)") print(f"Input object: {input_obj_path} ({input_obj_path.stat().st_size} bytes)") print(f"Output object: {output_obj_path}") print(f"Symbol name: {name}") # Create a temporary directory for intermediate files with tempfile.TemporaryDirectory() as tmpdir: tmp_path = Path(tmpdir) # Copy CUBIN to temp directory with a predictable name for symbol generation # The filename affects the auto-generated symbol names from ld temp_cubin = tmp_path / f"embedded_{name}.cubin" temp_cubin.write_bytes(cubin_path.read_bytes()) # Step 1: Create intermediate CUBIN object file with ld -r -b binary cubin_obj = tmp_path / f"cubin_{name}.o" ld_binary_cmd = ["ld", "-r", "-b", "binary", "-o", str(cubin_obj), str(temp_cubin.name)] if verbose: print("\n[Step 1/4] Creating CUBIN object file with ld") print(f"Command: {' '.join(ld_binary_cmd)}") print(f"Working directory: {tmp_path}") result = subprocess.run(ld_binary_cmd, cwd=str(tmp_path), capture_output=True, check=False) if result.returncode != 0: stderr = result.stderr.decode("utf-8") stdout = result.stdout.decode("utf-8") raise RuntimeError( f"ld (binary) failed with status {result.returncode}\n" f"Command: {' '.join(ld_binary_cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) if verbose and result.stdout: print(f"stdout: {result.stdout.decode('utf-8')}") if verbose and result.stderr: print(f"stderr: {result.stderr.decode('utf-8')}") # Step 1.5: Add .note.GNU-stack section to CUBIN object for security # This marks the stack as non-executable and prevents linker warnings # We do this before renaming so the section is preserved in the final output objcopy_stack_cmd = [ "objcopy", "--add-section", ".note.GNU-stack=/dev/null", "--set-section-flags", ".note.GNU-stack=noload,readonly", str(cubin_obj), ] if verbose: print("\n[Step 1.5/4] Adding .note.GNU-stack section to CUBIN object") print(f"Command: {' '.join(objcopy_stack_cmd)}") result = subprocess.run(objcopy_stack_cmd, capture_output=True, check=False) if result.returncode != 0: stderr = result.stderr.decode("utf-8") stdout = result.stdout.decode("utf-8") raise RuntimeError( f"objcopy (add stack section) failed with status {result.returncode}\n" f"Command: {' '.join(objcopy_stack_cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) if verbose and result.stdout: print(f"stdout: {result.stdout.decode('utf-8')}") if verbose and result.stderr: print(f"stderr: {result.stderr.decode('utf-8')}") # Step 2: Rename symbols with objcopy # The ld command creates symbols like: # _binary_embedded__cubin_start # _binary_embedded__cubin_end # We rename them to: # __tvm_ffi__cubin_ # __tvm_ffi__cubin__end # Note: We don't localize yet - that happens after merging cubin_obj_renamed = tmp_path / f"cubin_{name}_renamed.o" objcopy_cmd = [ "objcopy", "--rename-section", ".data=.rodata,alloc,load,readonly,data,contents", "--redefine-sym", f"_binary_embedded_{name}_cubin_start=__tvm_ffi__cubin_{name}", "--redefine-sym", f"_binary_embedded_{name}_cubin_end=__tvm_ffi__cubin_{name}_end", str(cubin_obj), str(cubin_obj_renamed), ] if verbose: print("\n[Step 2/4] Renaming CUBIN symbols with objcopy") print(f"Command: {' '.join(objcopy_cmd)}") result = subprocess.run(objcopy_cmd, capture_output=True, check=False) if result.returncode != 0: stderr = result.stderr.decode("utf-8") stdout = result.stdout.decode("utf-8") raise RuntimeError( f"objcopy failed with status {result.returncode}\n" f"Command: {' '.join(objcopy_cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) if verbose and result.stdout: print(f"stdout: {result.stdout.decode('utf-8')}") if verbose and result.stderr: print(f"stderr: {result.stderr.decode('utf-8')}") # Step 3: Merge input object file with CUBIN object using ld -r (relocatable link) ld_merge_cmd = [ "ld", "-r", "-o", str(output_obj_path), str(input_obj_path), str(cubin_obj_renamed), ] if verbose: print("\n[Step 3/4] Merging objects with ld (relocatable link)") print(f"Command: {' '.join(ld_merge_cmd)}") result = subprocess.run(ld_merge_cmd, capture_output=True, check=False) if result.returncode != 0: stderr = result.stderr.decode("utf-8") stdout = result.stdout.decode("utf-8") raise RuntimeError( f"ld (merge) failed with status {result.returncode}\n" f"Command: {' '.join(ld_merge_cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) if verbose and result.stdout: print(f"stdout: {result.stdout.decode('utf-8')}") if verbose and result.stderr: print(f"stderr: {result.stderr.decode('utf-8')}") # Step 4: Localize CUBIN symbols to prevent conflicts across object files # We do this after merging so the C++ code can reference the symbols during the link objcopy_localize_cmd = [ "objcopy", "--localize-symbol", f"__tvm_ffi__cubin_{name}", "--localize-symbol", f"__tvm_ffi__cubin_{name}_end", str(output_obj_path), ] if verbose: print("\n[Step 4/4] Localizing CUBIN symbols") print(f"Command: {' '.join(objcopy_localize_cmd)}") result = subprocess.run(objcopy_localize_cmd, capture_output=True, check=False) if result.returncode != 0: stderr = result.stderr.decode("utf-8") stdout = result.stdout.decode("utf-8") raise RuntimeError( f"objcopy (localize symbols) failed with status {result.returncode}\n" f"Command: {' '.join(objcopy_localize_cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) if verbose and result.stdout: print(f"stdout: {result.stdout.decode('utf-8')}") if verbose and result.stderr: print(f"stderr: {result.stderr.decode('utf-8')}") if verbose: print(f"\n✓ Successfully created combined object file: {output_obj_path}") print(" Contains:") print(f" - Original code from {input_obj_path.name}") print(" - Embedded CUBIN data with local symbols:") print(f" __tvm_ffi__cubin_{name} (local)") print(f" __tvm_ffi__cubin_{name}_end (local)") print(" - .note.GNU-stack section (non-executable stack)") def main() -> int: """Main entry point for the embed_cubin utility.""" # noqa: D401 parser = argparse.ArgumentParser( prog="python -m tvm_ffi.utils.embed_cubin", description="Embed CUBIN data into existing object files that use TVM_FFI_EMBED_CUBIN macro", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Basic usage python -m tvm_ffi.utils.embed_cubin \\ --output-obj new.o \\ --input-obj old.o \\ --cubin kernel.cubin \\ --name my_kernels # With verbose output python -m tvm_ffi.utils.embed_cubin \\ --output-obj new.o \\ --input-obj old.o \\ --cubin kernel.cubin \\ --name my_kernels \\ --verbose Workflow: 1. Compile C++ code that uses TVM_FFI_EMBED_CUBIN to create old.o 2. Compile CUDA kernel to CUBIN (e.g., using nvcc or NVRTC) 3. Use this tool to merge them into new.o 4. Link new.o into your final shared library Usage in C++ code (source compiled to old.o): TVM_FFI_EMBED_CUBIN(my_kernels); auto kernel = TVM_FFI_EMBED_CUBIN_GET_KERNEL(my_kernels, "kernel_name"); Requirements: - GNU binutils (ld and objcopy) must be available in PATH - Linux/Unix platform (Windows uses different embedding mechanisms) """, ) parser.add_argument( "--output-obj", type=str, required=True, metavar="PATH", help="Output object file path (e.g., new.o)", ) parser.add_argument( "--input-obj", type=str, required=True, metavar="PATH", help="Input object file path containing TVM_FFI_EMBED_CUBIN usage (e.g., old.o)", ) parser.add_argument( "--cubin", type=str, required=True, metavar="PATH", help="Input CUBIN file path (e.g., kernel.cubin)", ) parser.add_argument( "--name", type=str, required=True, metavar="NAME", help="Name used in TVM_FFI_EMBED_CUBIN macro (e.g., my_kernels)", ) parser.add_argument( "-v", "--verbose", action="store_true", help="Print detailed command output", ) args = parser.parse_args() try: cubin_path = Path(args.cubin).resolve() input_obj_path = Path(args.input_obj).resolve() output_obj_path = Path(args.output_obj).resolve() # Create output directory if it doesn't exist output_obj_path.parent.mkdir(parents=True, exist_ok=True) embed_cubin(cubin_path, input_obj_path, output_obj_path, args.name, args.verbose) if not args.verbose: print(f"✓ Created {output_obj_path}") return 0 except FileNotFoundError as e: print(f"Error: {e}", file=sys.stderr) return 1 except RuntimeError as e: print(f"Error: {e}", file=sys.stderr) return 2 except Exception as e: print(f"Unexpected error: {e}", file=sys.stderr) traceback.print_exc() return 3 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/python/tvm_ffi/utils/kwargs_wrapper.py000066400000000000000000000366371521067262500227130ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utilities for creating high-performance keyword argument wrapper functions. This module provides tools for wrapping positional-only callables with keyword argument support using code generation techniques. """ from __future__ import annotations import functools import inspect import keyword from typing import Any, Callable, Iterable from tvm_ffi.utils.unpack_dataclass import unpack_dataclass_to_tuple # Sentinel object for missing arguments MISSING = object() # Internal variable names used in generated code to avoid user argument conflicts _INTERNAL_TARGET_FUNC = "__i_target_func" _INTERNAL_MISSING = "__i_MISSING" _INTERNAL_DEFAULTS_DICT = "__i_arg_defaults" _INTERNAL_UNPACK = "__i_unpack" _INTERNAL_NAMES = { _INTERNAL_TARGET_FUNC, _INTERNAL_MISSING, _INTERNAL_DEFAULTS_DICT, _INTERNAL_UNPACK, } def _validate_argument_names(names: list[str], arg_type: str) -> None: """Validate that argument names are valid Python identifiers and unique. Parameters ---------- names List of argument names to validate. arg_type Description of the argument type (e.g., "Argument", "Keyword-only argument"). """ # Check for duplicate names if len(set(names)) != len(names): raise ValueError(f"Duplicate {arg_type.lower()} names found in: {names}") # Validate each name is a valid identifier for name in names: if not isinstance(name, str): raise TypeError( f"{arg_type} name must be a string, got {type(name).__name__}: {name!r}" ) if keyword.iskeyword(name): raise ValueError( f"Invalid {arg_type.lower()} name: {name!r} is a Python keyword and cannot be used as a parameter name" ) if not name.isidentifier(): raise ValueError( f"Invalid {arg_type.lower()} name: {name!r} is not a valid Python identifier" ) def _validate_wrapper_args( arg_names: list[str], arg_defaults: tuple, kwonly_names: list[str], kwonly_defaults: dict[str, Any], reserved_names: set[str], ) -> None: """Validate all input arguments for make_kwargs_wrapper. Parameters ---------- arg_names List of positional argument names. arg_defaults Tuple of default values for positional arguments. kwonly_names List of keyword-only argument names. kwonly_defaults Dictionary of default values for keyword-only arguments. reserved_names Set of reserved internal names that cannot be used as argument names. """ # Validate arg_names are valid Python identifiers and unique _validate_argument_names(arg_names, "Argument") # Validate arg_defaults is a tuple if not isinstance(arg_defaults, tuple): raise TypeError(f"arg_defaults must be a tuple, got {type(arg_defaults).__name__}") # Validate arg_defaults length doesn't exceed arg_names length if len(arg_defaults) > len(arg_names): raise ValueError( f"arg_defaults has {len(arg_defaults)} values but only " f"{len(arg_names)} positional arguments" ) # Validate kwonly_names are valid identifiers and unique _validate_argument_names(kwonly_names, "Keyword-only argument") # Validate kwonly_defaults keys are in kwonly_names kwonly_names_set = set(kwonly_names) for key in kwonly_defaults: if key not in kwonly_names_set: raise ValueError( f"Default provided for '{key}' which is not in kwonly_names: {kwonly_names}" ) # Validate no overlap between positional and keyword-only arguments arg_names_set = set(arg_names) overlap = arg_names_set & kwonly_names_set if overlap: raise ValueError(f"Arguments cannot be both positional and keyword-only: {overlap}") # Validate no conflict between user argument names and internal names all_user_names = arg_names_set | kwonly_names_set conflicts = all_user_names & reserved_names if conflicts: raise ValueError( f"Argument names conflict with internal names: {conflicts}. " f'Please avoid using names starting with "__i_"' ) def _build_wrapper_code( arg_names: list[str], arg_defaults: tuple, kwonly_names: list[str], kwonly_defaults: dict[str, Any], dc_to_tuple_set: set[str], ) -> tuple[str, dict[str, Any]]: """Build the generated wrapper code string and runtime defaults dict. Returns ------- A tuple of (code_str, runtime_defaults) where code_str is the generated wrapper function code and runtime_defaults maps arg names to their default values. """ # Build positional defaults dictionary (right-aligned) arg_defaults_dict = ( dict(zip(arg_names[-len(arg_defaults) :], arg_defaults)) if arg_defaults else {} ) arg_parts: list[str] = [] call_parts: list[str] = [] runtime_defaults: dict[str, Any] = {} def _transform_expr(name: str, expr: str) -> str: if name in dc_to_tuple_set: return f"{_INTERNAL_UNPACK}({expr})" return expr def _add_param_with_default(name: str, default_value: Any) -> None: # Directly embed None and bool defaults; use MISSING sentinel for others. if default_value is None: arg_parts.append(f"{name}=None") call_parts.append(_transform_expr(name, name)) elif type(default_value) is bool: default_value_str = "True" if default_value else "False" arg_parts.append(f"{name}={default_value_str}") call_parts.append(_transform_expr(name, name)) else: arg_parts.append(f"{name}={_INTERNAL_MISSING}") runtime_defaults[name] = default_value base_expr = ( f'{_INTERNAL_DEFAULTS_DICT}["{name}"] if {name} is {_INTERNAL_MISSING} else {name}' ) call_parts.append(_transform_expr(name, base_expr)) for name in arg_names: if name in arg_defaults_dict: _add_param_with_default(name, arg_defaults_dict[name]) else: arg_parts.append(name) call_parts.append(_transform_expr(name, name)) if kwonly_names: arg_parts.append("*") for name in kwonly_names: if name in kwonly_defaults: _add_param_with_default(name, kwonly_defaults[name]) else: arg_parts.append(name) call_parts.append(_transform_expr(name, name)) arg_list = ", ".join(arg_parts) call_list = ", ".join(call_parts) code_str = f""" def wrapper({arg_list}): return {_INTERNAL_TARGET_FUNC}({call_list}) """ return code_str, runtime_defaults def make_kwargs_wrapper( target_func: Callable, arg_names: list[str], arg_defaults: tuple = (), kwonly_names: list[str] | None = None, kwonly_defaults: dict[str, Any] | None = None, prototype: Callable | None = None, map_dataclass_to_tuple: list[str] | None = None, ) -> Callable: """Create a wrapper with kwargs support for a function that only accepts positional arguments. This function dynamically generates a wrapper using code generation to minimize overhead. Parameters ---------- target_func The underlying function to be called by the wrapper. This function must only accept positional arguments. arg_names A list of ALL positional argument names in order. These define the positional parameters that the wrapper will accept. Must not overlap with kwonly_names. arg_defaults A tuple of default values for positional arguments, right-aligned to arg_names (matching Python's __defaults__ behavior). The length of this tuple determines how many trailing arguments have defaults. Example: (10, 20) with arg_names=['a', 'b', 'c', 'd'] means c=10, d=20. Empty tuple () means no defaults. kwonly_names A list of keyword-only argument names. These arguments can only be passed by name, not positionally, and appear after a '*' separator in the signature. Can include both required and optional keyword-only arguments. Must not overlap with arg_names. Example: ['debug', 'timeout'] creates wrapper(..., *, debug, timeout). kwonly_defaults Optional dictionary of default values for keyword-only arguments (matching Python's __kwdefaults__ behavior). Keys must be a subset of kwonly_names. Keyword-only arguments not in this dict are required. Example: {'timeout': 30} with kwonly_names=['debug', 'timeout'] means 'debug' is required and 'timeout' is optional. prototype Optional prototype function to copy metadata (__name__, __doc__, __module__, __qualname__, __annotations__) from. If None, no metadata is copied. map_dataclass_to_tuple Optional list of argument names whose values should be converted from dataclass instances to tuples via ``unpack_dataclass_to_tuple``. Dataclass fields are unpacked to tuple and leaf fields are shallow copied. Nested dataclass fields are recursed automatically based on type annotations. Lists/tuples are recursed element-wise. Dicts are recursed on values. Non-dataclass leaves are passed through unchanged. Returns ------- A dynamically generated wrapper function with the specified signature Notes ----- The generated wrapper will directly embed default values for None and bool types and use a MISSING sentinel object to distinguish between explicitly passed arguments and those that should use default values for other types to ensure the generated code does not contain unexpected str repr. """ # Normalize inputs if kwonly_names is None: kwonly_names = [] if kwonly_defaults is None: kwonly_defaults = {} dc_to_tuple_set = set(map_dataclass_to_tuple) if map_dataclass_to_tuple else set() # Validate all input arguments _validate_wrapper_args(arg_names, arg_defaults, kwonly_names, kwonly_defaults, _INTERNAL_NAMES) # Build the generated wrapper code code_str, runtime_defaults = _build_wrapper_code( arg_names, arg_defaults, kwonly_names, kwonly_defaults, dc_to_tuple_set ) # Execute the generated code # Note: this is a limited use of exec that is safe. # We ensure generated code does not contain any untrusted input. # The argument names are validated and the default values are not part of generated code. # Instead default values are set to MISSING sentinel object and explicitly passed as exec_globals. # This is a practice adopted by `dataclasses` and `pydantic` exec_globals: dict[str, Any] = { _INTERNAL_TARGET_FUNC: target_func, _INTERNAL_MISSING: MISSING, _INTERNAL_DEFAULTS_DICT: runtime_defaults, } if dc_to_tuple_set: exec_globals[_INTERNAL_UNPACK] = unpack_dataclass_to_tuple namespace: dict[str, Any] = {} exec(code_str, exec_globals, namespace) new_func = namespace["wrapper"] # Copy metadata from prototype if provided if prototype is not None: functools.update_wrapper(new_func, prototype, updated=()) return new_func def make_kwargs_wrapper_from_signature( target_func: Callable, signature: inspect.Signature, prototype: Callable | None = None, exclude_arg_names: Iterable[str] | None = None, map_dataclass_to_tuple: list[str] | None = None, ) -> Callable: """Create a wrapper with kwargs support for a function that only accepts positional arguments. This is a convenience function that extracts parameter information from a signature object and calls make_kwargs_wrapper with the appropriate arguments. Supports both required and optional keyword-only arguments. Parameters ---------- target_func The underlying function to be called by the wrapper. signature An inspect.Signature object describing the desired wrapper signature. prototype Optional prototype function to copy metadata (__name__, __doc__, __module__, __qualname__, __annotations__) from. If None, no metadata is copied. exclude_arg_names Optional iterable of argument names to ignore when extracting parameters from the signature. These arguments will not be included in the generated wrapper. If a name in this iterable does not exist in the signature, it is silently ignored. map_dataclass_to_tuple Optional list of argument names to unpack via ``unpack_dataclass_to_tuple``. See ``make_kwargs_wrapper`` for details. Returns ------- A dynamically generated wrapper function with the specified signature. Raises ------ ValueError If the signature contains *args or **kwargs. """ # Normalize exclude_arg_names to a set for efficient lookup skip_set = set(exclude_arg_names) if exclude_arg_names is not None else set() # Extract positional and keyword-only parameters arg_names = [] arg_defaults_list = [] kwonly_names = [] kwonly_defaults = {} # Track when we start seeing defaults for positional args has_seen_positional_default = False for param_name, param in signature.parameters.items(): # Skip arguments that are in the skip list if param_name in skip_set: continue if param.kind == inspect.Parameter.VAR_POSITIONAL: raise ValueError("*args not supported in wrapper generation") elif param.kind == inspect.Parameter.VAR_KEYWORD: raise ValueError("**kwargs not supported in wrapper generation") elif param.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ): arg_names.append(param_name) if param.default is not inspect.Parameter.empty: has_seen_positional_default = True arg_defaults_list.append(param.default) elif has_seen_positional_default: # Required arg after optional arg (invalid in Python) raise ValueError( f"Required positional parameter '{param_name}' cannot follow " f"parameters with defaults" ) elif param.kind == inspect.Parameter.KEYWORD_ONLY: kwonly_names.append(param_name) if param.default is not inspect.Parameter.empty: kwonly_defaults[param_name] = param.default # Convert defaults list to tuple (right-aligned to arg_names) arg_defaults = tuple(arg_defaults_list) return make_kwargs_wrapper( target_func, arg_names, arg_defaults, kwonly_names, kwonly_defaults, prototype, map_dataclass_to_tuple, ) tvm-ffi-0.1.12/python/tvm_ffi/utils/lockfile.py000066400000000000000000000127461521067262500214400ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Simple cross-platform advisory file lock utilities.""" from __future__ import annotations import os import sys import time from typing import Any, Literal # Platform-specific imports for file locking if sys.platform == "win32": import msvcrt else: import fcntl class FileLock: """Provide a cross-platform file locking mechanism using Python's stdlib. This class implements an advisory lock, which must be respected by all cooperating processes. Please note that this lock does not prevent the same process from acquiring the lock multiple times; it is the caller's responsibility to manage this. Examples -------- .. code-block:: python from tvm_ffi.utils import FileLock with FileLock("/tmp/my.lock"): # Critical section guarded by the lock. # Other processes attempting to acquire the same lock will block # (or fail, if using ``acquire()``) until this context exits. do_work() """ def __init__(self, lock_file_path: str) -> None: """Initialize a file lock using the given lock file path.""" self.lock_file_path = lock_file_path self._file_descriptor: int | None = None def __enter__(self) -> FileLock: """Acquire the lock upon entering the context. This method blocks until the lock is acquired. """ self.blocking_acquire() return self def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Literal[False]: """Context manager protocol: release the lock upon exiting the 'with' block.""" self.release() return False # Propagate exceptions, if any def acquire(self) -> bool: """Acquire an exclusive, non-blocking lock on the file. Returns ------- ret: bool True if the lock was acquired, False otherwise. Raises ------ RuntimeError If an unexpected error occurs during lock acquisition. """ try: if self._file_descriptor is not None: return False # Lock is already held by this instance if sys.platform == "win32": self._file_descriptor = os.open( self.lock_file_path, os.O_RDWR | os.O_CREAT | os.O_BINARY ) msvcrt.locking(self._file_descriptor, msvcrt.LK_NBLCK, 1) else: # Unix-like systems self._file_descriptor = os.open(self.lock_file_path, os.O_WRONLY | os.O_CREAT) fcntl.flock(self._file_descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB) return True except (OSError, BlockingIOError): if self._file_descriptor is not None: os.close(self._file_descriptor) self._file_descriptor = None return False except Exception as e: if self._file_descriptor is not None: os.close(self._file_descriptor) self._file_descriptor = None raise RuntimeError(f"An unexpected error occurred: {e}") def blocking_acquire(self, timeout: float | None = None, poll_interval: float = 0.1) -> bool: """Wait until an exclusive lock can be acquired, with an optional timeout. Parameters ---------- timeout: float, optional The maximum time to wait for the lock in seconds. A value of None means wait indefinitely. poll_interval: float The time to wait between lock attempts in seconds. Returns ------- ret: bool True if the lock was acquired. Raises ------ TimeoutError If the lock is not acquired within the timeout period. RuntimeError If the lock is already held by this instance. """ if self._file_descriptor is not None: raise RuntimeError("Lock is already held by this instance.") start_time = time.time() while True: if self.acquire(): return True # Check for timeout if timeout is not None and (time.time() - start_time) > timeout: raise TimeoutError( f"Failed to acquire lock on '{self.lock_file_path}' after {timeout} seconds." ) time.sleep(poll_interval) def release(self) -> None: """Releases the lock and closes the file descriptor.""" if self._file_descriptor is not None: if sys.platform == "win32": msvcrt.locking(self._file_descriptor, msvcrt.LK_UNLCK, 1) else: fcntl.flock(self._file_descriptor, fcntl.LOCK_UN) os.close(self._file_descriptor) self._file_descriptor = None tvm-ffi-0.1.12/python/tvm_ffi/utils/unpack_dataclass.py000066400000000000000000000252211521067262500231400ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Fast dataclass-to-tuple conversion via JIT-compiled unpacking. This module provides ``unpack_dataclass_to_tuple``, a function that recursively converts dataclass instances to tuples. It JIT-compiles a per-class unpacker on first call and caches it per-thread for ~5-11x speedup over ``dataclasses.astuple`` with no deep-copy of leaf values. """ from __future__ import annotations import dataclasses import keyword import sys import threading import typing from typing import Any # Support both typing.Union and types.UnionType (PEP 604, Python 3.10+) if sys.version_info >= (3, 10): import types _LEAF_CONTAINER_ORIGINS = (tuple, dict, set, frozenset, typing.Union, types.UnionType) else: _LEAF_CONTAINER_ORIGINS = (tuple, dict, set, frozenset, typing.Union) # Type alias for dataclass-to-tuple schema (internal). # Schema values: # None -> leaf, direct attribute access (zero cost) # "unpack" -> dynamic dispatch via unpack_dataclass_to_tuple at runtime # dict -> nested struct, recurse inline # Example: {"x": None, "y": None} -> (__x.x, __x.y,) # Example: {"cfg": {"x": None, "y": None}, "data": "unpack"} # -> ((__x.cfg.x, __x.cfg.y,), __dispatch(__x.data),) DataclassToTupleSchema = dict # dict[str, None | str | DataclassToTupleSchema] # Sentinel value in schema: field should be dynamically dispatched UNPACK = "unpack" # Thread-local cache for JIT-compiled per-class unpack functions _tls = threading.local() # Types known to be safe leaves (never contain dataclass instances) _KNOWN_LEAF_TYPES: set[type] = {int, float, str, bool, bytes, complex, type(None)} def _is_known_leaf_type(tp: Any) -> bool: """Check if a type is definitely a leaf (no dataclass content or conversion needed). Note: list is NOT a leaf because it must be converted to a tuple per the unpack contract (matching dataclasses.astuple behavior). """ if isinstance(tp, type): return tp in _KNOWN_LEAF_TYPES if tp is Ellipsis: return True origin = typing.get_origin(tp) if origin is not None: # list is NOT a leaf — must be converted to tuple # tuple/dict/set/frozenset/Union are leaves if all args are leaves if origin in _LEAF_CONTAINER_ORIGINS: args = typing.get_args(tp) return bool(args) and all(_is_known_leaf_type(a) for a in args) return False def _classify_field_type( field_type: Any, memo: set[type] | None = None ) -> None | str | DataclassToTupleSchema: """Classify a resolved field type into a schema entry. Conservative: only returns None (leaf) when we are certain the type cannot contain a dataclass. Otherwise returns UNPACK (dynamic dispatch). """ if isinstance(field_type, str) or field_type is Any or field_type is object: return UNPACK if dataclasses.is_dataclass(field_type) and isinstance(field_type, type): # Guard against infinite recursion for self-referential dataclasses if memo is not None and field_type in memo: return UNPACK return _extract_dataclass_to_tuple_schema(field_type, memo=memo) # Known primitive types -> leaf if isinstance(field_type, type) and field_type in _KNOWN_LEAF_TYPES: return None # Generic containers: check element types # list always needs UNPACK (must be converted to tuple) # tuple/dict/set/frozenset/Union are leaves if all args are known leaves # Generic containers: list always UNPACK (must convert to tuple). # tuple/dict/set/frozenset/Union are leaves only if all args are known leaves. # Everything else (unknown type) -> UNPACK (conservative). origin = typing.get_origin(field_type) if origin in _LEAF_CONTAINER_ORIGINS: args = typing.get_args(field_type) if args and all(_is_known_leaf_type(a) for a in args): return None return UNPACK def _compile_dataclass_to_tuple_schema(prefix: str, schema: DataclassToTupleSchema) -> str: """Compile a DataclassToTupleSchema into an inline tuple expression. Parameters ---------- prefix The variable expression to unpack (e.g. "__x" or "__x.field"). schema The schema dict mapping field names to: - None: leaf, direct attribute access - "unpack": dynamic dispatch via __dispatch() at runtime - nested dict: recurse inline Returns ------- A string expression that evaluates to a tuple of the unpacked fields. """ parts: list[str] = [] for field_name, sub_schema in schema.items(): field_expr = f"{prefix}.{field_name}" if sub_schema is None: parts.append(field_expr) elif sub_schema == UNPACK: parts.append(f"__dispatch({field_expr})") else: parts.append(_compile_dataclass_to_tuple_schema(field_expr, sub_schema)) return "(" + ", ".join(parts) + (",)" if parts else ")") def _validate_dataclass_to_tuple_schema(schema: DataclassToTupleSchema) -> None: """Validate that a DataclassToTupleSchema contains only safe identifiers. This is critical for security since field names are embedded directly in generated code via exec(). The validation ensures: - Keys are strings (type check) - Keys pass str.isidentifier() — rejects any non-identifier characters - Keys are not Python keywords — rejects control flow injection - Values are only None, "unpack", or recursively-validated dicts Combined with the hardcoded prefix ("__x") and restricted exec_globals, this prevents any code injection through crafted field names. """ if not isinstance(schema, dict): raise TypeError(f"DataclassToTupleSchema must be a dict, got {type(schema).__name__}") for field_name, sub_schema in schema.items(): if not isinstance(field_name, str): raise TypeError(f"Schema field name must be a string, got {type(field_name).__name__}") if not field_name.isidentifier(): raise ValueError(f"Schema field name {field_name!r} is not a valid Python identifier") if keyword.iskeyword(field_name): raise ValueError(f"Schema field name {field_name!r} is a Python keyword") if sub_schema is not None and sub_schema != UNPACK: _validate_dataclass_to_tuple_schema(sub_schema) def _extract_dataclass_to_tuple_schema( cls: type, *, memo: set[type] | None = None ) -> DataclassToTupleSchema: """Extract a DataclassToTupleSchema from a dataclass class using type annotations. Classification per field (conservative: only leaf when certain): - Known dataclass type -> nested schema (recurse inline) - Known primitive type (int, float, str, bool, bytes, complex) -> None (leaf) - Container with only known-leaf args (list[int], dict[str, float]) -> None (leaf) - Container with dataclass/unknown args (list[Config]) -> "unpack" (dynamic dispatch) - Any, object, unresolved string annotation -> "unpack" (dynamic dispatch) - Unknown class -> "unpack" (dynamic dispatch) Uses typing.get_type_hints() to resolve PEP 563 string annotations. Uses memo set to prevent infinite recursion on self-referential dataclasses. """ if not dataclasses.is_dataclass(cls) or not isinstance(cls, type): raise TypeError(f"Expected a dataclass class, got {cls!r}") if memo is None: memo = set() memo.add(cls) try: type_hints = typing.get_type_hints(cls) except (NameError, TypeError, AttributeError): type_hints = {} schema: DataclassToTupleSchema = {} for f in dataclasses.fields(cls): field_type = type_hints.get(f.name, f.type) schema[f.name] = _classify_field_type(field_type, memo=memo) return schema def unpack_dataclass_to_tuple(x: Any) -> Any: """Fast recursively unpack a dataclass value to tuple representation. - Dataclass instances are unpacked to tuples of their field values. - Lists and tuples are recursed element-wise, returning a tuple. - Dicts are recursed on values, returning a new dict. - All other values are returned as-is (leaf passthrough). This function optimizes speed via JIT-compiling the conversion per dataclass class and caching it per-thread. It brings about 5-11x speedup vs ``dataclasses.astuple`` and does not deep-copy leaf values. Parameters ---------- x The value to unpack. Returns ------- The unpacked tuple representation, or ``x`` unchanged if it's a leaf. """ try: cache = _tls.cache except AttributeError: cache = _tls.cache = {} cls = type(x) fn = cache.get(cls) if fn is not None: return fn(x) # Cache miss — classify the type if dataclasses.is_dataclass(cls) and isinstance(cls, type): schema = _extract_dataclass_to_tuple_schema(cls) # Validate that all field names in the schema are safe Python identifiers. # This is critical: field names are embedded directly in the generated code string. # _validate_dataclass_to_tuple_schema ensures no code injection is possible via # crafted field names (isidentifier + iskeyword checks). _validate_dataclass_to_tuple_schema(schema) code_expr = _compile_dataclass_to_tuple_schema("__x", schema) code = f"def __unpack(__x): return {code_expr}" namespace: dict[str, Any] = {} # exec_globals only exposes __dispatch (our own function), no other capabilities. exec(code, {"__dispatch": unpack_dataclass_to_tuple}, namespace) fn = namespace["__unpack"] cache[cls] = fn return fn(x) if isinstance(x, (list, tuple)): return type(x)(unpack_dataclass_to_tuple(e) for e in x) if isinstance(x, dict): return {k: unpack_dataclass_to_tuple(v) for k, v in x.items()} # True leaf — cache identity so next call is just dict.get + return cache[cls] = _LEAF_IDENTITY return x def _LEAF_IDENTITY(x: Any) -> Any: """Identity function cached for known leaf types.""" return x tvm-ffi-0.1.12/rust/000077500000000000000000000000001521067262500141465ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/Cargo.toml000066400000000000000000000015451521067262500161030ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [workspace] members = ["tvm-ffi", "tvm-ffi-sys", "tvm-ffi-macros"] resolver = "2" tvm-ffi-0.1.12/rust/README.md000066400000000000000000000045761521067262500154410ustar00rootroot00000000000000 # Rust Packages (Experimental) Rust support for the `tvm-ffi` ABI. Currently, the rust support is in an experimental stage. This workspace contains three crates: - `tvm-ffi`: Safe, ergonomic Rust bindings over the ABI. - `tvm-ffi-sys`: Low-level exposure of raw C ABIs. - `tvm-ffi-macros`: Procedural macros used by `tvm-ffi` (derive/object helpers and exported function helpers). The overall project focuses on low-level, direct access to the ABI when possible for maximum efficiency while maintaining interoperability. ## Installation The Rust support depends on `libtvm_ffi`. Please install the `tvm-ffi` pip package by running: ```bash pip install -v .. ``` Confirm that `tvm-ffi-config` is available with: ```bash tvm-ffi-config --libdir ``` Then build the workspace with: ```bash cargo build ``` The build will: - Query `tvm-ffi-config --libdir` to add the appropriate link search path. - Link against `tvm_ffi`. - Update the appropriate dynamic loader path environment variable for `cargo run` and `cargo test`. For running downstream applications, you need to set the `LD_LIBRARY_PATH` so `libtvm_ffi` is available in the path. ```bash export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:`tvm-ffi-config --libdir` ``` ## Running Examples You can run an optional library-loading example similar to the quick_start examples in [examples/quick_start](../examples/quick_start/). ```bash cargo run --example load_library --features example ``` Check out the [load_library.rs](tvm-ffi/examples/load_library.rs) for details. tvm-ffi-0.1.12/rust/tvm-ffi-macros/000077500000000000000000000000001521067262500170005ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi-macros/Cargo.toml000066400000000000000000000021371521067262500207330ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [package] name = "tvm-ffi-macros" description = "Procedural macro crate for tvm-ffi" version = "0.1.0-alpha.0" edition = "2021" license = "Apache-2.0" [lib] proc-macro = true [dependencies] proc-macro2 = "^1.0" quote = "^1.0" syn = { version = "1.0.48", features = ["full", "parsing", "extra-traits"] } proc-macro-error = "^1.0" tvm-ffi-0.1.12/rust/tvm-ffi-macros/src/000077500000000000000000000000001521067262500175675ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi-macros/src/lib.rs000066400000000000000000000024561521067262500207120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use proc_macro::TokenStream; use proc_macro_error::proc_macro_error; mod object_macros; mod utils; #[proc_macro_error] #[proc_macro_derive(Object, attributes(type_key, type_index))] pub fn derive_object(input: TokenStream) -> TokenStream { TokenStream::from(object_macros::derive_object(input)) } #[proc_macro_error] #[proc_macro_derive(ObjectRef, attributes(type_key, type_index))] pub fn derive_object_ref(input: TokenStream) -> TokenStream { TokenStream::from(object_macros::derive_object_ref(input)) } tvm-ffi-0.1.12/rust/tvm-ffi-macros/src/object_macros.rs000066400000000000000000000240101521067262500227440ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use proc_macro::TokenStream; use quote::quote; use syn::DeriveInput; use crate::utils::*; /// Derive Object trait for a struct to generate boilerplate code pub fn derive_object(input: proc_macro::TokenStream) -> TokenStream { let tvm_ffi_crate = get_tvm_ffi_crate(); let derive_input = syn::parse_macro_input!(input as DeriveInput); let struct_name = derive_input.ident.clone(); let type_key = get_attr(&derive_input, "type_key") .map(attr_to_str) .expect("Expect #[type_key = \"\"] attribute"); // type index can be optional // for now we make it required for static index let type_index_tokens = match get_attr(&derive_input, "type_index").map(attr_to_expr) { Some(type_index) => { let type_index_expr = type_index.expect("Expect #[type_index(TypeIndex::)] attribute"); quote! { #[inline] fn type_index() -> i32 { #type_index_expr as i32 } } } None => { quote! { #[inline] fn type_index() -> i32 { static TYPE_INDEX: std::sync::LazyLock = std::sync::LazyLock::new(|| unsafe { let type_key_arg = #tvm_ffi_crate::tvm_ffi_sys::TVMFFIByteArray::from_str(#type_key); let mut tindex = 0; let ret = #tvm_ffi_crate::tvm_ffi_sys::TVMFFITypeKeyToIndex( &type_key_arg, &mut tindex ); if ret != 0 { proc_macro_error::abort!("Failed to get type index for type key: {}", #type_key); } tindex } ); *TYPE_INDEX } } } }; // search for field name base and derive the base type // we expect base always to be the first field let base_def_tokens = match &derive_input.data { syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| { let (base_id, base_ty) = (f.ident.clone()?, f.ty.clone()); // The transitive case of subtyping Some(quote! { #[inline] unsafe fn object_header_mut( this: &mut Self ) -> &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject { const _: () = { fn assert_impl() {} let _ = assert_impl::<#base_ty>; }; #base_ty::object_header_mut(&mut this.#base_id) } }) }), _ => panic!("First field must be `: `"), }; let expanded = quote! { unsafe impl #tvm_ffi_crate::object::ObjectCore for #struct_name { const TYPE_KEY: &'static str = #type_key; #type_index_tokens #base_def_tokens } }; TokenStream::from(expanded) } /// Derive ObjectRef trait for a struct to generate boilerplate code pub fn derive_object_ref(input: proc_macro::TokenStream) -> TokenStream { let tvm_ffi_crate = get_tvm_ffi_crate(); let derive_input = syn::parse_macro_input!(input as DeriveInput); let struct_name = derive_input.ident.clone(); // search for field name base and derive the base type // we expect base always to be the first field let data_ty = match &derive_input.data { syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| { let (data_id, data_ty) = (f.ident.clone()?, f.ty.clone()); if data_id == "data" { // The transitive case of subtyping Some(data_ty) } else { None } }), _ => panic!("derive only works for structs"), } .expect("First field must be `data: ObjectArc`"); let mut expanded = quote! { unsafe impl #tvm_ffi_crate::object::ObjectRefCore for #struct_name { type ContainerType = <#data_ty as std::ops::Deref>::Target; #[inline] fn data(this: &Self) -> &ObjectArc { &this.data } #[inline] fn into_data(this: Self) -> ObjectArc { this.data } #[inline] fn from_data(data: ObjectArc) -> Self { Self { data} } } // implement AnyCompatible for #struct_name unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for #struct_name { fn type_str() -> String { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; ::TYPE_KEY.into() } unsafe fn copy_to_any_view( src: &Self, data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny ) { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let type_index = ::type_index(); data.type_index = type_index as i32; data.small_str_len = 0; let data_ptr = #tvm_ffi_crate::object::ObjectArc::::as_raw( &src.data ); data.data_union.v_obj = data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject; } unsafe fn check_any_strict(data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny) -> bool { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let type_index = ::type_index(); data.type_index == type_index as i32 } unsafe fn copy_from_any_view_after_check( data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny ) -> Self { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let data_ptr = data.data_union.v_obj; // need to increase ref because original weak ptr // do not own the code #tvm_ffi_crate::object::unsafe_::inc_ref( data_ptr as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject ); Self { data : #tvm_ffi_crate::object::ObjectArc::from_raw( data_ptr as *mut ContainerType ) } } unsafe fn move_to_any( src: Self, data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny ) { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let type_index = ::type_index(); data.type_index = type_index as i32; data.small_str_len = 0; let data_ptr = #tvm_ffi_crate::object::ObjectArc::into_raw( src.data ); data.data_union.v_obj = data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject; } unsafe fn move_from_any_after_check( data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny ) -> Self { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let data_ptr = data.data_union.v_obj as *mut ContainerType; Self { data : #tvm_ffi_crate::object::ObjectArc::::from_raw(data_ptr) } } unsafe fn try_cast_from_any_view( data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny ) -> Result { type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore> ::ContainerType; let type_index = ::type_index(); if data.type_index == type_index as i32 { Ok(Self::copy_from_any_view_after_check(data)) } else { Err(()) } } } }; // skip ObjectRef since it can create circular dependency with any.rs if struct_name != "ObjectRef" { expanded.extend(quote! { #tvm_ffi_crate::impl_try_from_any!(#struct_name); #tvm_ffi_crate::impl_arg_into_ref!(#struct_name); #tvm_ffi_crate::impl_into_arg_holder_default!(#struct_name); }); } TokenStream::from(expanded) } tvm-ffi-0.1.12/rust/tvm-ffi-macros/src/utils.rs000066400000000000000000000047021521067262500213000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use proc_macro2::TokenStream; use quote::quote; use std::env; /// Get the tvm-rt crate name /// \return The tvm-rt crate name pub(crate) fn get_tvm_ffi_crate() -> TokenStream { if env::var("CARGO_PKG_NAME").unwrap() == "tvm-ffi" { quote!(crate) } else { quote!(tvm_ffi) } } /// Get an attribute by name from a derive input /// /// # Arguments /// * `derive_input` - The derive input to get the attribute from /// * `name` - The name of the attribute to get /// /// # Returns /// * `Option<&syn::Attribute>` - The attribute if it exists pub(crate) fn get_attr<'a>( derive_input: &'a syn::DeriveInput, name: &str, ) -> Option<&'a syn::Attribute> { derive_input.attrs.iter().find(|a| a.path.is_ident(name)) } /// Convert an attribute to a string /// /// # Arguments /// * `attr` - The attribute to convert /// /// # Returns /// * `syn::LitStr` - The string value of the attribute pub(crate) fn attr_to_str(attr: &syn::Attribute) -> syn::LitStr { match attr.parse_meta() { Ok(syn::Meta::NameValue(syn::MetaNameValue { lit: syn::Lit::Str(s), .. })) => s, Ok(_m) => panic!("Expected a string literal, got"), Err(e) => panic!("{}", e), } } /// Convert an attribute to an integer /// /// # Arguments /// * `attr` - The attribute to convert /// /// # Returns /// * `syn::Result` - The integer value of the attribute pub(crate) fn attr_to_expr(attr: &syn::Attribute) -> syn::Result { let parser = |input: syn::parse::ParseStream| { input.parse::() // parse expression after '=' }; syn::parse::Parser::parse2(parser, attr.tokens.clone()) } tvm-ffi-0.1.12/rust/tvm-ffi-sys/000077500000000000000000000000001521067262500163325ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi-sys/Cargo.toml000066400000000000000000000017261521067262500202700ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [package] name = "tvm-ffi-sys" description = "Low-level sys crate for tvm-ffi" version = "0.1.0-alpha.0" edition = "2021" license = "Apache-2.0" [lib] name = "tvm_ffi_sys" crate-type = ["lib"] tvm-ffi-0.1.12/rust/tvm-ffi-sys/build.rs000066400000000000000000000066061521067262500200070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::env; use std::process::Command; /// Update the LD_LIBRARY_PATH environment variable /// so cargo run/test can directly pick it up /// note that it won't always work later consumption of the /// library, and we still need to figure out linking by setting LD_LIBRARY_PATH fn update_ld_library_path(lib_dir: &str) { let os_env_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() { Ok("windows") => "PATH", Ok("macos") => "DYLD_LIBRARY_PATH", Ok("linux") => "LD_LIBRARY_PATH", _ => "", }; if os_env_var.is_empty() { return; } // Get the current value of the environment variable at build time (if any) let current_val = env::var(os_env_var).unwrap_or_else(|_| String::new()); // Use platform-specific separator let separator = if os_env_var == "PATH" { ";" } else { ":" }; let new_ld_path = if current_val.is_empty() { lib_dir.to_string() } else { format!("{}{}{}", current_val, separator, lib_dir) }; // this env is only used for cargo run/test println!("cargo:rustc-env={}={}", os_env_var, new_ld_path); } fn main() { // When building documentation, we may not need actual linking // Check if this is a rustdoc build let is_rustdoc = env::var("RUSTDOC").is_ok() || env::var("CARGO_CFG_DOC").is_ok(); // Run `mylib-config --libdir` to get the library path let config_result = Command::new("tvm-ffi-config").arg("--libdir").output(); let lib_dir = match config_result { Ok(output) if output.status.success() => String::from_utf8(output.stdout) .unwrap_or_default() .trim() .to_string(), _ if is_rustdoc => { // For rustdoc builds, we can proceed without the library eprintln!("Warning: tvm-ffi-config not available, skipping for documentation build"); return; } _ => { panic!("Failed to run tvm-ffi-config. Make sure tvm-ffi is installed."); } }; if lib_dir.is_empty() { if is_rustdoc { eprintln!( "Warning: Empty lib_dir from tvm-ffi-config, skipping for documentation build" ); return; } panic!("tvm-ffi-config returned empty library path"); } // add the library directory to the linker search path println!("cargo:rustc-link-search=native={}", lib_dir); // link the library println!("cargo:rustc-link-lib=dylib=tvm_ffi"); println!("cargo:rustc-link-lib=dylib=tvm_ffi_testing"); // update the LD_LIBRARY_PATH environment variable update_ld_library_path(&lib_dir); } tvm-ffi-0.1.12/rust/tvm-ffi-sys/src/000077500000000000000000000000001521067262500171215ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi-sys/src/c_api.rs000066400000000000000000000354221521067262500205500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOTE: we manually write the C ABI as they are reasonably minimal // and we need to ensure clear control of the atomic access etc. #![allow(non_camel_case_types)] use std::ffi::c_void; use std::sync::atomic::AtomicU64; use crate::dlpack::DLDataType; use crate::dlpack::DLDevice; /// The index type of the FFI objects #[repr(i32)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum TVMFFITypeIndex { /// None/nullptr value kTVMFFINone = 0, /// POD int value kTVMFFIInt = 1, /// POD bool value kTVMFFIBool = 2, /// POD float value kTVMFFIFloat = 3, /// Opaque pointer object kTVMFFIOpaquePtr = 4, /// DLDataType kTVMFFIDataType = 5, /// DLDevice kTVMFFIDevice = 6, /// DLTensor* kTVMFFIDLTensorPtr = 7, /// const char* kTVMFFIRawStr = 8, /// TVMFFIByteArray* kTVMFFIByteArrayPtr = 9, /// R-value reference to ObjectRef kTVMFFIObjectRValueRef = 10, /// Small string on stack kTVMFFISmallStr = 11, /// Small bytes on stack kTVMFFISmallBytes = 12, /// Start of statically defined objects. kTVMFFIStaticObjectBegin = 64, /// String object, layout = { TVMFFIObject, TVMFFIByteArray, ... } kTVMFFIStr = 65, /// Bytes object, layout = { TVMFFIObject, TVMFFIByteArray, ... } kTVMFFIBytes = 66, /// Error object. kTVMFFIError = 67, /// Function object. kTVMFFIFunction = 68, /// Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... } kTVMFFIShape = 69, /// Tensor object, layout = { TVMFFIObject, DLTensor, ... } kTVMFFITensor = 70, /// Array object. kTVMFFIArray = 71, //---------------------------------------------------------------- // more complex objects //---------------------------------------------------------------- /// Map object. kTVMFFIMap = 72, /// Runtime dynamic loaded module object. kTVMFFIModule = 73, /// Opaque python object. kTVMFFIOpaquePyObject = 74, } #[repr(i32)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum TVMFFIObjectDeleterFlagBitMask { kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0, kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1, kTVMFFIObjectDeleterFlagBitMaskBoth = (1 << 0) | (1 << 1), } /// Handle to Object from C API's pov pub type TVMFFIObjectHandle = *mut c_void; pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32); // constants for working with combined reference count pub const COMBINED_REF_COUNT_MASK_U32: u64 = (1u64 << 32) - 1; pub const COMBINED_REF_COUNT_STRONG_ONE: u64 = 1; pub const COMBINED_REF_COUNT_WEAK_ONE: u64 = 1u64 << 32; pub const COMBINED_REF_COUNT_BOTH_ONE: u64 = COMBINED_REF_COUNT_STRONG_ONE | COMBINED_REF_COUNT_WEAK_ONE; #[repr(C)] pub struct TVMFFIObject { pub combined_ref_count: AtomicU64, pub type_index: i32, pub __padding: u32, pub deleter: Option, // private padding to ensure 8 bytes alignment #[cfg(target_pointer_width = "32")] __padding: u32, } impl TVMFFIObject { pub fn new() -> Self { Self { combined_ref_count: AtomicU64::new(0), type_index: 0, __padding: 0, deleter: None, } } } /// Second union in TVMFFIAny - 8 bytes #[repr(C)] #[derive(Copy, Clone)] pub union TVMFFIAnyDataUnion { /// Integers pub v_int64: i64, /// Floating-point numbers pub v_float64: f64, /// Typeless pointers pub v_ptr: *mut c_void, /// Raw C-string pub v_c_str: *const i8, /// Ref counted objects pub v_obj: *mut TVMFFIObject, /// Data type pub v_dtype: DLDataType, /// Device pub v_device: DLDevice, /// Small string pub v_bytes: [u8; 8], /// uint64 repr mainly used for hashing pub v_uint64: u64, } /// TVM FFI Any value - a union type that can hold various data types #[repr(C)] #[derive(Copy, Clone)] pub struct TVMFFIAny { /// Type index of the object. /// The type index of Object and Any are shared in FFI. pub type_index: i32, /// small string length or zero padding pub small_str_len: u32, /// data union - 8 bytes pub data_union: TVMFFIAnyDataUnion, } impl TVMFFIAny { /// create a new instance of TVMFFIAny that represents None pub fn new() -> Self { Self { type_index: TVMFFITypeIndex::kTVMFFINone as i32, small_str_len: 0, data_union: TVMFFIAnyDataUnion { v_int64: 0 }, } } } /// Byte array data structure used by String and Bytes. #[repr(C)] pub struct TVMFFIByteArray { pub data: *const u8, pub size: usize, } impl TVMFFIByteArray { pub fn new(data: *const u8, size: usize) -> Self { Self { data, size } } /// Convert the TVMFFIByteArray to a str view /// /// # Arguments /// * `self` - The TVMFFIByteArray to convert. /// /// # Returns /// * `&str` - The converted str view. pub fn as_str(&self) -> &str { unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.data, self.size)) } } /// Unsafe function to create a TVMFFIByteArray from a string /// This function is unsafe as it does not check lifetime of the string /// the caller must ensure that the string is valid for the lifetime of the TVMFFIByteArray /// /// # Arguments /// * `data` - The string to create the TVMFFIByteArray from. /// /// # Returns /// * `TVMFFIByteArray` - The created TVMFFIByteArray. pub unsafe fn from_str(data: &str) -> Self { Self { data: data.as_ptr(), size: data.len(), } } } /// Safe call type for function ABI pub type TVMFFISafeCallType = unsafe extern "C" fn( handle: *mut c_void, args: *const TVMFFIAny, num_args: i32, result: *mut TVMFFIAny, ) -> i32; /// Function cell #[repr(C)] pub struct TVMFFIFunctionCell { /// A C API compatible call with exception catching. pub safe_call: TVMFFISafeCallType, pub cxx_call: *mut c_void, } unsafe impl Send for TVMFFIFunctionCell {} unsafe impl Sync for TVMFFIFunctionCell {} #[repr(i32)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum TVMFFIBacktraceUpdateMode { kTVMFFIBacktraceUpdateModeReplace = 0, kTVMFFIBacktraceUpdateModeAppend = 1, } /// Error cell used in error object following header. #[repr(C)] pub struct TVMFFIErrorCell { pub kind: TVMFFIByteArray, pub message: TVMFFIByteArray, pub backtrace: TVMFFIByteArray, pub update_backtrace: unsafe extern "C" fn( self_ptr: *mut c_void, backtrace: *const TVMFFIByteArray, update_mode: i32, ), } /// Shape cell used in shape object following header. #[repr(C)] pub struct TVMFFIShapeCell { pub data: *const i64, pub size: usize, } /// Field getter function pointer type pub type TVMFFIFieldGetter = unsafe extern "C" fn(field: *mut c_void, result: *mut TVMFFIAny) -> i32; /// Field setter function pointer type pub type TVMFFIFieldSetter = unsafe extern "C" fn(field: *mut c_void, value: *const TVMFFIAny) -> i32; /// Information support for optional object reflection #[repr(C)] pub struct TVMFFIFieldInfo { /// The name of the field pub name: TVMFFIByteArray, /// The docstring about the field pub doc: TVMFFIByteArray, /// The metadata of the field in JSON string pub metadata: TVMFFIByteArray, /// bitmask flags of the field pub flags: i64, /// The size of the field pub size: i64, /// The alignment of the field pub alignment: i64, /// The offset of the field pub offset: i64, /// The getter to access the field pub getter: Option, /// The setter to access the field. /// /// When kTVMFFIFieldFlagBitSetterIsFunctionObj is NOT set (default), /// this is a TVMFFIFieldSetter function pointer cast to *mut c_void. /// When kTVMFFIFieldFlagBitSetterIsFunctionObj IS set, /// this is a TVMFFIObjectHandle pointing to a FunctionObj. /// /// The setter is set even if the field is readonly for serialization. pub setter: *mut c_void, /// The default value or factory of the field, this field holds AnyView. /// Valid when flags set kTVMFFIFieldFlagBitMaskHasDefault. /// When kTVMFFIFieldFlagBitMaskDefaultFromFactory is also set, /// this is a callable factory function () -> Any. pub default_value_or_factory: TVMFFIAny, /// Records the compile-time static type kind of the field. pub field_static_type_index: i32, } /// Object creator function pointer type pub type TVMFFIObjectCreator = unsafe extern "C" fn(result: *mut TVMFFIObjectHandle) -> i32; /// Method information that can appear in reflection table #[repr(C)] pub struct TVMFFIMethodInfo { /// The name of the field pub name: TVMFFIByteArray, /// The docstring about the method pub doc: TVMFFIByteArray, /// Optional metadata of the method in JSON string pub metadata: TVMFFIByteArray, /// bitmask flags of the method pub flags: i64, /// The method wrapped as ffi::Function, stored as AnyView /// The first argument to the method is always the self for instance methods pub method: TVMFFIAny, } /// Extra information of object type that can be used for reflection /// /// This information is optional and can be used to enable reflection based /// creation of the object. #[repr(C)] pub struct TVMFFITypeMetadata { /// The docstring about the object pub doc: TVMFFIByteArray, /// An optional function that can create a new empty instance of the type pub creator: Option, /// Total size of the object struct, if it is fixed and known /// /// This field is set optional and set to 0 if not registered. pub total_size: i32, /// Optional meta-data for structural eq/hash pub structural_eq_hash_kind: i32, } /// Column array that stores extra attributes about types /// /// The attributes stored in a column array that can be looked up by type index. /// Note that the TypeAttr behaves like type_traits so column T so not contain /// attributes from base classes. #[repr(C)] pub struct TVMFFITypeAttrColumn { /// The data of the column, indexed by (type_index - begin_index). pub data: *const TVMFFIAny, /// The number of elements in the data array. /// The column covers type indices [begin_index, begin_index + size). pub size: i32, /// The starting type index of the column data. /// Lookup: if begin_index <= type_index < begin_index + size, /// the entry is data[(type_index - begin_index) as usize]. pub begin_index: i32, } /// Runtime type information for object type checking #[repr(C)] pub struct TVMFFITypeInfo { /// The runtime type index /// It can be allocated during runtime if the type is dynamic pub type_index: i32, /// number of parent types in the type hierachy pub type_depth: i32, /// the unique type key to identify the type pub type_key: TVMFFIByteArray, /// `type_acenstors[depth]` stores the type_index of the acenstors at depth level /// To keep things simple, we do not allow multiple inheritance so the /// hieracy stays as a tree pub type_acenstors: *const *const TVMFFITypeInfo, /// Cached hash value of the type key, used for consistent structural hashing pub type_key_hash: u64, /// number of reflection accessible fields pub num_fields: i32, /// number of reflection acccesible methods pub num_methods: i32, /// The reflection field information pub fields: *const TVMFFIFieldInfo, /// The reflection method pub methods: *const TVMFFIMethodInfo, /// The extra information of the type pub metadata: *const TVMFFITypeMetadata, } unsafe extern "C" { pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32; pub fn TVMFFIFunctionGetGlobal( name: *const TVMFFIByteArray, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFIFunctionSetGlobal( name: *const TVMFFIByteArray, f: TVMFFIObjectHandle, can_override: i32, ) -> i32; pub fn TVMFFIFunctionCreate( self_ptr: *mut c_void, safe_call: TVMFFISafeCallType, deleter: Option, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFIAnyViewToOwnedAny(any_view: *const TVMFFIAny, out: *mut TVMFFIAny) -> i32; pub fn TVMFFIFunctionCall( func: TVMFFIObjectHandle, args: *const TVMFFIAny, num_args: i32, result: *mut TVMFFIAny, ) -> i32; pub fn TVMFFIErrorMoveFromRaised(result: *mut TVMFFIObjectHandle); pub fn TVMFFIErrorSetRaised(error: TVMFFIObjectHandle); pub fn TVMFFIErrorSetRaisedFromCStr(kind: *const i8, message: *const i8); pub fn TVMFFIErrorCreate( kind: *const TVMFFIByteArray, message: *const TVMFFIByteArray, backtrace: *const TVMFFIByteArray, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFITensorFromDLPack( from: *mut c_void, require_alignment: i32, require_contiguous: i32, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFITensorToDLPack(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32; pub fn TVMFFITensorFromDLPackVersioned( from: *mut c_void, require_alignment: i32, require_contiguous: i32, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFITensorToDLPackVersioned(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32; pub fn TVMFFIStringFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32; pub fn TVMFFIBytesFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32; pub fn TVMFFIDataTypeFromString(str: *const TVMFFIByteArray, out: *mut DLDataType) -> i32; pub fn TVMFFIDataTypeToString(dtype: *const DLDataType, out: *mut TVMFFIAny) -> i32; pub fn TVMFFITraceback( filename: *const i8, lineno: i32, func: *const i8, cross_ffi_boundary: i32, ) -> *const TVMFFIByteArray; pub fn TVMFFIGetTypeInfo(type_index: i32) -> *const TVMFFITypeInfo; pub fn TVMFFITestingDummyTarget() -> i32; } tvm-ffi-0.1.12/rust/tvm-ffi-sys/src/c_env_api.rs000066400000000000000000000055371521067262500214240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // NOTE: we manually write the C ABI as they are reasonably minimal // and we need to ensure clear control of the atomic access etc. #![allow(non_camel_case_types)] use std::ffi::c_void; use std::os::raw::c_char; use crate::c_api::TVMFFIObjectHandle; use crate::dlpack::DLTensor; // ---------------------------------------------------------------------------- // Stream context // Focusing on minimalistic thread-local context recording stream being used. // We explicitly not handle allocation/de-allocation of stream here. // ---------------------------------------------------------------------------- /// The type of the stream handle. pub type TVMFFIStreamHandle = *mut c_void; /// DLPack tensor allocator function type pub type DLPackManagedTensorAllocator = unsafe extern "C" fn( prototype: *mut DLTensor, out: *mut *mut c_void, // DLManagedTensorVersioned** error_ctx: *mut c_void, set_error: unsafe extern "C" fn(*mut c_void, *const c_char, *const c_char), ) -> i32; unsafe extern "C" { pub fn TVMFFIEnvSetStream( device_type: i32, device_id: i32, stream: TVMFFIStreamHandle, opt_out_original_stream: *mut TVMFFIStreamHandle, ) -> i32; pub fn TVMFFIEnvGetStream(device_type: i32, device_id: i32) -> TVMFFIStreamHandle; pub fn TVMFFIEnvSetDLPackManagedTensorAllocator( allocator: DLPackManagedTensorAllocator, write_to_global_context: i32, opt_out_original_allocator: *mut DLPackManagedTensorAllocator, ) -> i32; pub fn TVMFFIEnvGetDLPackManagedTensorAllocator() -> DLPackManagedTensorAllocator; pub fn TVMFFIEnvCheckSignals() -> i32; pub fn TVMFFIEnvRegisterCAPI(name: *const c_char, symbol: *mut c_void) -> i32; pub fn TVMFFIEnvModLookupFromImports( library_ctx: TVMFFIObjectHandle, func_name: *const c_char, out: *mut TVMFFIObjectHandle, ) -> i32; pub fn TVMFFIEnvModRegisterContextSymbol(name: *const c_char, symbol: *mut c_void) -> i32; pub fn TVMFFIEnvModRegisterSystemLibSymbol(name: *const c_char, symbol: *mut c_void) -> i32; } tvm-ffi-0.1.12/rust/tvm-ffi-sys/src/dlpack.rs000066400000000000000000000064671521067262500207420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // DLPack C ABI declarations // NOTE: we manually write the C ABI as they are reasonably minimal // and we need to ensure clear control of the atomic access etc. #![allow(non_camel_case_types)] // DLPack related declarations #[repr(u32)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum DLDeviceType { kDLCPU = 1, kDLCUDA = 2, kDLCUDAHost = 3, kDLOpenCL = 4, kDLVulkan = 7, kDLMetal = 8, kDLVPI = 9, kDLROCM = 10, kDLROCMHost = 11, kDLExtDev = 12, kDLCUDAManaged = 13, kDLOneAPI = 14, kDLWebGPU = 15, kDLHexagon = 16, kDLMAIA = 17, kDLTrn = 18, } #[repr(C)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DLDevice { pub device_type: DLDeviceType, pub device_id: i32, } /// DLPack data type code enum #[repr(u8)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum DLDataTypeCode { kDLInt = 0, kDLUInt = 1, kDLFloat = 2, kDLBfloat = 4, kDLComplex = 5, kDLOpaqueHandle = 3, kDLBool = 6, kDLFloat8_e3m4 = 7, kDLFloat8_e4m3 = 8, kDLFloat8_e4m3b11fnuz = 9, kDLFloat8_e4m3fn = 10, kDLFloat8_e4m3fnuz = 11, kDLFloat8_e5m2 = 12, kDLFloat8_e5m2fnuz = 13, kDLFloat8_e8m0fnu = 14, kDLFloat6_e2m3fn = 15, kDLFloat6_e3m2fn = 16, kDLFloat4_e2m1fn = 17, } /// DLPack data type struct #[repr(C)] #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DLDataType { pub code: u8, pub bits: u8, pub lanes: u16, } /// DLPack tensor struct - plain C tensor object, does not manage memory #[repr(C)] #[derive(Debug, Clone, PartialEq, Eq)] pub struct DLTensor { /// The data pointer points to the allocated data pub data: *mut core::ffi::c_void, /// The device of the tensor pub device: DLDevice, /// Number of dimensions pub ndim: i32, /// The data type of the pointer pub dtype: DLDataType, /// The shape of the tensor pub shape: *mut i64, /// Strides of the tensor (in number of elements, not bytes) /// Can be NULL, indicating tensor is compact and row-majored pub strides: *mut i64, /// The offset in bytes to the beginning pointer to data pub byte_offset: u64, } impl DLDevice { pub fn new(device_type: DLDeviceType, device_id: i32) -> Self { Self { device_type: device_type, device_id: device_id, } } } impl DLDataType { pub fn new(code: DLDataTypeCode, bits: u8, lanes: u16) -> Self { Self { code: code as u8, bits: bits, lanes: lanes, } } } tvm-ffi-0.1.12/rust/tvm-ffi-sys/src/lib.rs000066400000000000000000000016531521067262500202420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ pub mod c_api; pub mod c_env_api; pub mod dlpack; pub use crate::c_api::*; pub use crate::c_env_api::*; pub use crate::dlpack::*; tvm-ffi-0.1.12/rust/tvm-ffi/000077500000000000000000000000001521067262500155165ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/Cargo.toml000066400000000000000000000024221521067262500174460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. [package] name = "tvm-ffi" description = "tvm-ffi rust support" version = "0.1.0-alpha.0" edition = "2021" license = "Apache-2.0" [lib] name = "tvm_ffi" crate-type = ["lib"] [dependencies] paste = "1.0" tvm-ffi-sys = { version = "0.1.0-alpha.0", path = "../tvm-ffi-sys" } tvm-ffi-macros = { version = "0.1.0-alpha.0", path = "../tvm-ffi-macros" } # only needed for example building [features] example = [] [[example]] name = "load_library" path = "examples/load_library.rs" required-features = ["example"] tvm-ffi-0.1.12/rust/tvm-ffi/build.rs000066400000000000000000000061121521067262500171630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::env; use std::process::Command; fn generate_example_lib() { let is_example_build = env::var("CARGO_FEATURE_EXAMPLE").is_ok(); if !is_example_build { return; } println!("Running optional build step to generate example library"); let output_dir = env::var("OUT_DIR").unwrap(); let _ = Command::new("python") .arg("scripts/generate_example_lib.py") .arg(output_dir) .output() .expect("Failed to generate example library"); println!("cargo:rerun-if-changed=scripts/generate_example_lib.py"); } /// Update the LD_LIBRARY_PATH environment variable /// so cargo run/test can directly pick it up /// note that it won't always work later consumption of the /// library, and we still need to figure out linking by setting LD_LIBRARY_PATH fn update_ld_library_path(lib_dir: &str) { let os_env_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() { Ok("windows") => "PATH", Ok("macos") => "DYLD_LIBRARY_PATH", Ok("linux") => "LD_LIBRARY_PATH", _ => "", }; if os_env_var.is_empty() { return; } // Get the current value of the environment variable at build time (if any) let current_val = env::var(os_env_var).unwrap_or_else(|_| String::new()); // Use platform-specific separator let separator = if os_env_var == "PATH" { ";" } else { ":" }; let new_ld_path = if current_val.is_empty() { lib_dir.to_string() } else { format!("{}{}{}", current_val, separator, lib_dir) }; // this env is only used for cargo run/test println!("cargo:rustc-env={}={}", os_env_var, new_ld_path); } fn main() { // Run `mylib-config --libdir` to get the library path let config_output = Command::new("tvm-ffi-config") .arg("--libdir") .output() .expect("Failed to run tvm-ffi-config"); let lib_dir = String::from_utf8(config_output.stdout) .expect("Invalid UTF-8 output from tvm-ffi-config") .trim() .to_string(); // update the LD_LIBRARY_PATH environment variable // note that we will also need to update ld_library_path for // the cases here besides the tvm-ffi-sys crate so cargo test works out of the box update_ld_library_path(&lib_dir); // generate the example library generate_example_lib(); } tvm-ffi-0.1.12/rust/tvm-ffi/examples/000077500000000000000000000000001521067262500173345ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/examples/load_library.rs000066400000000000000000000035111521067262500223450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::{Module, Result, Tensor}; fn main() { // lib_path is build using scripts/generate_example_lib.py [out_path] // invoked automatically through build.rs // you can replace it to any other ffi compatible library let lib_path = concat!(env!("OUT_DIR"), "/add_one_cpu.so"); let lib: tvm_ffi::Module = Module::load_from_file(lib_path).unwrap(); // get the function from the library let add_one_cpu: tvm_ffi::Function = lib.get_function("add_one_cpu").unwrap(); // load the function from the library // and cast into a typed version of the function let typed_add_one = tvm_ffi::into_typed_fn!( add_one_cpu, Fn(&Tensor, &Tensor) -> Result<()> ); // run the function let x_data: &[f32] = &[0.0, 1.0, 2.0, 3.0]; let x = Tensor::from_slice(x_data, &[4]).unwrap(); let y = Tensor::from_slice(&[0.0f32; 4], &[4]).unwrap(); println!("x: {:?}", x.data_as_slice::().unwrap()); typed_add_one(&x, &y).unwrap(); println!("y: {:?}", y.data_as_slice::().unwrap()); } tvm-ffi-0.1.12/rust/tvm-ffi/scripts/000077500000000000000000000000001521067262500172055ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/scripts/generate_example_lib.py000066400000000000000000000043121521067262500237120ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Script to generate add_one_cpu.so used in examples.""" import pathlib import shutil import sys import tvm_ffi.cpp def main() -> None: """Generate the FFI library for examples.""" if len(sys.argv) != 2: print("Usage: python generate_example_lib.py ") sys.exit(1) output_dir = sys.argv[1] output_lib_path = tvm_ffi.cpp.build_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions=["add_one_cpu"], ) target_lib_path = pathlib.Path(output_dir) / "add_one_cpu.so" print(f"Generated FFI library at {target_lib_path}") shutil.copy(output_lib_path, target_lib_path) if __name__ == "__main__": main() tvm-ffi-0.1.12/rust/tvm-ffi/src/000077500000000000000000000000001521067262500163055ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/src/any.rs000066400000000000000000000243251521067262500174500ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::error::Error; use crate::object; use crate::type_traits::AnyCompatible; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; use tvm_ffi_sys::{TVMFFIAny, TVMFFIAnyViewToOwnedAny}; /// Unmanaged Any that can hold reference to values #[derive(Copy, Clone)] #[repr(C)] pub struct AnyView<'a> { data: TVMFFIAny, /// needs to explicit mark lifetime to avoid lifetime mismatch _phantom: std::marker::PhantomData<&'a ()>, } /// Managed Any that can hold reference to values #[repr(C)] pub struct Any { data: TVMFFIAny, } //--------------------- // AnyView //--------------------- impl<'a> AnyView<'a> { pub fn new() -> Self { Self { data: TVMFFIAny::new(), _phantom: std::marker::PhantomData, } } #[inline] pub fn type_index(&self) -> i32 { self.data.type_index } /// More strict version than try_from/try_into /// /// This function will not try to cast the type /// and ensures invariance that the return value is only Some(T) /// Any::from(T) contains exactly the same value value /// /// will return Some(T) if the type is exactly compatible with T /// will return None if the type is not compatible with T #[inline] pub fn try_as(&self) -> Option where T: AnyCompatible, { unsafe { if T::check_any_strict(&self.data) { Some(T::copy_from_any_view_after_check(&self.data)) } else { None } } } /// Get the strong count of the underlying object for testing/debugging purposes /// /// If the underlying object is not ref counted, return None pub fn debug_strong_count(&self) -> Option { unsafe { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { Some(object::unsafe_::strong_count(self.data.data_union.v_obj)) } else { None } } } } impl<'a, T: AnyCompatible> From<&'a T> for AnyView<'a> { #[inline] fn from(value: &'a T) -> Self { unsafe { let mut data = TVMFFIAny::new(); T::copy_to_any_view(&value, &mut data); Self { data: data, _phantom: std::marker::PhantomData, } } } } impl Default for AnyView<'_> { fn default() -> Self { Self::new() } } /// Holder for Any value /// /// This is used to define try_from rule while conforming to orphan rule /// Users should not use this directly pub struct TryFromTemp { value: T, } impl TryFromTemp { /// Create a new holder for the value #[inline(always)] pub fn new(value: T) -> Self { Self { value } } /// Move the value out of the holder #[inline(always)] pub fn into_value(this: Self) -> T { this.value } } //--------------------- // Any //--------------------- impl Any { pub fn new() -> Self { Self { data: TVMFFIAny::new(), } } #[inline] pub fn type_index(&self) -> i32 { self.data.type_index } /// Try to query if stored typed in Any exactly matches the type T /// /// This function is fast in the case of failure and can be used to check /// if the type is compatible with T /// /// This function will not try to cast the type /// and ensures invariance that the return value is only Some(T) /// `Any::from(T)` contains exactly the same value value /// `Any::try_as()` contains exactly the same value value /// /// will return Some(T) if the type is exactly compatible with T /// will return None if the type is not compatible with T #[inline] pub fn try_as(&self) -> Option where T: AnyCompatible, { unsafe { if T::check_any_strict(&self.data) { Some(T::copy_from_any_view_after_check(&self.data)) } else { None } } } #[inline] pub unsafe fn as_data_ptr(&mut self) -> *mut TVMFFIAny { &mut self.data } #[inline] pub unsafe fn into_raw_ffi_any(this: Self) -> TVMFFIAny { let this = std::mem::ManuallyDrop::new(this); this.data } #[inline] pub unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self { Self { data } } /// Get the strong count of the underlying object for testing/debugging purposes /// /// If the underlying object is not ref counted, return None pub fn debug_strong_count(&self) -> Option { unsafe { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { Some(object::unsafe_::strong_count(self.data.data_union.v_obj)) } else { None } } } } impl Default for Any { fn default() -> Self { Self::new() } } impl Clone for Any { #[inline] fn clone(&self) -> Self { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { object::unsafe_::inc_ref(self.data.data_union.v_obj) } } Self { data: self.data } } } impl Drop for Any { #[inline] fn drop(&mut self) { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { object::unsafe_::dec_ref(self.data.data_union.v_obj) } } } } // convert Any ref to AnyView impl<'a> From<&'a Any> for AnyView<'a> { #[inline] fn from(value: &'a Any) -> Self { Self { data: value.data, _phantom: std::marker::PhantomData, } } } // convert AnyView to Any impl From> for Any { #[inline] fn from(value: AnyView<'_>) -> Self { unsafe { let mut data = TVMFFIAny::new(); crate::check_safe_call!(TVMFFIAnyViewToOwnedAny(&value.data, &mut data)).unwrap(); Self { data } } } } impl From for Any { #[inline] fn from(value: T) -> Self { unsafe { let mut data = TVMFFIAny::new(); T::move_to_any(value, &mut data); Self { data } } } } impl<'a, T: AnyCompatible> TryFrom> for TryFromTemp { type Error = crate::error::Error; #[inline] fn try_from(value: AnyView<'a>) -> Result { unsafe { if T::check_any_strict(&value.data) { Ok(TryFromTemp::new(T::copy_from_any_view_after_check( &value.data, ))) } else { T::try_cast_from_any_view(&value.data) .map_err(|_| { let msg = format!( "Cannot convert from type `{}` to `{}`", T::get_mismatch_type_info(&value.data), T::type_str() ); crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "") }) .map(TryFromTemp::new) } } } } impl TryFrom for TryFromTemp { type Error = crate::error::Error; #[inline] fn try_from(value: Any) -> Result { unsafe { if T::check_any_strict(&value.data) { let mut value = std::mem::ManuallyDrop::new(value); Ok(TryFromTemp::new(T::move_from_any_after_check( &mut value.data, ))) } else { T::try_cast_from_any_view(&value.data) .map_err(|_| { let msg = format!( "Cannot convert from type `{}` to `{}`", T::get_mismatch_type_info(&value.data), T::type_str() ); crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "") }) .map(TryFromTemp::new) } } } } crate::impl_try_from_any!( bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64, (), *mut core::ffi::c_void, crate::string::String, crate::string::Bytes, crate::object::ObjectRef, tvm_ffi_sys::dlpack::DLDataType, tvm_ffi_sys::dlpack::DLDevice, ); crate::impl_try_from_any_for_parametric!(Option); //------------------------------------------------------------ /// ArgTryFromAnyView: Helper for function argument passing ///----------------------------------------------------------- pub(crate) trait ArgTryFromAnyView: Sized { fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result; } impl ArgTryFromAnyView for T { fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result { unsafe { if T::check_any_strict(&value.data) { Ok(T::copy_from_any_view_after_check(&value.data)) } else { T::try_cast_from_any_view(&value.data).map_err(|_| { let msg = format!( "Argument #{}: Cannot convert from type `{}` to `{}`", arg_index, T::get_mismatch_type_info(&value.data), T::type_str() ); crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "") }) } } } } tvm-ffi-0.1.12/rust/tvm-ffi/src/collections/000077500000000000000000000000001521067262500206235ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/src/collections/array.rs000066400000000000000000000237151521067262500223170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::fmt::Debug; use std::marker::PhantomData; use std::ops::Deref; use crate::any::TryFromTemp; use crate::derive::Object; use crate::object::{Object, ObjectArc}; use crate::{Any, AnyCompatible, AnyView, ObjectCoreWithExtraItems, ObjectRefCore}; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; use tvm_ffi_sys::{TVMFFIAny, TVMFFIObject}; #[repr(C)] #[derive(Object)] #[type_key = "ffi.Array"] #[type_index(TypeIndex::kTVMFFIArray)] pub struct ArrayObj { pub object: Object, /// Pointer to the start of the element buffer (AddressOf(0)). pub data: *mut core::ffi::c_void, pub size: i64, pub capacity: i64, /// Optional custom deleter for the data pointer. pub data_deleter: Option, } unsafe impl ObjectCoreWithExtraItems for ArrayObj { type ExtraItem = TVMFFIAny; fn extra_items_count(this: &Self) -> usize { this.size as usize } } #[repr(C)] #[derive(Clone)] pub struct Array { data: ObjectArc, _marker: PhantomData, } impl Debug for Array { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let full_name = std::any::type_name::(); let short_name = full_name.split("::").last().unwrap_or(full_name); write!(f, "Array<{}>[{}]", short_name, self.len()) } } impl Default for Array { fn default() -> Self { Self::new(vec![]) } } unsafe impl ObjectRefCore for Array { type ContainerType = ArrayObj; fn data(this: &Self) -> &ObjectArc { &this.data } fn into_data(this: Self) -> ObjectArc { this.data } fn from_data(data: ObjectArc) -> Self { Self { data, _marker: PhantomData, } } } impl Array { /// Creates a new Array from a vector of items. pub fn new(items: Vec) -> Self { let capacity = items.len(); Self::new_with_capacity(items, capacity) } /// Internal helper to allocate an ArrayObj with specific headroom. fn new_with_capacity(items: Vec, capacity: usize) -> Self { let size = items.len(); // Allocate with capacity let arc = ObjectArc::::new_with_extra_items(ArrayObj { object: Object::new(), data: core::ptr::null_mut(), size: size as i64, capacity: capacity as i64, data_deleter: None, }); unsafe { let raw_ptr = ObjectArc::as_raw(&arc) as *mut ArrayObj; let container = &mut *raw_ptr; let base_ptr = ArrayObj::extra_items_mut(container).as_ptr() as *mut TVMFFIAny; container.data = base_ptr as *mut _; for (i, item) in items.into_iter().enumerate() { let any: Any = Any::from(item); let raw = Any::into_raw_ffi_any(any); core::ptr::write(base_ptr.add(i), raw); } } Self::from_data(arc) } pub fn len(&self) -> usize { self.data.size as usize } pub fn is_empty(&self) -> bool { self.len() == 0 } /// Retrieves an item at the given index. pub fn get(&self, index: usize) -> Result { if index >= self.len() { crate::bail!(crate::error::INDEX_ERROR, "Array get index out of bound"); } unsafe { let container = self.data.deref(); let base_ptr = container.data as *const TVMFFIAny; let raw_any_ref = &*base_ptr.add(index); match T::try_cast_from_any_view(raw_any_ref) { Ok(val) => Ok(val), Err(_) => crate::bail!( crate::error::TYPE_ERROR, "Failed to cast element at {} to {}", index, T::type_str() ), } } } pub fn iter(&'_ self) -> ArrayIterator<'_, T> { ArrayIterator { array: self, index: 0, len: self.len(), } } #[inline] fn as_container(&self) -> &ArrayObj { unsafe { let ptr = ObjectArc::as_raw(&self.data) as *const ArrayObj; &*ptr } } } // --- Index Implementation --- impl std::ops::Index for Array { type Output = AnyView<'static>; fn index(&self, index: usize) -> &Self::Output { let container = self.as_container(); let len = container.size as usize; if index >= len { panic!( "Index out of bounds: the len is {} but the index is {}", len, index ); } unsafe { let ptr = (container.data as *const AnyView<'static>).add(index); &*ptr } } } // --- Iterator Implementations --- pub struct ArrayIterator<'a, T: AnyCompatible + Clone> { array: &'a Array, index: usize, len: usize, } impl<'a, T: AnyCompatible + Clone> Iterator for ArrayIterator<'a, T> { type Item = T; fn next(&mut self) -> Option { if self.index < self.len { let item = self.array.get(self.index).ok(); self.index += 1; item } else { None } } } impl<'a, T: AnyCompatible + Clone> IntoIterator for &'a Array { type Item = T; type IntoIter = ArrayIterator<'a, T>; fn into_iter(self) -> Self::IntoIter { self.iter() } } impl FromIterator for Array { fn from_iter>(iter: I) -> Self { let items: Vec = iter.into_iter().collect(); Self::new(items) } } // --- Any Type System Conversions --- unsafe impl AnyCompatible for Array where T: AnyCompatible + Clone + 'static, { fn type_str() -> String { format!("Array<{}>", T::type_str()) } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { if data.type_index != TypeIndex::kTVMFFIArray as i32 { return false; } if std::any::TypeId::of::() == std::any::TypeId::of::() { return true; } let container = &*(data.data_union.v_obj as *const ArrayObj); let base_ptr = container.data as *const TVMFFIAny; for i in 0..container.size { let elem_any = &*base_ptr.add(i as usize); if !T::check_any_strict(elem_any) { return false; } } true } unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIArray as i32; data.data_union.v_obj = ObjectArc::as_raw(Self::data(src)) as *mut TVMFFIObject; data.small_str_len = 0; } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIArray as i32; data.data_union.v_obj = ObjectArc::into_raw(Self::into_data(src)) as *mut TVMFFIObject; data.small_str_len = 0; } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { let ptr = data.data_union.v_obj as *const ArrayObj; crate::object::unsafe_::inc_ref(ptr as *mut TVMFFIObject); Self::from_data(ObjectArc::from_raw(ptr)) } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { let ptr = data.data_union.v_obj as *const ArrayObj; let obj = Self::from_data(ObjectArc::from_raw(ptr)); data.type_index = TypeIndex::kTVMFFINone as i32; data.data_union.v_int64 = 0; obj } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index != TypeIndex::kTVMFFIArray as i32 { return Err(()); } // Fast path: if types match exactly, we can just copy the reference. if Self::check_any_strict(data) { return Ok(Self::copy_from_any_view_after_check(data)); } // Slow path: try to convert element by element. let container = &*(data.data_union.v_obj as *const ArrayObj); let base_ptr = container.data as *const TVMFFIAny; let mut items = Vec::with_capacity(container.size as usize); for i in 0..container.size { let any_v = &*base_ptr.add(i as usize); if let Ok(item) = T::try_cast_from_any_view(any_v) { items.push(item); } else { return Err(()); } } Ok(Array::new(items)) } } impl TryFrom for Array where T: AnyCompatible + Clone + 'static, { type Error = crate::error::Error; fn try_from(value: Any) -> Result { let temp: TryFromTemp = TryFromTemp::try_from(value)?; Ok(TryFromTemp::into_value(temp)) } } impl<'a, T> TryFrom> for Array where T: AnyCompatible + Clone + 'static, { type Error = crate::error::Error; fn try_from(value: AnyView<'a>) -> Result { let temp: TryFromTemp = TryFromTemp::try_from(value)?; Ok(TryFromTemp::into_value(temp)) } } tvm-ffi-0.1.12/rust/tvm-ffi/src/collections/mod.rs000066400000000000000000000015531521067262500217540ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /// Collection types pub mod array; pub mod shape; pub mod tensor; tvm-ffi-0.1.12/rust/tvm-ffi/src/collections/shape.rs000066400000000000000000000074451521067262500223030ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::derive::{Object, ObjectRef}; use crate::object::{Object, ObjectArc, ObjectCoreWithExtraItems}; use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd}; use std::ops::Deref; use tvm_ffi_sys::TVMFFIShapeCell; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; //----------------------------------------------------- // Shape //----------------------------------------------------- // ShapeObj for heap-allocated shape #[repr(C)] #[derive(Object)] #[type_key = "ffi.Shape"] #[type_index(TypeIndex::kTVMFFIShape)] pub struct ShapeObj { object: Object, data: TVMFFIShapeCell, } /// ABI stable owned Shape for ffi #[repr(C)] #[derive(ObjectRef, Clone)] pub struct Shape { data: ObjectArc, } impl Shape { /// Create a new empty Shape pub fn new() -> Self { let shape_obj = ShapeObj { object: Object::new(), data: TVMFFIShapeCell { data: std::ptr::null(), size: 0, }, }; Self { data: ObjectArc::new(shape_obj), } } /// Get the shape as a slice pub fn as_slice(&self) -> &[i64] { unsafe { std::slice::from_raw_parts(self.data.data.data, self.data.data.size) } } /// Fill the strides from the shape pub fn fill_strides_from_shape(shape: T, strides: &mut [i64]) where T: AsRef<[i64]>, { let shape = shape.as_ref(); let mut stride = 1; for i in (0..shape.len()).rev() { strides[i] = stride; stride *= shape[i]; } } } unsafe impl ObjectCoreWithExtraItems for ShapeObj { type ExtraItem = i64; fn extra_items_count(this: &Self) -> usize { this.data.size } } impl From for Shape where T: AsRef<[i64]>, { fn from(value: T) -> Self { unsafe { let value_slice: &[i64] = value.as_ref(); let mut obj_arc = ObjectArc::new_with_extra_items(ShapeObj { object: Object::new(), data: TVMFFIShapeCell { data: std::ptr::null(), size: value_slice.len(), }, }); // reset the data ptr correctly after Arc is created obj_arc.data.data = ShapeObj::extra_items(&obj_arc).as_ptr(); let extra_items = ShapeObj::extra_items_mut(&mut obj_arc); extra_items.copy_from_slice(value_slice); Self { data: obj_arc } } } } impl Deref for Shape { type Target = [i64]; #[inline] fn deref(&self) -> &[i64] { self.as_slice() } } impl PartialEq for Shape { #[inline] fn eq(&self, other: &Self) -> bool { self.as_slice() == other.as_slice() } } impl Eq for Shape {} impl PartialOrd for Shape { #[inline] fn partial_cmp(&self, other: &Self) -> Option { self.as_slice().partial_cmp(other.as_slice()) } } impl Ord for Shape { #[inline] fn cmp(&self, other: &Self) -> Ordering { self.as_slice().cmp(other.as_slice()) } } tvm-ffi-0.1.12/rust/tvm-ffi/src/collections/tensor.rs000066400000000000000000000264671521067262500225220ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::collections::shape::Shape; use crate::derive::{Object, ObjectRef}; use crate::dtype::AsDLDataType; use crate::dtype::DLDataTypeExt; use crate::error::Result; use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems}; use tvm_ffi_sys::dlpack::{DLDataType, DLDevice, DLDeviceType, DLTensor}; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; //----------------------------------------------------- // NDAllocator Trait //----------------------------------------------------- /// Trait for n-dimensional array allocators pub unsafe trait NDAllocator: 'static { /// The minimum alignment of the data allocated by the allocator const MIN_ALIGN: usize; /// Allocate data for the given DLTensor /// /// # Arguments /// * `tensor` - The DLTensor to allocate data for /// /// This method should fill in the data pointer of the DLTensor. unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void; /// Free data for the given DLTensor /// /// # Arguments /// * `tensor` - The DLTensor to free data for /// /// This method should free the data pointer of the DLTensor. unsafe fn free_data(&mut self, tensor: &DLTensor); } /// DLTensorExt trait /// This trait provides methods to get the number of elements and the item size of a DLTensor pub trait DLTensorExt { fn numel(&self) -> usize; fn item_size(&self) -> usize; } impl DLTensorExt for DLTensor { fn numel(&self) -> usize { unsafe { std::slice::from_raw_parts(self.shape, self.ndim as usize) .iter() .product::() as usize } } fn item_size(&self) -> usize { (self.dtype.bits as usize * self.dtype.lanes as usize + 7) / 8 } } //----------------------------------------------------- // Shape //----------------------------------------------------- // ShapeObj for heap-allocated shape #[repr(C)] #[derive(Object)] #[type_key = "ffi.Tensor"] #[type_index(TypeIndex::kTVMFFITensor)] pub struct TensorObj { object: Object, dltensor: DLTensor, } /// ABI stable owned Shape for ffi #[repr(C)] #[derive(ObjectRef, Clone)] pub struct Tensor { data: ObjectArc, } impl Tensor { /// Get the data pointer of the Tensor /// /// # Returns /// * `*mut core::ffi::c_void` - The data pointer of the Tensor pub fn data_ptr(&self) -> *const core::ffi::c_void { self.data.dltensor.data } /// Get the data pointer of the Tensor /// /// # Returns /// * `*mut core::ffi::c_void` - The data pointer of the Tensor pub fn data_ptr_mut(&mut self) -> *mut core::ffi::c_void { self.data.dltensor.data } /// Check if the Tensor is contiguous /// /// # Returns /// * `bool` - True if the Tensor is contiguous, false otherwise pub fn is_contiguous(&self) -> bool { let strides = self.strides(); let shape = self.shape(); let mut expected_stride = 1; for i in (0..self.ndim()).rev() { if strides[i] != expected_stride { return false; } expected_stride *= shape[i]; } true } pub fn data_as_slice(&self) -> Result<&[T]> { let dtype = T::DL_DATA_TYPE; if self.dtype() != dtype { crate::bail!( crate::error::TYPE_ERROR, "Data type mismatch {} vs {}", self.dtype().to_string(), dtype.to_string() ); } if self.device().device_type != DLDeviceType::kDLCPU { crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU"); } crate::ensure!( self.is_contiguous(), crate::error::RUNTIME_ERROR, "Tensor is not contiguous" ); unsafe { Ok(std::slice::from_raw_parts( self.data.dltensor.data as *const T, self.numel(), )) } } /// Returns the tensor data as a mutable slice. /// /// This method takes `&self` rather than `&mut self` by design: like /// `std::fs::File::write`, the *metadata* of a Tensor (shape, dtype, /// device) is governed by Rust's ownership rules, but writing to the /// underlying data buffer (CPU memory or a GPU pointer) is a side-effect /// outside Rust's aliasing model. Most C/CUDA kernel APIs accept a /// non-mut Tensor and mutate its data content, so requiring `&mut self` /// here would force artificial mutability annotations throughout the /// deep-learning stack with no real safety benefit. /// /// # Safety contract (caller responsibility) /// If the `Tensor` has been cloned (via `ObjectArc`), the caller must /// ensure no other clone is concurrently reading the data. #[allow(clippy::wrong_self_convention)] pub fn data_as_slice_mut(&self) -> Result<&mut [T]> { let dtype = T::DL_DATA_TYPE; if self.dtype() != dtype { crate::bail!( crate::error::TYPE_ERROR, "Data type mismatch: expected {}, got {}", dtype.to_string(), self.dtype().to_string() ); } if self.device().device_type != DLDeviceType::kDLCPU { crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU"); } crate::ensure!( self.is_contiguous(), crate::error::RUNTIME_ERROR, "Tensor is not contiguous" ); unsafe { Ok(std::slice::from_raw_parts_mut( self.data.dltensor.data as *mut T, self.numel(), )) } } pub fn shape(&self) -> &[i64] { unsafe { std::slice::from_raw_parts(self.data.dltensor.shape, self.ndim()) } } pub fn ndim(&self) -> usize { self.data.dltensor.ndim as usize } pub fn numel(&self) -> usize { self.data.dltensor.numel() } pub fn strides(&self) -> &[i64] { unsafe { std::slice::from_raw_parts(self.data.dltensor.strides, self.ndim()) } } pub fn dtype(&self) -> DLDataType { self.data.dltensor.dtype } pub fn device(&self) -> DLDevice { self.data.dltensor.device } } struct TensorObjFromNDAlloc where TNDAlloc: NDAllocator, { base: TensorObj, alloc: TNDAlloc, } unsafe impl ObjectCore for TensorObjFromNDAlloc { const TYPE_KEY: &'static str = TensorObj::TYPE_KEY; fn type_index() -> i32 { TensorObj::type_index() } unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject { TensorObj::object_header_mut(&mut this.base) } } unsafe impl ObjectCoreWithExtraItems for TensorObjFromNDAlloc { type ExtraItem = i64; fn extra_items_count(this: &Self) -> usize { (this.base.dltensor.ndim * 2) as usize } } impl Drop for TensorObjFromNDAlloc { fn drop(&mut self) { unsafe { self.alloc.free_data(&self.base.dltensor); } } } impl Tensor { // Create a Tensor from a NDAllocator /// /// # Arguments /// * `alloc` - The NDAllocator /// * `shape` - The shape of the Tensor /// * `dtype` - The data type of the Tensor /// * `device` - The device of the Tensor /// /// # Returns /// * `Tensor` - The created Tensor pub fn from_nd_alloc( alloc: TNDAlloc, shape: &[i64], dtype: DLDataType, device: DLDevice, ) -> Self where TNDAlloc: NDAllocator, { let tensor_obj = TensorObjFromNDAlloc { base: TensorObj { object: Object::new(), dltensor: DLTensor { data: std::ptr::null_mut(), device: device, ndim: shape.len() as i32, dtype: dtype, shape: std::ptr::null_mut(), strides: std::ptr::null_mut(), byte_offset: 0, }, }, alloc: alloc, }; unsafe { let mut obj_arc = ObjectArc::new_with_extra_items(tensor_obj); obj_arc.base.dltensor.shape = TensorObjFromNDAlloc::extra_items(&obj_arc).as_ptr() as *mut i64; obj_arc.base.dltensor.strides = obj_arc.base.dltensor.shape.add(shape.len()); let extra_items = TensorObjFromNDAlloc::extra_items_mut(&mut obj_arc); extra_items[..shape.len()].copy_from_slice(shape); Shape::fill_strides_from_shape(shape, &mut extra_items[shape.len()..]); let dltensor_ptr = &obj_arc.base.dltensor as *const DLTensor; obj_arc.base.dltensor.data = obj_arc.alloc.alloc_data(&*dltensor_ptr); Self { data: ObjectArc::from_raw(ObjectArc::into_raw(obj_arc) as *mut TensorObj), } } } /// Create a Tensor from a slice /// /// # Arguments /// * `slice` - The slice to create the Tensor from /// * `shape` - The shape of the Tensor /// /// # Returns /// * `Tensor` - The created Tensor pub fn from_slice(slice: &[T], shape: &[i64]) -> Result { let dtype = T::DL_DATA_TYPE; let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, shape, dtype, device); if tensor.numel() != slice.len() { crate::bail!(crate::error::VALUE_ERROR, "Slice length mismatch"); } tensor.data_as_slice_mut::()?.copy_from_slice(slice); Ok(tensor) } } /// Example CPU NDAllocator /// This allocator allocates data on the CPU pub struct CPUNDAlloc {} unsafe impl NDAllocator for CPUNDAlloc { const MIN_ALIGN: usize = 64; unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void { let numel = prototype.numel() as usize; let item_size = prototype.item_size(); let size = numel * item_size as usize; let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap(); let ptr = std::alloc::alloc(layout); ptr as *mut core::ffi::c_void } unsafe fn free_data(&mut self, tensor: &DLTensor) { let numel = tensor.numel() as usize; let item_size = tensor.item_size(); let size = numel * item_size; let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap(); std::alloc::dealloc(tensor.data as *mut u8, layout); } } tvm-ffi-0.1.12/rust/tvm-ffi/src/derive.rs000066400000000000000000000015761521067262500201420ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /// namespace to re-export derive macros pub use tvm_ffi_macros::{Object, ObjectRef}; tvm-ffi-0.1.12/rust/tvm-ffi/src/device.rs000066400000000000000000000064231521067262500201170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::error::Result; use crate::type_traits::AnyCompatible; use tvm_ffi_sys::dlpack::DLDevice; use tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex as TypeIndex}; use tvm_ffi_sys::{TVMFFIEnvGetStream, TVMFFIEnvSetStream, TVMFFIStreamHandle}; /// Get the current stream for a device pub fn current_stream(device: &DLDevice) -> TVMFFIStreamHandle { unsafe { TVMFFIEnvGetStream(device.device_type as i32, device.device_id) } } /// Call `f` with the device stream temporarily set to `stream`. /// /// # Safety /// /// `stream` must be a valid stream handle for the given device, or null. pub unsafe fn with_stream( device: &DLDevice, stream: TVMFFIStreamHandle, f: impl FnOnce() -> Result, ) -> Result { let mut prev_stream: TVMFFIStreamHandle = std::ptr::null_mut(); unsafe { crate::check_safe_call!(TVMFFIEnvSetStream( device.device_type as i32, device.device_id, stream, &mut prev_stream as *mut TVMFFIStreamHandle ))?; } let result = f()?; unsafe { crate::check_safe_call!(TVMFFIEnvSetStream( device.device_type as i32, device.device_id, prev_stream, std::ptr::null_mut() ))?; } Ok(result) } /// AnyCompatible for DLDevice unsafe impl AnyCompatible for DLDevice { fn type_str() -> String { // make it consistent with c++ representation "Device".to_string() } unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIDevice as i32; data.small_str_len = 0; data.data_union.v_uint64 = 0; data.data_union.v_device = *src; } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIDevice as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; data.data_union.v_device = src; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return data.type_index == TypeIndex::kTVMFFIDevice as i32; } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_device } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_device } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIDevice as i32 { Ok(data.data_union.v_device) } else { Err(()) } } } tvm-ffi-0.1.12/rust/tvm-ffi/src/dtype.rs000066400000000000000000000162631521067262500200100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::error::Result; use crate::type_traits::AnyCompatible; /// Data type handling use tvm_ffi_sys::dlpack::{DLDataType, DLDataTypeCode}; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; use tvm_ffi_sys::{TVMFFIAny, TVMFFIByteArray, TVMFFIDataTypeFromString, TVMFFIDataTypeToString}; /// Extra methods for DLDataType pub trait DLDataTypeExt: Sized { /// Convert the DLDataType to a string representation /// /// # Returns /// A string representation of the data type (e.g., "int32", "float64", "bool") fn to_string(&self) -> crate::string::String; /// Parse a string representation into a DLDataType /// /// # Arguments /// * `dtype_str` - The string representation of the data type to parse /// /// # Returns /// * `Ok(DLDataType)` - Successfully parsed data type /// * `Err(Error)` - Failed to parse the string /// /// # Examples /// ``` /// use tvm_ffi::{DLDataType, DLDataTypeExt}; /// /// let dtype = DLDataType::try_from_str("int32").unwrap(); /// ``` fn try_from_str(dtype_str: &str) -> Result; } impl DLDataTypeExt for DLDataType { fn to_string(&self) -> crate::string::String { unsafe { let mut ffi_any = TVMFFIAny::new(); crate::check_safe_call!(TVMFFIDataTypeToString(&*self, &mut ffi_any)).unwrap(); crate::any::Any::from_raw_ffi_any(ffi_any) .try_into() .unwrap() } } fn try_from_str(dtype_str: &str) -> Result { let mut dtype = DLDataType { code: DLDataTypeCode::kDLOpaqueHandle as u8, bits: 0, lanes: 0, }; unsafe { let dtype_byte_array = TVMFFIByteArray::from_str(dtype_str); crate::check_safe_call!(TVMFFIDataTypeFromString(&dtype_byte_array, &mut dtype))?; } Ok(dtype) } } /// AnyCompatible implementation for DLDataType /// /// This implementation allows DLDataType to be used with the TVM FFI Any system, /// enabling type-safe conversion between DLDataType and the generic Any type. unsafe impl AnyCompatible for DLDataType { /// Get the type string identifier for DLDataType /// /// # Returns /// The string "DataType" to match the C++ representation fn type_str() -> String { // make it consistent with c++ representation "DataType".to_string() } /// Copy a DLDataType to an Any view /// /// # Arguments /// * `src` - The DLDataType to copy from /// * `data` - The Any view to copy to unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIDataType as i32; data.small_str_len = 0; data.data_union.v_uint64 = 0; data.data_union.v_dtype = *src; } /// Move a DLDataType into an Any /// /// # Arguments /// * `src` - The DLDataType to move from /// * `data` - The Any to move into unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIDataType as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; data.data_union.v_dtype = src; } /// Check if an Any contains a DLDataType /// /// # Arguments /// * `data` - The Any to check /// /// # Returns /// `true` if the Any contains a DLDataType, `false` otherwise unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return data.type_index == TypeIndex::kTVMFFIDataType as i32; } /// Copy a DLDataType from an Any view (after type check) /// /// # Arguments /// * `data` - The Any view to copy from /// /// # Returns /// The copied DLDataType /// /// # Safety /// The caller must ensure that `data` contains a DLDataType unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_dtype } /// Move a DLDataType from an Any (after type check) /// /// # Arguments /// * `data` - The Any to move from /// /// # Returns /// The moved DLDataType /// /// # Safety /// The caller must ensure that `data` contains a DLDataType unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_dtype } /// Try to cast an Any view to a DLDataType /// /// This method supports both direct DLDataType conversion and string parsing. /// /// # Arguments /// * `data` - The Any view to cast from /// /// # Returns /// * `Ok(DLDataType)` - Successfully cast to DLDataType /// * `Err(())` - Failed to cast (wrong type or invalid string) unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIDataType as i32 { Ok(data.data_union.v_dtype) } else if let Ok(string) = crate::string::String::try_cast_from_any_view(data) { DLDataType::try_from_str(string.as_str()).map_err(|_| ()) } else { Err(()) } } } /// Trait to convert standard data types to DLDataType /// /// This trait provides a way to get the corresponding DLDataType for standard Rust types. /// It's implemented for common integer, unsigned integer, and floating-point types. pub trait AsDLDataType: Copy { /// The corresponding DLDataType for this type const DL_DATA_TYPE: DLDataType; } /// Macro to implement AsDLDataType for standard types /// /// This macro generates implementations of the AsDLDataType trait for standard Rust types. /// It takes the type, the DLPack data type code, and the number of bits. macro_rules! impl_as_dl_data_type { ($type: ty, $code: expr, $bits: expr) => { impl AsDLDataType for $type { const DL_DATA_TYPE: DLDataType = DLDataType { code: $code as u8, bits: $bits as u8, lanes: 1, }; } }; } impl_as_dl_data_type!(i8, DLDataTypeCode::kDLInt, 8); impl_as_dl_data_type!(i16, DLDataTypeCode::kDLInt, 16); impl_as_dl_data_type!(i32, DLDataTypeCode::kDLInt, 32); impl_as_dl_data_type!(i64, DLDataTypeCode::kDLInt, 64); impl_as_dl_data_type!(u8, DLDataTypeCode::kDLUInt, 8); impl_as_dl_data_type!(u16, DLDataTypeCode::kDLUInt, 16); impl_as_dl_data_type!(u32, DLDataTypeCode::kDLUInt, 32); impl_as_dl_data_type!(u64, DLDataTypeCode::kDLUInt, 64); impl_as_dl_data_type!(f32, DLDataTypeCode::kDLFloat, 32); impl_as_dl_data_type!(f64, DLDataTypeCode::kDLFloat, 64); tvm-ffi-0.1.12/rust/tvm-ffi/src/error.rs000066400000000000000000000151631521067262500200120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::derive::{Object, ObjectRef}; use crate::object::{Object, ObjectArc}; use std::ffi::c_void; use tvm_ffi_sys::TVMFFIBacktraceUpdateMode::kTVMFFIBacktraceUpdateModeAppend; use tvm_ffi_sys::{ TVMFFIByteArray, TVMFFIErrorCell, TVMFFIErrorCreate, TVMFFIErrorMoveFromRaised, TVMFFIErrorSetRaised, TVMFFIObjectHandle, TVMFFITypeIndex, }; /// Error kind, wraps in a struct to be explicit #[derive(Debug, Clone, PartialEq, Eq)] pub struct ErrorKind<'a>(&'a str); impl<'a> ErrorKind<'a> { pub fn as_str(&self) -> &str { self.0 } } impl<'a> std::fmt::Display for ErrorKind<'a> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) } } pub const VALUE_ERROR: ErrorKind = ErrorKind("ValueError"); pub const TYPE_ERROR: ErrorKind = ErrorKind("TypeError"); pub const RUNTIME_ERROR: ErrorKind = ErrorKind("RuntimeError"); pub const ATTRIBUTE_ERROR: ErrorKind = ErrorKind("AttributeError"); pub const KEY_ERROR: ErrorKind = ErrorKind("KeyError"); pub const INDEX_ERROR: ErrorKind = ErrorKind("IndexError"); /// error object #[repr(C)] #[derive(Object)] #[type_key = "ffi.Error"] #[type_index(TVMFFITypeIndex::kTVMFFIError)] pub struct ErrorObj { object: Object, cell: TVMFFIErrorCell, } /// Error reference class #[derive(Clone, ObjectRef)] pub struct Error { data: ObjectArc, } /// Default result that uses Error as the error type pub type Result = std::result::Result; impl Error { pub fn new(kind: ErrorKind<'_>, message: &str, traceback: &str) -> Self { unsafe { let kind_data = TVMFFIByteArray::from_str(kind.as_str()); let message_data = TVMFFIByteArray::from_str(message); let traceback_data = TVMFFIByteArray::from_str(traceback); let mut error_handle: TVMFFIObjectHandle = std::ptr::null_mut(); let ret = TVMFFIErrorCreate( &kind_data, &message_data, &traceback_data, &mut error_handle, ); assert_eq!(ret, 0, "Failed to create error object"); let error_obj = ObjectArc::from_raw(error_handle as *const ErrorObj); Self { data: error_obj } } } /// Create a new error by moving from raised error /// /// # Returns /// The error from the raised error pub fn from_raised() -> Self { unsafe { let mut error_handle: TVMFFIObjectHandle = std::ptr::null_mut(); TVMFFIErrorMoveFromRaised(&mut error_handle as *mut TVMFFIObjectHandle); assert!( !error_handle.is_null(), "Calling Error::from_raised but no error was raised" ); let error_obj = ObjectArc::from_raw(error_handle as *const ErrorObj); Self { data: error_obj } } } /// Set the error as raised /// /// # Arguments /// * `error` - The error to set as raised pub fn set_raised(error: &Self) { unsafe { TVMFFIErrorSetRaised(ObjectArc::as_raw(&error.data) as TVMFFIObjectHandle); } } /// Get the kind of the error /// /// # Returns /// The kind of the error pub fn kind(&self) -> ErrorKind<'_> { ErrorKind(&self.data.cell.kind.as_str()) } /// Get the message of the error /// /// # Returns /// The message of the error pub fn message(&self) -> &str { self.data.cell.message.as_str() } /// Get the backtrace of the error /// /// # Returns /// The backtrace of the error pub fn backtrace(&self) -> &str { self.data.cell.backtrace.as_str() } /// Get the traceback of the error in the order of most recent call last /// /// # Returns /// The traceback of the error pub fn traceback_most_recent_call_last(&self) -> String { let backtrace = self.backtrace(); let backtrace_lines = backtrace.split('\n'); let mut traceback = String::new(); for line in backtrace_lines.rev() { traceback.push_str(line); traceback.push('\n'); } traceback } /// Append the backtrace to the error /// /// # Arguments /// * `this` - The error to append the backtrace to /// * `backtrace` - The backtrace to append /// /// # Returns /// The error with the appended backtrace pub fn with_appended_backtrace(this: Self, backtrace: &str) -> Self { if ObjectArc::strong_count(&this.data) == 1 { // this is the only reference to the error // we can safely mutate the error unsafe { let backtrace_data = TVMFFIByteArray::from_str(backtrace); (this.data.cell.update_backtrace)( ObjectArc::as_raw(&this.data) as *mut ErrorObj as *mut c_void, &backtrace_data, kTVMFFIBacktraceUpdateModeAppend as i32, ); this } } else { // we need to create a new error because there is more than one unique reference // to the error let mut new_backtrace = String::new(); new_backtrace.push_str(this.backtrace()); new_backtrace.push_str(backtrace); return Error::new(this.kind(), this.message(), &new_backtrace); } } } impl std::fmt::Display for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, "Traceback (most recent call last):\n{}{}: {}", self.traceback_most_recent_call_last(), self.kind().as_str(), self.message() ) } } impl std::fmt::Debug for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { std::fmt::Display::fmt(self, f) } } impl std::error::Error for Error {} tvm-ffi-0.1.12/rust/tvm-ffi/src/extra/000077500000000000000000000000001521067262500174305ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/src/extra/mod.rs000066400000000000000000000014701521067262500205570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ pub mod module; tvm-ffi-0.1.12/rust/tvm-ffi/src/extra/module.rs000066400000000000000000000052771521067262500212760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::derive::{Object, ObjectRef}; use crate::error::Result; use crate::function::Function; use crate::object::{Object, ObjectArc}; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; //----------------------------------------------------- // Module //----------------------------------------------------- /// A TVM FFI Module for loading dynamic libraries and retrieving functions. #[repr(C)] #[derive(Object)] #[type_key = "ffi.Module"] #[type_index(TypeIndex::kTVMFFIModule)] pub struct ModuleObj { object: Object, } /// ABI-stable owned Module for FFI operations. #[repr(C)] #[derive(ObjectRef, Clone)] pub struct Module { data: ObjectArc, } impl Module { /// Load a module from a dynamic library file. /// /// # Arguments /// * `file_name` - Path to the dynamic library file to load /// /// # Returns /// * `Result` - A `Module` instance on success pub fn load_from_file>(file_name: Str) -> Result { static API_FUNC: std::sync::LazyLock = std::sync::LazyLock::new(|| Function::get_global("ffi.ModuleLoadFromFile").unwrap()); let file_name = crate::string::String::from(file_name); (*API_FUNC) .call_tuple_with_len::<1, _>((file_name,))? .try_into() } /// Get a function from the module by name. /// /// # Arguments /// * `name` - The name of the function to retrieve /// /// # Returns /// * `Result` - A `Function` instance on success pub fn get_function>(&self, name: Str) -> Result { static API_FUNC: std::sync::LazyLock = std::sync::LazyLock::new(|| Function::get_global("ffi.ModuleGetFunction").unwrap()); let name = crate::string::String::from(name); (*API_FUNC) .call_tuple_with_len::<3, _>((self, name, true))? .try_into() } } tvm-ffi-0.1.12/rust/tvm-ffi/src/function.rs000066400000000000000000000235521521067262500205070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::any::{Any, AnyView}; use crate::derive::{Object, ObjectRef}; use crate::error::{Error, Result}; use crate::function_internal::{AsPackedCallable, TupleAsPackedArgs}; use crate::object::{Object, ObjectArc, ObjectCore}; use tvm_ffi_sys::{ TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate, TVMFFIFunctionGetGlobal, TVMFFIFunctionSetGlobal, TVMFFIObjectHandle, TVMFFISafeCallType, TVMFFITypeIndex, }; /// function object #[repr(C)] #[derive(Object)] #[type_key = "ffi.Function"] #[type_index(TVMFFITypeIndex::kTVMFFIFunction)] pub struct FunctionObj { object: Object, cell: TVMFFIFunctionCell, } /// Error reference class #[derive(Clone, ObjectRef)] pub struct Function { data: ObjectArc, } //------------------------------------------------------------------------ // CallbackFunctionObjImpl //------------------------------------------------------------------------ /// Special helper class to hold a generic callback state as Object /// Logically this Impl can be viewed as a FunctionObj /// We can create an ObjectArc> so the deleter /// can correctly delete the entire object including callback part /// then we will convert to ObjectArc to be used as function #[repr(C)] struct CallbackFunctionObjImpl Result + 'static> { function: FunctionObj, callback: F, } impl Result + 'static> CallbackFunctionObjImpl { pub fn from_callback(callback: F) -> Self { Self { function: FunctionObj { object: Object::new(), cell: TVMFFIFunctionCell { // specfic callback for F safe_call: Self::invoke_callback, cxx_call: std::ptr::null_mut(), }, }, callback, } } unsafe extern "C" fn invoke_callback( handle: *mut std::ffi::c_void, args: *const TVMFFIAny, num_args: i32, result: *mut TVMFFIAny, ) -> i32 { let this = &*(handle as *mut Self); let packed_args = std::slice::from_raw_parts(args as *const AnyView, num_args as usize); let ret_value = (this.callback)(packed_args); match ret_value { Ok(value) => { *result = Any::into_raw_ffi_any(value); 0 } Err(error) => { Error::set_raised(&error); -1 } } } } unsafe impl Result + 'static> ObjectCore for CallbackFunctionObjImpl { const TYPE_KEY: &'static str = FunctionObj::TYPE_KEY; fn type_index() -> i32 { FunctionObj::type_index() } unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject { FunctionObj::object_header_mut(&mut this.function) } } impl Function { /// Call the function in packed format. pub fn call_packed(&self, packed_args: &[AnyView]) -> Result { unsafe { let packed_args_ptr = packed_args.as_ptr() as *const TVMFFIAny; let mut result = Any::new(); let ret_code = (self.data.cell.safe_call)( ObjectArc::as_raw(&self.data) as *mut FunctionObj as *mut std::ffi::c_void, packed_args_ptr, packed_args.len() as i32, Any::as_data_ptr(&mut result), ); if ret_code == 0 { Ok(result) } else { Err(Error::from_raised()) } } } pub fn call_tuple(&self, tuple_args: TupleType) -> Result where TupleType: TupleAsPackedArgs, { // This is a workaround for Rust's requirement that stack allocation size // must be known at compile time for generic types. // While we know args_len is a constant, Rust doesn't allow us to directly // declare [AnyView::new(); args_len] in generic contexts. // // We use a small vector optimization pattern: // 1. First allocate a small stack buffer (stack_args) // 2. If args_len exceeds STACK_LEN, allocate a heap buffer (heap_args) // 3. Use the appropriate buffer based on size // // Since args_len is a compile-time constant, the compiler should optimize // away the unused branch, making this approach efficient. const STACK_LEN: usize = 4; let mut stack_args = [AnyView::new(); STACK_LEN]; let mut heap_args = Vec::::new(); let args_len = ::LEN; // get packed arguments let packed_args: &mut [AnyView] = if args_len <= STACK_LEN { &mut stack_args[..args_len] } else { heap_args.resize(args_len, AnyView::new()); &mut heap_args[..args_len] }; (&tuple_args).fill_any_view(packed_args); self.call_packed(packed_args) } /// Call function with compile-time known argument count /// This is an optimized version of call_tuple for when the argument count /// is known at compile time, avoiding the small vector optimization overhead. /// /// # Arguments /// * `tuple_args` - The tuple arguments /// /// # Returns /// * `Any` - The result pub fn call_tuple_with_len( &self, tuple_args: TupleType, ) -> Result where TupleType: TupleAsPackedArgs, { let mut packed_args = [AnyView::new(); LEN]; (&tuple_args).fill_any_view(&mut packed_args); self.call_packed(&packed_args) } /// Get global function by name /// This function will throw an error if the function is not found. /// /// # Arguments /// * `name` - The name of the function /// /// # Returns /// * `Function` - The global function pub fn get_global(name: &str) -> Result { unsafe { let name_arg = TVMFFIByteArray::from_str(name); let mut result: TVMFFIObjectHandle = ::std::ptr::null_mut(); crate::check_safe_call!(TVMFFIFunctionGetGlobal(&name_arg, &mut result))?; if result.is_null() { crate::bail!(crate::error::RUNTIME_ERROR, "Function {} not found", name); } Ok(Self { data: ObjectArc::::from_raw(result as *mut FunctionObj), }) } } /// Register a function as a global function /// # Arguments /// * `name` - The name of the function /// * `func` - The function to register /// /// # Returns /// * `Result<()>` - The result of the registration pub fn register_global(name: &str, func: Function) -> Result<()> { unsafe { let name_arg = TVMFFIByteArray::from_str(name); let can_override = 0; crate::check_safe_call!(TVMFFIFunctionSetGlobal( &name_arg, ObjectArc::as_raw(&func.data) as *mut FunctionObj as TVMFFIObjectHandle, can_override ))?; Ok(()) } } /// Construct a function from a packed function /// # Arguments /// * `func` - The packed function in signature of `Fn(&[AnyView]) -> Result` /// /// # Returns /// * `Function` - The function pub fn from_packed(func: F) -> Self where F: Fn(&[AnyView]) -> Result + 'static, { unsafe { let callback_arc = ObjectArc::new(CallbackFunctionObjImpl::from_callback(func)); let func_arc = ObjectArc::::from_raw( ObjectArc::into_raw(callback_arc) as *mut FunctionObj ); Self { data: func_arc } } } /// Construct a function from a typed function /// # Arguments /// * `func` - The typed function with function signature of `F(T0, T1, ...) -> Result` /// /// # Returns /// * `Function` - The function pub fn from_typed(func: F) -> Self where F: AsPackedCallable + 'static, { let closure = move |packed_args: &[AnyView]| -> Result { let ret_value = func.call_packed(packed_args)?; Ok(ret_value) }; Self::from_packed(closure) } /// # Safety /// /// `handle` must be a valid pointer (or null) that is compatible with /// `safe_call` and `deleter`. The caller must ensure the handle outlives /// the returned `Function` (or that `deleter` properly frees it). pub unsafe fn from_extern_c( handle: *mut std::ffi::c_void, safe_call: TVMFFISafeCallType, deleter: Option, ) -> Self { unsafe { let mut out_handle: TVMFFIObjectHandle = std::ptr::null_mut(); crate::check_safe_call!(TVMFFIFunctionCreate( handle, safe_call, deleter, &mut out_handle )) .unwrap(); Self { data: ObjectArc::::from_raw(out_handle as *mut FunctionObj), } } } } tvm-ffi-0.1.12/rust/tvm-ffi/src/function_internal.rs000066400000000000000000000152511521067262500224000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::any::{Any, AnyView, ArgTryFromAnyView}; use crate::error::Result; use crate::string::{Bytes, String}; use crate::type_traits::AnyCompatible; //------------------------------------------------------------------------ // PackedCallable //------------------------------------------------------------------------ pub trait AsPackedCallable { // Call the function in packed convention fn call_packed(&self, packed_args: &[AnyView]) -> Result; } #[inline] pub fn call_packed_callable(func: Fun, packed_args: &[AnyView]) -> Result where Fun: AsPackedCallable, { func.call_packed(packed_args) } macro_rules! impl_as_packed_callable { ($len:literal; $($t:ident),*) => { impl AsPackedCallable<($($t,)*), Out> for Fun where Fun: Fn($($t,)*) -> Result + 'static, Any: From, $($t: ArgTryFromAnyView),* { fn call_packed(&self, packed_args: &[AnyView]) -> Result { crate::ensure!( packed_args.len() == $len, crate::error::VALUE_ERROR, "Expected {} arguments, got {}", $len, packed_args.len() ); // Expand the function call, consuming the iterator. let mut _arg_iter = packed_args.iter().enumerate(); let ret_value = self( $({ // unwrap is safe due to the length check above let (i, view) = _arg_iter.next().unwrap(); $t::try_from_any_view(view, i)? }),* )?; Ok(Any::from(ret_value)) } } } } impl_as_packed_callable!(0;); impl_as_packed_callable!(1; T0); impl_as_packed_callable!(2; T0, T1); impl_as_packed_callable!(3; T0, T1, T2); impl_as_packed_callable!(4; T0, T1, T2, T3); impl_as_packed_callable!(5; T0, T1, T2, T3, T4); impl_as_packed_callable!(6; T0, T1, T2, T3, T4, T5); impl_as_packed_callable!(7; T0, T1, T2, T3, T4, T5, T6); impl_as_packed_callable!(8; T0, T1, T2, T3, T4, T5, T6, T7); //-------------------------------------------------------------- // IntoArgHolder, helper to convert to canonical holding type // // This is needed sometimes for reference types that may need to // be converted to value types. //-------------------------------------------------------------- pub trait IntoArgHolder { type Target; fn into_arg_holder(self) -> Self::Target; } crate::impl_into_arg_holder_default!( bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64, String, Bytes ); // string will be converted to String for argument passing impl IntoArgHolder for &str { type Target = String; fn into_arg_holder(self) -> Self::Target { String::from(self) } } // string will be converted to String for argument passing impl IntoArgHolder for &[u8] { type Target = Bytes; fn into_arg_holder(self) -> Self::Target { Bytes::from(self) } } // helper trait to implement IntoArgHolderTuple to apply into_arg_holder to each element pub trait IntoArgHolderTuple { type Target; fn into_arg_holder_tuple(self) -> Self::Target; } macro_rules! impl_into_arg_holder_tuple { ( $($T:ident),* ; $($idx:tt),* ) => { impl<$($T),*> $crate::function_internal::IntoArgHolderTuple for ($($T,)*) where $($T: IntoArgHolder),* { type Target = ($($T::Target,)*); fn into_arg_holder_tuple(self) -> Self::Target { ($(self.$idx.into_arg_holder(),)*) } } }; } impl_into_arg_holder_tuple!(;); impl_into_arg_holder_tuple!(T0; 0); impl_into_arg_holder_tuple!(T0, T1; 0, 1); impl_into_arg_holder_tuple!(T0, T1, T2; 0, 1, 2); impl_into_arg_holder_tuple!(T0, T1, T2, T3; 0, 1, 2, 3); impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4; 0, 1, 2, 3, 4); impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5); impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6); impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7); //------------------------------------------------------------ // ArgIntoRef // // Helper to turn argument type to reference type // This is effectively AsRef but removes the need of T //----------------------------------------------------------- pub trait ArgIntoRef { type Target; fn to_ref(&self) -> &Self::Target; } crate::impl_arg_into_ref!( bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64, String, Bytes ); //----------------------------------------------------------- // TupleAsPackedArgs // // Helper to turn tuple type to packed arguments //----------------------------------------------------------- pub trait TupleAsPackedArgs { const LEN: usize; fn fill_any_view<'a>(&'a self, any_view: &mut [AnyView<'a>]); } macro_rules! impl_tuple_as_packed_args { ( $len:expr; $($T:ident),* ; $($idx:tt),* ) => { impl<$($T),*> TupleAsPackedArgs for ($($T,)*) where $( $T: ArgIntoRef, $T::Target: AnyCompatible, )* { const LEN: usize = $len; fn fill_any_view<'a>(&'a self, _any_view: &mut [AnyView<'a>]) { $( _any_view[$idx] = AnyView::from(self.$idx.to_ref()); )* } } }; } impl_tuple_as_packed_args!(0;;); impl_tuple_as_packed_args!(1; T0; 0); impl_tuple_as_packed_args!(2; T0, T1; 0, 1); impl_tuple_as_packed_args!(3; T0, T1, T2; 0, 1, 2); impl_tuple_as_packed_args!(4; T0, T1, T2, T3; 0, 1, 2, 3); impl_tuple_as_packed_args!(5; T0, T1, T2, T3, T4; 0, 1, 2, 3, 4); impl_tuple_as_packed_args!(6; T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5); impl_tuple_as_packed_args!(7; T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6); impl_tuple_as_packed_args!(8; T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7); tvm-ffi-0.1.12/rust/tvm-ffi/src/lib.rs000066400000000000000000000036131521067262500174240ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ pub mod any; pub mod collections; pub mod derive; pub mod device; pub mod dtype; pub mod error; pub mod extra; pub mod function; pub mod function_internal; pub mod macros; pub mod object; pub mod string; pub mod type_traits; pub use tvm_ffi_sys; pub use crate::any::{Any, AnyView}; pub use crate::collections::array::Array; pub use crate::collections::shape::Shape; pub use crate::collections::tensor::{CPUNDAlloc, NDAllocator, Tensor}; pub use crate::device::{current_stream, with_stream}; pub use crate::dtype::DLDataTypeExt; pub use crate::error::{Error, ErrorKind, Result}; pub use crate::error::{ ATTRIBUTE_ERROR, INDEX_ERROR, KEY_ERROR, RUNTIME_ERROR, TYPE_ERROR, VALUE_ERROR, }; pub use crate::extra::module::Module; pub use crate::function::Function; pub use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems, ObjectRefCore}; pub use crate::string::{Bytes, String}; pub use crate::type_traits::AnyCompatible; pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; pub use tvm_ffi_sys::{ DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, TVMFFIAny, TVMFFIObject, TVMFFIStreamHandle, }; tvm-ffi-0.1.12/rust/tvm-ffi/src/macros.rs000066400000000000000000000357211521067262500201470ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // rexport paste under macro namespace so downstream do not need to specify dep pub use paste; // ---------------------------------------------------------------------------- // Macros for error handling // ---------------------------------------------------------------------------- /// Macro gto get the name of the function /// /// # Usage /// Usage: function_name!() #[macro_export] macro_rules! function_name { () => {{ // dummy function to get the name of the function fn f() {} fn type_name_of(_: T) -> &'static str { std::any::type_name::() } let name = type_name_of(f); // remove the f() from the name &name[..name.len() - 3] }}; } /// Check the return code of the safe call /// /// # Arguments /// * `ret_code` - The return code of the safe call /// /// # Returns /// * `Result<(), Error>` - The result of the safe call /// Macro to check safe calls and automatically update traceback with file/line info /// /// Usage: check_safe_call!(function(args))?; #[macro_export] macro_rules! check_safe_call { ($expr:expr) => {{ let ret_code = $expr; if ret_code == 0 { Ok(()) } else { let error = $crate::error::Error::from_raised(); Err(error) } }}; } /// Create a new error with file/line info attached /// /// This macro automatically appends file/line info to the traceback /// /// # Arguments /// * `error_kind` - The kind of the error /// * `msg` - The message of the error /// * `args` - The posisble format arguments /// /// # Returns /// * `Result<(), Error>` - The result of the safe call #[macro_export] macro_rules! bail { ($error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{ let context = format!( " File \"{}\", line {}, in {}\n", file!(), line!(), $crate::function_name!() ); return Err($crate::error::Error::new($error_kind, &format!($fmt $(, $args)*), &context)); }}; } /// Create a new error with file/line info attached /// /// This macro automatically appends file/line info to the traceback /// /// # Arguments /// * `kind` - The kind of the error /// * `msg` - The message of the error /// * `args` - The posisble format arguments /// /// # Returns /// * `Result<(), Error>` - The result of the safe call #[macro_export] macro_rules! ensure { ($cond:expr, $error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{ if !$cond { $crate::bail!($error_kind, $fmt $(, $args)*); } }}; } /// Attach a context to a result if it is error /// /// This macro automatically appends file/line info to the traceback /// /// # Arguments /// * `error` - The error to attach the context to /// * `msg` - The message of the error /// /// # Returns /// * `Result<(), Error>` - The result of the safe call #[macro_export] macro_rules! attach_context { ($error:expr) => {{ match $error { Ok(value) => Ok(value), Err(error) => { let context = format!( " File \"{}\", line {}, in {}\n", file!(), line!(), $crate::function_name!() ); Err(Error::with_appended_backtrace(error, &context)) } } }}; } // ---------------------------------------------------------------------------- // Macros for any definitions // ---------------------------------------------------------------------------- // implements try from any for all integer types /// Macro to implement `TryFrom` and `TryFrom` for a list of types #[macro_export] macro_rules! impl_try_from_any { ($($t:ty),* $(,)?) => { $( impl<'a> TryFrom<$crate::any::AnyView<'a>> for $t { type Error = $crate::error::Error; #[inline(always)] fn try_from( value: $crate::any::AnyView<'a> ) -> Result { type TryFromTemp = $crate::any::TryFromTemp<$t>; return TryFromTemp::try_from(value).map(TryFromTemp::into_value); } } impl TryFrom<$crate::any::Any> for $t { type Error = $crate::error::Error; #[inline(always)] fn try_from( value: $crate::any::Any ) -> Result { type TryFromTemp = $crate::any::TryFromTemp<$t>; return TryFromTemp::try_from(value).map(TryFromTemp::into_value); } } )* }; } /// Macro to implement `TryFrom` and `TryFrom` for generic types like `Option` #[macro_export] macro_rules! impl_try_from_any_for_parametric { ($generic_type:ident<$param:ident>) => { impl<'a, $param: AnyCompatible> TryFrom<$crate::any::AnyView<'a>> for $generic_type<$param> { type Error = $crate::error::Error; #[inline(always)] fn try_from(value: $crate::any::AnyView<'a>) -> Result { type TryFromTemp = $crate::any::TryFromTemp<$generic_type<$param>>; return TryFromTemp::::try_from(value).map(TryFromTemp::::into_value); } } impl<$param: AnyCompatible> TryFrom<$crate::any::Any> for $generic_type<$param> { type Error = $crate::error::Error; #[inline(always)] fn try_from(value: $crate::any::Any) -> Result { type TryFromTemp = $crate::any::TryFromTemp<$generic_type<$param>>; return TryFromTemp::::try_from(value).map(TryFromTemp::::into_value); } } }; } /// Macro to implement IntoArgHolder for a list of types #[macro_export] macro_rules! impl_into_arg_holder_default { ($($t:ty),*) => { $( impl $crate::function_internal::IntoArgHolder for $t { type Target = $t; fn into_arg_holder(self) -> Self::Target { self } } impl<'a> $crate::function_internal::IntoArgHolder for &'a $t { type Target = &'a $t; fn into_arg_holder(self) -> Self::Target { self } } )* }; } /// Macro to implement ArgIntoRef for a list of types #[macro_export] macro_rules! impl_arg_into_ref { ($($t:ty),*) => { $( impl $crate::function_internal::ArgIntoRef for $t { type Target = $t; fn to_ref(&self) -> &Self::Target { &self } } impl<'a> $crate::function_internal::ArgIntoRef for &'a $t { type Target = $t; fn to_ref(&self) -> &Self::Target { &self } } )* } } // ---------------------------------------------------------------------------- // Macros for function definitions // ---------------------------------------------------------------------------- /// Macro to export a typed function as a C symbol that follows the tvm-ffi ABI /// /// # Arguments /// * `$name` - The name of the function /// * `$func` - The function to export /// /// # Example /// ```rust /// use tvm_ffi::*; /// /// fn add_one(x: i32) -> Result { Ok(x + 1) } /// /// tvm_ffi_dll_export_typed_func!(add_one, add_one); /// ``` #[macro_export] macro_rules! tvm_ffi_dll_export_typed_func { ($name:ident, $func:expr) => { $crate::macros::paste::paste! { // `#[no_mangle]` is required so the symbol is preserved in a // `cdylib` and matches the `__tvm_ffi_` naming convention // that `ffi.Module.load_from_file.` looks up via // `GetSymbolWithSymbolPrefix`. Without it, the linker strips the // function from the output `.so`. // // Using plain `#[no_mangle]` (rather than `#[unsafe(no_mangle)]`, // which would require rustc >= 1.82) keeps the crate buildable // on older toolchains. Edition-2024 callers will see a // deprecation warning, which is harmless. // // The path-qualified `$crate::tvm_ffi_sys::…` reference (rather // than a bare `tvm_ffi_sys::…`) lets downstream crates use the // macro without having to add `tvm-ffi-sys` to their own // `[dependencies]`. #[no_mangle] pub unsafe extern "C" fn [<__tvm_ffi_ $name>]( _handle: *mut std::ffi::c_void, args: *const $crate::tvm_ffi_sys::TVMFFIAny, num_args: i32, result: *mut $crate::tvm_ffi_sys::TVMFFIAny, ) -> i32 { let packed_args = std::slice::from_raw_parts(args as *const $crate::any::AnyView, num_args as usize); let ret_value = $crate::function_internal::call_packed_callable($func, packed_args); match ret_value { Ok(value) => { *result = $crate::any::Any::into_raw_ffi_any(value); 0 } Err(error) => { $crate::error::Error::set_raised(&error); -1 } } } } }; } ///----------------------------------------------------------- /// into_typed_fn /// /// Converts a generic `Function` into a typed function with compile-time /// argument count and type checking. This macro provides a convenient way /// to create type-safe wrappers around TVM functions. /// /// # Arguments /// * `$f` - The function identifier to convert /// * `$trait` - The trait type (typically `Fn`) /// * `($t0, $t1, ...)` - The argument types /// * `$ret_ty` - The return type /// /// # Example /// ```rust /// use tvm_ffi::*; /// /// let func = Function::from_typed(|x: i32, y: i32| -> Result { Ok(x + y) }); /// let typed_func = into_typed_fn!(func, Fn(i32, &i32) -> Result); /// let result = typed_func(10, &20).unwrap(); // Returns 30 /// assert_eq!(result, 30); /// ``` /// Note that the `into_typed_fn!` macro can specify arguments to be passed either /// by reference or by value in the argument list. /// We recommend passing by reference for ObjectRef types such as Tensor. /// Since the ffi mechanism requires us to pass arguments by reference. /// /// # Supported Argument Counts /// This macro supports functions with 0 to 8 arguments. ///----------------------------------------------------------- #[macro_export] macro_rules! into_typed_fn { // Case for 0 arguments ($f:expr, $trait:ident() -> $ret_ty:ty) => {{ let _f = $f; move || -> $ret_ty { Ok(_f.call_tuple_with_len::<0, _>(())?.try_into()?) } }}; // Case for 1 argument ($f:expr, $trait:ident($t0:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0,).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<1, _>(tuple_args)?.try_into()?) } }}; // Case for 2 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<2, _>(tuple_args)?.try_into()?) } }}; // Case for 3 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<3, _>(tuple_args)?.try_into()?) } }}; // Case for 4 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2, a3: $t3| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2, a3).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<4, _>(tuple_args)?.try_into()?) } }}; // Case for 5 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2, a3, a4).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<5, _>(tuple_args)?.try_into()?) } }}; // Case for 6 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2, a3, a4, a5).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<6, _>(tuple_args)?.try_into()?) } }}; // Case for 7 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2, a3, a4, a5, a6).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<7, _>(tuple_args)?.try_into()?) } }}; // Case for 8 arguments ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty, $t7:ty) -> $ret_ty:ty) => {{ let _f = $f; move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6, a7: $t7| -> $ret_ty { use $crate::function_internal::IntoArgHolderTuple; let tuple_args = (a0, a1, a2, a3, a4, a5, a6, a7).into_arg_holder_tuple(); Ok(_f.call_tuple_with_len::<8, _>(tuple_args)?.try_into()?) } }}; } tvm-ffi-0.1.12/rust/tvm-ffi/src/object.rs000066400000000000000000000402711521067262500201250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::ops::{Deref, DerefMut}; use std::sync::atomic::AtomicU64; use crate::derive::ObjectRef; pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; /// Object related ABI handling use tvm_ffi_sys::{TVMFFIObject, COMBINED_REF_COUNT_BOTH_ONE}; /// Object type is by default the TVMFFIObject #[repr(C)] pub struct Object { /// example implementation of the object header: TVMFFIObject, } /// Arc-like wrapper for Object that allows shared ownership /// /// \tparam T The type of the object to be wrapped #[repr(C)] pub struct ObjectArc { ptr: std::ptr::NonNull, _phantom: std::marker::PhantomData, } unsafe impl Send for ObjectArc {} unsafe impl Sync for ObjectArc {} /// Traits that can be used to check if a type is an object /// /// This trait is unsafe because it is used to access the object header /// and the object header is unsafe to access pub unsafe trait ObjectCore: Sized + 'static { /// the type key of the object const TYPE_KEY: &'static str; // return the type index of the object fn type_index() -> i32; /// Return the object header /// This function is implemented as a static function so /// /// # Arguments /// * `this` - The object to get the header /// /// # Returns /// * `&mut TVMFFIObject` - The object header /// \return The object header unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject; } /// Traits for objects with extra items that follows the object /// /// This extra trait can be helpful to implement array types and string types pub unsafe trait ObjectCoreWithExtraItems: ObjectCore { /// type of extra items storage that follows the object type ExtraItem; /// Return the number of extra items fn extra_items_count(this: &Self) -> usize; /// Return the extra items data pointer unsafe fn extra_items(this: &Self) -> &[Self::ExtraItem] { let extra_items_ptr = (this as *const Self as *const u8).add(std::mem::size_of::()); std::slice::from_raw_parts( extra_items_ptr as *const Self::ExtraItem, Self::extra_items_count(this), ) } /// Return the extra items data pointer unsafe fn extra_items_mut(this: &mut Self) -> &mut [Self::ExtraItem] { let extra_items_ptr = (this as *mut Self as *mut u8).add(std::mem::size_of::()); std::slice::from_raw_parts_mut( extra_items_ptr as *mut Self::ExtraItem, Self::extra_items_count(this), ) } } /// Traits to specify core operations of ObjectRef /// /// used by the ffi Any system and not user facing /// /// We mark as unsafe since it moves out the internal of the ObjectRef pub unsafe trait ObjectRefCore: Sized + Clone { type ContainerType: ObjectCore; fn data(this: &Self) -> &ObjectArc; fn into_data(this: Self) -> ObjectArc; fn from_data(data: ObjectArc) -> Self; } /// Base class for ObjectRef /// /// This class is used to store the data of the ObjectRef #[repr(C)] #[derive(ObjectRef, Clone)] pub struct ObjectRef { data: ObjectArc, } /// Unsafe operations on object pub(crate) mod unsafe_ { use tvm_ffi_sys::{ COMBINED_REF_COUNT_BOTH_ONE, COMBINED_REF_COUNT_MASK_U32, COMBINED_REF_COUNT_STRONG_ONE, COMBINED_REF_COUNT_WEAK_ONE, }; use std::ffi::c_void; use std::sync::atomic::{fence, Ordering}; use tvm_ffi_sys::TVMFFIObject; use tvm_ffi_sys::TVMFFIObjectDeleterFlagBitMask::{ kTVMFFIObjectDeleterFlagBitMaskBoth, kTVMFFIObjectDeleterFlagBitMaskStrong, kTVMFFIObjectDeleterFlagBitMaskWeak, }; /// Increase the strong reference count of the object /// /// This function is same as TVMFFIObjectIncRef but implemented natively in Rust /// /// # Arguments /// * `obj` - The object to increase the reference count #[inline] pub unsafe fn inc_ref(handle: *mut TVMFFIObject) { let obj = &mut *handle; obj.combined_ref_count.fetch_add(1, Ordering::Relaxed); } /// Decrease the strong reference count of the object /// /// This function is same as TVMFFIObjectDecRef but implemented natively in Rust /// /// # Arguments /// * `obj` - The object to decrease the reference count #[inline] pub unsafe fn dec_ref(handle: *mut TVMFFIObject) { let obj = &mut *handle; let old_combined_count = obj .combined_ref_count .fetch_sub(COMBINED_REF_COUNT_STRONG_ONE, Ordering::Relaxed); if old_combined_count == COMBINED_REF_COUNT_BOTH_ONE { if let Some(deleter) = obj.deleter { fence(Ordering::Acquire); deleter( obj as *mut TVMFFIObject as *mut c_void, kTVMFFIObjectDeleterFlagBitMaskBoth as i32, ); } } else if (old_combined_count & COMBINED_REF_COUNT_MASK_U32) == COMBINED_REF_COUNT_STRONG_ONE { // slow path, there is still a weak reference left // need to run two phase decrement fence(Ordering::Acquire); if let Some(deleter) = obj.deleter { deleter( obj as *mut TVMFFIObject as *mut c_void, kTVMFFIObjectDeleterFlagBitMaskStrong as i32, ); } let old_weak_count = obj .combined_ref_count .fetch_sub(COMBINED_REF_COUNT_WEAK_ONE, Ordering::Release); if old_weak_count == COMBINED_REF_COUNT_WEAK_ONE { fence(Ordering::Acquire); if let Some(deleter) = obj.deleter { deleter( obj as *mut TVMFFIObject as *mut c_void, kTVMFFIObjectDeleterFlagBitMaskWeak as i32, ); } } } } #[inline] pub unsafe fn strong_count(handle: *mut TVMFFIObject) -> usize { let obj = &mut *handle; (obj.combined_ref_count.load(Ordering::Relaxed) & COMBINED_REF_COUNT_MASK_U32) as usize } #[inline] pub unsafe fn weak_count(handle: *mut TVMFFIObject) -> usize { let obj = &mut *handle; (obj.combined_ref_count.load(Ordering::Relaxed) >> 32) as usize } /// Generic object deleter that works for object allocated from Box then into_raw pub unsafe extern "C" fn object_deleter_for_new(ptr: *mut c_void, flags: i32) where T: super::ObjectCore, { let obj = ptr as *mut T; if flags & kTVMFFIObjectDeleterFlagBitMaskStrong as i32 != 0 { // calling destructor of the object, does not free the memory std::ptr::drop_in_place(obj); } if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 { // free the memory std::alloc::dealloc(ptr as *mut u8, std::alloc::Layout::new::()); } } pub unsafe extern "C" fn object_deleter_for_new_with_extra_items( ptr: *mut c_void, flags: i32, ) where T: super::ObjectCoreWithExtraItems, { let obj = ptr as *mut T; if flags == kTVMFFIObjectDeleterFlagBitMaskBoth as i32 { // must get extra items count before dropping the object let extra_items_count = T::extra_items_count(&(*obj)); std::ptr::drop_in_place(obj); let layout = std::alloc::Layout::from_size_align( std::mem::size_of::() + extra_items_count * std::mem::size_of::(), std::mem::align_of::(), ) .unwrap(); // free the memory std::alloc::dealloc(ptr as *mut u8, layout); } else { assert_eq!(std::mem::size_of::() % std::mem::size_of::(), 0); if flags & kTVMFFIObjectDeleterFlagBitMaskStrong as i32 != 0 { // must get extra items count before dropping the object let extra_items_count = T::extra_items_count(&(*obj)); // calling destructor of the object, does not free the memory std::ptr::drop_in_place(obj); // record extra count in the original memory std::ptr::write(obj as *mut u64, extra_items_count as u64); } if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 { // read extra items count from the original memory // note we can no longer read it by calling T::extra_items_count(&(*obj)) // because the object is already dropped let extra_items_count = std::ptr::read(obj as *mut u64) as usize; let layout = std::alloc::Layout::from_size_align( std::mem::size_of::() + extra_items_count * std::mem::size_of::(), std::mem::align_of::(), ) .unwrap(); // free the memory std::alloc::dealloc(ptr as *mut u8, layout); } } } } //--------------------- // Object //--------------------- impl Object { pub fn new() -> Self { Self { header: TVMFFIObject::new(), } } } unsafe impl ObjectCore for Object { const TYPE_KEY: &'static str = "ffi.Object"; #[inline] fn type_index() -> i32 { TypeIndex::kTVMFFIStaticObjectBegin as i32 } #[inline] unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject { &mut this.header } } //--------------------- // ObjectArc //--------------------- impl ObjectArc { pub fn new(data: T) -> Self { unsafe { let layout = std::alloc::Layout::new::(); let raw_data_ptr = std::alloc::alloc(layout); if raw_data_ptr.is_null() { std::alloc::handle_alloc_error(layout); } let ptr = raw_data_ptr as *mut T; std::ptr::write(ptr, data); // now override the header directly std::ptr::write( ptr as *mut TVMFFIObject, TVMFFIObject { combined_ref_count: AtomicU64::new(COMBINED_REF_COUNT_BOTH_ONE), type_index: T::type_index(), __padding: 0, deleter: Some(unsafe_::object_deleter_for_new::), }, ); // move into the object arc ptr Self { ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T), _phantom: std::marker::PhantomData, } } } pub fn new_with_extra_items(data: T) -> Self where T: ObjectCoreWithExtraItems, { unsafe { // ensure strict alignment requirements // so we can have { T, U*extra_items } layout assert_eq!(std::mem::align_of::() % std::mem::align_of::(), 0); assert_eq!(std::mem::size_of::() % std::mem::align_of::(), 0); let extra_items_count = T::extra_items_count(&data); let layout = std::alloc::Layout::from_size_align( std::mem::size_of::() + extra_items_count * std::mem::size_of::(), std::mem::align_of::(), ) .unwrap(); let raw_data_ptr = std::alloc::alloc(layout); if raw_data_ptr.is_null() { std::alloc::handle_alloc_error(layout); } let ptr = raw_data_ptr as *mut T; std::ptr::write(ptr, data); // now override the header directly std::ptr::write( ptr as *mut TVMFFIObject, TVMFFIObject { combined_ref_count: AtomicU64::new(COMBINED_REF_COUNT_BOTH_ONE), type_index: T::type_index(), __padding: 0, deleter: Some(unsafe_::object_deleter_for_new_with_extra_items::), }, ); // move into the object arc ptr Self { ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T), _phantom: std::marker::PhantomData, } } } /// Move a previously allocated object into the ObjectArc /// /// # Arguments /// * `ptr` - The raw pointer to move into the ObjectArc /// /// # Returns /// * `ObjectArc` - The ObjectArc /// \return The ObjectArc #[inline] pub unsafe fn from_raw(ptr: *const T) -> Self { Self { ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T), _phantom: std::marker::PhantomData, } } /// Move the ObjectArc into a raw pointer /// /// # Arguments /// * `this` - The ObjectArc to move into a raw pointer /// /// # Returns /// * `*const T` - The raw pointer #[inline] pub unsafe fn into_raw(this: Self) -> *const T { let droped_this = std::mem::ManuallyDrop::new(this); droped_this.ptr.as_ptr() as *const T } /// Get the raw pointer from the ObjectArc /// /// Caller should view this as a non-owning reference /// /// # Arguments /// * `this` - The ObjectArc to get the raw pointer /// /// # Returns /// * `*const T` - The raw pointer /// \return The raw pointer #[inline] pub unsafe fn as_raw(this: &Self) -> *const T { this.ptr.as_ptr() as *const T } /// Get the raw mutable pointer from the ObjectArc /// /// Caller should view this as a non-owning reference /// /// # Arguments /// * `this` - The ObjectArc to get the raw pointer /// /// # Returns /// * `*mut T` - The raw pointer #[inline] pub unsafe fn as_raw_mut(this: &mut Self) -> *mut T { this.ptr.as_mut() } /// Get the strong reference count of the ObjectArc /// /// # Arguments /// * `this` - The ObjectArc to get the strong reference count /// /// # Returns /// * `usize` - The strong reference count #[inline] pub fn strong_count(this: &Self) -> usize { unsafe { unsafe_::strong_count(this.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject) } } /// Get the weak reference count of the ObjectArc /// /// # Arguments /// * `this` - The ObjectArc to get the weak reference count /// /// # Returns /// * `usize` - The weak reference count #[inline] pub fn weak_count(this: &Self) -> usize { unsafe { unsafe_::weak_count(this.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject) } } } // implement Deref for ObjectArc impl Deref for ObjectArc { type Target = T; #[inline] fn deref(&self) -> &Self::Target { unsafe { self.ptr.as_ref() } } } // implement DerefMut for ObjectArc impl DerefMut for ObjectArc { #[inline] fn deref_mut(&mut self) -> &mut Self::Target { unsafe { self.ptr.as_mut() } } } // implement Drop for ObjectArc impl Drop for ObjectArc { fn drop(&mut self) { unsafe { unsafe_::dec_ref(self.ptr.as_mut() as *mut T as *mut TVMFFIObject) } } } // implement Clone for ObjectArc impl Clone for ObjectArc { #[inline] fn clone(&self) -> Self { unsafe { unsafe_::inc_ref(self.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject) } Self { ptr: self.ptr, _phantom: std::marker::PhantomData, } } } tvm-ffi-0.1.12/rust/tvm-ffi/src/string.rs000066400000000000000000000370371521067262500201730ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use crate::derive::Object; use crate::object::{unsafe_, Object, ObjectArc, ObjectCoreWithExtraItems}; use crate::type_traits::AnyCompatible; use std::cmp::Ordering; use std::fmt::{Debug, Display}; use std::hash::{Hash, Hasher}; use std::ops::Deref; use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; use tvm_ffi_sys::{TVMFFIAny, TVMFFIAnyDataUnion, TVMFFIByteArray, TVMFFIObject}; //----------------------------------------------------- // Bytes //----------------------------------------------------- /// ABI stable Bytes container for ffi #[repr(C)] pub struct Bytes { data: TVMFFIAny, } // BytesObj for heap-allocated bytes #[repr(C)] #[derive(Object)] #[type_key = "ffi.Bytes"] #[type_index(TypeIndex::kTVMFFIBytes)] pub(crate) struct BytesObj { object: Object, data: TVMFFIByteArray, } impl Bytes { /// Create a new empty Bytes container pub fn new() -> Self { Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFISmallBytes as i32, small_str_len: 0, data_union: TVMFFIAnyDataUnion { v_int64: 0 }, }, } } /// Get the length of the bytes pub fn len(&self) -> usize { self.as_slice().len() } /// Get the bytes as a slice pub fn as_slice(&self) -> &[u8] { unsafe { if self.data.type_index == TypeIndex::kTVMFFISmallBytes as i32 { std::slice::from_raw_parts( self.data.data_union.v_bytes.as_ptr(), self.data.small_str_len as usize, ) } else { let str_obj: &BytesObj = &*(self.data.data_union.v_obj as *const BytesObj); std::slice::from_raw_parts(str_obj.data.data, str_obj.data.size) } } } } impl Default for Bytes { #[inline] fn default() -> Self { Self::new() } } unsafe impl ObjectCoreWithExtraItems for BytesObj { type ExtraItem = u8; // extra item is the trailing \0 for ffi compatibility #[inline] /// Get the count of extra items (trailing null byte for FFI compatibility) fn extra_items_count(this: &Self) -> usize { return this.data.size + 1; } } impl From for Bytes where T: AsRef<[u8]>, { #[inline] /// Create Bytes from any type that can be converted to a byte slice fn from(src: T) -> Self { let value: &[u8] = src.as_ref(); // to be compatible with normal c++ const MAX_SMALL_BYTES_LEN: usize = 7; unsafe { if value.len() <= MAX_SMALL_BYTES_LEN { let mut data_union = TVMFFIAnyDataUnion { v_int64: 0 }; data_union.v_bytes[..value.len()].copy_from_slice(value); // small bytes Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFISmallBytes as i32, small_str_len: value.len() as u32, data_union: data_union, }, } } else { // large bytes let mut obj_arc = ObjectArc::new_with_extra_items(BytesObj { object: Object::new(), data: TVMFFIByteArray { data: std::ptr::null(), size: value.len(), }, }); // reset the data ptr correctly after Arc is created obj_arc.data.data = BytesObj::extra_items(&obj_arc).as_ptr(); let extra_items = BytesObj::extra_items_mut(&mut obj_arc); extra_items[..value.len()].copy_from_slice(value); // write the trailing \0 for ffi compatibility extra_items[value.len()] = 0; Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFIBytes as i32, small_str_len: 0, data_union: TVMFFIAnyDataUnion { v_obj: ObjectArc::into_raw(obj_arc) as *mut BytesObj as *mut TVMFFIObject, }, }, } } } } } impl Deref for Bytes { type Target = [u8]; #[inline] fn deref(&self) -> &[u8] { self.as_slice() } } impl Clone for Bytes { #[inline] fn clone(&self) -> Self { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::inc_ref(self.data.data_union.v_obj) } } Self { data: self.data } } } impl Drop for Bytes { #[inline] fn drop(&mut self) { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::dec_ref(self.data.data_union.v_obj) } } } } impl PartialEq for Bytes { #[inline] fn eq(&self, other: &Self) -> bool { self.as_slice() == other.as_slice() } } impl Eq for Bytes {} impl PartialOrd for Bytes { #[inline] fn partial_cmp(&self, other: &Self) -> Option { self.as_slice().partial_cmp(other.as_slice()) } } impl Ord for Bytes { #[inline] fn cmp(&self, other: &Self) -> Ordering { self.as_slice().cmp(other.as_slice()) } } impl Hash for Bytes { #[inline] fn hash(&self, state: &mut H) { self.as_slice().hash(state); } } impl Display for Bytes { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { Debug::fmt(&self.as_slice(), f) } } impl Debug for Bytes { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ffi.Bytes") .field("data", &self.as_slice()) .finish() } } //----------------------------------------------------- // String //----------------------------------------------------- /// ABI stable String container for ffi #[repr(C)] pub struct String { data: TVMFFIAny, } // StringObj for heap-allocated string #[repr(C)] #[derive(Object)] #[type_key = "ffi.String"] #[type_index(TypeIndex::kTVMFFIStr)] pub(crate) struct StringObj { object: Object, data: TVMFFIByteArray, } unsafe impl ObjectCoreWithExtraItems for StringObj { type ExtraItem = u8; #[inline] /// Get the count of extra items (trailing null byte for FFI compatibility) fn extra_items_count(this: &Self) -> usize { // extra item is the trailing \0 for ffi compatibility return this.data.size + 1; } } impl String { /// Create a new empty String container pub fn new() -> Self { Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFISmallStr as i32, small_str_len: 0, data_union: TVMFFIAnyDataUnion { v_int64: 0 }, }, } } /// Get the length of the string in bytes pub fn len(&self) -> usize { self.as_bytes().len() } /// Get the string as a byte slice pub fn as_bytes(&self) -> &[u8] { unsafe { if self.data.type_index == TypeIndex::kTVMFFISmallStr as i32 { std::slice::from_raw_parts( self.data.data_union.v_bytes.as_ptr(), self.data.small_str_len as usize, ) } else { let str_obj: &StringObj = &*(self.data.data_union.v_obj as *const StringObj); std::slice::from_raw_parts(str_obj.data.data, str_obj.data.size) } } } /// Get the string as a str slice pub fn as_str(&self) -> &str { unsafe { std::str::from_utf8_unchecked(self.as_bytes()) } } } impl From for String where T: AsRef, { #[inline] /// Create String from any type that can be converted to a string slice fn from(src: T) -> Self { unsafe { let value: &str = src.as_ref(); let bytes = value.as_bytes(); const MAX_SMALL_BYTES_LEN: usize = 7; if bytes.len() <= MAX_SMALL_BYTES_LEN { let mut data_union = TVMFFIAnyDataUnion { v_int64: 0 }; data_union.v_bytes[..bytes.len()].copy_from_slice(bytes); Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFISmallStr as i32, small_str_len: bytes.len() as u32, data_union: data_union, }, } } else { let mut obj_arc = ObjectArc::new_with_extra_items(StringObj { object: Object::new(), data: TVMFFIByteArray { data: std::ptr::null(), size: bytes.len(), }, }); obj_arc.data.data = StringObj::extra_items(&obj_arc).as_ptr(); let extra_items = StringObj::extra_items_mut(&mut obj_arc); extra_items[..bytes.len()].copy_from_slice(bytes); // write the trailing \0 for ffi compatibility extra_items[bytes.len()] = 0; Self { data: TVMFFIAny { type_index: TypeIndex::kTVMFFIStr as i32, small_str_len: 0, data_union: TVMFFIAnyDataUnion { v_obj: ObjectArc::into_raw(obj_arc) as *mut StringObj as *mut TVMFFIObject, }, }, } } } } } impl Default for String { #[inline] fn default() -> Self { Self::new() } } impl Deref for String { type Target = str; #[inline] fn deref(&self) -> &str { self.as_str() } } impl Clone for String { #[inline] fn clone(&self) -> Self { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::inc_ref(self.data.data_union.v_obj) } } Self { data: self.data } } } impl Drop for String { #[inline] fn drop(&mut self) { if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::dec_ref(self.data.data_union.v_obj) } } } } impl PartialEq for String { #[inline] fn eq(&self, other: &Self) -> bool { self.as_bytes() == other.as_bytes() } } // Allows `my_string == "hello"` impl PartialEq for String where T: AsRef, { #[inline] fn eq(&self, other: &T) -> bool { self.as_str() == other.as_ref() } } impl Eq for String {} impl PartialEq for Bytes where T: AsRef<[u8]>, { #[inline] fn eq(&self, other: &T) -> bool { self.as_slice() == other.as_ref() } } impl PartialOrd for String { #[inline] fn partial_cmp(&self, other: &Self) -> Option { self.as_str().partial_cmp(other.as_str()) } } impl Ord for String { #[inline] fn cmp(&self, other: &Self) -> Ordering { self.as_str().cmp(other.as_str()) } } impl Hash for String { #[inline] fn hash(&self, state: &mut H) { self.as_str().hash(state); } } impl Display for String { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { Debug::fmt(&self.as_str(), f) } } impl Debug for String { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ffi.String") .field("data", &self.as_str()) .finish() } } //----------------------------------------------------- // AnyCompatible implementation for Bytes and String //----------------------------------------------------- unsafe impl AnyCompatible for Bytes { fn type_str() -> std::string::String { "ffi.Bytes".to_string() } unsafe fn copy_to_any_view(this: &Self, data: &mut TVMFFIAny) { *data = this.data; } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { *data = src.data; std::mem::forget(src); } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return data.type_index == TypeIndex::kTVMFFISmallBytes as i32 || data.type_index == TypeIndex::kTVMFFIBytes as i32; } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { if data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::inc_ref(data.data_union.v_obj) } } Self { data: *data } } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { Self { data: *data } } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIByteArrayPtr as i32 { // deep copy from bytearray ptr let bytes = &*(data.data_union.v_ptr as *const TVMFFIByteArray); Ok(Self::from(std::slice::from_raw_parts( bytes.data, bytes.size, ))) } else if data.type_index == TypeIndex::kTVMFFISmallBytes as i32 { Ok(Self { data: *data }) } else if data.type_index == TypeIndex::kTVMFFIBytes as i32 { unsafe { unsafe_::inc_ref(data.data_union.v_obj) } Ok(Self { data: *data }) } else { Err(()) } } } unsafe impl AnyCompatible for String { fn type_str() -> std::string::String { "ffi.String".to_string() } unsafe fn copy_to_any_view(this: &Self, data: &mut TVMFFIAny) { *data = this.data; } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { *data = src.data; std::mem::forget(src); } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return data.type_index == TypeIndex::kTVMFFISmallStr as i32 || data.type_index == TypeIndex::kTVMFFIStr as i32; } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { if data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 { unsafe { unsafe_::inc_ref(data.data_union.v_obj) } } Self { data: *data } } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { Self { data: *data } } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIRawStr as i32 { // 1. Create a CStr wrapper from the raw pointer. let c_str = std::ffi::CStr::from_ptr(data.data_union.v_c_str as *const std::os::raw::c_char); Ok(Self::from(c_str.to_str().expect("Invalid UTF-8"))) } else if data.type_index == TypeIndex::kTVMFFISmallStr as i32 { Ok(Self { data: *data }) } else if data.type_index == TypeIndex::kTVMFFIStr as i32 { unsafe { unsafe_::inc_ref(data.data_union.v_obj) } Ok(Self { data: *data }) } else { Err(()) } } } tvm-ffi-0.1.12/rust/tvm-ffi/src/type_traits.rs000066400000000000000000000272111521067262500212250ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; use tvm_ffi_sys::{TVMFFIAny, TVMFFIGetTypeInfo}; //----------------------------------------------------- // AnyCompatible //----------------------------------------------------- /// Trait to enable a value to be compatible with Any /// Enables TryFrom/Into AnyView/Any pub unsafe trait AnyCompatible: Sized { /// the value to copy to TVMFFIAny unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny); /// consume the value to move to Any unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny); // check if the value is compatible with the type unsafe fn check_any_strict(data: &TVMFFIAny) -> bool; /// Copy value from TVMFFIAny after checking /// caller must ensure that the value is compatible with the type unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self; /// the value to move from TVMFFIAny /// NOTE: pay very careful attention to avoid memory leak! /// - When calling from managed Any, remember to use std::mem::ManuallyDrop unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self; /// try to cast the value from AnyView unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result; /// Get the type key of a type when TryCastFromAnyView fails. fn get_mismatch_type_info(data: &TVMFFIAny) -> String { unsafe { let info = TVMFFIGetTypeInfo(data.type_index); assert!(!info.is_null(), "TVMFFIGetTypeInfo returned null"); (*info).type_key.as_str().to_string() } } /// the type string of the type fn type_str() -> String; } /// AnyCompatible for bool unsafe impl AnyCompatible for bool { unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIBool as i32; data.small_str_len = 0; data.data_union.v_int64 = *src as i64; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { data.type_index == TypeIndex::kTVMFFIBool as i32 } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_int64 != 0 } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIBool as i32; data.small_str_len = 0; data.data_union.v_int64 = src as i64; } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_int64 != 0 } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIBool as i32 || data.type_index == TypeIndex::kTVMFFIInt as i32 { unsafe { Ok(data.data_union.v_int64 != 0) } } else { Err(()) } } fn type_str() -> String { "bool".to_string() } } /// Macro to implement AnyCompatible for integer types macro_rules! impl_any_compatible_for_int { ($($int_type:ty),* $(,)?) => { $( unsafe impl AnyCompatible for $int_type { fn type_str() -> String { "int".to_string() } unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIInt as i32; data.small_str_len = 0; data.data_union.v_int64 = *src as i64; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { data.type_index == TypeIndex::kTVMFFIInt as i32 } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_int64 as $int_type } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIInt as i32; data.small_str_len = 0; data.data_union.v_int64 = src as i64; } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_int64 as $int_type } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIInt as i32 || data.type_index == TypeIndex::kTVMFFIBool as i32 { Ok(data.data_union.v_int64 as $int_type) } else { Err(()) } } } )* }; } // Implement AnyCompatible for all integer types impl_any_compatible_for_int!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize); // Implement AnyCompatible for `Option` unsafe impl AnyCompatible for Option { fn type_str() -> String { // make it consistent with c++ representation "Optional<".to_string() + T::type_str().as_str() + ">" } unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { if let Some(ref value) = src { T::copy_to_any_view(value, data); } else { data.type_index = TypeIndex::kTVMFFINone as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; } } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { if let Some(value) = src { T::move_to_any(value, data); } else { data.type_index = TypeIndex::kTVMFFINone as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; } } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return T::check_any_strict(data) || data.type_index == TypeIndex::kTVMFFINone as i32; } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { if data.type_index == TypeIndex::kTVMFFINone as i32 { None } else { Some(T::copy_from_any_view_after_check(data)) } } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { if data.type_index == TypeIndex::kTVMFFINone as i32 { None } else { Some(T::move_from_any_after_check(data)) } } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFINone as i32 { Ok(None) } else { T::try_cast_from_any_view(data).map(Some) } } } /// AnyCompatible for void* unsafe impl AnyCompatible for *mut core::ffi::c_void { unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIOpaquePtr as i32; data.small_str_len = 0; data.data_union.v_ptr = *src; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { data.type_index == TypeIndex::kTVMFFIOpaquePtr as i32 } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_ptr } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIOpaquePtr as i32; data.small_str_len = 0; data.data_union.v_int64 = src as i64; } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_ptr } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFINone as i32 { Ok(std::ptr::null_mut()) } else if data.type_index == TypeIndex::kTVMFFIOpaquePtr as i32 { unsafe { Ok(data.data_union.v_ptr) } } else { Err(()) } } fn type_str() -> String { "void*".to_string() } } /// Macro to implement AnyCompatible for integer types macro_rules! impl_any_compatible_for_float { ($($float_type:ty),* $(,)?) => { $( unsafe impl AnyCompatible for $float_type { fn type_str() -> String { "float".to_string() } unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIFloat as i32; data.small_str_len = 0; data.data_union.v_float64 = *src as f64; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { data.type_index == TypeIndex::kTVMFFIFloat as i32 } unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { data.data_union.v_float64 as $float_type } unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFIFloat as i32; data.small_str_len = 0; data.data_union.v_float64 = src as f64; } unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { data.data_union.v_float64 as $float_type } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFIFloat as i32 { Ok(data.data_union.v_float64 as $float_type) } else if data.type_index == TypeIndex::kTVMFFIInt as i32 || data.type_index == TypeIndex::kTVMFFIBool as i32 { Ok(data.data_union.v_int64 as $float_type) } else { Err(()) } } } )* }; } // Implement AnyCompatible for all float types impl_any_compatible_for_float!(f32, f64); // Special rule: we convert () to None // This allows us to effectively pass () around as null value // note that this is indeed a bit relaxation of the type // but it is necessary for us to enable void/none interoperability unsafe impl AnyCompatible for () { fn type_str() -> String { // make it consistent with c++ representation "None".to_string() } unsafe fn copy_to_any_view(_src: &Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFINone as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; } unsafe fn move_to_any(_src: Self, data: &mut TVMFFIAny) { data.type_index = TypeIndex::kTVMFFINone as i32; data.small_str_len = 0; data.data_union.v_int64 = 0; } unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { return data.type_index == TypeIndex::kTVMFFINone as i32; } unsafe fn copy_from_any_view_after_check(_data: &TVMFFIAny) -> Self { () } unsafe fn move_from_any_after_check(_data: &mut TVMFFIAny) -> Self { () } unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result { if data.type_index == TypeIndex::kTVMFFINone as i32 { Ok(()) } else { Err(()) } } } tvm-ffi-0.1.12/rust/tvm-ffi/tests/000077500000000000000000000000001521067262500166605ustar00rootroot00000000000000tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_any.rs000066400000000000000000000312111521067262500210520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; /// Macro to test integer types with both Any and AnyView macro_rules! test_any_int { ($($type:ty: $value:expr, $any_name:ident, $any_view_name:ident),* $(,)?) => { $( #[test] fn $any_name() { let any = Any::from($value); assert_eq!(any.type_index(), TypeIndex::kTVMFFIInt as i32); assert_eq!(<$type>::try_from(any).unwrap(), $value); } #[test] fn $any_view_name() { let value = $value; let any_view = AnyView::from(&value); assert_eq!(any_view.type_index(), TypeIndex::kTVMFFIInt as i32); assert_eq!(<$type>::try_from(any_view).unwrap(), $value); } )* }; } // Test all integer types with both Any and AnyView test_any_int!( i8: 127i8, test_any_int_i8, test_any_view_int_i8, i16: 32767i16, test_any_int_i16, test_any_view_int_i16, i32: 42i32, test_any_int_i32, test_any_view_int_i32, i64: 123456789i64, test_any_int_i64, test_any_view_int_i64, isize: 42isize, test_any_int_isize, test_any_view_int_isize, u8: 255u8, test_any_int_u8, test_any_view_int_u8, u16: 65535u16, test_any_int_u16, test_any_view_int_u16, u32: 4294967295u32, test_any_int_u32, test_any_view_int_u32, u64: 18446744073709551615u64, test_any_int_u64, test_any_view_int_u64, usize: 42usize, test_any_int_usize, test_any_view_int_usize ); /// Macro to test float types with both Any and AnyView macro_rules! test_any_float { ($($type:ty: $value:expr, $any_name:ident, $any_view_name:ident),* $(,)?) => { $( #[test] fn $any_name() { let any = Any::from($value); assert_eq!(any.type_index(), TypeIndex::kTVMFFIFloat as i32); assert_eq!(<$type>::try_from(any).unwrap(), $value); } )* }; } // Test all float types with both Any and AnyView test_any_float!( f32: 3.14f32, test_any_float_f32, test_any_view_float_f32, f64: 3.14f64, test_any_float_f64, test_any_view_float_f64 ); #[test] fn test_any_bool() { // Test both Any and AnyView with both true and false let any_true = Any::from(true); let any_false = Any::from(false); let value_true = true; let value_false = false; let any_view_true = AnyView::from(&value_true); let any_view_false = AnyView::from(&value_false); // Test Any assert_eq!(any_true.type_index(), TypeIndex::kTVMFFIBool as i32); assert_eq!(any_false.type_index(), TypeIndex::kTVMFFIBool as i32); assert_eq!(bool::try_from(any_true).unwrap(), true); assert_eq!(bool::try_from(any_false).unwrap(), false); // Test AnyView assert_eq!(any_view_true.type_index(), TypeIndex::kTVMFFIBool as i32); assert_eq!(any_view_false.type_index(), TypeIndex::kTVMFFIBool as i32); assert_eq!(bool::try_from(any_view_true).unwrap(), true); assert_eq!(bool::try_from(any_view_false).unwrap(), false); } #[test] fn test_type_mismatch_error() { let any = Any::from(42i32); // try from will try casting when possible // Try to convert int to bool - should be ok let result: Result = bool::try_from(any); assert!(result.is_ok()); // try_as is more strict and do not allow conversion from int to bool let any = Any::from(42i32); let opt_try_as = any.try_as::(); assert!(opt_try_as.is_none()); } #[test] fn test_type_mismatch_error_message() { let any_none = Any::from(Option::::from(None)); // try from will try casting when possible // Try to convert int to bool - should be ok let err = i32::try_from(any_none).unwrap_err(); assert_eq!(err.kind(), crate::error::TYPE_ERROR); assert!(err.message().contains("`None` to `int`")); let any_float = Any::from(1.2f32); let err = Option::::try_from(any_float).unwrap_err(); assert_eq!(err.kind(), crate::error::TYPE_ERROR); assert!(err.message().contains("`float` to `Optional`")); } #[test] fn test_cross_type_conversion_int_to_bool() { // Test that integers can be converted to bool (0 = false, non-zero = true) let any_true = Any::from(42i32); let any_false = Any::from(0i32); // This should work because our macro allows conversion from int to bool let bool_true: Result = bool::try_from(any_true); let bool_false: Result = bool::try_from(any_false); // Note: This test might fail depending on how try_cast_from_any_view is implemented // If it fails, it means the conversion isn't implemented yet assert_eq!(bool_true.unwrap(), true); assert_eq!(bool_false.unwrap(), false); } #[test] fn test_max_min_values() { // Test maximum and minimum values for i64 let any_i64_max = Any::from(i64::MAX); let any_i64_min = Any::from(i64::MIN); assert_eq!(i64::try_from(any_i64_max).unwrap(), i64::MAX); assert_eq!(i64::try_from(any_i64_min).unwrap(), i64::MIN); } #[test] fn test_any_void_ptr() { let any = Any::from(std::ptr::null_mut()); assert_eq!(any.type_index(), TypeIndex::kTVMFFIOpaquePtr as i32); assert_eq!( <*mut core::ffi::c_void>::try_from(any).unwrap(), std::ptr::null_mut() ); } #[test] fn test_any_option_i32() { // Test Option with Some value let some_value = Some(42i32); let any_some = Any::from(some_value); let any_view_some = AnyView::from(&some_value); // Test Any with Some value assert_eq!(any_some.type_index(), TypeIndex::kTVMFFIInt as i32); assert_eq!(Option::::try_from(any_some).unwrap(), Some(42i32)); // Test AnyView with Some value assert_eq!(any_view_some.type_index(), TypeIndex::kTVMFFIInt as i32); assert_eq!(Option::::try_from(any_view_some).unwrap(), Some(42i32)); // Test Option with None value let none_value: Option = None; let any_none = Any::from(none_value); let any_view_none = AnyView::from(&none_value); // Test Any with None value assert_eq!(any_none.type_index(), TypeIndex::kTVMFFINone as i32); assert_eq!(Option::::try_from(any_none).unwrap(), None::); // Test AnyView with None value assert_eq!(any_view_none.type_index(), TypeIndex::kTVMFFINone as i32); assert_eq!(Option::::try_from(any_view_none).unwrap(), None::); let any_float = Any::from(1.2f32); assert_eq!(any_float.type_index(), TypeIndex::kTVMFFIFloat as i32); assert_eq!( Option::::try_from(any_float).unwrap_err().kind(), TYPE_ERROR ); } #[test] fn test_any_string() { let any = Any::from(String::from("hello")); assert_eq!(any.type_index(), TypeIndex::kTVMFFISmallStr as i32); // Check reference count - small strings are not ref counted assert_eq!(any.debug_strong_count(), None); let out_string = String::try_from(any).unwrap(); assert_eq!(out_string, "hello"); } #[test] fn test_any_bytes() { let data_arr: &[u8; 3] = &[1, 2, 3]; let any = Any::from(Bytes::from(data_arr)); assert_eq!(any.type_index(), TypeIndex::kTVMFFISmallBytes as i32); // Check reference count - small bytes are not ref counted assert_eq!(any.debug_strong_count(), None); let out_bytes = Bytes::try_from(any).unwrap(); assert_eq!(out_bytes, data_arr); } #[test] fn test_any_big_string() { // Use a string longer than 7 characters to trigger big string allocation let long_string = "hello world this is a long string"; let input_str = String::from(long_string); let any = Any::from(input_str.clone()); assert_eq!(any.type_index(), TypeIndex::kTVMFFIStr as i32); // Check reference count - big strings are ref counted assert_eq!(any.debug_strong_count(), Some(2)); let out_string = String::try_from(any).unwrap(); assert_eq!(out_string, long_string); drop(out_string); assert_eq!(AnyView::from(&input_str).debug_strong_count(), Some(1)); } #[test] fn test_any_big_bytes() { // Use bytes longer than 7 bytes to trigger big bytes allocation let data_arr: &[u8; 10] = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; let any = Any::from(Bytes::from(data_arr)); assert_eq!(any.type_index(), TypeIndex::kTVMFFIBytes as i32); // Check reference count - big bytes are ref counted assert_eq!(any.debug_strong_count(), Some(1)); let out_bytes = Bytes::try_from(any).unwrap(); assert_eq!(out_bytes, data_arr); assert_eq!(AnyView::from(&out_bytes).debug_strong_count(), Some(1)); } #[test] fn test_anyview_any_conversion_and_cloning() { let input_str = String::from("hello world this is a long string"); // Test AnyView behavior - should not increment reference count let any_view = AnyView::from(&input_str); assert_eq!(any_view.debug_strong_count(), Some(1)); // Test AnyView cloning - should not affect reference count let any_view_clone = any_view.clone(); assert_eq!(any_view.debug_strong_count(), Some(1)); assert_eq!(any_view_clone.debug_strong_count(), Some(1)); // Test AnyView to Any conversion - should increment reference count let any1 = Any::from(any_view); assert_eq!(any1.debug_strong_count(), Some(2)); // Test Any cloning - should increment reference count let any2 = any1.clone(); assert_eq!(any1.debug_strong_count(), Some(3)); assert_eq!(any2.debug_strong_count(), Some(3)); // Test Any to AnyView conversion - should not change reference count let any_view_from_any = AnyView::from(&any1); assert_eq!(any_view_from_any.debug_strong_count(), Some(3)); // Test AnyView to Any conversion again let any3 = Any::from(any_view_clone); assert_eq!(any3.debug_strong_count(), Some(4)); // Test data integrity through conversions let out_string1 = String::try_from(any1.clone()).unwrap(); let out_string2 = String::try_from(any2.clone()).unwrap(); let out_string3 = String::try_from(any3.clone()).unwrap(); let out_string_view = String::try_from(any_view_from_any).unwrap(); assert_eq!(out_string1, input_str); assert_eq!(out_string2, input_str); assert_eq!(out_string3, input_str); assert_eq!(out_string_view, input_str); // Test reference count cleanup drop(any1); assert_eq!(any3.debug_strong_count(), Some(7)); drop(any2); assert_eq!(any3.debug_strong_count(), Some(6)); drop(any3); assert_eq!(AnyView::from(&input_str).debug_strong_count(), Some(5)); } #[test] fn test_any_dl_device() { use tvm_ffi::{DLDevice, DLDeviceType}; // Test DLDevice with CPU device let cpu_device = DLDevice::new(DLDeviceType::kDLCPU, 0); let any_cpu = Any::from(cpu_device); let any_view_cpu = AnyView::from(&cpu_device); // Test Any with CPU device assert_eq!(any_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32); let converted_cpu = DLDevice::try_from(any_cpu).unwrap(); assert_eq!(converted_cpu.device_type, DLDeviceType::kDLCPU); assert_eq!(converted_cpu.device_id, 0); // Test AnyView with CPU device assert_eq!(any_view_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32); let converted_cpu_view = DLDevice::try_from(any_view_cpu).unwrap(); assert_eq!(converted_cpu_view.device_type, DLDeviceType::kDLCPU); assert_eq!(converted_cpu_view.device_id, 0); // Test DLDevice with CUDA device let cuda_device = DLDevice::new(DLDeviceType::kDLCUDA, 1); let any_cuda = Any::from(cuda_device); let any_view_cuda = AnyView::from(&cuda_device); // Test Any with CUDA device assert_eq!(any_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32); let converted_cuda = DLDevice::try_from(any_cuda).unwrap(); assert_eq!(converted_cuda.device_type, DLDeviceType::kDLCUDA); assert_eq!(converted_cuda.device_id, 1); // Test AnyView with CUDA device assert_eq!(any_view_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32); let converted_cuda_view = DLDevice::try_from(any_view_cuda).unwrap(); assert_eq!(converted_cuda_view.device_type, DLDeviceType::kDLCUDA); assert_eq!(converted_cuda_view.device_id, 1); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_array.rs000066400000000000000000000105161521067262500214060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; /// Helper to create a Tensor with a specific float value and shape fn create_tensor(val: f32, shape: &[i64]) -> Tensor { let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1); let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, shape, dtype, device); if let Ok(slice) = tensor.data_as_slice_mut::() { slice[0] = val; } tensor } /// Helper to extract the first float value from a Tensor fn get_val(tensor: &Tensor) -> f32 { tensor .data_as_slice::() .expect("Type mismatch or null")[0] } #[test] fn test_array_core_and_iteration() { let t1 = create_tensor(10.0, &[1, 2]); let t2 = create_tensor(20.0, &[3, 4, 5]); let array = Array::new(vec![t1.clone(), t2.clone()]); // Core Accessors assert_eq!(array.len(), 2); assert!(!array.is_empty()); // Value Integrity assert_eq!(get_val(&Tensor::try_from(array[0]).unwrap()), 10.0); assert_eq!(Tensor::try_from(array[0]).unwrap().ndim(), 2); assert_eq!(Tensor::try_from(array[1]).unwrap().ndim(), 3); // Iteration let vals: Vec = array.iter().map(|t| get_val(&t)).collect(); assert_eq!(vals, vec![10.0, 20.0]); } #[test] fn test_array_any_conversions() { let array = Array::new(vec![ create_tensor(1.0, &[1]), create_tensor(2.0, &[1]), create_tensor(3.0, &[1]), ]); // Test Any/AnyView Roundtrip (Verifies AnyCompatible and Trait Bounds) let any = Any::from(array); assert_eq!(any.type_index(), TypeIndex::kTVMFFIArray as i32); let back: Array = Array::try_from(any).expect("Any -> Array failed"); assert_eq!(back.len(), 3); assert_eq!(get_val(&back.get(2).unwrap()), 3.0); let view = AnyView::from(&back); let back_from_view: Array = Array::try_from(view).expect("AnyView -> Array failed"); assert_eq!(back_from_view.len(), 3); } #[test] fn test_array_recursive_type_checking() { // 1. Create an Array of Shapes let shape_array = Array::new(vec![Shape::from(vec![1, 2]), Shape::from(vec![3])]); // 2. Wrap it in Any let any_val = Any::from(shape_array); // 3. Try to convert Any (containing Shapes) into Array // This should FAIL because T::check_any_strict (Tensor) will fail on Shape elements let tensor_cast = Array::::try_from(any_val.clone()); assert!( tensor_cast.is_err(), "Should not be able to cast Array to Array" ); // 4. Verify valid cast works let shape_cast = Array::::try_from(any_val); assert!( shape_cast.is_ok(), "Should be able to cast back to correct type" ); } #[test] fn test_array_parametric_heterogeneity() { // Verify Array works with different ObjectRefCore types let shape_array = Array::new(vec![Shape::from(vec![1, 2, 3]), Shape::from(vec![10])]); assert_eq!(shape_array.get(0).unwrap().as_slice(), &[1, 2, 3]); assert_eq!(shape_array.get(1).unwrap().as_slice(), &[10]); let function_array = Array::new(vec![ Function::get_global("ffi.String").unwrap(), Function::get_global("ffi.Bytes").unwrap(), ]); assert_eq!( into_typed_fn!( function_array.get(0).unwrap(), Fn(String) -> Result )("hello".into()) .unwrap(), "hello" ); assert_eq!( into_typed_fn!( function_array.get(1).unwrap(), Fn(Bytes) -> Result )([1, 2, 3].into()) .unwrap(), &[1, 2, 3] ); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_device.rs000066400000000000000000000032451521067262500215300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; #[test] fn test_device_stream() { let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let dummy_stream = 2 as TVMFFIStreamHandle; // SAFETY: dummy_stream is a test value; CPU device accepts any stream handle. unsafe { with_stream(&device, dummy_stream, || { assert_eq!(current_stream(&device), dummy_stream); Ok(()) }) } .unwrap(); assert_eq!(current_stream(&device), std::ptr::null_mut()); } #[test] fn test_device_any_conversion() { let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let any_device: Any = device.into(); assert_eq!(any_device.type_index(), TypeIndex::kTVMFFIDevice as i32); let converted_device: DLDevice = any_device.try_into().unwrap(); assert_eq!(converted_device.device_type, device.device_type); assert_eq!(converted_device.device_id, device.device_id); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_dtype.rs000066400000000000000000000060111521067262500214100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; #[test] fn test_dl_datatype_from_string() { // Test valid dtype strings let test_cases = vec![ ( "int32", DLDataType { code: DLDataTypeCode::kDLInt as u8, bits: 32, lanes: 1, }, ), ( "float32", DLDataType { code: DLDataTypeCode::kDLFloat as u8, bits: 32, lanes: 1, }, ), ( "int8", DLDataType { code: DLDataTypeCode::kDLInt as u8, bits: 8, lanes: 1, }, ), ( "float64", DLDataType { code: DLDataTypeCode::kDLFloat as u8, bits: 64, lanes: 1, }, ), ( "uint32", DLDataType { code: DLDataTypeCode::kDLUInt as u8, bits: 32, lanes: 1, }, ), ]; for item in test_cases { let dtype = DLDataType::try_from_str(item.0).unwrap(); assert_eq!(dtype, item.1); } } #[test] fn test_dl_datatype_any_conversion() { let original_dtype = DLDataType { code: DLDataTypeCode::kDLInt as u8, bits: 32, lanes: 1, }; // Test conversion to Any let any_dtype: Any = original_dtype.into(); // Test conversion back from Any let converted_dtype: DLDataType = any_dtype.try_into().unwrap(); assert_eq!(converted_dtype.code, original_dtype.code); assert_eq!(converted_dtype.bits, original_dtype.bits); assert_eq!(converted_dtype.lanes, original_dtype.lanes); } #[test] fn test_dl_datatype_any_with_string_conversion() { // Test that DLDataType can be converted from string via Any let dtype_str = "int32"; let dtype = DLDataType::try_from_str(dtype_str).unwrap(); // Convert to Any let any_dtype: Any = dtype.into(); // Convert back from Any let converted_dtype: DLDataType = any_dtype.try_into().unwrap(); assert_eq!(converted_dtype.code, dtype.code); assert_eq!(converted_dtype.bits, dtype.bits); assert_eq!(converted_dtype.lanes, dtype.lanes); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_error.rs000066400000000000000000000057551521067262500214320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; #[test] fn test_error_new() { let error = Error::new(RUNTIME_ERROR, "test error", "test traceback"); assert_eq!(error.kind(), RUNTIME_ERROR); assert_eq!(error.message(), "test error"); assert_eq!(error.backtrace(), "test traceback"); } #[test] fn test_error_from_raised() { let error0 = Error::new(RUNTIME_ERROR, "test error", "test traceback"); Error::set_raised(&error0); let error1 = Error::from_raised(); assert_eq!(error1.kind(), RUNTIME_ERROR); assert_eq!(error1.message(), "test error"); assert_eq!(error1.backtrace(), "test traceback"); // we have two references because one in error, and another one in the function call assert_eq!(ObjectArc::strong_count(Error::data(&error1)), 2); } fn error_fn0(flag: bool) -> Result { // if flag is true, throw an error ensure!(!flag, RUNTIME_ERROR, "test error {flag}"); Ok(1) } fn error_fn1(flag: bool) -> Result { attach_context!(error_fn0(flag)) } #[test] fn test_error_with_context() { let error0 = attach_context!(error_fn1(true)).unwrap_err(); assert_eq!(error0.kind(), RUNTIME_ERROR); assert_eq!(error0.message(), "test error true"); assert!(error0.backtrace().contains("test_error.rs")); assert!(error0.backtrace().contains("error_fn0")); assert!(error0.backtrace().contains("error_fn1")); assert!(error0.backtrace().contains("test_error_with_context")); let result = error_fn1(false).unwrap(); assert_eq!(result, 1); } #[test] fn test_error_any_convert() { let error = Error::new(RUNTIME_ERROR, "test error", "test traceback"); let any = Any::from(error.clone()); let error2 = any.clone().try_as::().unwrap(); assert_eq!(error2.kind(), RUNTIME_ERROR); assert_eq!(error2.message(), "test error"); assert_eq!(error2.backtrace(), "test traceback"); let res = i32::try_from(any); assert!(res.is_err()); assert!(res.unwrap_err().message().contains("`ffi.Error` to `int`")); let anyview = AnyView::from(&error); let error3 = Error::try_from(anyview).unwrap(); assert_eq!(error3.kind(), RUNTIME_ERROR); assert_eq!(error3.message(), "test error"); assert_eq!(error3.backtrace(), "test traceback"); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_function.rs000066400000000000000000000162351521067262500221210ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; #[test] fn test_function_dummpy_c_api() { let ret = unsafe { tvm_ffi_sys::TVMFFITestingDummyTarget() }; assert_eq!(ret, 0); } #[test] fn test_function_get_global_required() { let fecho = Function::get_global("testing.echo").unwrap(); let a = 1; let args = [AnyView::from(&a)]; let result = fecho.call_packed(&args).unwrap(); assert_eq!(i32::try_from(result).unwrap(), 1); } #[test] fn test_function_from_packed() { let value = 2; let v2 = 4; let check_and_add_value = Function::from_packed(move |args: &[AnyView]| -> Result { ensure!( args.len() == 1, VALUE_ERROR, "Expected 1 argument, got {}", args.len() ); let v0 = i32::try_from(args[0])?; ensure!(v0 == value, VALUE_ERROR, "Expected {}, got {}", value, v0); Ok(Any::from(v0 + v2)) }); let args = [AnyView::from(&value)]; let result = check_and_add_value.call_packed(&args).unwrap(); assert_eq!(i32::try_from(result).unwrap(), 6); } #[test] fn test_function_from_typed() { let offset = 2; // test one argument let sum1 = Function::from_typed(move |x: i32| -> Result { Ok(x + offset) }); let result = sum1.call_packed(&[AnyView::from(&1)]).unwrap(); assert_eq!(i32::try_from(result).unwrap(), 1 + offset); // test two arguments let sum2 = Function::from_typed(move |x: i32, y: i32| -> Result { Ok(x + y) }); let result = sum2 .call_packed(&[AnyView::from(&1), AnyView::from(&2)]) .unwrap(); assert_eq!(i32::try_from(result).unwrap(), 3); // test three arguments let sum3f = Function::from_typed(move |x: i32, y: i32, z: f32| -> Result { Ok((x + y) as f32 + z) }); let result = sum3f .call_packed(&[AnyView::from(&1), AnyView::from(&2), AnyView::from(&3)]) .unwrap(); assert_eq!(f32::try_from(result).unwrap(), 6.0); } #[test] fn test_function_call_tuple() { let offset = 2; // test one argument let sum1 = Function::from_typed(move |x: i32| -> Result { Ok(x + offset) }); let result = sum1.call_tuple((1,)).unwrap(); assert_eq!(i32::try_from(result).unwrap(), 1 + offset); // test pass by reference let result = sum1.call_tuple_with_len::<1, _>((&1,)).unwrap(); assert_eq!(i32::try_from(result).unwrap(), 1 + offset); let typed_fn = |x: &i32| -> Result { Ok(sum1.call_tuple((x,))?.try_into()?) }; let result = typed_fn(&1); assert_eq!(result.unwrap(), 1 + offset); } #[test] fn test_function_into_typed_fn() { let offset = 2; let typed_sum1 = into_typed_fn!( Function::from_typed(move |x: i32| -> Result { Ok(x + offset) }), Fn(&i32) -> Result); assert_eq!(typed_sum1(&1).unwrap(), 1 + offset); // try to box the resulting function let sum2 = Function::from_typed(move |x: i32, y: i32| -> Result { Ok(x + y) }); let typed_sum2 = Box::new(into_typed_fn!(sum2, Fn(&i32, i32) -> Result)); assert_eq!(typed_sum2(&1, 2).unwrap(), 3); // test three arguments let sum3 = Function::from_typed(move |x: i32, y: i32, z: f32| -> Result { Ok((x + y) as f32 + z) }); let typed_sum3 = Box::new(into_typed_fn!(sum3, Fn(&i32, i32, f32) -> Result)); assert_eq!(typed_sum3(&1, 2, 3.0).unwrap(), 6.0); } #[test] fn test_function_echo_tensor_typed() { let echo = into_typed_fn!( Function::get_global("testing.echo").unwrap(), Fn(&Tensor) -> Result ); let data: &[f32] = &[1.0, 2.0, 3.0, 4.0]; let tensor = Tensor::from_slice(data, &[1, 2, 2]).unwrap(); // write tensor content here let result = echo(&tensor).unwrap(); assert_eq!(result.data_ptr(), tensor.data_ptr()); assert_eq!(AnyView::from(&result).debug_strong_count(), Some(2)); let result_data = result.data_as_slice::().unwrap(); assert_eq!(result_data.len(), 4); assert_eq!(result_data[0], 1.0); assert_eq!(result_data[1], 2.0); assert_eq!(result_data[2], 3.0); assert_eq!(result_data[3], 4.0); } fn testing_add_one(x: i32) -> Result { Ok(x + 1) } tvm_ffi_dll_export_typed_func!(testing_add_one, testing_add_one); #[test] fn test_function_from_extern_c() { // SAFETY: null handle is valid for testing_add_one which doesn't use the handle. let add_one = unsafe { Function::from_extern_c(std::ptr::null_mut(), __tvm_ffi_testing_add_one, None) }; let typed_add_one = into_typed_fn!(add_one, Fn(i32) -> Result); assert_eq!(typed_add_one(1).unwrap(), 2); } #[test] fn test_function_echo_string_bytes() { let echo = Function::get_global("testing.echo").unwrap(); let echo_str = into_typed_fn!( echo.clone(), Fn(&str) -> Result ); let result = echo_str("hello").unwrap(); assert_eq!(result, "hello"); let echo_bytes = into_typed_fn!( echo.clone(), Fn(&[u8]) -> Result ); let result = echo_bytes(b"hello").unwrap(); assert_eq!(result, b"hello"); } #[test] fn test_function_apply() { let add_one = Function::from_typed(|x: i32| -> Result { Ok(x + 1) }); let fapply = into_typed_fn!( Function::get_global("testing.apply").unwrap(), Fn(Function, i32) -> Result ); let result = fapply(add_one, 3).unwrap(); assert_eq!(result, 4); } fn test_add_one_tensor(x: tvm_ffi::Tensor, y: tvm_ffi::Tensor) -> Result<()> { let x_data = x.data_as_slice::()?; let y_data = y.data_as_slice_mut::()?; for i in 0..x_data.len() { y_data[i] = x_data[i] + 1.0; } Ok(()) } tvm_ffi_dll_export_typed_func!(test_add_one_tensor, test_add_one_tensor); #[test] fn test_function_call_tensor_fn() { // SAFETY: null handle is valid for test_add_one_tensor which doesn't use the handle. let add_one = unsafe { Function::from_extern_c(std::ptr::null_mut(), __tvm_ffi_test_add_one_tensor, None) }; let typed_add_one = into_typed_fn!(add_one, Fn(&Tensor, &Tensor) -> Result<()>); let x_data: &[f32] = &[0.0, 1.0, 2.0, 3.0]; let x = Tensor::from_slice(x_data, &[2, 2]).unwrap(); let y_data: &[f32] = &[0.0, 0.0, 0.0, 0.0]; let y = Tensor::from_slice(y_data, &[2, 2]).unwrap(); typed_add_one(&x, &y).unwrap(); let y_data = y.data_as_slice::().unwrap(); assert_eq!(y_data[0], 1.0); assert_eq!(y_data[1], 2.0); assert_eq!(y_data[2], 3.0); assert_eq!(y_data[3], 4.0); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_object.rs000066400000000000000000000115261521067262500215400ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; use tvm_ffi::*; // must have repr(C) for the object header stays in the same position #[repr(C)] struct TestIntObj { object: Object, pub value: i64, // counter for recording the number of times the object is deleted delete_counter: Arc, pub extra_item_count: u64, } impl TestIntObj { pub fn new(value: i64, delete_counter: Arc, extra_item_count: u64) -> Self { Self { object: Object::new(), value, delete_counter, extra_item_count, } } } impl Drop for TestIntObj { fn drop(&mut self) { self.delete_counter.fetch_add(1, Ordering::Relaxed); } } unsafe impl ObjectCore for TestIntObj { const TYPE_KEY: &'static str = Object::TYPE_KEY; #[inline] fn type_index() -> i32 { Object::type_index() } #[inline] unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject { Object::object_header_mut(&mut this.object) } } unsafe impl ObjectCoreWithExtraItems for TestIntObj { type ExtraItem = u64; #[inline] fn extra_items_count(this: &Self) -> usize { this.extra_item_count as usize } } #[test] fn test_object_arc() { let delete_counter = Arc::new(AtomicU32::new(0)); let obj_arc = ObjectArc::new(TestIntObj::new(11, delete_counter.clone(), 0)); assert_eq!(obj_arc.value, 11); assert_eq!(ObjectArc::strong_count(&obj_arc), 1); assert_eq!(ObjectArc::weak_count(&obj_arc), 1); let ref1 = obj_arc.clone(); assert_eq!(ObjectArc::strong_count(&obj_arc), 2); assert_eq!(ObjectArc::weak_count(&obj_arc), 1); let ref2 = obj_arc.clone(); assert_eq!(ObjectArc::strong_count(&obj_arc), 3); assert_eq!(ObjectArc::weak_count(&obj_arc), 1); assert_eq!(ref1.value, 11); // drop obj_arc drop(obj_arc); assert_eq!(ObjectArc::strong_count(&ref1), 2); assert_eq!(ObjectArc::weak_count(&ref1), 1); assert_eq!(delete_counter.load(Ordering::Relaxed), 0); // drop ref1 drop(ref1); assert_eq!(ObjectArc::strong_count(&ref2), 1); assert_eq!(ObjectArc::weak_count(&ref2), 1); assert_eq!(delete_counter.load(Ordering::Relaxed), 0); // drop ref2 drop(ref2); assert_eq!(delete_counter.load(Ordering::Relaxed), 1); } #[test] fn test_object_arc_with_extra_items() { let delete_counter = Arc::new(AtomicU32::new(0)); let mut obj_arc = ObjectArc::new_with_extra_items(TestIntObj::new(12, delete_counter.clone(), 10)); assert_eq!(obj_arc.value, 12); assert_eq!(ObjectArc::strong_count(&obj_arc), 1); assert_eq!(ObjectArc::weak_count(&obj_arc), 1); assert_eq!(delete_counter.load(Ordering::Relaxed), 0); unsafe { // layout check of extra items assert_eq!(TestIntObj::extra_items_count(&obj_arc), 10); assert_eq!(TestIntObj::extra_items(&obj_arc).len(), 10); assert_eq!(TestIntObj::extra_items_mut(&mut obj_arc).len(), 10); assert_eq!( TestIntObj::extra_items_mut(&mut obj_arc).as_ptr() as *mut u8, (ObjectArc::as_raw_mut(&mut obj_arc) as *mut u8).add(std::mem::size_of::()) ); } drop(obj_arc); assert_eq!(delete_counter.load(Ordering::Relaxed), 1); } #[test] fn test_object_arc_from_raw() { unsafe { let delete_counter = Arc::new(AtomicU32::new(0)); let obj_arc = ObjectArc::new(TestIntObj::new(11, delete_counter.clone(), 0)); let raw_ptr = ObjectArc::into_raw(obj_arc); let obj_arc2 = ObjectArc::from_raw(raw_ptr); assert_eq!(obj_arc2.value, 11); assert_eq!(ObjectArc::strong_count(&obj_arc2), 1); assert_eq!(ObjectArc::weak_count(&obj_arc2), 1); assert_eq!(delete_counter.load(Ordering::Relaxed), 0); // drop obj_arc2 drop(obj_arc2); assert_eq!(delete_counter.load(Ordering::Relaxed), 1); } } #[test] fn test_object_arc_option_size() { assert_eq!( std::mem::size_of::>>(), std::mem::size_of::>() ); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_shape.rs000066400000000000000000000027641521067262500213760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; // ============================================================================ // Shape Tests // ============================================================================ #[test] fn test_any_shape() { let shape = Shape::from(vec![1, 2, 3, 4]); let any = Any::from(shape.clone()); let any_view = AnyView::from(&shape); assert_eq!(any.type_index(), TypeIndex::kTVMFFIShape as i32); let converted = Shape::try_from(any).unwrap(); assert_eq!(converted.as_slice(), &[1, 2, 3, 4]); assert_eq!(any_view.type_index(), TypeIndex::kTVMFFIShape as i32); let converted_view = Shape::try_from(any_view).unwrap(); assert_eq!(converted_view.as_slice(), &[1, 2, 3, 4]); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_string.rs000066400000000000000000000243421521067262500216000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; use tvm_ffi::*; // ============================================================================ // Bytes Tests // ============================================================================ #[test] fn test_bytes_new() { let bytes = Bytes::new(); assert_eq!(bytes.len(), 0); assert_eq!(bytes.as_slice(), &[]); } #[test] fn test_bytes_default() { let bytes = Bytes::default(); assert_eq!(bytes.len(), 0); assert_eq!(bytes.as_slice(), &[]); } #[test] fn test_bytes_from_small() { let data = b"hello"; let bytes = Bytes::from(data.as_slice()); assert_eq!(bytes.len(), 5); assert_eq!(bytes.as_slice(), data); } #[test] fn test_bytes_from_large() { let data = b"this is a very long string that exceeds the small bytes limit"; let bytes = Bytes::from(data.as_slice()); assert_eq!( AnyView::from(&bytes).type_index(), TypeIndex::kTVMFFIBytes as i32 ); assert_eq!(bytes.len(), data.len()); assert_eq!(bytes.as_slice(), data); } #[test] fn test_bytes_deref() { let data = b"test"; let bytes = Bytes::from(data.as_slice()); assert_eq!(&*bytes, data); } #[test] fn test_bytes_clone_basic() { let data = b"clone test"; let bytes1 = Bytes::from(data.as_slice()); let bytes2 = bytes1.clone(); assert_eq!(bytes1, bytes2); assert_eq!(bytes1.as_slice(), bytes2.as_slice()); } #[test] fn test_bytes_eq() { let data = b"equal"; let bytes1 = Bytes::from(data.as_slice()); let bytes2 = Bytes::from(data.as_slice()); assert_eq!(bytes1, bytes2); let bytes3 = Bytes::from(b"different"); assert_ne!(bytes1, bytes3); } #[test] fn test_bytes_ord() { let bytes1 = Bytes::from(b"abc"); let bytes2 = Bytes::from(b"def"); let bytes3 = Bytes::from(b"abc"); assert!(bytes1 < bytes2); assert!(bytes1 <= bytes2); assert!(bytes2 > bytes1); assert!(bytes2 >= bytes1); assert_eq!(bytes1.cmp(&bytes3), std::cmp::Ordering::Equal); } #[test] fn test_bytes_hash() { let data = b"hash test"; let bytes1 = Bytes::from(data.as_slice()); let bytes2 = Bytes::from(data.as_slice()); let mut hasher1 = DefaultHasher::new(); let mut hasher2 = DefaultHasher::new(); bytes1.hash(&mut hasher1); bytes2.hash(&mut hasher2); assert_eq!(hasher1.finish(), hasher2.finish()); } // ============================================================================ // String Tests // ============================================================================ #[test] fn test_string_new() { let s = String::new(); assert_eq!(s.len(), 0); assert_eq!(s.as_str(), ""); assert_eq!(s.as_bytes(), &[]); } #[test] fn test_string_default() { let s = String::default(); assert_eq!(s.len(), 0); assert_eq!(s.as_str(), ""); } #[test] fn test_string_from_small() { let input = "hello"; let s = String::from(input); assert_eq!(s.len(), 5); assert_eq!(s.as_str(), input); assert_eq!(s.as_bytes(), input.as_bytes()); } #[test] fn test_string_from_large() { let input = "this is a very long string that exceeds the small string limit"; let s = String::from(input); assert_eq!(s.len(), input.len()); assert_eq!(s.as_str(), input); assert_eq!(s.as_bytes(), input.as_bytes()); } #[test] fn test_string_deref() { let input = "deref test"; let s = String::from(input); assert_eq!(&*s, input); } #[test] fn test_string_clone_basic() { let input = "clone test"; let s1 = String::from(input); let s2 = s1.clone(); assert_eq!(s1, s2); assert_eq!(s1.as_str(), s2.as_str()); } #[test] fn test_string_eq() { let input = "equal"; let s1 = String::from(input); let s2 = String::from(input); assert_eq!(s1, s2); let s3 = String::from("different"); assert_ne!(s1, s3); } #[test] fn test_string_ord() { let s1 = String::from("abc"); let s2 = String::from("def"); let s3 = String::from("abc"); assert!(s1 < s2); assert!(s1 <= s2); assert!(s2 > s1); assert!(s2 >= s1); assert_eq!(s1.cmp(&s3), std::cmp::Ordering::Equal); } #[test] fn test_string_hash() { let input = "hash test"; let s1 = String::from(input); let s2 = String::from(input); let mut hasher1 = DefaultHasher::new(); let mut hasher2 = DefaultHasher::new(); s1.hash(&mut hasher1); s2.hash(&mut hasher2); assert_eq!(hasher1.finish(), hasher2.finish()); } // ============================================================================ // Edge Cases // ============================================================================ #[test] fn test_empty_strings() { let empty_str = String::from(""); assert_eq!(empty_str.len(), 0); assert_eq!(empty_str.as_str(), ""); let empty_bytes = Bytes::from(&[]); assert_eq!(empty_bytes.len(), 0); assert_eq!(empty_bytes.as_slice(), &[]); } #[test] fn test_unicode_strings() { let unicode = "\u{1F680} Hello \u{4E16}\u{754C} \u{1F30D}"; let s = String::from(unicode); assert_eq!(s.as_str(), unicode); assert_eq!(s.len(), unicode.len()); } #[test] fn test_small_vs_large_boundary() { // Test exactly at the boundary (7 bytes) let boundary = "1234567"; // 7 characters = 7 bytes let s = String::from(boundary); assert_eq!(s.as_str(), boundary); // Test just over the boundary (8 bytes) let over_boundary = "12345678"; // 8 characters = 8 bytes let s2 = String::from(over_boundary); assert_eq!(s2.as_str(), over_boundary); } #[test] fn test_bytes_small_vs_large_boundary() { // Test exactly at the boundary (7 bytes) let boundary = b"1234567"; // 7 bytes let bytes = Bytes::from(boundary.as_slice()); assert_eq!(bytes.as_slice(), boundary); // Test just over the boundary (8 bytes) let over_boundary = b"12345678"; // 8 bytes let bytes2 = Bytes::from(over_boundary.as_slice()); assert_eq!(bytes2.as_slice(), over_boundary); } #[test] fn test_mixed_comparisons() { let s1 = String::from("test"); let s2 = String::from("test"); let s3 = String::from("different"); // Test all comparison operators assert!(s1 == s2); assert!(s1 != s3); assert!(s1 <= s2); assert!(s1 >= s2); assert!(s1 > s3); assert!(s3 < s1); } #[test] fn test_bytes_mixed_comparisons() { let b1 = Bytes::from(b"test"); let b2 = Bytes::from(b"test"); let b3 = Bytes::from(b"different"); // Test all comparison operators assert!(b1 == b2); assert!(b1 != b3); assert!(b1 <= b2); assert!(b1 >= b2); assert!(b1 > b3); assert!(b3 < b1); } #[test] fn test_bytes_debug() { let data = b"hello"; let bytes = Bytes::from(data.as_slice()); let debug_str = format!("{:?}", bytes); assert!(debug_str.contains("Bytes")); } #[test] fn test_string_debug() { let input = "world"; let s = String::from(input); let debug_str = format!("{:?}", s); assert!(debug_str.contains("String")); assert!(debug_str.contains("world")); } #[test] fn test_bytes_clone_small() { let data = b"small"; // 5 bytes - small case let bytes1 = Bytes::from(data.as_slice()); let bytes2 = bytes1.clone(); assert_eq!(bytes1, bytes2); assert_eq!(bytes1.as_slice(), bytes2.as_slice()); assert_eq!(bytes1.len(), bytes2.len()); // For small bytes, the underlying data should be different (copied) // We can verify this by checking that the slices have different pointers let slice1 = bytes1.as_slice(); let slice2 = bytes2.as_slice(); assert_ne!(slice1.as_ptr(), slice2.as_ptr()); } #[test] fn test_bytes_clone_large() { let data = b"this is a very long string that exceeds the small bytes limit"; // large case let bytes1 = Bytes::from(data.as_slice()); let bytes2 = bytes1.clone(); assert_eq!(bytes1, bytes2); assert_eq!(bytes1.as_slice(), bytes2.as_slice()); assert_eq!(bytes1.len(), bytes2.len()); // For large bytes, the underlying data should be the same (shared) // We can verify this by checking that the slices have the same pointer let slice1 = bytes1.as_slice(); let slice2 = bytes2.as_slice(); assert_eq!(slice1.as_ptr(), slice2.as_ptr()); } #[test] fn test_string_clone_small() { let input = "small"; // 5 bytes - small case let s1 = String::from(input); assert!(AnyView::from(&s1).debug_strong_count().is_none()); let s2 = s1.clone(); assert_eq!(s1, s2); assert_eq!(s1.as_str(), s2.as_str()); assert_eq!(s1.len(), s2.len()); // For small strings, the underlying data should be different (copied) // We can verify this by checking that the byte slices have different pointers let bytes1 = s1.as_bytes(); let bytes2 = s2.as_bytes(); assert_ne!(bytes1.as_ptr(), bytes2.as_ptr()); } #[test] fn test_string_clone_large() { // large case let input = "this is a very long string that exceeds the small string limit"; let s1 = String::from(input); assert_eq!(AnyView::from(&s1).debug_strong_count().unwrap(), 1); let s2 = s1.clone(); assert_eq!(AnyView::from(&s1).debug_strong_count().unwrap(), 2); assert_eq!(AnyView::from(&s2).debug_strong_count().unwrap(), 2); assert_eq!(s1, s2); assert_eq!(s1.as_str(), s2.as_str()); assert_eq!(s1.len(), s2.len()); // For large strings, the underlying data should be the same (shared) // We can verify this by checking that the byte slices have the same pointer let bytes1 = s1.as_bytes(); let bytes2 = s2.as_bytes(); assert_eq!(bytes1.as_ptr(), bytes2.as_ptr()); } tvm-ffi-0.1.12/rust/tvm-ffi/tests/test_tensor.rs000066400000000000000000000103771521067262500216070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ use tvm_ffi::*; // ============================================================================ // Tensor Tests // ============================================================================ #[test] fn test_tensor_basic() { // Create a tensor using CPUNDAlloc let shape = [2, 3, 4]; let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1); let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device); // Test accessor methods assert_eq!(tensor.shape(), &shape); assert_eq!(tensor.ndim(), 3); assert_eq!(tensor.dtype().code, DLDataTypeCode::kDLFloat as u8); assert_eq!(tensor.dtype().bits, 32 as u8); assert_eq!(tensor.device().device_type, DLDeviceType::kDLCPU); // Test strides (should be calculated correctly for row-major layout) let strides = tensor.strides(); assert_eq!(strides.len(), 3); // For shape [2, 3, 4], strides should be [12, 4, 1] (row-major) assert_eq!(strides[0], 12); // 3 * 4 assert_eq!(strides[1], 4); // 4 assert_eq!(strides[2], 1); // 1 } #[test] fn test_tensor_data_as_slice_f32() { // Create a tensor using CPUNDAlloc with f32 data type let shape = [2, 3, 4]; let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1); let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device); // Test data_as_slice for f32 let data_slice = tensor.data_as_slice::().unwrap(); assert_eq!(data_slice.len(), 24); // 2 * 3 * 4 = 24 elements // Test data_as_slice_mut for f32 let data_slice_mut = tensor.data_as_slice_mut::().unwrap(); assert_eq!(data_slice_mut.len(), 24); // Test that we can write to the tensor data for i in 0..data_slice_mut.len() { data_slice_mut[i] = i as f32; } // Test that we can read the written data let data_slice_read = tensor.data_as_slice::().unwrap(); for i in 0..data_slice_read.len() { assert_eq!(data_slice_read[i], i as f32); } } #[test] fn test_tensor_data_as_slice_type_mismatch() { // Create a tensor with f32 data type let shape = [2, 3]; let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1); let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device); // Test that trying to access as f64 (wrong type) fails let result = tensor.data_as_slice::(); assert!(result.is_err()); // Test that trying to access as i32 (wrong type) fails let result = tensor.data_as_slice::(); assert!(result.is_err()); } #[test] fn test_any_tensor() { let shape = [2, 3, 4]; let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1); let device = DLDevice::new(DLDeviceType::kDLCPU, 0); let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device); let any = Any::from(tensor.clone()); let any_view = AnyView::from(&tensor); assert_eq!(any.type_index(), TypeIndex::kTVMFFITensor as i32); let converted = Tensor::try_from(any).unwrap(); assert_eq!(converted.shape(), &shape); assert_eq!(converted.dtype().code, DLDataTypeCode::kDLFloat as u8); assert_eq!(any_view.type_index(), TypeIndex::kTVMFFITensor as i32); let converted_view = Tensor::try_from(any_view).unwrap(); assert_eq!(converted_view.shape(), &shape); assert_eq!(converted_view.dtype().code, DLDataTypeCode::kDLFloat as u8); } tvm-ffi-0.1.12/src/000077500000000000000000000000001521067262500137405ustar00rootroot00000000000000tvm-ffi-0.1.12/src/ffi/000077500000000000000000000000001521067262500145045ustar00rootroot00000000000000tvm-ffi-0.1.12/src/ffi/backtrace.cc000066400000000000000000000145631521067262500167430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file backtrace.cc * \brief Backtrace implementation on non-windows platforms * \note We use the term "backtrace" to be consistent with python naming convention. */ #ifndef _MSC_VER #include #include #include "./backtrace_utils.h" #if TVM_FFI_USE_LIBBACKTRACE #include #include #include #include #include #include #if TVM_FFI_BACKTRACE_ON_SEGFAULT #include #endif namespace tvm { namespace ffi { namespace { void BacktraceCreateErrorCallback(void*, const char* msg, int) { std::cerr << "Could not initialize backtrace state: " << msg << std::endl; } backtrace_state* BacktraceCreate() { return backtrace_create_state(nullptr, 1, BacktraceCreateErrorCallback, nullptr); } static backtrace_state* _bt_state = BacktraceCreate(); std::string DemangleName(std::string name) { int status = 0; size_t length = name.size(); char* demangled_name = abi::__cxa_demangle(name.c_str(), nullptr, &length, &status); if (demangled_name && status == 0 && length > 0) { name = demangled_name; } if (demangled_name) { std::free(demangled_name); } return name; } void BacktraceErrorCallback(void*, const char*, int) { // do nothing } void BacktraceSyminfoCallback(void* data, uintptr_t pc, const char* symname, uintptr_t, uintptr_t) { auto str = reinterpret_cast(data); if (symname != nullptr) { *str = DemangleName(symname); } else { std::ostringstream s; s << "0x" << std::setfill('0') << std::setw(sizeof(uintptr_t) * 2) << std::hex << pc; *str = s.str(); } } int BacktraceFullCallback(void* data, uintptr_t pc, const char* filename, int lineno, const char* symbol) { auto stack_trace = reinterpret_cast(data); std::string symbol_str = ""; if (symbol) { symbol_str = DemangleName(symbol); } else { // see if syminfo gives anything backtrace_syminfo(_bt_state, pc, BacktraceSyminfoCallback, BacktraceErrorCallback, &symbol_str); } symbol = symbol_str.data(); if (stack_trace->ExceedBacktraceLimit()) { return 1; } if (stack_trace->stop_at_boundary && DetectFFIBoundary(filename, symbol)) { return 1; } // skip extra frames if (stack_trace->skip_frame_count > 0) { stack_trace->skip_frame_count--; return 0; } if (ShouldExcludeFrame(filename, symbol)) { return 0; } stack_trace->Append(filename, symbol, lineno); return 0; } } // namespace } // namespace ffi } // namespace tvm const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const char* func, int cross_ffi_boundary) { // We collapse the backtrace into a single function // to simplify the backtrace detection handling (since we need to detect TVMFFIBacktrace) static thread_local std::string backtrace_str; static thread_local TVMFFIByteArray backtrace_array; // pass in current line as here so last line of backtrace is always accurate tvm::ffi::BacktraceStorage backtrace; backtrace.stop_at_boundary = cross_ffi_boundary == 0; if (filename != nullptr && func != nullptr) { // need to skip TVMFFIBacktrace and the caller function // which is already included in filename and func backtrace.skip_frame_count = 2; if (!tvm::ffi::ShouldExcludeFrame(filename, func)) { backtrace.Append(filename, func, lineno); } } // libbacktrace eats memory if run on multiple threads at the same time, so we guard against it if (tvm::ffi::_bt_state != nullptr) { static std::mutex m; std::scoped_lock lock(m); backtrace_full(tvm::ffi::_bt_state, 0, tvm::ffi::BacktraceFullCallback, tvm::ffi::BacktraceErrorCallback, &backtrace); } backtrace_str = backtrace.GetBacktrace(); backtrace_array.data = backtrace_str.data(); backtrace_array.size = backtrace_str.size(); return &backtrace_array; } #if TVM_FFI_BACKTRACE_ON_SEGFAULT TVM_FFI_COLD_CODE void TVMFFISegFaultHandler(int sig) { // Technically we shouldn't do any allocation in a signal handler, but // Backtrace may allocate. What's the worst it could do? We're already // crashing. const TVMFFIByteArray* backtrace = TVMFFIBacktrace(nullptr, 0, nullptr, 1); std::cerr << "!!!!!!! Segfault encountered !!!!!!!\n" << std::string(backtrace->data, backtrace->size) << std::endl; // Re-raise signal with default handler struct sigaction act; std::memset(&act, 0, sizeof(struct sigaction)); act.sa_flags = SA_RESETHAND; act.sa_handler = SIG_DFL; sigaction(sig, &act, nullptr); raise(sig); } TVM_FFI_COLD_CODE __attribute__((constructor)) void TVMFFIInstallSignalHandler() { // this may override already installed signal handlers std::signal(SIGSEGV, TVMFFISegFaultHandler); } #endif // TVM_FFI_BACKTRACE_ON_SEGFAULT #else // fallback implementation simply print out the last trace const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const char* func, int cross_ffi_boundary) { static thread_local std::string backtrace_str; static thread_local TVMFFIByteArray backtrace_array; std::ostringstream backtrace_stream; if (filename != nullptr && func != nullptr) { // python style backtrace backtrace_stream << " File \"" << filename << "\", line " << lineno << ", in " << func << std::endl; } backtrace_str = backtrace_stream.str(); backtrace_array.data = backtrace_str.data(); backtrace_array.size = backtrace_str.size(); return &backtrace_array; } #endif // TVM_FFI_USE_LIBBACKTRACE #endif // _MSC_VER tvm-ffi-0.1.12/src/ffi/backtrace_utils.h000066400000000000000000000123261521067262500200200ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file backtrace.h * \brief Common headers for backtrace. * \note We use the term "backtrace" to be consistent with python naming convention. */ #ifndef TVM_FFI_TRACEBACK_H_ #define TVM_FFI_TRACEBACK_H_ #include #include #include #include #include namespace tvm { namespace ffi { #ifdef _MSC_VER #pragma warning(push) #pragma warning(disable : 4996) // std::getenv is unsafe #endif inline int32_t GetBacktraceLimit() { if (const char* env = std::getenv("TVM_TRACEBACK_LIMIT")) { return std::stoi(env); } return 512; } #ifdef _MSC_VER #pragma warning(pop) #endif /*! * \brief List frame patterns that should be excluded as they contain less information */ inline bool ShouldExcludeFrame(const char* filename, const char* symbol) { if (symbol != nullptr) { if (strncmp(symbol, "tvm::ffi::Function", 18) == 0) { return true; } if (strncmp(symbol, "tvm::ffi::details::", 19) == 0) { return true; } if (strncmp(symbol, "TVMFFIBacktrace", 15) == 0) { return true; } if (strncmp(symbol, "TVMFFIErrorSetRaisedFromCStr", 28) == 0) { return true; } if (strncmp(symbol, "TVMFFIErrorSetRaisedFromCStrParts", 33) == 0) { return true; } // C++ stdlib frames if (strncmp(symbol, "__libc_", 7) == 0) { return true; } // libffi.so stack frames. These may also show up as numeric // addresses with no symbol name. This could be improved in the // future by using dladdr() to check whether an address is contained // in libffi.so if (strncmp(symbol, "ffi_call_", 9) == 0) { return true; } } if (filename) { // Stack frames for TVM FFI if (strstr(filename, "include/tvm/ffi/error.h") != nullptr) { return true; } if (strstr(filename, "include/tvm/ffi/function_details.h") != nullptr) { return true; } if (strstr(filename, "include/tvm/ffi/function.h") != nullptr) { return true; } if (strstr(filename, "include/tvm/ffi/any.h") != nullptr) { return true; } // C++ stdlib frames if (strstr(filename, "include/c++/") != nullptr) { return true; } } return false; } /** * \brief List frames that should stop the backtrace. * \param filename The filename of the frame. * \param symbol The symbol name of the frame. * \return true if the frame should stop the backtrace. * \note We stop backtrace at the FFI boundary. */ inline bool DetectFFIBoundary(const char* filename, const char* symbol) { if (symbol != nullptr) { if (strncmp(symbol, "TVMFFIFunctionCall", 18) == 0) { return true; } // python ABI functions if (strncmp(symbol, "slot_tp_call", 12) == 0) { return true; } if (strncmp(symbol, "object_is_not_callable", 11) == 0) { return true; } // Python interpreter stack frames // we stop backtrace at the Python interpreter stack frames // since these frame will be handled from by the python side. if (strncmp(symbol, "_Py", 3) == 0 || strncmp(symbol, "PyObject", 8) == 0) { return true; } } return false; } /*! * \brief storage to store backtrace */ struct BacktraceStorage { /*! \brief The stream to store the backtrace. */ std::ostringstream backtrace_stream_; /*! \brief The number of lines in the backtrace. */ size_t line_count_ = 0; /*! \brief Maximum size of the backtrace. */ size_t max_frame_size = GetBacktraceLimit(); /*! \brief Number of frames to skip. */ size_t skip_frame_count = 0; /*! \brief Whether to stop at the ffi boundary. */ bool stop_at_boundary = true; void Append(const char* filename, const char* func, int lineno) { // skip frames with empty filename if (filename == nullptr) { if (func != nullptr) { if (strncmp(func, "0x0", 3) == 0) { return; } if (strncmp(func, "", 9) == 0) { return; } filename = ""; } else { return; } } backtrace_stream_ << " File \"" << filename << "\""; backtrace_stream_ << ", line " << lineno; backtrace_stream_ << ", in " << func << '\n'; line_count_++; } bool ExceedBacktraceLimit() const { return line_count_ >= max_frame_size; } // get backtrace in the order of most recent call last std::string GetBacktrace() const { return backtrace_stream_.str(); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_TRACEBACK_H_ tvm-ffi-0.1.12/src/ffi/backtrace_win.cc000066400000000000000000000117621521067262500176160ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file backtrace_win.cc * \brief Backtrace implementation on windows platform * \note We use the term "backtrace" to be consistent with python naming convention. */ #ifdef _MSC_VER // clang-format off #include #include // NOLINT(*) // clang-format on #include #include #include #include #include "./backtrace_utils.h" const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const char* func, int cross_ffi_boundary) { static thread_local std::string backtrace_str; static thread_local TVMFFIByteArray backtrace_array; // pass in current line as here so last line of backtrace is always accurate tvm::ffi::BacktraceStorage backtrace; backtrace.stop_at_boundary = cross_ffi_boundary == 0; if (filename != nullptr && func != nullptr) { // need to skip TVMFFIBacktrace and the caller function // which is already included in filename and func backtrace.skip_frame_count = 2; backtrace.Append(filename, func, lineno); } HANDLE process = GetCurrentProcess(); HANDLE thread = GetCurrentThread(); SymSetOptions(SYMOPT_LOAD_LINES | SYMOPT_UNDNAME); SymInitialize(process, NULL, TRUE); CONTEXT context = {}; RtlCaptureContext(&context); STACKFRAME64 stack = {}; DWORD machine_type; #ifdef _M_IX86 machine_type = IMAGE_FILE_MACHINE_I386; stack.AddrPC.Offset = context.Eip; stack.AddrStack.Offset = context.Esp; stack.AddrFrame.Offset = context.Ebp; #elif _M_X64 machine_type = IMAGE_FILE_MACHINE_AMD64; stack.AddrPC.Offset = context.Rip; stack.AddrStack.Offset = context.Rsp; stack.AddrFrame.Offset = context.Rbp; #elif _M_ARM64 machine_type = IMAGE_FILE_MACHINE_ARM64; stack.AddrPC.Offset = context.Pc; stack.AddrStack.Offset = context.Sp; stack.AddrFrame.Offset = context.Fp; #else #error "Unsupported architecture" #endif stack.AddrPC.Mode = AddrModeFlat; stack.AddrFrame.Mode = AddrModeFlat; stack.AddrStack.Mode = AddrModeFlat; while (!backtrace.ExceedBacktraceLimit()) { if (!StackWalk64(machine_type, process, thread, &stack, &context, nullptr, SymFunctionTableAccess64, SymGetModuleBase64, nullptr)) { break; } if (stack.AddrPC.Offset == 0) { break; } const char* filename = nullptr; const char* symbol = ""; int lineno = 0; // Get file and line number IMAGEHLP_LINE64 line_info; ZeroMemory(&line_info, sizeof(IMAGEHLP_LINE64)); line_info.SizeOfStruct = sizeof(IMAGEHLP_LINE64); DWORD displacement32 = 0; if (SymGetLineFromAddr64(process, stack.AddrPC.Offset, &displacement32, &line_info)) { filename = line_info.FileName; lineno = line_info.LineNumber; } // allocate symbol info that aligns to the SYMBOL_INFO // we use u64 here to be safe size_t total_symbol_bytes = sizeof(SYMBOL_INFO) + MAX_SYM_NAME * sizeof(TCHAR); size_t total_u64_words = (total_symbol_bytes + 7) / 8; static_assert(8 % alignof(SYMBOL_INFO) == 0); std::vector symbol_buffer(total_u64_words, 0); if (filename != nullptr) { // only run symbol translation if we have the file name // this is because SymFromAddr can return wrong symbol which becomes even more // confusing when pdb file do not exist PSYMBOL_INFO symbol_info = reinterpret_cast(symbol_buffer.data()); symbol_info->SizeOfStruct = sizeof(SYMBOL_INFO); symbol_info->MaxNameLen = MAX_SYM_NAME; DWORD64 displacement = 0; if (SymFromAddr(process, stack.AddrPC.Offset, &displacement, symbol_info)) { symbol = symbol_info->Name; } } if (backtrace.stop_at_boundary && tvm::ffi::DetectFFIBoundary(filename, symbol)) { break; } // skip extra frames if (backtrace.skip_frame_count > 0) { backtrace.skip_frame_count--; continue; } if (tvm::ffi::ShouldExcludeFrame(filename, symbol)) { continue; } backtrace.Append(filename, symbol, lineno); } SymCleanup(process); backtrace_str = backtrace.GetBacktrace(); backtrace_array.data = backtrace_str.data(); backtrace_array.size = backtrace_str.size(); return &backtrace_array; } #endif // _MSC_VER tvm-ffi-0.1.12/src/ffi/container.cc000066400000000000000000000225511521067262500170020ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/container.cc */ #include #include #include #include #include #include #include #include "object_internal.h" namespace tvm { namespace ffi { namespace { /*! * \brief Recursively scan an Any element for the first non-CPU tensor device. * \param elem The element to inspect. * \param out Output device; written only when a non-CPU tensor is found. * \return true if a non-CPU tensor was found. */ bool FindFirstNonCPUDevice(const Any& elem, DLDevice* out) { switch (elem.type_index()) { case TypeIndex::kTVMFFITensor: { const auto* tensor = elem.as(); if (tensor->device.device_type != kDLCPU) { *out = tensor->device; return true; } break; } case TypeIndex::kTVMFFIArray: case TypeIndex::kTVMFFIList: { const auto* seq = elem.as(); for (const auto& it : *seq) { if (FindFirstNonCPUDevice(it, out)) return true; } break; } case TypeIndex::kTVMFFIMap: case TypeIndex::kTVMFFIDict: { const auto* map = elem.as(); for (const auto& it : *map) { if (FindFirstNonCPUDevice(it.second, out)) return true; } break; } default: break; } return false; } } // namespace // Favor struct outside function scope as MSVC may have bug for in fn scope struct. class MapForwardIterFunctor { public: MapForwardIterFunctor(ffi::MapObj::iterator iter, ffi::MapObj::iterator end) : iter_(iter), end_(end) {} // 0 get current key // 1 get current value // 2 move to next: return true if success, false if end Any operator()(int command) const { if (command == 0) { return (*iter_).first; } else if (command == 1) { return (*iter_).second; } else { ++iter_; if (iter_ == end_) { return false; } return true; } } private: mutable ffi::MapObj::iterator iter_; ffi::MapObj::iterator end_; }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::EnsureTypeAttrColumn(refl::type_attr::kAnyHash); refl::EnsureTypeAttrColumn(refl::type_attr::kAnyEqual); refl::GlobalDef() .def_packed("ffi.Array", [](ffi::PackedArgs args, Any* ret) { *ret = Array(args.data(), args.data() + args.size()); }) .def("ffi.ArrayGetItem", [](const ffi::ArrayObj* n, int64_t i) -> Any { return n->at(i); }) .def("ffi.ArraySize", [](const ffi::ArrayObj* n) -> int64_t { return static_cast(n->size()); }) .def("ffi.ArrayContains", [](const ffi::ArrayObj* n, const Any& value) -> bool { AnyEqual eq; return std::any_of(n->begin(), n->end(), [&](const Any& elem) { return eq(elem, value); }); }) .def_packed("ffi.List", [](ffi::PackedArgs args, Any* ret) { *ret = List(args.data(), args.data() + args.size()); }) .def("ffi.ListGetItem", [](const ffi::ListObj* n, int64_t i) -> Any { return n->at(i); }) .def("ffi.ListSetItem", [](ffi::List n, int64_t i, Any value) -> void { n.Set(i, std::move(value)); }) .def("ffi.ListSize", [](const ffi::ListObj* n) -> int64_t { return static_cast(n->size()); }) .def("ffi.ListContains", [](const ffi::ListObj* n, const Any& value) -> bool { AnyEqual eq; return std::any_of(n->begin(), n->end(), [&](const Any& elem) { return eq(elem, value); }); }) .def("ffi.ListAppend", [](ffi::List n, const Any& value) -> void { n.push_back(value); }) .def("ffi.ListInsert", [](ffi::List n, int64_t i, const Any& value) -> void { n.insert(n.begin() + i, value); }) .def("ffi.ListPop", [](const ffi::List& n, int64_t i) -> Any { ffi::ListObj* obj = n.GetListObj(); Any value = obj->at(i); obj->erase(i); return value; }) .def("ffi.ListErase", [](const ffi::List& n, int64_t i) -> void { n.GetListObj()->erase(i); }) .def("ffi.ListEraseRange", [](const ffi::List& n, int64_t start, int64_t stop) -> void { n.GetListObj()->erase(start, stop); }) .def("ffi.ListReplaceSlice", [](ffi::List n, int64_t start, int64_t stop, const ffi::List& replacement) -> void { // Snapshot replacement before erasing in case n and replacement alias the same object. ffi::List rep_copy = n.same_as(replacement) ? ffi::List(replacement.begin(), replacement.end()) : replacement; n.GetListObj()->erase(start, stop); if (rep_copy.empty()) { return; } const ffi::ListObj* replacement_obj = rep_copy.GetListObj(); TVM_FFI_ICHECK(replacement_obj != nullptr); n.insert(n.begin() + start, replacement_obj->begin(), replacement_obj->end()); }) .def("ffi.ListReverse", [](const ffi::List& n) -> void { ffi::ListObj* obj = n.GetListObj(); if (obj != nullptr) { obj->SeqBaseObj::Reverse(); } }) .def("ffi.ListClear", [](ffi::List n) -> void { n.clear(); }) .def_packed("ffi.Map", [](ffi::PackedArgs args, Any* ret) { TVM_FFI_ICHECK_EQ(args.size() % 2, 0); Map data; for (int i = 0; i < args.size(); i += 2) { data.Set(args[i], args[i + 1]); } *ret = data; }) .def("ffi.MapSize", [](const ffi::MapObj* n) -> int64_t { return static_cast(n->size()); }) .def("ffi.MapGetItem", [](const ffi::MapObj* n, const Any& k) -> Any { return n->at(k); }) .def("ffi.MapCount", [](const ffi::MapObj* n, const Any& k) -> int64_t { return static_cast(n->count(k)); }) .def("ffi.MapForwardIterFunctor", [](const ffi::MapObj* n) -> ffi::Function { return ffi::Function::FromTyped(MapForwardIterFunctor(n->begin(), n->end())); }) .def("ffi.MapGetItemOrMissing", [](const ffi::MapObj* n, const Any& k) -> Any { auto it = n->find(k); if (it != n->end()) { return it->second; } return GetMissingObject(); }) .def_packed("ffi.Dict", [](ffi::PackedArgs args, Any* ret) { TVM_FFI_ICHECK_EQ(args.size() % 2, 0); Dict data; for (int i = 0; i < args.size(); i += 2) { data.Set(args[i], args[i + 1]); } *ret = data; }) .def("ffi.DictSize", [](const ffi::DictObj* n) -> int64_t { return static_cast(n->size()); }) .def("ffi.DictGetItem", [](const ffi::DictObj* n, const Any& k) -> Any { return n->at(k); }) .def("ffi.DictSetItem", [](ffi::Dict d, const Any& k, const Any& v) -> void { d.Set(k, v); }) .def("ffi.DictCount", [](const ffi::DictObj* n, const Any& k) -> int64_t { return static_cast(n->count(k)); }) .def("ffi.DictErase", [](ffi::Dict d, const Any& k) -> void { d.erase(k); }) .def("ffi.DictClear", [](ffi::Dict d) -> void { d.clear(); }) .def("ffi.DictForwardIterFunctor", [](const ffi::DictObj* n) -> ffi::Function { return ffi::Function::FromTyped(MapForwardIterFunctor(n->begin(), n->end())); }) .def("ffi.DictGetItemOrMissing", [](const ffi::DictObj* n, const Any& k) -> Any { auto it = n->find(k); if (it != n->end()) { return it->second; } return GetMissingObject(); }) .def("ffi.ContainerFindFirstNonCPUDevice", [](const Any& container) -> DLDevice { DLDevice result{kDLCPU, 0}; FindFirstNonCPUDevice(container, &result); return result; }); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/dtype.cc000066400000000000000000000304311521067262500161410ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include namespace tvm { namespace ffi { namespace details { /*! * \brief Get the custom type name for a given type code. */ inline String DLDataTypeCodeGetCustomTypeName(DLDataTypeCode type_code) { static Function fget_custom_type_name = Function::GetGlobalRequired("dtype.get_custom_type_name"); return fget_custom_type_name(static_cast(type_code)).cast(); } /*! * \brief Get the custom type name for a given type code. * \param str The string to parse. * \param scan The scan pointer. * \return The custom type name. */ inline int ParseCustomDataTypeCode(const std::string_view& str, const char** scan) { TVM_FFI_ICHECK(str.substr(0, 6) == "custom") << "Not a valid custom datatype string"; auto tmp = str.data(); TVM_FFI_ICHECK(str.data() == tmp); *scan = str.data() + 6; TVM_FFI_ICHECK(str.data() == tmp); if (**scan != '[') { TVM_FFI_THROW(ValueError) << "expected opening brace after 'custom' type in" << str; } TVM_FFI_ICHECK(str.data() == tmp); *scan += 1; TVM_FFI_ICHECK(str.data() == tmp); size_t custom_name_len = 0; TVM_FFI_ICHECK(str.data() == tmp); while (*scan + custom_name_len <= str.data() + str.length() && *(*scan + custom_name_len) != ']') { ++custom_name_len; } TVM_FFI_ICHECK(str.data() == tmp); if (*(*scan + custom_name_len) != ']') { TVM_FFI_THROW(ValueError) << "expected closing brace after 'custom' type in" << str; } TVM_FFI_ICHECK(str.data() == tmp); *scan += custom_name_len + 1; TVM_FFI_ICHECK(str.data() == tmp); auto type_name = str.substr(7, custom_name_len); TVM_FFI_ICHECK(str.data() == tmp); static Function fget_custom_type_code = Function::GetGlobalRequired("dtype.get_custom_type_code"); return fget_custom_type_code(std::string(type_name)).cast(); } /* * \brief Convert a DLDataTypeCode to a string. * \param os The output stream. * \param type_code The DLDataTypeCode to convert. */ inline void PrintDLDataTypeCodeAsStr(std::ostream& os, DLDataTypeCode type_code) { // NOLINT(*) switch (static_cast(type_code)) { case kDLInt: { os << "int"; break; } case kDLUInt: { os << "uint"; break; } case kDLFloat: { os << "float"; break; } case kDLOpaqueHandle: { os << "handle"; break; } case kDLBfloat: { os << "bfloat"; break; } case kDLFloat8_e3m4: { os << "float8_e3m4"; break; } case kDLFloat8_e4m3: { os << "float8_e4m3"; break; } case kDLFloat8_e4m3b11fnuz: { os << "float8_e4m3b11fnuz"; break; } case kDLFloat8_e4m3fn: { os << "float8_e4m3fn"; break; } case kDLFloat8_e4m3fnuz: { os << "float8_e4m3fnuz"; break; } case kDLFloat8_e5m2: { os << "float8_e5m2"; break; } case kDLFloat8_e5m2fnuz: { os << "float8_e5m2fnuz"; break; } case kDLFloat8_e8m0fnu: { os << "float8_e8m0fnu"; break; } case kDLFloat6_e2m3fn: { os << "float6_e2m3fn"; break; } case kDLFloat6_e3m2fn: { os << "float6_e3m2fn"; break; } case kDLFloat4_e2m1fn: { os << "float4_e2m1fn"; break; } default: { if (static_cast(type_code) >= static_cast(DLExtDataTypeCode::kDLExtCustomBegin)) { os << "custom[" << details::DLDataTypeCodeGetCustomTypeName(type_code) << "]"; } else { TVM_FFI_THROW(ValueError) << "DLDataType contains unknown type_code=" << static_cast(type_code); } TVM_FFI_UNREACHABLE(); } } } } // namespace details /*! * \brief Printer function for DLDataType. * \param os The output stream. * \param dtype The DLDataType to print. * \return The output stream. */ inline std::string DLDataTypeToString_(DLDataType dtype) { // NOLINT(*) if (dtype.bits == 8 && dtype.lanes == 1 && dtype.code == kDLBool) { return "bool"; } // specially handle void if (dtype.code == kDLOpaqueHandle && dtype.lanes == 0 && dtype.bits == 0) { return ""; } std::ostringstream os; if (dtype.code >= kDLExtCustomBegin) { os << "custom[" << details::DLDataTypeCodeGetCustomTypeName(static_cast(dtype.code)) << "]"; } else { os << details::DLDataTypeCodeAsCStr(static_cast(dtype.code)); } if (dtype.code == kDLOpaqueHandle) return os.str(); int16_t lanes = static_cast(dtype.lanes); if (dtype.code < kDLFloat8_e3m4 && (dtype.code != kDLBool || dtype.bits != 8)) { os << static_cast(dtype.bits); } if (lanes > 1) { os << 'x' << lanes; } else if (lanes < -1) { os << "xvscalex" << -lanes; } return os.str(); } /*! * \brief Parse a string to a DLDataType. * \param str The string to convert. * \return The corresponding DLDataType. */ inline DLDataType StringViewToDLDataType_(std::string_view str) { DLDataType dtype; // handle void type if (str.length() == 0 || str == "void") { dtype.code = kDLOpaqueHandle; dtype.bits = 0; dtype.lanes = 0; return dtype; } // set the default values; dtype.bits = 32; dtype.lanes = 1; const char* scan; const char* str_end = str.data() + str.length(); // Helper lambda to parse decimal digits from a bounded string_view // Returns the parsed value and updates *ptr to point past the last digit auto parse_digits = [](const char** ptr, const char* end) -> uint32_t { uint64_t value = 0; const char* start_ptr = *ptr; while (*ptr < end && **ptr >= '0' && **ptr <= '9') { value = value * 10 + (**ptr - '0'); (*ptr)++; } if (value > UINT32_MAX) { TVM_FFI_THROW(ValueError) << "Integer value in dtype string '" << std::string_view(start_ptr, *ptr - start_ptr) << "' is out of range for uint32_t"; } return static_cast(value); }; // Helper lambda to parse lanes specification (e.g., "x16" or "xvscalex4") // Returns the parsed lanes value and updates *ptr to point past the lanes specification // Supports scalable vectors with the "xvscale" prefix (represented as negative lanes) auto parse_lanes = [&](const char** ptr, const char* end, const std::string_view& dtype_str, bool allow_scalable = false) -> uint16_t { int multiplier = 1; // Check for "xvscale" prefix for scalable vectors if (allow_scalable && (end - *ptr >= 7) && strncmp(*ptr, "xvscale", 7) == 0) { multiplier = -1; *ptr += 7; } if (*ptr >= end || **ptr != 'x') { return 1; // No lanes specification, default to 1 } (*ptr)++; // Skip 'x' const char* digits_start = *ptr; uint32_t lanes_val = parse_digits(ptr, end); if (*ptr == digits_start || lanes_val == 0) { TVM_FFI_THROW(ValueError) << "Invalid lanes specification in dtype '" << dtype_str << "'. Lanes must be a positive integer."; } if (lanes_val > UINT16_MAX) { TVM_FFI_THROW(ValueError) << "Lanes value " << lanes_val << " is out of range for uint16_t in dtype '" << dtype_str << "'"; } return static_cast(multiplier * lanes_val); }; auto parse_float = [&](const std::string_view& str, int offset, int code, int bits) { dtype.code = static_cast(code); dtype.bits = static_cast(bits); scan = str.data() + offset; const char* endpt = scan; dtype.lanes = parse_lanes(&endpt, str_end, str); scan = endpt; if (scan != str_end) { TVM_FFI_THROW(ValueError) << "unknown dtype `" << str << '`'; } return dtype; }; if (str.compare(0, 3, "int") == 0) { dtype.code = kDLInt; scan = str.data() + 3; } else if (str.compare(0, 4, "uint") == 0) { dtype.code = kDLUInt; scan = str.data() + 4; } else if (str.compare(0, 4, "bool") == 0) { dtype.code = kDLBool; dtype.bits = 8; scan = str.data() + 4; } else if (str.compare(0, 5, "float") == 0) { if (str.compare(5, 2, "8_") == 0) { if (str.compare(7, 4, "e3m4") == 0) { return parse_float(str, 11, kDLFloat8_e3m4, 8); } else if (str.compare(7, 4, "e4m3") == 0) { if (str.compare(11, 7, "b11fnuz") == 0) { return parse_float(str, 18, kDLFloat8_e4m3b11fnuz, 8); } else if (str.compare(11, 2, "fn") == 0) { if (str.compare(13, 2, "uz") == 0) { return parse_float(str, 15, kDLFloat8_e4m3fnuz, 8); } else { return parse_float(str, 13, kDLFloat8_e4m3fn, 8); } } else { return parse_float(str, 11, kDLFloat8_e4m3, 8); } } else if (str.compare(7, 8, "e5m2fnuz") == 0) { return parse_float(str, 15, kDLFloat8_e5m2fnuz, 8); } else if (str.compare(7, 4, "e5m2") == 0) { return parse_float(str, 11, kDLFloat8_e5m2, 8); } else if (str.compare(7, 7, "e8m0fnu") == 0) { return parse_float(str, 14, kDLFloat8_e8m0fnu, 8); } else { TVM_FFI_THROW(ValueError) << "unknown float8 type `" << str << '`'; TVM_FFI_UNREACHABLE(); } } else if (str.compare(5, 2, "6_") == 0) { if (str.compare(7, 6, "e2m3fn") == 0) { return parse_float(str, 13, kDLFloat6_e2m3fn, 6); } else if (str.compare(7, 6, "e3m2fn") == 0) { return parse_float(str, 13, kDLFloat6_e3m2fn, 6); } else { TVM_FFI_THROW(ValueError) << "unknown float6 type `" << str << '`'; TVM_FFI_UNREACHABLE(); } } else if (str.compare(5, 2, "4_") == 0) { // kFloat4_e2m1fn if (str.compare(7, 6, "e2m1fn") == 0) { return parse_float(str, 13, kDLFloat4_e2m1fn, 4); } else { TVM_FFI_THROW(ValueError) << "unknown float4 type `" << str << '`'; TVM_FFI_UNREACHABLE(); } } else { dtype.code = kDLFloat; scan = str.data() + 5; } } else if (str.compare(0, 6, "handle") == 0) { dtype.code = kDLOpaqueHandle; dtype.bits = 64; // handle uses 64 bit by default. scan = str.data() + 6; } else if (str.compare(0, 6, "bfloat") == 0) { dtype.code = kDLBfloat; dtype.bits = 16; scan = str.data() + 6; } else if (str.compare(0, 6, "custom") == 0) { dtype.code = static_cast(details::ParseCustomDataTypeCode(str, &scan)); } else { scan = str.data(); TVM_FFI_THROW(ValueError) << "unknown dtype `" << str << '`'; } // Parse bits manually to handle non-null-terminated string_view const char* xdelim = scan; uint32_t bits_val = parse_digits(&xdelim, str_end); if (bits_val > UINT8_MAX) { TVM_FFI_THROW(ValueError) << "Bits value " << bits_val << " is out of range for uint8_t in dtype '" << str << "'"; } uint8_t bits = static_cast(bits_val); if (bits != 0) dtype.bits = bits; const char* endpt = xdelim; dtype.lanes = parse_lanes(&endpt, str_end, str, /*allow_scalable=*/true); if (endpt != str_end) { TVM_FFI_THROW(ValueError) << "unknown dtype `" << str << '`'; } return dtype; } } // namespace ffi } // namespace tvm int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType* out) { TVM_FFI_SAFE_CALL_BEGIN(); *out = tvm::ffi::StringViewToDLDataType_(std::string_view(str->data, str->size)); TVM_FFI_SAFE_CALL_END(); } int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIAny* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::String out_str(tvm::ffi::DLDataTypeToString_(*dtype)); tvm::ffi::TypeTraits::MoveToAny(std::move(out_str), out); TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/error.cc000066400000000000000000000127471521067262500161570ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/error.cc * \brief Error handling implementation */ #include #include #include #include namespace tvm { namespace ffi { class SafeCallContext { public: void SetRaised(TVMFFIObjectHandle error) { last_error_ = details::ObjectUnsafe::ObjectPtrFromUnowned(static_cast(error)); } void SetRaisedByCstr(const char* kind, const char* message, const TVMFFIByteArray* backtrace) { Error error(kind, message, backtrace); last_error_ = details::ObjectUnsafe::ObjectPtrFromObjectRef(std::move(error)); } void SetRaisedByCstrParts(const char* kind, const char** message_parts, int32_t num_parts, const TVMFFIByteArray* backtrace) { std::string message; size_t total_len = 0; for (int i = 0; i < num_parts; ++i) { if (message_parts[i] != nullptr) { total_len += std::strlen(message_parts[i]); } } message.reserve(total_len); for (int i = 0; i < num_parts; ++i) { if (message_parts[i] != nullptr) { message.append(message_parts[i]); } } Error error(kind, message, backtrace); last_error_ = details::ObjectUnsafe::ObjectPtrFromObjectRef(std::move(error)); } void MoveFromRaised(TVMFFIObjectHandle* result) { result[0] = details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(last_error_)); } static SafeCallContext* ThreadLocal() { static thread_local SafeCallContext ctx; return &ctx; } private: ObjectPtr last_error_; }; } // namespace ffi } // namespace tvm void TVMFFIErrorSetRaisedFromCStr(const char* kind, const char* message) { // NOTE: run backtrace here to simplify the depth of tracekback tvm::ffi::SafeCallContext::ThreadLocal()->SetRaisedByCstr( kind, message, TVMFFIBacktrace(nullptr, 0, nullptr, 0)); } void TVMFFIErrorSetRaisedFromCStrParts(const char* kind, const char** message_parts, int32_t num_parts) { // NOTE: run backtrace here to simplify the depth of tracekback tvm::ffi::SafeCallContext::ThreadLocal()->SetRaisedByCstrParts( kind, message_parts, num_parts, TVMFFIBacktrace(nullptr, 0, nullptr, 0)); } void TVMFFIErrorSetRaised(TVMFFIObjectHandle error) { tvm::ffi::SafeCallContext::ThreadLocal()->SetRaised(error); } void TVMFFIErrorMoveFromRaised(TVMFFIObjectHandle* result) { tvm::ffi::SafeCallContext::ThreadLocal()->MoveFromRaised(result); } int TVMFFIErrorCreate(const TVMFFIByteArray* kind, const TVMFFIByteArray* message, const TVMFFIByteArray* backtrace, TVMFFIObjectHandle* out) { // log other errors to the logger TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); try { tvm::ffi::Error error(std::string(kind->data, kind->size), std::string(message->data, message->size), std::string(backtrace->data, backtrace->size)); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(error)); return 0; } catch (const std::bad_alloc& e) { return -1; } TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIErrorCreate); } int TVMFFIErrorCreateWithCauseAndExtraContext( const TVMFFIByteArray* kind, const TVMFFIByteArray* message, const TVMFFIByteArray* backtrace, TVMFFIObjectHandle cause_chain, TVMFFIObjectHandle extra_context, TVMFFIObjectHandle* out) { // log other errors to the logger TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); try { std::optional cause_chain_error; if (cause_chain != nullptr) { cause_chain_error = tvm::ffi::details::ObjectUnsafe::ObjectRefFromObjectPtr( tvm::ffi::details::ObjectUnsafe::ObjectPtrFromUnowned( static_cast(cause_chain))); } std::optional extra_context_ref; if (extra_context != nullptr) { extra_context_ref = tvm::ffi::details::ObjectUnsafe::ObjectRefFromObjectPtr( tvm::ffi::details::ObjectUnsafe::ObjectPtrFromUnowned( static_cast(extra_context))); } tvm::ffi::Error error(std::string(kind->data, kind->size), std::string(message->data, message->size), std::string(backtrace->data, backtrace->size), std::move(cause_chain_error), std::move(extra_context_ref)); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(error)); return 0; } catch (const std::bad_alloc& e) { return -1; } TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIErrorCreateWithCauseAndExtraContext); } tvm-ffi-0.1.12/src/ffi/extra/000077500000000000000000000000001521067262500156275ustar00rootroot00000000000000tvm-ffi-0.1.12/src/ffi/extra/buffer_stream.h000066400000000000000000000077201521067262500206320ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file buffer_stream.h * \brief Internal minimal stream helper to read from a buffer. */ #ifndef TVM_FFI_EXTRA_BUFFER_STREAM_H_ #define TVM_FFI_EXTRA_BUFFER_STREAM_H_ #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Lightweight stream helper to read from a buffer. */ class BufferInStream { public: /*! * \brief constructor * \param p_buffer the head pointer of the memory region. * \param buffer_size the size of the memorybuffer */ BufferInStream(const void* data, size_t size) : data_(reinterpret_cast(data)), size_(size) {} /*! * \brief Reads raw from stream. * \param ptr pointer to the data to be read * \param size the size of the data to be read * \return the number of bytes read */ size_t Read(void* ptr, size_t size) { size_t nread = std::min(size_ - curr_ptr_, size); if (nread != 0) std::memcpy(ptr, data_ + curr_ptr_, nread); curr_ptr_ += nread; return nread; } /*! * \brief Reads arithmetic data from stream in endian-aware manner. * \param data data to be read * \tparam T the data type to be read * \return whether the read was successful */ template >> bool Read(T* data) { bool ret = Read(static_cast(data), sizeof(T)) == sizeof(T); // NOLINT(*) if (!TVM_FFI_IO_NO_ENDIAN_SWAP) { ByteSwap(static_cast(data), sizeof(T), 1); } return ret; } /*! * \brief Reads an array of data from stream in endian-aware manner. * \param data data to be read * \param size the size of the data to be read * \return whether the read was successful */ template >> bool ReadArray(T* data, size_t size) { bool ret = this->Read(static_cast(data), sizeof(T) * size) == sizeof(T) * size; // NOLINT(*) if (!TVM_FFI_IO_NO_ENDIAN_SWAP) { ByteSwap(data, sizeof(T), size); } return ret; } /*! * \brief Reads a string from stream. * \param data data to be read * \return whether the read was successful */ bool Read(std::string* data) { // use uint64_t to ensure platform independent size uint64_t size = 0; if (!this->Read(&size)) return false; data->resize(size); if (!this->Read(data->data(), size)) return false; return true; } /*! * \brief Reads a vector of data from stream in endian-aware manner. * \param data data to be read * \return whether the read was successful */ template >> bool Read(std::vector* data) { uint64_t size = 0; if (!this->Read(&size)) return false; data->resize(size); return this->ReadArray(data->data(), size); } private: /*! \brief in memory buffer */ const char* data_; /*! \brief size of the buffer */ size_t size_; /*! \brief current pointer */ size_t curr_ptr_{0}; }; // class BytesInStream } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_BUFFER_STREAM_H_ tvm-ffi-0.1.12/src/ffi/extra/dataclass.cc000066400000000000000000002231221521067262500200770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file src/ffi/extra/dataclass.cc * \brief Reflection-based dataclass operations: * deep copy, repr printing, recursive hash, recursive compare. */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "../object_internal.h" namespace tvm { namespace ffi { namespace refl = ::tvm::ffi::reflection; // ============================================================================ // Shared utilities, CRTP base, and all operation classes // ============================================================================ namespace { // ---------- Shared traversal utilities ---------- /*! \brief Maximum traversal stack depth for iterative traversals. */ constexpr size_t kMaxTraversalStackDepth = 1 << 20; /*! \brief Return true for Str or SmallStr type indices. */ bool IsStringType(int32_t ti) { return ti == TypeIndex::kTVMFFIStr || ti == TypeIndex::kTVMFFISmallStr; } /*! \brief Return true for Bytes or SmallBytes type indices. */ bool IsBytesType(int32_t ti) { return ti == TypeIndex::kTVMFFIBytes || ti == TypeIndex::kTVMFFISmallBytes; } /*! \brief Extract raw pointer and length from a String or SmallStr value. */ void GetStringData(const Any& val, const TVMFFIAny* data, int32_t ti, const char** out_ptr, size_t* out_len) { if (ti == TypeIndex::kTVMFFISmallStr) { *out_ptr = data->v_bytes; *out_len = data->small_str_len; } else { const auto* obj = details::AnyUnsafe::CopyFromAnyViewAfterCheck(val); *out_ptr = obj->data; *out_len = obj->size; } } /*! \brief Extract raw pointer and length from a Bytes or SmallBytes value. */ void GetBytesData(const Any& val, const TVMFFIAny* data, int32_t ti, const char** out_ptr, size_t* out_len) { if (ti == TypeIndex::kTVMFFISmallBytes) { *out_ptr = data->v_bytes; *out_len = data->small_str_len; } else { const auto* obj = details::AnyUnsafe::CopyFromAnyViewAfterCheck(val); *out_ptr = obj->data; *out_len = obj->size; } } // ---------- CRTP base for object-graph DFS ---------- /*! * \brief Common frame base for single-value DFS traversals (hash, repr, copy). */ struct FrameBase { enum Kind : uint8_t { kSequence, kMap, kObject }; Kind kind; int32_t type_index; const Object* obj; std::vector children; std::vector field_infos; // kObject only size_t child_idx = 0; size_t container_size = 0; size_t NumChildren() const { return children.size(); } }; /*! * \brief CRTP base that owns the iterative DFS RunLoop and container dispatch. * * \tparam Derived Concrete operation class. * \tparam FrameT Frame type (must provide child_idx and NumChildren()). * \tparam ResultT Result type produced per frame. */ template class ObjectGraphDFS { friend Derived; ObjectGraphDFS() = default; protected: std::vector stack_; Derived& self() { return static_cast(*this); } ResultT RunLoop() { while (!stack_.empty()) { auto& f = stack_.back(); bool pushed = false; while (f.child_idx < f.NumChildren()) { size_t idx = f.child_idx++; ResultT r{}; if (self().TryVisitChild(f, idx, &r)) { if (self().FeedChild(f, r)) { return self().OnTerminate(std::move(r)); } } else { auto maybe = self().PushChildFrame(f, idx); if (maybe.has_value()) { if (self().FeedChild(f, *maybe)) { return self().OnTerminate(std::move(*maybe)); } } else { pushed = true; break; } } } if (pushed) continue; ResultT result = self().FinalizeFrame(f); self().OnFrameComplete(f); stack_.pop_back(); if (stack_.empty()) return result; if (self().FeedChild(stack_.back(), result)) { return self().OnTerminate(std::move(result)); } } TVM_FFI_UNREACHABLE(); } void EnumerateChildren(FrameBase& frame, const Any& value, const Object* obj, int32_t ti) { using details::AnyUnsafe; switch (ti) { case TypeIndex::kTVMFFIArray: { auto seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(value); frame.kind = FrameBase::kSequence; frame.container_size = seq.size(); frame.children.reserve(seq.size()); for (const auto& elem : seq) frame.children.push_back(elem); return; } case TypeIndex::kTVMFFIList: { auto seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(value); frame.kind = FrameBase::kSequence; frame.container_size = seq.size(); frame.children.reserve(seq.size()); for (const auto& elem : seq) frame.children.push_back(elem); return; } case TypeIndex::kTVMFFIMap: { auto map = AnyUnsafe::CopyFromAnyViewAfterCheck>(value); frame.kind = FrameBase::kMap; frame.container_size = map.size(); frame.children.reserve(map.size() * 2); for (const auto& kv : map) { frame.children.push_back(kv.first); frame.children.push_back(kv.second); } return; } case TypeIndex::kTVMFFIDict: { auto map = AnyUnsafe::CopyFromAnyViewAfterCheck>(value); frame.kind = FrameBase::kMap; frame.container_size = map.size(); frame.children.reserve(map.size() * 2); for (const auto& kv : map) { frame.children.push_back(kv.first); frame.children.push_back(kv.second); } return; } default: { frame.kind = FrameBase::kObject; const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti); uint32_t skip = self().GetFieldSkipMask(); refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { if (finfo->flags & skip) return; refl::FieldGetter getter(finfo); frame.children.push_back(getter(obj)); frame.field_infos.push_back(finfo); }); frame.container_size = frame.field_infos.size(); return; } } } void PushFrame(const Any& value) { if (this->stack_.size() >= kMaxTraversalStackDepth) { TVM_FFI_THROW(ValueError) << "ObjectGraphDFS: maximum stack depth (" << kMaxTraversalStackDepth << ") exceeded"; } const Object* obj = static_cast(value.as()); int32_t ti = obj->type_index(); self().OnEnter(obj); FrameT frame; frame.obj = obj; frame.type_index = ti; EnumerateChildren(frame, value, obj, ti); self().OnFrameInit(frame); this->stack_.push_back(std::move(frame)); } }; // ---------- Deep Copy ---------- struct CopyFrame : FrameBase { Any copy; std::vector resolved; // Array/Map: accumulated resolved children bool is_key = true; // Dict/Map: tracking key vs value Any current_key; // Dict: last resolved key size_t feed_idx = 0; // Object: field setter index }; /*! * \brief Iterative DFS deep copier (CRTP-based). * * - Mutable containers (List/Dict) and reflected objects register their copy * in copy_map_ before resolving children (enables cyclic back-references). * - Immutable containers (Array/Map) mark in_progress_ and build from resolved * children in FinalizeFrame. * - Deferred fixup pass replaces stale placeholder references in mutable * containers after all copies are fully constructed. */ class ObjectDeepCopier : public ObjectGraphDFS { public: explicit ObjectDeepCopier(refl::TypeAttrColumn* column) : column_(column) {} Any Run(const Any& value) { if (value.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) return value; Any result; if (TryCopyImmediate(value, &result)) { // Immediate (no children) } else { this->PushFrame(value); result = this->RunLoop(); } if (has_deferred_) { FixupDeferredReferences(); } return result; } // ---------- CRTP customization points ---------- uint32_t GetFieldSkipMask() { return 0; } void OnEnter(const Object* obj) { // Tracking is done in OnFrameInit per container type. } void OnFrameInit(CopyFrame& f) { int32_t ti = f.type_index; switch (ti) { case TypeIndex::kTVMFFIArray: { in_progress_.insert(f.obj); break; } case TypeIndex::kTVMFFIList: { List new_list; new_list.reserve(static_cast(f.container_size)); f.copy = new_list; copy_map_[f.obj] = f.copy; break; } case TypeIndex::kTVMFFIMap: { in_progress_.insert(f.obj); break; } case TypeIndex::kTVMFFIDict: { Dict new_dict; f.copy = new_dict; copy_map_[f.obj] = f.copy; break; } default: { // Reflected object: shallow-copy and register const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti); TVM_FFI_ICHECK((*column_)[ti] != nullptr) << "Cannot deep copy object of type \"" << std::string_view(type_info->type_key.data, type_info->type_key.size) << "\" because it is not copy-constructible"; Function copy_fn = (*column_)[ti].cast(); f.copy = copy_fn(f.obj); copy_map_[f.obj] = f.copy; break; } } } bool TryVisitChild(CopyFrame& f, size_t idx, Any* out) { return TryCopyImmediate(f.children[idx], out); } std::optional PushChildFrame(CopyFrame& f, size_t idx) { this->PushFrame(f.children[idx]); return std::nullopt; } bool FeedChild(CopyFrame& f, Any resolved) { int32_t ti = f.type_index; switch (ti) { case TypeIndex::kTVMFFIList: { f.copy.cast>().push_back(resolved); break; } case TypeIndex::kTVMFFIArray: { f.resolved.push_back(std::move(resolved)); break; } case TypeIndex::kTVMFFIDict: { if (f.is_key) { f.current_key = std::move(resolved); f.is_key = false; } else { f.copy.cast>().Set(f.current_key, resolved); f.is_key = true; } break; } case TypeIndex::kTVMFFIMap: { f.resolved.push_back(std::move(resolved)); break; } default: { // Reflected object: set field if changed const Any& original = f.children[f.feed_idx]; if (!original.same_as(resolved)) { refl::FieldSetter setter(f.field_infos[f.feed_idx]); setter(f.copy.as(), resolved); } f.feed_idx++; break; } } return false; } Any FinalizeFrame(CopyFrame& f) { int32_t ti = f.type_index; if (ti == TypeIndex::kTVMFFIArray) { Array new_arr; new_arr.reserve(static_cast(f.resolved.size())); for (const auto& elem : f.resolved) { new_arr.push_back(elem); } copy_map_[f.obj] = new_arr; return new_arr; } if (ti == TypeIndex::kTVMFFIMap) { Map new_map; for (size_t i = 0; i + 1 < f.resolved.size(); i += 2) { new_map.Set(f.resolved[i], f.resolved[i + 1]); } copy_map_[f.obj] = new_map; return new_map; } return f.copy; } void OnFrameComplete(CopyFrame& f) { int32_t ti = f.type_index; if (ti == TypeIndex::kTVMFFIArray || ti == TypeIndex::kTVMFFIMap) { in_progress_.erase(f.obj); } } Any OnTerminate(Any r) { return r; } private: refl::TypeAttrColumn* column_; std::unordered_map copy_map_; std::unordered_set in_progress_; bool has_deferred_ = false; bool TryCopyImmediate(const Any& value, Any* out) { if (value.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) { *out = value; return true; } const Object* obj = value.as(); if (obj == nullptr) { *out = value; return true; } // Already copied auto it = copy_map_.find(obj); if (it != copy_map_.end()) { *out = it->second; return true; } // In-progress immutable container: return original as placeholder if (in_progress_.count(obj)) { has_deferred_ = true; *out = value; return true; } int32_t ti = obj->type_index(); // Immutable leaf objects if (ti == TypeIndex::kTVMFFIStr || ti == TypeIndex::kTVMFFIBytes || ti == TypeIndex::kTVMFFIShape) { *out = value; return true; } return false; } void FixupDeferredReferences() { for (auto& [orig_ptr, copy_any] : copy_map_) { const Object* copy_obj = copy_any.as(); if (!copy_obj) continue; int32_t ti = copy_obj->type_index(); if (ti == TypeIndex::kTVMFFIList) { FixupList(copy_any); } else if (ti == TypeIndex::kTVMFFIDict) { FixupDict(copy_any); } else if (ti >= TypeIndex::kTVMFFIStaticObjectEnd) { FixupObject(copy_obj, ti); } } } void FixupList(const Any& list_any) { List list = list_any.cast>(); int64_t n = static_cast(list.size()); for (int64_t i = 0; i < n; ++i) { const Any& elem = list[i]; if (elem.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) continue; const Object* elem_obj = elem.as(); if (!elem_obj) continue; auto it = copy_map_.find(elem_obj); if (it != copy_map_.end()) { list.Set(i, it->second); } } } void FixupDict(const Any& dict_any) { Dict dict = dict_any.cast>(); const DictObj* dict_obj = dict_any.as(); // Collect entries that need key or value fixup. // old_key, new_key, new_val std::vector> fixups; for (const auto& [k, v] : *dict_obj) { Any new_key = k; Any new_val = v; bool changed = false; if (k.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) { const Object* k_obj = k.as(); if (k_obj) { auto it = copy_map_.find(k_obj); if (it != copy_map_.end()) { new_key = it->second; changed = true; } } } if (v.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) { const Object* v_obj = v.as(); if (v_obj) { auto it = copy_map_.find(v_obj); if (it != copy_map_.end()) { new_val = it->second; changed = true; } } } if (changed) { fixups.emplace_back(k, new_key, new_val); } } for (auto& [old_key, new_key, new_val] : fixups) { dict.erase(old_key); dict.Set(new_key, new_val); } } void FixupObject(const Object* copy_obj, int32_t ti) { const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti); refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { refl::FieldGetter getter(finfo); Any field_val = getter(copy_obj); if (field_val.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) return; const Object* field_obj = field_val.as(); if (!field_obj) return; auto it = copy_map_.find(field_obj); if (it != copy_map_.end() && !it->second.same_as(field_val)) { refl::FieldSetter setter(finfo); setter(copy_obj, it->second); } }); } }; // ---------- Repr helpers ---------- /*! * \brief Convert a DLDeviceType to a short name string. */ const char* DeviceTypeName(int device_type) { switch (device_type) { case kDLCPU: return "cpu"; case kDLCUDA: return "cuda"; case kDLCUDAHost: return "cuda_host"; case kDLOpenCL: return "opencl"; case kDLVulkan: return "vulkan"; case kDLMetal: return "metal"; case kDLVPI: return "vpi"; case kDLROCM: return "rocm"; case kDLROCMHost: return "rocm_host"; case kDLExtDev: return "ext_dev"; case kDLCUDAManaged: return "cuda_managed"; case kDLOneAPI: return "oneapi"; case kDLWebGPU: return "webgpu"; case kDLHexagon: return "hexagon"; case kDLMAIA: return "maia"; case kDLTrn: return "trn"; default: return "unknown"; } } /*! * \brief Format a DLDevice as "device_name:device_id". */ std::string DeviceToString(DLDevice device) { std::ostringstream os; os << DeviceTypeName(device.device_type) << ":" << device.device_id; return os.str(); } /*! * \brief Format raw bytes as a Python-style bytes literal: b"...". */ std::string FormatBytes(const char* data, size_t size) { std::ostringstream os; os << "b\""; for (size_t i = 0; i < size; ++i) { unsigned char c = static_cast(data[i]); if (c >= 32 && c < 127 && c != '\"' && c != '\\') { os << static_cast(c); } else { os << "\\x" << std::hex << std::setw(2) << std::setfill('0') << static_cast(c); } } os << "\""; return os.str(); } /*! * \brief Format an object address as a hex string. */ std::string AddressStr(const Object* obj) { std::ostringstream os; os << "0x" << std::hex << reinterpret_cast(obj); return os.str(); } struct ReprFrame : FrameBase { std::string header; std::vector child_reprs; }; /*! * \brief Iterative DFS-based repr printer (CRTP-based). * * Uses ObjectGraphDFS to iterate the object graph with an explicit stack. * Handles String, Bytes, Tensor, Shape, Array, List, Map, Dict formatting * directly (no registered built-in hooks). * Custom __ffi_repr__ hooks on user types are still supported. * * Address display is controlled by the TVM_FFI_REPR_WITH_ADDR environment variable. */ class ReprPrinter : public ObjectGraphDFS { public: String Run(const Any& value) { const char* env = std::getenv("TVM_FFI_REPR_WITH_ADDR"); show_addr_ = env != nullptr && std::string_view(env) == "1"; std::string result; if (TryReprImmediate(value, &result)) return String(result); this->PushFrame(value); return String(this->RunLoop()); } // ---------- CRTP customization points ---------- uint32_t GetFieldSkipMask() { return kTVMFFIFieldFlagBitMaskReprOff; } void OnEnter(const Object* obj) { state_[obj] = State::kInProgress; } void OnFrameInit(ReprFrame& f) { int32_t ti = f.type_index; const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti); std::string type_key(type_info->type_key.data, type_info->type_key.size); f.header = show_addr_ ? (type_key + "@" + AddressStr(f.obj)) : type_key; } bool TryVisitChild(ReprFrame& f, size_t idx, std::string* out) { return TryReprImmediate(f.children[idx], out); } std::optional PushChildFrame(ReprFrame& f, size_t idx) { this->PushFrame(f.children[idx]); return std::nullopt; } bool FeedChild(ReprFrame& f, std::string repr) { f.child_reprs.push_back(std::move(repr)); return false; } std::string FinalizeFrame(ReprFrame& f) { std::string result; switch (f.kind) { case FrameBase::kSequence: { if (f.type_index == TypeIndex::kTVMFFIArray) { result = "("; for (size_t i = 0; i < f.child_reprs.size(); ++i) { if (i > 0) result += ", "; result += f.child_reprs[i]; } if (f.child_reprs.size() == 1) result += ","; result += ")"; } else { result = "["; for (size_t i = 0; i < f.child_reprs.size(); ++i) { if (i > 0) result += ", "; result += f.child_reprs[i]; } result += "]"; } break; } case FrameBase::kMap: { result = "{"; for (size_t i = 0; i + 1 < f.child_reprs.size(); i += 2) { if (i > 0) result += ", "; result += f.child_reprs[i] + ": " + f.child_reprs[i + 1]; } result += "}"; break; } case FrameBase::kObject: { if (f.child_reprs.empty()) { result = f.header; } else { result = f.header + "("; std::unordered_set seen_names; bool first = true; size_t repr_idx = 0; for (size_t i = 0; i < f.field_infos.size(); ++i) { std::string_view name(f.field_infos[i]->name.data, f.field_infos[i]->name.size); if (!seen_names.insert(name).second) { repr_idx++; continue; } if (!first) result += ", "; first = false; result += std::string(name) + "=" + f.child_reprs[repr_idx++]; } result += ")"; } break; } } if (show_addr_ && (f.type_index == TypeIndex::kTVMFFIArray || f.type_index == TypeIndex::kTVMFFIList || f.type_index == TypeIndex::kTVMFFIMap || f.type_index == TypeIndex::kTVMFFIDict)) { result += "@" + AddressStr(f.obj); } state_[f.obj] = State::kDone; repr_cache_[f.obj] = result; return result; } void OnFrameComplete(ReprFrame&) {} std::string OnTerminate(std::string r) { return r; } private: enum class State : int8_t { kNotVisited = 0, kInProgress = 1, kDone = 2 }; std::unordered_map state_; std::unordered_map repr_cache_; bool show_addr_ = false; // ---------- Immediate repr ---------- bool TryReprImmediate(const Any& value, std::string* out) { int32_t ti = value.type_index(); switch (ti) { case TypeIndex::kTVMFFINone: *out = "None"; return true; case TypeIndex::kTVMFFIBool: *out = value.cast() ? "True" : "False"; return true; case TypeIndex::kTVMFFIInt: *out = std::to_string(value.cast()); return true; case TypeIndex::kTVMFFIFloat: { std::ostringstream os; os << value.cast(); *out = os.str(); return true; } case TypeIndex::kTVMFFIDataType: { String s = DLDataTypeToString(value.cast()); *out = std::string(s.data(), s.size()); return true; } case TypeIndex::kTVMFFIDevice: { *out = DeviceToString(value.cast()); return true; } default: break; } if (ti == TypeIndex::kTVMFFISmallStr) { String s = value.cast(); String escaped = EscapedStringPy(s); *out = std::string(escaped.data(), escaped.size()); return true; } if (ti == TypeIndex::kTVMFFISmallBytes) { Bytes b = value.cast(); *out = FormatBytes(b.data(), b.size()); return true; } if (ti < TypeIndex::kTVMFFIStaticObjectBegin) { *out = value.GetTypeKey(); return true; } const Object* obj = static_cast(value.as()); if (obj == nullptr) { *out = "None"; return true; } // Check cache / cycle auto it = state_.find(obj); if (it != state_.end()) { if (it->second == State::kDone) { *out = repr_cache_[obj]; return true; } if (it->second == State::kInProgress) { *out = show_addr_ ? ("...@" + AddressStr(obj)) : "..."; return true; } } // Built-in sentinel singletons — pointer-identity dispatch ahead of any // type-keyed lookups so the generic ``ffi.Object`` framing is skipped. static const Object* missing_ptr = GetMissingObject().get(); static const Object* kwargs_ptr = GetKwargsObject().get(); if (obj == missing_ptr) { *out = ""; return true; } if (obj == kwargs_ptr) { *out = ""; return true; } // String/Bytes on heap if (ti == TypeIndex::kTVMFFIStr) { String s = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); String escaped = EscapedStringPy(s); *out = std::string(escaped.data(), escaped.size()); return true; } if (ti == TypeIndex::kTVMFFIBytes) { Bytes b = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); *out = FormatBytes(b.data(), b.size()); return true; } // Tensor if (ti == TypeIndex::kTVMFFITensor) { const TensorObj* t = value.as(); std::ostringstream os; os << DLDataTypeToString(t->dtype) << "["; for (int i = 0; i < t->ndim; ++i) { if (i > 0) os << ", "; os << t->shape[i]; } os << "]@" << DeviceToString(t->device); if (show_addr_) os << "@" << AddressStr(obj); *out = os.str(); return true; } // Shape if (ti == TypeIndex::kTVMFFIShape) { const ShapeObj* s = value.as(); std::ostringstream os; os << "Shape("; for (size_t i = 0; i < s->size; ++i) { if (i > 0) os << ", "; os << s->data[i]; } os << ")"; *out = os.str(); return true; } // Custom __ffi_repr__ hook static refl::TypeAttrColumn repr_column(refl::type_attr::kRepr); AnyView custom_repr = repr_column[ti]; if (custom_repr != nullptr) { state_[obj] = State::kInProgress; Function repr_fn = custom_repr.cast(); Function fn_repr = CreateFnRepr(); String r = repr_fn(obj, fn_repr).cast(); std::string result(r.data(), r.size()); if (show_addr_ && (ti == TypeIndex::kTVMFFIArray || ti == TypeIndex::kTVMFFIList || ti == TypeIndex::kTVMFFIMap || ti == TypeIndex::kTVMFFIDict)) { result += "@" + AddressStr(obj); } state_[obj] = State::kDone; repr_cache_[obj] = result; *out = result; return true; } // Default repr for EnumObj subclasses: ``.<_name>``. Reached only // when no user-registered ``__ffi_repr__`` hook has claimed this type, so // explicit repr overrides on specific enum subclasses still take precedence. if (obj->IsInstance()) { const EnumObj* enum_obj = static_cast(obj); const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti); std::string result(type_info->type_key.data, type_info->type_key.size); result += '.'; result.append(enum_obj->_name.data(), enum_obj->_name.size()); if (show_addr_) result += "@" + AddressStr(obj); state_[obj] = State::kDone; repr_cache_[obj] = result; *out = result; return true; } // Needs a frame return false; } // ---------- Custom hook callback ---------- Function CreateFnRepr() { return Function::FromTyped([this](AnyView value) -> String { Any v(value); std::string result; if (TryReprImmediate(v, &result)) return String(result); std::vector saved; saved.swap(this->stack_); this->PushFrame(v); result = this->RunLoop(); this->stack_.swap(saved); return String(result); }); } }; // ---------- Section 3: Recursive Hash ---------- /*! * \brief Iterative reflection-based recursive hasher (CRTP-based). * * Uses ObjectGraphDFS to iterate the object graph with an explicit stack. * Supports custom __ffi_hash__ hooks. * * Computes a deterministic hash consistent with RecursiveEq: * RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b) */ struct HashFrame : FrameBase { uint64_t hash = 0; size_t seq_index = 0; std::vector entry_hashes; bool in_key = true; uint64_t key_hash = 0; }; class RecursiveHasher : public ObjectGraphDFS { public: uint64_t HashAny(const Any& value) { uint64_t h; if (TryHashImmediate(value, &h)) return h; this->PushFrame(value); return this->RunLoop(); } // ---------- CRTP customization points ---------- uint32_t GetFieldSkipMask() { return kTVMFFIFieldFlagBitMaskHashOff | kTVMFFIFieldFlagBitMaskCompareOff; } void OnEnter(const Object* obj) { on_stack_.insert(obj); } void OnFrameInit(HashFrame& f) { const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(f.type_index); if (f.kind == FrameBase::kObject) { // Reflected objects: use type_key_hash alone (matches old behavior). // Containers already combine type_key_hash with container_size. f.hash = type_info->type_key_hash; } else { f.hash = details::StableHashCombine(type_info->type_key_hash, f.container_size); } if (f.kind == FrameBase::kMap) { f.entry_hashes.reserve(f.container_size); } } bool TryVisitChild(HashFrame& f, size_t idx, uint64_t* out) { return TryHashImmediate(f.children[idx], out); } std::optional PushChildFrame(HashFrame& f, size_t idx) { this->PushFrame(f.children[idx]); return std::nullopt; } bool FeedChild(HashFrame& f, uint64_t h) { switch (f.kind) { case FrameBase::kSequence: { f.hash = details::StableHashCombine(f.hash, details::StableHashCombine(h, f.seq_index++)); break; } case FrameBase::kMap: { if (f.in_key) { f.key_hash = h; f.in_key = false; } else { f.entry_hashes.push_back(details::StableHashCombine(f.key_hash, h)); f.in_key = true; } break; } case FrameBase::kObject: { f.hash = details::StableHashCombine(f.hash, h); break; } } return false; // never terminates early } uint64_t FinalizeFrame(HashFrame& f) { if (f.kind == FrameBase::kMap) { std::sort(f.entry_hashes.begin(), f.entry_hashes.end()); for (uint64_t eh : f.entry_hashes) { f.hash = details::StableHashCombine(f.hash, eh); } } return f.hash; } void OnFrameComplete(HashFrame& f) { if (f.obj != nullptr) { memo_[f.obj] = f.hash; on_stack_.erase(f.obj); } } uint64_t OnTerminate(uint64_t r) { return r; } private: std::unordered_set on_stack_; std::unordered_map memo_; // ---------- Immediate (non-recursive) hashing ---------- bool TryHashImmediate(const Any& value, uint64_t* out) { using details::AnyUnsafe; const TVMFFIAny* data = AnyUnsafe::TVMFFIAnyPtrFromAny(value); int32_t ti = data->type_index; // None if (ti == TypeIndex::kTVMFFINone) { *out = details::StableHashCombine(uint64_t{0}, uint64_t{0}); return true; } // String (Str/SmallStr cross-variant) if (IsStringType(ti)) { *out = HashString(value, data, ti); return true; } // Bytes (Bytes/SmallBytes cross-variant) if (IsBytesType(ti)) { *out = HashBytes(value, data, ti); return true; } // POD types if (ti < TypeIndex::kTVMFFIStaticObjectBegin) { *out = HashPOD(value, data, ti); return true; } // Object types const Object* obj = static_cast(value.as()); if (obj == nullptr) { *out = details::StableHashCombine(uint64_t{0}, uint64_t{0}); return true; } // Return memoized hash if already fully hashed. auto memo_it = memo_.find(obj); if (memo_it != memo_.end()) { *out = memo_it->second; return true; } // Cycle detection: if on the call stack, return sentinel. if (on_stack_.count(obj)) { *out = TVMFFIGetTypeInfo(obj->type_index())->type_key_hash; return true; } // Shape is always immediate (no children) if (ti == TypeIndex::kTVMFFIShape) { uint64_t h = HashShape(AnyUnsafe::CopyFromAnyViewAfterCheck(value)); memo_[obj] = h; *out = h; return true; } // Check for custom __ffi_hash__ hook static refl::TypeAttrColumn hash_column(refl::type_attr::kHash); AnyView custom = hash_column[obj->type_index()]; if (custom != nullptr) { on_stack_.insert(obj); Function hook = custom.cast(); Function fn_hash = CreateFnHash(); int64_t r = hook(obj, fn_hash).cast(); uint64_t h = static_cast(r); memo_[obj] = h; on_stack_.erase(obj); *out = h; return true; } // For reflected types (not built-in containers), error if the type has // __ffi_eq__ or __ffi_compare__ but no __ffi_hash__. if (ti >= TypeIndex::kTVMFFIStaticObjectEnd) { static refl::TypeAttrColumn eq_column(refl::type_attr::kEq); static refl::TypeAttrColumn cmp_column(refl::type_attr::kCompare); if (eq_column[obj->type_index()] != nullptr || cmp_column[obj->type_index()] != nullptr) { const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(ti); TVM_FFI_THROW(ValueError) << "RecursiveHash: type '" << String(info->type_key) << "' defines __ffi_eq__ or __ffi_compare__ but not __ffi_hash__. " << "Add a __ffi_hash__ hook to maintain the invariant " << "RecursiveEq(a,b) => RecursiveHash(a)==RecursiveHash(b)."; } } // Needs a frame (sequence, map, or reflected object) return false; } // ---------- Custom __ffi_hash__ callback ---------- Function CreateFnHash() { return Function::FromTyped([this](AnyView value) -> int64_t { Any v(value); uint64_t h; if (TryHashImmediate(v, &h)) return static_cast(h); std::vector saved; saved.swap(this->stack_); this->PushFrame(v); h = this->RunLoop(); this->stack_.swap(saved); return static_cast(h); }); } // ---------- POD hashing ---------- static uint64_t HashPOD(const Any& value, const TVMFFIAny* data, int32_t ti) { switch (ti) { case TypeIndex::kTVMFFIBool: { uint64_t v = data->v_int64 != 0 ? 1 : 0; return details::StableHashCombine(static_cast(ti), v); } case TypeIndex::kTVMFFIInt: { return details::StableHashCombine(static_cast(ti), static_cast(data->v_int64)); } case TypeIndex::kTVMFFIFloat: { double v = data->v_float64; uint64_t bits; if (std::isnan(v)) { double canonical = std::numeric_limits::quiet_NaN(); std::memcpy(&bits, &canonical, sizeof(bits)); } else if (v == 0.0) { double pos_zero = 0.0; std::memcpy(&bits, &pos_zero, sizeof(bits)); } else { std::memcpy(&bits, &v, sizeof(bits)); } return details::StableHashCombine(static_cast(ti), bits); } case TypeIndex::kTVMFFIDataType: { DLDataType dt = data->v_dtype; uint64_t h = details::StableHashCombine(static_cast(ti), static_cast(dt.code)); h = details::StableHashCombine(h, static_cast(dt.bits)); h = details::StableHashCombine(h, static_cast(dt.lanes)); return h; } case TypeIndex::kTVMFFIDevice: { DLDevice dev = data->v_device; uint64_t h = details::StableHashCombine(static_cast(ti), static_cast(dev.device_type)); h = details::StableHashCombine(h, static_cast(dev.device_id)); return h; } default: { return details::StableHashCombine(static_cast(ti), static_cast(data->v_uint64)); } } } // ---------- String hashing ---------- static uint64_t HashString(const Any& value, const TVMFFIAny* data, int32_t ti) { const char* ptr; size_t len; GetStringData(value, data, ti, &ptr, &len); return details::StableHashCombine(static_cast(TypeIndex::kTVMFFIStr), details::StableHashBytes(ptr, len)); } // ---------- Bytes hashing ---------- static uint64_t HashBytes(const Any& value, const TVMFFIAny* data, int32_t ti) { const char* ptr; size_t len; GetBytesData(value, data, ti, &ptr, &len); return details::StableHashCombine(static_cast(TypeIndex::kTVMFFIBytes), details::StableHashBytes(ptr, len)); } // ---------- Shape hashing ---------- static uint64_t HashShape(const Shape& shape) { uint64_t h = details::StableHashCombine(shape->GetTypeKeyHash(), shape.size()); for (int64_t dim : shape) { h = details::StableHashCombine(h, static_cast(dim)); } return h; } }; // ---------- Recursive Compare ---------- struct CompareFrame { enum Kind : uint8_t { kSequence, kMap, kObject }; Kind kind; std::vector> children; size_t child_idx = 0; size_t lhs_size = 0; size_t rhs_size = 0; const Object* lhs_obj = nullptr; const Object* rhs_obj = nullptr; size_t NumChildren() const { return children.size(); } }; /*! * \brief Iterative three-way recursive comparer (CRTP-based). * * Returns int32_t: -1 (lhs < rhs), 0 (equal), +1 (lhs > rhs). * Uses ObjectGraphDFS RunLoop with pair-based CompareFrame (not FrameBase). * Supports custom __ffi_eq__ and __ffi_compare__ hooks. */ class RecursiveComparer : public ObjectGraphDFS { public: explicit RecursiveComparer(bool eq_only) : eq_only_(eq_only) {} int32_t CompareAny(const Any& lhs, const Any& rhs) { int32_t cmp; if (TryCompareImmediate(lhs, rhs, &cmp)) return cmp; auto eager = PushPairFrame(lhs, rhs); if (eager.has_value()) return *eager; return this->RunLoop(); } // ---------- CRTP customization points ---------- uint32_t GetFieldSkipMask() { return kTVMFFIFieldFlagBitMaskCompareOff; } void OnEnter(const Object*) {} void OnFrameInit(CompareFrame&) {} bool TryVisitChild(CompareFrame& f, size_t idx, int32_t* out) { auto& [child_lhs, child_rhs] = f.children[idx]; return TryCompareImmediate(child_lhs, child_rhs, out); } std::optional PushChildFrame(CompareFrame& f, size_t idx) { auto& [child_lhs, child_rhs] = f.children[idx]; return PushPairFrame(child_lhs, child_rhs); } bool FeedChild(CompareFrame&, int32_t result) { return result != 0; } int32_t FinalizeFrame(CompareFrame& f) { if (f.kind == CompareFrame::kSequence) { if (f.lhs_size < f.rhs_size) return -1; if (f.lhs_size > f.rhs_size) return 1; } return 0; } void OnFrameComplete(CompareFrame& f) { if (f.lhs_obj != nullptr && f.rhs_obj != nullptr) { on_stack_.erase(std::make_pair(f.lhs_obj, f.rhs_obj)); } } int32_t OnTerminate(int32_t result) { return PropagateNonZero(result); } private: struct PairHash { size_t operator()(std::pair p) const { auto h1 = std::hash()(p.first); auto h2 = std::hash()(p.second); return h1 ^ (h2 * 0x9e3779b97f4a7c15ULL + 0x9e3779b9 + (h1 << 6) + (h1 >> 2)); } }; bool eq_only_; std::unordered_set, PairHash> on_stack_; // ---------- Immediate (non-recursive) comparison ---------- bool TryCompareImmediate(const Any& lhs, const Any& rhs, int32_t* out) { using details::AnyUnsafe; const TVMFFIAny* lhs_data = AnyUnsafe::TVMFFIAnyPtrFromAny(lhs); const TVMFFIAny* rhs_data = AnyUnsafe::TVMFFIAnyPtrFromAny(rhs); int32_t lti = lhs_data->type_index; int32_t rti = rhs_data->type_index; if (lti == TypeIndex::kTVMFFINone && rti == TypeIndex::kTVMFFINone) { *out = 0; return true; } if (lti == TypeIndex::kTVMFFINone) { *out = -1; return true; } if (rti == TypeIndex::kTVMFFINone) { *out = 1; return true; } if (IsStringType(lti) && IsStringType(rti)) { *out = CompareString(lhs, rhs, lhs_data, rhs_data, lti, rti); return true; } if (IsBytesType(lti) && IsBytesType(rti)) { *out = CompareBytes(lhs, rhs, lhs_data, rhs_data, lti, rti); return true; } if (lti != rti) { if (eq_only_) { *out = 1; return true; } TVM_FFI_THROW(TypeError) << "Cannot compare values of different types: " << lhs.GetTypeKey() << " vs " << rhs.GetTypeKey(); } if (lti < TypeIndex::kTVMFFIStaticObjectBegin) { *out = ComparePOD(lhs, rhs, lhs_data, rhs_data, lti); return true; } const Object* lhs_obj = static_cast(lhs.as()); const Object* rhs_obj = static_cast(rhs.as()); if (lhs_obj == rhs_obj) { *out = 0; return true; } if (lhs_obj == nullptr) { *out = -1; return true; } if (rhs_obj == nullptr) { *out = 1; return true; } if (lti == TypeIndex::kTVMFFIShape) { *out = CompareShape(AnyUnsafe::CopyFromAnyViewAfterCheck(lhs), AnyUnsafe::CopyFromAnyViewAfterCheck(rhs)); return true; } auto pair = std::make_pair(lhs_obj, rhs_obj); if (on_stack_.count(pair)) { if (eq_only_) { *out = 0; return true; } TVM_FFI_THROW(ValueError) << "RecursiveCompare: cyclic reference detected in ordering"; } if (lti >= TypeIndex::kTVMFFIStaticObjectEnd) { return TryCustomHook(lhs_obj, rhs_obj, out); } return false; } // ---------- Custom hook dispatch ---------- bool TryCustomHook(const Object* lhs_obj, const Object* rhs_obj, int32_t* out) { if (lhs_obj->type_index() != rhs_obj->type_index()) { if (eq_only_) { *out = 1; return true; } const TVMFFITypeInfo* lhs_info = TVMFFIGetTypeInfo(lhs_obj->type_index()); const TVMFFITypeInfo* rhs_info = TVMFFIGetTypeInfo(rhs_obj->type_index()); TVM_FFI_THROW(TypeError) << "Cannot compare objects of different types: " << String(lhs_info->type_key) << " vs " << String(rhs_info->type_key); } static refl::TypeAttrColumn eq_column(refl::type_attr::kEq); static refl::TypeAttrColumn cmp_column(refl::type_attr::kCompare); int32_t ti = lhs_obj->type_index(); AnyView custom_eq = eq_column[ti]; AnyView custom_cmp = cmp_column[ti]; if (eq_only_) { if (custom_eq != nullptr) { auto pair = std::make_pair(lhs_obj, rhs_obj); on_stack_.insert(pair); Function hook = custom_eq.cast(); Function fn_eq = CreateFnEq(); bool result = hook(lhs_obj, rhs_obj, fn_eq).cast(); on_stack_.erase(pair); *out = result ? 0 : 1; return true; } if (custom_cmp != nullptr) { auto pair = std::make_pair(lhs_obj, rhs_obj); on_stack_.insert(pair); Function hook = custom_cmp.cast(); Function fn_cmp = CreateFnCompare(); int32_t result = hook(lhs_obj, rhs_obj, fn_cmp).cast(); on_stack_.erase(pair); *out = result; return true; } } else { if (custom_cmp != nullptr) { auto pair = std::make_pair(lhs_obj, rhs_obj); on_stack_.insert(pair); Function hook = custom_cmp.cast(); Function fn_cmp = CreateFnCompare(); int32_t result = hook(lhs_obj, rhs_obj, fn_cmp).cast(); on_stack_.erase(pair); *out = result; return true; } } return false; } // ---------- Pair-based frame creation ---------- /*! * \brief Push a pair frame or return an eager result (map mismatch). */ std::optional PushPairFrame(const Any& lhs, const Any& rhs) { if (this->stack_.size() >= kMaxTraversalStackDepth) { TVM_FFI_THROW(ValueError) << "RecursiveCompare: maximum stack depth (" << kMaxTraversalStackDepth << ") exceeded"; } using details::AnyUnsafe; const TVMFFIAny* lhs_data = AnyUnsafe::TVMFFIAnyPtrFromAny(lhs); int32_t lti = lhs_data->type_index; const Object* lhs_obj = static_cast(lhs.as()); const Object* rhs_obj = static_cast(rhs.as()); auto pair = std::make_pair(lhs_obj, rhs_obj); on_stack_.insert(pair); switch (lti) { case TypeIndex::kTVMFFIArray: { auto lhs_seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(lhs); auto rhs_seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(rhs); PushSequenceFrame(lhs_seq, rhs_seq, lhs_obj, rhs_obj); return std::nullopt; } case TypeIndex::kTVMFFIList: { auto lhs_seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(lhs); auto rhs_seq = AnyUnsafe::CopyFromAnyViewAfterCheck>(rhs); PushSequenceFrame(lhs_seq, rhs_seq, lhs_obj, rhs_obj); return std::nullopt; } case TypeIndex::kTVMFFIMap: { auto lhs_map = AnyUnsafe::CopyFromAnyViewAfterCheck>(lhs); auto rhs_map = AnyUnsafe::CopyFromAnyViewAfterCheck>(rhs); return PushMapFrame(lhs_map, rhs_map, lhs_obj, rhs_obj); } case TypeIndex::kTVMFFIDict: { auto lhs_map = AnyUnsafe::CopyFromAnyViewAfterCheck>(lhs); auto rhs_map = AnyUnsafe::CopyFromAnyViewAfterCheck>(rhs); return PushMapFrame(lhs_map, rhs_map, lhs_obj, rhs_obj); } default: { return PushObjectFrame(lhs_obj, rhs_obj); } } } template void PushSequenceFrame(const SeqType& lhs, const SeqType& rhs, const Object* lhs_obj, const Object* rhs_obj) { CompareFrame frame; frame.kind = CompareFrame::kSequence; frame.lhs_size = lhs.size(); frame.rhs_size = rhs.size(); frame.lhs_obj = lhs_obj; frame.rhs_obj = rhs_obj; size_t min_len = std::min(lhs.size(), rhs.size()); frame.children.reserve(min_len); for (size_t i = 0; i < min_len; ++i) { frame.children.emplace_back(lhs[i], rhs[i]); } this->stack_.push_back(std::move(frame)); } template std::optional PushMapFrame(const MapType& lhs, const MapType& rhs, const Object* lhs_obj, const Object* rhs_obj) { if (lhs.size() != rhs.size()) { on_stack_.erase(std::make_pair(lhs_obj, rhs_obj)); if (eq_only_) return 1; TVM_FFI_THROW(TypeError) << "Cannot order Map/Dict values"; } CompareFrame frame; frame.kind = CompareFrame::kMap; frame.lhs_size = lhs.size(); frame.rhs_size = rhs.size(); frame.lhs_obj = lhs_obj; frame.rhs_obj = rhs_obj; frame.children.reserve(lhs.size()); for (const auto& kv : lhs) { auto it = rhs.find(kv.first); if (it == rhs.end()) { on_stack_.erase(std::make_pair(lhs_obj, rhs_obj)); if (eq_only_) return 1; TVM_FFI_THROW(TypeError) << "Cannot order Map/Dict values"; } frame.children.emplace_back(kv.second, (*it).second); } this->stack_.push_back(std::move(frame)); return std::nullopt; } std::optional PushObjectFrame(const Object* lhs, const Object* rhs) { if (lhs->type_index() != rhs->type_index()) { auto pair = std::make_pair(lhs, rhs); on_stack_.erase(pair); if (eq_only_) return 1; const TVMFFITypeInfo* lhs_info = TVMFFIGetTypeInfo(lhs->type_index()); const TVMFFITypeInfo* rhs_info = TVMFFIGetTypeInfo(rhs->type_index()); TVM_FFI_THROW(TypeError) << "Cannot compare objects of different types: " << String(lhs_info->type_key) << " vs " << String(rhs_info->type_key); } const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(lhs->type_index()); CompareFrame frame; frame.kind = CompareFrame::kObject; frame.lhs_obj = lhs; frame.rhs_obj = rhs; refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { if (finfo->flags & kTVMFFIFieldFlagBitMaskCompareOff) return; refl::FieldGetter getter(finfo); frame.children.emplace_back(getter(lhs), getter(rhs)); }); this->stack_.push_back(std::move(frame)); return std::nullopt; } // ---------- Propagation ---------- int32_t PropagateNonZero(int32_t result) { if (!eq_only_) { for (auto it = this->stack_.rbegin(); it != this->stack_.rend(); ++it) { if (it->kind == CompareFrame::kMap) { CleanupAllFrames(); TVM_FFI_THROW(TypeError) << "Cannot order Map/Dict values"; } } } CleanupAllFrames(); return result; } void CleanupAllFrames() { for (auto& f : this->stack_) { if (f.lhs_obj != nullptr && f.rhs_obj != nullptr) { on_stack_.erase(std::make_pair(f.lhs_obj, f.rhs_obj)); } } this->stack_.clear(); } // ---------- Custom hook callbacks ---------- Function CreateFnEq() { return Function::FromTyped([this](AnyView l, AnyView r) -> bool { Any lhs(l), rhs(r); int32_t cmp; if (TryCompareImmediate(lhs, rhs, &cmp)) return cmp == 0; std::vector saved; saved.swap(this->stack_); auto eager = PushPairFrame(lhs, rhs); if (eager.has_value()) { this->stack_.swap(saved); return *eager == 0; } cmp = this->RunLoop(); this->stack_.swap(saved); return cmp == 0; }); } Function CreateFnCompare() { return Function::FromTyped([this](AnyView l, AnyView r) -> int32_t { Any lhs(l), rhs(r); int32_t cmp; if (TryCompareImmediate(lhs, rhs, &cmp)) return cmp; std::vector saved; saved.swap(this->stack_); auto eager = PushPairFrame(lhs, rhs); if (eager.has_value()) { this->stack_.swap(saved); return *eager; } cmp = this->RunLoop(); this->stack_.swap(saved); return cmp; }); } // ---------- POD comparison ---------- int32_t ComparePOD(const Any& lhs, const Any& rhs, const TVMFFIAny* lhs_data, const TVMFFIAny* rhs_data, int32_t ti) { switch (ti) { case TypeIndex::kTVMFFIBool: { bool a = lhs_data->v_int64 != 0; bool b = rhs_data->v_int64 != 0; return static_cast(a) - static_cast(b); } case TypeIndex::kTVMFFIInt: { int64_t a = lhs_data->v_int64; int64_t b = rhs_data->v_int64; if (a < b) return -1; if (a > b) return 1; return 0; } case TypeIndex::kTVMFFIFloat: { double a = lhs_data->v_float64; double b = rhs_data->v_float64; if (std::isnan(a) && std::isnan(b)) { if (eq_only_) return 0; TVM_FFI_THROW(TypeError) << "Cannot order NaN values"; } if (std::isnan(a) || std::isnan(b)) { if (eq_only_) return 1; TVM_FFI_THROW(TypeError) << "Cannot order NaN values"; } if (a < b) return -1; if (a > b) return 1; return 0; } case TypeIndex::kTVMFFIDataType: { DLDataType a = lhs_data->v_dtype; DLDataType b = rhs_data->v_dtype; if (a.code != b.code) return (a.code < b.code) ? -1 : 1; if (a.bits != b.bits) return (a.bits < b.bits) ? -1 : 1; if (a.lanes != b.lanes) return (a.lanes < b.lanes) ? -1 : 1; return 0; } case TypeIndex::kTVMFFIDevice: { DLDevice a = lhs_data->v_device; DLDevice b = rhs_data->v_device; if (a.device_type != b.device_type) return (a.device_type < b.device_type) ? -1 : 1; if (a.device_id != b.device_id) return (a.device_id < b.device_id) ? -1 : 1; return 0; } default: { if (lhs_data->zero_padding == rhs_data->zero_padding && lhs_data->v_int64 == rhs_data->v_int64) { return 0; } if (eq_only_) return 1; TVM_FFI_THROW(TypeError) << "Cannot order values of type " << lhs.GetTypeKey(); } } TVM_FFI_UNREACHABLE(); } // ---------- String comparison ---------- static int32_t CompareString(const Any& lhs, const Any& rhs, const TVMFFIAny* lhs_data, const TVMFFIAny* rhs_data, int32_t lti, int32_t rti) { const char* lhs_ptr; size_t lhs_len; const char* rhs_ptr; size_t rhs_len; GetStringData(lhs, lhs_data, lti, &lhs_ptr, &lhs_len); GetStringData(rhs, rhs_data, rti, &rhs_ptr, &rhs_len); return SignFromMemncmp(Bytes::memncmp(lhs_ptr, rhs_ptr, lhs_len, rhs_len)); } // ---------- Bytes comparison ---------- static int32_t CompareBytes(const Any& lhs, const Any& rhs, const TVMFFIAny* lhs_data, const TVMFFIAny* rhs_data, int32_t lti, int32_t rti) { const char* lhs_ptr; size_t lhs_len; const char* rhs_ptr; size_t rhs_len; GetBytesData(lhs, lhs_data, lti, &lhs_ptr, &lhs_len); GetBytesData(rhs, rhs_data, rti, &rhs_ptr, &rhs_len); return SignFromMemncmp(Bytes::memncmp(lhs_ptr, rhs_ptr, lhs_len, rhs_len)); } static int32_t SignFromMemncmp(int v) { if (v < 0) return -1; if (v > 0) return 1; return 0; } // ---------- Shape comparison ---------- static int32_t CompareShape(const Shape& lhs, const Shape& rhs) { size_t min_len = std::min(lhs.size(), rhs.size()); for (size_t i = 0; i < min_len; ++i) { if (lhs[i] < rhs[i]) return -1; if (lhs[i] > rhs[i]) return 1; } if (lhs.size() < rhs.size()) return -1; if (lhs.size() > rhs.size()) return 1; return 0; } }; // ---------- Python-defined type support ---------- /*! * \brief Deleter for objects whose layout is defined from Python via Field descriptors. * * For the "strong" phase, iterates all reflected fields and destructs * Any/ObjectRef values in-place (to release references). For the "weak" * phase, frees the underlying calloc'd memory. */ void PyClassDeleter(void* self_void, int flags) { TVMFFIObject* self = static_cast(self_void); if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) { const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(self->type_index); refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { void* field_addr = reinterpret_cast(self) + finfo->offset; int32_t ti = finfo->field_static_type_index; if (ti == TypeIndex::kTVMFFIAny) { // Any field: call destructor to release owned references reinterpret_cast(field_addr)->~Any(); } else if (ti >= TypeIndex::kTVMFFIStaticObjectBegin) { // ObjectRef field: call destructor to DecRef reinterpret_cast(field_addr)->~ObjectRef(); } // POD types (int, float, bool, etc.): no cleanup needed }); } if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) { std::free(self_void); } } /*! * \brief Generic field getter for Python-defined types. * * Reads a value of type T from the given field address and packs it into * a TVMFFIAny result. * * \tparam T The C++ type stored at the field address. */ template int PyClassFieldGetter(void* field, TVMFFIAny* result) { TVM_FFI_SAFE_CALL_BEGIN(); *result = details::AnyUnsafe::MoveAnyToTVMFFIAny(Any(*reinterpret_cast(field))); TVM_FFI_SAFE_CALL_END(); } /*! * \brief Return the TVMFFIFieldGetter function pointer for a given field type index. * * \param field_type_index The static type index of the field. * \return The function pointer as int64_t for FFI transport. */ int64_t GetFieldGetter(int32_t field_type_index) { TVMFFIFieldGetter getter = nullptr; switch (field_type_index) { case TypeIndex::kTVMFFIInt: getter = &PyClassFieldGetter; break; case TypeIndex::kTVMFFIFloat: getter = &PyClassFieldGetter; break; case TypeIndex::kTVMFFIBool: getter = &PyClassFieldGetter; break; case TypeIndex::kTVMFFIOpaquePtr: getter = &PyClassFieldGetter; break; case TypeIndex::kTVMFFIDataType: getter = &PyClassFieldGetter; break; case TypeIndex::kTVMFFIDevice: getter = &PyClassFieldGetter; break; default: if (field_type_index == TypeIndex::kTVMFFIAny || field_type_index == TypeIndex::kTVMFFINone) { getter = &PyClassFieldGetter; } else if (field_type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { getter = &PyClassFieldGetter; } else { TVM_FFI_THROW(ValueError) << "Unsupported field type index for getter: " << field_type_index; } break; } return reinterpret_cast(getter); } /*! * \brief Write a converted value to a field of the appropriate C++ type. * * Dispatches on field_type_index to reinterpret the destination address and * assign from the converted Any value. */ void WriteFieldValue(void* field_addr, int32_t field_type_index, Any value) { switch (field_type_index) { case TypeIndex::kTVMFFIInt: *reinterpret_cast(field_addr) = value.cast(); return; case TypeIndex::kTVMFFIFloat: *reinterpret_cast(field_addr) = value.cast(); return; case TypeIndex::kTVMFFIBool: *reinterpret_cast(field_addr) = value.cast(); return; case TypeIndex::kTVMFFIOpaquePtr: *reinterpret_cast(field_addr) = value.cast(); return; case TypeIndex::kTVMFFIDataType: *reinterpret_cast(field_addr) = value.cast(); return; case TypeIndex::kTVMFFIDevice: *reinterpret_cast(field_addr) = value.cast(); return; default: break; } if (field_type_index == TypeIndex::kTVMFFIAny || field_type_index == TypeIndex::kTVMFFINone) { *reinterpret_cast(field_addr) = std::move(value); } else if (field_type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { *reinterpret_cast(field_addr) = value.cast(); } else { TVM_FFI_THROW(ValueError) << "Unsupported field type index for setter: " << field_type_index; } } /*! * \brief Create a FunctionObj setter for a Python-defined field. * * The returned Function accepts (OpaquePtr field_addr, AnyView value), * calls f_convert to coerce the value via the type_converter, and writes * the result to the field. * * \param field_type_index The static type index of the field. * \param type_converter_int Opaque pointer (as int64_t) to the Python _TypeConverter (borrowed). * \param f_convert_int C function pointer (as int64_t): int(void*, const TVMFFIAny*, TVMFFIAny*). * Returns 0 on success, -1 on error (error stored in TLS). * \return A packed Function suitable for use as a FunctionObj setter. */ Function MakeFieldSetter(int32_t field_type_index, int64_t type_converter_int, int64_t f_convert_int) { // NOLINTNEXTLINE(performance-no-int-to-ptr) void* type_converter = reinterpret_cast(type_converter_int); using FConvert = int (*)(void*, const TVMFFIAny*, TVMFFIAny*); // NOLINTNEXTLINE(performance-no-int-to-ptr) auto f_convert = reinterpret_cast(f_convert_int); return Function::FromPacked([field_type_index, type_converter, f_convert]( const AnyView* args, int32_t num_args, Any* rv) { void* field_addr = args[0].cast(); // Call the Cython-level type converter via C function pointer. TVMFFIAny converted; converted.type_index = TypeIndex::kTVMFFINone; converted.v_int64 = 0; int err = f_convert(type_converter, reinterpret_cast(&args[1]), &converted); if (err != 0) { throw details::MoveFromSafeCallRaised(); } // Take ownership of the converted value and write to the field. Any owned = details::AnyUnsafe::MoveTVMFFIAnyToAny(&converted); WriteFieldValue(field_addr, field_type_index, std::move(owned)); }); } // ============================================================================ // Shared helpers for MakeInit // ============================================================================ /*! * \brief Return the lazily initialized KWARGS sentinel object. */ const ObjectRef& GetKwargsSentinel() { static ObjectRef kwargs_sentinel = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); return kwargs_sentinel; } /*! \brief Pre-computed field analysis for auto-generated init. */ struct AutoInitInfo { struct Entry { const TVMFFIFieldInfo* info; bool init; bool kw_only; bool has_default; }; std::vector all_fields; std::vector init_indices; std::vector pos_indices; std::unordered_map name_to_index; std::string_view type_key; }; /*! * \brief Build AutoInitInfo by analysing reflected fields for a type. */ std::shared_ptr BuildAutoInitInfo(int32_t type_index) { const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index); auto info = std::make_shared(); info->type_key = std::string_view(type_info->type_key.data, type_info->type_key.size); refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* fi) { bool is_init = (fi->flags & kTVMFFIFieldFlagBitMaskInitOff) == 0; bool is_kw = (fi->flags & kTVMFFIFieldFlagBitMaskKwOnly) != 0; bool has_def = (fi->flags & kTVMFFIFieldFlagBitMaskHasDefault) != 0; info->all_fields.push_back({fi, is_init, is_kw, has_def}); size_t idx = info->all_fields.size() - 1; if (is_init) { info->init_indices.push_back(idx); info->name_to_index[std::string_view(fi->name.data, fi->name.size)] = idx; if (!is_kw) { info->pos_indices.push_back(idx); } } }); // Required fields come before optional ones in positional list. std::stable_partition(info->pos_indices.begin(), info->pos_indices.end(), [&](size_t idx) { return !info->all_fields[idx].has_default; }); return info; } /*! * \brief Bind packed arguments to fields on an existing object. * * Handles both positional-only and KWARGS calling conventions, sets fields * via ``CallFieldSetter``, fills defaults for unbound fields, and raises * ``TypeError`` for missing required arguments (with ``keyword-only`` * distinction when applicable). */ void BindFieldArgs(Object* obj, const AutoInitInfo& info, const TVMFFIAny* raw_args, int num_args) { const ObjectRef& kwargs_sentinel = GetKwargsSentinel(); std::vector field_set(info.all_fields.size(), false); auto set_field = [&](size_t fi, const TVMFFIAny* value) { void* addr = reinterpret_cast(obj) + info.all_fields[fi].info->offset; TVM_FFI_CHECK_SAFE_CALL(refl::CallFieldSetter(info.all_fields[fi].info, addr, value)); field_set[fi] = true; }; // ---- 1. Find KWARGS sentinel position ------------------------------------ int kwargs_pos = -1; for (int i = 0; i < num_args; ++i) { auto opt = AnyView::CopyFromTVMFFIAny(raw_args[i]).as(); if (opt.has_value() && opt.value().same_as(kwargs_sentinel)) { kwargs_pos = i; break; } } // ---- 2. Bind arguments to fields ----------------------------------------- if (kwargs_pos >= 0) { // --- 2a. KWARGS mode ---------------------------------------------------- int pos_arg = 0; for (size_t fi : info.pos_indices) { if (pos_arg < kwargs_pos) { set_field(fi, &raw_args[pos_arg]); ++pos_arg; } } if (pos_arg < kwargs_pos) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() takes at most " << info.pos_indices.size() << " positional argument(s), but " << kwargs_pos << " were given"; } int kv_count = num_args - kwargs_pos - 1; if (kv_count % 2 != 0) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() KWARGS requires an even number of key-value arguments"; } for (int i = kwargs_pos + 1; i < num_args; i += 2) { String key = AnyView::CopyFromTVMFFIAny(raw_args[i]).cast(); std::string_view key_sv(key.data(), key.size()); auto it = info.name_to_index.find(key_sv); if (it == info.name_to_index.end()) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() got an unexpected keyword argument '" << key << "'"; } size_t idx = it->second; if (field_set[idx]) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() got multiple values " << "for argument '" << key << "'"; } set_field(idx, &raw_args[i + 1]); } } else { // --- 2b. Positional-only mode ------------------------------------------- if (static_cast(num_args) > info.pos_indices.size()) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() takes at most " << info.pos_indices.size() << " positional argument(s), but " << num_args << " were given"; } for (int i = 0; i < num_args; ++i) { set_field(info.pos_indices[i], &raw_args[i]); } } // ---- 3. Fill defaults and check required fields -------------------------- for (size_t fi = 0; fi < info.all_fields.size(); ++fi) { if (field_set[fi]) continue; if (info.all_fields[fi].has_default) { void* addr = reinterpret_cast(obj) + info.all_fields[fi].info->offset; refl::SetFieldToDefault(info.all_fields[fi].info, addr); } else if (info.all_fields[fi].init) { auto fname = std::string_view(info.all_fields[fi].info->name.data, info.all_fields[fi].info->name.size); if (info.all_fields[fi].kw_only) { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() missing required keyword-only argument: '" << fname << "'"; } else { TVM_FFI_THROW(TypeError) << info.type_key << ".__ffi_init__() missing required argument: '" << fname << "'"; } } // init=False without default: leave at creator default. } } Function MakeInit(int32_t type_index) { auto info = BuildAutoInitInfo(type_index); const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index); return Function::FromPacked([info, type_info](PackedArgs args, Any* rv) { ObjectPtr obj_ptr = CreateEmptyObject(type_info); const auto raw_args = reinterpret_cast(args.data()); BindFieldArgs(obj_ptr.get(), *info, raw_args, args.size()); *rv = ObjectRef(obj_ptr); }); } void RegisterFFIInit(int32_t type_index) { namespace refl = ::tvm::ffi::reflection; Function auto_init_fn = MakeInit(type_index); TVMFFIByteArray attr_name = refl::AsByteArray(refl::type_attr::kInit); TVMFFIAny attr_value = AnyView(auto_init_fn).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &attr_name, &attr_value)); } /*! * \brief Combined registration for Python-defined types: * ``__ffi_init__``, ``__ffi_new__``, ``__ffi_shallow_copy__`` */ void PyClassRegisterTypeAttrColumns(int32_t type_index, int32_t total_size) { namespace refl = ::tvm::ffi::reflection; const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index); // Step 1. Register `__ffi_init__` RegisterFFIInit(type_index); // Step 2. Register `__ffi_new__` Function new_fn = Function::FromTyped([type_index, total_size]() -> ObjectRef { void* obj_ptr = std::calloc(1, static_cast(total_size)); if (!obj_ptr) { TVM_FFI_THROW(RuntimeError) << "Failed to allocate " << total_size << " bytes for type " << TypeIndexToTypeKey(type_index); } TVMFFIObject* ffi_obj = reinterpret_cast(obj_ptr); ffi_obj->type_index = type_index; ffi_obj->combined_ref_count = details::kCombinedRefCountBothOne; ffi_obj->deleter = PyClassDeleter; // calloc zero-initializes all bytes. For non-trivial field types: // - Any: zero state is {type_index=kTVMFFINone, v_int64=0}, representing None. // - ObjectRef: zero state is a null pointer. // Both are valid initial states whose destructors and assignment operators // handle correctly, so no placement construction is needed. Object* obj = reinterpret_cast(obj_ptr); return ObjectRef(details::ObjectUnsafe::ObjectPtrFromOwned(obj)); }); TVMFFIByteArray attr_name = refl::AsByteArray(refl::type_attr::kNew); TVMFFIAny attr_value = AnyView(new_fn).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &attr_name, &attr_value)); // Step 3. Register `__ffi_shallow_copy__` Function copy_fn = Function::FromTyped([type_index, total_size, type_info](const Object* src) -> ObjectRef { void* obj_ptr = std::calloc(1, static_cast(total_size)); if (!obj_ptr) { TVM_FFI_THROW(RuntimeError) << "Failed to allocate for shallow copy"; } TVMFFIObject* ffi_obj = reinterpret_cast(obj_ptr); ffi_obj->type_index = type_index; ffi_obj->combined_ref_count = details::kCombinedRefCountBothOne; ffi_obj->deleter = PyClassDeleter; refl::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* finfo) { void* dst = reinterpret_cast(obj_ptr) + finfo->offset; const void* field_src = reinterpret_cast(src) + finfo->offset; int32_t ti = finfo->field_static_type_index; if (ti == TypeIndex::kTVMFFIAny) { new (dst) Any(*reinterpret_cast(field_src)); } else if (ti >= TypeIndex::kTVMFFIStaticObjectBegin) { new (dst) ObjectRef(*reinterpret_cast(field_src)); } else { // POD: memcpy std::memcpy(dst, field_src, static_cast(finfo->size)); } }); Object* obj = reinterpret_cast(obj_ptr); return ObjectRef(details::ObjectUnsafe::ObjectPtrFromOwned(obj)); }); TVMFFIByteArray copy_attr_name = refl::AsByteArray(refl::type_attr::kShallowCopy); TVMFFIAny copy_attr_value = AnyView(copy_fn).CopyToTVMFFIAny(); TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, ©_attr_name, ©_attr_value)); } } // namespace // ============================================================================ // Deep Copy — public API // ============================================================================ Any DeepCopy(const Any& value) { static refl::TypeAttrColumn column(refl::type_attr::kShallowCopy); ObjectDeepCopier copier(&column); return copier.Run(value); } // ============================================================================ // Repr Printing — public API // ============================================================================ String ReprPrint(const Any& value) { ReprPrinter printer; return printer.Run(value); } // ============================================================================ // Recursive Hash / Compare — public API // ============================================================================ int64_t RecursiveHash(const Any& value) { RecursiveHasher hasher; return static_cast(hasher.HashAny(value)); } bool RecursiveEq(const Any& lhs, const Any& rhs) { RecursiveComparer cmp(/*eq_only=*/true); return cmp.CompareAny(lhs, rhs) == 0; } bool RecursiveLt(const Any& lhs, const Any& rhs) { RecursiveComparer cmp(/*eq_only=*/false); return cmp.CompareAny(lhs, rhs) < 0; } bool RecursiveLe(const Any& lhs, const Any& rhs) { RecursiveComparer cmp(/*eq_only=*/false); return cmp.CompareAny(lhs, rhs) <= 0; } bool RecursiveGt(const Any& lhs, const Any& rhs) { RecursiveComparer cmp(/*eq_only=*/false); return cmp.CompareAny(lhs, rhs) > 0; } bool RecursiveGe(const Any& lhs, const Any& rhs) { RecursiveComparer cmp(/*eq_only=*/false); return cmp.CompareAny(lhs, rhs) >= 0; } // ============================================================================ // Registration // ============================================================================ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::EnsureTypeAttrColumn(refl::type_attr::kShallowCopy); refl::EnsureTypeAttrColumn(refl::type_attr::kRepr); refl::EnsureTypeAttrColumn(refl::type_attr::kHash); refl::EnsureTypeAttrColumn(refl::type_attr::kEq); refl::EnsureTypeAttrColumn(refl::type_attr::kCompare); refl::EnsureTypeAttrColumn(refl::type_attr::kNew); refl::EnsureTypeAttrColumn(refl::type_attr::kInit); refl::GlobalDef() .def("ffi._RegisterFFIInit", RegisterFFIInit) .def("ffi.MakeFieldSetter", MakeFieldSetter) .def("ffi._PyClassRegisterTypeAttrColumns", PyClassRegisterTypeAttrColumns) .def("ffi.DeepCopy", DeepCopy) .def("ffi.ReprPrint", ReprPrint) .def("ffi.RecursiveHash", RecursiveHash) .def("ffi.MakeFieldGetter", GetFieldGetter) .def("ffi.RecursiveEq", RecursiveEq) .def("ffi.RecursiveLt", RecursiveLt) .def("ffi.RecursiveLe", RecursiveLe) .def("ffi.RecursiveGt", RecursiveGt) .def("ffi.RecursiveGe", RecursiveGe); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/env_c_api.cc000066400000000000000000000113531521067262500200640ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/env_c_api.cc * \brief Environment C API implementation. */ #include #include namespace tvm { namespace ffi { /*! * \brief Execution environment specific API registry. * * This registry stores C API function pointers about * execution environment(e.g. python) specific API function that * we need for specific low-level handling(e.g. signal checking). * * We only stores the C API function when absolutely necessary (e.g. when signal handler * cannot trap back into python). Always consider use the Function FFI when possible * in other cases. */ class EnvCAPIRegistry { public: /*! * \brief Callback to check if signals have been sent to the process and * if so invoke the registered signal handler in the frontend environment. * * When running FFI in another language (Python), the signal handler * may not be immediately executed, but instead the signal is marked * in the interpreter state (to ensure non-blocking of the signal handler). * * \return 0 if no error happens, -1 if error happens. */ using F_PyErr_CheckSignals = int (*)(); /*! \brief Callback to increment/decrement the python ref count */ using F_Py_IncDefRef = void (*)(void*); /*! * \brief PyErr_CheckSignal function */ F_PyErr_CheckSignals pyerr_check_signals = nullptr; /*! \brief PyGILState_Ensure function */ void* (*py_gil_state_ensure)() = nullptr; /*! \brief PyGILState_Release function */ void (*py_gil_state_release)(void*) = nullptr; static EnvCAPIRegistry* Global() { static EnvCAPIRegistry* inst = new EnvCAPIRegistry(); return inst; } // register environment(e.g. python) specific api functions void Register(const String& symbol_name, void* fptr) { if (symbol_name == "PyErr_CheckSignals") { Update(symbol_name, &pyerr_check_signals, fptr); } else if (symbol_name == "PyGILState_Ensure") { Update(symbol_name, &py_gil_state_ensure, fptr); } else if (symbol_name == "PyGILState_Release") { Update(symbol_name, &py_gil_state_release, fptr); } else { TVM_FFI_THROW(ValueError) << "Unknown env API " + symbol_name; } } int EnvCheckSignals() { // check python signal to see if there are exception raised if (pyerr_check_signals != nullptr) { // The C++ env comes without gil, so we need to grab gil here WithGIL context(this); if ((*pyerr_check_signals)() != 0) { // The error will let FFI know that the frontend environment // already set an error. return -1; } } return 0; } private: // update the internal API table template void Update(const String& symbol_name, FType* target, void* ptr) { FType ptr_casted = reinterpret_cast(ptr); target[0] = ptr_casted; } struct WithGIL { explicit WithGIL(EnvCAPIRegistry* self) : self(self) { TVM_FFI_ICHECK(self->py_gil_state_ensure); TVM_FFI_ICHECK(self->py_gil_state_release); gil_state = self->py_gil_state_ensure(); } ~WithGIL() { if (self && gil_state) { self->py_gil_state_release(gil_state); } } WithGIL(const WithGIL&) = delete; WithGIL(WithGIL&&) = delete; WithGIL& operator=(const WithGIL&) = delete; WithGIL& operator=(WithGIL&&) = delete; EnvCAPIRegistry* self = nullptr; void* gil_state = nullptr; }; }; } // namespace ffi } // namespace tvm int TVMFFIEnvCheckSignals() { return tvm::ffi::EnvCAPIRegistry::Global()->EnvCheckSignals(); } /*! * \brief Register a symbol into the from the surrounding env. * \param name The name of the symbol. * \param symbol The symbol to register. * \return 0 when success, nonzero when failure happens */ int TVMFFIEnvRegisterCAPI(const char* name, void* symbol) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::String s_name(name); tvm::ffi::EnvCAPIRegistry::Global()->Register(s_name, symbol); TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/extra/env_context.cc000066400000000000000000000141301521067262500204710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/env_context.cc * * \brief A minimalistic env context based on ffi values. */ #include #include #include #include namespace tvm { namespace ffi { class EnvContext { public: void SetStream(int32_t device_type, int32_t device_id, TVMFFIStreamHandle stream, TVMFFIStreamHandle* out_original_stream) { if (static_cast(device_type) >= stream_table_.size()) { stream_table_.resize(device_type + 1); } if (static_cast(device_id) >= stream_table_[device_type].size()) { stream_table_[device_type].resize(device_id + 1, nullptr); } if (out_original_stream != nullptr) { *out_original_stream = stream_table_[device_type][device_id]; } stream_table_[device_type][device_id] = stream; } TVMFFIStreamHandle GetStream(int32_t device_type, int32_t device_id) { if (static_cast(device_type) < stream_table_.size() && static_cast(device_id) < stream_table_[device_type].size()) { return stream_table_[device_type][device_id]; } return nullptr; } DLPackManagedTensorAllocator GetDLPackManagedTensorAllocator() { if (dlpack_allocator_ != nullptr) { return dlpack_allocator_; } return GlobalTensorAllocator(); } void SetDLPackManagedTensorAllocator(DLPackManagedTensorAllocator allocator, int write_to_global_context, DLPackManagedTensorAllocator* opt_out_original_allocator) { if (opt_out_original_allocator != nullptr) { // only returns the cached local allocator and ignore global allocator *opt_out_original_allocator = dlpack_allocator_; } if (write_to_global_context != 0) { GlobalTensorAllocator() = allocator; } dlpack_allocator_ = allocator; } static EnvContext* ThreadLocal() { static thread_local EnvContext inst; return &inst; } private: // use static function to avoid static initialization order issue static DLPackManagedTensorAllocator& GlobalTensorAllocator() { // NOLINT(*) static DLPackManagedTensorAllocator allocator = nullptr; return allocator; } std::vector> stream_table_; DLPackManagedTensorAllocator dlpack_allocator_ = nullptr; }; } // namespace ffi } // namespace tvm int TVMFFIEnvSetStream(int32_t device_type, int32_t device_id, TVMFFIStreamHandle stream, TVMFFIStreamHandle* out_original_stream) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::EnvContext::ThreadLocal()->SetStream(device_type, device_id, stream, out_original_stream); TVM_FFI_SAFE_CALL_END(); } TVMFFIStreamHandle TVMFFIEnvGetStream(int32_t device_type, int32_t device_id) { TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); return tvm::ffi::EnvContext::ThreadLocal()->GetStream(device_type, device_id); TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIEnvGetStream); } int TVMFFIEnvSetDLPackManagedTensorAllocator( DLPackManagedTensorAllocator allocator, int write_to_global_context, DLPackManagedTensorAllocator* opt_out_original_allocator) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::EnvContext::ThreadLocal()->SetDLPackManagedTensorAllocator( allocator, write_to_global_context, opt_out_original_allocator); TVM_FFI_SAFE_CALL_END(); } DLPackManagedTensorAllocator TVMFFIEnvGetDLPackManagedTensorAllocator() { TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); return tvm::ffi::EnvContext::ThreadLocal()->GetDLPackManagedTensorAllocator(); TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIEnvGetDLPackManagedTensorAllocator); } void TVMFFIEnvTensorAllocSetError(void* error_ctx, const char* kind, const char* message) { TVMFFIErrorSetRaisedFromCStr(kind, message); } int TVMFFIEnvTensorAlloc(DLTensor* prototype, TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); DLPackManagedTensorAllocator dlpack_alloc = tvm::ffi::EnvContext::ThreadLocal()->GetDLPackManagedTensorAllocator(); if (dlpack_alloc == nullptr) { TVMFFIErrorSetRaisedFromCStr( "RuntimeError", "TVMFFIEnvTensorAlloc: allocator is nullptr, " "likely because TVMFFIEnvSetDLPackManagedTensorAllocator has not been called."); return -1; } DLManagedTensorVersioned* dlpack_tensor = nullptr; int ret = (*dlpack_alloc)(prototype, &dlpack_tensor, nullptr, TVMFFIEnvTensorAllocSetError); if (ret != 0) return ret; // need to first check ret so we can properly propagate the error from caller side. TVM_FFI_CHECK(dlpack_tensor != nullptr, RuntimeError) << "TVMFFIEnvTensorAlloc: failed to allocate a tensor from the allocator"; if (dlpack_tensor->dl_tensor.strides != nullptr || dlpack_tensor->dl_tensor.ndim == 0) { *out = tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr( tvm::ffi::make_object>( dlpack_tensor, /*extra_strides_at_tail=*/false)); } else { *out = tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr( tvm::ffi::make_inplace_array_object< tvm::ffi::details::TensorObjFromDLPack, int64_t>( dlpack_tensor->dl_tensor.ndim, dlpack_tensor, /*extra_strides_at_tail=*/true)); } TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/extra/json_parser.cc000066400000000000000000000577571521067262500205100ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/json/parser.cc * * \brief A minimalistic JSON parser based on ffi values. */ #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace json { /*! * \brief Helper class to parse a JSON string. * * Keep leaf level string/number parse also in context. */ class JSONParserContext { public: JSONParserContext(const char* begin, const char* end) : begin_(begin), cur_(begin), end_(end) { last_line_begin_ = cur_; } /*! * \brief Peek the current character. * \return The current character, or -1 if the end of the string is reached. */ int Peek() const { return (cur_ != end_ ? static_cast(*reinterpret_cast(cur_)) : -1); } /*! * \brief Skip the next char that we know is not a space * * \note Caller must explicitly call SkipSpaces first or use * Peek already that confirms char is not any space char. */ void SkipNextAssumeNoSpace() { ++cur_; } /*! * \brief Get the current position. * \return The current position. */ const char* GetCurrentPos() const { return cur_; } /*! * \brief Set the current position for better error message * \param pos The new position. * \note implementation can do it as no-op if needed */ void SetCurrentPosForBetterErrorMsg(const char* pos) { cur_ = pos; } /*! * \brief Skip the space characters. * \note This function does not check if the end of the string is reached. */ void SkipSpaces() { while (cur_ != end_) { if (!(*cur_ == ' ' || *cur_ == '\t' || *cur_ == '\n' || *cur_ == '\r')) { break; } if (*cur_ == '\n') { ++line_counter_; last_line_begin_ = cur_ + 1; } ++cur_; } } /*! * \brief Check if the next characters match the given string. * \param str The string to match. * \param len The length of the string. * \return True if the next characters match the given string, false otherwise. */ bool MatchLiteral(const char* pattern, int len) { const char* pend = pattern + len; const char* ptr = pattern; for (; ptr != pend && cur_ != end_; ++ptr, ++cur_) { if (*ptr != *cur_) { return false; } } // we get to the end of the pattern and match is successful return ptr == pend; } /* * \brief Parse the next strin starting with a double quote. * \param out The output string. * \return Whether the next string parsing is successful. */ bool NextString(json::Value* out) { // NOTE: we keep string parsing logic here to allow some special // optimizations for simple string that do not e const char* start_pos = cur_; TVM_FFI_ICHECK(*cur_ == '\"'); // skip first double quote ++cur_; // the loop focuses on simple string without escape characters for (; cur_ != end_; ++cur_) { if (*cur_ == '\"') { *out = String(start_pos + 1, cur_ - start_pos - 1); ++cur_; return true; } // Use uint8_t: on platforms where char is signed, raw UTF-8 bytes (>= 0x80) // would be negative and incorrectly treated as control characters. if (*reinterpret_cast(cur_) < ' ' || *cur_ == '\\') { // fallback to full string handling return this->NextStringWithFullHandling(out, start_pos); } } this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorUnterminatedString(); return false; } /*! * \brief Parse the next number. * \param out The output number. * \return Whether the next number parsing is successful. */ bool NextNumber(json::Value* out) { const char* start_pos = cur_; if (cur_ == end_) { this->SetErrorExpectingValue(); return false; } // JSON number grammar: // // number = [ minus ] int [ frac ] [ exp ] // decimal-point = %x2E ; . // digit1-9 = %x31-39 ; 1-9 // e = %x65 / %x45 ; e E // exp = e [ minus / plus ] 1*DIGIT // frac = decimal-point 1*DIGIT std::string temp_buffer; bool maybe_int = true; // parse [minus], cross check for Infinity/NaN/-Infinity if (*cur_ == '-') { temp_buffer.push_back('-'); ++cur_; if (cur_ != end_ && *cur_ == 'I') { if (this->MatchLiteral("Infinity", 8)) { *out = FastMathSafeNegInf(); return true; } else { this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorExpectingValue(); return false; } } } else if (*cur_ == 'I') { if (this->MatchLiteral("Infinity", 8)) { *out = FastMathSafePosInf(); return true; } else { this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorExpectingValue(); return false; } } else if (*cur_ == 'N') { if (this->MatchLiteral("NaN", 3)) { *out = FastMathSafeNaN(); return true; } else { this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorExpectingValue(); return false; } } // read in all parts that are possibly part of a number while (cur_ != end_) { char next_char = *cur_; if ((next_char >= '0' && next_char <= '9') || next_char == 'e' || next_char == 'E' || next_char == '+' || next_char == '-' || next_char == '.') { temp_buffer.push_back(next_char); if (next_char == '.' || next_char == 'e' || next_char == 'E') { maybe_int = false; } ++cur_; } else { break; } } if (temp_buffer.empty()) { this->SetErrorExpectingValue(); return false; } // parse from temp_buffer_ if (maybe_int) { // now try to parse the number as int64 char* end_ptr; errno = 0; intmax_t int_val = strtoimax(temp_buffer.data(), &end_ptr, 10); if (errno == 0 && int_val >= std::numeric_limits::min() && int_val <= std::numeric_limits::max() && end_ptr == temp_buffer.data() + temp_buffer.size()) { *out = static_cast(int_val); return true; } } { // now try to parse number as double char* end_ptr; errno = 0; double double_val = strtod(temp_buffer.data(), &end_ptr); if (errno == 0 && end_ptr == temp_buffer.data() + temp_buffer.size()) { *out = double_val; return true; } else { this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorExpectingValue(); return false; } } } /*! * \brief Get the current line context. * \return The current line context. */ String GetSyntaxErrorContext(std::string err_prefix) const { int64_t column = static_cast(cur_ - last_line_begin_) + 1; int64_t char_pos = static_cast(cur_ - begin_); if (err_prefix.empty()) { err_prefix = "Syntax error"; } err_prefix += ": line " + std::to_string(line_counter_) + " column " + std::to_string(column) + " (char " + std::to_string(char_pos) + ")"; return String(err_prefix); } std::string FinalizeErrorMsg() { if (error_msg_.empty()) { SetErrorDefault(); } return std::string(error_msg_); } void SetErrorDefault() { error_msg_ = GetSyntaxErrorContext("Syntax error near"); } void SetErrorExpectingValue() { error_msg_ = GetSyntaxErrorContext("Expecting value"); } void SetErrorInvalidControlCharacter() { error_msg_ = GetSyntaxErrorContext("Invalid control character at"); } void SetErrorUnterminatedString() { error_msg_ = GetSyntaxErrorContext("Unterminated string starting at"); } void SetErrorInvalidUnicodeEscape() { error_msg_ = GetSyntaxErrorContext("Invalid \\uXXXX escape"); } void SetErrorInvalidSurrogatePair() { error_msg_ = GetSyntaxErrorContext("Invalid surrogate pair of \\uXXXX escapes"); } void SetErrorInvalidEscape() { error_msg_ = GetSyntaxErrorContext("Invalid \\escape"); } void SetErrorExtraData() { error_msg_ = GetSyntaxErrorContext("Extra data"); } void SetErrorExpectingPropertyName() { error_msg_ = GetSyntaxErrorContext("Expecting property name enclosed in double quotes"); } void SetErrorExpectingColon() { error_msg_ = GetSyntaxErrorContext("Expecting \':\' delimiter"); } void SetErrorExpectingComma() { error_msg_ = GetSyntaxErrorContext("Expecting \',\' delimiter"); } private: static double FastMathSafePosInf() { #ifdef __FAST_MATH__ union { uint64_t from; double to; } u; u.from = 0x7FF0000000000000ULL; // write "from", read "to" return u.to; #else return std::numeric_limits::infinity(); #endif } static double FastMathSafeNegInf() { #ifdef __FAST_MATH__ union { uint64_t from; double to; } u; u.from = 0xFFF0000000000000ULL; // write "from", read "to" return u.to; #else return -std::numeric_limits::infinity(); #endif } static double FastMathSafeNaN() { #ifdef __FAST_MATH__ union { uint64_t from; double to; } u; u.from = 0x7FF8000000000000ULL; // write "from", read "to" return u.to; #else return std::numeric_limits::quiet_NaN(); #endif } // Full string parsing with escape and unicode handling bool NextStringWithFullHandling(Any* out, const char* start_pos) { // copy over the prefix that was already parsed std::string out_str(start_pos + 1, cur_ - start_pos - 1); while (cur_ != end_) { // Use uint8_t: on platforms where char is signed, raw UTF-8 bytes (>= 0x80) // would be negative and incorrectly treated as control characters. if (*reinterpret_cast(cur_) < ' ') { this->SetErrorInvalidControlCharacter(); return false; } if (*cur_ == '\"') { *out = String(std::move(out_str)); ++cur_; return true; } if (*cur_ == '\\') { ++cur_; switch (*cur_) { // handle escape characters per JSON spec(RFC 8259) #define HANDLE_ESCAPE_CHAR(pattern, val) \ case pattern: \ ++cur_; \ out_str.push_back(val); \ break HANDLE_ESCAPE_CHAR('\"', '\"'); HANDLE_ESCAPE_CHAR('\\', '\\'); HANDLE_ESCAPE_CHAR('/', '/'); HANDLE_ESCAPE_CHAR('b', '\b'); HANDLE_ESCAPE_CHAR('f', '\f'); HANDLE_ESCAPE_CHAR('n', '\n'); HANDLE_ESCAPE_CHAR('r', '\r'); HANDLE_ESCAPE_CHAR('t', '\t'); #undef HANDLE_ESCAPE_CHAR case 'u': { const char* escape_pos = cur_; // handle unicode code point ++cur_; int32_t first_i16, code_point = 0; if (!Parse4Hex(&first_i16)) { this->SetCurrentPosForBetterErrorMsg(escape_pos); this->SetErrorInvalidUnicodeEscape(); return false; } // Check if the first i16 is a UTF-16 surrogate pair // // Surrogate pair encoding rule: // U' = yyyyyyyyyyxxxxxxxxxx // U - 0x10000 // W1 = 110110yyyyyyyyyy // 0xD800 + yyyyyyyyyy // W2 = 110111xxxxxxxxxx // 0xDC00 + xxxxxxxxxx // // Range of W1 and W2: // 0xD800 - 0xDBFF for W1 // 0xDC00 - 0xDFFF for W2 // both W1 and W2 fit into 0xD800 - 0xDFFF // Detect if the first i16 fit into range of W1/W2 if (first_i16 >= 0xD800 && first_i16 <= 0xDFFF) { // we are in the surrogate pair range if (first_i16 >= 0xDC00) { this->SetCurrentPosForBetterErrorMsg(escape_pos); this->SetErrorInvalidSurrogatePair(); // we need to return false instead because this range is for W2 return false; } if (!this->MatchLiteral("\\u", 2)) { this->SetCurrentPosForBetterErrorMsg(escape_pos); this->SetErrorInvalidSurrogatePair(); return false; } escape_pos = cur_; // get the value of the W2 (second i16) int32_t second_i16; if (!Parse4Hex(&second_i16)) { this->SetCurrentPosForBetterErrorMsg(escape_pos); this->SetErrorInvalidUnicodeEscape(); return false; } if (!(second_i16 >= 0xDC00 && second_i16 <= 0xDFFF)) { this->SetCurrentPosForBetterErrorMsg(escape_pos); this->SetErrorInvalidSurrogatePair(); return false; } // recover the code point code_point = ((first_i16 - 0xD800) << 10) + (second_i16 - 0xDC00) + 0x10000; } else { // not a surrogate case, just assign as code point code_point = first_i16; } // now need to push back the string based on UTF-8 encoding // UTF-8 encoding rule: four cases // ------------------------------------------------------------ // Pattern | code point range // ------------------------------------------------------------ // 0xxxxxxx | 0x0 - 0x7F // 110xxxxx 10xxxxxx | 0x80 - 0x7FF // 1110xxxx 10xxxxxx 10xxxxxx | 0x800 - 0xFFFF // 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx | 0x10000 - end // ------------------------------------------------------------ if (code_point < 0x80) { out_str.push_back(static_cast(code_point)); } else if (code_point < 0x800) { // first byte: 110xxxxx (5 effective bits) // second byte: 10xxxxxx (6 effecive bits) // shift by 6 bits to get the first bytes out_str.push_back(static_cast(0xC0 | (code_point >> 6))); // mask by 6 effective bits out_str.push_back(static_cast(0x80 | (code_point & 0x3F))); } else if (code_point < 0x10000) { // first byte: 1110xxxx (4 effective bits) // second byte: 10xxxxxx (6 effecive bits) // third byte: 10xxxxxx (6 effecive bits) // shift by 12 bits to get the first bytes out_str.push_back(static_cast(0xE0 | (code_point >> 12))); // shift by 6 bits to get the second bytes, mask by 6 effective bits out_str.push_back(static_cast(0x80 | ((code_point >> 6) & 0x3F))); // mask by 6 effective bits out_str.push_back(static_cast(0x80 | (code_point & 0x3F))); } else { // first byte: 11110xxx (3 effective bits) // second byte: 10xxxxxx (6 effecive bits) // third byte: 10xxxxxx (6 effecive bits) // fourth byte: 10xxxxxx (6 effecive bits) // shift by 18 bits to get the first bytes out_str.push_back(static_cast(0xF0 | (code_point >> 18))); // shift by 12 bits to get the second bytes, mask by 6 effective bits out_str.push_back(static_cast(0x80 | ((code_point >> 12) & 0x3F))); // shift by 6 bits to get the third bytes, mask by 6 effective bits out_str.push_back(static_cast(0x80 | ((code_point >> 6) & 0x3F))); // mask by 6 effective bits out_str.push_back(static_cast(0x80 | (code_point & 0x3F))); } break; } default: { this->SetErrorInvalidEscape(); return false; } } } else { out_str.push_back(*cur_); ++cur_; } } this->SetCurrentPosForBetterErrorMsg(start_pos); this->SetErrorUnterminatedString(); return false; } /*! * \brief Parse the four hex digits of a unicode code point per json spec. * \param out_i16 The output i16 number * \return True if four hex digits are parsed successfully, false otherwise. */ bool Parse4Hex(int32_t* out_i16) { int32_t result = 0; for (int i = 0; i < 4; ++i, ++cur_) { int hex_val = *reinterpret_cast(cur_); if (hex_val >= '0' && hex_val <= '9') { hex_val -= '0'; } else if (hex_val >= 'a' && hex_val <= 'f') { hex_val -= 'a' - 0xa; } else if (hex_val >= 'A' && hex_val <= 'F') { hex_val -= 'A' - 0xa; } else { return false; } result = result * 16 + hex_val; } *out_i16 = result; return true; } /*! \brief The beginning of the string */ const char* begin_; /*! \brief The current pointer */ const char* cur_; /*! \brief End of the string */ const char* end_; /*! \brief The beginning of the last line */ const char* last_line_begin_; /*! \brief The error message */ std::string error_msg_; /*! \brief The line counter */ int64_t line_counter_{1}; }; class JSONParser { public: static json::Value Parse(const String& json_str, String* error_msg) { JSONParser parser(json_str); json::Value result; if (parser.ParseValue(&result) && parser.ParseTail()) { if (error_msg != nullptr) { *error_msg = String(""); } return result; } if (error_msg != nullptr) { *error_msg = parser.ctx_.FinalizeErrorMsg(); TVM_FFI_ICHECK(!error_msg->empty()); } else { TVM_FFI_THROW(ValueError) << parser.ctx_.FinalizeErrorMsg(); } // note that when we don't throw, error msg is set to indicate // an error happens return nullptr; } private: explicit JSONParser(String json_str) : storage_(std::move(json_str)), ctx_(storage_.data(), storage_.data() + storage_.size()) {} bool ParseTail() { ctx_.SkipSpaces(); // there are extra data in the tail if (ctx_.Peek() != -1) { ctx_.SetErrorExtraData(); return false; } return true; } bool ParseValue(json::Value* out) { ctx_.SkipSpaces(); // record start pos for cases where we might need to reset // current position for better error message auto start_pos = ctx_.GetCurrentPos(); // check if the end of the string is reached switch (ctx_.Peek()) { case -1: { ctx_.SetErrorExpectingValue(); return false; } case '{': { return ParseObject(out); } case '[': { return ParseArray(out); } case '\"': { return ctx_.NextString(out); } case 't': { ctx_.SkipNextAssumeNoSpace(); if (ctx_.MatchLiteral("rue", 3)) { *out = true; return true; } else { ctx_.SetCurrentPosForBetterErrorMsg(start_pos); ctx_.SetErrorExpectingValue(); return false; } } case 'f': { ctx_.SkipNextAssumeNoSpace(); if (ctx_.MatchLiteral("alse", 4)) { *out = false; return true; } else { ctx_.SetCurrentPosForBetterErrorMsg(start_pos); ctx_.SetErrorExpectingValue(); return false; } } case 'n': { ctx_.SkipNextAssumeNoSpace(); if (ctx_.MatchLiteral("ull", 3)) { *out = nullptr; return true; } else { ctx_.SetCurrentPosForBetterErrorMsg(start_pos); ctx_.SetErrorExpectingValue(); return false; } } default: { return ctx_.NextNumber(out); } } return false; } bool ParseObject(json::Value* out) { size_t stack_top = object_temp_stack_.size(); json::Object result; ctx_.SkipNextAssumeNoSpace(); ctx_.SkipSpaces(); int next_char = ctx_.Peek(); if (next_char == -1) { ctx_.SetErrorExpectingPropertyName(); return false; } // empty object if (next_char == '}') { ctx_.SkipNextAssumeNoSpace(); *out = json::Object(); return true; } // non-empty object while ((next_char = ctx_.Peek()) != -1) { if (next_char != '\"') { ctx_.SetErrorExpectingPropertyName(); return false; } json::Value key; if (!ctx_.NextString(&key)) return false; ctx_.SkipSpaces(); if (ctx_.Peek() != ':') { ctx_.SetErrorExpectingColon(); return false; } ctx_.SkipNextAssumeNoSpace(); json::Value value; if (!ParseValue(&value)) return false; object_temp_stack_.emplace_back(key, value); // result.Set(key, value); ctx_.SkipSpaces(); if (ctx_.Peek() == '}') { ctx_.SkipNextAssumeNoSpace(); *out = json::Object(object_temp_stack_.begin() + static_cast(stack_top), object_temp_stack_.end()); // recover the stack to original state object_temp_stack_.resize(stack_top); return true; } else if (ctx_.Peek() == ',') { ctx_.SkipNextAssumeNoSpace(); // must skip space so next iteration do not have to do so ctx_.SkipSpaces(); } else { ctx_.SetErrorExpectingComma(); return false; } } return false; } bool ParseArray(json::Value* out) { size_t stack_top = array_temp_stack_.size(); ctx_.SkipNextAssumeNoSpace(); ctx_.SkipSpaces(); int next_char = ctx_.Peek(); if (next_char == -1) { ctx_.SetErrorExpectingValue(); return false; } // empty array if (next_char == ']') { ctx_.SkipNextAssumeNoSpace(); *out = json::Array(); return true; } // non-empty array while ((next_char = ctx_.Peek()) != -1) { // NOLINT(clang-analyzer-deadcode.DeadStores) json::Value value; // no need to skip space here because we already skipped space // at the beginning or in previous iteration if (!ParseValue(&value)) return false; array_temp_stack_.emplace_back(std::move(value)); ctx_.SkipSpaces(); next_char = ctx_.Peek(); if (next_char == ',') { ctx_.SkipNextAssumeNoSpace(); // must skip space so next iteration do not have to do so ctx_.SkipSpaces(); } else if (next_char == ']') { ctx_.SkipNextAssumeNoSpace(); *out = json::Array(array_temp_stack_.begin() + static_cast(stack_top), array_temp_stack_.end()); // recover the stack array_temp_stack_.resize(stack_top); return true; } else { ctx_.SetErrorExpectingComma(); return false; } } return false; } String storage_; JSONParserContext ctx_; // Temp stack for intermediate values // we first create a persistent stack to store the parsed values // then create the final array/object object with the precise size std::vector array_temp_stack_; std::vector> object_temp_stack_; }; json::Value Parse(const String& json_str, String* error_msg) { return JSONParser::Parse(json_str, error_msg); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("ffi.json.Parse", [](const String& json_str) { return json::Parse(json_str); }); } } // namespace json } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/json_writer.cc000066400000000000000000000211721521067262500205060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/json/writer.cc * * \brief A minimalistic JSON writer based on ffi values. */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace json { class JSONWriter { public: // NOLINTNEXTLINE(performance-unnecessary-value-param) static String Stringify(const json::Value& value, Optional indent) { JSONWriter writer(indent.value_or(0)); writer.WriteValue(value); return String(std::move(writer.result_)); } private: explicit JSONWriter(int indent) : indent_(indent), out_iter_(result_) {} static bool FastMathSafeIsNaN(double x) { #ifdef __FAST_MATH__ // Bit-level NaN detection (IEEE 754 double) // IEEE 754 standard: https://en.wikipedia.org/wiki/IEEE_754 // NaN is encoded as all 1s in the exponent and non-zero in the mantissa static_assert(sizeof(double) == sizeof(uint64_t), "Unexpected double size"); union { double from; uint64_t to; } u; u.from = x; // write "from", read "to" uint64_t bits = u.to; uint64_t exponent = (bits >> 52) & 0x7FF; uint64_t mantissa = bits & 0xFFFFFFFFFFFFFull; return (exponent == 0x7FF) && (mantissa != 0); #else // Safe to use std::isnan when fast-math is off return std::isnan(x); #endif } static bool FastMathSafeIsInf(double x) { #ifdef __FAST_MATH__ // IEEE 754 standard: https://en.wikipedia.org/wiki/IEEE_754 // Inf is encoded as all 1s in the exponent and zero in the mantissa static_assert(sizeof(double) == sizeof(uint64_t), "Unexpected double size"); union { double from; uint64_t to; } u; u.from = x; // write "from", read "to" uint64_t bits = u.to; uint64_t exponent = (bits >> 52) & 0x7FF; uint64_t mantissa = bits & 0xFFFFFFFFFFFFFull; // inf is encoded as all 1s in the exponent and zero in the mantissa return (exponent == 0x7FF) && (mantissa == 0); #else return std::isinf(x); #endif } void WriteValue(const json::Value& value) { switch (value.type_index()) { case TypeIndex::kTVMFFINone: { WriteLiteral("null", 4); break; } case TypeIndex::kTVMFFIBool: { bool bool_value = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); if (bool_value) { WriteLiteral("true", 4); } else { WriteLiteral("false", 5); } break; } case TypeIndex::kTVMFFIInt: { WriteInt(details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIFloat: { WriteFloat(details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFISmallStr: case TypeIndex::kTVMFFIStr: { WriteString(details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIArray: { WriteArray(details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIList: { WriteList(details::AnyUnsafe::CopyFromAnyViewAfterCheck>(value)); break; } case TypeIndex::kTVMFFIMap: { WriteObject(details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } default: { TVM_FFI_THROW(ValueError) << "Unsupported type: `" << value.GetTypeKey() << "`"; TVM_FFI_UNREACHABLE(); } } } void WriteLiteral(const char* literal, int size) { for (int i = 0; i < size; ++i) { *out_iter_++ = literal[i]; } } void WriteInt(int64_t value) { // the biggest possible string representation of -INT64_MIN char buffer[sizeof("-9223372036854775808") + 1]; int size = TVM_FFI_SNPRINTF(buffer, sizeof(buffer), "%" PRId64, value); WriteLiteral(buffer, size); } void WriteFloat(double value) { // largest possible string representation of a double is around 24 chars plus // one null terminator keep 32 to be safe char buffer[32]; if (FastMathSafeIsNaN(value)) { WriteLiteral("NaN", 3); } else if (FastMathSafeIsInf(value)) { if (value < 0) { WriteLiteral("-Infinity", 9); } else { WriteLiteral("Infinity", 8); } } else { double int_part; // if the value can be represented as integer if (std::fabs(value) < (1ULL << 53) && std::modf(value, &int_part) == 0) { // always print an extra .0 for integer so integer numbers are printed as floats // this helps us to distinguish between integer and float, which is not necessary // but helps to ensure roundtrip property of the parser/printer in terms of int/float types int size = TVM_FFI_SNPRINTF(buffer, sizeof(buffer), "%.1f", int_part); WriteLiteral(buffer, size); } else { // Save 17 decimal digits to avoid loss during loading JSON // this is the maximum precision that can be represented in a double int size = TVM_FFI_SNPRINTF(buffer, sizeof(buffer), "%.17g", value); WriteLiteral(buffer, size); } } } void WriteString(const String& value) { String escaped = EscapeStringJSON(value); std::copy(escaped.data(), escaped.data() + escaped.size(), out_iter_); } void WriteArray(const json::Array& value) { WriteSequence(value); } void WriteList(const ffi::List& value) { const void* ptr = static_cast(value.get()); if (active_lists_.count(ptr)) { TVM_FFI_THROW(ValueError) << "Cycle detected: List contains itself"; } active_lists_.insert(ptr); WriteSequence(value); active_lists_.erase(ptr); } template void WriteSequence(const SeqType& value) { *out_iter_++ = '['; if (indent_ != 0) { total_indent_ += indent_; } for (size_t i = 0; i < value.size(); ++i) { if (i != 0) { *out_iter_++ = ','; } if (indent_ != 0) { WriteIndent(); } WriteValue(value[static_cast(i)]); } if (indent_ != 0) { total_indent_ -= indent_; WriteIndent(); } *out_iter_++ = ']'; } void WriteObject(const json::Object& value) { *out_iter_++ = '{'; if (indent_ != 0) { total_indent_ += indent_; } int counter = 0; for (const auto& [key, value] : value) { if (counter++ != 0) { *out_iter_++ = ','; } if (indent_ != 0) { WriteIndent(); } auto opt_key = key.as(); if (!opt_key.has_value()) { TVM_FFI_THROW(ValueError) << "Expect key to be string, got `" << key.GetTypeKey() << "`"; } WriteString(*opt_key); *out_iter_++ = ':'; if (indent_ != 0) { *out_iter_++ = ' '; } WriteValue(value); } if (indent_ != 0) { total_indent_ -= indent_; WriteIndent(); } *out_iter_++ = '}'; } // Write a newline and indent the current level void WriteIndent() { *out_iter_++ = '\n'; for (int i = 0; i < total_indent_; ++i) { *out_iter_++ = ' '; } } int indent_ = 0; int total_indent_ = 0; std::string result_; std::back_insert_iterator out_iter_; std::unordered_set active_lists_; }; String Stringify(const json::Value& value, Optional indent) { return JSONWriter::Stringify(value, indent); // NOLINT(performance-unnecessary-value-param) } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("ffi.json.Stringify", Stringify); } } // namespace json } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/library_module.cc000066400000000000000000000210451521067262500211510ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/library_module.cc * * \brief Library module implementation. */ #include #include #include #include #include "buffer_stream.h" #include "module_internal.h" namespace tvm { namespace ffi { class LibraryModuleObj final : public ModuleObj { public: explicit LibraryModuleObj(ObjectPtr lib) : lib_(std::move(lib)) {} const char* kind() const final { return "library"; } /*! \brief Get the property of the runtime module .*/ int GetPropertyMask() const final { return Module::kBinarySerializable | Module::kRunnable; }; Optional GetFunction(const String& name) final { TVMFFISafeCallType faddr; faddr = reinterpret_cast(lib_->GetSymbolWithSymbolPrefix(name)); // ensure the function keeps the Library Module alive Module self_strong_ref = GetRef(this); if (faddr != nullptr) { return ffi::Function::FromPacked([faddr, self_strong_ref](ffi::PackedArgs args, ffi::Any* rv) { TVM_FFI_ICHECK_LT(rv->type_index(), ffi::TypeIndex::kTVMFFIStaticObjectBegin); TVM_FFI_CHECK_SAFE_CALL((*faddr)(nullptr, reinterpret_cast(args.data()), args.size(), reinterpret_cast(rv))); }); } return std::nullopt; } Optional GetFunctionMetadata(const String& name) final { // Look for __tvm_ffi__metadata_ symbol String metadata_symbol = symbol::tvm_ffi_metadata_prefix + name; void* symbol = lib_->GetSymbol(metadata_symbol); if (symbol != nullptr) { using MetadataGetter = int (*)(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); auto metadata_getter = reinterpret_cast(symbol); return Function::InvokeExternC(nullptr, metadata_getter).cast(); } return std::nullopt; } Optional GetFunctionDoc(const String& name) final { // Look for __tvm_ffi__doc_ symbol String doc_symbol = symbol::tvm_ffi_doc_prefix + name; void* symbol = lib_->GetSymbol(doc_symbol); if (symbol != nullptr) { using DocGetter = int (*)(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); auto doc_getter = reinterpret_cast(symbol); return Function::InvokeExternC(nullptr, doc_getter).cast(); } return std::nullopt; } private: ObjectPtr lib_; }; Module LoadModuleFromBytes(const std::string& kind, const Bytes& bytes) { std::string loader_key = "ffi.Module.load_from_bytes." + kind; const auto floader = tvm::ffi::Function::GetGlobal(loader_key); if (!floader.has_value()) { TVM_FFI_THROW(RuntimeError) << "Library binary was created using {" << kind << "} but a loader of that name is not registered. " << "Make sure to have runtime that registers " << loader_key; } return (*floader)(bytes).cast(); } /*! * \brief Process libary binary to recover binary-serialized modules * \param library_bin The binary embedded in the library. * \param opt_lib The library, can be nullptr in which case we expect to deserialize * all binary-serialized modules * \param library_ctx_addr the pointer to library module as ctx addr * \return the root module * */ Module ProcessLibraryBin(const char* library_bin, const ObjectPtr& opt_lib, void** library_ctx_addr = nullptr) { // Layout of the library binary: // [] [] ... // key can be: "_lib", or a module kind // - "_lib" indicate this location places the library module // - other keys are module kind, in which case [val: bytes] contains // the serialized bytes from the custom module // Import tree structure (CSR structure of child indices): // = > > TVM_FFI_ICHECK(library_bin != nullptr); uint64_t nbytes = 0; for (size_t i = 0; i < sizeof(nbytes); ++i) { uint64_t c = static_cast(library_bin[i]); nbytes |= (c & 0xffUL) << (i * 8); } BufferInStream stream(library_bin + sizeof(nbytes), static_cast(nbytes)); std::vector import_tree_indptr; std::vector import_tree_child_indices; TVM_FFI_ICHECK(stream.Read(&import_tree_indptr)); TVM_FFI_ICHECK(stream.Read(&import_tree_child_indices)); size_t num_modules = import_tree_indptr.size() - 1; std::vector modules; modules.reserve(num_modules); for (uint64_t i = 0; i < num_modules; ++i) { std::string kind; TVM_FFI_ICHECK(stream.Read(&kind)); // "_lib" serves as a placeholder in the module import tree to indicate where // to place the DSOModule if (kind == "_lib") { TVM_FFI_ICHECK(opt_lib != nullptr) << "_lib is not allowed during module serialization"; auto lib_mod_ptr = make_object(opt_lib); if (library_ctx_addr) { *library_ctx_addr = lib_mod_ptr.get(); } modules.emplace_back(lib_mod_ptr); } else { std::string module_bytes; TVM_FFI_ICHECK(stream.Read(&module_bytes)); Module m = LoadModuleFromBytes(kind, Bytes(module_bytes)); modules.emplace_back(m); } } for (size_t i = 0; i < modules.size(); ++i) { for (size_t j = import_tree_indptr[i]; j < import_tree_indptr[i + 1]; ++j) { Array* module_imports = ModuleObj::InternalUnsafe::GetImports(modules[i].operator->()); auto child_index = import_tree_child_indices[j]; TVM_FFI_ICHECK(child_index < modules.size()); module_imports->emplace_back(modules[child_index]); } } return modules[0]; } // registry to store context symbols class ContextSymbolRegistry { public: void InitContextSymbols(const ObjectPtr& lib) { for (const auto& [name, symbol] : context_symbols_) { if (void** symbol_addr = reinterpret_cast(lib->GetSymbol(name))) { *symbol_addr = symbol; } } } void VisitContextSymbols(const ffi::TypedFunction& callback) { for (const auto& [name, symbol] : context_symbols_) { callback(name, symbol); } } void Register(String name, void* symbol) { context_symbols_.emplace_back(std::move(name), symbol); } static ContextSymbolRegistry* Global() { static ContextSymbolRegistry* inst = new ContextSymbolRegistry(); return inst; } private: std::vector> context_symbols_; }; void Module::VisitContextSymbols(const ffi::TypedFunction& callback) { ContextSymbolRegistry::Global()->VisitContextSymbols(callback); } Module CreateLibraryModule(ObjectPtr lib) { // NOLINT(performance-unnecessary-value-param) const char* library_bin = reinterpret_cast(lib->GetSymbol(ffi::symbol::tvm_ffi_library_bin)); void** library_ctx_addr = reinterpret_cast(lib->GetSymbol(ffi::symbol::tvm_ffi_library_ctx)); ContextSymbolRegistry::Global()->InitContextSymbols(lib); if (library_bin != nullptr) { // we have embedded binaries that needs to be deserialized return ProcessLibraryBin(library_bin, lib, library_ctx_addr); } else { // Only have one single DSO Module auto lib_mod_ptr = make_object(lib); Module root_mod = Module(lib_mod_ptr); if (library_ctx_addr) { *library_ctx_addr = root_mod.operator->(); } return root_mod; } } } // namespace ffi } // namespace tvm int TVMFFIEnvModRegisterContextSymbol(const char* name, void* symbol) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::String s_name(name); tvm::ffi::ContextSymbolRegistry::Global()->Register(s_name, symbol); TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/extra/library_module_dynamic_lib.cc000066400000000000000000000064151521067262500235070ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file library_module_dynamic_lib.cc * \brief Create library module to load from dynamic shared library. */ #include #include #include #include "module_internal.h" #if defined(_WIN32) #include #else #include #endif #if defined(__hexagon__) extern "C" { #include } #endif namespace tvm { namespace ffi { class DSOLibrary final : public Library { public: explicit DSOLibrary(const String& name) { Load(name); } ~DSOLibrary() final { if (lib_handle_) Unload(); } void* GetSymbol(const String& name) final { return GetSymbol_(name.c_str()); } private: // private system dependent implementation void* GetSymbol_(const char* name); void Load(const String& name); void Unload(); #if defined(_WIN32) //! \brief Windows library handle HMODULE lib_handle_{nullptr}; #else // \brief Linux library handle void* lib_handle_{nullptr}; #endif }; #if defined(_WIN32) void* DSOLibrary::GetSymbol_(const char* name) { return reinterpret_cast(GetProcAddress(lib_handle_, (LPCSTR)name)); // NOLINT(*) } void DSOLibrary::Load(const String& name) { // use wstring version that is needed by LLVM. std::wstring wname(name.data(), name.data() + name.size()); lib_handle_ = LoadLibraryW(wname.c_str()); TVM_FFI_ICHECK(lib_handle_ != nullptr) << "Failed to load dynamic shared library " << name; } void DSOLibrary::Unload() { FreeLibrary(lib_handle_); lib_handle_ = nullptr; } #else void DSOLibrary::Load(const String& name) { lib_handle_ = dlopen(name.c_str(), RTLD_LAZY | RTLD_LOCAL); TVM_FFI_ICHECK(lib_handle_ != nullptr) << "Failed to load dynamic shared library " << name << " " << dlerror(); #if defined(__hexagon__) int p; int rc = dlinfo(lib_handle_, RTLD_DI_LOAD_ADDR, &p); if (rc) FARF(ERROR, "error getting model .so start address : %u", rc); else FARF(ALWAYS, "Model .so Start Address : %x", p); #endif } void* DSOLibrary::GetSymbol_(const char* name) { return dlsym(lib_handle_, name); } void DSOLibrary::Unload() { dlclose(lib_handle_); lib_handle_ = nullptr; } #endif TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("ffi.Module.load_from_file.so", [](const String& library_path, const String&) { return CreateLibraryModule(make_object(library_path)); }); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/library_module_system_lib.cc000066400000000000000000000106241521067262500234040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file system_library.cc * \brief Create library module that directly get symbol from the system lib. */ #include #include #include #include #include #include #include #include "module_internal.h" namespace tvm { namespace ffi { class SystemLibSymbolRegistry { public: void RegisterSymbol(const std::string& name, void* ptr) { auto it = symbol_table_.find(name); if (it != symbol_table_.end() && ptr != (*it).second) { std::cerr << "Warning:SystemLib symbol " << name << " get overriden to a different address " << ptr << "->" << (*it).second << std::endl; } symbol_table_.Set(name, ptr); } void* GetSymbol(const String& name) { auto it = symbol_table_.find(name); if (it != symbol_table_.end()) { return (*it).second; } else { return nullptr; } } static SystemLibSymbolRegistry* Global() { static SystemLibSymbolRegistry* inst = new SystemLibSymbolRegistry(); return inst; } private: // Internal symbol table Map symbol_table_; }; class SystemLibrary final : public Library { public: explicit SystemLibrary(String symbol_prefix) : symbol_prefix_(std::move(symbol_prefix)) {} void* GetSymbol(const String& name) final { // The `name` might or might not already contain the symbol prefix. // Therefore, we check both with and without the prefix. String name_with_prefix = symbol_prefix_ + name; void* symbol = reg_->GetSymbol(name_with_prefix); if (symbol != nullptr) { return symbol; } return reg_->GetSymbol(name); } void* GetSymbolWithSymbolPrefix(const String& name) final { // The `name` might or might not already contain the symbol prefix. // Therefore, we check both with and without the prefix. String name_with_prefix = symbol::tvm_ffi_symbol_prefix + symbol_prefix_ + name; void* symbol = reg_->GetSymbol(name_with_prefix); if (symbol != nullptr) { return symbol; } name_with_prefix = symbol::tvm_ffi_symbol_prefix + name; return reg_->GetSymbol(name_with_prefix); } private: SystemLibSymbolRegistry* reg_ = SystemLibSymbolRegistry::Global(); String symbol_prefix_; }; class SystemLibModuleRegistry { public: Module GetOrCreateModule(const String& symbol_prefix) { std::scoped_lock lock(mutex_); auto it = lib_map_.find(symbol_prefix); if (it != lib_map_.end()) { return (*it).second; } else { Module mod = CreateLibraryModule(make_object(symbol_prefix)); lib_map_.Set(symbol_prefix, mod); return mod; } } static SystemLibModuleRegistry* Global() { static SystemLibModuleRegistry* inst = new SystemLibModuleRegistry(); return inst; } private: // Internal mutex std::mutex mutex_; // maps prefix to the library module // we need to make sure each lib map have an unique // copy through out the entire lifetime of the process Map lib_map_; }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_packed("ffi.SystemLib", [](ffi::PackedArgs args, ffi::Any* rv) { String symbol_prefix = ""; if (args.size() != 0) { symbol_prefix = args[0].cast(); } *rv = SystemLibModuleRegistry::Global()->GetOrCreateModule(symbol_prefix); }); } } // namespace ffi } // namespace tvm int TVMFFIEnvModRegisterSystemLibSymbol(const char* name, void* ptr) { tvm::ffi::SystemLibSymbolRegistry::Global()->RegisterSymbol(name, ptr); return 0; } tvm-ffi-0.1.12/src/ffi/extra/module.cc000066400000000000000000000160401521067262500174240ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include "module_internal.h" namespace tvm { namespace ffi { /*! * \brief Global modules, i.e. modules that are owned by the runtime and should not be unloaded. * On the frontend, a module is added to the registry if `keep_alive=True` when `load_module` is * called. */ class ModuleGlobals { public: void Add(const Module& m) { std::scoped_lock lock(mutex_); modules_.Set(m, 1); } void Remove(const Module& m) { std::scoped_lock lock(mutex_); modules_.erase(m); } static ModuleGlobals* Get() { static ModuleGlobals instance; return &instance; } private: Map modules_; std::mutex mutex_; }; Optional ModuleObj::GetFunction(const String& name, bool query_imports) { if (auto opt_func = this->GetFunction(name)) { return opt_func; } if (query_imports) { for (const Any& import : imports_) { if (auto opt_func = import.cast()->GetFunction(name, query_imports)) { return *opt_func; } } } return std::nullopt; } Optional ModuleObj::GetFunctionMetadata(const String& name, bool query_imports) { if (auto opt_metadata = this->GetFunctionMetadata(name)) { return opt_metadata; } if (query_imports) { for (const Any& import : imports_) { if (auto opt_metadata = import.cast()->GetFunctionMetadata(name, query_imports)) { return *opt_metadata; } } } return std::nullopt; } Optional ModuleObj::GetFunctionDoc(const String& name, bool query_imports) { if (auto opt_str = this->GetFunctionDoc(name)) { return opt_str; } if (query_imports) { for (const Any& import : imports_) { if (auto opt_str = import.cast()->GetFunctionDoc(name, query_imports)) { return *opt_str; } } } return std::nullopt; } void ModuleObj::ImportModule(const Module& other) { std::unordered_set visited{other.operator->()}; std::vector stack{other.operator->()}; while (!stack.empty()) { const ModuleObj* n = stack.back(); stack.pop_back(); for (const Any& m : n->imports_) { const ModuleObj* next = m.cast(); if (visited.count(next)) continue; visited.insert(next); stack.push_back(next); } } if (visited.count(this)) { TVM_FFI_THROW(RuntimeError) << "Cyclic dependency detected during import"; } imports_.push_back(other); } void ModuleObj::ClearImports() { imports_.clear(); } bool ModuleObj::ImplementsFunction(const String& name, bool query_imports) { if (this->ImplementsFunction(name)) { return true; } if (query_imports) { for (const Any& import : imports_) { if (import.cast()->ImplementsFunction(name, query_imports)) { return true; } } } return false; } Module Module::LoadFromFile(const String& file_name) { String format = [&file_name]() -> String { const char* data = file_name.data(); for (size_t i = file_name.size(); i > 0; i--) { if (data[i - 1] == '.') { return String(data + i, file_name.size() - i); } } TVM_FFI_THROW(RuntimeError) << "Failed to get file format from " << file_name; TVM_FFI_UNREACHABLE(); }(); if (format == "dll" || format == "dylib" || format == "dso") { format = "so"; } String loader_name = "ffi.Module.load_from_file." + format; const auto floader = tvm::ffi::Function::GetGlobal(loader_name); if (!floader.has_value()) { TVM_FFI_THROW(RuntimeError) << "Loader for `." << format << "` files is not registered," << " resolved to (" << loader_name << ") in the global registry." << "Ensure that you have loaded the correct runtime code, and" << "that you are on the correct hardware architecture."; } return (*floader)(file_name, format).cast(); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; ModuleObj::InternalUnsafe::RegisterReflection(); refl::GlobalDef() .def("ffi.ModuleLoadFromFile", &Module::LoadFromFile) .def_method("ffi.ModuleImplementsFunction", [](const Module& mod, const String& name, bool query_imports) { return mod->ImplementsFunction(name, query_imports); }) .def_method("ffi.ModuleGetFunctionMetadata", [](const Module& mod, const String& name, bool query_imports) { return mod->GetFunctionMetadata(name, query_imports); }) .def_method("ffi.ModuleGetFunctionDoc", [](const Module& mod, const String& name, bool query_imports) { return mod->GetFunctionDoc(name, query_imports); }) .def_method("ffi.ModuleGetFunction", [](const Module& mod, const String& name, bool query_imports) { return mod->GetFunction(name, query_imports); }) .def_method("ffi.ModuleGetPropertyMask", &ModuleObj::GetPropertyMask) .def_method("ffi.ModuleInspectSource", &ModuleObj::InspectSource) .def_method("ffi.ModuleGetKind", [](const Module& mod) -> String { return mod->kind(); }) .def_method("ffi.ModuleGetWriteFormats", &ModuleObj::GetWriteFormats) .def_method("ffi.ModuleWriteToFile", &ModuleObj::WriteToFile) .def_method("ffi.ModuleImportModule", &ModuleObj::ImportModule) .def_method("ffi.ModuleClearImports", &ModuleObj::ClearImports) .def_method("ffi.ModuleGlobalsAdd", [](const Module& mod) { ModuleGlobals::Get()->Add(mod); }) .def_method("ffi.ModuleGlobalsRemove", [](const Module& mod) { ModuleGlobals::Get()->Remove(mod); }); } } // namespace ffi } // namespace tvm int TVMFFIEnvModLookupFromImports(TVMFFIObjectHandle library_ctx, const char* func_name, TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); *out = tvm::ffi::ModuleObj::InternalUnsafe::GetFunctionFromImports( reinterpret_cast(library_ctx), func_name); TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/extra/module_internal.h000066400000000000000000000073061521067262500211670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file library_module.h * \brief Module that builds from a libary of symbols. */ #ifndef TVM_FFI_EXTRA_MODULE_INTERNAL_H_ #define TVM_FFI_EXTRA_MODULE_INTERNAL_H_ #include #include #include namespace tvm { namespace ffi { /*! * \brief Library is the common interface * for storing data in the form of shared libaries. * * \sa src/ffi/extra/dso_library.cc * \sa src/ffi/extra/system_library.cc */ class Library : public Object { public: // destructor. virtual ~Library() = default; /*! * \brief Get the symbol address for a given name. * \param name The name of the symbol. * \return The symbol. */ virtual void* GetSymbol(const String& name) = 0; /*! * \brief Get the symbol address for a given name with the tvm ffi symbol prefix. * \param name The name of the symbol. * \return The symbol. * \note This function will be overloaded by systemlib implementation. */ virtual void* GetSymbolWithSymbolPrefix(const String& name) { String name_with_prefix = symbol::tvm_ffi_symbol_prefix + name; return GetSymbol(name_with_prefix); } // NOTE: we do not explicitly create an type index and type_key here for libary. // This is because we do not need dynamic type downcasting and only need to use the refcounting }; struct ModuleObj::InternalUnsafe { static Array* GetImports(ModuleObj* module) { return &(module->imports_); } static void* GetFunctionFromImports(ModuleObj* module, const char* name) { // backend implementation for TVMFFIEnvModLookupFromImports static std::mutex mutex_; std::scoped_lock lock(mutex_); String s_name(name); auto it = module->import_lookup_cache_.find(s_name); if (it != module->import_lookup_cache_.end()) { return const_cast((*it).second.operator->()); } auto opt_func = [&]() -> std::optional { for (const Any& import : module->imports_) { if (auto opt_func = import.cast()->GetFunction(s_name, true)) { return *opt_func; } } // try global at last return tvm::ffi::Function::GetGlobal(s_name); }(); if (!opt_func.has_value()) { TVM_FFI_THROW(RuntimeError) << "Cannot find function " << name << " in the imported modules or global registry."; } module->import_lookup_cache_.Set(s_name, *opt_func); return const_cast((*opt_func).operator->()); } static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef().def_ro("imports_", &ModuleObj::imports_); } }; /*! * \brief Create a library module from a given library. * * \param lib The library. * * \return The corresponding loaded module. */ Module CreateLibraryModule(ObjectPtr lib); } // namespace ffi } // namespace tvm #endif // TVM_FFI_EXTRA_MODULE_INTERNAL_H_ tvm-ffi-0.1.12/src/ffi/extra/reflection_extra.cc000066400000000000000000000207331521067262500215000ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/reflection_extra.cc * * \brief Extra reflection registrations. * */ #include #include #include #include #include namespace tvm { namespace ffi { namespace reflection { void MakeObjectFromPackedArgs(ffi::PackedArgs args, Any* ret) { int32_t type_index; if (auto opt_type_index = args[0].try_cast()) { type_index = *opt_type_index; } else { String type_key = args[0].cast(); TVMFFIByteArray type_key_array = TVMFFIByteArray{type_key.data(), type_key.size()}; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_array, &type_index)); } TVM_FFI_ICHECK(args.size() % 2 == 1); const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index); ObjectPtr ptr = CreateEmptyObject(type_info); std::vector keys; std::vector keys_found; for (int i = 1; i < args.size(); i += 2) { keys.push_back(args[i].cast()); } keys_found.resize(keys.size(), false); auto search_field = [&](const TVMFFIByteArray& field_name) { for (size_t i = 0; i < keys.size(); ++i) { if (keys_found[i]) continue; if (keys[i].compare(field_name) == 0) { return i; } } return keys.size(); }; auto update_fields = [&](const TVMFFITypeInfo* tinfo) { for (int i = 0; i < tinfo->num_fields; ++i) { const TVMFFIFieldInfo* field_info = tinfo->fields + i; size_t arg_index = search_field(field_info->name); void* field_addr = reinterpret_cast(ptr.get()) + field_info->offset; if (arg_index < keys.size()) { AnyView field_value = args[static_cast(arg_index * 2 + 2)]; reflection::CallFieldSetter(field_info, field_addr, reinterpret_cast(&field_value)); keys_found[arg_index] = true; } else if (field_info->flags & kTVMFFIFieldFlagBitMaskHasDefault) { reflection::SetFieldToDefault(field_info, field_addr); } else { TVM_FFI_THROW(TypeError) << "Required field `" << String(field_info->name.data, field_info->name.size) << "` not set in type `" << TypeIndexToTypeKey(type_index) << "`"; } } }; // iterate through acenstors in parent to child order // skip the first one since it is always the root object for (int i = 1; i < type_info->type_depth; ++i) { update_fields(type_info->type_ancestors[i]); } update_fields(type_info); for (size_t i = 0; i < keys.size(); ++i) { if (!keys_found[i]) { TVM_FFI_THROW(TypeError) << "Type `" << TypeIndexToTypeKey(type_index) << "` does not have field `" << keys[i] << "`"; } } *ret = ObjectRef(ptr); } inline void AccessStepRegisterReflection() { // register access step reflection here since it is only needed for bindings namespace refl = tvm::ffi::reflection; refl::ObjectDef(refl::init(false)) .def_ro("kind", &AccessStepObj::kind) .def_ro("key", &AccessStepObj::key); // Register __ffi_repr__ for AccessStep: format one step fragment. // kAttr -> ".name" // kArrayItem -> "[index]" // kMapItem -> "[]" (string keys are quoted via fn_repr) // kAttrMissing -> "[]" // kArrayItemMissing -> "[]" // kMapItemMissing -> "[>]" refl::TypeAttrDef().def( refl::type_attr::kRepr, [](const AccessStep& step, const ffi::Function& fn_repr) -> ffi::String { std::ostringstream os; switch (step->kind) { case AccessKind::kAttr: os << "." << step->key.cast(); break; case AccessKind::kArrayItem: os << "[" << step->key.cast() << "]"; break; case AccessKind::kMapItem: os << "[" << fn_repr(step->key).cast() << "]"; break; case AccessKind::kAttrMissing: os << "[key).cast() << ">]"; break; case AccessKind::kArrayItemMissing: os << "[key.cast() << ">]"; break; case AccessKind::kMapItemMissing: os << "[key).cast() << ">]"; break; default: TVM_FFI_UNREACHABLE(); } return os.str(); }); } inline void AccessPathRegisterReflection() { // register access path reflection here since it is only needed for bindings namespace refl = tvm::ffi::reflection; refl::ObjectDef(refl::init(false)) .def_ro("parent", &AccessPathObj::parent) .def_ro("step", &AccessPathObj::step) .def_ro("depth", &AccessPathObj::depth) .def_static("_root", &AccessPath::Root) .def("_extend", &AccessPathObj::Extend) .def("_attr", &AccessPathObj::Attr) .def("_array_item", &AccessPathObj::ArrayItem) .def("_map_item", &AccessPathObj::MapItem) .def("_attr_missing", &AccessPathObj::AttrMissing) .def("_array_item_missing", &AccessPathObj::ArrayItemMissing) .def("_map_item_missing", &AccessPathObj::MapItemMissing) .def("_is_prefix_of", &AccessPathObj::IsPrefixOf) .def("_to_steps", &AccessPathObj::ToSteps) .def("_path_equal", [](const AccessPath& self, const AccessPath& other) { return self->PathEqual(other); }); // Register __ffi_repr__ for AccessPath: flatten via ToSteps() and walk the resulting vector. // Root (depth == 0) emits ""; non-root nodes concatenate each step's fragment. refl::TypeAttrDef().def( refl::type_attr::kRepr, [](const AccessPath& path, const ffi::Function& fn_repr) -> ffi::String { Array steps = path->ToSteps(); std::ostringstream os; os << ""; for (const AccessStep& step : steps) { os << fn_repr(step).cast(); } return os.str(); }); } int64_t StructuralKeyHash(const Any& key) { const auto* key_obj = key.as(); TVM_FFI_ICHECK_NOTNULL(key_obj); return key_obj->hash_i64; } bool StructuralKeyEqual(const Any& lhs, const Any& rhs) { if (lhs.same_as(rhs)) { return true; } const auto* lhs_obj = lhs.as(); const auto* rhs_obj = rhs.as(); TVM_FFI_ICHECK_NOTNULL(lhs_obj); TVM_FFI_ICHECK_NOTNULL(rhs_obj); if (lhs_obj->hash_i64 != rhs_obj->hash_i64) { return false; } return StructuralEqual::Equal(lhs_obj->key, rhs_obj->key); } inline void StructuralKeyRegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def(refl::init(), "Constructor that takes any key and computes its structural hash") .def_ro("key", &StructuralKeyObj::key) .def_ro("hash_i64", &StructuralKeyObj::hash_i64); refl::TypeAttrDef() .attr(type_attr::kAnyHash, reinterpret_cast(&StructuralKeyHash)) .attr(type_attr::kAnyEqual, reinterpret_cast(&StructuralKeyEqual)); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; AccessStepRegisterReflection(); AccessPathRegisterReflection(); StructuralKeyRegisterReflection(); refl::GlobalDef() .def_packed("ffi.MakeObjectFromPackedArgs", MakeObjectFromPackedArgs) .def("ffi.StructuralKey", [](Any key) { return StructuralKey(std::move(key)); }) .def("ffi.StructuralKeyEqual", StructuralKeyEqual); } } // namespace reflection } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/serialization.cc000066400000000000000000000425741521067262500210270ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/serialization.cc * * \brief Reflection-based serialization utilities. */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { class ObjectGraphSerializer { public: static json::Value Serialize(const Any& value, const Any& metadata) { ObjectGraphSerializer serializer; json::Object result; result.Set("root_index", serializer.GetOrCreateNodeIndex(value)); result.Set("nodes", std::move(serializer.nodes_)); if (metadata != nullptr) { result.Set("metadata", metadata); } return result; } private: ObjectGraphSerializer() = default; int64_t GetOrCreateNodeIndex(const Any& value) { // already mapped value, return the index auto it = node_index_map_.find(value); if (it != node_index_map_.end()) { return (*it).second; } json::Object node; switch (value.type_index()) { case TypeIndex::kTVMFFINone: { node.Set("type", ffi::StaticTypeKey::kTVMFFINone); break; } case TypeIndex::kTVMFFIBool: { node.Set("type", ffi::StaticTypeKey::kTVMFFIBool); node.Set("data", details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIInt: { node.Set("type", ffi::StaticTypeKey::kTVMFFIInt); node.Set("data", details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIFloat: { node.Set("type", ffi::StaticTypeKey::kTVMFFIFloat); node.Set("data", details::AnyUnsafe::CopyFromAnyViewAfterCheck(value)); break; } case TypeIndex::kTVMFFIDataType: { DLDataType dtype = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIDataType); node.Set("data", DLDataTypeToString(dtype)); break; } case TypeIndex::kTVMFFIDevice: { DLDevice device = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIDevice); node.Set("data", json::Array{ static_cast(device.device_type), static_cast(device.device_id), }); break; } case TypeIndex::kTVMFFISmallStr: case TypeIndex::kTVMFFIStr: { String str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIStr); node.Set("data", str); break; } case TypeIndex::kTVMFFISmallBytes: case TypeIndex::kTVMFFIBytes: { Bytes bytes = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIBytes); node.Set("data", Base64Encode(bytes)); break; } case TypeIndex::kTVMFFIArray: { Array array = details::AnyUnsafe::CopyFromAnyViewAfterCheck>(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIArray); node.Set("data", CreateSequenceData(array)); break; } case TypeIndex::kTVMFFIList: { List list = details::AnyUnsafe::CopyFromAnyViewAfterCheck>(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIList); const void* list_ptr = static_cast(list.get()); if (!active_lists_.insert(list_ptr).second) { TVM_FFI_THROW(ValueError) << "Cycle detected during serialization: a List contains itself"; } node.Set("data", CreateSequenceData(list)); active_lists_.erase(list_ptr); break; } case TypeIndex::kTVMFFIMap: { Map map = details::AnyUnsafe::CopyFromAnyViewAfterCheck>(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIMap); node.Set("data", CreateMapBaseData(static_cast(map.get()))); break; } case TypeIndex::kTVMFFIDict: { Dict dict = details::AnyUnsafe::CopyFromAnyViewAfterCheck>(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIDict); const void* dict_ptr = static_cast(dict.get()); if (!active_lists_.insert(dict_ptr).second) { TVM_FFI_THROW(ValueError) << "Cycle detected during serialization: a Dict contains itself"; } node.Set("data", CreateMapBaseData(static_cast(dict.get()))); active_lists_.erase(dict_ptr); break; } case TypeIndex::kTVMFFIShape: { ffi::Shape shape = details::AnyUnsafe::CopyFromAnyViewAfterCheck(value); node.Set("type", ffi::StaticTypeKey::kTVMFFIShape); node.Set("data", Array(shape->data, shape->data + shape->size)); break; } default: { if (value.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) { // serialize type key since type index is runtime dependent node.Set("type", value.GetTypeKey()); node.Set("data", CreateObjectData(value)); } else { TVM_FFI_THROW(RuntimeError) << "Cannot serialize type `" << value.GetTypeKey() << "`"; TVM_FFI_UNREACHABLE(); } } } int64_t node_index = static_cast(nodes_.size()); nodes_.push_back(node); node_index_map_.emplace(value, node_index); return node_index; } template json::Array CreateSequenceData(const SeqType& value) { json::Array data; data.reserve(static_cast(value.size())); for (const Any& item : value) { data.push_back(GetOrCreateNodeIndex(item)); } return data; } json::Array CreateMapBaseData(const MapBaseObj* value) { json::Array data; data.reserve(static_cast(value->size()) * 2); for (const auto& [key, val] : *value) { data.push_back(GetOrCreateNodeIndex(key)); data.push_back(GetOrCreateNodeIndex(val)); } return data; } // create the data for the object, if the type has a custom data to json function, // use it. otherwise, we go over the fields and create the data. json::Value CreateObjectData(const Any& value) { static reflection::TypeAttrColumn data_to_json = reflection::TypeAttrColumn(reflection::type_attr::kDataToJson); if (data_to_json[value.type_index()] != nullptr) { return data_to_json[value.type_index()].cast()(value); } // NOTE: invariant: lhs and rhs are already the same type const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(value.type_index()); if (!HasCreator(type_info)) { TVM_FFI_THROW(TypeError) << "Type `" << String(type_info->type_key) << "` does not support ToJSONGraph " << "(no native creator or __ffi_new__ type attr)"; } const Object* obj = value.cast(); json::Object data; // go over the content and hash the fields reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* field_info) { // get the field value from both side reflection::FieldGetter getter(field_info); Any field_value = getter(obj); int field_static_type_index = field_info->field_static_type_index; String field_name(field_info->name); // for static field index that are known, we can directly set the field value. switch (field_static_type_index) { case TypeIndex::kTVMFFINone: { data.Set(field_name, nullptr); break; } case TypeIndex::kTVMFFIBool: { data.Set(field_name, details::AnyUnsafe::CopyFromAnyViewAfterCheck(field_value)); break; } case TypeIndex::kTVMFFIInt: { data.Set(field_name, details::AnyUnsafe::CopyFromAnyViewAfterCheck(field_value)); break; } case TypeIndex::kTVMFFIFloat: { data.Set(field_name, details::AnyUnsafe::CopyFromAnyViewAfterCheck(field_value)); break; } case TypeIndex::kTVMFFIDataType: { DLDataType dtype = details::AnyUnsafe::CopyFromAnyViewAfterCheck(field_value); data.Set(field_name, DLDataTypeToString(dtype)); break; } default: { // for dynamic field index, we need need to put them onto nodes int64_t node_index = GetOrCreateNodeIndex(field_value); data.Set(field_name, node_index); break; } } }); return data; } // maps the original value to the index of the node in the nodes_ array std::unordered_map node_index_map_; // records nodes that are serialized json::Array nodes_; // tracks List nodes currently being serialized (for cycle detection) std::unordered_set active_lists_; }; json::Value ToJSONGraph(const Any& value, const Any& metadata) { return ObjectGraphSerializer::Serialize(value, metadata); } class ObjectGraphDeserializer { public: static Any Deserialize(const json::Value& value) { ObjectGraphDeserializer deserializer(value); return deserializer.GetOrDecodeNode(deserializer.root_index_); } Any GetOrDecodeNode(int64_t node_index) { // already decoded null index if (node_index == decoded_null_index_) { return Any(nullptr); } // already decoded if (decoded_nodes_[node_index] != nullptr) { return decoded_nodes_[node_index]; } // now decode the node Any value = DecodeNode(node_index, nodes_[node_index].cast()); decoded_nodes_[node_index] = value; if (value == nullptr) { decoded_null_index_ = node_index; } return value; } private: Any DecodeNode(int64_t node_index, const json::Object& node) { String type_key = node["type"].cast(); TVMFFIByteArray type_key_arr{type_key.data(), type_key.length()}; int32_t type_index; TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_arr, &type_index)); switch (type_index) { case TypeIndex::kTVMFFINone: { return nullptr; } case TypeIndex::kTVMFFIBool: { return node["data"].cast(); } case TypeIndex::kTVMFFIInt: { return node["data"].cast(); } case TypeIndex::kTVMFFIFloat: { return node["data"].cast(); } case TypeIndex::kTVMFFIDataType: { return StringToDLDataType(node["data"].cast()); } case TypeIndex::kTVMFFIDevice: { Array data = node["data"].cast>(); return DLDevice{static_cast(data[0]), data[1]}; } case TypeIndex::kTVMFFIStr: { return node["data"].cast(); } case TypeIndex::kTVMFFIBytes: { return Base64Decode(node["data"].cast()); } case TypeIndex::kTVMFFIMap: { return DecodeMapLikeData>(node["data"].cast()); } case TypeIndex::kTVMFFIDict: { return DecodeMapLikeData>(node["data"].cast()); } case TypeIndex::kTVMFFIArray: { return DecodeSequenceData>(node["data"].cast()); } case TypeIndex::kTVMFFIList: { return DecodeSequenceData>(node["data"].cast()); } case TypeIndex::kTVMFFIShape: { Array data = node["data"].cast>(); return ffi::Shape(data); } default: { return DecodeObjectData(type_index, node["data"]); } } } template SeqType DecodeSequenceData(const json::Array& data) { SeqType sequence; sequence.reserve(static_cast(data.size())); for (const auto& elem : data) { sequence.push_back(GetOrDecodeNode(elem.cast())); } return sequence; } template MapType DecodeMapLikeData(const json::Array& data) { MapType result; const int64_t n = static_cast(data.size()); for (int64_t i = 0; i < n; i += 2) { int64_t key_index = data[i].cast(); int64_t value_index = data[i + 1].cast(); result.Set(GetOrDecodeNode(key_index), GetOrDecodeNode(value_index)); } return result; } Any DecodeObjectData(int32_t type_index, const json::Value& data) { static reflection::TypeAttrColumn data_from_json = reflection::TypeAttrColumn(reflection::type_attr::kDataFromJson); if (data_from_json[type_index] != nullptr) { return data_from_json[type_index].cast()(data); } // otherwise, we go over the fields and create the data. const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index); ObjectPtr ptr = CreateEmptyObject(type_info); auto decode_field_value = [&](const TVMFFIFieldInfo* field_info, const json::Value& data) -> Any { switch (field_info->field_static_type_index) { case TypeIndex::kTVMFFINone: { return nullptr; } case TypeIndex::kTVMFFIBool: { return data.cast(); } case TypeIndex::kTVMFFIInt: { return data.cast(); } case TypeIndex::kTVMFFIFloat: { return data.cast(); } case TypeIndex::kTVMFFIDataType: { return StringToDLDataType(data.cast()); } default: { return GetOrDecodeNode(data.cast()); } } }; json::Object data_object = data.cast(); reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* field_info) { String field_name(field_info->name); void* field_addr = reinterpret_cast(ptr.get()) + field_info->offset; if (data_object.count(field_name) != 0) { Any field_value = decode_field_value(field_info, data_object[field_name]); reflection::CallFieldSetter(field_info, field_addr, reinterpret_cast(&field_value)); } else if (field_info->flags & kTVMFFIFieldFlagBitMaskHasDefault) { reflection::SetFieldToDefault(field_info, field_addr); } else { TVM_FFI_THROW(TypeError) << "Required field `" << String(field_info->name.data, field_info->name.size) << "` not set in type `" << TypeIndexToTypeKey(type_index) << "`"; } }); return ObjectRef(ptr); } explicit ObjectGraphDeserializer(const json::Value& serialized) { if (!serialized.as()) { TVM_FFI_THROW(ValueError) << "Invalid JSON Object Graph, expected an object"; } json::Object encoded_object = serialized.cast(); if (encoded_object.count("root_index") == 0 || !encoded_object["root_index"].as()) { TVM_FFI_THROW(ValueError) << "Invalid JSON Object Graph, expected `root_index` integer field"; } if (encoded_object.count("nodes") == 0 || !encoded_object["nodes"].as()) { TVM_FFI_THROW(ValueError) << "Invalid JSON Object Graph, expected `nodes` array field"; } root_index_ = encoded_object["root_index"].cast(); nodes_ = encoded_object["nodes"].cast(); decoded_nodes_.resize(nodes_.size(), Any(nullptr)); } // nodes json::Array nodes_; // root index int64_t root_index_; // null index if already created int64_t decoded_null_index_{-1}; // decoded nodes std::vector decoded_nodes_; }; Any FromJSONGraph(const json::Value& value) { return ObjectGraphDeserializer::Deserialize(value); } // string version of the api Any FromJSONGraphString(const String& value) { return FromJSONGraph(json::Parse(value)); } String ToJSONGraphString(const Any& value, const Any& metadata) { return json::Stringify(ToJSONGraph(value, metadata)); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("ffi.ToJSONGraph", ToJSONGraph) .def("ffi.ToJSONGraphString", ToJSONGraphString) .def("ffi.FromJSONGraph", FromJSONGraph) .def("ffi.FromJSONGraphString", FromJSONGraphString); refl::EnsureTypeAttrColumn(refl::type_attr::kDataToJson); refl::EnsureTypeAttrColumn(refl::type_attr::kDataFromJson); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/structural_equal.cc000066400000000000000000000530731521067262500215450ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/reflection/structural_equal.cc * * \brief Structural equal implementation. */ #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /** * \brief Internal Handler class for structural equal comparison. */ class StructEqualHandler { public: StructEqualHandler() = default; bool CompareAny(ffi::Any lhs, ffi::Any rhs) { using ffi::details::AnyUnsafe; const TVMFFIAny* lhs_data = AnyUnsafe::TVMFFIAnyPtrFromAny(lhs); const TVMFFIAny* rhs_data = AnyUnsafe::TVMFFIAnyPtrFromAny(rhs); if (lhs_data->type_index != rhs_data->type_index) { // type_index mismatch, if index is not string, return false if (lhs_data->type_index != kTVMFFIStr && lhs_data->type_index != kTVMFFISmallStr && lhs_data->type_index != kTVMFFISmallBytes && lhs_data->type_index != kTVMFFIBytes) { return false; } // small string and normal string comparison if (lhs_data->type_index == kTVMFFIStr && rhs_data->type_index == kTVMFFISmallStr) { const details::BytesObjBase* lhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); return Bytes::memequal(lhs_str->data, rhs_data->v_bytes, lhs_str->size, rhs_data->small_str_len); } if (lhs_data->type_index == kTVMFFISmallStr && rhs_data->type_index == kTVMFFIStr) { const details::BytesObjBase* rhs_str = details::AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs_data->v_bytes, rhs_str->data, lhs_data->small_str_len, rhs_str->size); } if (lhs_data->type_index == kTVMFFIBytes && rhs_data->type_index == kTVMFFISmallBytes) { const details::BytesObjBase* lhs_bytes = details::AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); return Bytes::memequal(lhs_bytes->data, rhs_data->v_bytes, lhs_bytes->size, rhs_data->small_str_len); } if (lhs_data->type_index == kTVMFFISmallBytes && rhs_data->type_index == kTVMFFIBytes) { const details::BytesObjBase* rhs_bytes = details::AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs_data->v_bytes, rhs_bytes->data, lhs_data->small_str_len, rhs_bytes->size); } return false; } if (lhs_data->type_index < TypeIndex::kTVMFFIStaticObjectBegin) { // specially handle nan for float, as there can be multiple representations of nan if (lhs_data->type_index == TypeIndex::kTVMFFIFloat && std::isnan(lhs_data->v_float64)) { return std::isnan(rhs_data->v_float64); } // this is POD data, we can just compare the value return lhs_data->zero_padding == rhs_data->zero_padding && lhs_data->v_int64 == rhs_data->v_int64; } switch (lhs_data->type_index) { case TypeIndex::kTVMFFIStr: case TypeIndex::kTVMFFIBytes: { // compare bytes const details::BytesObjBase* lhs_str = AnyUnsafe::CopyFromAnyViewAfterCheck(lhs); const details::BytesObjBase* rhs_str = AnyUnsafe::CopyFromAnyViewAfterCheck(rhs); return Bytes::memequal(lhs_str->data, rhs_str->data, lhs_str->size, rhs_str->size); } case TypeIndex::kTVMFFIArray: { return CompareArray(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck>(std::move(rhs))); } case TypeIndex::kTVMFFIList: { return CompareList(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck>(std::move(rhs))); } case TypeIndex::kTVMFFIMap: { return CompareMap(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck>(std::move(rhs))); } case TypeIndex::kTVMFFIDict: { return CompareMap(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck>(std::move(rhs))); } case TypeIndex::kTVMFFIShape: { return CompareShape(AnyUnsafe::MoveFromAnyAfterCheck(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck(std::move(rhs))); } case TypeIndex::kTVMFFITensor: { return CompareTensor(AnyUnsafe::MoveFromAnyAfterCheck(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck(std::move(rhs))); } default: { return CompareObject(AnyUnsafe::MoveFromAnyAfterCheck(std::move(lhs)), AnyUnsafe::MoveFromAnyAfterCheck(std::move(rhs))); } } } bool CompareObject(const ObjectRef& lhs, const ObjectRef& rhs) { // NOTE: invariant: lhs and rhs are already the same type const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(lhs->type_index()); if (type_info->metadata == nullptr) { TVM_FFI_THROW(TypeError) << "Type metadata is not set for type `" << String(type_info->type_key) << "`, so StructuralHash is not supported for this type"; } if (type_info->metadata->structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { TVM_FFI_THROW(TypeError) << "_type_s_eq_hash_kind is not set for type `" << String(type_info->type_key) << "`, so StructuralHash is not supported for this type"; } auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind; if (structural_eq_hash_kind == kTVMFFISEqHashKindUniqueInstance) { // use pointer comparison return lhs.same_as(rhs); } if (structural_eq_hash_kind == kTVMFFISEqHashKindConstTreeNode) { // fast path: constant tree node, pointer equality indicate equality and avoid content // comparison if false, we should still run content comparison if (lhs.same_as(rhs)) return true; } // check recorded mapping for DAG and fre var if (structural_eq_hash_kind == kTVMFFISEqHashKindDAGNode || structural_eq_hash_kind == kTVMFFISEqHashKindFreeVar) { // if there is pre-recorded mapping, need to cross check the pointer equality after mapping auto it = equal_map_lhs_.find(lhs); if (it != equal_map_lhs_.end()) { return it->second.same_as(rhs); } // if rhs is mapped but lhs is not, it means lhs is a free var, return false if (equal_map_rhs_.count(rhs)) { return false; } } if (structural_eq_hash_kind != kTVMFFISEqHashKindFreeVar) { bool success = CompareFields(lhs, rhs, type_info); if (success && structural_eq_hash_kind == kTVMFFISEqHashKindDAGNode) { // record the equality mapping for DAG nodes equal_map_lhs_[lhs] = rhs; equal_map_rhs_[rhs] = lhs; } return success; } // FreeVar path. In a non-recursive def region the FreeVar's own // sub-fields are walked outside the def region (nested free vars // there must resolve against an outer binding, not rebind), so we // clamp ``def_region_kind_`` to ``kNone`` around the CompareFields // call and restore before the binding decision below. TVMFFIDefRegionKind saved_def_region_kind = def_region_kind_; if (def_region_kind_ == kTVMFFIDefRegionKindNonRecursive) { def_region_kind_ = kTVMFFIDefRegionKindNone; } bool success = CompareFields(lhs, rhs, type_info); def_region_kind_ = saved_def_region_kind; if (!success) return false; // FreeVar that is not yet mapped: bind it iff identity-equal or we // are inside a def region. if (lhs.same_as(rhs) || def_region_kind_ != kTVMFFIDefRegionKindNone) { equal_map_lhs_[lhs] = rhs; equal_map_rhs_[rhs] = lhs; return true; } return false; } // Compare an object's fields (generic walk or via the custom // __ffi_s_equal__ callback). Does not touch the FreeVar def-region // clamp — the caller (CompareObject) handles that for the FreeVar case. bool CompareFields(const ObjectRef& lhs, const ObjectRef& rhs, const TVMFFITypeInfo* type_info) { static reflection::TypeAttrColumn custom_s_equal = reflection::TypeAttrColumn(reflection::type_attr::kSEqual); bool success = true; if (custom_s_equal[type_info->type_index] == nullptr) { // We recursively compare the fields the object reflection::ForEachFieldInfoWithEarlyStop(type_info, [&](const TVMFFIFieldInfo* field_info) { // skip fields that are marked as structural eq hash ignore if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) return false; // get the field value from both side reflection::FieldGetter getter(field_info); Any lhs_value = getter(lhs); Any rhs_value = getter(rhs); // Dispatch on the def-region flags. constexpr int64_t kSEqHashDefAny = kTVMFFIFieldFlagBitMaskSEqHashDefRecursive | kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive; if (field_info->flags & kSEqHashDefAny) { TVMFFIDefRegionKind new_kind = (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive) ? kTVMFFIDefRegionKindNonRecursive : kTVMFFIDefRegionKindRecursive; std::swap(new_kind, def_region_kind_); success = CompareAny(lhs_value, rhs_value); std::swap(new_kind, def_region_kind_); } else { success = CompareAny(lhs_value, rhs_value); } if (!success) { // record the first mismatching field if we sub-rountine compare failed if (mismatch_lhs_reverse_path_ != nullptr) { mismatch_lhs_reverse_path_->emplace_back( reflection::AccessStep::Attr(String(field_info->name))); mismatch_rhs_reverse_path_->emplace_back( reflection::AccessStep::Attr(String(field_info->name))); } // return true to indicate early stop return true; } else { // return false to continue checking other fields return false; } }); } else { // run custom equal function defined via __s_equal__ type attribute if (s_equal_callback_ == nullptr) { s_equal_callback_ = ffi::Function::FromTyped( // Third parameter is a ``TVMFFIDefRegionKind`` (passed on the wire // as ``int`` to keep the FFI signature stable across language // boundaries). [this](AnyView inner_lhs, AnyView inner_rhs, int def_region_kind, AnyView field_name) { // NOTE: we explicitly make field_name as AnyView to avoid copy overhead initially // and only cast to string if mismatch happens TVMFFIDefRegionKind new_kind = (def_region_kind == kTVMFFIDefRegionKindNone) ? def_region_kind_ : static_cast(def_region_kind); std::swap(new_kind, def_region_kind_); bool sub_success = CompareAny(inner_lhs, inner_rhs); std::swap(new_kind, def_region_kind_); if (!sub_success) { if (mismatch_lhs_reverse_path_ != nullptr) { String field_name_str = field_name.cast(); mismatch_lhs_reverse_path_->emplace_back( reflection::AccessStep::Attr(field_name_str)); mismatch_rhs_reverse_path_->emplace_back( reflection::AccessStep::Attr(field_name_str)); } } return sub_success; }); } success = custom_s_equal[type_info->type_index] .cast()(lhs, rhs, s_equal_callback_) .cast(); } return success; } template bool CompareMap(const MapType& lhs, const MapType& rhs) { if (lhs.size() != rhs.size()) { // size mismatch, and there is no path tracing // return false since we don't need informative error message if (mismatch_lhs_reverse_path_ == nullptr) return false; } // compare key and value pair by pair for (auto kv : lhs) { Any rhs_key = this->MapLhsToRhs(kv.first); auto it = rhs.find(rhs_key); if (it == rhs.end()) { if (mismatch_lhs_reverse_path_ != nullptr) { mismatch_lhs_reverse_path_->emplace_back(reflection::AccessStep::MapItem(kv.first)); mismatch_rhs_reverse_path_->emplace_back(reflection::AccessStep::MapItemMissing(rhs_key)); } return false; } // now recursively compare value if (!CompareAny(kv.second, (*it).second)) { if (mismatch_lhs_reverse_path_ != nullptr) { mismatch_lhs_reverse_path_->emplace_back(reflection::AccessStep::MapItem(kv.first)); mismatch_rhs_reverse_path_->emplace_back(reflection::AccessStep::MapItem(rhs_key)); } return false; } } // fast path, all contents equals to each other if (lhs.size() == rhs.size()) return true; // slow path, cross check every key from rhs in lhs to find the missing // key for better error reporting for (auto kv : rhs) { Any lhs_key = this->MapRhsToLhs(kv.first); auto it = lhs.find(lhs_key); if (it == lhs.end()) { if (mismatch_lhs_reverse_path_ != nullptr) { mismatch_lhs_reverse_path_->emplace_back(reflection::AccessStep::MapItemMissing(lhs_key)); mismatch_rhs_reverse_path_->emplace_back(reflection::AccessStep::MapItem(kv.first)); } return false; } } return false; } // NOLINTNEXTLINE(performance-unnecessary-value-param) bool CompareArray(ffi::Array lhs, ffi::Array rhs) { return CompareSequence(std::move(lhs), std::move(rhs)); } // NOLINTNEXTLINE(performance-unnecessary-value-param) bool CompareList(ffi::List lhs, ffi::List rhs) { return CompareSequence(std::move(lhs), std::move(rhs)); } template // NOLINTNEXTLINE(performance-unnecessary-value-param) bool CompareSequence(SeqType lhs, SeqType rhs) { if (lhs.size() != rhs.size()) { // fast path, size mismatch, and there is no path tracing // return false since we don't need informative error message if (mismatch_lhs_reverse_path_ == nullptr) return false; } for (int64_t i = 0, n = static_cast(std::min(lhs.size(), rhs.size())); i < n; ++i) { if (!CompareAny(lhs[i], rhs[i])) { if (mismatch_lhs_reverse_path_ != nullptr) { mismatch_lhs_reverse_path_->emplace_back(reflection::AccessStep::ArrayItem(i)); mismatch_rhs_reverse_path_->emplace_back(reflection::AccessStep::ArrayItem(i)); } return false; } } if (lhs.size() == rhs.size()) return true; if (mismatch_lhs_reverse_path_ != nullptr) { if (lhs.size() > rhs.size()) { mismatch_lhs_reverse_path_->emplace_back( reflection::AccessStep::ArrayItem(static_cast(rhs.size()))); mismatch_rhs_reverse_path_->emplace_back( reflection::AccessStep::ArrayItemMissing(static_cast(rhs.size()))); } else { mismatch_lhs_reverse_path_->emplace_back( reflection::AccessStep::ArrayItemMissing(static_cast(lhs.size()))); mismatch_rhs_reverse_path_->emplace_back( reflection::AccessStep::ArrayItem(static_cast(lhs.size()))); } } return false; } // NOLINTNEXTLINE(performance-unnecessary-value-param) bool CompareShape(Shape lhs, Shape rhs) { if (lhs.size() != rhs.size()) { return false; } for (size_t i = 0; i < lhs.size(); ++i) { if (lhs[i] != rhs[i]) { return false; } } return true; } // NOLINTNEXTLINE(performance-unnecessary-value-param) bool CompareTensor(Tensor lhs, Tensor rhs) { if (lhs.same_as(rhs)) return true; if (lhs.ndim() != rhs.ndim()) return false; for (int i = 0; i < lhs.ndim(); ++i) { if (lhs.size(i) != rhs.size(i)) return false; } if (lhs.dtype() != rhs.dtype()) return false; if (!skip_tensor_content_) { TVM_FFI_ICHECK_EQ(lhs.device().device_type, kDLCPU) << "can only compare CPU tensor"; TVM_FFI_ICHECK_EQ(rhs.device().device_type, kDLCPU) << "can only compare CPU tensor"; TVM_FFI_ICHECK(lhs.IsContiguous()) << "Can only compare contiguous tensor"; TVM_FFI_ICHECK(rhs.IsContiguous()) << "Can only compare contiguous tensor"; size_t data_size = GetDataSize(lhs); return std::memcmp(lhs.data_ptr(), rhs.data_ptr(), data_size) == 0; } else { return true; } } Any MapLhsToRhs(Any lhs) const { if (lhs.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) { return lhs; } ObjectRef lhs_obj = ffi::details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(lhs)); auto it = equal_map_lhs_.find(lhs_obj); if (it != equal_map_lhs_.end()) { return it->second; } return lhs_obj; } Any MapRhsToLhs(Any rhs) const { if (rhs.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) { return rhs; } ObjectRef rhs_obj = ffi::details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(rhs)); auto it = equal_map_rhs_.find(rhs_obj); if (it != equal_map_rhs_.end()) { return it->second; } return rhs_obj; } // Current def-region kind. ``kNone`` means we are not in a def region; // free vars discovered here do not bind (they must already be bound by an // outer scope or comparison falls back to pointer identity). ``kRecursive`` // and ``kNonRecursive`` enable binding for the field-flag-driven walk and // for the custom-callback path respectively (see CompareObject). TVMFFIDefRegionKind def_region_kind_{kTVMFFIDefRegionKindNone}; // whether we compare tensor data bool skip_tensor_content_{false}; // the root lhs for result printing std::vector* mismatch_lhs_reverse_path_ = nullptr; std::vector* mismatch_rhs_reverse_path_ = nullptr; // lazily initialize custom equal function ffi::Function s_equal_callback_ = nullptr; // map from lhs to rhs std::unordered_map equal_map_lhs_; // map from rhs to lhs std::unordered_map equal_map_rhs_; }; bool StructuralEqual::Equal(const Any& lhs, const Any& rhs, bool map_free_vars, bool skip_tensor_content) { StructEqualHandler handler; handler.def_region_kind_ = map_free_vars ? kTVMFFIDefRegionKindRecursive : kTVMFFIDefRegionKindNone; handler.skip_tensor_content_ = skip_tensor_content; return handler.CompareAny(lhs, rhs); } Optional StructuralEqual::GetFirstMismatch(const Any& lhs, const Any& rhs, bool map_free_vars, bool skip_tensor_content) { StructEqualHandler handler; handler.def_region_kind_ = map_free_vars ? kTVMFFIDefRegionKindRecursive : kTVMFFIDefRegionKindNone; handler.skip_tensor_content_ = skip_tensor_content; std::vector lhs_reverse_path; std::vector rhs_reverse_path; handler.mismatch_lhs_reverse_path_ = &lhs_reverse_path; handler.mismatch_rhs_reverse_path_ = &rhs_reverse_path; if (handler.CompareAny(lhs, rhs)) { return std::nullopt; } using reflection::AccessPath; reflection::AccessPath lhs_path = AccessPath::FromSteps(lhs_reverse_path.rbegin(), lhs_reverse_path.rend()); reflection::AccessPath rhs_path = AccessPath::FromSteps(rhs_reverse_path.rbegin(), rhs_reverse_path.rend()); return reflection::AccessPathPair(lhs_path, rhs_path); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("ffi.StructuralEqual", StructuralEqual::Equal) .def("ffi.GetFirstStructuralMismatch", StructuralEqual::GetFirstMismatch); // ensure the type attribute column is presented in the system even if it is empty. refl::EnsureTypeAttrColumn(refl::type_attr::kSEqual); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/structural_hash.cc000066400000000000000000000421041521067262500213520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/reflection/structural_equal.cc * * \brief Structural equal implementation. */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /** * \brief Internal Handler class for structural hash. */ class StructuralHashHandler { public: StructuralHashHandler() = default; uint64_t HashAny(ffi::Any src) { using ffi::details::AnyUnsafe; const TVMFFIAny* src_data = AnyUnsafe::TVMFFIAnyPtrFromAny(src); if (src_data->type_index < TypeIndex::kTVMFFIStaticObjectBegin) { // specially handle nan for float, as there can be multiple representations of nan // make sure they map to the same hash value if (src_data->type_index == TypeIndex::kTVMFFIFloat && std::isnan(src_data->v_float64)) { TVMFFIAny temp = *src_data; temp.v_float64 = std::numeric_limits::quiet_NaN(); return details::StableHashCombine(temp.type_index, temp.v_uint64); } if (src_data->type_index == TypeIndex::kTVMFFISmallStr) { // for small string, we use the same type key hash as normal string // so heap allocated string and on stack string will have the same hash return details::StableHashCombine(TypeIndex::kTVMFFIStr, details::StableHashSmallStrBytes(src_data)); } // this is POD data, we can just hash the value return details::StableHashCombine(src_data->type_index, src_data->v_uint64); } switch (src_data->type_index) { case TypeIndex::kTVMFFIStr: case TypeIndex::kTVMFFIBytes: { // return same hash as AnyHash const details::BytesObjBase* src_str = AnyUnsafe::CopyFromAnyViewAfterCheck(src); return details::StableHashCombine(src_data->type_index, details::StableHashBytes(src_str->data, src_str->size)); } case TypeIndex::kTVMFFIArray: { return HashArray(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(src))); } case TypeIndex::kTVMFFIList: { return HashList(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(src))); } case TypeIndex::kTVMFFIMap: { return HashMap(AnyUnsafe::MoveFromAnyAfterCheck>(std::move(src))); } case TypeIndex::kTVMFFIDict: { Dict dict = AnyUnsafe::MoveFromAnyAfterCheck>(std::move(src)); return HashMapBase(static_cast(dict.get())); } case TypeIndex::kTVMFFIShape: { return HashShape(AnyUnsafe::MoveFromAnyAfterCheck(std::move(src))); } case TypeIndex::kTVMFFITensor: { return HashTensor(AnyUnsafe::MoveFromAnyAfterCheck(std::move(src))); } default: { return HashObject(AnyUnsafe::MoveFromAnyAfterCheck(std::move(src))); } } } uint64_t HashObject(const ObjectRef& obj) { // NOTE: invariant: lhs and rhs are already the same type const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index()); if (type_info->metadata == nullptr) { TVM_FFI_THROW(TypeError) << "Type metadata is not set for type `" << String(type_info->type_key) << "`, so StructuralHash is not supported for this type"; } if (type_info->metadata->structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { TVM_FFI_THROW(TypeError) << "_type_s_eq_hash_kind is not set for type `" << String(type_info->type_key) << "`, so StructuralHash is not supported for this type"; } auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind; if (structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { // Fallback to pointer hash return std::hash()(obj.get()); } // return recored hash value if it is already computed auto it = hash_memo_.find(obj); if (it != hash_memo_.end()) { return it->second; } uint64_t hash_value; if (structural_eq_hash_kind != kTVMFFISEqHashKindFreeVar) { hash_value = HashFields(obj, type_info, obj->GetTypeKeyHash()); } else { // FreeVar path. In a non-recursive def region the FreeVar's own // sub-fields are walked outside the def region (nested free vars // there hash by pointer, matching use semantics), so we clamp // ``def_region_kind_`` to ``kNone`` around the HashFields call and // restore before the FreeVar-level injection below. // // We always call HashFields, even in use mode where the returned // ``hash_value`` is discarded by the pointer-hash fallback. The // walk's side effect on ``free_var_counter_`` (incremented for // every nested FreeVar reached via SEqHashDef-tagged sub-fields) // is observable to FreeVars hashed later in the same traversal; // skipping the walk would silently change those subsequent hashes. TVMFFIDefRegionKind saved_def_region_kind = def_region_kind_; if (def_region_kind_ == kTVMFFIDefRegionKindNonRecursive) { def_region_kind_ = kTVMFFIDefRegionKindNone; } hash_value = HashFields(obj, type_info, obj->GetTypeKeyHash()); def_region_kind_ = saved_def_region_kind; if (def_region_kind_ != kTVMFFIDefRegionKindNone) { // use lexical order of free var and its type hash_value = details::StableHashCombine(hash_value, free_var_counter_++); } else { // Fallback to pointer hash; we are not in a def region. hash_value = std::hash()(obj.get()); } } // if it is a DAG node, also record the lexical order of graph counter // this helps to distinguish DAG from trees. if (structural_eq_hash_kind == kTVMFFISEqHashKindDAGNode) { hash_value = details::StableHashCombine(hash_value, graph_node_counter_++); } // record the hash value for this object hash_memo_[obj] = hash_value; return hash_value; } // Hash an object's fields (generic walk or via the custom __ffi_s_hash__ // callback). Does not touch the FreeVar def-region clamp — that lives // inline in HashObject's FreeVar branch, which wraps this helper. uint64_t HashFields(const ObjectRef& obj, const TVMFFITypeInfo* type_info, uint64_t init_hash) { static reflection::TypeAttrColumn custom_s_hash = reflection::TypeAttrColumn(reflection::type_attr::kSHash); if (custom_s_hash[type_info->type_index] == nullptr) { // go over the content and hash the fields reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* field_info) { // skip fields that are marked as structural eq hash ignore if (!(field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore)) { reflection::FieldGetter getter(field_info); Any field_value = getter(obj); // Dispatch on the def-region flags (mirror of the equality side). constexpr int64_t kSEqHashDefAny = kTVMFFIFieldFlagBitMaskSEqHashDefRecursive | kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive; if (field_info->flags & kSEqHashDefAny) { TVMFFIDefRegionKind new_kind = (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive) ? kTVMFFIDefRegionKindNonRecursive : kTVMFFIDefRegionKindRecursive; std::swap(new_kind, def_region_kind_); init_hash = details::StableHashCombine(init_hash, HashAny(field_value)); std::swap(new_kind, def_region_kind_); } else { init_hash = details::StableHashCombine(init_hash, HashAny(field_value)); } } }); } else { if (s_hash_callback_ == nullptr) { s_hash_callback_ = // Third parameter is a ``TVMFFIDefRegionKind`` (passed on the wire // as ``int`` to keep the FFI signature stable across language // boundaries). ffi::Function::FromTyped([this](AnyView val, uint64_t inner_init, int def_region_kind) { TVMFFIDefRegionKind new_kind = (def_region_kind == kTVMFFIDefRegionKindNone) ? def_region_kind_ : static_cast(def_region_kind); std::swap(new_kind, def_region_kind_); uint64_t hv = HashAny(val); std::swap(new_kind, def_region_kind_); // we explicitly bitcast the result from `uint64_t` to `int64_t`. // The range of `uint64_t` is too large to fit as `int64_t`, so if we don't bitcast, // it will trigger an overflow error in `uint64_t` -> `Any` conversion. return static_cast(details::StableHashCombine(inner_init, hv)); }); } init_hash = custom_s_hash[type_info->type_index] .cast()(obj, static_cast(init_hash), s_hash_callback_) .cast(); } return init_hash; } // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashArray(Array arr) { return HashSequence(std::move(arr)); } // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashList(List list) { return HashSequence(std::move(list)); } template // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashSequence(SeqType seq) { uint64_t hash_value = details::StableHashCombine(seq->GetTypeKeyHash(), seq.size()); for (const auto& elem : seq) { hash_value = details::StableHashCombine(hash_value, HashAny(elem)); } return hash_value; } // Find an order independent hash value for a given Any. // Order independent hash value means the hash value will remain stable independent // of the order we hash the content at the current context. // This property is needed to support stable hash for map. std::optional FindOrderIndependentHash(const Any& src) { using ffi::details::AnyUnsafe; const TVMFFIAny* src_data = AnyUnsafe::TVMFFIAnyPtrFromAny(src); if (src_data->type_index < TypeIndex::kTVMFFIStaticObjectBegin) { if (src_data->type_index == TypeIndex::kTVMFFISmallStr) { // for small string, we use the same type key hash as normal string // so heap allocated string and on stack string will have the same hash return details::StableHashCombine( TypeIndex::kTVMFFIStr, details::StableHashBytes(src_data->v_bytes, src_data->small_str_len)); } // this is POD data, we can just hash the value return details::StableHashCombine(src_data->type_index, src_data->v_uint64); } else { if (src_data->type_index == TypeIndex::kTVMFFIStr || src_data->type_index == TypeIndex::kTVMFFIBytes) { const details::BytesObjBase* src_str = AnyUnsafe::CopyFromAnyViewAfterCheck(src); // return same hash as AnyHash return details::StableHashCombine(src_data->type_index, details::StableHashBytes(src_str->data, src_str->size)); } else { // if the hash of the object is already computed, return it auto it = hash_memo_.find(src.cast()); if (it != hash_memo_.end()) { return it->second; } return std::nullopt; } } } // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashMap(Map map) { return HashMapBase(static_cast(map.get())); } uint64_t HashMapBase(const MapBaseObj* map) { // Compute a deterministic hash value for the map. uint64_t hash_value = details::StableHashCombine(map->GetTypeKeyHash(), map->size()); std::vector> items; for (const auto& [key, value] : *map) { // if we cannot find order independent hash, we skip the key if (auto hash_key = FindOrderIndependentHash(key)) { items.emplace_back(*hash_key, value); } } // sort the items by the hash key, so the hash value is deterministic // and independent of the order of insertion std::sort(items.begin(), items.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); for (size_t i = 0; i < items.size();) { size_t k = i + 1; for (; k < items.size() && items[k].first == items[i].first; ++k) { } // detect ties, which are rare, but we need to skip value hash during ties // to make sure that the hash value is deterministic. if (k == i + 1) { // no ties, we just hash the key and value hash_value = details::StableHashCombine(hash_value, items[i].first); hash_value = details::StableHashCombine(hash_value, HashAny(items[i].second)); } else { // ties occur, we skip the value hash to make sure that the hash value is deterministic. hash_value = details::StableHashCombine(hash_value, items[i].first); } i = k; } return hash_value; } // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashShape(Shape shape) { uint64_t hash_value = details::StableHashCombine(shape->GetTypeKeyHash(), shape.size()); for (int64_t i : shape) { hash_value = details::StableHashCombine(hash_value, i); } return hash_value; } // NOLINTNEXTLINE(performance-unnecessary-value-param) uint64_t HashTensor(Tensor tensor) { uint64_t hash_value = details::StableHashCombine(tensor->GetTypeKeyHash(), tensor.ndim()); for (int i = 0; i < tensor.ndim(); ++i) { hash_value = details::StableHashCombine(hash_value, tensor.size(i)); } TVMFFIAny temp; temp.v_uint64 = 0; temp.v_dtype = tensor.dtype(); hash_value = details::StableHashCombine(hash_value, temp.v_int64); if (!skip_tensor_content_) { TVM_FFI_ICHECK_EQ(tensor.device().device_type, kDLCPU) << "can only hash CPU tensor"; TVM_FFI_ICHECK(tensor.IsContiguous()) << "Can only hash contiguous tensor"; size_t data_size = GetDataSize(tensor.numel(), tensor.dtype()); uint64_t data_hash = details::StableHashBytes(static_cast(tensor.data_ptr()), data_size); hash_value = details::StableHashCombine(hash_value, data_hash); } return hash_value; } // Current def-region kind. ``kNone`` means we are not in a def region; free // vars hash by pointer. ``kRecursive`` and ``kNonRecursive`` enable // ``free_var_counter_``-based hashing for the field-flag-driven walk and // for the custom-callback path respectively (see HashObject). TVMFFIDefRegionKind def_region_kind_{kTVMFFIDefRegionKindNone}; bool skip_tensor_content_{false}; // free var counter. uint32_t free_var_counter_{0}; // graph node counter. uint32_t graph_node_counter_{0}; // lazily initialize custom hash function ffi::Function s_hash_callback_ = nullptr; // map from lhs to rhs std::unordered_map hash_memo_; }; uint64_t StructuralHash::Hash(const Any& value, bool map_free_vars, bool skip_tensor_content) { StructuralHashHandler handler; handler.def_region_kind_ = map_free_vars ? kTVMFFIDefRegionKindRecursive : kTVMFFIDefRegionKindNone; handler.skip_tensor_content_ = skip_tensor_content; return handler.HashAny(value); } static int64_t FFIStructuralHash(const Any& value, bool map_free_vars, bool skip_tensor_content) { uint64_t result = StructuralHash::Hash(value, map_free_vars, skip_tensor_content); // we explicitly bitcast the result from `uint64_t` to `int64_t`. // The range of `uint64_t` is too large to fit as `int64_t`, so if we don't bitcast, // it will trigger an overflow error in `uint64_t` -> `Any` conversion. return static_cast(result); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("ffi.StructuralHash", FFIStructuralHash); refl::EnsureTypeAttrColumn(refl::type_attr::kSHash); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/extra/visit_error_context.cc000066400000000000000000000340461521067262500222600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/extra/visit_error_context.cc * \brief VisitErrorContext implementation — breadcrumb-trail collection and access-path * extraction for recursive Object visits. */ #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /** * \brief Internal handler for finding all access paths in a root ObjectRef * that match the breadcrumb pattern stored in a VisitErrorContext. */ class VisitErrorAccessPathFinder { public: TVM_FFI_COLD_CODE explicit VisitErrorAccessPathFinder(VisitErrorContext context, bool allow_prefix_match) : context_(std::move(context)), allow_prefix_match_(allow_prefix_match) { // Normalize the breadcrumb pattern before matching. Two kinds of noise can // accumulate in reverse_visit_pattern at recording time: // // 1. Consecutive duplicates of the same ObjectRef. This is expected when // TVM_FFI_VISIT_THROW(kind, node) and an enclosing // TVM_FFI_VISIT_END(node) at the same level both record // the same `node` — the throw site seeds the context with `node`, and // the catch handler immediately above appends `node` again. Treat any // run of identical adjacent entries as a single frame. // // 2. Null / undefined ObjectRefs. These can leak in when a visited node // was mutated or torn down between recording and matching, leaving a // stale slot. We can never match a null pointer against a live tree, // so drop them silently rather than letting them break the chain. // // The matcher operates on `records_` from here on; the original // `context_->reverse_visit_pattern` is left untouched so callers that // inspect the context for debugging still see the raw history. const List& raw = context_->reverse_visit_pattern; records_.reserve(raw.size()); for (const ObjectRef& entry : raw) { if (!entry.defined()) continue; if (!records_.empty() && entry.same_as(records_.back())) continue; records_.push_back(entry); } } TVM_FFI_COLD_CODE Array Find(const ObjectRef& root) { this->VisitObject(root); return Array(results_.begin(), results_.end()); } private: /** * \brief Stack-allocated mirror of AccessStepObj used during the descent hot path. * For kAttr: stores const TVMFFIFieldInfo* encoded as void* in key (via * kTVMFFIOpaquePtr). String allocation is deferred to ToAccessStep(). * Not an Object — pure struct, no Object header / refcount. */ struct TempAccessStep { reflection::AccessKind kind; Any key{}; // For kAttr: FieldInfo pointer encoded as void* (kTVMFFIOpaquePtr). // For kArrayItem: int64 index. // For kMapItem: the user's key. static TempAccessStep Attr(const TVMFFIFieldInfo* fi) { TempAccessStep s; s.kind = reflection::AccessKind::kAttr; // We store `fi` as an OPAQUE POINTER inside `Any` (kTVMFFIOpaquePtr). // `Any` is used here purely as a value-carrying slot for the bits of // the pointer — it does NOT take ownership, does NOT dereference // through it, and does NOT mutate the pointee. The FieldInfo struct // is read-only metadata owned by the type-info registry and lives // for the lifetime of the type info, so the underlying object is // unaffected by the const_cast below. // // The const_cast is needed only because `Any` provides TypeTraits for // `void*` but not for `const void*`. On retrieval (ToAccessStep() // below) we cast back to `const TVMFFIFieldInfo*` before any use. s.key = Any(const_cast(static_cast(fi))); return s; } static TempAccessStep ArrayItem(int64_t index) { TempAccessStep s; s.kind = reflection::AccessKind::kArrayItem; s.key = Any(index); return s; } static TempAccessStep MapItem(Any k) { TempAccessStep s; s.kind = reflection::AccessKind::kMapItem; s.key = std::move(k); return s; } /*! \brief Materialize the heap-allocated AccessStep. Called only at match time. * The String allocation for fi->name happens HERE, once. */ reflection::AccessStep ToAccessStep() const { if (kind == reflection::AccessKind::kAttr) { // Recover the original `const TVMFFIFieldInfo*` from the opaque // pointer stored in `key`. The bits round-trip unchanged; the // const-qualifier is restored here, immediately before any use. const TVMFFIFieldInfo* fi = static_cast(key.cast()); return reflection::AccessStep::Attr(String(fi->name.data, fi->name.size)); } else if (kind == reflection::AccessKind::kArrayItem) { return reflection::AccessStep::ArrayItem(key.cast()); } else { // kMapItem return reflection::AccessStep::MapItem(key); } } }; void VisitAny(Any value) { // Skip null Any silently — error-path defensive. if (value == nullptr) return; const int32_t type_index = value.type_index(); if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) { // Primitive — cannot hold an ObjectRef chain entry. return; } switch (type_index) { case TypeIndex::kTVMFFIArray: this->VisitSequence( details::AnyUnsafe::MoveFromAnyAfterCheck>(std::move(value))); break; case TypeIndex::kTVMFFIList: this->VisitSequence(details::AnyUnsafe::MoveFromAnyAfterCheck>(std::move(value))); break; case TypeIndex::kTVMFFIMap: this->VisitMap(details::AnyUnsafe::MoveFromAnyAfterCheck>(std::move(value))); break; case TypeIndex::kTVMFFIDict: this->VisitMap(details::AnyUnsafe::MoveFromAnyAfterCheck>(std::move(value))); break; default: if (type_index >= TypeIndex::kTVMFFIStaticObjectBegin) { ObjectRef obj = details::AnyUnsafe::MoveFromAnyAfterCheck(std::move(value)); this->VisitObject(obj); } break; } } void VisitObject(const ObjectRef& node) { // Defensive: error path; never throw. if (!node.defined()) return; const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(node->type_index()); if (type_info == nullptr || type_info->metadata == nullptr) return; bool matched_step = num_pattern_step_matched_ < records_.size() && node.same_as(records_[records_.size() - 1 - num_pattern_step_matched_]); if (matched_step) { ++num_pattern_step_matched_; if (num_pattern_step_matched_ == records_.size()) { // Full match — materialize the AccessPath and record. results_.push_back(this->MaterializeAccessPath()); --num_pattern_step_matched_; return; } } this->VisitChildrenFields(node, type_info); if (matched_step) --num_pattern_step_matched_; } void VisitChildrenFields(ObjectRef node, const TVMFFITypeInfo* type_info) { // Snapshot results_.size() before descent to detect any inner full or // prefix match recorded by a deeper call. If results_ grew, some inner // node was recorded — we do not record again here. size_t results_before = results_.size(); reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo* field_info) { reflection::FieldGetter getter(field_info); Any child_val = getter(node); this->PushStep(TempAccessStep::Attr(field_info)); this->VisitAny(std::move(child_val)); this->PopStep(); }); // Leaf-with-prefix-match: this node contributed a match step, but no // inner result was recorded from its subtree. Record the current node's // path as the best-effort prefix. path_steps_ reflects the path TO this // node (the step leading here was pushed by our caller), so // MaterializeAccessPath() yields the correct prefix path. // Skip when path_steps_ is empty — AccessPath::Root() gives no useful info. if (allow_prefix_match_ && num_pattern_step_matched_ > 0 && num_pattern_step_matched_ < records_.size() && !path_steps_.empty() && results_.size() == results_before) { results_.push_back(this->MaterializeAccessPath()); } } template void VisitSequence(const SeqType& seq) { for (size_t i = 0; i < seq.size(); ++i) { Any item = seq[static_cast(i)]; if (item == nullptr) continue; this->PushStep(TempAccessStep::ArrayItem(static_cast(i))); this->VisitAny(std::move(item)); this->PopStep(); } } template void VisitMap(const MapType& m) { for (const std::pair& kv : m) { if (kv.first == nullptr || kv.second == nullptr) continue; this->PushStep(TempAccessStep::MapItem(kv.first)); this->VisitAny(kv.second); this->PopStep(); } } /*! \brief Append a TempAccessStep to the descent stack; cache unchanged. */ void PushStep(TempAccessStep step) { path_steps_.push_back(std::move(step)); } /*! \brief Pop the top TempAccessStep; truncate materialized_paths_ to match. */ void PopStep() { path_steps_.pop_back(); if (materialized_paths_.size() > path_steps_.size()) { materialized_paths_.erase( materialized_paths_.begin() + static_cast(path_steps_.size()), materialized_paths_.end()); } } // Cache invariant maintained jointly with PushStep / PopStep: // materialized_paths_[k] (when present) is the AccessPath built // from path_steps_[0..k+1] as they currently exist, and // materialized_paths_.size() <= path_steps_.size(). // // - PushStep: appends to path_steps_; cache unchanged. // - PopStep: pops path_steps_; truncates cache to match. // - MaterializeAccessPath: extends cache up to path_steps_.size(), // chaining new AccessPath nodes via Extend. // // Lazy materialization avoids per-descent AccessPath allocation; // cache amortizes across consecutive matches with shared prefix // (LCA sharing). /*! \brief Materialize the AccessPath at the current descent depth. */ reflection::AccessPath MaterializeAccessPath() { // Root-itself match: matching fired before any descent step was pushed, // so the access path is just Root(). Returning early also avoids reading // back() from an empty materialized_paths_. if (path_steps_.empty()) return reflection::AccessPath::Root(); for (size_t idx = materialized_paths_.size(); idx < path_steps_.size(); ++idx) { reflection::AccessPath parent = (idx == 0) ? reflection::AccessPath::Root() : materialized_paths_[idx - 1]; materialized_paths_.push_back(parent->Extend(path_steps_[idx].ToAccessStep())); } return materialized_paths_.back(); } // The visit error context whose pattern we're matching against. VisitErrorContext context_; // When true, record prefix matches at leaves (partial pattern match). bool allow_prefix_match_; // Normalized breadcrumb pattern derived from context_->reverse_visit_pattern // by dropping null/undefined entries and collapsing consecutive duplicates. // See the constructor for rationale. std::vector records_; // Count of pattern entries matched so far (root-closest first). Incremented on // match, decremented on unwind; full match when equal to pattern size. size_t num_pattern_step_matched_{0}; // Current descent path — entries pushed on field/index/key descent, popped on unwind. std::vector path_steps_; // Lazy cache of materialized AccessPath nodes, parallel to path_steps_. // materialized_paths_[k] corresponds to path_steps_[0..k+1] as they currently // stand. Extended on match; truncated by PopStep. size <= path_steps_.size(). std::vector materialized_paths_; // Recorded full-match paths. std::vector results_; }; // --------------------------------------------------------------------------- // VisitErrorContext::FindAccessPaths // --------------------------------------------------------------------------- TVM_FFI_COLD_CODE Array VisitErrorContext::FindAccessPaths( const ObjectRef& root, const VisitErrorContext& visit_context, bool allow_prefix_match) { VisitErrorAccessPathFinder finder(visit_context, allow_prefix_match); return finder.Find(root); } // --------------------------------------------------------------------------- // Registration // --------------------------------------------------------------------------- TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("reverse_visit_pattern", &VisitErrorContextObj::reverse_visit_pattern) .def_ro("prev_error_context", &VisitErrorContextObj::prev_error_context); refl::GlobalDef().def("ffi.VisitErrorContext.FindAccessPaths", &VisitErrorContext::FindAccessPaths); } } // namespace ffi } // namespace tvm tvm-ffi-0.1.12/src/ffi/function.cc000066400000000000000000000222771521067262500166520ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/function.cc * \brief Function call registry and safecall context */ #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { /*! * \brief Global function table. * * \note We do not use mutex to guard updating of GlobalFunctionTable * * The assumption is that updating of GlobalFunctionTable will be done * in the main thread during initialization or loading, or * explicitly locked from the caller. * * Then the followup code will leverage the information */ class GlobalFunctionTable { public: // Note: this class is hidden from the public API, so we just // use it as a private class as ObjectRef class Entry : public Object, public TVMFFIMethodInfo { public: String name_data; String doc_data; String metadata_data; ffi::Function func_data; explicit Entry(const TVMFFIMethodInfo* method_info) { // make copy of the metadata name_data = String(method_info->name.data, method_info->name.size); doc_data = String(method_info->doc.data, method_info->doc.size); metadata_data = String(method_info->metadata.data, method_info->metadata.size); func_data = AnyView::CopyFromTVMFFIAny(method_info->method).cast(); this->SyncMethodInfo(method_info->flags); } explicit Entry(String name, ffi::Function func) : name_data(std::move(name)), func_data(std::move(func)) { this->SyncMethodInfo(kTVMFFIFieldFlagBitMaskIsStaticMethod); } private: void SyncMethodInfo(int64_t flags) { this->method = AnyView(func_data).CopyToTVMFFIAny(); this->flags = flags; this->name = TVMFFIByteArray{name_data.data(), name_data.size()}; this->doc = TVMFFIByteArray{doc_data.data(), doc_data.size()}; this->metadata = TVMFFIByteArray{metadata_data.data(), metadata_data.size()}; } }; void Update(const String& name, Function func, bool can_override) { if (TVM_FFI_PREDICT_FALSE(table_.count(name) != 0)) { if (!can_override) { TVM_FFI_THROW(RuntimeError) << "Global Function `" << name << "` is already registered"; } } table_.Set(name, ObjectRef(make_object(name, std::move(func)))); } void Update(const TVMFFIMethodInfo* method_info, bool can_override) { String name(method_info->name.data, method_info->name.size); if (TVM_FFI_PREDICT_FALSE(table_.count(name) != 0)) { if (!can_override) { TVM_FFI_LOG_AND_THROW(RuntimeError) << "Global Function `" << name << "` is already registered, possible causes:\n" << "- Two GlobalDef().def registrations for the same function \n" << "Please remove the duplicate registration."; } } table_.Set(name, ObjectRef(make_object(method_info))); } bool Remove(const String& name) { auto it = table_.find(name); if (it == table_.end()) return false; table_.erase(name); return true; } const Entry* Get(const String& name) { auto it = table_.find(name); if (it == table_.end()) return nullptr; const Object* obj = (*it).second.cast(); return static_cast(obj); } Array ListNames() const { Array names; names.reserve(static_cast(table_.size())); for (const auto& kv : table_) { names.push_back(kv.first); } return names; } static GlobalFunctionTable* Global() { // We deliberately create a new instance via raw new // This is because GlobalFunctionTable can contain callbacks into // the host language (Python) and the resource can become invalid // indeterministic order of destruction and forking. // The resources will only be recycled during program exit. static GlobalFunctionTable* inst = new GlobalFunctionTable(); return inst; } private: Map table_; }; } // namespace ffi } // namespace tvm int TVMFFIFunctionCreate(void* self, TVMFFISafeCallType safe_call, void (*deleter)(void* self), TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::Function func = tvm::ffi::Function::FromExternC(self, safe_call, deleter); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(func)); TVM_FFI_SAFE_CALL_END(); } int TVMFFIAnyViewToOwnedAny(const TVMFFIAny* any_view, TVMFFIAny* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::Any result(*reinterpret_cast(any_view)); *out = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(result)); TVM_FFI_SAFE_CALL_END(); } int TVMFFIFunctionSetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle f, int override) { using namespace tvm::ffi; TVM_FFI_SAFE_CALL_BEGIN(); String name_str(name->data, name->size); GlobalFunctionTable::Global()->Update(name_str, GetRef(static_cast(f)), override != 0); TVM_FFI_SAFE_CALL_END(); } int TVMFFIFunctionSetGlobalFromMethodInfo(const TVMFFIMethodInfo* method_info, int override) { using namespace tvm::ffi; TVM_FFI_SAFE_CALL_BEGIN(); GlobalFunctionTable::Global()->Update(method_info, override != 0); TVM_FFI_SAFE_CALL_END(); } int TVMFFIFunctionGetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle* out) { using namespace tvm::ffi; TVM_FFI_SAFE_CALL_BEGIN(); String name_str(name->data, name->size); const GlobalFunctionTable::Entry* fp = GlobalFunctionTable::Global()->Get(name_str); if (fp != nullptr) { tvm::ffi::Function func(fp->func_data); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(func)); } else { *out = nullptr; } TVM_FFI_SAFE_CALL_END(); } int TVMFFIFunctionCall(TVMFFIObjectHandle func, TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { using namespace tvm::ffi; #ifdef _MSC_VER // Avoid tail call optimization // in MSVC many cases python symbols are hidden, so we need this function symbol // to be in the call frame to reliably detect the ffi boundary volatile int ret = reinterpret_cast(func)->safe_call(func, args, num_args, result); return ret; #else // NOTE: this is a tail call return reinterpret_cast(func)->safe_call(func, args, num_args, result); #endif } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("ffi.FunctionRemoveGlobal", [](const tvm::ffi::String& name) -> bool { return tvm::ffi::GlobalFunctionTable::Global()->Remove(name); }) .def("ffi.FunctionListGlobalNamesFunctor", []() { // NOTE: we return functor instead of array // so list global function names do not need to depend on array // this is because list global function names usually is a core api that happens // before array ffi functions are available. tvm::ffi::Array names = tvm::ffi::GlobalFunctionTable::Global()->ListNames(); auto return_functor = [names](int64_t i) -> tvm::ffi::Any { if (i < 0) { return names.size(); } else { return names[i]; } }; return tvm::ffi::Function::FromTyped(return_functor); }) .def("ffi.String", [](tvm::ffi::String val) -> tvm::ffi::String { return val; }) .def("ffi.Bytes", [](tvm::ffi::Bytes val) -> tvm::ffi::Bytes { return val; }) .def("ffi.GetGlobalFuncMetadata", [](const tvm::ffi::String& name) -> tvm::ffi::String { const auto* f = tvm::ffi::GlobalFunctionTable::Global()->Get(name); if (f == nullptr) { TVM_FFI_THROW(RuntimeError) << "Global Function is not found: " << name; } return f->metadata_data; }) .def("ffi.FunctionFromExternC", [](void* self, void* safe_call, void* deleter) -> tvm::ffi::Function { return tvm::ffi::Function::FromExternC(self, reinterpret_cast(safe_call), reinterpret_cast(deleter)); }); } tvm-ffi-0.1.12/src/ffi/init_once.cc000066400000000000000000000063271521067262500167720ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/init_once.cc * \brief Handle Init Once C API implementation. */ #include #include #include #ifdef _MSC_VER #include #endif namespace { TVM_FFI_INLINE void* AtomicLoadHandleAcquire(void** src_addr) { #ifdef _MSC_VER #if defined(_WIN32_WINNT) && (_WIN32_WINNT >= 0x0602) return InterlockedCompareExchangePointerAcquire(reinterpret_cast(src_addr), // nullptr, nullptr); #else return InterlockedCompareExchangePointer(reinterpret_cast(src_addr), // nullptr, nullptr); #endif #else return __atomic_load_n(src_addr, __ATOMIC_ACQUIRE); #endif } TVM_FFI_INLINE void AtomicStoreHandleRelease(void** dst_addr, void* src) { #ifdef _MSC_VER _InterlockedExchangePointer(reinterpret_cast(dst_addr), src); #else __atomic_store_n(dst_addr, src, __ATOMIC_RELEASE); #endif } } // namespace int TVMFFIHandleInitOnce(void** handle_addr, int (*init_func)(void** result)) { // fast path: handle is already initialized // we still need atomic load acquire to ensure the handle is not initialized if (AtomicLoadHandleAcquire(handle_addr) != nullptr) return 0; // slow path: handle is not initialized, call initialization function // note: slow path is not meant to be called frequently, so we use a simple mutex static std::mutex mutex; std::scoped_lock lock(mutex); // must check again here, because another thread may have initialized the // handle before we acquired the lock if (*handle_addr != nullptr) return 0; void* result = nullptr; int ret = init_func(&result); if (ret != 0) return ret; if (result == nullptr) { TVMFFIErrorSetRaisedFromCStr("RuntimeError", "init_func must return a non-NULL handle"); return -1; } // NOTE: we must use atomic store release to ensure the result is // visible to other thread's atomic load acquire AtomicStoreHandleRelease(handle_addr, result); return 0; } int TVMFFIHandleDeinitOnce(void** handle_addr, int (*deinit_func)(void* handle)) { #ifdef _MSC_VER void* old_handle = _InterlockedExchangePointer(reinterpret_cast(handle_addr), nullptr); #else void* old_handle = __atomic_exchange_n(handle_addr, nullptr, __ATOMIC_ACQ_REL); #endif if (old_handle != nullptr) { return (*deinit_func)(old_handle); } return 0; } tvm-ffi-0.1.12/src/ffi/object.cc000066400000000000000000000701541521067262500162700ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/object.cc * \brief Registry to record dynamic types */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "object_internal.h" namespace tvm { namespace ffi { /*! * \brief Global registry that manages * * \note We do not use mutex to guard updating of TypeTable * * The assumption is that updating of TypeTable will be done * in the main thread during initialization or loading, or * explicitly locked from the caller. * * Then the followup code will leverage the information */ class TypeTable { public: /*! \brief Type information */ struct Entry : public TypeInfo { /*! \brief stored type key */ String type_key_data; /*! \brief acenstor information */ std::vector type_ancestors_data; /*! \brief type fields informaton */ std::vector type_fields_data; /*! \brief type methods informaton */ std::vector type_methods_data; /*! \brief extra information */ TVMFFITypeMetadata metadata_data; // NOTE: the indices in [index, index + num_reserved_slots) are // reserved for the child-class of this type. /*! \brief Total number of slots reserved for the type and its children. */ int32_t num_slots; /*! \brief number of allocated child slots. */ int32_t allocated_slots; /*! \brief Whether child can overflow. */ bool child_slots_can_overflow{true}; Entry(int32_t type_index, int32_t type_depth, String type_key, int32_t num_slots, bool child_slots_can_overflow, const Entry* parent) { // setup fields in the class this->type_key_data = std::move(type_key); this->num_slots = num_slots; this->allocated_slots = 1; this->child_slots_can_overflow = child_slots_can_overflow; // set up type acenstors information if (type_depth != 0) { TVM_FFI_ICHECK_NOTNULL(parent); TVM_FFI_ICHECK_EQ(type_depth, parent->type_depth + 1); type_ancestors_data.resize(type_depth); // copy over parent's type information for (int32_t i = 0; i < parent->type_depth; ++i) { type_ancestors_data[i] = parent->type_ancestors[i]; } // set last type information to be parent type_ancestors_data[parent->type_depth] = parent; } // initialize type info: no change to type_key and type_ancestors fields // after this line this->type_index = type_index; this->type_depth = type_depth; this->type_key = TVMFFIByteArray{this->type_key_data.data(), this->type_key_data.length()}; this->type_key_hash = std::hash()(this->type_key_data); this->type_ancestors = type_ancestors_data.data(); // initialize the reflection information this->num_fields = 0; this->num_methods = 0; this->fields = nullptr; this->methods = nullptr; this->metadata = nullptr; } }; struct TypeAttrColumnData : public TVMFFITypeAttrColumn { std::vector data_; }; int32_t GetOrAllocTypeIndex(String type_key, int32_t static_type_index, int32_t type_depth, int32_t num_child_slots, bool child_slots_can_overflow, int32_t parent_type_index) { auto it = type_key2index_.find(type_key); if (it != type_key2index_.end()) { return type_table_[(*it).second]->type_index; } // get parent's entry Entry* parent = [&]() -> Entry* { if (parent_type_index < 0) return nullptr; // try to allocate from parent's type table. TVM_FFI_ICHECK_LT(parent_type_index, type_table_.size()) << " type_key=" << type_key << ", static_index=" << static_type_index; return type_table_[parent_type_index].get(); }(); // get allocated index int32_t allocated_tindex = [&]() { // Step 0: static allocation if (static_type_index >= 0) { TVM_FFI_ICHECK_LT(static_type_index, type_table_.size()); TVM_FFI_ICHECK(type_table_[static_type_index] == nullptr) << "Conflicting static index " << static_type_index << " between " << ToStringView(type_table_[static_type_index]->type_key) << " and " << type_key; return static_type_index; } TVM_FFI_ICHECK_NOTNULL(parent); int num_slots = num_child_slots + 1; if (parent->allocated_slots + num_slots <= parent->num_slots) { // allocate the slot from parent's reserved pool int32_t allocated_tindex = parent->type_index + parent->allocated_slots; // update parent's state parent->allocated_slots += num_slots; return allocated_tindex; } // Step 2: allocate from overflow TVM_FFI_ICHECK(parent->child_slots_can_overflow) << "Reach maximum number of sub-classes for " << ToStringView(parent->type_key); // allocate new entries. int32_t allocated_tindex = static_cast(type_counter_); type_counter_ += num_slots; TVM_FFI_ICHECK_LE(type_table_.size(), type_counter_); type_table_.reserve(type_counter_); // resize type table while (static_cast(type_table_.size()) < type_counter_) { type_table_.emplace_back(nullptr); } return allocated_tindex; }(); // if parent cannot overflow, then this class cannot. if (parent != nullptr && !(parent->child_slots_can_overflow)) { child_slots_can_overflow = false; } // total number of slots include the type itself. if (parent != nullptr) { TVM_FFI_ICHECK_GT(allocated_tindex, parent->type_index); } type_table_[allocated_tindex] = std::make_unique(allocated_tindex, type_depth, type_key, num_child_slots + 1, child_slots_can_overflow, parent); // update the key2index mapping. type_key2index_.Set(type_key, allocated_tindex); return allocated_tindex; } int32_t TypeKeyToIndex(const TVMFFIByteArray* type_key) { String type_key_str(type_key->data, type_key->size); auto it = type_key2index_.find(type_key_str); TVM_FFI_ICHECK(it != type_key2index_.end()) << "Cannot find type `" << type_key_str << "`"; return static_cast((*it).second); } Entry* GetTypeEntry(int32_t type_index) { Entry* entry = nullptr; if (type_index >= 0 && static_cast(type_index) < type_table_.size()) { entry = type_table_[type_index].get(); } TVM_FFI_ICHECK(entry != nullptr) << "Cannot find type info for type_index=" << type_index; return entry; } Array GetRegisteredTypeKeys() const { Array ret; for (const auto& entry : type_table_) { if (entry) { ret.push_back(entry->type_key_data); } } return ret; } void RegisterTypeField(int32_t type_index, const TVMFFIFieldInfo* info) { Entry* entry = GetTypeEntry(type_index); std::string_view new_name(info->name.data, info->name.size); std::string_view type_key(entry->type_key.data, entry->type_key.size); // Check: no duplicate field name within this type's own fields. for (const auto& existing : entry->type_fields_data) { TVM_FFI_ICHECK(std::string_view(existing.name.data, existing.name.size) != new_name) << "Duplicate field name \"" << new_name << "\" in type \"" << type_key << "\""; } // Warn: field name should not shadow any ancestor field. for (int32_t d = 0; d < entry->type_depth; ++d) { const TVMFFITypeInfo* ancestor = entry->type_ancestors[d]; for (int32_t i = 0; i < ancestor->num_fields; ++i) { if (std::string_view(ancestor->fields[i].name.data, ancestor->fields[i].name.size) == new_name) { std::cerr << "[WARNING] Field \"" << new_name << "\" in type \"" << type_key << "\" duplicates an ancestor field in \"" << std::string_view(ancestor->type_key.data, ancestor->type_key.size) << "\". Child types should not re-register inherited fields." << std::endl; } } } TVMFFIFieldInfo field_data = *info; // Retain FunctionObj setter via any_pool_ so it outlives the Entry. if ((field_data.flags & kTVMFFIFieldFlagBitSetterIsFunctionObj) && field_data.setter != nullptr) { TVMFFIAny setter_ref; setter_ref.type_index = kTVMFFIFunction; setter_ref.zero_padding = 0; setter_ref.v_obj = static_cast(field_data.setter); any_pool_.emplace_back(AnyView::CopyFromTVMFFIAny(setter_ref)); } field_data.name = this->CopyString(info->name); field_data.doc = this->CopyString(info->doc); field_data.metadata = this->CopyString(info->metadata); if (info->flags & kTVMFFIFieldFlagBitMaskHasDefault) { field_data.default_value_or_factory = this->CopyAny(AnyView::CopyFromTVMFFIAny(info->default_value_or_factory)) .CopyToTVMFFIAny(); } else { field_data.default_value_or_factory = AnyView(nullptr).CopyToTVMFFIAny(); } entry->type_fields_data.push_back(field_data); // refresh ptr as the data can change entry->fields = entry->type_fields_data.data(); entry->num_fields = static_cast(entry->type_fields_data.size()); } void RegisterTypeMethod(int32_t type_index, const TVMFFIMethodInfo* info) { Entry* entry = GetTypeEntry(type_index); TVMFFIMethodInfo method_data = *info; method_data.name = this->CopyString(info->name); method_data.doc = this->CopyString(info->doc); method_data.metadata = this->CopyString(info->metadata); method_data.method = this->CopyAny(AnyView::CopyFromTVMFFIAny(info->method)).CopyToTVMFFIAny(); entry->type_methods_data.push_back(method_data); entry->methods = entry->type_methods_data.data(); entry->num_methods = static_cast(entry->type_methods_data.size()); } void RegisterTypeMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata) { Entry* entry = GetTypeEntry(type_index); if (entry->metadata != nullptr) { TVM_FFI_LOG_AND_THROW(RuntimeError) << "Overriding " << ToStringView(entry->type_key) << ", possible causes:\n" << "- two ObjectDef() calls for the same T \n" << "- when we forget to assign _type_key to ObjectRef that inherits from T\n" << "- another type with the same key is already registered\n" << "Cross check the reflection registration."; } entry->metadata_data = *metadata; entry->metadata_data.doc = this->CopyString(metadata->doc); entry->metadata = &(entry->metadata_data); } void RegisterTypeAttr(int32_t type_index, const TVMFFIByteArray* name, const TVMFFIAny* value) { AnyView value_view = AnyView::CopyFromTVMFFIAny(*value); String name_str(*name); size_t column_index = 0; auto it = type_attr_name_to_column_index_.find(name_str); if (it == type_attr_name_to_column_index_.end()) { column_index = type_attr_columns_.size(); type_attr_columns_.emplace_back(std::make_unique()); type_attr_name_to_column_index_.Set(name_str, static_cast(column_index)); } else { column_index = (*it).second; } TypeAttrColumnData* column = type_attr_columns_[column_index].get(); if (type_index == kTVMFFINone) { // Sentinel: just ensure the column exists without registering a value. if (column->data_.empty()) { column->data = reinterpret_cast(column->data_.data()); column->size = 0; column->begin_index = 0; } return; } // TODO(1.0): set begin_index to first registered type_index for sparse column storage // For now, begin_index is always 0 (resize from index 0). if (static_cast(type_index) >= column->data_.size()) { // Extend back from index 0. column->data_.resize(static_cast(type_index) + 1, Any(nullptr)); } column->data = reinterpret_cast(column->data_.data()); column->size = static_cast(column->data_.size()); column->begin_index = 0; Any& slot = column->data_[type_index - column->begin_index]; if (slot.type_index() != kTVMFFINone) { TVM_FFI_THROW(RuntimeError) << "TypeAttr `" << name_str << "` is already registered for type index " << type_index << ". To update the stored value, register a mutable container (e.g., Dict/List) " << "once and mutate it in place on subsequent calls."; } slot = value_view; } const TVMFFITypeAttrColumn* GetTypeAttrColumn(const TVMFFIByteArray* name) { String name_str(*name); auto it = type_attr_name_to_column_index_.find(name_str); if (it == type_attr_name_to_column_index_.end()) return nullptr; return type_attr_columns_[(*it).second].get(); } void Dump(int min_children_count) { std::vector num_children(type_table_.size(), 0); // expected child slots compute the expected slots // based on the current child slot setting std::vector expected_child_slots(type_table_.size(), 0); // reverse accumulation so we can get total counts in a bottom-up manner. for (auto it = type_table_.rbegin(); it != type_table_.rend(); ++it) { const Entry* ptr = it->get(); if (ptr != nullptr && ptr->type_depth != 0) { int parent_index = ptr->type_ancestors[ptr->type_depth - 1]->type_index; num_children[parent_index] += num_children[ptr->type_index] + 1; if (expected_child_slots[ptr->type_index] + 1 < ptr->num_slots) { expected_child_slots[ptr->type_index] = ptr->num_slots - 1; } expected_child_slots[parent_index] += expected_child_slots[ptr->type_index] + 1; } } for (const auto& ptr : type_table_) { if (ptr != nullptr && num_children[ptr->type_index] >= min_children_count) { std::cerr << '[' << ptr->type_index << "]\t" << ToStringView(ptr->type_key); if (ptr->type_depth != 0) { int32_t parent_index = ptr->type_ancestors[ptr->type_depth - 1]->type_index; std::cerr << "\tparent=" << ToStringView(type_table_[parent_index]->type_key); } else { std::cerr << "\tparent=root"; } std::cerr << "\tnum_child_slots=" << ptr->num_slots - 1 << "\tnum_children=" << num_children[ptr->type_index] << "\texpected_child_slots=" << expected_child_slots[ptr->type_index] << std::endl; } } } static TypeTable* Global() { // deliberately create a new instance via raw new // to ensure table lives longer in case unloading // still need the table info // memory will be recycled by the OS at program exit static TypeTable* inst = new TypeTable(); return inst; } private: TypeTable() : type_attr_columns_{}, type_attr_name_to_column_index_{} { type_table_.reserve(TypeIndex::kTVMFFIDynObjectBegin); for (int32_t i = 0; i < TypeIndex::kTVMFFIDynObjectBegin; ++i) { type_table_.emplace_back(nullptr); } // initialize the entry for object this->GetOrAllocTypeIndex(String(Object::_type_key), Object::_type_index, Object::_type_depth, Object::_type_child_slots, Object::_type_child_slots_can_overflow, -1); TVMFFITypeMetadata info; info.total_size = sizeof(Object); info.structural_eq_hash_kind = kTVMFFISEqHashKindUnsupported; info.creator = nullptr; info.doc = TVMFFIByteArray{nullptr, 0}; RegisterTypeMetadata(Object::_type_index, &info); // reserve the static types ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFINone, TypeIndex::kTVMFFINone); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIInt, TypeIndex::kTVMFFIInt); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIBool, TypeIndex::kTVMFFIBool); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIFloat, TypeIndex::kTVMFFIFloat); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIOpaquePtr, TypeIndex::kTVMFFIOpaquePtr); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDataType, TypeIndex::kTVMFFIDataType); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDevice, TypeIndex::kTVMFFIDevice); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDLTensorPtr, TypeIndex::kTVMFFIDLTensorPtr); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIRawStr, TypeIndex::kTVMFFIRawStr); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIByteArrayPtr, TypeIndex::kTVMFFIByteArrayPtr); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIObjectRValueRef, TypeIndex::kTVMFFIObjectRValueRef); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFISmallStr, TypeIndex::kTVMFFISmallStr); ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFISmallBytes, TypeIndex::kTVMFFISmallBytes); // reserved static type indices for depth 1 object types ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIStr, TypeIndex::kTVMFFIStr); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIBytes, TypeIndex::kTVMFFIBytes); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIError, TypeIndex::kTVMFFIError); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIFunction, TypeIndex::kTVMFFIFunction); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIShape, TypeIndex::kTVMFFIShape); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFITensor, TypeIndex::kTVMFFITensor); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIArray, TypeIndex::kTVMFFIArray); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIList, TypeIndex::kTVMFFIList); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIMap, TypeIndex::kTVMFFIMap); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIDict, TypeIndex::kTVMFFIDict); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIModule, TypeIndex::kTVMFFIModule); ReserveDepthOneObjectTypeIndex(StaticTypeKey::kTVMFFIOpaquePyObject, TypeIndex::kTVMFFIOpaquePyObject); } void ReserveBuiltinTypeIndex(const char* type_key, int32_t static_type_index) { this->GetOrAllocTypeIndex(String(type_key), static_type_index, 0, 0, false, -1); } void ReserveDepthOneObjectTypeIndex(const char* type_key, int32_t static_type_index) { this->GetOrAllocTypeIndex(String(type_key), static_type_index, 1, 0, false, TypeIndex::kTVMFFIObject); } static ObjectPtr MakeInplaceString(const char* data, size_t length) { ObjectPtr p = make_inplace_array_object(length + 1); static_assert(alignof(details::StringObj) % alignof(char) == 0); static_assert(sizeof(details::StringObj) % alignof(char) == 0); char* dest_data = reinterpret_cast(p.get()) + sizeof(details::StringObj); p->data = dest_data; p->size = length; std::memcpy(dest_data, data, length); dest_data[length] = '\0'; return p; } TVMFFIByteArray CopyString(TVMFFIByteArray str) { if (str.size == 0) { return TVMFFIByteArray{nullptr, 0}; } // use explicit object creation to ensure the space pointer to not move auto str_obj = MakeInplaceString(str.data, str.size); TVMFFIByteArray c_val{str_obj->data, str_obj->size}; any_pool_.emplace_back(ObjectRef(std::move(str_obj))); return c_val; } AnyView CopyAny(Any val) { AnyView view = AnyView(val); any_pool_.emplace_back(std::move(val)); return view; } int64_t type_counter_{TypeIndex::kTVMFFIDynObjectBegin}; std::vector> type_table_; Map type_key2index_; std::vector any_pool_; // type attribute columns std::vector> type_attr_columns_; Map type_attr_name_to_column_index_; }; /** * \brief Opaque implementation */ class OpaqueObjectImpl : public Object, public TVMFFIOpaqueObjectCell { public: OpaqueObjectImpl(void* handle, void (*deleter)(void* handle)) : deleter_(deleter) { this->handle = handle; } void SetTypeIndex(int32_t type_index) { details::ObjectUnsafe::GetHeader(this)->type_index = type_index; } ~OpaqueObjectImpl() { if (deleter_ != nullptr) { deleter_(handle); } } private: void (*deleter_)(void* handle); }; ObjectRef GetMissingObject() { static ObjectRef missing_obj(make_object()); return missing_obj; } ObjectRef GetKwargsObject() { static ObjectRef kwargs_obj(make_object()); return kwargs_obj; } } // namespace ffi } // namespace tvm void TVMFFIGetVersion(TVMFFIVersion* out_version) { out_version->major = TVM_FFI_VERSION_MAJOR; out_version->minor = TVM_FFI_VERSION_MINOR; out_version->patch = TVM_FFI_VERSION_PATCH; } int TVMFFIObjectDecRef(TVMFFIObjectHandle handle) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::details::ObjectUnsafe::DecRefObjectHandle(handle); TVM_FFI_SAFE_CALL_END(); } int TVMFFIObjectIncRef(TVMFFIObjectHandle handle) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::details::ObjectUnsafe::IncRefObjectHandle(handle); TVM_FFI_SAFE_CALL_END(); } int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index, void (*deleter)(void* handle), TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); if (type_index != kTVMFFIOpaquePyObject) { TVM_FFI_THROW(RuntimeError) << "Only kTVMFFIOpaquePyObject is supported for now"; } // create initial opaque object tvm::ffi::ObjectPtr p = tvm::ffi::make_object(handle, deleter); // need to set the type index after creation, because the set to RuntimeTypeIndex() // happens after the constructor is called p->SetTypeIndex(type_index); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(p)); TVM_FFI_SAFE_CALL_END(); } int TVMFFITypeKeyToIndex(const TVMFFIByteArray* type_key, int32_t* out_tindex) { TVM_FFI_SAFE_CALL_BEGIN(); out_tindex[0] = tvm::ffi::TypeTable::Global()->TypeKeyToIndex(type_key); TVM_FFI_SAFE_CALL_END(); } int TVMFFITypeRegisterField(int32_t type_index, const TVMFFIFieldInfo* info) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::TypeTable::Global()->RegisterTypeField(type_index, info); TVM_FFI_SAFE_CALL_END(); } int TVMFFITypeRegisterMethod(int32_t type_index, const TVMFFIMethodInfo* info) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::TypeTable::Global()->RegisterTypeMethod(type_index, info); TVM_FFI_SAFE_CALL_END(); } int TVMFFITypeRegisterMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::TypeTable::Global()->RegisterTypeMetadata(type_index, metadata); TVM_FFI_SAFE_CALL_END(); } int TVMFFITypeRegisterAttr(int32_t type_index, const TVMFFIByteArray* name, const TVMFFIAny* value) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::TypeTable::Global()->RegisterTypeAttr(type_index, name, value); TVM_FFI_SAFE_CALL_END(); } const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray* name) { TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); return tvm::ffi::TypeTable::Global()->GetTypeAttrColumn(name); TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIGetTypeAttrColumn); } int32_t TVMFFITypeGetOrAllocIndex(const TVMFFIByteArray* type_key, int32_t static_type_index, int32_t type_depth, int32_t num_child_slots, int32_t child_slots_can_overflow, int32_t parent_type_index) { TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); tvm::ffi::String s_type_key(type_key->data, type_key->size); return tvm::ffi::TypeTable::Global()->GetOrAllocTypeIndex( s_type_key, static_type_index, type_depth, num_child_slots, child_slots_can_overflow, parent_type_index); TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFITypeGetOrAllocIndex); } const TVMFFITypeInfo* TVMFFIGetTypeInfo(int32_t type_index) { TVM_FFI_LOG_EXCEPTION_CALL_BEGIN(); return tvm::ffi::TypeTable::Global()->GetTypeEntry(type_index); TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIGetTypeInfo); } // string APIs, we blend into object.cc to keep things simple int TVMFFIStringFromByteArray(const TVMFFIByteArray* input, TVMFFIAny* out) { TVM_FFI_SAFE_CALL_BEGIN(); // must set to none first out->type_index = kTVMFFINone; tvm::ffi::TypeTraits::MoveToAny(tvm::ffi::String(input->data, input->size), out); TVM_FFI_SAFE_CALL_END(); } int TVMFFIBytesFromByteArray(const TVMFFIByteArray* input, TVMFFIAny* out) { TVM_FFI_SAFE_CALL_BEGIN(); // must set to none first out->type_index = kTVMFFINone; tvm::ffi::TypeTraits::MoveToAny(tvm::ffi::Bytes(input->data, input->size), out); TVM_FFI_SAFE_CALL_END(); } namespace { TVM_FFI_STATIC_INIT_BLOCK() { namespace ffi = ::tvm::ffi; namespace refl = ::tvm::ffi::reflection; refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef>); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef>); // Skipped: TypeIndex::kTVMFFIModule // Skipped: TypeIndex::kTVMFFIOpaquePyObject refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef>); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef>); refl::ObjectDef(refl::init(false)) .def_ro("_value", &ffi::EnumObj::_value, "Ordinal assigned at registration.", refl::AttachFieldFlag::SEqHashIgnore()) .def_ro("_name", &ffi::EnumObj::_name, "Instance name."); refl::EnsureTypeAttrColumn(refl::type_attr::kEnumEntries); refl::EnsureTypeAttrColumn(refl::type_attr::kEnumAttrs); refl::EnsureTypeAttrColumn(refl::type_attr::kEnumValueEntries); refl::GlobalDef() .def_method("ffi.GetRegisteredTypeKeys", []() -> ffi::Array { return ffi::TypeTable::Global()->GetRegisteredTypeKeys(); }) .def("ffi.GetInvalidObject", ffi::GetMissingObject) .def("ffi.GetKwargsObject", ffi::GetKwargsObject); } } // namespace tvm-ffi-0.1.12/src/ffi/object_internal.h000066400000000000000000000024221521067262500200170ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file object_internal.h * \brief Internal declarations for object sentinel helpers. */ #ifndef TVM_FFI_OBJECT_INTERNAL_H_ #define TVM_FFI_OBJECT_INTERNAL_H_ #include namespace tvm { namespace ffi { /*! \brief Return the singleton "missing" sentinel object. */ ObjectRef GetMissingObject(); /*! \brief Return the singleton KWARGS sentinel object. */ ObjectRef GetKwargsObject(); } // namespace ffi } // namespace tvm #endif // TVM_FFI_OBJECT_INTERNAL_H_ tvm-ffi-0.1.12/src/ffi/tensor.cc000066400000000000000000000106061521067262500163300ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /* * \file src/ffi/tensor.cc * \brief Tensor C API implementation */ #include #include #include #include namespace tvm { namespace ffi { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_packed("ffi.Shape", [](ffi::PackedArgs args, Any* ret) { int64_t* mutable_data; ObjectPtr shape = details::MakeEmptyShape(args.size(), &mutable_data); for (int i = 0; i < args.size(); ++i) { if (auto opt_int = args[i].try_cast()) { mutable_data[i] = *opt_int; } else { TVM_FFI_THROW(ValueError) << "Expect shape to take list of int arguments"; } } *ret = details::ObjectUnsafe::ObjectRefFromObjectPtr(shape); }); } } // namespace ffi } // namespace tvm int TVMFFITensorCreateUnsafeView(TVMFFIObjectHandle source, const DLTensor* prototype, TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::ObjectPtr source_tensor = tvm::ffi::details::ObjectUnsafe::ObjectPtrFromUnowned( static_cast(source)); class ViewNDAlloc { public: explicit ViewNDAlloc(tvm::ffi::ObjectPtr tensor) : tensor_(std::move(tensor)) {} void AllocData(DLTensor* tensor, void* data_ptr) { tensor->data = data_ptr; } void FreeData(DLTensor* tensor) {} private: tvm::ffi::ObjectPtr tensor_; }; void* source_data_ptr = prototype->data; size_t num_extra_i64_at_tail = static_cast(prototype->ndim) * 2; ViewNDAlloc alloc(source_tensor); tvm::ffi::Tensor ret_tensor( tvm::ffi::make_inplace_array_object, int64_t>(num_extra_i64_at_tail, alloc, *prototype, source_data_ptr)); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(ret_tensor)); TVM_FFI_SAFE_CALL_END(); } int TVMFFITensorFromDLPack(DLManagedTensor* from, int32_t min_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::Tensor tensor = tvm::ffi::Tensor::FromDLPack(from, static_cast(min_alignment), require_contiguous); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(tensor)); TVM_FFI_SAFE_CALL_END(); } int TVMFFITensorFromDLPackVersioned(DLManagedTensorVersioned* from, int32_t min_alignment, int32_t require_contiguous, TVMFFIObjectHandle* out) { TVM_FFI_SAFE_CALL_BEGIN(); tvm::ffi::Tensor tensor = tvm::ffi::Tensor::FromDLPackVersioned( from, static_cast(min_alignment), require_contiguous); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(tensor)); TVM_FFI_SAFE_CALL_END(); } int TVMFFITensorToDLPack(TVMFFIObjectHandle from, DLManagedTensor** out) { TVM_FFI_SAFE_CALL_BEGIN(); *out = tvm::ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned( static_cast(from)) ->ToDLPack(); TVM_FFI_SAFE_CALL_END(); } int TVMFFITensorToDLPackVersioned(TVMFFIObjectHandle from, DLManagedTensorVersioned** out) { TVM_FFI_SAFE_CALL_BEGIN(); *out = tvm::ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned( static_cast(from)) ->ToDLPackVersioned(); TVM_FFI_SAFE_CALL_END(); } tvm-ffi-0.1.12/src/ffi/testing/000077500000000000000000000000001521067262500161615ustar00rootroot00000000000000tvm-ffi-0.1.12/src/ffi/testing/testing.cc000066400000000000000000001131201521067262500201430ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ // This file is used for testing the FFI API. // NOTE: TVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1 is set via CMake target_compile_definitions #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { // Step 1: Define the object class (stores the actual data) class TestIntPairObj : public tvm::ffi::Object { public: int64_t a; int64_t b; TestIntPairObj() = default; TestIntPairObj(int64_t a, int64_t b) : a(a), b(b) {} // Required: declare type information TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestIntPair", TestIntPairObj, tvm::ffi::Object); }; // Step 2: Define the reference wrapper (user-facing interface) class TestIntPair : public tvm::ffi::ObjectRef { public: // Constructor explicit TestIntPair(int64_t a, int64_t b) { data_ = tvm::ffi::make_object(a, b); } int64_t Sum() const { return get()->a + get()->b; } // Required: define object reference methods TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestIntPair, tvm::ffi::ObjectRef, TestIntPairObj); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def(refl::init()) .def_ro("a", &TestIntPairObj::a, "Field `a`") .def_ro("b", &TestIntPairObj::b, "Field `b`") .def("sum", &TestIntPair::Sum, "Method to compute sum of a and b"); refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); } // C++-backed enum used by the Python ``Enum`` tests to exercise both // ``EnumDef``-registered entries and the Python ``ClassVar``-based binding. class TestEnumVariantObj : public tvm::ffi::EnumObj { public: TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestEnumVariant", TestEnumVariantObj, tvm::ffi::EnumObj); }; class TestEnumVariant : public tvm::ffi::Enum { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestEnumVariant, tvm::ffi::Enum, TestEnumVariantObj); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; // ObjectDef registers the type on destruction, so the temporary is intentional; // silence clang-tidy's bugprone-unused-raii since the RAII finalisation is the point. refl::ObjectDef(refl::init(false)); // NOLINT(bugprone-unused-raii) refl::TypeAttrDef().def( refl::type_attr::kConvert, &refl::details::FFIConvertFromAnyViewToObjectRef); refl::EnumDef("Alpha").set_attr("code", int64_t{10}); refl::EnumDef("Beta").set_attr("code", int64_t{20}); } class TestObjectBase : public Object { public: int64_t v_i64; double v_f64; String v_str; int64_t AddI64(int64_t other) const { return v_i64 + other; } // declare as one slot, with float as overflow static constexpr bool _type_mutable = true; static constexpr uint32_t _type_child_slots = 1; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestObjectBase", TestObjectBase, Object); }; class TestObjectDerived : public TestObjectBase { public: Map v_map; Array v_array; // declare as one slot, with float as overflow TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestObjectDerived", TestObjectDerived, TestObjectBase); }; class TestCxxClassBase : public Object { public: int64_t v_i64 = 0; int32_t v_i32 = 0; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxClassBase", TestCxxClassBase, Object); }; class TestCxxClassDerived : public TestCxxClassBase { public: double v_f64 = 0.0; float v_f32 = 0.0f; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxClassDerived", TestCxxClassDerived, TestCxxClassBase); }; class TestCxxClassDerivedDerived : public TestCxxClassDerived { public: String v_str; bool v_bool = false; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxClassDerivedDerived", TestCxxClassDerivedDerived, TestCxxClassDerived); }; class TestCxxInitSubsetObj : public Object { public: int64_t required_field = 0; int64_t optional_field = -1; String note; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxInitSubset", TestCxxInitSubsetObj, Object); }; class TestCxxKwOnly : public Object { public: int64_t x = 0; int64_t y = 0; int64_t z = 0; int64_t w = 0; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxKwOnly", TestCxxKwOnly, Object); }; // Test: auto-generated __ffi_init__ with init(false) and KwOnly(true) per-field. // No explicit refl::init<>() — the auto-init is generated in ObjectDef's destructor. class TestCxxAutoInitObj : public Object { public: int64_t a; ///< init=True, positional (default) int64_t b; ///< init=False, has default int64_t c; ///< init=True, kw_only int64_t d; ///< init=True, positional, has default static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxAutoInit", TestCxxAutoInitObj, Object); }; // Test: auto-generated init with all fields init=True (no Init/KwOnly traits). class TestCxxAutoInitSimpleObj : public Object { public: int64_t x; int64_t y; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxAutoInitSimple", TestCxxAutoInitSimpleObj, Object); }; // Test: auto-generated init with all fields excluded from init. class TestCxxAutoInitAllInitOffObj : public Object { public: int64_t x = 7; ///< init=False, has reflection default int64_t y = 0; ///< init=False, has reflection default int64_t z = 1234; ///< init=False, no reflection default (creator default is kept) static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxAutoInitAllInitOff", TestCxxAutoInitAllInitOffObj, Object); }; // Test: mixed positional + kw_only + defaults + init=False field. class TestCxxAutoInitKwOnlyDefaultsObj : public Object { public: int64_t p_required; ///< init=True, positional, required int64_t p_default; ///< init=True, positional, default=11 int64_t k_required; ///< init=True, kw_only, required int64_t k_default; ///< init=True, kw_only, default=22 int64_t hidden; ///< init=False, default=33 static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxAutoInitKwOnlyDefaults", TestCxxAutoInitKwOnlyDefaultsObj, Object); }; // Test: inheritance + auto-generated init. class TestCxxAutoInitParentObj : public Object { public: int64_t parent_required; int64_t parent_default; static constexpr bool _type_mutable = true; static constexpr uint32_t _type_child_slots = 1; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxAutoInitParent", TestCxxAutoInitParentObj, Object); }; class TestCxxAutoInitChildObj : public TestCxxAutoInitParentObj { public: int64_t child_required; int64_t child_kw_only; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestCxxAutoInitChild", TestCxxAutoInitChildObj, TestCxxAutoInitParentObj); }; // Test: init(false) at class level suppresses auto-generated __ffi_init__. class TestCxxNoAutoInitObj : public Object { public: int64_t x; int64_t y; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestCxxNoAutoInit", TestCxxNoAutoInitObj, Object); }; class TestDeepCopyEdgesObj : public Object { public: Any v_any; ObjectRef v_obj; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestDeepCopyEdges", TestDeepCopyEdgesObj, Object); }; class TestNonCopyable : public Object { public: int64_t value; explicit TestNonCopyable(int64_t value) : value(value) {} TestNonCopyable(const TestNonCopyable&) = delete; TestNonCopyable& operator=(const TestNonCopyable&) = delete; TVM_FFI_DECLARE_OBJECT_INFO("testing.TestNonCopyable", TestNonCopyable, Object); }; class TestUnregisteredBaseObject : public Object { public: int64_t v1; explicit TestUnregisteredBaseObject(int64_t v1) : v1(v1) {} int64_t GetV1PlusOne() const { return v1 + 1; } TVM_FFI_DECLARE_OBJECT_INFO("testing.TestUnregisteredBaseObject", TestUnregisteredBaseObject, Object); }; class TestUnregisteredObject : public TestUnregisteredBaseObject { public: int64_t v2; explicit TestUnregisteredObject(int64_t v1, int64_t v2) : TestUnregisteredBaseObject(v1), v2(v2) {} int64_t GetV2PlusTwo() const { return v2 + 2; } TVM_FFI_DECLARE_OBJECT_INFO("testing.TestUnregisteredObject", TestUnregisteredObject, TestUnregisteredBaseObject); }; class TestCompareObj : public Object { public: int64_t key; String name; int64_t ignored_field; TestCompareObj() = default; TestCompareObj(int64_t key, String name, int64_t ignored_field) : key(key), name(std::move(name)), ignored_field(ignored_field) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestCompare", TestCompareObj, Object); }; class TestCompare : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestCompare, ObjectRef, TestCompareObj); }; class TestHashObj : public Object { public: int64_t key; String name; int64_t hash_ignored; TestHashObj() = default; TestHashObj(int64_t key, String name, int64_t hash_ignored) : key(key), name(std::move(name)), hash_ignored(hash_ignored) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestHash", TestHashObj, Object); }; class TestHash : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestHash, ObjectRef, TestHashObj); }; class TestCustomHashObj : public Object { public: int64_t key; String label; TestCustomHashObj() = default; TestCustomHashObj(int64_t key, String label) : key(key), label(std::move(label)) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestCustomHash", TestCustomHashObj, Object); }; class TestCustomHash : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestCustomHash, ObjectRef, TestCustomHashObj); }; class TestCustomCompareObj : public Object { public: int64_t key; String label; TestCustomCompareObj() = default; TestCustomCompareObj(int64_t key, String label) : key(key), label(std::move(label)) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestCustomCompare", TestCustomCompareObj, Object); }; class TestCustomCompare : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestCustomCompare, ObjectRef, TestCustomCompareObj); }; // Test object with __ffi_eq__ but deliberately no __ffi_hash__. // Used to verify the RecursiveHash guard that rejects eq-without-hash types. class TestEqWithoutHashObj : public Object { public: int64_t key; String label; TestEqWithoutHashObj() = default; TestEqWithoutHashObj(int64_t key, String label) : key(key), label(std::move(label)) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestEqWithoutHash", TestEqWithoutHashObj, Object); }; class TestEqWithoutHash : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestEqWithoutHash, ObjectRef, TestEqWithoutHashObj); }; // NOLINTNEXTLINE(performance-unnecessary-value-param) TVM_FFI_NO_INLINE void TestRaiseError(String kind, String msg) { // keep name and no liner for testing backtrace throw ffi::Error(kind, msg, TVMFFIBacktrace(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0)); } TVM_FFI_NO_INLINE void TestApply(PackedArgs args, Any* ret) { // keep name and no liner for testing backtrace auto f = args[0].cast(); f.CallPacked(args.Slice(1), ret); } // NOLINTNEXTLINE(bugprone-reserved-identifier) int __add_one_c_symbol(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* ret) { TVM_FFI_SAFE_CALL_BEGIN(); int x = reinterpret_cast(args)[0].cast(); reinterpret_cast(ret)[0] = x + 1; TVM_FFI_SAFE_CALL_END(); } void _mlir_add_one_c_symbol(void** packed_args) { void* handle = *reinterpret_cast(packed_args[0]); const TVMFFIAny* args = *reinterpret_cast(packed_args[1]); int32_t num_args = *reinterpret_cast(packed_args[2]); TVMFFIAny* rv = *reinterpret_cast(packed_args[3]); int* ret_code = reinterpret_cast(packed_args[4]); *ret_code = __add_one_c_symbol(handle, args, num_args, rv); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_rw("v_i64", &TestObjectBase::v_i64, refl::default_value(10), "i64 field") .def_ro("v_f64", &TestObjectBase::v_f64, refl::default_value(10.0)) .def_rw("v_str", &TestObjectBase::v_str, refl::default_value("hello")) .def("add_i64", &TestObjectBase::AddI64, "add_i64 method"); refl::ObjectDef() .def_rw("v_map", &TestObjectDerived::v_map) .def_rw("v_array", &TestObjectDerived::v_array); refl::ObjectDef() .def_rw("v_i64", &TestCxxClassBase::v_i64, refl::repr(false)) .def_rw("v_i32", &TestCxxClassBase::v_i32, refl::repr(false)); refl::ObjectDef() .def_rw("v_f64", &TestCxxClassDerived::v_f64) .def_rw("v_f32", &TestCxxClassDerived::v_f32, refl::default_value(float{8.0f})); refl::ObjectDef() .def_rw("v_str", &TestCxxClassDerivedDerived::v_str, refl::default_value(String("default"))) .def_rw("v_bool", &TestCxxClassDerivedDerived::v_bool); refl::ObjectDef() .def_rw("required_field", &TestCxxInitSubsetObj::required_field) .def_rw("optional_field", &TestCxxInitSubsetObj::optional_field, refl::init(false), refl::default_value(int64_t{-1})) .def_rw("note", &TestCxxInitSubsetObj::note, refl::init(false), refl::default_value(String("default"))); refl::ObjectDef() .def_rw("x", &TestCxxKwOnly::x, refl::kw_only(true)) .def_rw("y", &TestCxxKwOnly::y, refl::kw_only(true)) .def_rw("z", &TestCxxKwOnly::z, refl::kw_only(true)) .def_rw("w", &TestCxxKwOnly::w, refl::kw_only(true), refl::default_value(int64_t{100})); // No refl::init<>() — auto-generates __ffi_init__ in ObjectDef destructor. refl::ObjectDef() .def_rw("a", &TestCxxAutoInitObj::a) .def_rw("b", &TestCxxAutoInitObj::b, refl::init(false), refl::default_value(int64_t{42})) .def_rw("c", &TestCxxAutoInitObj::c, refl::kw_only(true)) .def_rw("d", &TestCxxAutoInitObj::d, refl::default_value(int64_t{99})); refl::ObjectDef() .def_rw("x", &TestCxxAutoInitSimpleObj::x) .def_rw("y", &TestCxxAutoInitSimpleObj::y); refl::ObjectDef() .def_rw("x", &TestCxxAutoInitAllInitOffObj::x, refl::init(false), refl::default_value(int64_t{7})) .def_rw("y", &TestCxxAutoInitAllInitOffObj::y, refl::init(false), refl::default_value(int64_t{9})) .def_rw("z", &TestCxxAutoInitAllInitOffObj::z, refl::init(false)); refl::ObjectDef() .def_rw("p_required", &TestCxxAutoInitKwOnlyDefaultsObj::p_required) .def_rw("p_default", &TestCxxAutoInitKwOnlyDefaultsObj::p_default, refl::default_value(int64_t{11})) .def_rw("k_required", &TestCxxAutoInitKwOnlyDefaultsObj::k_required, refl::kw_only(true)) .def_rw("k_default", &TestCxxAutoInitKwOnlyDefaultsObj::k_default, refl::kw_only(true), refl::default_value(int64_t{22})) .def_rw("hidden", &TestCxxAutoInitKwOnlyDefaultsObj::hidden, refl::init(false), refl::default_value(int64_t{33})); refl::ObjectDef() .def_rw("parent_required", &TestCxxAutoInitParentObj::parent_required) .def_rw("parent_default", &TestCxxAutoInitParentObj::parent_default, refl::default_value(int64_t{5})); refl::ObjectDef() .def_rw("child_required", &TestCxxAutoInitChildObj::child_required) .def_rw("child_kw_only", &TestCxxAutoInitChildObj::child_kw_only, refl::kw_only(true)); // init(false) at class level: has fields, has creator, but no __ffi_init__. refl::ObjectDef(refl::init(false)) .def_rw("x", &TestCxxNoAutoInitObj::x) .def_rw("y", &TestCxxNoAutoInitObj::y); refl::ObjectDef() .def_rw("v_any", &TestDeepCopyEdgesObj::v_any) .def_rw("v_obj", &TestDeepCopyEdgesObj::v_obj); refl::ObjectDef() .def(refl::init()) .def_ro("value", &TestNonCopyable::value); refl::ObjectDef() .def(refl::init(), "Constructor of TestUnregisteredBaseObject") .def_ro("v1", &TestUnregisteredBaseObject::v1) .def("get_v1_plus_one", &TestUnregisteredBaseObject::GetV1PlusOne, "Get (v1 + 1) from TestUnregisteredBaseObject"); refl::ObjectDef() .def(refl::init(), "Constructor of TestUnregisteredObject") .def_ro("v2", &TestUnregisteredObject::v2) .def("get_v2_plus_two", &TestUnregisteredObject::GetV2PlusTwo, "Get (v2 + 2) from TestUnregisteredObject"); refl::ObjectDef() .def(refl::init()) .def_ro("key", &TestCompareObj::key) .def_ro("name", &TestCompareObj::name) .def_ro("ignored_field", &TestCompareObj::ignored_field, refl::compare(false)); refl::ObjectDef() .def(refl::init()) .def_ro("key", &TestHashObj::key) .def_ro("name", &TestHashObj::name) .def_ro("hash_ignored", &TestHashObj::hash_ignored, refl::hash(false)); refl::ObjectDef() .def(refl::init()) .def_ro("key", &TestCustomHashObj::key) .def_ro("label", &TestCustomHashObj::label); refl::ObjectDef() .def(refl::init()) .def_ro("key", &TestCustomCompareObj::key) .def_ro("label", &TestCustomCompareObj::label); refl::ObjectDef() .def(refl::init()) .def_ro("key", &TestEqWithoutHashObj::key) .def_ro("label", &TestEqWithoutHashObj::label); refl::GlobalDef() .def("testing.test_raise_error", TestRaiseError) .def("testing.add_one", [](int x) { return x + 1; }) .def_packed("testing.nop", [](PackedArgs args, Any* ret) {}) .def_packed("testing.echo", [](PackedArgs args, Any* ret) { *ret = args[0]; }) .def_packed("testing.apply", TestApply) .def("testing.run_check_signal", [](int nsec) { for (int i = 0; i < nsec; ++i) { if (TVMFFIEnvCheckSignals() != 0) { throw ffi::EnvErrorAlreadySet(); } std::this_thread::sleep_for(std::chrono::seconds(1)); } std::cout << "Function finished without catching signal" << std::endl; }) .def("testing.object_use_count", [](const Object* obj) { return obj->use_count(); }) .def("testing.make_unregistered_object", []() { return ObjectRef(make_object(41, 42)); }) .def("testing.get_add_one_c_symbol", []() { TVMFFISafeCallType symbol = __add_one_c_symbol; // NOLINTNEXTLINE(bugprone-casting-through-void) return reinterpret_cast(reinterpret_cast(symbol)); }) .def("testing.get_mlir_add_one_c_symbol", []() { // NOLINTNEXTLINE(bugprone-casting-through-void) return reinterpret_cast(reinterpret_cast(_mlir_add_one_c_symbol)); }) .def("testing.optional_tensor_view_has_value", [](const Optional& t) { return t.has_value(); }) .def("testing.enum_variant_get", [](const String& name) -> Enum { return EnumObj::Get(name); }) .def_method("testing.TestIntPairSum", &TestIntPair::Sum, "Get sum of the pair") // Container-with-tensor test helpers for DLPack container conversion // NOLINTBEGIN(performance-unnecessary-value-param) .def("testing.make_array_with_tensor", [](Tensor t) -> Array { return {std::move(t)}; }) .def("testing.make_array_with_mixed", [](Tensor t, int64_t x) -> Array { return {std::move(t), x, String("hello")}; }) .def("testing.make_nested_array_with_tensor", [](const Tensor& t) -> Array { Array inner{t, 42}; return {std::move(inner), t}; }) .def("testing.make_list_with_tensor", [](Tensor t, int64_t x) -> List { List result; result.push_back(std::move(t)); result.push_back(x); return result; }) .def("testing.make_map_with_tensor", [](Tensor t) -> Map { Map result; result.Set("tensor", std::move(t)); result.Set("value", 42); result.Set("name", String("test")); return result; }) .def("testing.make_dict_with_tensor", [](Tensor t) -> Dict { Dict result; result.Set("tensor", std::move(t)); result.Set("value", 42); return result; }) .def("testing.make_empty_array_with_tensor_input", [](const Tensor& t) -> Array { return Array(); }) .def("testing.make_nested_map_with_tensor", [](const Tensor& t1, Tensor t2) -> Map { Array arr{t1, std::move(t2)}; Map inner; inner.Set("t", t1); Map result; result.Set("array", std::move(arr)); result.Set("map", std::move(inner)); result.Set("scalar", 99); return result; }); // NOLINTEND(performance-unnecessary-value-param) } } // namespace ffi } // namespace tvm // ----------------------------------------------------------------------------- // Additional comprehensive schema coverage // ----------------------------------------------------------------------------- namespace tvm { namespace ffi { // ----------------------------------------------------------------------------- // Implementation functions for schema testing // ----------------------------------------------------------------------------- namespace schema_test_impl { // Simple types int64_t schema_id_int(int64_t x) { return x; } double schema_id_float(double x) { return x; } bool schema_id_bool(bool x) { return x; } DLDevice schema_id_device(DLDevice d) { return d; } DLDataType schema_id_dtype(DLDataType dt) { return dt; } String schema_id_string(String s) { return s; } Bytes schema_id_bytes(Bytes b) { return b; } Function schema_id_func(Function f) { return f; } TypedFunction schema_id_func_typed( TypedFunction f) { return f; } Any schema_id_any(Any a) { return a; } ObjectRef schema_id_object(ObjectRef o) { return o; } DLTensor* schema_id_dltensor(DLTensor* t) { return t; } Tensor schema_id_tensor(Tensor t) { return t; } void schema_tensor_view_input(TensorView t) {} // Optional types Optional schema_id_opt_int(Optional o) { return o; } Optional schema_id_opt_str(Optional o) { return o; } Optional schema_id_opt_obj(Optional o) { return o; } // Array types Array schema_id_arr_int(Array arr) { return arr; } Array schema_id_arr_str(Array arr) { return arr; } Array schema_id_arr_obj(Array arr) { return arr; } const ArrayObj* schema_id_arr(const ArrayObj* arr) { return arr; } // Map types Map schema_id_map_str_int(Map m) { return m; } Map schema_id_map_str_str(Map m) { return m; } Map schema_id_map_str_obj(Map m) { return m; } const MapObj* schema_id_map(const MapObj* m) { return m; } // Variant types Variant schema_id_variant_int_str(Variant v) { return v; } Variant> schema_variant_mix( Variant> v) { return v; } // List types List schema_id_list_int(List lst) { return lst; } List schema_id_list_str(List lst) { return lst; } List schema_id_list_obj(List lst) { return lst; } // Dict types Dict schema_id_dict_str_int(Dict d) { return d; } Dict schema_id_dict_str_str(Dict d) { return d; } // Complex nested types Map> schema_arr_map_opt(const Array>& arr, Map> mp, const Optional& os) { // no-op combine if (os.has_value()) { Array extra; for (const auto& i : arr) { if (i.has_value()) extra.push_back(i.value()); } mp.Set(os.value(), extra); } return mp; } // Edge cases int64_t schema_no_args() { return 1; } void schema_no_return(int64_t x) {} void schema_no_args_no_return() {} // Member function pattern int64_t test_int_pair_sum_wrapper(const TestIntPair& target) { return target.Sum(); } // Documentation export int64_t test_add_with_docstring(int64_t a, int64_t b) { return a + b; } } // namespace schema_test_impl // A class with a wide variety of field types and method signatures class SchemaAllTypesObj : public Object { public: // POD and builtin types bool v_bool{true}; int64_t v_int; double v_float; DLDevice v_device; DLDataType v_dtype; // Atomic object types String v_string; Bytes v_bytes; // Containers and combinations Optional v_opt_int; Optional v_opt_str; Array v_arr_int; Array v_arr_str; Map v_map_str_int; Map> v_map_str_arr_int; Variant, Map> v_variant; Optional>> v_opt_arr_variant; // Constructor used by refl::init in make_with SchemaAllTypesObj(int64_t vi, double vf, String s) // NOLINT(*): explicit not necessary here : v_int(vi), v_float(vf), v_device(TVMFFIDLDeviceFromIntPair(kDLCPU, 0)), v_dtype(StringToDLDataType("float32")), v_string(std::move(s)), v_variant(String("v")) {} // Some methods to exercise RegisterMethod int64_t AddInt(int64_t x) const { return v_int + x; } Array AppendInt(Array xs, int64_t y) const { xs.push_back(y); return xs; } Optional MaybeConcat(Optional a, Optional b) const { if (a.has_value() && b.has_value()) return String(a.value() + b.value()); if (a.has_value()) return a; if (b.has_value()) return b; return Optional(std::nullopt); } Map> MergeMap(Map> lhs, // NOLINTNEXTLINE(performance-unnecessary-value-param) Map> rhs) const { for (const auto& kv : rhs) { if (!lhs.count(kv.first)) { lhs.Set(kv.first, kv.second); } else { Array arr = lhs[kv.first]; for (const auto& v : kv.second) arr.push_back(v); lhs.Set(kv.first, arr); } } return lhs; } static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("testing.SchemaAllTypes", SchemaAllTypesObj, Object); }; class SchemaAllTypes : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SchemaAllTypes, ObjectRef, SchemaAllTypesObj); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; // Register fields of various types (RegisterField usage) refl::ObjectDef() .def_rw("v_bool", &SchemaAllTypesObj::v_bool, "bool field", refl::Metadata{{"bool_attr", true}, // {"int_attr", 1}, // {"str_attr", "hello"}}) .def_rw("v_int", &SchemaAllTypesObj::v_int, refl::default_value(0), "int field") .def_rw("v_float", &SchemaAllTypesObj::v_float, refl::default_value(0.0), "float field") .def_rw("v_device", &SchemaAllTypesObj::v_device, "device field") .def_rw("v_dtype", &SchemaAllTypesObj::v_dtype, "dtype field") .def_rw("v_string", &SchemaAllTypesObj::v_string, refl::default_value("s"), "string field") .def_rw("v_bytes", &SchemaAllTypesObj::v_bytes, "bytes field") .def_rw("v_opt_int", &SchemaAllTypesObj::v_opt_int, "optional int") .def_rw("v_opt_str", &SchemaAllTypesObj::v_opt_str, "optional str") .def_rw("v_arr_int", &SchemaAllTypesObj::v_arr_int, "array") .def_rw("v_arr_str", &SchemaAllTypesObj::v_arr_str, "array") .def_rw("v_map_str_int", &SchemaAllTypesObj::v_map_str_int, "map") .def_rw("v_map_str_arr_int", &SchemaAllTypesObj::v_map_str_arr_int, "map>") .def_rw("v_variant", &SchemaAllTypesObj::v_variant, "variant,map>") .def_rw("v_opt_arr_variant", &SchemaAllTypesObj::v_opt_arr_variant, "optional>>") // Register methods (RegisterMethod usage) .def("add_int", &SchemaAllTypesObj::AddInt, "add int method", refl::Metadata{{"bool_attr", true}, // {"int_attr", 1}, // {"str_attr", "hello"}}) .def("append_int", &SchemaAllTypesObj::AppendInt, "append int to array") .def("maybe_concat", &SchemaAllTypesObj::MaybeConcat, "optional concat") .def("merge_map", &SchemaAllTypesObj::MergeMap, "merge maps") // Register static creator (also a static method) .def_static( "make_with", [](int64_t vi, double vf, String s) { return SchemaAllTypes(make_object(vi, vf, std::move(s))); }, "Constructor with subset of fields"); // Global typed functions to exercise RegisterFunc with various schemas refl::GlobalDef() .def("testing.schema_id_int", schema_test_impl::schema_id_int, refl::Metadata{{"bool_attr", true}, // {"int_attr", 1}, // {"str_attr", "hello"}}) .def("testing.schema_id_float", schema_test_impl::schema_id_float) .def("testing.schema_id_bool", schema_test_impl::schema_id_bool) .def("testing.schema_id_device", schema_test_impl::schema_id_device) .def("testing.schema_id_dtype", schema_test_impl::schema_id_dtype) .def("testing.schema_id_string", schema_test_impl::schema_id_string) .def("testing.schema_id_bytes", schema_test_impl::schema_id_bytes) .def("testing.schema_id_func", schema_test_impl::schema_id_func) .def("testing.schema_id_func_typed", schema_test_impl::schema_id_func_typed) .def("testing.schema_id_any", schema_test_impl::schema_id_any) .def("testing.schema_id_object", schema_test_impl::schema_id_object) .def("testing.schema_id_dltensor", schema_test_impl::schema_id_dltensor) .def("testing.schema_id_tensor", schema_test_impl::schema_id_tensor) .def("testing.schema_tensor_view_input", schema_test_impl::schema_tensor_view_input) .def("testing.schema_id_opt_int", schema_test_impl::schema_id_opt_int) .def("testing.schema_id_opt_str", schema_test_impl::schema_id_opt_str) .def("testing.schema_id_opt_obj", schema_test_impl::schema_id_opt_obj) .def("testing.schema_id_arr_int", schema_test_impl::schema_id_arr_int) .def("testing.schema_id_arr_str", schema_test_impl::schema_id_arr_str) .def("testing.schema_id_arr_obj", schema_test_impl::schema_id_arr_obj) .def("testing.schema_id_arr", schema_test_impl::schema_id_arr) .def("testing.schema_id_list_int", schema_test_impl::schema_id_list_int) .def("testing.schema_id_list_str", schema_test_impl::schema_id_list_str) .def("testing.schema_id_list_obj", schema_test_impl::schema_id_list_obj) .def("testing.schema_id_map_str_int", schema_test_impl::schema_id_map_str_int) .def("testing.schema_id_map_str_str", schema_test_impl::schema_id_map_str_str) .def("testing.schema_id_map_str_obj", schema_test_impl::schema_id_map_str_obj) .def("testing.schema_id_map", schema_test_impl::schema_id_map) .def("testing.schema_id_dict_str_int", schema_test_impl::schema_id_dict_str_int) .def("testing.schema_id_dict_str_str", schema_test_impl::schema_id_dict_str_str) .def("testing.schema_id_variant_int_str", schema_test_impl::schema_id_variant_int_str) .def_packed("testing.schema_packed", [](PackedArgs args, Any* ret) {}) .def("testing.schema_arr_map_opt", schema_test_impl::schema_arr_map_opt) .def("testing.schema_variant_mix", schema_test_impl::schema_variant_mix, "variant passthrough") .def("testing.schema_no_args", schema_test_impl::schema_no_args) .def("testing.schema_no_return", schema_test_impl::schema_no_return) .def("testing.schema_no_args_no_return", schema_test_impl::schema_no_args_no_return); TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one", reinterpret_cast(__add_one_c_symbol)); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; // Register __ffi_hash__ for TestCustomHash: only hashes `key`, ignores `label`. refl::TypeAttrDef().def( refl::type_attr::kHash, [](const Object* self, const Function& fn_hash) -> int64_t { auto* obj = static_cast(self); return fn_hash(AnyView(obj->key)).cast(); }); // Register __ffi_hash__ for TestCustomCompare: only hashes `key`, consistent with eq/compare. refl::TypeAttrDef().def( refl::type_attr::kHash, [](const Object* self, const Function& fn_hash) -> int64_t { auto* obj = static_cast(self); return fn_hash(AnyView(obj->key)).cast(); }); // Register __ffi_eq__ for TestCustomCompare: compares only `key`. refl::TypeAttrDef().def( refl::type_attr::kEq, [](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool { auto* a = static_cast(lhs); auto* b = static_cast(rhs); return fn_eq(AnyView(a->key), AnyView(b->key)).cast(); }); // Register __ffi_compare__ for TestCustomCompare: three-way ordering on `key`. refl::TypeAttrDef().def( refl::type_attr::kCompare, [](const Object* lhs, const Object* rhs, const Function& fn_cmp) -> int32_t { auto* a = static_cast(lhs); auto* b = static_cast(rhs); return fn_cmp(AnyView(a->key), AnyView(b->key)).cast(); }); // Register __ffi_eq__ for TestEqWithoutHash: deliberately no __ffi_hash__. // This exercises the RecursiveHash guard that rejects eq-without-hash types. refl::TypeAttrDef().def( refl::type_attr::kEq, [](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool { auto* a = static_cast(lhs); auto* b = static_cast(rhs); return fn_eq(AnyView(a->key), AnyView(b->key)).cast(); }); } } // namespace ffi } // namespace tvm // ----------------------------------------------------------------------------- // Exported symbols for metadata testing on DLL-exported functions // ----------------------------------------------------------------------------- // We keep minimal DLL exports here to verify the export mechanism. TVM_FFI_DLL_EXPORT_TYPED_FUNC(testing_dll_schema_id_int, tvm::ffi::schema_test_impl::schema_id_int) // Documentation export TVM_FFI_DLL_EXPORT_TYPED_FUNC(testing_dll_test_add_with_docstring, tvm::ffi::schema_test_impl::test_add_with_docstring); TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(testing_dll_test_add_with_docstring, R"(Add two integers and return the sum. Parameters ---------- a : int First integer b : int Second integer Returns ------- result : int Sum of a and b)"); extern "C" TVM_FFI_DLL_EXPORT int TVMFFITestingDummyTarget() { return 0; } tvm-ffi-0.1.12/tests/000077500000000000000000000000001521067262500143135ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/cpp/000077500000000000000000000000001521067262500150755ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/cpp/CMakeLists.txt000066400000000000000000000017451521067262500176440ustar00rootroot00000000000000file(GLOB _test_sources "${CMAKE_CURRENT_SOURCE_DIR}/test*.cc") file(GLOB _test_extra_sources "${CMAKE_CURRENT_SOURCE_DIR}/extra/test*.cc") if (TVM_FFI_USE_EXTRA_CXX_API) list(APPEND _test_sources ${_test_extra_sources}) endif () add_executable(tvm_ffi_tests ${_test_sources}) set_target_properties( tvm_ffi_tests PROPERTIES CXX_STANDARD 17 CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) tvm_ffi_add_cxx_warning(tvm_ffi_tests) add_sanitizer_address(tvm_ffi_tests) tvm_ffi_add_apple_dsymutil(tvm_ffi_tests) tvm_ffi_add_msvc_flags(tvm_ffi_tests) target_link_libraries(tvm_ffi_tests PRIVATE tvm_ffi_shared) target_link_libraries(tvm_ffi_tests PRIVATE tvm_ffi_testing) tvm_ffi_add_gtest(tvm_ffi_tests) if (MSVC) target_link_options(tvm_ffi_tests PRIVATE /DEBUG) endif () tvm-ffi-0.1.12/tests/cpp/extra/000077500000000000000000000000001521067262500162205ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/cpp/extra/test_access_path_repr.cc000066400000000000000000000030411521067262500230710ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include namespace { TEST(AccessPathRepr, StreamlinedFormat) { using namespace tvm::ffi; using namespace tvm::ffi::reflection; AccessPath p = AccessPath::Root() ->Attr("config") ->Attr("layers") ->ArrayItem(0) ->MapItem(String("name")) ->AttrMissing("weights") ->ArrayItemMissing(3) ->MapItemMissing(String("bias")); EXPECT_EQ(std::string(ReprPrint(Any(p)).c_str()), ".config.layers[0][\"name\"]" "[][][]"); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_c_env_api.cc000066400000000000000000000300521521067262500215110ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include namespace { using namespace tvm::ffi; struct CPUNDAlloc { void AllocData(DLTensor* tensor) { tensor->data = malloc(GetDataSize(*tensor)); } void FreeData(DLTensor* tensor) { free(tensor->data); } }; inline Tensor Empty(const Shape& shape, DLDataType dtype, DLDevice device) { return Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); } int TestDLPackManagedTensorAllocator(DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx, void (*SetError)(void* error_ctx, const char* kind, const char* message)) { Shape shape(prototype->shape, prototype->shape + prototype->ndim); Tensor nd = Empty(shape, prototype->dtype, prototype->device); *out = nd.ToDLPackVersioned(); return 0; } int TestDLPackManagedTensorAllocatorError(DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx, void (*SetError)(void* error_ctx, const char* kind, const char* message)) { SetError(error_ctx, "MemoryError", "TestDLPackManagedTensorAllocatorError"); return -1; } TEST(CEnvAPI, TVMFFIEnvSetDLPackManagedTensorAllocator) { auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator(); DLPackManagedTensorAllocator pre_allocator; TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocator, 0, &pre_allocator); EXPECT_EQ(old_allocator, pre_allocator); TVMFFIEnvSetDLPackManagedTensorAllocator(old_allocator, 0, nullptr); } TEST(CEnvAPI, TVMFFIEnvTensorAlloc) { auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator(); TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocator, 0, nullptr); Tensor tensor = Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); EXPECT_EQ(tensor.use_count(), 1); EXPECT_EQ(tensor.shape().size(), 3); EXPECT_EQ(tensor.size(0), 1); EXPECT_EQ(tensor.size(1), 2); EXPECT_EQ(tensor.size(2), 3); EXPECT_EQ(tensor.dtype().code, kDLFloat); EXPECT_EQ(tensor.dtype().bits, 32); EXPECT_EQ(tensor.dtype().lanes, 1); EXPECT_EQ(tensor.device().device_type, kDLCPU); EXPECT_EQ(tensor.device().device_id, 0); EXPECT_NE(tensor.data_ptr(), nullptr); TVMFFIEnvSetDLPackManagedTensorAllocator(old_allocator, 0, nullptr); } TEST(CEnvAPI, TVMFFIEnvTensorAllocError) { auto old_allocator = TVMFFIEnvGetDLPackManagedTensorAllocator(); TVMFFIEnvSetDLPackManagedTensorAllocator(TestDLPackManagedTensorAllocatorError, 0, nullptr); EXPECT_THROW( { try { Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); } catch (const tvm::ffi::Error& e) { EXPECT_EQ(e.kind(), "MemoryError"); EXPECT_EQ(e.message(), "TestDLPackManagedTensorAllocatorError"); throw; } }, tvm::ffi::Error); TVMFFIEnvSetDLPackManagedTensorAllocator(old_allocator, 0, nullptr); } // Helper functions for TVMFFIHandleInitDeinitOnce test static int InitSuccess(void** handle_addr) { *handle_addr = new int(42); return 0; } static int InitShouldNotBeCalled(void** handle_addr) { *handle_addr = new int(999); return 0; } static int DeinitSuccess(void* h) { delete static_cast(h); return 0; } static int DeinitShouldNotBeCalled(void* h) { // Should not be called when handle is already null return -1; } static int InitWithError(void** handle_addr) { TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Initialization failed"); return -1; } static int InitReturnsNull(void** handle_addr) { *handle_addr = nullptr; // Invalid: must return non-null handle return 0; } static int InitForDeinitError(void** handle_addr) { *handle_addr = new int(100); return 0; } static int DeinitWithError(void* h) { delete static_cast(h); TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Deinitialization failed"); return -1; } static int InitValue123(void** handle_addr) { *handle_addr = new int(123); return 0; } static int InitValue456(void** handle_addr) { *handle_addr = new int(456); return 0; } TEST(CEnvAPI, TVMFFIHandleInitDeinitOnce) { // Test 1: Successful initialization void* handle = nullptr; int ret = TVMFFIHandleInitOnce(&handle, InitSuccess); EXPECT_EQ(ret, 0); EXPECT_NE(handle, nullptr); EXPECT_EQ(*static_cast(handle), 42); // Test 2: Multiple init calls should not re-initialize (idempotent) void* original_handle = handle; ret = TVMFFIHandleInitOnce(&handle, InitShouldNotBeCalled); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, original_handle); // Handle should remain unchanged EXPECT_EQ(*static_cast(handle), 42); // Value should still be 42 // Test 3: Successful deinitialization ret = TVMFFIHandleDeinitOnce(&handle, DeinitSuccess); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, nullptr); // Test 4: Multiple deinit calls should be safe (idempotent) ret = TVMFFIHandleDeinitOnce(&handle, DeinitShouldNotBeCalled); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, nullptr); // Test 5: Init error - init_func returns error code void* handle2 = nullptr; ret = TVMFFIHandleInitOnce(&handle2, InitWithError); EXPECT_NE(ret, 0); EXPECT_EQ(handle2, nullptr); // Test 6: Init error - init_func returns nullptr (invalid) void* handle3 = nullptr; ret = TVMFFIHandleInitOnce(&handle3, InitReturnsNull); EXPECT_NE(ret, 0); EXPECT_EQ(handle3, nullptr); // Test 7: Deinit error - deinit_func returns error void* handle4 = nullptr; ret = TVMFFIHandleInitOnce(&handle4, InitForDeinitError); EXPECT_EQ(ret, 0); EXPECT_NE(handle4, nullptr); ret = TVMFFIHandleDeinitOnce(&handle4, DeinitWithError); EXPECT_NE(ret, 0); EXPECT_EQ(handle4, nullptr); // Handle should still be set to nullptr // Test 8: Init-deinit lifecycle void* handle5 = nullptr; ret = TVMFFIHandleInitOnce(&handle5, InitValue123); EXPECT_EQ(ret, 0); EXPECT_NE(handle5, nullptr); EXPECT_EQ(*static_cast(handle5), 123); ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess); EXPECT_EQ(ret, 0); EXPECT_EQ(handle5, nullptr); // Test 9: Ensure subsequent init after deinit works ret = TVMFFIHandleInitOnce(&handle5, InitValue456); EXPECT_EQ(ret, 0); EXPECT_NE(handle5, nullptr); EXPECT_EQ(*static_cast(handle5), 456); // Clean up ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess); EXPECT_EQ(ret, 0); } // Helper functions and data for multithreaded test struct ThreadSafeCounter { int value; std::atomic* init_count_ptr; std::atomic* deinit_count_ptr; ThreadSafeCounter(int v, std::atomic* init_ptr, std::atomic* deinit_ptr) : value(v), init_count_ptr(init_ptr), deinit_count_ptr(deinit_ptr) {} }; // Global pointers for the current test counters static std::atomic* g_init_count = nullptr; static std::atomic* g_deinit_count = nullptr; static int InitWithCounter(void** handle_addr) { auto* counter = new ThreadSafeCounter(42, g_init_count, g_deinit_count); if (counter->init_count_ptr) { counter->init_count_ptr->fetch_add(1, std::memory_order_relaxed); } // Small delay to increase the race window std::this_thread::sleep_for(std::chrono::microseconds(100)); *handle_addr = counter; return 0; } static int DeinitWithCounter(void* h) { auto* counter = static_cast(h); if (counter->deinit_count_ptr) { counter->deinit_count_ptr->fetch_add(1, std::memory_order_relaxed); } // Small delay to increase the race window std::this_thread::sleep_for(std::chrono::microseconds(100)); delete counter; return 0; } TEST(CEnvAPI, TVMFFIHandleInitDeinitOnceMultithreaded) { // Test 1: Multiple threads calling InitOnce - should initialize only once { void* handle = nullptr; const int num_threads = 4; std::vector threads; threads.reserve(num_threads); std::vector results(num_threads); std::mutex mtx; std::condition_variable cv; bool ready = false; std::atomic init_count{0}; // Set global counter pointers g_init_count = &init_count; g_deinit_count = nullptr; // Create threads that all try to initialize simultaneously for (int i = 0; i < num_threads; ++i) { threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() { // Wait for all threads to be ready std::unique_lock lock(mtx); cv.wait(lock, [&ready] { return ready; }); lock.unlock(); results[i] = TVMFFIHandleInitOnce(&handle, InitWithCounter); }); } // Signal all threads to start { std::scoped_lock lock(mtx); ready = true; } cv.notify_all(); // Wait for all threads to complete for (auto& t : threads) { t.join(); } // All threads should succeed for (int i = 0; i < num_threads; ++i) { EXPECT_EQ(results[i], 0); } // Handle should be initialized EXPECT_NE(handle, nullptr); auto* counter = static_cast(handle); EXPECT_EQ(counter->value, 42); // Init should have been called exactly once EXPECT_EQ(init_count.load(), 1); // Clean up int ret = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter); EXPECT_EQ(ret, 0); // Reset global pointers g_init_count = nullptr; } // Test 2: Multiple threads calling DeinitOnce - should deinitialize only once { void* handle = nullptr; std::atomic init_count{0}; std::atomic deinit_count{0}; // Set global counter pointers g_init_count = &init_count; g_deinit_count = &deinit_count; // Initialize first int ret = TVMFFIHandleInitOnce(&handle, InitWithCounter); EXPECT_EQ(ret, 0); EXPECT_NE(handle, nullptr); const int num_threads = 4; std::vector threads; threads.reserve(num_threads); std::vector results(num_threads); std::mutex mtx; std::condition_variable cv; bool ready = false; // Create threads that all try to deinitialize simultaneously for (int i = 0; i < num_threads; ++i) { threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() { // Wait for all threads to be ready std::unique_lock lock(mtx); cv.wait(lock, [&ready] { return ready; }); lock.unlock(); results[i] = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter); }); } // Signal all threads to start { std::scoped_lock lock(mtx); ready = true; } cv.notify_all(); // Wait for all threads to complete for (auto& t : threads) { t.join(); } // All threads should succeed for (int i = 0; i < num_threads; ++i) { EXPECT_EQ(results[i], 0); } // Handle should be null EXPECT_EQ(handle, nullptr); // Deinit should have been called exactly once EXPECT_EQ(deinit_count.load(), 1); // Reset global pointers g_init_count = nullptr; g_deinit_count = nullptr; } } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_dataclass.cc000066400000000000000000000117521521067262500215330ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include "../testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; // TIntObj::RegisterReflection() is called from test_reflection.cc (same binary). // --------------------------------------------------------------------------- // Any << overload // --------------------------------------------------------------------------- TEST(OstreamAny, Primitives) { { std::ostringstream os; os << Any(int64_t{42}); EXPECT_EQ(os.str(), "42"); } { std::ostringstream os; os << Any(3.14); EXPECT_EQ(os.str(), "3.14"); } { std::ostringstream os; os << Any(true); EXPECT_EQ(os.str(), "True"); } { std::ostringstream os; os << Any(nullptr); EXPECT_EQ(os.str(), "None"); } { std::ostringstream os; os << Any(String("hello")); EXPECT_EQ(os.str(), "\"hello\""); } } TEST(OstreamAny, ObjectRef) { std::ostringstream os; os << Any(TInt(5)); EXPECT_EQ(os.str(), "test.Int(value=5)"); } // --------------------------------------------------------------------------- // ObjectRef << overload // --------------------------------------------------------------------------- TEST(OstreamObjectRef, Direct) { std::ostringstream os; os << TInt(7); EXPECT_EQ(os.str(), "test.Int(value=7)"); } TEST(OstreamObjectRef, DerivedSlicing) { // Slice derived TInt into ObjectRef base — runtime type must survive TInt ti(7); const ObjectRef& base = ti; std::ostringstream os; os << base; EXPECT_EQ(os.str(), "test.Int(value=7)"); } TEST(OstreamObjectRef, NullObjectRef) { TNumber n; // default-constructed, null std::ostringstream os; os << n; EXPECT_EQ(os.str(), "None"); } TEST(OstreamObjectRef, ExistingShape) { // Shape's own operator<< must still win — otherwise repr would include type prefix Shape s({1, 2, 3}); std::ostringstream os; os << s; // Shape's operator<< writes [1, 2, 3] EXPECT_EQ(os.str(), "[1, 2, 3]"); } TEST(OstreamObjectRef, ExistingString) { // String's own operator<< writes raw content (no quotes), not repr-quoted String str("abc"); std::ostringstream os; os << str; EXPECT_EQ(os.str(), "abc"); } // --------------------------------------------------------------------------- // Variant << overload // --------------------------------------------------------------------------- TEST(OstreamVariant, IntAlternative) { Variant v(int64_t{42}); std::ostringstream os; os << v; EXPECT_EQ(os.str(), "42"); } TEST(OstreamVariant, StringAlternative) { Variant v(String("hi")); std::ostringstream os; os << v; // Goes through Any/ReprPrint so strings are repr-quoted EXPECT_EQ(os.str(), "\"hi\""); } TEST(OstreamVariant, ObjectRefAlternative) { Variant v(TInt(9)); std::ostringstream os; os << v; EXPECT_EQ(os.str(), "test.Int(value=9)"); } // --------------------------------------------------------------------------- // Optional << overload // --------------------------------------------------------------------------- TEST(OstreamOptional, IntPresent) { Optional o(int64_t{42}); std::ostringstream os; os << o; EXPECT_EQ(os.str(), "42"); } TEST(OstreamOptional, IntEmpty) { Optional o(std::nullopt); std::ostringstream os; os << o; EXPECT_EQ(os.str(), "None"); } TEST(OstreamOptional, StringPresent) { Optional o(String("x")); std::ostringstream os; os << o; // Goes through Any/ReprPrint so strings are repr-quoted EXPECT_EQ(os.str(), "\"x\""); } TEST(OstreamOptional, StringEmpty) { Optional o(std::nullopt); std::ostringstream os; os << o; EXPECT_EQ(os.str(), "None"); } TEST(OstreamOptional, ObjectRefPresent) { Optional o(TInt(3)); std::ostringstream os; os << o; EXPECT_EQ(os.str(), "test.Int(value=3)"); } TEST(OstreamOptional, ObjectRefEmpty) { Optional o; std::ostringstream os; os << o; EXPECT_EQ(os.str(), "None"); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_json_parser.cc000066400000000000000000000354141521067262500221220ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include namespace { using namespace tvm::ffi; inline bool FastMathSafeIsNaN(double x) { #ifdef __FAST_MATH__ // Bit-level NaN detection (IEEE 754 double) // IEEE 754 standard: https://en.wikipedia.org/wiki/IEEE_754 // NaN is encoded as all 1s in the exponent and non-zero in the mantissa static_assert(sizeof(double) == sizeof(uint64_t), "Unexpected double size"); uint64_t bits = *reinterpret_cast(&x); uint64_t exponent = (bits >> 52) & 0x7FF; uint64_t mantissa = bits & 0xFFFFFFFFFFFFFull; return (exponent == 0x7FF) && (mantissa != 0); #else // Safe to use std::isnan when fast-math is off return std::isnan(x); #endif } inline bool FastMathSafeIsInf(double x) { #ifdef __FAST_MATH__ // IEEE 754 standard: https://en.wikipedia.org/wiki/IEEE_754 // Inf is encoded as all 1s in the exponent and zero in the mantissa static_assert(sizeof(double) == sizeof(uint64_t), "Unexpected double size"); uint64_t bits = *reinterpret_cast(&x); uint64_t exponent = (bits >> 52) & 0x7FF; uint64_t mantissa = bits & 0xFFFFFFFFFFFFFull; // inf is encoded as all 1s in the exponent and zero in the mantissa return (exponent == 0x7FF) && (mantissa == 0); #else return std::isinf(x); #endif } TEST(JSONParser, BoolNull) { // boolean value EXPECT_EQ(json::Parse("true").cast(), true); EXPECT_EQ(json::Parse("false").cast(), false); EXPECT_EQ(json::Parse("null"), nullptr); } TEST(JSONParser, WrongBoolNull) { String error_msg; EXPECT_EQ(json::Parse("nul", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("fals", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("\n\nfx", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 3 column 1 (char 2)"); EXPECT_EQ(json::Parse("fx", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("n1", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("t1", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("f1", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); } TEST(JSONParser, Number) { // number EXPECT_EQ(json::Parse("123").cast(), 123); EXPECT_EQ(json::Parse("-124").cast(), -124); EXPECT_EQ(json::Parse("123.456").cast(), 123.456); // parsing scientific notation EXPECT_EQ(json::Parse("1.456e12").cast(), 1.456e12); // NaN EXPECT_EQ(FastMathSafeIsNaN(json::Parse("NaN").cast()), true); // Infinity EXPECT_EQ(FastMathSafeIsInf(json::Parse("Infinity").cast()), true); // -Infinity EXPECT_EQ(FastMathSafeIsInf(-json::Parse("-Infinity").cast()), true); // Test zero variants EXPECT_EQ(json::Parse("0").cast(), 0); EXPECT_EQ(json::Parse("-0").cast(), -0.0); EXPECT_EQ(json::Parse("0.0").cast(), 0.0); // Test very large numbers EXPECT_EQ(json::Parse("9223372036854775807").cast(), std::numeric_limits::max()); EXPECT_EQ(json::Parse("-9223372036854775808").cast(), std::numeric_limits::min()); // Test very small decimals EXPECT_EQ(json::Parse("1e-10").cast(), 1e-10); EXPECT_EQ(json::Parse("-1e-10").cast(), -1e-10); // Test scientific notation edge cases EXPECT_EQ(json::Parse("1E+10").cast(), 1E+10); EXPECT_EQ(json::Parse("1e+10").cast(), 1e+10); EXPECT_EQ(json::Parse("1E-10").cast(), 1E-10); EXPECT_EQ(json::Parse("123.456E+10").cast(), 123.456E+10); } TEST(JSONParser, WrongNumber) { String error_msg; EXPECT_EQ(json::Parse("123.456.789", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); // Test invalid number formats EXPECT_EQ(json::Parse("123e", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("123e+", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("123E-", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); } TEST(JSONParser, String) { EXPECT_EQ(json::Parse("\"hello\"").cast(), "hello"); EXPECT_EQ(json::Parse("\n\t \"hello\"\n\r").cast(), "hello"); EXPECT_EQ(json::Parse("\"hello\\nworld\"").cast(), "hello\nworld"); EXPECT_EQ(json::Parse("\"\"").cast(), ""); // test escape characters EXPECT_EQ(json::Parse("\"\\ta\\n\\/\\f\\\"\\\\\"").cast(), "\ta\n/\f\"\\"); // test unicode code point EXPECT_EQ(json::Parse("\"\\u0041\"").cast(), "A"); // test unicode surrogate pair EXPECT_EQ(json::Parse("\"\\uD83D\\uDE04hello\"").cast(), u8"\U0001F604hello"); } TEST(JSONParser, WrongString) { String error_msg; EXPECT_EQ(json::Parse("\"hello", &error_msg), nullptr); EXPECT_EQ(error_msg, "Unterminated string starting at: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("\"hello\x01\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid control character at: line 1 column 7 (char 6)"); EXPECT_EQ(json::Parse("\"hello\\uxx\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid \\uXXXX escape: line 1 column 8 (char 7)"); EXPECT_EQ(json::Parse("\"hello\\uDC00\\uDE04\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid surrogate pair of \\uXXXX escapes: line 1 column 8 (char 7)"); EXPECT_EQ(json::Parse("\"hello\\uD800\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid surrogate pair of \\uXXXX escapes: line 1 column 8 (char 7)"); EXPECT_EQ(json::Parse("\"hello\\uD800\\uxx\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid \\uXXXX escape: line 1 column 15 (char 14)"); EXPECT_EQ(json::Parse("\"hello\\a\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Invalid \\escape: line 1 column 8 (char 7)"); } TEST(JSONParser, Array) { EXPECT_TRUE(StructuralEqual()(json::Parse("[]"), json::Array{})); EXPECT_TRUE(StructuralEqual()(json::Parse("[1, 2,\n\t\"a\"]"), json::Array{1, 2, "a"})); } TEST(JSONParser, WrongArray) { String error_msg; EXPECT_EQ(json::Parse("]", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); EXPECT_EQ(json::Parse("[1,]", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 4 (char 3)"); EXPECT_EQ(json::Parse("[", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 2 (char 1)"); EXPECT_EQ(json::Parse("[1a", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting ',' delimiter: line 1 column 3 (char 2)"); EXPECT_EQ(json::Parse("[1,2,3", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting ',' delimiter: line 1 column 7 (char 6)"); EXPECT_EQ(json::Parse("[1] a", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 6 (char 5)"); } TEST(JSONParser, Object) { EXPECT_TRUE(StructuralEqual()(json::Parse("{}"), json::Object{})); EXPECT_TRUE(StructuralEqual()(json::Parse("{\"a\": 1, \n\"b\": \t\"c\"} "), json::Object{{"a", 1}, {"b", "c"}})); } TEST(JSONParser, ObjectOrderPreserving) { auto obj = json::Parse(R"({"c": 1, "a": 2, "b": 3} )"); json::Array keys; for (auto& [key, value] : obj.cast()) { keys.push_back(key); } EXPECT_TRUE(StructuralEqual()(keys, json::Array{"c", "a", "b"})); } TEST(JSONParser, WrongObject) { String error_msg; EXPECT_EQ(json::Parse("{\"a\":", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 6 (char 5)"); EXPECT_EQ(json::Parse("{", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting property name enclosed in double quotes: line 1 column 2 (char 1)"); // Test incomplete structures EXPECT_EQ(json::Parse("{\"incomplete\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting ':' delimiter: line 1 column 14 (char 13)"); } TEST(JSONParser, NestedObject) { EXPECT_TRUE( StructuralEqual()(json::Parse("{\"a\": \t{\"b\": 1}, \n\"c\": [1, 2, 3]}"), json::Object{{"a", json::Object{{"b", 1}}}, {"c", json::Array{1, 2, 3}}})); EXPECT_TRUE(StructuralEqual()( json::Parse("{\"a\": \t{\"b\": 1}, \n\"c\": [1, null, Infinity]}"), json::Object{{"a", json::Object{{"b", 1}}}, {"c", json::Array{1, nullptr, std::numeric_limits::infinity()}}})); EXPECT_TRUE(StructuralEqual()( json::Parse("[{}, {\"a\": [1.1, 1000000]}]"), json::Array{json::Object{}, json::Object{{"a", json::Array{1.1, 1000000}}}})); } TEST(JSONParser, WrongNestedObject) { String error_msg; EXPECT_EQ(json::Parse("{\"a\":\n\n[1]", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting ',' delimiter: line 3 column 4 (char 10)"); EXPECT_EQ(json::Parse("{\"a\":\n\n[abc]}", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 3 column 2 (char 8)"); } // edge cases TEST(JSONParser, WhitespaceHandling) { // Test various whitespace characters EXPECT_EQ(json::Parse(" \t\n\r true \t\n\r ").cast(), true); EXPECT_EQ(json::Parse("\n\n\n123\n\n\n").cast(), 123); EXPECT_EQ(json::Parse(" \"hello world\" ").cast(), "hello world"); // Test whitespace in arrays and objects EXPECT_TRUE(StructuralEqual()(json::Parse(" [ 1 , 2 , 3 ] "), json::Array{1, 2, 3})); EXPECT_TRUE(StructuralEqual()(json::Parse(" { \"a\" : 1 , \"b\" : 2 } "), json::Object{{"a", 1}, {"b", 2}})); } TEST(JSONParser, WrongEmptyAndMinimalInputs) { String error_msg; // Test empty string EXPECT_EQ(json::Parse("", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 1 column 1 (char 0)"); // Test only whitespace EXPECT_EQ(json::Parse(" \t\n ", &error_msg), nullptr); EXPECT_EQ(error_msg, "Expecting value: line 2 column 5 (char 9)"); } TEST(JSONParser, UnicodeEdgeCases) { // Test various unicode characters EXPECT_EQ(json::Parse("\"\\u0000\"").cast(), std::string("\0", 1)); // replace using \U to avoid encoding issues EXPECT_EQ(json::Parse("\"\\u00FF\"").cast(), u8"\U000000FF"); EXPECT_EQ(json::Parse("\"\\u4E2D\\u6587\"").cast(), u8"\U00004E2D\U00006587"); // Test multiple surrogate pairs EXPECT_EQ(json::Parse("\"\\uD83D\\uDE00\\uD83D\\uDE01\"").cast(), u8"\U0001F600\U0001F601"); } TEST(JSONParser, RawUTF8Bytes) { // Regression test: raw UTF-8 bytes (>= 0x80) must not be rejected as control characters. // This failed when the parser used signed char comparison: *cur_ < ' ' EXPECT_EQ(json::Parse("\"\xE4\xB8\xAD\xE6\x96\x87\"").cast(), "\xE4\xB8\xAD\xE6\x96\x87"); } TEST(JSONParser, LargeInputs) { // Test large array std::string large_array = "["; for (int i = 0; i < 1000; ++i) { if (i > 0) large_array += ","; large_array += std::to_string(i); } large_array += "]"; auto result = json::Parse(large_array); EXPECT_TRUE(result != nullptr); EXPECT_EQ(result.cast().size(), 1000); // Test large object std::string large_object = "{"; for (int i = 0; i < 500; ++i) { if (i > 0) large_object += ","; large_object += "\"key" + std::to_string(i) + "\":" + std::to_string(i); } large_object += "}"; result = json::Parse(large_object); EXPECT_TRUE(result != nullptr); EXPECT_EQ(result.cast().size(), 500); } TEST(JSONParser, MixedDataTypes) { // Test complex nested structure with all data types std::string complex_json = R"({ "null_value": null, "boolean_true": true, "boolean_false": false, "integer": 42, "negative_integer": -42, "float": 3.14159, "scientific": 1.23e-4, "string": "hello world", "unicode_string": "Hello \u4e16\u754c \ud83c\udf0d", "empty_string": "", "empty_array": [], "empty_object": {}, "number_array": [1, 2, 3, 4, 5], "mixed_array": [1, "two", true, null, 3.14], "nested_object": { "level1": { "level2": { "data": [1, 2, {"nested_array": [true, false]}] } } } })"; auto result = json::Parse(complex_json); // Create expected structure for comparison json::Object expected{ {"null_value", nullptr}, {"boolean_true", true}, {"boolean_false", false}, {"integer", 42}, {"negative_integer", -42}, {"float", 3.14159}, {"scientific", 1.23e-4}, {"string", "hello world"}, {"unicode_string", u8"Hello \U00004E16\U0000754C \U0001F30D"}, {"empty_string", ""}, {"empty_array", json::Array{}}, {"empty_object", json::Object{}}, {"number_array", json::Array{1, 2, 3, 4, 5}}, {"mixed_array", json::Array{1, "two", true, nullptr, 3.14}}, {"nested_object", json::Object{ {"level1", json::Object{ {"level2", json::Object{ {"data", json::Array{1, 2, json::Object{{"nested_array", json::Array{true, false}}}}}}}}}}}}; EXPECT_TRUE(StructuralEqual()(result, expected)); } TEST(JSONParser, WrongExtraData) { String error_msg; EXPECT_EQ(json::Parse("truee", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 5 (char 4)"); EXPECT_EQ(json::Parse("true false", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 6 (char 5)"); EXPECT_EQ(json::Parse("123 456", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 5 (char 4)"); EXPECT_EQ(json::Parse("\"hello\" \"world\"", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 9 (char 8)"); EXPECT_EQ(json::Parse("{} []", &error_msg), nullptr); EXPECT_EQ(error_msg, "Extra data: line 1 column 4 (char 3)"); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_json_writer.cc000066400000000000000000000174451521067262500221460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include namespace { using namespace tvm::ffi; TEST(JSONWriter, BoolNull) { // boolean value EXPECT_EQ(json::Stringify(json::Value(true)), "true"); EXPECT_EQ(json::Stringify(json::Value(false)), "false"); EXPECT_EQ(json::Stringify(json::Value(nullptr)), "null"); } TEST(JSONWriter, Integer) { // positive integer EXPECT_EQ(json::Stringify(json::Value(42)), "42"); // negative integer EXPECT_EQ(json::Stringify(json::Value(-123)), "-123"); // zero EXPECT_EQ(json::Stringify(json::Value(0)), "0"); // large positive integer EXPECT_EQ(json::Stringify(json::Value(std::numeric_limits::max())), "9223372036854775807"); // large negative integer EXPECT_EQ(json::Stringify(json::Value(std::numeric_limits::min())), "-9223372036854775808"); } TEST(JSONWriter, Float) { // regular float EXPECT_EQ(json::Stringify(json::Value(2.5)), "2.5"); // integer-like float (should have .0 suffix) EXPECT_EQ(json::Stringify(json::Value(5.0)), "5.0"); EXPECT_EQ(json::Stringify(json::Value(-10.0)), "-10.0"); // zero float EXPECT_EQ(json::Stringify(json::Value(0.0)), "0.0"); // scientific notation for very small numbers EXPECT_EQ(json::Stringify(json::Value(-7.89e-15)), "-7.89e-15"); // short scientific notation (shorter than fixed-point) EXPECT_EQ(json::Stringify(json::Value(2e-8)), "2e-08"); // NaN EXPECT_EQ(json::Stringify(json::Value(std::numeric_limits::quiet_NaN())), "NaN"); // positive infinity EXPECT_EQ(json::Stringify(json::Value(std::numeric_limits::infinity())), "Infinity"); // negative infinity EXPECT_EQ(json::Stringify(json::Value(-std::numeric_limits::infinity())), "-Infinity"); } TEST(JSONWriter, String) { // simple string EXPECT_EQ(json::Stringify(json::Value(String("hello"))), "\"hello\""); // empty string EXPECT_EQ(json::Stringify(json::Value(String(""))), "\"\""); // string with escaped characters EXPECT_EQ(json::Stringify(json::Value(String("\"quoted\""))), "\"\\\"quoted\\\"\""); EXPECT_EQ(json::Stringify(json::Value(String("backslash\\"))), "\"backslash\\\\\""); EXPECT_EQ(json::Stringify(json::Value(String("forward/slash"))), "\"forward\\/slash\""); EXPECT_EQ(json::Stringify(json::Value(String("line\nbreak"))), "\"line\\nbreak\""); EXPECT_EQ(json::Stringify(json::Value(String("tab\there"))), "\"tab\\there\""); EXPECT_EQ(json::Stringify(json::Value(String("carriage\rreturn"))), "\"carriage\\rreturn\""); // string with control character EXPECT_EQ(json::Stringify(json::Value(String(std::string("\x01", 1) + "control"))), "\"\\u0001control\""); } TEST(JSONWriter, Array) { // empty array json::Array empty_array; EXPECT_EQ(json::Stringify(empty_array), "[]"); // single element array json::Array single_array{42}; EXPECT_EQ(json::Stringify(single_array), "[42]"); // multiple elements array json::Array multi_array{1, "hello", true}; EXPECT_EQ(json::Stringify(multi_array), "[1,\"hello\",true]"); // nested array json::Array nested_array{json::Array{1, 2}, 3}; EXPECT_EQ(json::Stringify(nested_array), "[[1,2],3]"); } TEST(JSONWriter, Object) { // empty object json::Object empty_object; EXPECT_EQ(json::Stringify(empty_object), "{}"); // single key-value pair json::Object single_object{{String("key"), String("value")}}; EXPECT_EQ(json::Stringify(single_object), "{\"key\":\"value\"}"); // multiple key-value pairs - insertion order preservation json::Object multi_object{{"name", "Alice"}, {"age", 30}, {"active", true}, {"score", 95.5}}; EXPECT_EQ(json::Stringify(multi_object), "{\"name\":\"Alice\",\"age\":30,\"active\":true,\"score\":95.5}"); } TEST(JSONWriter, InsertionOrderPreservation) { // test that objects preserve insertion order json::Object ordered_object{ {"zebra", "last"}, {"alpha", "first"}, {"beta", "middle"}, {"gamma", 123}, {"delta", true}}; EXPECT_EQ( json::Stringify(ordered_object), "{\"zebra\":\"last\",\"alpha\":\"first\",\"beta\":\"middle\",\"gamma\":123,\"delta\":true}"); // test with indentation to verify order is preserved std::string ordered_indented = json::Stringify(ordered_object, 2); EXPECT_EQ(ordered_indented, String(R"({ "zebra": "last", "alpha": "first", "beta": "middle", "gamma": 123, "delta": true })")); // test nested objects also preserve order json::Object nested_ordered{ {"outer1", json::Object{{"inner_z", "z_value"}, {"inner_a", "a_value"}, {"inner_m", "m_value"}}}, {"outer2", json::Object{{"third", 3}, {"first", 1}, {"second", 2}}}}; std::string nested_ordered_indented = json::Stringify(nested_ordered, 2); EXPECT_EQ(nested_ordered_indented, String(R"({ "outer1": { "inner_z": "z_value", "inner_a": "a_value", "inner_m": "m_value" }, "outer2": { "third": 3, "first": 1, "second": 2 } })")); } TEST(JSONWriter, NestedStructures) { // object containing array json::Object obj_with_array{{String("numbers"), json::Array{1, 2, 3}}}; EXPECT_EQ(json::Stringify(obj_with_array), "{\"numbers\":[1,2,3]}"); // array containing object json::Array arr_with_obj{json::Object{{String("key"), String("value")}}}; EXPECT_EQ(json::Stringify(arr_with_obj), "[{\"key\":\"value\"}]"); // deeply nested structure json::Object nested_obj{ {String("nested"), json::Array{json::Object{{String("deep"), String("value")}}}}}; EXPECT_EQ(json::Stringify(nested_obj), "{\"nested\":[{\"deep\":\"value\"}]}"); } TEST(JSONWriter, Indentation) { // test with indentation json::Array arr{1, 2}; std::string indented = json::Stringify(arr, 2); EXPECT_EQ(indented, String(R"([ 1, 2 ])")); // object with indentation json::Object obj{{"key", "value"}}; std::string indented_obj = json::Stringify(obj, 2); EXPECT_EQ(indented_obj, String(R"({ "key": "value" })")); // complex nested structure with multiple data types // keep double as .5 so output is deterministic as they exactly rounds to power of 2 json::Object complex_nested{ {"name", "test"}, {"count", 42}, {"price", 3.5}, {"active", true}, {"metadata", nullptr}, {"numbers", json::Array{1, 2, 3}}, {"config", json::Object{{"enabled", false}, {"timeout", 30.5}, {"tags", json::Array{"production", "critical", nullptr}}}}, {"matrix", json::Array{json::Array{1, 2}, json::Array{3.5, 4.5}, json::Array{"a", "b"}}}}; std::string complex_indented = json::Stringify(complex_nested, 2); EXPECT_EQ(complex_indented, String(R"({ "name": "test", "count": 42, "price": 3.5, "active": true, "metadata": null, "numbers": [ 1, 2, 3 ], "config": { "enabled": false, "timeout": 30.5, "tags": [ "production", "critical", null ] }, "matrix": [ [ 1, 2 ], [ 3.5, 4.5 ], [ "a", "b" ] ] })")); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_serialization.cc000066400000000000000000001030721521067262500224460ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include #include #include "../testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Serialization, BoolNull) { json::Object expected_null = json::Object{{"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "None"}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(nullptr), expected_null)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_null), nullptr)); json::Object expected_true = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "bool"}, {"data", true}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(true), expected_true)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_true), true)); json::Object expected_false = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "bool"}, {"data", false}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(false), expected_false)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_false), false)); } TEST(Serialization, IntegerTypes) { // Test positive integer json::Object expected_int = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "int"}, {"data", 42}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(static_cast(42)), expected_int)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_int), static_cast(42))); } TEST(Serialization, FloatTypes) { // Test positive float json::Object expected_float = json::Object{{"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "float"}, {"data", 3.14159}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(3.14159), expected_float)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_float), 3.14159)); } TEST(Serialization, StringTypes) { // Test short string json::Object expected_short = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", String("hello")}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(String("hello")), expected_short)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_short), String("hello"))); // Test long string std::string long_str(1000, 'x'); json::Object expected_long = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", String(long_str)}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(String(long_str)), expected_long)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_long), String(long_str))); // Test string with special characters json::Object expected_special = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", String("hello\nworld\t\"quotes\"")}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(String("hello\nworld\t\"quotes\"")), expected_special)); EXPECT_TRUE( StructuralEqual()(FromJSONGraph(expected_special), String("hello\nworld\t\"quotes\""))); } TEST(Serialization, Bytes) { // Test empty bytes Bytes empty_bytes; json::Object expected_empty = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Bytes"}, {"data", ""}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_bytes), expected_empty)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty), empty_bytes)); // Test bytes with that encoded as base64 Bytes bytes_content = Bytes("abcd"); json::Object expected_encoded = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Bytes"}, {"data", "YWJjZA=="}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(bytes_content), expected_encoded)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_encoded), bytes_content)); // Test bytes with that encoded as base64, that contains control characters via utf-8 char bytes_v2_content[] = {0x01, 0x02, 0x03, 0x04, 0x01, 0x0b}; Bytes bytes_v2 = Bytes(bytes_v2_content, sizeof(bytes_v2_content)); json::Object expected_encoded_v2 = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Bytes"}, {"data", "AQIDBAEL"}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(bytes_v2), expected_encoded_v2)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_encoded_v2), bytes_v2)); } TEST(Serialization, DataTypes) { // Test int32 dtype DLDataType int32_dtype; int32_dtype.code = kDLInt; int32_dtype.bits = 32; int32_dtype.lanes = 1; json::Object expected_int32 = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "DataType"}, {"data", String("int32")}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(int32_dtype), expected_int32)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_int32), int32_dtype)); // Test float64 dtype DLDataType float64_dtype; float64_dtype.code = kDLFloat; float64_dtype.bits = 64; float64_dtype.lanes = 1; json::Object expected_float64 = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "DataType"}, {"data", String("float64")}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(float64_dtype), expected_float64)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_float64), float64_dtype)); // Test vector dtype DLDataType vector_dtype; vector_dtype.code = kDLFloat; vector_dtype.bits = 32; vector_dtype.lanes = 4; json::Object expected_vector = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "DataType"}, {"data", String("float32x4")}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(vector_dtype), expected_vector)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_vector), vector_dtype)); } TEST(Serialization, DeviceTypes) { // Test CPU device DLDevice cpu_device; cpu_device.device_type = kDLCPU; cpu_device.device_id = 0; json::Object expected_cpu = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "Device"}, {"data", json::Array{static_cast(kDLCPU), static_cast(0)}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(cpu_device), expected_cpu)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_cpu), cpu_device)); // Test GPU device DLDevice gpu_device; gpu_device.device_type = kDLCUDA; gpu_device.device_id = 1; json::Object expected_gpu = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{ {"type", "Device"}, {"data", json::Array{static_cast(kDLCUDA), 1}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(gpu_device), expected_gpu)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_gpu), gpu_device)); } TEST(Serialization, Arrays) { // Test empty array Array empty_array; json::Object expected_empty = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Array"}, {"data", json::Array{}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_array), expected_empty)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty), empty_array)); // Test single element array Array single_array; single_array.push_back(Any(42)); json::Object expected_single = json::Object{{"root_index", 1}, {"nodes", json::Array{ json::Object{{"type", "int"}, {"data", static_cast(42)}}, json::Object{{"type", "ffi.Array"}, {"data", json::Array{0}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(single_array), expected_single)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_single), single_array)); // Test duplicated element array Array duplicated_array; duplicated_array.push_back(42); duplicated_array.push_back(42); json::Object expected_duplicated = json::Object{{"root_index", 1}, {"nodes", json::Array{ json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.Array"}, {"data", json::Array{0, 0}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(duplicated_array), expected_duplicated)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_duplicated), duplicated_array)); // Test mixed element array, note that 42 and "hello" are duplicated and will // be indexed as 0 and 1 Array mixed_array; mixed_array.push_back(42); mixed_array.push_back(String("hello")); mixed_array.push_back(true); mixed_array.push_back(nullptr); mixed_array.push_back(42); mixed_array.push_back(String("hello")); json::Object expected_mixed = json::Object{ {"root_index", 4}, {"nodes", json::Array{ json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.String"}, {"data", String("hello")}}, json::Object{{"type", "bool"}, {"data", true}}, json::Object{{"type", "None"}}, json::Object{{"type", "ffi.Array"}, {"data", json::Array{0, 1, 2, 3, 0, 1}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(mixed_array), expected_mixed)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_mixed), mixed_array)); } TEST(Serialization, Maps) { // Test empty map Map empty_map; json::Object expected_empty = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Map"}, {"data", json::Array{}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_map), expected_empty)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty), empty_map)); // Test single element map Map single_map{{"key", 42}}; json::Object expected_single = json::Object{ {"root_index", 2}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", String("key")}}, json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.Map"}, {"data", json::Array{0, 1}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(single_map), expected_single)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_single), single_map)); // Test duplicated element map Map duplicated_map{{"b", 42}, {"a", 42}}; json::Object expected_duplicated = json::Object{ {"root_index", 3}, {"nodes", json::Array{ json::Object{{"type", "ffi.String"}, {"data", "b"}}, json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.String"}, {"data", "a"}}, json::Object{{"type", "ffi.Map"}, {"data", json::Array{0, 1, 2, 1}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(duplicated_map), expected_duplicated)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_duplicated), duplicated_map)); } TEST(Serialization, Dicts) { // Test empty dict Dict empty_dict; json::Object expected_empty = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Dict"}, {"data", json::Array{}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_dict), expected_empty)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty), empty_dict)); // Test single element dict Dict single_dict{{"key", 42}}; json::Object expected_single = json::Object{ {"root_index", 2}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", String("key")}}, json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.Dict"}, {"data", json::Array{0, 1}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(single_dict), expected_single)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_single), single_dict)); // Test duplicated element dict Dict duplicated_dict{{"b", 42}, {"a", 42}}; json::Object expected_duplicated = json::Object{ {"root_index", 3}, {"nodes", json::Array{ json::Object{{"type", "ffi.String"}, {"data", "b"}}, json::Object{{"type", "int"}, {"data", 42}}, json::Object{{"type", "ffi.String"}, {"data", "a"}}, json::Object{{"type", "ffi.Dict"}, {"data", json::Array{0, 1, 2, 1}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(duplicated_dict), expected_duplicated)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_duplicated), duplicated_dict)); } TEST(Serialization, DictWithIntKeys) { Dict dict; dict.Set(static_cast(1), String("one")); dict.Set(static_cast(2), String("two")); json::Value serialized = ToJSONGraph(dict); Any deserialized = FromJSONGraph(serialized); Dict result = deserialized.cast>(); EXPECT_EQ(result.size(), 2); EXPECT_EQ(std::string(result[1].cast()), "one"); EXPECT_EQ(std::string(result[2].cast()), "two"); } TEST(Serialization, DictWithArrayValues) { Array arr; arr.push_back(10); arr.push_back(20); Dict dict{{"nums", arr}}; json::Value serialized = ToJSONGraph(dict); Any deserialized = FromJSONGraph(serialized); Dict result = deserialized.cast>(); Array result_arr = result["nums"].cast>(); EXPECT_EQ(result_arr.size(), 2); EXPECT_EQ(result_arr[0].cast(), 10); EXPECT_EQ(result_arr[1].cast(), 20); } TEST(Serialization, DictOfObjects) { TVar x("x"); Dict dict{{"var", x}}; json::Value serialized = ToJSONGraph(dict); Any deserialized = FromJSONGraph(serialized); Dict result = deserialized.cast>(); EXPECT_EQ(std::string(result["var"].cast()->name), "x"); } TEST(Serialization, Shapes) { Shape empty_shape; json::Object expected_empty_shape = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Shape"}, {"data", json::Array{}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_shape), expected_empty_shape)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty_shape), empty_shape)); Shape shape({1, 2, 3}); json::Object expected_shape = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.Shape"}, {"data", json::Array{1, 2, 3}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(shape), expected_shape)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_shape), shape)); } TEST(Serialization, TestObjectVar) { TVar x = TVar("x"); json::Object expected_x = json::Object{ {"root_index", 1}, {"nodes", json::Array{json::Object{{"type", "ffi.String"}, {"data", "x"}}, json::Object{{"type", "test.Var"}, {"data", json::Object{{"name", 0}}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(x), expected_x)); EXPECT_TRUE(StructuralEqual::Equal(FromJSONGraph(expected_x), x, /*map_free_vars=*/true)); } TEST(Serialization, TestObjectIntCustomToJSON) { TInt value = TInt(42); json::Object expected_i = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "test.Int"}, {"data", json::Object{{"value", 42}}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(value), expected_i)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_i), value)); } TEST(Serialization, TestObjectFunc) { TVar x = TVar("x"); // comment fields are ignored TFunc fa = TFunc({x}, {x, x}, String("comment a")); json::Object expected_fa = json::Object{ {"root_index", 5}, {"nodes", json::Array{ json::Object{{"type", "ffi.String"}, {"data", "x"}}, // string "x" json::Object{{"type", "test.Var"}, {"data", json::Object{{"name", 0}}}}, // var x json::Object{{"type", "ffi.Array"}, {"data", json::Array{1}}}, // array [x] json::Object{{"type", "ffi.Array"}, {"data", json::Array{1, 1}}}, // array [x, x] json::Object{{"type", "ffi.String"}, {"data", "comment a"}}, // "comment a" json::Object{{"type", "test.Func"}, {"data", json::Object{{"params", 2}, {"body", 3}, {"comment", 4}}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(fa), expected_fa)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_fa), fa)); TFunc fb = TFunc({}, {}, std::nullopt); json::Object expected_fb = json::Object{ {"root_index", 3}, {"nodes", json::Array{ json::Object{{"type", "ffi.Array"}, {"data", json::Array{}}}, json::Object{{"type", "ffi.Array"}, {"data", json::Array{}}}, json::Object{{"type", "None"}}, json::Object{{"type", "test.Func"}, {"data", json::Object{{"params", 0}, {"body", 1}, {"comment", 2}}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(fb), expected_fb)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_fb), fb)); } TEST(Serialization, AttachMetadata) { bool value = true; json::Object metadata{{"version", "1.0"}}; json::Object expected = json::Object{{"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "bool"}, {"data", true}}}}, {"metadata", metadata}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(value, metadata), expected)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected), value)); } TEST(Serialization, ListBasic) { // Test empty list List empty_list; json::Object expected_empty = json::Object{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "ffi.List"}, {"data", json::Array{}}}}}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(empty_list), expected_empty)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_empty), empty_list)); // Test single element list List single_list; single_list.push_back(Any(42)); json::Object expected_single = json::Object{{"root_index", 1}, {"nodes", json::Array{ json::Object{{"type", "int"}, {"data", static_cast(42)}}, json::Object{{"type", "ffi.List"}, {"data", json::Array{0}}}, }}}; EXPECT_TRUE(StructuralEqual()(ToJSONGraph(single_list), expected_single)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_single), single_list)); } TEST(Serialization, ListRoundTrip) { // Test roundtrip for nested list List nested; nested.push_back(1); nested.push_back(String("hello")); nested.push_back(true); json::Value serialized = ToJSONGraph(nested); Any deserialized = FromJSONGraph(serialized); EXPECT_TRUE(StructuralEqual()(deserialized, nested)); } TEST(Serialization, DISABLED_ListCycleDetection) { List lst; lst.push_back(42); lst.push_back(lst); // creates a cycle via shared mutable reference EXPECT_ANY_THROW(ToJSONGraph(lst)); } TEST(Serialization, ShuffleNodeOrder) { // the FromJSONGraph is agnostic to the node order // so we can shuffle the node order as it reads nodes lazily Map duplicated_map{{"b", 42}, {"a", 42}}; json::Object expected_shuffled = json::Object{ {"root_index", 0}, {"nodes", json::Array{ json::Object{{"type", "ffi.Map"}, {"data", json::Array{2, 3, 1, 3}}}, json::Object{{"type", "ffi.String"}, {"data", "a"}}, json::Object{{"type", "ffi.String"}, {"data", "b"}}, json::Object{{"type", "int"}, {"data", 42}}, }}}; EXPECT_TRUE(StructuralEqual()(FromJSONGraph(expected_shuffled), duplicated_map)); } // --------------------------------------------------------------------------- // Integer edge cases // --------------------------------------------------------------------------- TEST(Serialization, IntegerEdgeCases) { // zero EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(static_cast(0))), static_cast(0))); // negative EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(static_cast(-1))), static_cast(-1))); // large positive int64_t large = 1000000000000LL; EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(large)), large)); // large negative int64_t large_neg = -999999999999LL; EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(large_neg)), large_neg)); // INT64_MIN and INT64_MAX int64_t imin = std::numeric_limits::min(); int64_t imax = std::numeric_limits::max(); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(imin)), imin)); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(imax)), imax)); } // --------------------------------------------------------------------------- // Float edge cases // --------------------------------------------------------------------------- TEST(Serialization, FloatEdgeCases) { // zero EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(0.0)), 0.0)); // negative EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(-1.5)), -1.5)); // very large EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(1e300)), 1e300)); // very small EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(1e-300)), 1e-300)); } // --------------------------------------------------------------------------- // String edge cases // --------------------------------------------------------------------------- TEST(Serialization, EmptyString) { String empty(""); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(empty)), empty)); } TEST(Serialization, UnicodeString) { String unicode("hello 世界 🌍"); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(unicode)), unicode)); } TEST(Serialization, NullCharInString) { // String with embedded null characters std::string with_null("ab\0cd", 5); String s(with_null); EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(s)), s)); } // --------------------------------------------------------------------------- // Object with all POD field types (exercises node-graph path for POD fields) // --------------------------------------------------------------------------- TEST(Serialization, AllFieldsObject) { DLDataType dtype; dtype.code = kDLFloat; dtype.bits = 32; dtype.lanes = 1; DLDevice device; device.device_type = kDLCUDA; device.device_id = 3; Array arr; arr.push_back(1); arr.push_back(String("two")); Map map{{"k", 99}}; TAllFields obj(true, -7, 2.5, dtype, device, String("hello"), String("opt"), arr, map); json::Value serialized = ToJSONGraph(obj); Any deserialized = FromJSONGraph(serialized); // verify each field TAllFields result = deserialized.cast(); EXPECT_EQ(result->v_bool, true); EXPECT_EQ(result->v_int, -7); EXPECT_DOUBLE_EQ(result->v_float, 2.5); EXPECT_EQ(result->v_dtype.code, kDLFloat); EXPECT_EQ(result->v_dtype.bits, 32); EXPECT_EQ(result->v_dtype.lanes, 1); EXPECT_EQ(result->v_device.device_type, kDLCUDA); EXPECT_EQ(result->v_device.device_id, 3); EXPECT_EQ(std::string(result->v_str), "hello"); EXPECT_TRUE(result->v_opt_str.has_value()); EXPECT_EQ(std::string(result->v_opt_str.value()), "opt"); EXPECT_EQ(result->v_array.size(), 2); EXPECT_EQ(result->v_map.size(), 1); } TEST(Serialization, AllFieldsObjectOptionalNone) { DLDataType dtype; dtype.code = kDLInt; dtype.bits = 64; dtype.lanes = 1; DLDevice device; device.device_type = kDLCPU; device.device_id = 0; TAllFields obj(false, 0, 0.0, dtype, device, String(""), std::nullopt, Array(), Map()); json::Value serialized = ToJSONGraph(obj); Any deserialized = FromJSONGraph(serialized); TAllFields result = deserialized.cast(); EXPECT_EQ(result->v_bool, false); EXPECT_EQ(result->v_int, 0); EXPECT_DOUBLE_EQ(result->v_float, 0.0); EXPECT_EQ(std::string(result->v_str), ""); EXPECT_FALSE(result->v_opt_str.has_value()); EXPECT_EQ(result->v_array.size(), 0); EXPECT_EQ(result->v_map.size(), 0); } // --------------------------------------------------------------------------- // Default field values during deserialization // --------------------------------------------------------------------------- TEST(Serialization, DefaultFieldValues) { // serialize a TWithDefaults, then deserialize from JSON with missing default fields TWithDefaults original(100, 42, "default", true); json::Value serialized = ToJSONGraph(original); // roundtrip should work Any deserialized = FromJSONGraph(serialized); TWithDefaults result = deserialized.cast(); EXPECT_EQ(result->required_val, 100); EXPECT_EQ(result->default_int, 42); EXPECT_EQ(std::string(result->default_str), "default"); EXPECT_EQ(result->default_bool, true); } TEST(Serialization, DefaultFieldValuesMissing) { // manually construct JSON with only required field, defaults should kick in // required_val is int64_t so it is inlined directly (POD field) json::Object data; data.Set("required_val", static_cast(999)); json::Object graph{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "test.WithDefaults"}, {"data", data}}}}}; Any result = FromJSONGraph(graph); TWithDefaults obj = result.cast(); EXPECT_EQ(obj->required_val, 999); EXPECT_EQ(obj->default_int, 42); EXPECT_EQ(std::string(obj->default_str), "default"); EXPECT_EQ(obj->default_bool, true); } // --------------------------------------------------------------------------- // Shared object references // --------------------------------------------------------------------------- TEST(Serialization, SharedObjectReferences) { TVar shared_var("shared"); // two funcs share the same var TFunc f1({shared_var}, {shared_var, shared_var}, std::nullopt); json::Value serialized = ToJSONGraph(f1); Any deserialized = FromJSONGraph(serialized); TFunc result = deserialized.cast(); // all references to "shared" should be the same object after deserialization // via the node dedup mechanism EXPECT_EQ(result->params.size(), 1); EXPECT_EQ(result->body.size(), 2); // the params[0] and body[0] and body[1] should all refer to the same object EXPECT_EQ(result->params[0].get(), result->body[0].get()); EXPECT_EQ(result->body[0].get(), result->body[1].get()); } // --------------------------------------------------------------------------- // Nested objects // --------------------------------------------------------------------------- TEST(Serialization, NestedObjects) { TVar x("x"); TVar y("y"); TFunc inner({x}, {x}, String("inner")); // put the inner func as a body element of the outer func TFunc outer({y}, {inner}, String("outer")); json::Value serialized = ToJSONGraph(outer); Any deserialized = FromJSONGraph(serialized); TFunc result = deserialized.cast(); EXPECT_EQ(result->comment.value(), "outer"); TFunc inner_result = Any(result->body[0]).cast(); EXPECT_EQ(inner_result->comment.value(), "inner"); EXPECT_EQ(std::string(Any(inner_result->params[0]).cast()->name), "x"); } // --------------------------------------------------------------------------- // Map with integer keys // --------------------------------------------------------------------------- TEST(Serialization, MapWithIntKeys) { Map map; map.Set(static_cast(1), String("one")); map.Set(static_cast(2), String("two")); json::Value serialized = ToJSONGraph(map); Any deserialized = FromJSONGraph(serialized); Map result = deserialized.cast>(); EXPECT_EQ(result.size(), 2); EXPECT_EQ(std::string(result[1].cast()), "one"); EXPECT_EQ(std::string(result[2].cast()), "two"); } // --------------------------------------------------------------------------- // Nested containers // --------------------------------------------------------------------------- TEST(Serialization, NestedArrays) { Array inner1; inner1.push_back(1); inner1.push_back(2); Array inner2; inner2.push_back(3); Array outer; outer.push_back(inner1); outer.push_back(inner2); json::Value serialized = ToJSONGraph(outer); Any deserialized = FromJSONGraph(serialized); Array result = deserialized.cast>(); EXPECT_EQ(result.size(), 2); Array r1 = result[0].cast>(); Array r2 = result[1].cast>(); EXPECT_EQ(r1.size(), 2); EXPECT_EQ(r1[0].cast(), 1); EXPECT_EQ(r1[1].cast(), 2); EXPECT_EQ(r2.size(), 1); EXPECT_EQ(r2[0].cast(), 3); } TEST(Serialization, MapWithArrayValues) { Array arr; arr.push_back(10); arr.push_back(20); Map map{{"nums", arr}}; json::Value serialized = ToJSONGraph(map); Any deserialized = FromJSONGraph(serialized); Map result = deserialized.cast>(); Array result_arr = result["nums"].cast>(); EXPECT_EQ(result_arr.size(), 2); EXPECT_EQ(result_arr[0].cast(), 10); EXPECT_EQ(result_arr[1].cast(), 20); } // --------------------------------------------------------------------------- // Array and Map with objects // --------------------------------------------------------------------------- TEST(Serialization, ArrayOfObjects) { TVar x("x"); TVar y("y"); Array arr; arr.push_back(x); arr.push_back(y); json::Value serialized = ToJSONGraph(arr); Any deserialized = FromJSONGraph(serialized); Array result = deserialized.cast>(); EXPECT_EQ(result.size(), 2); EXPECT_EQ(std::string(result[0].cast()->name), "x"); EXPECT_EQ(std::string(result[1].cast()->name), "y"); } TEST(Serialization, MapOfObjects) { TVar x("x"); Map map{{"var", x}}; json::Value serialized = ToJSONGraph(map); Any deserialized = FromJSONGraph(serialized); Map result = deserialized.cast>(); EXPECT_EQ(std::string(result["var"].cast()->name), "x"); } // --------------------------------------------------------------------------- // Mixed-type array (exercises runtime type dispatch for each element) // --------------------------------------------------------------------------- TEST(Serialization, MixedTypeArrayRoundTrip) { DLDataType dtype; dtype.code = kDLInt; dtype.bits = 32; dtype.lanes = 1; DLDevice device; device.device_type = kDLCPU; device.device_id = 0; Array arr; arr.push_back(nullptr); arr.push_back(true); arr.push_back(false); arr.push_back(static_cast(42)); arr.push_back(3.14); arr.push_back(String("hello")); arr.push_back(dtype); arr.push_back(device); // roundtrip and verify structural equality EXPECT_TRUE(StructuralEqual()(FromJSONGraph(ToJSONGraph(arr)), arr)); } // --------------------------------------------------------------------------- // Error cases // --------------------------------------------------------------------------- TEST(Serialization, ErrorMissingRequiredField) { // required_val is required but not provided json::Object data; json::Object graph{ {"root_index", 0}, {"nodes", json::Array{json::Object{{"type", "test.WithDefaults"}, {"data", data}}}}}; EXPECT_ANY_THROW(FromJSONGraph(graph)); } TEST(Serialization, ErrorInvalidRootStructure) { // not an object EXPECT_ANY_THROW(FromJSONGraph(json::Value(42))); } TEST(Serialization, ErrorMissingRootIndex) { json::Object graph{{"nodes", json::Array{json::Object{{"type", "None"}}}}}; EXPECT_ANY_THROW(FromJSONGraph(graph)); } TEST(Serialization, ErrorMissingNodes) { json::Object graph{{"root_index", 0}}; EXPECT_ANY_THROW(FromJSONGraph(graph)); } // --------------------------------------------------------------------------- // String serialization roundtrip (json::Stringify / json::Parse) // --------------------------------------------------------------------------- TEST(Serialization, StringRoundTrip) { TVar x("x"); TFunc f({x}, {x}, String("comment")); String json_str = json::Stringify(ToJSONGraph(f)); Any deserialized = FromJSONGraph(json::Parse(json_str)); EXPECT_TRUE(StructuralEqual::Equal(deserialized, f, /*map_free_vars=*/true)); } TEST(Serialization, StringRoundTripPrimitives) { auto rt = [](const Any& v) { return FromJSONGraph(json::Parse(json::Stringify(ToJSONGraph(v)))); }; // int EXPECT_TRUE(StructuralEqual()(rt(static_cast(123)), 123)); // bool EXPECT_TRUE(StructuralEqual()(rt(true), true)); // float EXPECT_TRUE(StructuralEqual()(rt(2.718), 2.718)); // string EXPECT_TRUE(StructuralEqual()(rt(String("test")), String("test"))); // null EXPECT_TRUE(StructuralEqual()(rt(nullptr), nullptr)); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_structural_equal_hash.cc000066400000000000000000000342551521067262500242010ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include #include "../testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; namespace refl = tvm::ffi::reflection; TEST(StructuralEqualHash, Array) { Array a = {1, 2, 3}; Array b = {1, 2, 3}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); Array c = {1, 3}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); auto diff_a_c = StructuralEqual::GetFirstMismatch(a, c); // first directly interepret diff, EXPECT_TRUE(diff_a_c.has_value()); auto lhs_steps = (*diff_a_c).get<0>()->ToSteps(); auto rhs_steps = (*diff_a_c).get<1>()->ToSteps(); EXPECT_EQ(lhs_steps[0]->kind, refl::AccessKind::kArrayItem); EXPECT_EQ(rhs_steps[0]->kind, refl::AccessKind::kArrayItem); EXPECT_EQ(lhs_steps[0]->key.cast(), 1); EXPECT_EQ(rhs_steps[0]->key.cast(), 1); EXPECT_EQ(lhs_steps.size(), 1); EXPECT_EQ(rhs_steps.size(), 1); // use structural equal for checking in future parts // given we have done some basic checks above by directly interepret diff, Array d = {1, 2}; auto diff_a_d = StructuralEqual::GetFirstMismatch(a, d); auto expected_diff_a_d = refl::AccessPathPair(refl::AccessPath::FromSteps({ refl::AccessStep::ArrayItem(2), }), refl::AccessPath::FromSteps({ refl::AccessStep::ArrayItemMissing(2), })); // then use structural equal to check it EXPECT_TRUE(StructuralEqual()(diff_a_d, expected_diff_a_d)); } TEST(StructuralEqualHash, Map) { // same map but different insertion order Map a = {{"a", 1}, {"b", 2}, {"c", 3}}; Map b = {{"b", 2}, {"c", 3}, {"a", 1}}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); Map c = {{"a", 1}, {"b", 2}, {"c", 4}}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); auto diff_a_c = StructuralEqual::GetFirstMismatch(a, c); auto expected_diff_a_c = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("c"), refl::AccessPath::Root()->MapItem("c")); EXPECT_TRUE(diff_a_c.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_c, expected_diff_a_c)); } TEST(StructuralEqualHash, NestedMapArray) { Map> a = {{"a", {1, 2, 3}}, {"b", {4, "hello", 6}}}; Map> b = {{"a", {1, 2, 3}}, {"b", {4, "hello", 6}}}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); Map> c = {{"a", {1, 2, 3}}, {"b", {4, "world", 6}}}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); auto diff_a_c = StructuralEqual::GetFirstMismatch(a, c); auto expected_diff_a_c = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("b")->ArrayItem(1), refl::AccessPath::Root()->MapItem("b")->ArrayItem(1)); EXPECT_TRUE(diff_a_c.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_c, expected_diff_a_c)); Map> d = {{"a", {1, 2, 3}}}; auto diff_a_d = StructuralEqual::GetFirstMismatch(a, d); auto expected_diff_a_d = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("b"), refl::AccessPath::Root()->MapItemMissing("b")); EXPECT_TRUE(diff_a_d.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_d, expected_diff_a_d)); auto diff_d_a = StructuralEqual::GetFirstMismatch(d, a); auto expected_diff_d_a = refl::AccessPathPair(refl::AccessPath::Root()->MapItemMissing("b"), refl::AccessPath::Root()->MapItem("b")); } TEST(StructuralEqualHash, Dict) { // same dict but different insertion order Dict a = {{"a", 1}, {"b", 2}, {"c", 3}}; Dict b = {{"b", 2}, {"c", 3}, {"a", 1}}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); Dict c = {{"a", 1}, {"b", 2}, {"c", 4}}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); auto diff_a_c = StructuralEqual::GetFirstMismatch(a, c); auto expected_diff_a_c = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("c"), refl::AccessPath::Root()->MapItem("c")); EXPECT_TRUE(diff_a_c.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_c, expected_diff_a_c)); } TEST(StructuralEqualHash, NestedDictArray) { Dict> a = {{"a", {1, 2, 3}}, {"b", {4, "hello", 6}}}; Dict> b = {{"a", {1, 2, 3}}, {"b", {4, "hello", 6}}}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); Dict> c = {{"a", {1, 2, 3}}, {"b", {4, "world", 6}}}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); auto diff_a_c = StructuralEqual::GetFirstMismatch(a, c); auto expected_diff_a_c = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("b")->ArrayItem(1), refl::AccessPath::Root()->MapItem("b")->ArrayItem(1)); EXPECT_TRUE(diff_a_c.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_c, expected_diff_a_c)); Dict> d = {{"a", {1, 2, 3}}}; auto diff_a_d = StructuralEqual::GetFirstMismatch(a, d); auto expected_diff_a_d = refl::AccessPathPair(refl::AccessPath::Root()->MapItem("b"), refl::AccessPath::Root()->MapItemMissing("b")); EXPECT_TRUE(diff_a_d.has_value()); EXPECT_TRUE(StructuralEqual()(diff_a_d, expected_diff_a_d)); } TEST(StructuralEqualHash, DictVsMapDifferentType) { Map m = {{"a", 1}, {"b", 2}}; Dict d = {{"a", 1}, {"b", 2}}; // Different type_index => not equal EXPECT_FALSE(StructuralEqual()(m, d)); // Different type_key_hash => different hash (very likely) EXPECT_NE(StructuralHash()(m), StructuralHash()(d)); } TEST(StructuralEqualHash, FreeVar) { TVar a = TVar("a"); TVar b = TVar("b"); EXPECT_TRUE(StructuralEqual::Equal(a, b, /*map_free_vars=*/true)); EXPECT_FALSE(StructuralEqual::Equal(a, b)); EXPECT_NE(StructuralHash()(a), StructuralHash()(b)); EXPECT_EQ(StructuralHash::Hash(a, /*map_free_vars=*/true), StructuralHash::Hash(b, /*map_free_vars=*/true)); } TEST(StructuralEqualHash, FuncDefAndIgnoreField) { TVar x = TVar("x"); TVar y = TVar("y"); // comment fields are ignored TFunc fa = TFunc({x}, {TInt(1), x}, String("comment a")); TFunc fb = TFunc({y}, {TInt(1), y}, String("comment b")); TFunc fc = TFunc({x}, {TInt(1), TInt(2)}, String("comment c")); EXPECT_TRUE(StructuralEqual()(fa, fb)); EXPECT_EQ(StructuralHash()(fa), StructuralHash()(fb)); EXPECT_FALSE(StructuralEqual()(fa, fc)); auto diff_fa_fc = StructuralEqual::GetFirstMismatch(fa, fc); auto expected_diff_fa_fc = refl::AccessPathPair(refl::AccessPath::FromSteps({ refl::AccessStep::Attr("body"), refl::AccessStep::ArrayItem(1), }), refl::AccessPath::FromSteps({ refl::AccessStep::Attr("body"), refl::AccessStep::ArrayItem(1), })); EXPECT_TRUE(diff_fa_fc.has_value()); EXPECT_TRUE(StructuralEqual()(diff_fa_fc, expected_diff_fa_fc)); } TEST(StructuralEqualHash, CustomTreeNode) { TVar x = TVar("x"); TVar y = TVar("y"); // comment fields are ignored TCustomFunc fa = TCustomFunc({x}, {TInt(1), x}, "comment a"); TCustomFunc fb = TCustomFunc({y}, {TInt(1), y}, "comment b"); TCustomFunc fc = TCustomFunc({x}, {TInt(1), TInt(2)}, "comment c"); EXPECT_TRUE(StructuralEqual()(fa, fb)); EXPECT_EQ(StructuralHash()(fa), StructuralHash()(fb)); EXPECT_FALSE(StructuralEqual()(fa, fc)); auto diff_fa_fc = StructuralEqual::GetFirstMismatch(fa, fc); auto expected_diff_fa_fc = refl::AccessPathPair(refl::AccessPath::Root()->Attr("body")->ArrayItem(1), refl::AccessPath::Root()->Attr("body")->ArrayItem(1)); EXPECT_TRUE(diff_fa_fc.has_value()); EXPECT_TRUE(StructuralEqual()(diff_fa_fc, expected_diff_fa_fc)); } // Regression tests for the SEqHashDefRecursive vs SEqHashDefNonRecursive // distinction. ``TDefHolder`` has two sibling fields: // - ``def_recursive`` tagged AttachFieldFlag::SEqHashDefRecursive() // - ``def_non_recursive`` tagged AttachFieldFlag::SEqHashDefNonRecursive() // each holding a ``TVarWithDep`` (a FreeVar with a sub-field ``dep`` that // can itself reference another FreeVar). The four sub-cases below cover // the observable behaviors of the two flags. TEST(StructuralEqualHash, NonRecursiveDef) { { // (a) Recursive flag rebinds nested FreeVars transitively. // ``def_non_recursive`` is the same object on both sides so it equates // by pointer; the case isolates the recursive field's rebinding. SCOPED_TRACE("recursive flag rebinds nested FreeVars"); TVarWithDep a("a", TVar("m")); TVarWithDep b("b", TVar("n")); TDefHolder lhs(/*def_recursive=*/a, /*def_non_recursive=*/a); TDefHolder rhs(/*def_recursive=*/b, /*def_non_recursive=*/b); EXPECT_TRUE(StructuralEqual()(lhs, rhs)); EXPECT_EQ(StructuralHash::Hash(lhs, /*map_free_vars=*/true), StructuralHash::Hash(rhs, /*map_free_vars=*/true)); } { // (b) Non-recursive flag does NOT rebind nested FreeVars: the top-level // FreeVar binds but the nested ``dep`` is clamped out of the def region. // With no outer binding for "p"/"q", equality must fail. SCOPED_TRACE("non-recursive flag does not rebind nested FreeVars"); TVarWithDep shared("shared", std::nullopt); TVarWithDep c_with_dep("c", TVar("p")); TVarWithDep d_with_dep("d", TVar("q")); TDefHolder lhs(/*def_recursive=*/shared, /*def_non_recursive=*/c_with_dep); TDefHolder rhs(/*def_recursive=*/shared, /*def_non_recursive=*/d_with_dep); EXPECT_FALSE(StructuralEqual::Equal(lhs, rhs, /*map_free_vars=*/false)); } { // (c) Non-recursive flag works if nested FreeVars resolve via an outer // binding — here we cheat by wiring the same pointer, so the nested // FreeVar passes the same-as pointer check without needing the def // region to be on inside its sub-field walk. SCOPED_TRACE("nested FreeVars resolve via outer pointer identity"); TVar shared_dep("dep"); TVarWithDep c_with_dep("c", shared_dep); TVarWithDep d_with_dep("d", shared_dep); TVarWithDep shared("shared", std::nullopt); TDefHolder lhs(/*def_recursive=*/shared, /*def_non_recursive=*/c_with_dep); TDefHolder rhs(/*def_recursive=*/shared, /*def_non_recursive=*/d_with_dep); EXPECT_TRUE(StructuralEqual()(lhs, rhs)); EXPECT_EQ(StructuralHash()(lhs), StructuralHash()(rhs)); } { // (d) Top-level FreeVar still binds under non-recursive — only the // FreeVar's sub-fields are clamped out; the binding step itself for // the immediate FreeVar is not suppressed. SCOPED_TRACE("top-level FreeVar still binds under non-recursive flag"); TVarWithDep shared("shared", std::nullopt); TVarWithDep c_no_dep("c", std::nullopt); TVarWithDep d_no_dep("d", std::nullopt); TDefHolder lhs(/*def_recursive=*/shared, /*def_non_recursive=*/c_no_dep); TDefHolder rhs(/*def_recursive=*/shared, /*def_non_recursive=*/d_no_dep); EXPECT_TRUE(StructuralEqual()(lhs, rhs)); EXPECT_EQ(StructuralHash::Hash(lhs, /*map_free_vars=*/true), StructuralHash::Hash(rhs, /*map_free_vars=*/true)); } } TEST(StructuralEqualHash, List) { List a = {1, 2, 3}; List b = {1, 2, 3}; EXPECT_TRUE(StructuralEqual()(a, b)); EXPECT_EQ(StructuralHash()(a), StructuralHash()(b)); List c = {1, 3}; EXPECT_FALSE(StructuralEqual()(a, c)); EXPECT_NE(StructuralHash()(a), StructuralHash()(c)); } TEST(StructuralEqualHash, ListVsArrayDifferentType) { Array arr = {1, 2, 3}; List lst = {1, 2, 3}; // Different type_index => not equal EXPECT_FALSE(StructuralEqual()(arr, lst)); // Different type_key_hash => different hash (very likely) EXPECT_NE(StructuralHash()(arr), StructuralHash()(lst)); } TEST(StructuralEqualHash, DISABLED_ListCycleDetection) { List lst; lst.push_back(42); lst.push_back(lst); // creates a cycle EXPECT_ANY_THROW(StructuralHash()(lst)); EXPECT_ANY_THROW(StructuralEqual()(lst, lst)); } TEST(StructuralEqualHash, ArraySelfInsertProducesSnapshot) { Array arr; arr.push_back(arr); Array snapshot = arr[0].cast>(); EXPECT_TRUE(snapshot.empty()); EXPECT_FALSE(snapshot.same_as(arr)); EXPECT_TRUE(StructuralEqual()(arr, arr)); EXPECT_EQ(StructuralHash()(arr), StructuralHash()(arr)); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_structural_key.cc000066400000000000000000000046761521067262500226630ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include namespace { using namespace tvm::ffi; TEST(StructuralKey, EqualityAndStdHash) { StructuralKey k1(Array{1, 2, 3}); StructuralKey k2(Array{1, 2, 3}); StructuralKey k3(Array{1, 2, 4}); EXPECT_FALSE(k1.same_as(k2)); EXPECT_EQ(k1->hash_i64, static_cast(StructuralHash::Hash(k1->key))); EXPECT_EQ(k2->hash_i64, static_cast(StructuralHash::Hash(k2->key))); EXPECT_TRUE(k1 == k2); EXPECT_FALSE(k1 != k2); EXPECT_FALSE(k1 == k3); EXPECT_TRUE(k1 != k3); EXPECT_EQ(std::hash()(k1), std::hash()(k2)); EXPECT_NE(std::hash()(k1), std::hash()(k3)); std::unordered_map map; map.emplace(k1, 10); map[k2] = 20; map[k3] = 30; EXPECT_EQ(map.size(), 2U); EXPECT_EQ(map.at(k1), 20); EXPECT_EQ(map.at(k2), 20); EXPECT_EQ(map.at(k3), 30); } TEST(StructuralKey, DefaultCtorInitializesHash) { StructuralKeyObj obj; EXPECT_EQ(obj.hash_i64, 0); } TEST(StructuralKey, MapAnyKeyUsesStructuralAttrs) { Map map; StructuralKey k1(Array{1, 2, 3}); StructuralKey k2(Array{1, 2, 3}); StructuralKey k3(Array{1, 3, 5}); map.Set(k1, 1); map.Set(k2, 2); map.Set(k3, 3); // k1 and k2 are structurally equal and should occupy the same map key. EXPECT_EQ(map.size(), 2U); EXPECT_EQ(map.at(k1).cast(), 2); EXPECT_EQ(map.at(k2).cast(), 2); EXPECT_EQ(map.at(k3).cast(), 3); } } // namespace tvm-ffi-0.1.12/tests/cpp/extra/test_visit_error_context.cc000066400000000000000000000617331521067262500237130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include namespace { using namespace tvm::ffi; namespace refl = tvm::ffi::reflection; // --------------------------------------------------------------------------- // Test fixture: TPair — a minimal two-field Object with lhs and rhs slots. // // Covers kAttr AccessKind via its named fields. // Array and Map are used alongside TPair in individual // tests to cover kArrayItem and kMapItem. // --------------------------------------------------------------------------- class TPair : public Object { public: ObjectRef lhs; ObjectRef rhs; static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.TPair", TPair, Object); }; class TPairRef : public ObjectRef { public: TPairRef(ObjectRef lhs, ObjectRef rhs) { ObjectPtr n = make_object(); n->lhs = std::move(lhs); n->rhs = std::move(rhs); data_ = std::move(n); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TPairRef, ObjectRef, TPair); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace r = tvm::ffi::reflection; r::ObjectDef().def_ro("lhs", &TPair::lhs).def_ro("rhs", &TPair::rhs); } // --------------------------------------------------------------------------- // Helper: Visit — recursively walks a tree and throws when it reaches // throw_target, wrapping each recursion level with the visit // macros so the error chain accumulates on unwind. // --------------------------------------------------------------------------- void Visit(const ObjectRef& node, const ObjectRef& throw_target); void Visit(const ObjectRef& node, const ObjectRef& throw_target) { TVM_FFI_VISIT_BEGIN(); if (node.same_as(throw_target)) { TVM_FFI_THROW(ValueError) << "boom"; } else if (Optional pair = node.as()) { Visit(pair.value()->lhs, throw_target); Visit(pair.value()->rhs, throw_target); } else if (Optional> arr = node.as>()) { for (const ObjectRef& child : *arr) { Visit(child, throw_target); } } else if (Optional> m = node.as>()) { for (const std::pair& kv : *m) { Visit(kv.second.cast(), throw_target); } } TVM_FFI_VISIT_END(node); } // =========================================================================== // Layer A — Macro → Chain // // Verify that TVM_FFI_VISIT_BEGIN/_END correctly builds the // reverse_visit_pattern as an exception unwinds through visit levels. // =========================================================================== // --------------------------------------------------------------------------- // MacroBuildsChain: throw deep in a two-level visit; confirm the chain // records nodes in innermost-first order (throw site first, root last). // --------------------------------------------------------------------------- TEST(VisitErrorContext, MacroBuildsChain) { // Tree: root = (lhs=leaf, rhs=empty) // Visiting root -> descends into leaf -> throws. Array empty; TPairRef leaf(empty, empty); TPairRef root(leaf, empty); Error caught("RuntimeError", "", ""); bool did_catch = false; try { Visit(root, leaf); } catch (Error& err) { caught = err; did_catch = true; } EXPECT_TRUE(did_catch); Optional visit_context = VisitErrorContext::TryGetFromError(caught); ASSERT_TRUE(visit_context.has_value()); // reverse_visit_pattern is innermost-first: leaf pushed first (inner catch // fires first during unwind), root pushed last. const List& chain = visit_context.value()->reverse_visit_pattern; EXPECT_GE(chain.size(), 2u); // Innermost entry is leaf (the throw site). EXPECT_TRUE(chain[0].same_as(leaf)); // Outermost entry is root. EXPECT_TRUE(chain[chain.size() - 1].same_as(root)); // Verify order structurally: [leaf, ..., root]. List expected_chain = {leaf, root}; EXPECT_TRUE(StructuralEqual::Equal(expected_chain, chain)); } // --------------------------------------------------------------------------- // MacroPreExistingPayloadWrap: when UpdateVisitErrorContext is called on // an error that already has a non-context extra_context payload, the existing // payload is stashed in prev_error_context and the visit chain starts // fresh. Subsequent pushes append to the same context object; prev_error_context // is not modified. // --------------------------------------------------------------------------- TEST(VisitErrorContext, MacroPreExistingPayloadWrap) { // Stand-in for a pre-existing extra_context payload. Array empty; TPairRef existing_payload(empty, empty); TPairRef leaf(empty, empty); TPairRef root(leaf, empty); // Construct an error that already carries a non-context extra_context payload. Error err("RuntimeError", "test", "", std::nullopt, std::optional(existing_payload)); ASSERT_TRUE(err.extra_context().has_value()); // First push: existing payload is not a VisitErrorContext → wrapped. tvm::ffi::details::UpdateVisitErrorContext(err, leaf); Optional visit_context = VisitErrorContext::TryGetFromError(err); ASSERT_TRUE(visit_context.has_value()); ASSERT_EQ(visit_context.value()->reverse_visit_pattern.size(), 1u); EXPECT_TRUE(visit_context.value()->reverse_visit_pattern[0].same_as(leaf)); ASSERT_TRUE(visit_context.value()->prev_error_context.has_value()); EXPECT_TRUE(visit_context.value()->prev_error_context.value().same_as(existing_payload)); // Second push: context already exists → node appended; prev_error_context unchanged. tvm::ffi::details::UpdateVisitErrorContext(err, root); Optional visit_context_after = VisitErrorContext::TryGetFromError(err); ASSERT_TRUE(visit_context_after.has_value()); List expected_chain = {leaf, root}; EXPECT_TRUE( StructuralEqual::Equal(expected_chain, visit_context_after.value()->reverse_visit_pattern)); ASSERT_TRUE(visit_context_after.value()->prev_error_context.has_value()); EXPECT_TRUE(visit_context_after.value()->prev_error_context.value().same_as(existing_payload)); } // --------------------------------------------------------------------------- // MacroThrowBuildsChain: TVM_FFI_VISIT_THROW seeds the // VisitErrorContext with the throw-site node directly (no separate // UpdateVisitErrorContext call). Enclosing VISIT_BEGIN/END frames then // append parent nodes on rethrow. // // Sub-scenario A — standalone: THROW outside any BEGIN/END wrapper produces // a context whose reverse_visit_pattern is exactly [throw_site]. // // Sub-scenario B — nested under BEGIN/END: when THROW fires inside a // VISIT_BEGIN/END(root) pair, the enclosing END appends root. The throw site // itself is recorded by THROW. If the throw_target happens to equal `node` // at the surrounding level (a common case), FindAccessPaths' cleanup // collapses the consecutive duplicate; tested separately in // FindAccessPaths.RecordsCleanup. // --------------------------------------------------------------------------- TEST(VisitErrorContext, MacroThrowBuildsChain) { Array empty; // Sub-scenario A — standalone THROW (no surrounding BEGIN/END). { TPairRef throw_site(empty, empty); bool did_catch = false; Error caught("placeholder", "", ""); try { TVM_FFI_VISIT_THROW(ValueError, throw_site) << "boom"; } catch (Error& err) { caught = err; did_catch = true; } ASSERT_TRUE(did_catch); Optional visit_context = VisitErrorContext::TryGetFromError(caught); ASSERT_TRUE(visit_context.has_value()); const List& chain = visit_context.value()->reverse_visit_pattern; ASSERT_EQ(chain.size(), 1u); EXPECT_TRUE(chain[0].same_as(throw_site)); } // Sub-scenario B — THROW inside a BEGIN/END(root) wrapper. { TPairRef throw_site(empty, empty); TPairRef root(throw_site, empty); bool did_catch = false; Error caught("placeholder", "", ""); try { TVM_FFI_VISIT_BEGIN(); TVM_FFI_VISIT_THROW(ValueError, throw_site) << "boom"; TVM_FFI_VISIT_END(root); } catch (Error& err) { caught = err; did_catch = true; } ASSERT_TRUE(did_catch); Optional visit_context = VisitErrorContext::TryGetFromError(caught); ASSERT_TRUE(visit_context.has_value()); const List& chain = visit_context.value()->reverse_visit_pattern; // THROW seeds with throw_site; END appends root. List expected = {throw_site, root}; EXPECT_TRUE(StructuralEqual::Equal(expected, chain)); } } // =========================================================================== // Layer B — Chain → AccessPaths // // Verify that FindAccessPaths(root, ctx) resolves a manually-constructed // VisitErrorContext against a tree, producing AccessPaths that describe // where matched nodes live in the tree. // // Each test: // 1. Builds a root tree. // 2. Constructs a VisitErrorContext with a specific reverse_visit_pattern. // 3. Calls FindAccessPaths(root, ctx). // 4. Asserts on the result using StructuralEqual::Equal on AccessPath // (mirrors the pattern used in test_structural_equal_hash.cc). // =========================================================================== // --------------------------------------------------------------------------- // BasicMatch: unambiguous single-node path via two Attr steps; and CSE // multi-match when the same pointer appears in two sibling slots. // // Sub-scenario A — SingleMatch: // Tree: root.lhs = mid, mid.lhs = leaf // Pattern: [leaf, mid, root] (innermost-first) // Expected path: Root->Attr("lhs")->Attr("lhs") // // Sub-scenario B — CSEMultiMatch: // Tree: root.lhs = shared, root.rhs = shared (identical pointer) // Pattern: [shared, root] // Expected: two paths, one for each slot (Attr("lhs") and Attr("rhs")). // --------------------------------------------------------------------------- TEST(FindAccessPaths, BasicMatch) { Array empty; // Sub-scenario A: single unambiguous path. { TPairRef leaf(empty, empty); TPairRef mid(leaf, empty); TPairRef root(mid, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{leaf, mid, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs")->Attr("lhs"); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } // Sub-scenario B: same pointer in two slots → two paths. { TPairRef shared(empty, empty); TPairRef root(shared, shared); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{shared, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 2u); refl::AccessPath expected_lhs = refl::AccessPath::Root()->Attr("lhs"); refl::AccessPath expected_rhs = refl::AccessPath::Root()->Attr("rhs"); bool found_lhs = false; bool found_rhs = false; for (const refl::AccessPath& p : paths) { if (StructuralEqual::Equal(p, expected_lhs)) found_lhs = true; if (StructuralEqual::Equal(p, expected_rhs)) found_rhs = true; } EXPECT_TRUE(found_lhs); EXPECT_TRUE(found_rhs); } } // --------------------------------------------------------------------------- // AccessKindCoverage: exercises Attr (TPair fields), ArrayItem // (Array), and MapItem (Map) access kinds. // // Sub-scenario A — Attr (kAttr): // Tree: root.lhs = mid, mid.lhs = leaf // Pattern: [leaf, root] // Expected path: Root->Attr("lhs")->Attr("lhs") // // Sub-scenario B — ArrayItem (kArrayItem): // Tree: root.lhs = [leaf1, leaf2] (Array in ObjectRef slot) // Pattern: [leaf1, root] // Expected path: Root->Attr("lhs")->ArrayItem(0) // // Sub-scenario C — MapItem (kMapItem): // Tree: root.lhs = {"key" -> target} (Map in ObjectRef slot) // Pattern: [target, root] // Expected path: Root->Attr("lhs")->MapItem("key") // --------------------------------------------------------------------------- TEST(FindAccessPaths, AccessKindCoverage) { Array empty; // Sub-scenario A: Attr steps. { TPairRef leaf(empty, empty); TPairRef mid(leaf, empty); TPairRef root(mid, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{leaf, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs")->Attr("lhs"); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } // Sub-scenario B: ArrayItem step. { TPairRef leaf1(empty, empty); TPairRef leaf2(empty, empty); Array arr = {leaf1, leaf2}; TPairRef root(arr, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{leaf1, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs")->ArrayItem(0); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } // Sub-scenario C: MapItem step. { TPairRef target(empty, empty); Map m; m.Set(String("key"), target); TPairRef root(m, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{target, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs")->MapItem(String("key")); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } } // --------------------------------------------------------------------------- // SparsePatternAnchors: demonstrates that intermediate breadcrumb anchors in // the pattern disambiguate when an innermost target appears in multiple branches. // // Tree layout: // root // ├─ branch_a: tpa → m_inner_a (m_inner shared, no m_outer ancestor) // └─ branch_b: tpb → m_outer → tpc → m_inner_b (same m_inner pointer) // // m_inner is a POINTER-IDENTICAL ObjectRef (ref-count shared) used as both // m_inner_a (under branch_a) and m_inner_b (under branch_b via m_outer). // m_outer appears only under branch_b. // // Scenario 1: pattern with only innermost — 2 candidate paths. // Pattern: [m_inner, root] // Expected: paths.size() == 2 // // Scenario 2: anchor narrows to branch_b only. // Pattern: [m_inner, m_outer, root] // Expected: paths.size() == 1, path through branch_b // --------------------------------------------------------------------------- TEST(FindAccessPaths, SparsePatternAnchors) { Array empty; // Shared inner node — identical pointer in both branches. TPairRef m_inner(empty, empty); // branch_a: tpa.lhs = m_inner (no m_outer ancestor) TPairRef tpa(m_inner, empty); // branch_b: tpb.lhs = m_outer, m_outer.lhs = tpc, tpc.lhs = m_inner TPairRef tpc(m_inner, empty); TPairRef m_outer(tpc, empty); TPairRef tpb(m_outer, empty); // root: lhs = tpa (branch_a), rhs = tpb (branch_b) TPairRef root(tpa, tpb); // Scenario 1: pattern = [m_inner, root] → 2 paths (one per branch). { ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{m_inner, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); EXPECT_EQ(paths.size(), 2u); } // Scenario 2: pattern = [m_inner, m_outer, root] → 1 path through branch_b. { ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{m_inner, m_outer, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); // Path must go through branch_b (rhs), then m_outer (lhs), then tpc (lhs), then m_inner (lhs). refl::AccessPath expected = refl::AccessPath::Root()->Attr("rhs")->Attr("lhs")->Attr("lhs")->Attr("lhs"); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } } // --------------------------------------------------------------------------- // PartialChain: strict mode returns nothing when the innermost pattern entry // is unreachable; prefix-match mode returns the deepest matched prefix path. // // Sub-scenario A — strict (allow_prefix_match=false): // Pattern: [unreachable, child, root] // Expected: 0 results (unreachable breaks full match) // // Sub-scenario B — prefix match (allow_prefix_match=true): // Same pattern, same tree. // Expected: at least one result at the path to child (Root->Attr("lhs")). // --------------------------------------------------------------------------- TEST(FindAccessPaths, PartialChain) { Array empty; // unreachable is not present anywhere in the tree below root. TPairRef unreachable(empty, empty); TPairRef child(empty, empty); TPairRef root(child, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{unreachable, child, root}; VisitErrorContext ctx(std::move(ctx_obj)); // Sub-scenario A: strict mode — no full match. { Array paths = VisitErrorContext::FindAccessPaths(root, ctx); EXPECT_EQ(paths.size(), 0u); } // Sub-scenario B: prefix-match mode — path to child reported. { Array paths = VisitErrorContext::FindAccessPaths(root, ctx, /*allow_prefix_match=*/true); ASSERT_GE(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs"); bool found_expected = false; for (const refl::AccessPath& p : paths) { if (StructuralEqual::Equal(p, expected)) { found_expected = true; break; } } EXPECT_TRUE(found_expected); } } // --------------------------------------------------------------------------- // EdgeCases: empty pattern produces no results; null entries in the pattern // never match any live node and do not crash. // --------------------------------------------------------------------------- TEST(FindAccessPaths, EdgeCases) { Array empty; // Sub-scenario A: empty pattern. { TPairRef root(empty, empty); ObjectPtr ctx_obj = make_object(); // reverse_visit_pattern left empty. VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); EXPECT_EQ(paths.size(), 0u); } // Sub-scenario A2: root-itself match — pattern is [root], the throw fired // before any descent. Expected: a single AccessPath equal to Root() (no // crash from materializing on an empty descent stack). { TPairRef root(empty, empty); ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); EXPECT_TRUE(StructuralEqual::Equal(refl::AccessPath::Root(), paths[0])); } // Sub-scenario B: null entries in pattern. { TPairRef root(empty, empty); ObjectPtr ctx_obj = make_object(); ObjectRef null_ref; ctx_obj->reverse_visit_pattern = List{null_ref, null_ref}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths; ASSERT_NO_THROW({ paths = VisitErrorContext::FindAccessPaths(root, ctx); }); EXPECT_EQ(paths.size(), 0u); } } // --------------------------------------------------------------------------- // RecordsCleanup: FindAccessPaths normalizes reverse_visit_pattern before // matching by (a) dropping null entries and (b) collapsing runs of // consecutive duplicates. Both arise naturally — null from stale/torn-down // nodes; duplicates when TVM_FFI_VISIT_THROW(node) is wrapped by // a TVM_FFI_VISIT_END(node) at the same level. Without cleanup // the redundant frame would over-constrain the match and yield zero paths. // --------------------------------------------------------------------------- TEST(FindAccessPaths, RecordsCleanup) { Array empty; TPairRef leaf(empty, empty); TPairRef root(leaf, empty); ObjectRef null_ref; refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs"); // Sub-scenario A — consecutive duplicates collapse. // Raw pattern: [leaf, leaf, root, root]; effective: [leaf, root]. { ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{leaf, leaf, root, root}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } // Sub-scenario B — nulls are skipped, surrounding entries still match. // Raw pattern: [null, leaf, null, root, null]; effective: [leaf, root]. { ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{null_ref, leaf, null_ref, root, null_ref}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } // Sub-scenario C — combined: nulls plus dup-consecutive. // Raw pattern: [leaf, null, leaf, root, root, null]; effective: [leaf, root]. { ObjectPtr ctx_obj = make_object(); ctx_obj->reverse_visit_pattern = List{leaf, null_ref, leaf, root, root, null_ref}; VisitErrorContext ctx(std::move(ctx_obj)); Array paths = VisitErrorContext::FindAccessPaths(root, ctx); ASSERT_EQ(paths.size(), 1u); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } } // --------------------------------------------------------------------------- // TryGetFromError: returns NullOpt when absent, the context when present. // Composes correctly with FindAccessPaths. // --------------------------------------------------------------------------- TEST(VisitErrorContext, TryGetFromError) { Array empty; TPairRef leaf(empty, empty); TPairRef root(leaf, empty); Error err("RuntimeError", "test", ""); // No context attached → TryGetFromError returns NullOpt. Optional no_context = VisitErrorContext::TryGetFromError(err); EXPECT_FALSE(no_context.has_value()); // Attach context via UpdateVisitErrorContext. tvm::ffi::details::UpdateVisitErrorContext(err, leaf); tvm::ffi::details::UpdateVisitErrorContext(err, root); Optional visit_context = VisitErrorContext::TryGetFromError(err); ASSERT_TRUE(visit_context.has_value()); // Compose TryGetFromError + FindAccessPaths → should resolve to root.lhs. Array paths = VisitErrorContext::FindAccessPaths(root, visit_context.value()); ASSERT_EQ(paths.size(), 1u); refl::AccessPath expected = refl::AccessPath::Root()->Attr("lhs"); EXPECT_TRUE(StructuralEqual::Equal(expected, paths[0])); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_any.cc000066400000000000000000000322051521067262500172340ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Any, Int) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `int`"), std::string::npos); throw; } }, ::tvm::ffi::Error); AnyView view1 = 1; EXPECT_EQ(view1.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIInt); EXPECT_EQ(view1.CopyToTVMFFIAny().v_int64, 1); auto int_v1 = view1.cast(); EXPECT_EQ(int_v1, 1); int64_t v1 = 2; view0 = v1; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIInt); EXPECT_EQ(view0.CopyToTVMFFIAny().v_int64, 2); uint64_t v2 = static_cast(std::numeric_limits::max()) + 1; EXPECT_THROW( { try { view0 = v2; } catch (const Error& error) { EXPECT_EQ(error.kind(), "OverflowError"); std::string what = error.what(); EXPECT_NE(what.find("is too large to fit in int64_t"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(Any, Enum) { enum class ENum : int { A = 1, B = 2, }; AnyView view0; Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); AnyView view1 = ENum::A; EXPECT_EQ(view1.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIInt); EXPECT_EQ(view1.CopyToTVMFFIAny().v_int64, 1); ENum v1 = view1.cast(); EXPECT_EQ(v1, ENum::A); } TEST(Any, bool) { AnyView view0; Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `bool`"), std::string::npos); throw; } }, ::tvm::ffi::Error); AnyView view1 = true; EXPECT_EQ(view1.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIBool); EXPECT_EQ(view1.CopyToTVMFFIAny().v_int64, 1); auto int_v1 = view1.cast(); EXPECT_EQ(int_v1, 1); bool v1 = false; view0 = v1; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIBool); EXPECT_EQ(view0.CopyToTVMFFIAny().v_int64, 0); } TEST(Any, nullptrcmp) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); EXPECT_TRUE(view0 == nullptr); EXPECT_FALSE(view0 != nullptr); view0 = 1; EXPECT_TRUE(view0 != nullptr); EXPECT_FALSE(view0 == nullptr); Any any0 = view0; EXPECT_TRUE(any0 != nullptr); EXPECT_FALSE(any0 == nullptr); any0 = nullptr; EXPECT_TRUE(any0 == nullptr); EXPECT_FALSE(any0 != nullptr); } TEST(Any, Float) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `float`"), std::string::npos); throw; } }, ::tvm::ffi::Error); AnyView view1_int = 1; auto float_v1 = view1_int.cast(); EXPECT_EQ(float_v1, 1); AnyView view2 = 2.2; EXPECT_EQ(view2.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIFloat); EXPECT_EQ(view2.CopyToTVMFFIAny().v_float64, 2.2); float v1 = 2; view0 = v1; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIFloat); EXPECT_EQ(view0.CopyToTVMFFIAny().v_float64, 2); } TEST(Any, Device) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `Device`"), std::string::npos); throw; } }, ::tvm::ffi::Error); DLDevice device{kDLCUDA, 1}; AnyView view1_device = device; auto dtype_v1 = view1_device.cast(); EXPECT_EQ(dtype_v1.device_type, kDLCUDA); EXPECT_EQ(dtype_v1.device_id, 1); Any any2 = DLDevice{kDLCPU, 0}; TVMFFIAny ffi_v2 = details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(any2)); EXPECT_EQ(ffi_v2.type_index, TypeIndex::kTVMFFIDevice); EXPECT_EQ(ffi_v2.v_device.device_type, kDLCPU); EXPECT_EQ(ffi_v2.v_device.device_id, 0); } TEST(Any, DLTensor) { AnyView view0; Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `DLTensor*`"), std::string::npos); throw; } }, ::tvm::ffi::Error); DLTensor dltensor; AnyView view1_dl = &dltensor; auto dl_v1 = view1_dl.cast(); EXPECT_EQ(dl_v1, &dltensor); } TEST(Any, Object) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); // int object is not nullable Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); TInt v1(11); EXPECT_EQ(v1.use_count(), 1); // view won't increase refcount AnyView view1 = v1; EXPECT_EQ(v1.use_count(), 1); // any will trigger ref count increase Any any1 = v1; EXPECT_EQ(v1.use_count(), 2); // copy to another view AnyView view2 = any1; EXPECT_EQ(v1.use_count(), 2); // convert to weak raw object ptr const TIntObj* v1_ptr = view2.cast(); EXPECT_EQ(v1.use_count(), 2); EXPECT_EQ(v1_ptr->value, 11); Any any2 = v1_ptr; EXPECT_EQ(v1.use_count(), 3); EXPECT_TRUE(any2.as().has_value()); any2 = const_cast(v1_ptr); EXPECT_TRUE(any2.as().has_value()); // convert to raw opaque ptr void* raw_v1_ptr = const_cast(v1_ptr); any2 = raw_v1_ptr; EXPECT_TRUE(any2.as().value() == v1_ptr); // NOLINT(bugprone-unchecked-optional-access) // convert to ObjectRef { auto v1_obj_ref = view2.cast(); EXPECT_EQ(v1.use_count(), 3); any2 = v1_obj_ref; EXPECT_EQ(v1.use_count(), 4); EXPECT_TRUE(any2.as().has_value()); any2.reset(); } // convert that triggers error EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view1.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); std::cout << what; EXPECT_NE(what.find("Cannot convert from type `test.Int` to `test.Float`"), std::string::npos); throw; } }, ::tvm::ffi::Error); // Try to convert to number auto number0 = any1.cast(); EXPECT_EQ(v1.use_count(), 3); EXPECT_TRUE(number0.as()); EXPECT_EQ(number0.as()->value, 11); EXPECT_TRUE(!any1.as().has_value()); auto int1 = view2.cast(); EXPECT_EQ(v1.use_count(), 4); any1.reset(); EXPECT_EQ(v1.use_count(), 3); } TEST(Any, ObjectRefWithFallbackTraits) { // Test case for TPrimExpr fallback from Any Any any1 = TPrimExpr("float32", 3.14); auto v0 = any1.cast(); EXPECT_EQ(v0->value, 3.14); EXPECT_EQ(v0->dtype, "float32"); any1 = true; auto v1 = any1.cast(); EXPECT_EQ(v1->value, 1); EXPECT_EQ(v1->dtype, "bool"); any1 = static_cast(42); auto v2 = any1.cast(); EXPECT_EQ(v2->value, 42); EXPECT_EQ(v2->dtype, "int64"); any1 = 2.718; auto v3 = any1.cast(); EXPECT_EQ(v3->value, 2.718); EXPECT_EQ(v3->dtype, "float32"); // Test case for TPrimExpr fallback from AnyView TPrimExpr texpr1("float32", 3.14); AnyView view1 = texpr1; auto v4 = view1.cast(); EXPECT_EQ(v4->value, 3.14); EXPECT_EQ(v4->dtype, "float32"); view1 = true; auto v5 = view1.cast(); EXPECT_EQ(v5->value, 1); EXPECT_EQ(v5->dtype, "bool"); view1 = static_cast(42); auto v6 = view1.cast(); EXPECT_EQ(v6->value, 42); EXPECT_EQ(v6->dtype, "int64"); view1 = 2.718; auto v7 = view1.cast(); EXPECT_EQ(v7->value, 2.718); EXPECT_EQ(v7->dtype, "float32"); // Test case for TPrimExpr fallback from Any with String any1 = std::string("test_string"); auto v8 = any1.cast(); EXPECT_EQ(v8->dtype, "test_string"); EXPECT_EQ(v8->value, 0); // Test case for TPrimExpr fallback from AnyView with String view1 = "test_string"; auto v9 = view1.cast(); EXPECT_EQ(v9->dtype, "test_string"); EXPECT_EQ(v9->value, 0); } TEST(Any, CastVsAs) { AnyView view0 = 1; // as only runs strict check auto opt_v0 = view0.as(); EXPECT_TRUE(opt_v0.has_value()); EXPECT_EQ(opt_v0.value(), 1); // NOLINT(bugprone-unchecked-optional-access) auto opt_v1 = view0.as(); EXPECT_TRUE(!opt_v1.has_value()); auto opt_v2 = view0.as(); EXPECT_TRUE(!opt_v2.has_value()); // try_cast will try run the conversion. auto opt_v3 = view0.try_cast(); EXPECT_TRUE(opt_v3.has_value()); EXPECT_EQ(opt_v3.value(), 1); // NOLINT(bugprone-unchecked-optional-access) auto opt_v4 = view0.try_cast(); EXPECT_TRUE(opt_v4.has_value()); EXPECT_EQ(opt_v4.value(), 1); // NOLINT(bugprone-unchecked-optional-access) Any any1 = true; auto opt_v5 = any1.as(); EXPECT_TRUE(opt_v5.has_value()); EXPECT_EQ(opt_v5.value(), 1); // NOLINT(bugprone-unchecked-optional-access) auto opt_v6 = any1.try_cast(); EXPECT_TRUE(opt_v6.has_value()); EXPECT_EQ(opt_v6.value(), 1); // NOLINT(bugprone-unchecked-optional-access) auto opt_v7 = any1.try_cast(); EXPECT_TRUE(opt_v7.has_value()); } TEST(Any, ObjectMove) { Any any1 = TPrimExpr("float32", 3.14); auto v0 = std::move(any1).cast(); EXPECT_EQ(v0->value, 3.14); EXPECT_EQ(v0.use_count(), 1); EXPECT_TRUE(any1 == nullptr); // NOLINT(bugprone-use-after-move) } TEST(Any, AnyEqualHash) { // small string Any a = "a1"; // on heap allocated string Any b = String(std::string("a1")); EXPECT_EQ(a.type_index(), TypeIndex::kTVMFFISmallStr); EXPECT_EQ(b.type_index(), TypeIndex::kTVMFFIStr); EXPECT_TRUE(AnyEqual()(a, b)); EXPECT_EQ(AnyHash()(a), AnyHash()(b)); Any c = Bytes("a1", 2); Any d = Bytes(std::string("a1")); EXPECT_EQ(c.type_index(), TypeIndex::kTVMFFISmallBytes); EXPECT_EQ(d.type_index(), TypeIndex::kTVMFFIBytes); EXPECT_TRUE(AnyEqual()(c, d)); EXPECT_EQ(AnyHash()(c), AnyHash()(d)); } TEST(Any, CustomAnyHash) { // Covers the OpaquePtr custom hash branch. Any int_src = TInt(7); uint64_t int_expected = details::StableHashCombine( int_src.type_index(), static_cast(TInt::CustomAnyHash(int_src))); EXPECT_EQ(AnyHash()(int_src), int_expected); // Covers the ffi.Function custom hash branch. Any float_src = TFloat(3.5); uint64_t float_expected = details::StableHashCombine( float_src.type_index(), static_cast(TFloat::CustomAnyHash(float_src))); EXPECT_EQ(AnyHash()(float_src), float_expected); } TEST(Any, CustomAnyEqual) { // Covers the OpaquePtr custom equal branch. Any int_lhs = TInt(7); Any int_rhs = TInt(7); Any int_diff = TInt(8); EXPECT_TRUE(AnyEqual()(int_lhs, int_rhs)); EXPECT_FALSE(AnyEqual()(int_lhs, int_diff)); // Covers the ffi.Function custom equal branch. Any float_lhs = TFloat(3.5); Any float_rhs = TFloat(3.5); Any float_diff = TFloat(4.5); EXPECT_TRUE(AnyEqual()(float_lhs, float_rhs)); EXPECT_FALSE(AnyEqual()(float_lhs, float_diff)); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_array.cc000066400000000000000000000236001521067262500175620ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Array, Basic) { Array arr = {TInt(11), TInt(12)}; TInt v1 = arr[0]; EXPECT_EQ(v1->value, 11); EXPECT_EQ(v1.use_count(), 2); EXPECT_EQ(arr[1]->value, 12); } TEST(Array, COWSet) { Array arr = {TInt(11), TInt(12)}; Array arr2 = arr; EXPECT_EQ(arr.use_count(), 2); arr.Set(1, TInt(13)); EXPECT_EQ(arr.use_count(), 1); EXPECT_EQ(arr[1]->value, 13); EXPECT_EQ(arr2[1]->value, 12); } TEST(Array, MutateInPlaceForUniqueReference) { TInt x(1); Array arr{x, x}; EXPECT_TRUE(arr.unique()); auto* before = arr.get(); // NOLINTNEXTLINE(performance-unnecessary-value-param) arr.MutateByApply([](TInt) { return TInt(2); }); auto* after = arr.get(); EXPECT_EQ(before, after); } TEST(Array, CopyWhenMutatingNonUniqueReference) { TInt x(1); Array arr{x, x}; Array arr2 = arr; EXPECT_TRUE(!arr.unique()); auto* before = arr.get(); // NOLINTNEXTLINE(performance-unnecessary-value-param) arr.MutateByApply([](TInt) { return TInt(2); }); auto* after = arr.get(); EXPECT_NE(before, after); } TEST(Array, Map) { // Basic functionality TInt x(1), y(1); Array var_arr{x, y}; Array expr_arr = var_arr.Map([](TInt var) -> TNumber { // NOLINT(performance-unnecessary-value-param) return TFloat(static_cast(var->value + 1)); }); EXPECT_NE(var_arr.get(), expr_arr.get()); EXPECT_TRUE(expr_arr[0]->IsInstance()); EXPECT_TRUE(expr_arr[1]->IsInstance()); } TEST(Array, Iterator) { Array array{1, 2, 3}; std::vector vector(array.begin(), array.end()); EXPECT_EQ(vector[1], 2); } TEST(Array, PushPop) { Array a; std::vector b; for (int i = 0; i < 10; ++i) { a.push_back(i); b.push_back(i); ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), b.size()); int n = static_cast(a.size()); for (int j = 0; j < n; ++j) { ASSERT_EQ(a[j], b[j]); } } for (int i = 9; i >= 0; --i) { ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), b.size()); a.pop_back(); b.pop_back(); int n = static_cast(a.size()); for (int j = 0; j < n; ++j) { ASSERT_EQ(a[j], b[j]); } } ASSERT_EQ(a.empty(), true); } TEST(Array, ResizeReserveClear) { for (size_t n = 0; n < 10; ++n) { Array a; Array b; a.resize(static_cast(n)); b.reserve(static_cast(n)); ASSERT_EQ(a.size(), n); ASSERT_GE(a.capacity(), n); a.clear(); b.clear(); ASSERT_EQ(a.size(), 0); ASSERT_EQ(b.size(), 0); } } TEST(Array, InsertErase) { Array a; std::vector b; for (int n = 1; n <= 10; ++n) { a.insert(a.end(), n); b.insert(b.end(), n); for (int pos = 0; pos <= n; ++pos) { a.insert(a.begin() + pos, pos); b.insert(b.begin() + pos, pos); ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), n + 1); ASSERT_EQ(b.size(), n + 1); for (int k = 0; k <= n; ++k) { ASSERT_EQ(a[k], b[k]); } a.erase(a.begin() + pos); b.erase(b.begin() + pos); } ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), n); } } TEST(Array, InsertEraseRange) { Array range_a{-1, -2, -3, -4}; std::vector range_b{-1, -2, -3, -4}; Array a; std::vector b; static_assert(std::is_same_v); for (size_t n = 1; n <= 10; ++n) { a.insert(a.end(), static_cast(n)); b.insert(b.end(), static_cast(n)); for (size_t pos = 0; pos <= n; ++pos) { a.insert(a.begin() + static_cast(pos), range_a.begin(), range_a.end()); b.insert(b.begin() + static_cast(pos), range_b.begin(), range_b.end()); ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), n + range_a.size()); ASSERT_EQ(b.size(), n + range_b.size()); size_t m = n + range_a.size(); for (size_t k = 0; k < m; ++k) { ASSERT_EQ(a[k], b[k]); } a.erase(a.begin() + static_cast(pos), a.begin() + static_cast(pos + range_a.size())); b.erase(b.begin() + static_cast(pos), b.begin() + static_cast(pos + range_b.size())); } ASSERT_EQ(a.front(), b.front()); ASSERT_EQ(a.back(), b.back()); ASSERT_EQ(a.size(), n); } } TEST(Array, FuncArrayAnyArg) { // NOLINTNEXTLINE(performance-unnecessary-value-param) Function fadd_one = Function::FromTyped([](Array a) -> Any { return a[0].cast() + 1; }); EXPECT_EQ(fadd_one(Array{1}).cast(), 2); } TEST(Array, MapUniquePropogation) { // Basic functionality Array var_arr{TInt(1), TInt(2)}; // NOLINTNEXTLINE(performance-unnecessary-value-param) var_arr.MutateByApply([](TInt x) -> TInt { EXPECT_TRUE(x.unique()); return x; }); } TEST(Array, AnyImplicitConversion) { Array arr0_mixed = {11.1, 1}; EXPECT_EQ(arr0_mixed[1].cast(), 1); AnyView view0 = arr0_mixed; auto arr0_float = view0.cast>(); // they are not the same because arr_mixed // stores arr_mixed[1] as int but we need to convert to float EXPECT_TRUE(!arr0_float.same_as(arr0_mixed)); EXPECT_EQ(arr0_float[1], 1.0); Any any1 = arr0_float; // if storage check passes, the same array get returned auto arr1_float = any1.cast>(); EXPECT_TRUE(arr1_float.same_as(arr0_float)); // total count equals 3 include any1 EXPECT_EQ(arr1_float.use_count(), 3); // convert to Array do not need any conversion auto arr1_mixed = any1.cast>(); EXPECT_TRUE(arr1_mixed.same_as(arr1_float)); EXPECT_EQ(arr1_float.use_count(), 4); } TEST(Array, AnyConvertCheck) { Array arr = {11.1, 1}; EXPECT_EQ(arr[1].cast(), 1); AnyView view0 = arr; auto arr1 = view0.cast>(); EXPECT_EQ(arr1[0], 11.1); EXPECT_EQ(arr1[1], 1.0); Any any1 = arr; EXPECT_THROW( { try { [[maybe_unused]] auto arr2 = any1.cast>(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `Array[index 0: float]` to `Array`"), std::string::npos); throw; } }, ::tvm::ffi::Error); Array> arr_nested = {{}, {TInt(1), TFloat(2)}}; any1 = arr_nested; auto arr1_nested = any1.cast>>(); EXPECT_EQ(arr1_nested.use_count(), 3); EXPECT_THROW( { try { [[maybe_unused]] auto arr2 = any1.cast>>(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("`Array[index 1: Array[index 0: test.Int]]` to `Array>`"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(Array, Upcast) { Array a0 = {1, 2, 3}; Array a1 = a0; EXPECT_EQ(a1[0].cast(), 1); EXPECT_EQ(a1[1].cast(), 2); EXPECT_EQ(a1[2].cast(), 3); Array> a2 = {a0}; Array> a3 = a2; Array> a4 = a2; static_assert(details::type_contains_v, Array>); static_assert(details::type_contains_v>); } TEST(Array, Contains) { Function f = Function::GetGlobalRequired("ffi.ArrayContains"); Array arr = {1, 2, 3, 4, 5}; EXPECT_TRUE(f(arr, 3).cast()); EXPECT_TRUE(f(arr, 1).cast()); EXPECT_TRUE(f(arr, 5).cast()); EXPECT_FALSE(f(arr, 10).cast()); EXPECT_FALSE(f(arr, 0).cast()); Array empty_arr; EXPECT_FALSE(f(empty_arr, 1).cast()); Array str_arr = {String("hello"), String("world")}; EXPECT_TRUE(f(str_arr, String("hello")).cast()); EXPECT_TRUE(f(str_arr, String("world")).cast()); EXPECT_FALSE(f(str_arr, String("foo")).cast()); } TEST(Array, NegativeIndexThrows) { Array arr = {1, 2, 3}; // Directly test ArrayObj methods, which are the ones modified in this PR. // The Array wrapper methods already had negative index checks. ArrayObj* arr_obj = arr.GetArrayObj(); // Test ArrayObj::at (which calls operator[]) EXPECT_THROW( { try { [[maybe_unused]] const auto& val = arr_obj->at(-1); } catch (const Error& error) { EXPECT_EQ(error.kind(), "IndexError"); throw; } }, ::tvm::ffi::Error); // Test ArrayObj::SetItem EXPECT_THROW( { try { arr_obj->SetItem(-1, Any(42)); } catch (const Error& error) { EXPECT_EQ(error.kind(), "IndexError"); throw; } }, ::tvm::ffi::Error); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_base.cc000066400000000000000000000021431521067262500173550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include namespace { TEST(Base, Version) { TVMFFIVersion version; TVMFFIGetVersion(&version); EXPECT_EQ(version.major, TVM_FFI_VERSION_MAJOR); EXPECT_EQ(version.minor, TVM_FFI_VERSION_MINOR); EXPECT_EQ(version.patch, TVM_FFI_VERSION_PATCH); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_dict.cc000066400000000000000000000122051521067262500173660ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include namespace { using namespace tvm::ffi; TEST(Dict, Basic) { Dict d; d.Set("a", 1); d.Set("b", 2); EXPECT_EQ(d.size(), 2); EXPECT_EQ(d.at("a"), 1); EXPECT_EQ(d["b"], 2); EXPECT_EQ(d.count("a"), 1); EXPECT_EQ(d.count("c"), 0); EXPECT_FALSE(d.empty()); } TEST(Dict, FindAndGet) { Dict d; d.Set("x", 42); auto it = d.find("x"); EXPECT_TRUE(it != d.end()); EXPECT_EQ((*it).second, 42); auto it2 = d.find("y"); EXPECT_TRUE(it2 == d.end()); auto opt = d.Get("x"); ASSERT_TRUE(opt.has_value()); EXPECT_EQ(opt.value(), 42); // NOLINT(bugprone-unchecked-optional-access) auto opt2 = d.Get("y"); EXPECT_FALSE(opt2.has_value()); } TEST(Dict, SharedMutation) { // Two handles point to the same DictObj Dict d1; d1.Set("a", 1); Dict d2 = d1; // shallow copy of ObjectRef // Mutate through d1 d1.Set("b", 2); // d2 should see the change (no COW) EXPECT_EQ(d2.size(), 2); EXPECT_EQ(d2["b"], 2); // Same underlying object EXPECT_EQ(d1.get(), d2.get()); } TEST(Dict, InplaceSwitchTo) { // Insert >4 elements to trigger transition from small to dense layout. // Verify the ObjectPtr address stays the same. Dict d1; d1.Set("a", 1); Dict d2 = d1; // alias const void* original_ptr = d1.get(); // Insert enough elements to trigger rehash d1.Set("b", 2); d1.Set("c", 3); d1.Set("d", 4); d1.Set("e", 5); d1.Set("f", 6); // ObjectPtr must be stable (InplaceSwitchTo) EXPECT_EQ(static_cast(d1.get()), original_ptr); // Alias must point to same object and see all elements EXPECT_EQ(static_cast(d2.get()), original_ptr); EXPECT_EQ(d2.size(), 6); EXPECT_EQ(d2["f"], 6); } TEST(Dict, ManyElements) { Dict d; for (int i = 0; i < 100; ++i) { d.Set(i, i * 10); } EXPECT_EQ(d.size(), 100); for (int i = 0; i < 100; ++i) { EXPECT_EQ(d[i], i * 10); } } TEST(Dict, Erase) { Dict d; d.Set("a", 1); d.Set("b", 2); d.Set("c", 3); d.erase("b"); EXPECT_EQ(d.size(), 2); EXPECT_EQ(d.count("b"), 0); EXPECT_EQ(d["a"], 1); EXPECT_EQ(d["c"], 3); } TEST(Dict, Clear) { Dict d; d.Set("a", 1); d.Set("b", 2); d.clear(); EXPECT_EQ(d.size(), 0); EXPECT_TRUE(d.empty()); } TEST(Dict, Iteration) { Dict d; d.Set("x", 10); d.Set("y", 20); int count = 0; for (auto [k, v] : d) { ++count; if (k == "x") { EXPECT_EQ(v, 10); } if (k == "y") { EXPECT_EQ(v, 20); } } EXPECT_EQ(count, 2); } TEST(Dict, PODKeys) { Dict d; d.Set(1, "one"); d.Set(2, "two"); EXPECT_EQ(d[1], "one"); EXPECT_EQ(d[2], "two"); } TEST(Dict, AnyConversion) { Dict d; d.Set(String("key"), 42); Any any_d = d; auto d2 = any_d.cast>(); EXPECT_EQ(d2.size(), 1); } TEST(Dict, InitializerList) { Dict d{{"a", 1}, {"b", 2}}; EXPECT_EQ(d.size(), 2); EXPECT_EQ(d["a"], 1); EXPECT_EQ(d["b"], 2); } TEST(Dict, UpdateExistingKey) { Dict d; d.Set("a", 1); d.Set("a", 2); EXPECT_EQ(d.size(), 1); EXPECT_EQ(d["a"], 2); } TEST(Dict, DefaultConstruction) { Dict d; EXPECT_EQ(d.size(), 0); EXPECT_TRUE(d.empty()); // Set on default-constructed should work d.Set("a", 1); EXPECT_EQ(d.size(), 1); } TEST(Dict, CrossConvMapToDict) { Map m{{"a", 1}, {"b", 2}}; Any any_m = m; // Cast Map to Dict via Any — triggers cross-conversion auto d = any_m.cast>(); EXPECT_EQ(d.size(), 2); EXPECT_EQ(d["a"], 1); EXPECT_EQ(d["b"], 2); } TEST(Dict, CrossConvDictToMap) { Dict d{{"x", 10}, {"y", 20}}; Any any_d = d; // Cast Dict to Map via Any — triggers cross-conversion auto m = any_d.cast>(); EXPECT_EQ(m.size(), 2); EXPECT_EQ(m["x"], 10); EXPECT_EQ(m["y"], 20); } TEST(Dict, CrossConvEmptyMapToDict) { Map m; Any any_m = m; auto d = any_m.cast>(); EXPECT_EQ(d.size(), 0); EXPECT_TRUE(d.empty()); } TEST(Dict, CrossConvEmptyDictToMap) { Dict d; Any any_d = d; auto m = any_d.cast>(); EXPECT_EQ(m.size(), 0); EXPECT_TRUE(m.empty()); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_dtype.cc000066400000000000000000000201471521067262500175740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include namespace { using namespace tvm::ffi; TEST(DType, StringConversion) { DLDataType dtype = DLDataType{kDLFloat, 32, 1}; EXPECT_EQ(DLDataTypeToString(dtype), "float32"); EXPECT_EQ(StringToDLDataType("float32"), dtype); dtype = DLDataType{kDLInt, 16, 2}; EXPECT_EQ(DLDataTypeToString(dtype), "int16x2"); EXPECT_EQ(StringToDLDataType("int16x2"), dtype); dtype = DLDataType{kDLOpaqueHandle, 0, 0}; EXPECT_EQ(DLDataTypeToString(dtype), ""); EXPECT_EQ(StringToDLDataType("void"), dtype); // test bfloat with lanes dtype = DLDataType{kDLBfloat, 16, 2}; EXPECT_EQ(DLDataTypeToString(dtype), "bfloat16x2"); EXPECT_EQ(StringToDLDataType("bfloat16x2"), dtype); // test float8 dtype = DLDataType{kDLFloat8_e4m3fn, 8, 2}; EXPECT_EQ(DLDataTypeToString(dtype), "float8_e4m3fnx2"); EXPECT_EQ(StringToDLDataType("float8_e4m3fnx2"), dtype); } TEST(DType, StringConversionAllDLPackTypes) { std::vector> test_cases = { {DLDataType{kDLFloat, 32, 1}, "float32"}, {DLDataType{kDLInt, 16, 1}, "int16"}, {DLDataType{kDLUInt, 16, 1}, "uint16"}, {DLDataType{kDLBfloat, 16, 1}, "bfloat16"}, {DLDataType{kDLFloat8_e3m4, 8, 1}, "float8_e3m4"}, {DLDataType{kDLFloat8_e4m3, 8, 1}, "float8_e4m3"}, {DLDataType{kDLFloat8_e4m3b11fnuz, 8, 1}, "float8_e4m3b11fnuz"}, {DLDataType{kDLFloat8_e4m3fn, 8, 1}, "float8_e4m3fn"}, {DLDataType{kDLFloat8_e4m3fnuz, 8, 1}, "float8_e4m3fnuz"}, {DLDataType{kDLFloat8_e5m2, 8, 1}, "float8_e5m2"}, {DLDataType{kDLFloat8_e5m2fnuz, 8, 1}, "float8_e5m2fnuz"}, {DLDataType{kDLFloat8_e8m0fnu, 8, 1}, "float8_e8m0fnu"}, {DLDataType{kDLFloat6_e2m3fn, 6, 1}, "float6_e2m3fn"}, {DLDataType{kDLFloat6_e3m2fn, 6, 1}, "float6_e3m2fn"}, {DLDataType{kDLFloat4_e2m1fn, 4, 1}, "float4_e2m1fn"}, }; for (const auto& [dtype, str] : test_cases) { EXPECT_EQ(DLDataTypeToString(dtype), str); EXPECT_EQ(StringToDLDataType(str), dtype); } } TEST(DataType, AnyConversion) { AnyView view0; EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone); Optional opt_v0 = view0.as(); EXPECT_TRUE(!opt_v0.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto v0 = view0.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `None` to `DataType`"), std::string::npos); throw; } }, ::tvm::ffi::Error); DLDataType dtype{kDLFloat, 32, 1}; AnyView view1_dtype = dtype; auto dtype_v1 = view1_dtype.cast(); EXPECT_EQ(dtype_v1.code, kDLFloat); EXPECT_EQ(dtype_v1.bits, 32); EXPECT_EQ(dtype_v1.lanes, 1); Any any2 = DLDataType{kDLInt, 16, 2}; TVMFFIAny ffi_v2 = details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(any2)); EXPECT_EQ(ffi_v2.type_index, TypeIndex::kTVMFFIDataType); EXPECT_EQ(ffi_v2.v_dtype.code, kDLInt); EXPECT_EQ(ffi_v2.v_dtype.bits, 16); EXPECT_EQ(ffi_v2.v_dtype.lanes, 2); } // String can be automatically converted to DLDataType TEST(DataType, AnyConversionWithString) { AnyView view0 = "float32"; Optional opt_v0 = view0.try_cast(); DLDataType dtype_v0 = opt_v0.value(); EXPECT_EQ(dtype_v0.code, kDLFloat); EXPECT_EQ(dtype_v0.bits, 32); EXPECT_EQ(dtype_v0.lanes, 1); Any any = String("bfloat16x2"); Optional opt_v1 = any.try_cast(); EXPECT_EQ(opt_v1.value().code, kDLBfloat); EXPECT_EQ(opt_v1.value().bits, 16); EXPECT_EQ(opt_v1.value().lanes, 2); } TEST(DType, NonNullTerminatedStringView) { // Simulate memory scenario similar to Electron where memory after string // contains garbage data (digits from previous strings) // // We test by calling TVMFFIDataTypeFromString directly with TVMFFIByteArray // to bypass String's automatic null-termination // Helper lambda to test with raw byte array (no null terminator) auto test_dtype_from_bytes = [](const char* data, size_t size) -> DLDataType { TVMFFIByteArray byte_array{data, size}; DLDataType dtype; int ret = TVMFFIDataTypeFromString(&byte_array, &dtype); EXPECT_EQ(ret, 0) << "TVMFFIDataTypeFromString failed"; return dtype; }; // Test 1: "float16" followed by digit garbage char buffer1[] = "float16999888777"; DLDataType dtype1 = test_dtype_from_bytes(buffer1, 7); // Only "float16" EXPECT_EQ(dtype1.code, kDLFloat); EXPECT_EQ(dtype1.bits, 16); // Should be 16, not 16999888777! EXPECT_EQ(dtype1.lanes, 1); // Test 2: "int32" followed by "x4" from previous leftover char buffer2[] = "int32x4extradata"; DLDataType dtype2 = test_dtype_from_bytes(buffer2, 5); // Only "int32" EXPECT_EQ(dtype2.code, kDLInt); EXPECT_EQ(dtype2.bits, 32); // Should be 32, not parse the 'x4' EXPECT_EQ(dtype2.lanes, 1); // Should be 1, not 4 // Test 3: "uint8" followed by more digits char buffer3[] = "uint8192"; DLDataType dtype3 = test_dtype_from_bytes(buffer3, 5); // Only "uint8" EXPECT_EQ(dtype3.code, kDLUInt); EXPECT_EQ(dtype3.bits, 8); // Should be 8, not 8192 EXPECT_EQ(dtype3.lanes, 1); // Test 4: "bfloat16" followed by "x2" garbage char buffer4[] = "bfloat16x2garbage"; DLDataType dtype4 = test_dtype_from_bytes(buffer4, 8); // Only "bfloat16" EXPECT_EQ(dtype4.code, kDLBfloat); EXPECT_EQ(dtype4.bits, 16); EXPECT_EQ(dtype4.lanes, 1); // Should be 1, not 2 // Test 5: "bfloat16x2" - lanes within bounds (should work) DLDataType dtype5 = test_dtype_from_bytes(buffer4, 10); // "bfloat16x2" EXPECT_EQ(dtype5.code, kDLBfloat); EXPECT_EQ(dtype5.bits, 16); EXPECT_EQ(dtype5.lanes, 2); // Should correctly parse x2 // Test 6: Truly non-null-terminated - overwrite null byte char buffer6[] = "float64AAAAA"; buffer6[7] = 'X'; // Ensure no null terminator at position 7 DLDataType dtype6 = test_dtype_from_bytes(buffer6, 7); // "float64" EXPECT_EQ(dtype6.code, kDLFloat); EXPECT_EQ(dtype6.bits, 64); EXPECT_EQ(dtype6.lanes, 1); // Test 7: "int8" followed by "x16" pattern char buffer7[] = "int8x16leftovers"; DLDataType dtype7 = test_dtype_from_bytes(buffer7, 4); // Only "int8" EXPECT_EQ(dtype7.code, kDLInt); EXPECT_EQ(dtype7.bits, 8); EXPECT_EQ(dtype7.lanes, 1); // Should be 1, not 16 // Test 8: With actual x specification that should parse DLDataType dtype8 = test_dtype_from_bytes(buffer7, 7); // "int8x16" EXPECT_EQ(dtype8.code, kDLInt); EXPECT_EQ(dtype8.bits, 8); EXPECT_EQ(dtype8.lanes, 16); // Should correctly parse x16 // Test 9: Scalable vector - "int32xvscalex4" char buffer9[] = "int32xvscalex4extra"; DLDataType dtype9 = test_dtype_from_bytes(buffer9, 14); // "int32xvscalex4" EXPECT_EQ(dtype9.code, kDLInt); EXPECT_EQ(dtype9.bits, 32); EXPECT_EQ(dtype9.lanes, static_cast(-4)); // Scalable: -4 // Test 10: Scalable vector with garbage after char buffer10[] = "float16xvscalex8999"; DLDataType dtype10 = test_dtype_from_bytes(buffer10, 16); // "float16xvscalex8" EXPECT_EQ(dtype10.code, kDLFloat); EXPECT_EQ(dtype10.bits, 16); EXPECT_EQ(dtype10.lanes, static_cast(-8)); // Should be -8, not parse "999" } } // namespace tvm-ffi-0.1.12/tests/cpp/test_error.cc000066400000000000000000000173671521067262500176120ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include namespace { using namespace tvm::ffi; void ThrowRuntimeError() { TVM_FFI_THROW(RuntimeError) << "test0"; } TEST(Error, Backtrace) { EXPECT_THROW( { try { ThrowRuntimeError(); } catch (const Error& error) { EXPECT_EQ(error.message(), "test0"); EXPECT_EQ(error.kind(), "RuntimeError"); std::string full_message = error.FullMessage(); EXPECT_NE(full_message.find("line"), std::string::npos); EXPECT_NE(full_message.find("ThrowRuntimeError"), std::string::npos); EXPECT_NE(full_message.find("RuntimeError: test0"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(CheckError, Backtrace) { EXPECT_THROW( { try { TVM_FFI_ICHECK_GT(2, 3); } catch (const Error& error) { EXPECT_EQ(error.kind(), "InternalError"); std::string full_message = error.FullMessage(); EXPECT_NE(full_message.find("line"), std::string::npos); EXPECT_NE(full_message.find("2 > 3"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(CheckError, ValueError) { int value = -5; EXPECT_THROW( { try { TVM_FFI_CHECK(value >= 0, ValueError) << "Value must be non-negative, got " << value; } catch (const Error& error) { EXPECT_EQ(error.kind(), "ValueError"); std::string full_message = error.FullMessage(); EXPECT_NE(full_message.find("line"), std::string::npos); EXPECT_NE(full_message.find("Check failed: (value >= 0) is false"), std::string::npos); EXPECT_NE(full_message.find("Value must be non-negative, got -5"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(CheckError, IndexError) { int index = 10; int array_size = 5; EXPECT_THROW( { try { TVM_FFI_CHECK(index < array_size, IndexError) << "Index " << index << " out of bounds for array of size " << array_size; } catch (const Error& error) { EXPECT_EQ(error.kind(), "IndexError"); std::string full_message = error.FullMessage(); EXPECT_NE(full_message.find("line"), std::string::npos); EXPECT_NE(full_message.find("Check failed: (index < array_size) is false"), std::string::npos); EXPECT_NE(full_message.find("Index 10 out of bounds for array of size 5"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(CheckError, PassingCondition) { // This should not throw EXPECT_NO_THROW(TVM_FFI_CHECK(true, ValueError)); EXPECT_NO_THROW(TVM_FFI_CHECK(5 < 10, IndexError)); } // Helper: expect that a throwing statement throws Error with the given kind and // a message containing the expected substring. #define EXPECT_CHECK_ERROR(stmt, expected_kind, expected_substr) \ { \ bool caught = false; \ try { \ stmt; \ } catch (const ::tvm::ffi::Error& e) { \ caught = true; \ EXPECT_EQ(e.kind(), expected_kind); \ EXPECT_NE(e.message().find(expected_substr), std::string::npos) \ << "message: " << e.message(); \ } \ EXPECT_TRUE(caught) << "Expected " #stmt " to throw"; \ } TEST(CheckError, CheckBinaryOps) { // EQ: pass and fail TVM_FFI_CHECK_EQ(3, 3, ValueError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_EQ(3, 5, ValueError), "ValueError", "3 vs. 5"); // NE: pass and fail TVM_FFI_CHECK_NE(3, 4, ValueError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_NE(3, 3, ValueError), "ValueError", "3 vs. 3"); // LT: pass and fail TVM_FFI_CHECK_LT(3, 4, IndexError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_LT(5, 3, IndexError), "IndexError", "5 vs. 3"); // GT: pass and fail TVM_FFI_CHECK_GT(4, 3, IndexError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_GT(3, 5, IndexError), "IndexError", "3 vs. 5"); // LE: pass and fail TVM_FFI_CHECK_LE(3, 3, TypeError); TVM_FFI_CHECK_LE(2, 3, TypeError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_LE(5, 3, TypeError), "TypeError", "5 vs. 3"); // GE: pass and fail TVM_FFI_CHECK_GE(4, 3, TypeError); TVM_FFI_CHECK_GE(3, 3, TypeError); EXPECT_CHECK_ERROR(TVM_FFI_CHECK_GE(2, 5, TypeError), "TypeError", "2 vs. 5"); // NOTNULL: pass and fail int x = 42; int* p = &x; EXPECT_EQ(TVM_FFI_CHECK_NOTNULL(p, ValueError), p); int* q = nullptr; EXPECT_CHECK_ERROR((void)TVM_FFI_CHECK_NOTNULL(q, ValueError), "ValueError", "Check not null"); } TEST(CheckError, DCheck) { #ifdef NDEBUG // Release: failing conditions are no-ops. TVM_FFI_DCHECK(false); TVM_FFI_DCHECK_EQ(1, 2); TVM_FFI_DCHECK_NE(1, 1); TVM_FFI_DCHECK_LT(5, 3); TVM_FFI_DCHECK_GT(3, 5); TVM_FFI_DCHECK_LE(5, 3); TVM_FFI_DCHECK_GE(3, 5); int* q = nullptr; int* r = TVM_FFI_DCHECK_NOTNULL(q); EXPECT_EQ(r, nullptr); #else // Debug: passing conditions succeed. TVM_FFI_DCHECK(true); TVM_FFI_DCHECK_EQ(3, 3); TVM_FFI_DCHECK_NE(3, 4); TVM_FFI_DCHECK_LT(3, 4); TVM_FFI_DCHECK_GT(4, 3); TVM_FFI_DCHECK_LE(3, 3); TVM_FFI_DCHECK_GE(4, 3); int x = 42; int* p = &x; EXPECT_EQ(TVM_FFI_DCHECK_NOTNULL(p), p); // Debug: failing conditions throw InternalError. EXPECT_CHECK_ERROR(TVM_FFI_DCHECK(false), "InternalError", "Check failed"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_EQ(1, 2), "InternalError", "1 vs. 2"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_NE(1, 1), "InternalError", "1 vs. 1"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_LT(5, 3), "InternalError", "5 vs. 3"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_GT(3, 5), "InternalError", "3 vs. 5"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_LE(5, 3), "InternalError", "5 vs. 3"); EXPECT_CHECK_ERROR(TVM_FFI_DCHECK_GE(3, 5), "InternalError", "3 vs. 5"); #endif } TEST(Error, AnyConvert) { Any any = Error("TypeError", "here", "test0"); Optional opt_err = any.as(); EXPECT_EQ(opt_err.value().kind(), "TypeError"); EXPECT_EQ(opt_err.value().message(), "here"); } TEST(Error, TracebackMostRecentCallLast) { Error error("TypeError", "here", "test0\ntest1\ntest2\n"); EXPECT_EQ(error.TracebackMostRecentCallLast(), "test2\ntest1\ntest0\n"); } TEST(Error, CauseChain) { Error original_error("TypeError", "here", "test0"); Error cause_chain("ValueError", "cause", "test1", original_error, std::nullopt); auto opt_cause = cause_chain.cause_chain(); EXPECT_TRUE(opt_cause.has_value()); if (opt_cause.has_value()) { EXPECT_EQ(opt_cause->kind(), "TypeError"); } EXPECT_TRUE(!cause_chain.extra_context().has_value()); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_example.cc000066400000000000000000000203421521067262500200770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include // test-cases used in example code namespace { void ExampleAny() { namespace ffi = tvm::ffi; // Create an Any from various types ffi::Any int_value = 42; ffi::Any float_value = 3.14; ffi::Any string_value = "hello world"; // AnyView provides a lightweight view without ownership ffi::AnyView view = int_value; // we can cast Any/AnyView to a specific type int extracted = view.cast(); EXPECT_EQ(extracted, 42); // If we are not sure about the type // we can use as to get an optional value std::optional maybe_int = view.as(); if (maybe_int.has_value()) { EXPECT_EQ(maybe_int.value(), 42); } // Try cast is another version that will try to run the type // conversion even if the type does not exactly match std::optional maybe_int_try = view.try_cast(); if (maybe_int_try.has_value()) { EXPECT_EQ(maybe_int_try.value(), 42); } } TEST(Example, Any) { ExampleAny(); } void ExampleFunctionFromPacked() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fadd1 = ffi::Function::FromPacked([](const ffi::AnyView* args, int32_t num_args, ffi::Any* rv) { TVM_FFI_ICHECK_EQ(num_args, 1); int a = args[0].cast(); *rv = a + 1; }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); } void ExampleFunctionFromTyped() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fadd1 = ffi::Function::FromTyped([](const int a) -> int { return a + 1; }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); } void ExampleFunctionPassFunction() { namespace ffi = tvm::ffi; // Create a function from a typed lambda ffi::Function fapply = ffi::Function::FromTyped( // NOLINTNEXTLINE(performance-unnecessary-value-param) [](const ffi::Function f, ffi::Any param) { return f(param.cast()); }); ffi::Function fadd1 = ffi::Function::FromTyped( // [](const int a) -> int { return a + 1; }); int b = fapply(fadd1, 2).cast(); EXPECT_EQ(b, 3); } void ExamplegGlobalFunctionRegistry() { namespace ffi = tvm::ffi; ffi::reflection::GlobalDef().def("xyz.add1", [](const int a) -> int { return a + 1; }); ffi::Function fadd1 = ffi::Function::GetGlobalRequired("xyz.add1"); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); } void FuncThrowError() { namespace ffi = tvm::ffi; TVM_FFI_THROW(TypeError) << "test0"; } void ExampleErrorHandling() { namespace ffi = tvm::ffi; try { FuncThrowError(); } catch (const ffi::Error& e) { EXPECT_EQ(e.kind(), "TypeError"); EXPECT_EQ(e.message(), "test0"); std::cout << e.TracebackMostRecentCallLast() << std::endl; } } TEST(Example, Function) { ExampleFunctionFromPacked(); ExampleFunctionFromTyped(); ExampleFunctionPassFunction(); ExamplegGlobalFunctionRegistry(); ExampleErrorHandling(); } struct CPUNDAlloc { void AllocData(DLTensor* tensor) { tensor->data = malloc(tvm::ffi::GetDataSize(*tensor)); } void FreeData(DLTensor* tensor) { free(tensor->data); } }; void ExampleTensor() { namespace ffi = tvm::ffi; ffi::Shape shape = {1, 2, 3}; DLDataType dtype = {kDLFloat, 32, 1}; DLDevice device = {kDLCPU, 0}; ffi::Tensor tensor = ffi::Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); } void ExampleTensorDLPack() { namespace ffi = tvm::ffi; ffi::Shape shape = {1, 2, 3}; DLDataType dtype = {kDLFloat, 32, 1}; DLDevice device = {kDLCPU, 0}; ffi::Tensor tensor = ffi::Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); // convert to DLManagedTensorVersioned DLManagedTensorVersioned* dlpack = tensor.ToDLPackVersioned(); // load back from DLManagedTensorVersioned ffi::Tensor tensor2 = ffi::Tensor::FromDLPackVersioned(dlpack); } TEST(Example, Tensor) { ExampleTensor(); ExampleTensorDLPack(); } void ExampleString() { namespace ffi = tvm::ffi; ffi::String str = "hello world"; EXPECT_EQ(str.size(), 11); std::string std_str = str; EXPECT_EQ(std_str, "hello world"); } TEST(Example, String) { ExampleString(); } void ExampleArray() { namespace ffi = tvm::ffi; ffi::Array numbers = {1, 2, 3}; EXPECT_EQ(numbers.size(), 3); EXPECT_EQ(numbers[0], 1); // NOLINTNEXTLINE(performance-unnecessary-value-param) ffi::Function head = ffi::Function::FromTyped([](const ffi::Array a) { return a[0]; }); EXPECT_EQ(head(numbers).cast(), 1); try { // throw an error because 2.2 is not int head(ffi::Array({1, 2.2})); } catch (const ffi::Error& e) { EXPECT_EQ(e.kind(), "TypeError"); } } void ExampleTuple() { namespace ffi = tvm::ffi; ffi::Tuple tup(42, "hello", true); EXPECT_EQ(tup.get<0>(), 42); EXPECT_EQ(tup.get<1>(), "hello"); EXPECT_EQ(tup.get<2>(), true); } TEST(Example, Array) { ExampleArray(); ExampleTuple(); } void ExampleMap() { namespace ffi = tvm::ffi; ffi::Map map0 = {{"Alice", 100}, {"Bob", 95}}; EXPECT_EQ(map0.size(), 2); EXPECT_EQ(map0.at("Alice"), 100); EXPECT_EQ(map0.count("Alice"), 1); } TEST(Example, Map) { ExampleMap(); } void ExampleOptional() { namespace ffi = tvm::ffi; ffi::Optional opt0 = 100; EXPECT_EQ(opt0.has_value(), true); EXPECT_EQ(opt0.value(), 100); ffi::Optional opt1; EXPECT_EQ(opt1.has_value(), false); EXPECT_EQ(opt1.value_or("default"), "default"); } TEST(Example, Optional) { ExampleOptional(); } void ExampleVariant() { namespace ffi = tvm::ffi; ffi::Variant var0 = 100; EXPECT_EQ(var0.get(), 100); var0 = ffi::String("hello"); std::optional maybe_str = var0.as(); EXPECT_EQ(maybe_str.value(), "hello"); // NOLINT(bugprone-unchecked-optional-access) std::optional maybe_int2 = var0.as(); EXPECT_EQ(maybe_int2.has_value(), false); } TEST(Example, Variant) { ExampleVariant(); } // Step 1: Define the object class (stores the actual data) class MyIntPairObj : public tvm::ffi::Object { public: int64_t a; int64_t b; MyIntPairObj() = default; MyIntPairObj(int64_t a, int64_t b) : a(a), b(b) {} // Required: declare type information TVM_FFI_DECLARE_OBJECT_INFO_FINAL("example.MyIntPair", MyIntPairObj, tvm::ffi::Object); }; // Step 2: Define the reference wrapper (user-facing interface) class MyIntPair : public tvm::ffi::ObjectRef { public: // Constructor explicit MyIntPair(int64_t a, int64_t b) { data_ = tvm::ffi::make_object(a, b); } // Required: define object reference methods TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MyIntPair, tvm::ffi::ObjectRef, MyIntPairObj); }; void ExampleObjectPtr() { namespace ffi = tvm::ffi; ffi::ObjectPtr obj = ffi::make_object(100, 200); EXPECT_EQ(obj->a, 100); EXPECT_EQ(obj->b, 200); } void ExampleObjectRef() { namespace ffi = tvm::ffi; MyIntPair pair(100, 200); EXPECT_EQ(pair->a, 100); EXPECT_EQ(pair->b, 200); } void ExampleObjectRefAny() { namespace ffi = tvm::ffi; MyIntPair pair(100, 200); ffi::Any any = pair; MyIntPair pair2 = any.cast(); EXPECT_EQ(pair2->a, 100); EXPECT_EQ(pair2->b, 200); } TEST(Example, ObjectPtr) { ExampleObjectPtr(); ExampleObjectRef(); ExampleObjectRefAny(); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_expected.cc000066400000000000000000000273411521067262500202530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; // Test implicit construction from success value TEST(Expected, BasicOk) { Expected result = 42; EXPECT_TRUE(result.is_ok()); EXPECT_FALSE(result.is_err()); EXPECT_TRUE(result.has_value()); EXPECT_EQ(result.value(), 42); EXPECT_EQ(result.value_or(0), 42); } // Test implicit construction from error TEST(Expected, BasicErr) { Expected result = Error("RuntimeError", "test error", ""); EXPECT_FALSE(result.is_ok()); EXPECT_TRUE(result.is_err()); EXPECT_FALSE(result.has_value()); Error err = result.error(); EXPECT_EQ(err.kind(), "RuntimeError"); EXPECT_EQ(err.message(), "test error"); } // Test explicit construction via Unexpected wrapper TEST(Expected, UnexpectedWrapper) { Expected result = Unexpected(Error("RuntimeError", "unexpected error", "")); EXPECT_FALSE(result.is_ok()); EXPECT_TRUE(result.is_err()); EXPECT_EQ(result.error().kind(), "RuntimeError"); EXPECT_EQ(result.error().message(), "unexpected error"); } // Test Unexpected deduction guide and error() accessor TEST(Expected, UnexpectedAPI) { auto u = Unexpected(Error("TypeError", "type mismatch", "")); static_assert(std::is_same_v>); EXPECT_EQ(u.error().kind(), "TypeError"); EXPECT_EQ(u.error().message(), "type mismatch"); Error moved_err = std::move(u).error(); EXPECT_EQ(moved_err.kind(), "TypeError"); } // Test value_or with error TEST(Expected, ValueOrWithError) { Expected result = Error("RuntimeError", "test error", ""); EXPECT_EQ(result.value_or(99), 99); } // Test with ObjectRef types TEST(Expected, ObjectRefType) { Expected result = TInt(123); EXPECT_TRUE(result.is_ok()); EXPECT_EQ(result.value()->value, 123); } // Test with String type TEST(Expected, StringType) { Expected result = String("hello"); EXPECT_TRUE(result.is_ok()); EXPECT_EQ(result.value(), "hello"); Expected err_result = Error("ValueError", "bad string", ""); EXPECT_TRUE(err_result.is_err()); } // Test TypeTraits conversion: Expected -> Any -> Expected TEST(Expected, TypeTraitsRoundtrip) { Expected original = 42; // Convert to Any (should unwrap to int) Any any_value = original; EXPECT_EQ(any_value.cast(), 42); // Convert back to Expected (should reconstruct as Ok) Expected recovered = any_value.cast>(); EXPECT_TRUE(recovered.is_ok()); EXPECT_EQ(recovered.value(), 42); } // Test TypeTraits conversion with Error TEST(Expected, TypeTraitsErrorRoundtrip) { Expected original = Error("TypeError", "conversion failed", ""); // Convert to Any (should unwrap to Error) Any any_value = original; EXPECT_TRUE(any_value.as().has_value()); // Convert back to Expected (should reconstruct as Err) Expected recovered = any_value.cast>(); EXPECT_TRUE(recovered.is_err()); EXPECT_EQ(recovered.error().kind(), "TypeError"); } // Test move semantics TEST(Expected, MoveSemantics) { Expected result = String("test"); EXPECT_TRUE(result.is_ok()); String value = std::move(result).value(); EXPECT_EQ(value, "test"); } // Test CallExpected with normal function TEST(Expected, CallExpectedNormal) { auto safe_add = [](int a, int b) { return a + b; }; Function func = Function::FromTyped(safe_add); Expected result = func.CallExpected(5, 3); EXPECT_TRUE(result.is_ok()); EXPECT_EQ(result.value(), 8); } // Test CallExpected with throwing function TEST(Expected, CallExpectedThrowing) { auto throwing_func = [](int a) -> int { if (a < 0) { TVM_FFI_THROW(ValueError) << "Negative value not allowed"; } return a * 2; }; Function func = Function::FromTyped(throwing_func); // Normal case Expected result_ok = func.CallExpected(5); EXPECT_TRUE(result_ok.is_ok()); EXPECT_EQ(result_ok.value(), 10); // Error case Expected result_err = func.CallExpected(-1); EXPECT_TRUE(result_err.is_err()); EXPECT_EQ(result_err.error().kind(), "ValueError"); } // Test that lambda returning Expected works directly TEST(Expected, LambdaDirectCall) { auto safe_divide = [](int a, int b) -> Expected { if (b == 0) { return Error("ValueError", "Division by zero", ""); } return a / b; }; // Direct call to lambda should work Expected result = safe_divide(10, 2); EXPECT_TRUE(result.is_ok()); EXPECT_EQ(result.value(), 5); // Check the value can be extracted int val = result.value(); EXPECT_EQ(val, 5); // Check assigning to Any works Any any_val = result.value(); EXPECT_EQ(any_val.cast(), 5); } // Test registering function that returns Expected TEST(Expected, RegisterExpectedReturning) { auto safe_divide = [](int a, int b) -> Expected { if (b == 0) { return Error("ValueError", "Division by zero", ""); } return a / b; }; // Verify the FunctionInfo extracts Expected as return type using FuncInfo = tvm::ffi::details::FunctionInfo; static_assert(std::is_same_v>, "Return type should be Expected"); Function::SetGlobal("test.safe_divide3", Function::FromTyped(safe_divide)); Function func = Function::GetGlobalRequired("test.safe_divide3"); // Normal call should throw when function returns Err EXPECT_THROW({ func(10, 0).cast(); }, Error); // Normal call should succeed when function returns Ok int result = func(10, 2).cast(); EXPECT_EQ(result, 5); // CallExpected should return Expected Expected exp_ok = func.CallExpected(10, 2); EXPECT_TRUE(exp_ok.is_ok()); EXPECT_EQ(exp_ok.value(), 5); Expected exp_err = func.CallExpected(10, 0); EXPECT_TRUE(exp_err.is_err()); EXPECT_EQ(exp_err.error().message(), "Division by zero"); } // Test Expected with Optional (nested types) TEST(Expected, NestedOptional) { Expected> result = Optional(42); EXPECT_TRUE(result.is_ok()); EXPECT_TRUE(result.value().has_value()); EXPECT_EQ(result.value().value(), 42); Expected> result_none = Optional(std::nullopt); EXPECT_TRUE(result_none.is_ok()); EXPECT_FALSE(result_none.value().has_value()); } // Test Expected with Array TEST(Expected, ArrayType) { Array arr{1, 2, 3}; Expected> result = arr; EXPECT_TRUE(result.is_ok()); EXPECT_EQ(result.value().size(), 3); EXPECT_EQ(result.value()[0], 1); } // Test complex example: function returning Expected> TEST(Expected, ComplexExample) { auto parse_csv = [](const String& input) -> Expected> { if (input.size() == 0) { return Error("ValueError", "Empty input", ""); } // Simple split by comma Array result; result.push_back(input); // Simplified for test return result; }; Function::SetGlobal("test.parse_csv", Function::FromTyped(parse_csv)); Function func = Function::GetGlobalRequired("test.parse_csv"); Expected> result_ok = func.CallExpected>(String("a,b,c")); EXPECT_TRUE(result_ok.is_ok()); Expected> result_err = func.CallExpected>(String("")); EXPECT_TRUE(result_err.is_err()); EXPECT_EQ(result_err.error().message(), "Empty input"); } // Test bad access throws the original error TEST(Expected, BadAccessThrowsOriginalError) { Expected result = Error("CustomError", "original message", ""); try { result.value(); FAIL() << "Expected Error to be thrown"; } catch (const Error& e) { // Verify the original error is preserved EXPECT_EQ(e.kind(), "CustomError"); EXPECT_EQ(e.message(), "original message"); } } // Test error() throws RuntimeError on bad access TEST(Expected, ErrorBadAccessThrows) { Expected result = 42; EXPECT_THROW({ result.error(); }, Error); } // Test rvalue overload for value() TEST(Expected, RvalueValueAccess) { auto get_expected = []() -> Expected { return String("rvalue test"); }; // Call value() on rvalue String val = get_expected().value(); EXPECT_EQ(val, "rvalue test"); } // Test rvalue overload for error() TEST(Expected, RvalueErrorAccess) { auto get_expected = []() -> Expected { return Error("TestError", "rvalue error", ""); }; // Call error() on rvalue Error err = get_expected().error(); EXPECT_EQ(err.kind(), "TestError"); EXPECT_EQ(err.message(), "rvalue error"); } // Test Expected with inheritance (Error is a subclass of ObjectRef) // This ensures is_ok() and is_err() work correctly for ObjectRef types TEST(Expected, ObjectRefInheritance) { // Expected with an actual ObjectRef value ObjectRef obj = TInt(100); Expected result_ok(obj); EXPECT_TRUE(result_ok.is_ok()); EXPECT_FALSE(result_ok.is_err()); EXPECT_TRUE(result_ok.value().defined()); // Expected with an error Expected result_err = Error("TestError", "test", ""); EXPECT_FALSE(result_err.is_ok()); EXPECT_TRUE(result_err.is_err()); EXPECT_EQ(result_err.error().kind(), "TestError"); } // Test TryCastFromAnyView with incompatible type TEST(Expected, TryCastIncompatible) { Any any_str = String("hello"); auto result = any_str.try_cast>(); EXPECT_FALSE(result.has_value()); // Cannot convert String to Expected } // Test that Expected::value() && compiles and runs correctly. // Requires TypeTraits::MoveFromAnyAfterCheck to be defined. TEST(ExpectedRvalueMove, DLDataTypeMoveCompiles) { Expected e = DLDataType{kDLFloat, 32, 1}; DLDataType moved = std::move(e).value(); EXPECT_EQ(moved.code, kDLFloat); EXPECT_EQ(moved.bits, 32); EXPECT_EQ(moved.lanes, 1); } // Test that value_or() && moves rather than copies for Object types. TEST(ExpectedRvalueMove, ValueOrMovesNotCopies) { Expected e = String("hello"); String moved = std::move(e).value_or(String("default")); EXPECT_EQ(moved, "hello"); } // Test value_or() && on error path returns default. TEST(ExpectedRvalueMove, ValueOrRvalueErrorPath) { Expected e = Error("ValueError", "oops", ""); String result = std::move(e).value_or(String("fallback")); EXPECT_EQ(result, "fallback"); } // Test POD types compile and run correctly with rvalue value(). TEST(ExpectedRvalueMove, PodTypesCompile) { EXPECT_EQ(std::move(Expected(42)).value(), 42); EXPECT_EQ(std::move(Expected(3.14)).value(), 3.14); EXPECT_EQ(std::move(Expected(true)).value(), true); DLDataType dtype{kDLInt, 64, 1}; EXPECT_EQ(std::move(Expected(dtype)).value().code, kDLInt); EXPECT_EQ(std::move(Expected(dtype)).value().bits, 64); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_function.cc000066400000000000000000000237371521067262500203040ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1 #include #include #include #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Func, FromPacked) { Function fadd1 = Function::FromPacked([](const AnyView* args, int32_t num_args, Any* rv) { EXPECT_EQ(num_args, 1); int32_t a = args[0].cast(); *rv = a + 1; }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); Function fadd2 = Function::FromPacked([](const AnyView* args, int32_t num_args, Any* rv) { EXPECT_EQ(num_args, 1); auto a = args[0].cast(); EXPECT_EQ(a.use_count(), 2); *rv = a->value + 1; }); EXPECT_EQ(fadd2(TInt(12)).cast(), 13); } TEST(Func, PackedArgs) { Function fadd1 = Function::FromPacked([](PackedArgs args, Any* rv) { EXPECT_EQ(args.size(), 1); int32_t a = args[0].cast(); *rv = a + 1; }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); Function fadd2 = Function::FromPacked([](PackedArgs args, Any* rv) { EXPECT_EQ(args.size(), 1); TInt a = args[0].cast(); EXPECT_EQ(a.use_count(), 2); *rv = a->value + 1; }); EXPECT_EQ(fadd2(TInt(12)).cast(), 13); TInt v(12); AnyView data[3]; PackedArgs::Fill(data, 3, 1, v); EXPECT_EQ(data[0].cast(), 3); EXPECT_EQ(data[1].cast(), 1); EXPECT_EQ(data[2].cast()->value, 12); } TEST(Func, FromTyped) { // try decution Function fadd1 = Function::FromTyped([](const int32_t& a) -> int { return a + 1; }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); // convert that triggers error EXPECT_THROW( { try { fadd1(1.1); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: int) -> int`. " "Expected `int` but got `float`"); throw; } }, ::tvm::ffi::Error); // convert that triggers error EXPECT_THROW( { try { fadd1(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched number of arguments when calling: `(0: int) -> int`. " "Expected 1 but got 0 arguments"); throw; } }, ::tvm::ffi::Error); // convert with DLTensor* triggers error EXPECT_THROW( { try { DLTensor dltensor; fadd1(&dltensor); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: int) -> int`. " "Expected `int` but got `DLTensor*`"); throw; } }, ::tvm::ffi::Error); // try decution Function fpass_and_return = Function::FromTyped( // NOLINTNEXTLINE(performance-unnecessary-value-param) [](TInt x, int value, AnyView z) -> Function { EXPECT_EQ(x.use_count(), 2); EXPECT_EQ(x->value, value); if (auto opt = z.as()) { EXPECT_EQ(value, *opt); } return Function::FromTyped([value](int x) -> int { return x + value; }); }, "fpass_and_return"); TInt a(11); auto fret = fpass_and_return(std::move(a), 11, 11).cast(); EXPECT_EQ(fret(12).cast(), 23); EXPECT_THROW( { try { fpass_and_return(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched number of arguments when calling: " "`fpass_and_return(0: test.Int, 1: int, 2: AnyView) -> ffi.Function`. " "Expected 3 but got 0 arguments"); throw; } }, ::tvm::ffi::Error); Function fconcact = Function::FromTyped([](const String& a, const String& b) -> String { return a + b; }); EXPECT_EQ(fconcact("abc", "def").cast(), "abcdef"); } TEST(Func, PassReturnAny) { // NOLINTNEXTLINE(performance-unnecessary-value-param) Function fadd_one = Function::FromTyped([](Any a) -> Any { return a.cast() + 1; }); EXPECT_EQ(fadd_one(1).cast(), 2); } TEST(Func, Global) { Function::SetGlobal("testing.add1", Function::FromTyped([](const int32_t& a) -> int { return a + 1; })); auto fadd1 = Function::GetGlobalRequired("testing.add1"); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); auto fnot_exist = Function::GetGlobal("testing.not_existing_func"); EXPECT_TRUE(!fnot_exist); auto fname_functor = // NOLINTNEXTLINE(bugprone-unchecked-optional-access) Function::GetGlobal("ffi.FunctionListGlobalNamesFunctor").value()().cast(); Array names; int len = fname_functor(-1).cast(); for (int i = 0; i < len; ++i) { names.push_back(fname_functor(i).cast()); } EXPECT_TRUE(std::find(names.begin(), names.end(), "testing.add1") != names.end()); } TEST(Func, TypedFunction) { TypedFunction fadd1 = [](int a) -> int { return a + 1; }; EXPECT_EQ(fadd1(1), 2); TypedFunction fadd2([](int a) -> int { return a + 2; }); EXPECT_EQ(fadd2(1), 3); EXPECT_EQ(fadd2.packed()(1).cast(), 3); TypedFunction fcheck_int; EXPECT_TRUE(fcheck_int == nullptr); fcheck_int = [](int a) -> void { EXPECT_EQ(a, 1); }; fcheck_int(1); } TEST(Func, TypedFunctionAsAny) { TypedFunction fadd1 = [](int a) -> int { return a + 1; }; Any fany(std::move(fadd1)); EXPECT_TRUE(fadd1 == nullptr); // NOLINT(bugprone-use-after-move,clang-analyzer-cplusplus.Move) auto fadd1_dup = fany.cast>(); EXPECT_EQ(fadd1_dup(1), 2); } TEST(Func, TypedFunctionAsAnyView) { TypedFunction fadd2 = [](int a) -> int { return a + 2; }; AnyView fview(fadd2); auto fadd2_dup = fview.cast>(); EXPECT_EQ(fadd2_dup(1), 3); } TEST(Func, ObjectRefWithFallbackTraits) { // test cases to test automatic type conversion via ObjectRefWithFallbackTraits // through TPrimExpr Function freturn_primexpr = Function::FromTyped([](TPrimExpr a) -> TPrimExpr { return a; }); auto result_int = freturn_primexpr(1).cast(); EXPECT_EQ(result_int->dtype, "int64"); EXPECT_EQ(result_int->value, 1); // Test case for float auto result_float = freturn_primexpr(2.5).cast(); EXPECT_EQ(result_float->dtype, "float32"); EXPECT_EQ(result_float->value, 2.5); // Test case for bool auto result_bool = freturn_primexpr(true).cast(); EXPECT_EQ(result_bool->dtype, "bool"); EXPECT_EQ(result_bool->value, 1); // Test case for string auto result_string = freturn_primexpr("test_string").cast(); EXPECT_EQ(result_string->dtype, "test_string"); EXPECT_EQ(result_string->value, 0); EXPECT_THROW( { try { freturn_primexpr(TInt(1)); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ( error.message(), "Mismatched type on argument #0 when calling: `(0: test.PrimExpr) -> test.PrimExpr`. " "Expected `test.PrimExpr` but got `test.Int`"); throw; } }, ::tvm::ffi::Error); } TEST(SetRaisedFromCStr, ValueError) { TVMFFIErrorSetRaisedFromCStr("ValueError", "Value must be non-negative, got -5"); Error error = tvm::ffi::details::MoveFromSafeCallRaised(); EXPECT_EQ(error.kind(), "ValueError"); EXPECT_EQ(error.message(), "Value must be non-negative, got -5"); } TEST(SetRaisedFromCStrParts, TypeError) { const char* message_parts[] = {"Mismatched", nullptr, " got Tensor"}; TVMFFIErrorSetRaisedFromCStrParts("TypeError", message_parts, 3); Error error = tvm::ffi::details::MoveFromSafeCallRaised(); EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched got Tensor"); } int testing_add1(int x) { return x + 1; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(testing_add1, testing_add1) TEST(Func, FromExternC) { // this is the function abi convention Function fadd1 = Function::FromExternC(nullptr, __tvm_ffi_testing_add1, nullptr); EXPECT_EQ(fadd1(1).cast(), 2); } int invoke_testing_add1(int x) { return Function::InvokeExternC(nullptr, __tvm_ffi_testing_add1, x).cast(); } TEST(Func, InvokeExternC) { EXPECT_EQ(invoke_testing_add1(1), 2); } extern "C" int __tvm_ffi__metadata_testing_add1(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); TEST(Func, StaticLinkingMetadata) { String metadata_str = Function::InvokeExternC(nullptr, __tvm_ffi__metadata_testing_add1).cast(); Map metadata = json::Parse(metadata_str).cast>(); EXPECT_TRUE(metadata.count("type_schema")); std::string type_schema_str = metadata["type_schema"].cast(); EXPECT_TRUE(type_schema_str.find("int") != std::string::npos); } extern "C" TVM_FFI_DLL int TVMFFITestingDummyTarget(); TEST(Func, DummyCFunc) { int value = TVMFFITestingDummyTarget(); EXPECT_EQ(value, 0); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_init_once.cc000066400000000000000000000215531521067262500204200ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include namespace { // Helper functions for TVMFFIHandleInitDeinitOnce test static int InitSuccess(void** handle_addr) { *handle_addr = new int(42); return 0; } static int InitShouldNotBeCalled(void** handle_addr) { *handle_addr = new int(999); return 0; } static int DeinitSuccess(void* h) { delete static_cast(h); return 0; } static int DeinitShouldNotBeCalled(void* h) { // Should not be called when handle is already null return -1; } static int InitWithError(void** handle_addr) { TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Initialization failed"); return -1; } static int InitReturnsNull(void** handle_addr) { *handle_addr = nullptr; // Invalid: must return non-null handle return 0; } static int InitForDeinitError(void** handle_addr) { *handle_addr = new int(100); return 0; } static int DeinitWithError(void* h) { delete static_cast(h); TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Deinitialization failed"); return -1; } static int InitValue123(void** handle_addr) { *handle_addr = new int(123); return 0; } static int InitValue456(void** handle_addr) { *handle_addr = new int(456); return 0; } TEST(CEnvAPI, TVMFFIHandleInitDeinitOnce) { // Test 1: Successful initialization void* handle = nullptr; int ret = TVMFFIHandleInitOnce(&handle, InitSuccess); EXPECT_EQ(ret, 0); EXPECT_NE(handle, nullptr); EXPECT_EQ(*static_cast(handle), 42); // Test 2: Multiple init calls should not re-initialize (idempotent) void* original_handle = handle; ret = TVMFFIHandleInitOnce(&handle, InitShouldNotBeCalled); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, original_handle); // Handle should remain unchanged EXPECT_EQ(*static_cast(handle), 42); // Value should still be 42 // Test 3: Successful deinitialization ret = TVMFFIHandleDeinitOnce(&handle, DeinitSuccess); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, nullptr); // Test 4: Multiple deinit calls should be safe (idempotent) ret = TVMFFIHandleDeinitOnce(&handle, DeinitShouldNotBeCalled); EXPECT_EQ(ret, 0); EXPECT_EQ(handle, nullptr); // Test 5: Init error - init_func returns error code void* handle2 = nullptr; ret = TVMFFIHandleInitOnce(&handle2, InitWithError); EXPECT_NE(ret, 0); EXPECT_EQ(handle2, nullptr); // Test 6: Init error - init_func returns nullptr (invalid) void* handle3 = nullptr; ret = TVMFFIHandleInitOnce(&handle3, InitReturnsNull); EXPECT_NE(ret, 0); EXPECT_EQ(handle3, nullptr); // Test 7: Deinit error - deinit_func returns error void* handle4 = nullptr; ret = TVMFFIHandleInitOnce(&handle4, InitForDeinitError); EXPECT_EQ(ret, 0); EXPECT_NE(handle4, nullptr); ret = TVMFFIHandleDeinitOnce(&handle4, DeinitWithError); EXPECT_NE(ret, 0); EXPECT_EQ(handle4, nullptr); // Handle should still be set to nullptr // Test 8: Init-deinit lifecycle void* handle5 = nullptr; ret = TVMFFIHandleInitOnce(&handle5, InitValue123); EXPECT_EQ(ret, 0); EXPECT_NE(handle5, nullptr); EXPECT_EQ(*static_cast(handle5), 123); ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess); EXPECT_EQ(ret, 0); EXPECT_EQ(handle5, nullptr); // Test 9: Ensure subsequent init after deinit works ret = TVMFFIHandleInitOnce(&handle5, InitValue456); EXPECT_EQ(ret, 0); EXPECT_NE(handle5, nullptr); EXPECT_EQ(*static_cast(handle5), 456); // Clean up ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess); EXPECT_EQ(ret, 0); } // Helper functions and data for multithreaded test struct ThreadSafeCounter { int value; std::atomic* init_count_ptr; std::atomic* deinit_count_ptr; ThreadSafeCounter(int v, std::atomic* init_ptr, std::atomic* deinit_ptr) : value(v), init_count_ptr(init_ptr), deinit_count_ptr(deinit_ptr) {} }; // Global pointers for the current test counters static std::atomic* g_init_count = nullptr; static std::atomic* g_deinit_count = nullptr; static int InitWithCounter(void** handle_addr) { auto* counter = new ThreadSafeCounter(42, g_init_count, g_deinit_count); if (counter->init_count_ptr) { counter->init_count_ptr->fetch_add(1, std::memory_order_relaxed); } // Small delay to increase the race window std::this_thread::sleep_for(std::chrono::microseconds(100)); *handle_addr = counter; return 0; } static int DeinitWithCounter(void* h) { auto* counter = static_cast(h); if (counter->deinit_count_ptr) { counter->deinit_count_ptr->fetch_add(1, std::memory_order_relaxed); } // Small delay to increase the race window std::this_thread::sleep_for(std::chrono::microseconds(100)); delete counter; return 0; } TEST(CEnvAPI, TVMFFIHandleInitDeinitOnceMultithreaded) { // Test 1: Multiple threads calling InitOnce - should initialize only once { void* handle = nullptr; const int num_threads = 4; std::vector threads; threads.reserve(num_threads); std::vector results(num_threads); std::mutex mtx; std::condition_variable cv; bool ready = false; std::atomic init_count{0}; // Set global counter pointers g_init_count = &init_count; g_deinit_count = nullptr; // Create threads that all try to initialize simultaneously for (int i = 0; i < num_threads; ++i) { threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() { // Wait for all threads to be ready std::unique_lock lock(mtx); cv.wait(lock, [&ready] { return ready; }); lock.unlock(); results[i] = TVMFFIHandleInitOnce(&handle, InitWithCounter); }); } // Signal all threads to start { std::scoped_lock lock(mtx); ready = true; } cv.notify_all(); // Wait for all threads to complete for (auto& t : threads) { t.join(); } // All threads should succeed for (int i = 0; i < num_threads; ++i) { EXPECT_EQ(results[i], 0); } // Handle should be initialized EXPECT_NE(handle, nullptr); auto* counter = static_cast(handle); EXPECT_EQ(counter->value, 42); // Init should have been called exactly once EXPECT_EQ(init_count.load(), 1); // Clean up int ret = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter); EXPECT_EQ(ret, 0); // Reset global pointers g_init_count = nullptr; } // Test 2: Multiple threads calling DeinitOnce - should deinitialize only once { void* handle = nullptr; std::atomic init_count{0}; std::atomic deinit_count{0}; // Set global counter pointers g_init_count = &init_count; g_deinit_count = &deinit_count; // Initialize first int ret = TVMFFIHandleInitOnce(&handle, InitWithCounter); EXPECT_EQ(ret, 0); EXPECT_NE(handle, nullptr); const int num_threads = 4; std::vector threads; threads.reserve(num_threads); std::vector results(num_threads); std::mutex mtx; std::condition_variable cv; bool ready = false; // Create threads that all try to deinitialize simultaneously for (int i = 0; i < num_threads; ++i) { threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() { // Wait for all threads to be ready std::unique_lock lock(mtx); cv.wait(lock, [&ready] { return ready; }); lock.unlock(); results[i] = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter); }); } // Signal all threads to start { std::scoped_lock lock(mtx); ready = true; } cv.notify_all(); // Wait for all threads to complete for (auto& t : threads) { t.join(); } // All threads should succeed for (int i = 0; i < num_threads; ++i) { EXPECT_EQ(results[i], 0); } // Handle should be null EXPECT_EQ(handle, nullptr); // Deinit should have been called exactly once EXPECT_EQ(deinit_count.load(), 1); // Reset global pointers g_init_count = nullptr; g_deinit_count = nullptr; } } } // namespace tvm-ffi-0.1.12/tests/cpp/test_list.cc000066400000000000000000000160171521067262500174230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include namespace { using namespace tvm::ffi; TEST(List, Basic) { List list = {11, 12}; EXPECT_EQ(list.size(), 2U); EXPECT_EQ(list[0], 11); EXPECT_EQ(list[1], 12); } TEST(List, SharedMutation) { List list = {1, 2}; List alias = list; EXPECT_TRUE(list.same_as(alias)); list.Set(1, 3); EXPECT_EQ(alias[1], 3); alias.push_back(4); EXPECT_EQ(list.size(), 3U); EXPECT_EQ(list[2], 4); } TEST(List, AssignmentOperators) { List a = {1, 2}; List b; b = a; EXPECT_TRUE(a.same_as(b)); b.Set(0, 5); EXPECT_EQ(a[0], 5); List c; c = std::move(b); EXPECT_TRUE(c.same_as(a)); List d; d = c; EXPECT_EQ(d.size(), c.size()); } TEST(List, PushPopInsertErase) { List list; std::vector vector; for (int i = 0; i < 10; ++i) { list.push_back(i); vector.push_back(i); } EXPECT_EQ(list.size(), vector.size()); for (size_t i = 0; i < vector.size(); ++i) { EXPECT_EQ(list[static_cast(i)], vector[i]); } list.insert(list.begin() + 5, 100); vector.insert(vector.begin() + 5, 100); EXPECT_EQ(list.size(), vector.size()); for (size_t i = 0; i < vector.size(); ++i) { EXPECT_EQ(list[static_cast(i)], vector[i]); } list.erase(list.begin() + 3, list.begin() + 7); vector.erase(vector.begin() + 3, vector.begin() + 7); EXPECT_EQ(list.size(), vector.size()); for (size_t i = 0; i < vector.size(); ++i) { EXPECT_EQ(list[static_cast(i)], vector[i]); } list.pop_back(); vector.pop_back(); EXPECT_EQ(list.size(), vector.size()); for (size_t i = 0; i < vector.size(); ++i) { EXPECT_EQ(list[static_cast(i)], vector[i]); } } TEST(List, ReserveReallocationPreservesValues) { List list; for (int i = 0; i < 8; ++i) { list.push_back(i); } auto* before_obj = list.GetListObj(); size_t before_capacity = list.capacity(); int64_t reserve_target = static_cast(before_capacity) + 32; list.reserve(reserve_target); auto* after_obj = list.GetListObj(); EXPECT_EQ(before_obj, after_obj); EXPECT_GE(list.capacity(), before_capacity + 32); for (int i = 0; i < 8; ++i) { EXPECT_EQ(list[i], i); } list.reserve(1); EXPECT_GE(list.capacity(), before_capacity + 32); } TEST(List, AnyImplicitConversionFromArray) { Array array = {1, 2.5}; AnyView array_view = array; List list = array_view.cast>(); EXPECT_EQ(list.size(), 2U); EXPECT_EQ(list[0], 1.0); EXPECT_EQ(list[1], 2.5); EXPECT_FALSE(list.same_as(array)); list.Set(0, 99.0); EXPECT_EQ(array[0].cast(), 1); List list_any = {1, 2}; AnyView list_view = list_any; List list_any_roundtrip = list_view.cast>(); EXPECT_TRUE(list_any_roundtrip.same_as(list_any)); } TEST(List, AnyConvertCheck) { Any any = Array{String("x"), 1}; EXPECT_THROW( { try { [[maybe_unused]] auto value = any.cast>(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `List[index 0:"), std::string::npos); EXPECT_NE(what.find("to `List`"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(List, AnyImplicitConversionToArray) { List list = {10, 20, 30}; AnyView list_view = list; auto arr = list_view.cast>(); EXPECT_EQ(arr.size(), 3U); EXPECT_EQ(arr[0], 10); EXPECT_EQ(arr[1], 20); EXPECT_EQ(arr[2], 30); EXPECT_FALSE(arr.same_as(list)); } TEST(List, EmptyListDestructorDoesNotCrash) { { List empty; } { List filled = {1, 2, 3}; filled.clear(); } } TEST(List, SeqBaseObjPopBack) { List list = {10, 20, 30}; ListObj* obj = list.GetListObj(); obj->pop_back(); EXPECT_EQ(list.size(), 2U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 20); obj->pop_back(); obj->pop_back(); EXPECT_EQ(list.size(), 0U); EXPECT_THROW(obj->pop_back(), Error); } TEST(List, SeqBaseObjErase) { List list = {10, 20, 30, 40, 50}; ListObj* obj = list.GetListObj(); // Erase single element at index 2 (value 30) obj->erase(2); EXPECT_EQ(list.size(), 4U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 20); EXPECT_EQ(list[2], 40); EXPECT_EQ(list[3], 50); // Out of bounds EXPECT_THROW(obj->erase(4), Error); EXPECT_THROW(obj->erase(-1), Error); } TEST(List, SeqBaseObjEraseRange) { List list = {10, 20, 30, 40, 50}; ListObj* obj = list.GetListObj(); // Erase range [1, 3) -> removes 20, 30 obj->erase(int64_t{1}, int64_t{3}); EXPECT_EQ(list.size(), 3U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 40); EXPECT_EQ(list[2], 50); // No-op erase obj->erase(int64_t{1}, int64_t{1}); EXPECT_EQ(list.size(), 3U); // Invalid ranges EXPECT_THROW(obj->erase(int64_t{2}, int64_t{1}), Error); EXPECT_THROW(obj->erase(int64_t{-1}, int64_t{2}), Error); EXPECT_THROW(obj->erase(int64_t{0}, int64_t{4}), Error); } TEST(List, SeqBaseObjInsert) { List list; list.reserve(10); list.push_back(10); list.push_back(30); ListObj* obj = list.GetListObj(); // Insert 20 at index 1 obj->insert(1, Any(int64_t{20})); EXPECT_EQ(list.size(), 3U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 20); EXPECT_EQ(list[2], 30); // Insert at beginning obj->insert(0, Any(int64_t{5})); EXPECT_EQ(list[0], 5); EXPECT_EQ(list.size(), 4U); // Insert at end obj->insert(4, Any(int64_t{40})); EXPECT_EQ(list[4], 40); EXPECT_EQ(list.size(), 5U); // Out of bounds EXPECT_THROW(obj->insert(-1, Any(int64_t{0})), Error); EXPECT_THROW(obj->insert(6, Any(int64_t{0})), Error); } TEST(List, SeqBaseObjResize) { List list = {10, 20, 30}; list.reserve(10); ListObj* obj = list.GetListObj(); // Grow obj->resize(5); EXPECT_EQ(list.size(), 5U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 20); EXPECT_EQ(list[2], 30); // Shrink obj->resize(2); EXPECT_EQ(list.size(), 2U); EXPECT_EQ(list[0], 10); EXPECT_EQ(list[1], 20); // No-op obj->resize(2); EXPECT_EQ(list.size(), 2U); // Negative size EXPECT_THROW(obj->resize(-1), Error); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_map.cc000066400000000000000000000242501521067262500172230ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Map, Basic) { Map map0; TInt k0(0); map0.Set(k0, 1); EXPECT_EQ(map0.size(), 1); map0.Set(k0, 2); EXPECT_EQ(map0.size(), 1); auto it = map0.find(k0); EXPECT_TRUE(it != map0.end()); EXPECT_EQ((*it).second, 2); } TEST(Map, PODKey) { Map map0; // int as key map0.Set(1, 2); // float key is different map0.Set(1.1, 3); EXPECT_EQ(map0.size(), 2); auto it = map0.find(1.1); EXPECT_TRUE(it != map0.end()); EXPECT_EQ((*it).second.cast(), 3); } TEST(Map, Object) { TInt x(1); TInt z(100); TInt zz(1000); Map dict{{x, z}, {z, zz}}; EXPECT_EQ(dict.size(), 2); EXPECT_TRUE(dict[x].same_as(z)); EXPECT_TRUE(dict.count(z)); EXPECT_TRUE(!dict.count(zz)); } TEST(Map, Str) { TInt x(1); TInt z(100); Map dict{{"x", z}, {"z", z}}; EXPECT_EQ(dict.size(), 2); EXPECT_TRUE(dict["x"].same_as(z)); } TEST(Map, Mutate) { TInt x(1); TInt z(100); TInt zz(1000); Map dict{{x, z}, {z, zz}}; EXPECT_TRUE(dict[x].same_as(z)); dict.Set(x, zz); auto dict2 = dict; EXPECT_EQ(dict2.count(z), 1); dict.Set(zz, x); EXPECT_EQ(dict2.count(zz), 0); EXPECT_EQ(dict.count(zz), 1); auto it = dict.find(zz); EXPECT_TRUE(it != dict.end() && (*it).second.same_as(x)); it = dict2.find(zz); EXPECT_TRUE(it == dict2.end()); } TEST(Map, Clear) { TInt x(1); TInt z(100); Map dict{{x, z}, {z, z}}; EXPECT_EQ(dict.size(), 2); dict.clear(); EXPECT_EQ(dict.size(), 0); } TEST(Map, Insert) { auto check = [](const Map& result, std::unordered_map expected) { EXPECT_EQ(result.size(), expected.size()); for (const auto& kv : result) { EXPECT_TRUE(expected.count(kv.first)); EXPECT_EQ(expected[kv.first], kv.second); expected.erase(kv.first); } }; Map result; std::unordered_map expected; char key = 'a'; int64_t val = 1; for (int i = 0; i < 26; ++i, ++key, ++val) { std::string s(1, key); result.Set(s, val); expected[s] = val; check(result, expected); } } TEST(Map, Erase) { auto check = [](const Map& result, std::unordered_map expected) { EXPECT_EQ(result.size(), expected.size()); for (const auto& kv : result) { EXPECT_TRUE(expected.count(kv.first)); EXPECT_EQ(expected[kv.first], kv.second); expected.erase(kv.first); } }; Map map{{"a", 1}, {"b", 2}, {"c", 3}, {"d", 4}, {"e", 5}}; std::unordered_map stl; std::transform(map.begin(), map.end(), std::inserter(stl, stl.begin()), [](auto&& p) { return std::make_pair(p.first, p.second); }); for (char c = 'a'; c <= 'e'; ++c) { Map result = map; std::unordered_map expected(stl); std::string key(1, c); result.erase(key); expected.erase(key); check(result, expected); } } TEST(Map, AnyImplicitConversion) { Map map0; map0.Set(1, 2); map0.Set(2, 3.1); EXPECT_EQ(map0.size(), 2); // check will trigger copy AnyView view0 = map0; auto map1 = view0.cast>(); EXPECT_TRUE(!map1.same_as(map0)); EXPECT_EQ(map1[1], 2); EXPECT_EQ(map1[2], 3.1); EXPECT_EQ(map1.use_count(), 1); auto map2 = view0.cast>(); EXPECT_TRUE(map2.same_as(map0)); EXPECT_EQ(map2.use_count(), 2); auto map3 = view0.cast>(); EXPECT_TRUE(!map3.same_as(map0)); EXPECT_EQ(map3.use_count(), 1); Map map4{{"yes", 1.1}, {"no", 2.2}}; Any any1 = map4; auto map5 = any1.cast>(); EXPECT_TRUE(map5.same_as(map4)); EXPECT_EQ(map5.use_count(), 3); auto map6 = any1.cast>(); EXPECT_TRUE(map6.same_as(map4)); EXPECT_EQ(map6.use_count(), 4); EXPECT_EQ(map6["yes"].cast(), 1.1); EXPECT_EQ(map6["no"].cast(), 2.2); auto map7 = any1.cast>(); EXPECT_TRUE(map7.same_as(map4)); EXPECT_EQ(map7.use_count(), 5); auto map8 = any1.cast>(); EXPECT_TRUE(!map8.same_as(map4)); EXPECT_EQ(map8.use_count(), 1); EXPECT_EQ(map8["yes"]->value, 1.1); EXPECT_EQ(map8["no"]->value, 2.2); } TEST(Map, AnyConvertCheck) { Map map = {{11, 1.1}}; EXPECT_EQ(map[11].cast(), 1.1); AnyView view0 = map; auto arr1 = view0.cast>(); EXPECT_EQ(arr1[11], 1.1); Any any1 = map; using WrongMap = Map; EXPECT_THROW( { try { [[maybe_unused]] auto arr2 = any1.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE( what.find( "Cannot convert from type `Map[K, some value is float]` to `Map`"), std::string::npos); throw; } }, ::tvm::ffi::Error); using WrongMap2 = Map; EXPECT_THROW( { try { [[maybe_unused]] auto arr2 = any1.cast(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); EXPECT_NE(what.find("Cannot convert from type `Map[some key is int, V]` to " "`Map`"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(Map, FunctionGetItem) { Function f = Function::FromTyped([](const MapObj* n, const Any& k) -> Any { return n->at(k); }, "map_get_item"); Map map{{"x", 1}, {"y", 2}}; Any k("x"); Any v = f(map, k); EXPECT_EQ(v.cast(), 1); } TEST(Map, Upcast) { Map m0 = {{1, 2}, {3, 4}}; Map m1 = m0; EXPECT_EQ(m1[1].cast(), 2); EXPECT_EQ(m1[3].cast(), 4); static_assert(details::type_contains_v, Map>); Map> m2 = {{"x", {1}}, {"y", {2}}}; Map> m3 = m2; } template void PrintMap(const Map& m0) { std::cout << "{"; for (auto it = m0.begin(); it != m0.end(); ++it) { if (it != m0.begin()) { std::cout << ", "; } std::cout << (*it).first << ": " << (*it).second; } std::cout << "}" << std::endl; } TEST(Map, MapInsertOrder) { // test that map preserves the insertion order auto get_reverse_order = [](size_t size) { std::vector reverse_order; for (int i = static_cast(size); i != 0; --i) { reverse_order.push_back(i - 1); } return reverse_order; }; // NOLINTNEXTLINE(performance-unnecessary-value-param) auto check_map = [&](Map m0, size_t size, const std::vector& order) { auto lhs = m0.begin(); auto rhs = order.begin(); while (lhs != m0.end()) { TVM_FFI_ICHECK_EQ((*lhs).first, "hello" + std::to_string(*rhs)); TVM_FFI_ICHECK_EQ((*lhs).second, *rhs); ++lhs; ++rhs; } lhs = m0.end(); rhs = order.begin() + static_cast(size); do { --lhs; --rhs; TVM_FFI_ICHECK_EQ((*lhs).first, "hello" + std::to_string(*rhs)); TVM_FFI_ICHECK_EQ((*lhs).second, *rhs); } while (lhs != m0.begin()); }; auto check_order = [&](std::vector order) { Map m0; for (size_t i = 0; i < order.size(); ++i) { m0.Set("hello" + std::to_string(order[i]), order[i]); check_map(m0, i + 1, order); } check_map(m0, order.size(), order); // erase a few items m0.erase("hello" + std::to_string(order[0])); auto item0 = order[0]; order.erase(order.begin()); check_map(m0, order.size(), order); // erase the middle part if (order.size() > 1) { m0.erase("hello" + std::to_string(order[1])); order.erase(order.begin() + 1); check_map(m0, order.size(), order); } // erase the end m0.erase("hello" + std::to_string(order.back())); auto item2 = order.back(); order.erase(order.end() - 1); check_map(m0, order.size(), order); EXPECT_NE(m0.size(), 0); // put back some items order.push_back(item2); m0.Set("hello" + std::to_string(item2), item2); check_map(m0, order.size(), order); order.push_back(item0); m0.Set("hello" + std::to_string(item0), item0); check_map(m0, order.size(), order); }; // test with 17 items: DenseMapObj check_order(get_reverse_order(17)); // test with 4 items: SmallMapObj check_order(get_reverse_order(4)); } TEST(Map, EmptyIter) { Map m0; EXPECT_EQ(m0.begin(), m0.end()); // create a big map and then erase to keep a dense map empty for (int i = 0; i < 10; ++i) { m0.Set("hello" + std::to_string(i), i); } for (int i = 0; i < 10; ++i) { m0.erase("hello" + std::to_string(i)); } EXPECT_EQ(m0.size(), 0); // now m0 is dense map with all empty slots EXPECT_EQ(m0.begin(), m0.end()); } TEST(Map, DuplicatedKeysInit) { std::vector> data = {{"a", 1}, {"a", 2}, {"a", 3}}; Map map(data.begin(), data.end()); EXPECT_EQ(map.size(), 1); EXPECT_EQ(map["a"], 3); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_metadata.cc000066400000000000000000000320351521067262500202260ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include #include #include // Forward declarations for exported FFI functions extern "C" { // NOLINTNEXTLINE(bugprone-reserved-identifier) int __tvm_ffi__metadata_testing_dll_schema_id_int(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); // NOLINTNEXTLINE(bugprone-reserved-identifier) int __tvm_ffi__metadata_testing_dll_test_add_with_docstring(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); // NOLINTNEXTLINE(bugprone-reserved-identifier) int __tvm_ffi__doc_testing_dll_test_add_with_docstring(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); } namespace { using namespace tvm::ffi; using namespace tvm::ffi::reflection; // Helper to call metadata FFI functions and return the String result static String CallMetadataFunc(int (*func)(void*, const TVMFFIAny*, int32_t, TVMFFIAny*)) { return Function::InvokeExternC(nullptr, func).cast(); } static std::string ParseMetadataToSchema(const String& metadata) { return json::Parse(metadata) .cast>()["type_schema"] // .cast(); // } static std::string ParseMetadataToSchema(const TVMFFIByteArray& metadata) { return json::Parse(String(metadata)) .cast>()["type_schema"] // .cast(); // } TEST(Schema, GlobalFuncTypeSchema) { // Helper to fetch global function type schema via the exposed utility Function get_metadata = Function::GetGlobalRequired("ffi.GetGlobalFuncMetadata"); auto fetch = [&](const char* name) -> std::string { String metadata = get_metadata(String(name)).cast(); return ParseMetadataToSchema(metadata); }; // Simple IDs EXPECT_EQ(fetch("testing.schema_id_int"), R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"}]})"); EXPECT_EQ(fetch("testing.schema_id_float"), R"({"type":"ffi.Function","args":[{"type":"float"},{"type":"float"}]})"); EXPECT_EQ(fetch("testing.schema_id_bool"), R"({"type":"ffi.Function","args":[{"type":"bool"},{"type":"bool"}]})"); EXPECT_EQ(fetch("testing.schema_id_device"), R"({"type":"ffi.Function","args":[{"type":"Device"},{"type":"Device"}]})"); EXPECT_EQ(fetch("testing.schema_id_dtype"), R"({"type":"ffi.Function","args":[{"type":"DataType"},{"type":"DataType"}]})"); EXPECT_EQ(fetch("testing.schema_id_string"), R"({"type":"ffi.Function","args":[{"type":"ffi.String"},{"type":"ffi.String"}]})"); EXPECT_EQ(fetch("testing.schema_id_bytes"), R"({"type":"ffi.Function","args":[{"type":"ffi.Bytes"},{"type":"ffi.Bytes"}]})"); EXPECT_EQ(fetch("testing.schema_id_func"), R"({"type":"ffi.Function","args":[{"type":"ffi.Function"},{"type":"ffi.Function"}]})"); EXPECT_EQ( fetch("testing.schema_id_func_typed"), R"({"type":"ffi.Function","args":[{"type":"ffi.Function","args":[{"type":"None"},{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]},{"type":"ffi.Function","args":[{"type":"None"},{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]}]})"); EXPECT_EQ(fetch("testing.schema_id_any"), R"({"type":"ffi.Function","args":[{"type":"Any"},{"type":"Any"}]})"); EXPECT_EQ(fetch("testing.schema_id_object"), R"({"type":"ffi.Function","args":[{"type":"ffi.Object"},{"type":"ffi.Object"}]})"); EXPECT_EQ(fetch("testing.schema_id_dltensor"), R"({"type":"ffi.Function","args":[{"type":"DLTensor*"},{"type":"DLTensor*"}]})"); EXPECT_EQ(fetch("testing.schema_id_tensor"), R"({"type":"ffi.Function","args":[{"type":"ffi.Tensor"},{"type":"ffi.Tensor"}]})"); EXPECT_EQ(fetch("testing.schema_tensor_view_input"), R"({"type":"ffi.Function","args":[{"type":"None"},{"type":"DLTensor*"}]})"); EXPECT_EQ( fetch("testing.schema_id_opt_int"), R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"int"}]},{"type":"Optional","args":[{"type":"int"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_opt_str"), R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_opt_obj"), R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.Object"}]},{"type":"Optional","args":[{"type":"ffi.Object"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_arr_int"), R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"ffi.Array","args":[{"type":"int"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_arr_str"), R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"ffi.String"}]},{"type":"ffi.Array","args":[{"type":"ffi.String"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_arr_obj"), R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"ffi.Object"}]},{"type":"ffi.Array","args":[{"type":"ffi.Object"}]}]})"); EXPECT_EQ(fetch("testing.schema_id_arr"), R"({"type":"ffi.Function","args":[{"type":"ffi.Array"},{"type":"ffi.Array"}]})"); EXPECT_EQ( fetch("testing.schema_id_map_str_int"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_map_str_str"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]}]})"); EXPECT_EQ( fetch("testing.schema_id_map_str_obj"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]}]})"); EXPECT_EQ(fetch("testing.schema_id_map"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map"},{"type":"ffi.Map"}]})"); EXPECT_EQ( fetch("testing.schema_id_variant_int_str"), R"({"type":"ffi.Function","args":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]},{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]}]})"); // Packed function registered via def_packed: schema is plain ffi.Function EXPECT_EQ(fetch("testing.schema_packed"), R"({"type":"ffi.Function"})"); // Mixed containers and optionals EXPECT_EQ( fetch("testing.schema_arr_map_opt"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"ffi.Array","args":[{"type":"Optional","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})"); EXPECT_EQ( fetch("testing.schema_variant_mix"), R"({"type":"ffi.Function","args":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]})"); // No-arg and no-return combinations EXPECT_EQ(fetch("testing.schema_no_args"), R"({"type":"ffi.Function","args":[{"type":"int"}]})"); EXPECT_EQ(fetch("testing.schema_no_return"), R"({"type":"ffi.Function","args":[{"type":"None"},{"type":"int"}]})"); EXPECT_EQ(fetch("testing.schema_no_args_no_return"), R"({"type":"ffi.Function","args":[{"type":"None"}]})"); } TEST(Schema, FieldTypeSchemas) { // Validate type schema JSON on fields of testing.SchemaAllTypes const char* kTypeKey = "testing.SchemaAllTypes"; // Helper to fetch a field's type schema by name auto field_schema = [&](const char* field_name) -> std::string { const TVMFFIFieldInfo* info = GetFieldInfo(kTypeKey, field_name); return ParseMetadataToSchema(info->metadata); }; EXPECT_EQ(field_schema("v_bool"), R"({"type":"bool"})"); EXPECT_EQ(field_schema("v_int"), R"({"type":"int"})"); EXPECT_EQ(field_schema("v_float"), R"({"type":"float"})"); EXPECT_EQ(field_schema("v_device"), R"({"type":"Device"})"); EXPECT_EQ(field_schema("v_dtype"), R"({"type":"DataType"})"); EXPECT_EQ(field_schema("v_string"), R"({"type":"ffi.String"})"); EXPECT_EQ(field_schema("v_bytes"), R"({"type":"ffi.Bytes"})"); EXPECT_EQ(field_schema("v_opt_int"), R"({"type":"Optional","args":[{"type":"int"}]})"); EXPECT_EQ(field_schema("v_opt_str"), R"({"type":"Optional","args":[{"type":"ffi.String"}]})"); EXPECT_EQ(field_schema("v_arr_int"), R"({"type":"ffi.Array","args":[{"type":"int"}]})"); EXPECT_EQ(field_schema("v_arr_str"), R"({"type":"ffi.Array","args":[{"type":"ffi.String"}]})"); EXPECT_EQ(field_schema("v_map_str_int"), R"({"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]})"); EXPECT_EQ( field_schema("v_map_str_arr_int"), R"({"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]})"); EXPECT_EQ( field_schema("v_variant"), R"({"type":"Variant","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}]})"); EXPECT_EQ( field_schema("v_opt_arr_variant"), R"({"type":"Optional","args":[{"type":"ffi.Array","args":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]}]}]})"); } TEST(Schema, MethodTypeSchemas) { const char* kTypeKey = "testing.SchemaAllTypes"; auto method_schema = [&](const char* method_name) -> std::string { const TVMFFIMethodInfo* info = GetMethodInfo(kTypeKey, method_name); return ParseMetadataToSchema(info->metadata); }; // Instance methods EXPECT_EQ( method_schema("add_int"), R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"testing.SchemaAllTypes"},{"type":"int"}]})"); EXPECT_EQ( method_schema("append_int"), R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"testing.SchemaAllTypes"},{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"int"}]})"); EXPECT_EQ( method_schema("maybe_concat"), R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"testing.SchemaAllTypes"},{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})"); EXPECT_EQ( method_schema("merge_map"), R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"testing.SchemaAllTypes"},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]})"); // Static method make_with: return type is the object type itself. // Build expected JSON as ffi.Function with return type = type_key and args = (int, float, str) EXPECT_EQ( method_schema("make_with"), R"({"type":"ffi.Function","args":[{"type":"testing.SchemaAllTypes"},{"type":"int"},{"type":"float"},{"type":"ffi.String"}]})"); } TEST(Schema, DLLExportedFuncMetadata) { // Minimal sanity check that DLL export metadata mechanism works. EXPECT_EQ(ParseMetadataToSchema(CallMetadataFunc(__tvm_ffi__metadata_testing_dll_schema_id_int)), R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"}]})"); } TEST(Schema, DLLExportedFuncDocumentation) { EXPECT_EQ(ParseMetadataToSchema( CallMetadataFunc(__tvm_ffi__metadata_testing_dll_test_add_with_docstring)), R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"},{"type":"int"}]})"); String doc = CallMetadataFunc(__tvm_ffi__doc_testing_dll_test_add_with_docstring); std::string doc_str(doc); EXPECT_EQ(doc_str, R"(Add two integers and return the sum. Parameters ---------- a : int First integer b : int Second integer Returns ------- result : int Sum of a and b)"); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_object.cc000066400000000000000000000203201521067262500177060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Object, RefCounter) { ObjectPtr a = make_object(11); ObjectPtr b = a; EXPECT_EQ(a->value, 11); EXPECT_EQ(a.use_count(), 2); ObjectPtr aa = make_object(*a); EXPECT_EQ(aa.use_count(), 1); EXPECT_EQ(aa->value, 11); b.reset(); EXPECT_EQ(a.use_count(), 1); EXPECT_TRUE(b == nullptr); EXPECT_EQ(b.use_count(), 0); ObjectPtr c = std::move(a); EXPECT_EQ(c.use_count(), 1); EXPECT_TRUE(a == nullptr); // NOLINT(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_EQ(c->value, 11); } TEST(Object, TypeInfo) { const TypeInfo* info = TVMFFIGetTypeInfo(TIntObj::RuntimeTypeIndex()); EXPECT_TRUE(info != nullptr); EXPECT_EQ(info->type_index, TIntObj::RuntimeTypeIndex()); EXPECT_EQ(info->type_depth, 2); EXPECT_EQ(info->type_ancestors[0]->type_index, Object::RuntimeTypeIndex()); EXPECT_EQ(info->type_ancestors[1]->type_index, TNumberObj::RuntimeTypeIndex()); EXPECT_GE(info->type_index, TypeIndex::kTVMFFIDynObjectBegin); } TEST(Object, InstanceCheck) { ObjectPtr a = make_object(11); ObjectPtr b = make_object(11); EXPECT_TRUE(a->IsInstance()); EXPECT_TRUE(a->IsInstance()); EXPECT_TRUE(a->IsInstance()); EXPECT_TRUE(!a->IsInstance()); EXPECT_TRUE(a->IsInstance()); EXPECT_TRUE(b->IsInstance()); EXPECT_TRUE(!b->IsInstance()); EXPECT_TRUE(b->IsInstance()); } TEST(ObjectRef, as) { ObjectRef a = TInt(10); ObjectRef b = TFloat(20); // nullable object ObjectRef c(nullptr); EXPECT_TRUE(a.as() != nullptr); EXPECT_TRUE(a.as() == nullptr); EXPECT_TRUE(a.as() != nullptr); EXPECT_TRUE(b.as() == nullptr); EXPECT_TRUE(b.as() != nullptr); EXPECT_TRUE(b.as() != nullptr); EXPECT_TRUE(c.as() == nullptr); EXPECT_TRUE(c.as() == nullptr); EXPECT_TRUE(c.as() == nullptr); EXPECT_EQ(a.as()->value, 10); EXPECT_EQ(b.as()->value, 20); } TEST(ObjectRef, UnsafeInit) { ObjectRef a(UnsafeInit{}); EXPECT_TRUE(a.get() == nullptr); TInt b(UnsafeInit{}); EXPECT_TRUE(b.get() == nullptr); } TEST(Object, CAPIAccessor) { ObjectRef a = TInt(10); TVMFFIObjectHandle obj = details::ObjectUnsafe::RawObjectPtrFromObjectRef(a); int32_t type_index = TVMFFIObjectGetTypeIndex(obj); EXPECT_EQ(type_index, TIntObj::RuntimeTypeIndex()); } TEST(Object, WeakObjectPtr) { // Test basic construction from ObjectPtr ObjectPtr strong_ptr = make_object(42); WeakObjectPtr weak_ptr(strong_ptr); EXPECT_EQ(strong_ptr.use_count(), 1); EXPECT_FALSE(weak_ptr.expired()); EXPECT_EQ(weak_ptr.use_count(), 1); // Test lock() when object is still alive ObjectPtr locked_ptr = weak_ptr.lock(); EXPECT_TRUE(locked_ptr != nullptr); EXPECT_EQ(locked_ptr->value, 42); EXPECT_EQ(strong_ptr.use_count(), 2); EXPECT_EQ(weak_ptr.use_count(), 2); // Test lock() when object is expired strong_ptr.reset(); locked_ptr.reset(); EXPECT_TRUE(weak_ptr.expired()); EXPECT_EQ(weak_ptr.use_count(), 0); ObjectPtr expired_lock = weak_ptr.lock(); EXPECT_TRUE(expired_lock == nullptr); } TEST(Object, WeakObjectPtrAssignment) { // Test copy construction ObjectPtr new_strong = make_object(100); WeakObjectPtr weak1(new_strong); WeakObjectPtr weak2(weak1); EXPECT_EQ(new_strong.use_count(), 1); EXPECT_FALSE(weak1.expired()); EXPECT_FALSE(weak2.expired()); EXPECT_EQ(weak1.use_count(), 1); EXPECT_EQ(weak2.use_count(), 1); // Test move construction WeakObjectPtr weak3(std::move(weak1)); // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_TRUE(weak1.expired()); // weak1 should be moved from EXPECT_FALSE(weak3.expired()); EXPECT_EQ(weak3.use_count(), 1); // Test assignment WeakObjectPtr weak4; weak4 = weak2; EXPECT_FALSE(weak2.expired()); EXPECT_FALSE(weak4.expired()); EXPECT_EQ(weak2.use_count(), 1); EXPECT_EQ(weak4.use_count(), 1); // Test move assignment WeakObjectPtr weak5; weak5 = std::move(weak2); // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_TRUE(weak2.expired()); // weak2 should be moved from EXPECT_FALSE(weak5.expired()); EXPECT_EQ(weak5.use_count(), 1); // Test reset() weak3.reset(); EXPECT_TRUE(weak3.expired()); EXPECT_EQ(weak3.use_count(), 0); // Test swap() ObjectPtr strong_a = make_object(200); ObjectPtr strong_b = make_object(300); WeakObjectPtr weak_a(strong_a); WeakObjectPtr weak_b(strong_b); weak_a.swap(weak_b); EXPECT_EQ(weak_a.lock()->value, 300); EXPECT_EQ(weak_b.lock()->value, 200); // Test construction from nullptr WeakObjectPtr null_weak(nullptr); EXPECT_TRUE(null_weak.expired()); EXPECT_EQ(null_weak.use_count(), 0); EXPECT_TRUE(null_weak.lock() == nullptr); // Test inheritance compatibility ObjectPtr number_ptr = make_object(500); WeakObjectPtr number_weak(number_ptr); EXPECT_FALSE(number_weak.expired()); EXPECT_EQ(number_weak.use_count(), 1); // Test that weak references don't prevent object deletion ObjectPtr temp_strong = make_object(999); WeakObjectPtr temp_weak(temp_strong); EXPECT_FALSE(temp_weak.expired()); temp_strong.reset(); EXPECT_TRUE(temp_weak.expired()); EXPECT_TRUE(temp_weak.lock() == nullptr); // Test multiple weak references ObjectPtr multi_strong = make_object(777); WeakObjectPtr multi_weak1(multi_strong); WeakObjectPtr multi_weak2(multi_strong); WeakObjectPtr multi_weak3(multi_strong); EXPECT_EQ(multi_strong.use_count(), 1); EXPECT_FALSE(multi_weak1.expired()); EXPECT_FALSE(multi_weak2.expired()); EXPECT_FALSE(multi_weak3.expired()); // All weak references should be able to lock ObjectPtr lock1 = multi_weak1.lock(); ObjectPtr lock2 = multi_weak2.lock(); ObjectPtr lock3 = multi_weak3.lock(); EXPECT_EQ(multi_strong.use_count(), 4); EXPECT_EQ(lock1->value, 777); EXPECT_EQ(lock2->value, 777); EXPECT_EQ(lock3->value, 777); } TEST(Object, OpaqueObject) { thread_local int deleter_trigger_counter = 0; struct DummyOpaqueObject { int value; explicit DummyOpaqueObject(int value) : value(value) {} static void Deleter(void* handle) { deleter_trigger_counter++; delete static_cast(handle); } }; TVMFFIObjectHandle handle = nullptr; TVM_FFI_CHECK_SAFE_CALL(TVMFFIObjectCreateOpaque(new DummyOpaqueObject(10), kTVMFFIOpaquePyObject, DummyOpaqueObject::Deleter, &handle)); ObjectPtr a = details::ObjectUnsafe::ObjectPtrFromOwned(static_cast(handle)); EXPECT_EQ(a->type_index(), kTVMFFIOpaquePyObject); EXPECT_EQ(static_cast(TVMFFIOpaqueObjectGetCellPtr(a.get())->handle)->value, 10); EXPECT_EQ(a.use_count(), 1); EXPECT_EQ(deleter_trigger_counter, 0); a.reset(); EXPECT_EQ(deleter_trigger_counter, 1); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_optional.cc000066400000000000000000000143511521067262500202740ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Optional, TInt) { Optional x; Optional y = TInt(11); static_assert(sizeof(Optional) == sizeof(ObjectRef)); EXPECT_TRUE(!x.has_value()); EXPECT_EQ(x.value_or(TInt(12))->value, 12); EXPECT_TRUE(y.has_value()); EXPECT_EQ(y.value_or(TInt(12))->value, 11); Any z_any = std::move(y); EXPECT_TRUE(z_any != nullptr); EXPECT_EQ((z_any.cast())->value, 11); // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_TRUE(!y.has_value()); // move from any to optional auto y2 = std::move(z_any).cast>(); EXPECT_EQ(y2.use_count(), 1); EXPECT_TRUE(y2.has_value()); EXPECT_EQ(y2.value_or(TInt(12))->value, 11); } TEST(Optional, double) { Optional x; Optional y = 11.0; static_assert(sizeof(Optional) > sizeof(ObjectRef)); EXPECT_TRUE(!x.has_value()); EXPECT_EQ(x.value_or(12), 12); EXPECT_TRUE(x != 12); EXPECT_TRUE(y.has_value()); EXPECT_EQ(y.value_or(12), 11); EXPECT_TRUE(y == 11); EXPECT_TRUE(y != 12); } TEST(Optional, AnyConvertInt) { Optional opt_v0 = 1; EXPECT_EQ(opt_v0.value(), 1); EXPECT_TRUE(opt_v0.has_value()); AnyView view0 = opt_v0; EXPECT_EQ(view0.cast(), 1); Any any1; auto opt_v1 = std::move(any1).cast>(); EXPECT_TRUE(!opt_v1.has_value()); Optional opt_v2 = 11; Any any2 = std::move(opt_v2); EXPECT_EQ(any2.cast(), 11); } TEST(Optional, AnyConvertArray) { AnyView view0; Array> arr_nested = {{}, {TInt(1), TFloat(2)}}; view0 = arr_nested; auto opt_arr = view0.cast>>>(); EXPECT_EQ(arr_nested.use_count(), 2); auto arr1 = view0.cast>>>(); EXPECT_EQ(arr_nested.use_count(), 3); EXPECT_EQ(arr1.value()[1][1].as()->value, 2); Any any1; auto arr2 = any1.cast>>>(); EXPECT_TRUE(!arr2.has_value()); EXPECT_THROW( { try { [[maybe_unused]] auto arr2 = view0.cast>>>(); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); std::string what = error.what(); std::cout << what << std::endl; EXPECT_NE(what.find("to `Optional>>`"), std::string::npos); throw; } }, ::tvm::ffi::Error); } TEST(Optional, OptionalOfOptional) { // testcase of optional Optional> opt_opt_int; EXPECT_TRUE(!opt_opt_int.has_value()); Optional> opt_opt_int2 = Optional(std::nullopt); EXPECT_TRUE(opt_opt_int2.has_value()); EXPECT_TRUE(!opt_opt_int2.value().has_value()); // Optional> Optional> opt_opt_tint; EXPECT_TRUE(!opt_opt_tint.has_value()); Optional> opt_opt_tint2 = Optional(std::nullopt); EXPECT_TRUE(opt_opt_tint2.has_value()); EXPECT_TRUE(!opt_opt_tint2.value().has_value()); opt_opt_tint2 = std::nullopt; EXPECT_TRUE(!opt_opt_tint2.has_value()); Optional> opt_opt_tint3 = Optional(TInt(42)); EXPECT_TRUE(opt_opt_tint3.has_value()); EXPECT_TRUE(opt_opt_tint3.value().has_value()); EXPECT_EQ(opt_opt_tint3.value().value()->value, 42); } TEST(Optional, ValueMove) { Optional y = TInt(11); TInt x = std::move(y).value(); // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_TRUE(!y.has_value()); EXPECT_EQ(x->value, 11); Optional opt_tint = TInt(21); EXPECT_TRUE(opt_tint.has_value()); EXPECT_EQ((*opt_tint)->value, 21); TInt moved_tint = *std::move(opt_tint); EXPECT_EQ(moved_tint->value, 21); // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) EXPECT_TRUE(!opt_tint.has_value()); } TEST(Optional, OptionalInArray) { // This pattern plus iteration may cause memory leak // this is because arr[0] returns a temporary object // and further call arr[0].value() may return a reference to // the temporary object Array>> arr = {Array({TInt(0), TInt(1)})}; int counter = 0; for (const auto& x : arr[0].value()) { EXPECT_EQ(x->value, counter++); } Any any = arr; auto opt_arr = any.cast>>>(); EXPECT_EQ(opt_arr[0].value()[0]->value, 0); } TEST(Optional, String) { Optional opt_str; EXPECT_TRUE(!opt_str.has_value()); EXPECT_EQ(opt_str.value_or("default"), "default"); EXPECT_TRUE(opt_str != "default"); EXPECT_TRUE(opt_str != String("default")); EXPECT_TRUE(opt_str == std::nullopt); opt_str = "hello"; EXPECT_TRUE(opt_str.has_value()); EXPECT_EQ(opt_str.value(), "hello"); EXPECT_TRUE(opt_str == "hello"); EXPECT_TRUE(opt_str == String("hello")); EXPECT_TRUE(opt_str != std::nullopt); static_assert(sizeof(Optional) == sizeof(String)); } TEST(Optional, Bytes) { Optional opt_bytes; EXPECT_TRUE(!opt_bytes.has_value()); EXPECT_EQ(opt_bytes.value_or(std::string("default")), "default"); opt_bytes = std::string("hello"); EXPECT_TRUE(opt_bytes.has_value()); EXPECT_EQ(opt_bytes.value().operator std::string(), "hello"); EXPECT_TRUE(opt_bytes != std::nullopt); static_assert(sizeof(Optional) == sizeof(Bytes)); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_overload.cc000066400000000000000000000070111521067262500202550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include namespace { using namespace tvm::ffi; struct TestOverloadObj : public Object { explicit TestOverloadObj(int32_t x) : type(Type::INT) {} explicit TestOverloadObj(float y) : type(Type::FLOAT) {} static int AddOneInt(int x) { return x + 1; } static float AddOneFloat(float x) { return x + 1.0f; } template auto Holds(T) const { if constexpr (std::is_same_v) { return type == Type::INT; } else if constexpr (std::is_same_v) { return type == Type::FLOAT; } else { static_assert(sizeof(T) == 0, "Unsupported type"); } } enum class Type { INT, FLOAT } type; TVM_FFI_DECLARE_OBJECT_INFO("test.TestOverloadObj", TestOverloadObj, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::OverloadObjectDef() .def(refl::init()) .def(refl::init()) .def("hold_same_type", &TestOverloadObj::Holds) .def("hold_same_type", &TestOverloadObj::Holds) .def_static("add_one_static", &TestOverloadObj::AddOneInt) .def_static("add_one_static", &TestOverloadObj::AddOneFloat); } TEST(Reflection, CallOverloadedInitMethod) { Function init_method = reflection::GetMethod("test.TestOverloadObj", "__ffi_init__"); Any obj_a = init_method(10); // choose the int constructor EXPECT_TRUE(obj_a.as() != nullptr); EXPECT_EQ(obj_a.as()->type, TestOverloadObj::Type::INT); Any obj_b = init_method(3.14f); // choose the float constructor EXPECT_TRUE(obj_b.as() != nullptr); EXPECT_EQ(obj_b.as()->type, TestOverloadObj::Type::FLOAT); } TEST(Reflection, CallOverloadedMethod) { Function init_method = reflection::GetMethod("test.TestOverloadObj", "__ffi_init__"); Function hold_same_type = reflection::GetMethod("test.TestOverloadObj", "hold_same_type"); Any obj_a = init_method(10); // choose the int constructor Any res_a = hold_same_type(obj_a, 20); EXPECT_EQ(res_a.as(), true); Any res_b = hold_same_type(obj_a, 3.14f); EXPECT_EQ(res_b.as(), false); } TEST(Reflection, CallOverloadedStaticMethod) { Function add_one = reflection::GetMethod("test.TestOverloadObj", "add_one_static"); Any res_a = add_one(20); EXPECT_EQ(res_a.as(), 21); Any res_b = add_one(1.0f); static_assert(1.0f + 1.0f == 2.0f); EXPECT_EQ(res_b.as(), 2.0f); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_reflection.cc000066400000000000000000000677041521067262500206130ustar00rootroot00000000000000 /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include "./testing_object.h" /*! \brief Look up __ffi_init__ from the TypeAttrColumn (not the method table). */ static tvm::ffi::Function GetInitAttr(const char* type_key) { static tvm::ffi::reflection::TypeAttrColumn col(tvm::ffi::reflection::type_attr::kInit); return col[tvm::ffi::TypeKeyToIndex(type_key)].cast(); } namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; struct TestObjA : public Object { int64_t x; int64_t y; TestObjA(int64_t x, int64_t y) : x(x), y(y) {} static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("test.TestObjA", TestObjA, Object); }; struct TestObjADerived : public TestObjA { int64_t z; TestObjADerived(int64_t x, int64_t y, int64_t z) : TestObjA(x, y), z(z) {} TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.TestObjADerived", TestObjADerived, TestObjA); }; struct TestObjRefADerived : public ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestObjRefADerived, ObjectRef, TestObjADerived); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; TIntObj::RegisterReflection(); TFloatObj::RegisterReflection(); TPrimExprObj::RegisterReflection(); TVarObj::RegisterReflection(); TVarWithDepObj::RegisterReflection(); TDefHolderObj::RegisterReflection(); TFuncObj::RegisterReflection(); TCustomFuncObj::RegisterReflection(); TAllFieldsObj::RegisterReflection(); TWithDefaultsObj::RegisterReflection(); refl::ObjectDef() .def(refl::init()) .def_ro("x", &TestObjA::x) .def_rw("y", &TestObjA::y); refl::ObjectDef() .def(refl::init()) .def_ro("z", &TestObjADerived::z); } TEST(Reflection, GetFieldByteOffset) { EXPECT_EQ(reflection::GetFieldByteOffsetToObject(&TestObjA::x), sizeof(TVMFFIObject)); EXPECT_EQ(reflection::GetFieldByteOffsetToObject(&TestObjA::y), 8 + sizeof(TVMFFIObject)); EXPECT_EQ(reflection::GetFieldByteOffsetToObject(&TIntObj::value), sizeof(TVMFFIObject)); } TEST(Reflection, FieldGetter) { ObjectRef a = TInt(10); reflection::FieldGetter getter("test.Int", "value"); EXPECT_EQ(getter(a).cast(), 10); ObjectRef b = TFloat(10.0); reflection::FieldGetter getter_float("test.Float", "value"); EXPECT_EQ(getter_float(b).cast(), 10.0); } TEST(Reflection, FieldSetter) { ObjectRef a = TFloat(10.0); reflection::FieldSetter setter("test.Float", "value"); setter(a, 20.0); EXPECT_EQ(a.as()->value, 20.0); } TEST(Reflection, FieldInfo) { const TVMFFIFieldInfo* info_int = reflection::GetFieldInfo("test.Int", "value"); EXPECT_FALSE(info_int->flags & kTVMFFIFieldFlagBitMaskHasDefault); EXPECT_FALSE(info_int->flags & kTVMFFIFieldFlagBitMaskWritable); EXPECT_EQ(Bytes(info_int->doc).operator std::string(), ""); const TVMFFIFieldInfo* info_float = reflection::GetFieldInfo("test.Float", "value"); EXPECT_EQ(info_float->default_value_or_factory.v_float64, 10.0); EXPECT_TRUE(info_float->flags & kTVMFFIFieldFlagBitMaskHasDefault); EXPECT_FALSE(info_float->flags & kTVMFFIFieldFlagBitMaskWritable); EXPECT_EQ(Bytes(info_float->doc).operator std::string(), "float value field"); const TVMFFIFieldInfo* info_prim_expr_dtype = reflection::GetFieldInfo("test.PrimExpr", "dtype"); AnyView default_value = AnyView::CopyFromTVMFFIAny(info_prim_expr_dtype->default_value_or_factory); EXPECT_EQ(default_value.cast(), "float"); EXPECT_TRUE(info_prim_expr_dtype->flags & kTVMFFIFieldFlagBitMaskHasDefault); EXPECT_TRUE(info_prim_expr_dtype->flags & kTVMFFIFieldFlagBitMaskWritable); EXPECT_EQ(Bytes(info_prim_expr_dtype->doc).operator std::string(), "dtype field"); } TEST(Reflection, MethodInfo) { const TVMFFIMethodInfo* info_int_static_add = reflection::GetMethodInfo("test.Int", "static_add"); EXPECT_TRUE(info_int_static_add->flags & kTVMFFIFieldFlagBitMaskIsStaticMethod); EXPECT_EQ(Bytes(info_int_static_add->doc).operator std::string(), "static add method"); const TVMFFIMethodInfo* info_float_add = reflection::GetMethodInfo("test.Float", "add"); EXPECT_FALSE(info_float_add->flags & kTVMFFIFieldFlagBitMaskIsStaticMethod); EXPECT_EQ(Bytes(info_float_add->doc).operator std::string(), "add method"); const TVMFFIMethodInfo* info_float_sub = reflection::GetMethodInfo("test.Float", "sub"); EXPECT_FALSE(info_float_sub->flags & kTVMFFIFieldFlagBitMaskIsStaticMethod); EXPECT_EQ(Bytes(info_float_sub->doc).operator std::string(), ""); } TEST(Reflection, CallMethod) { Function static_int_add = reflection::GetMethod("test.Int", "static_add"); EXPECT_EQ(static_int_add(TInt(1), TInt(2)).cast()->value, 3); Function float_add = reflection::GetMethod("test.Float", "add"); EXPECT_EQ(float_add(TFloat(1), 2.0).cast(), 3.0); Function float_sub = reflection::GetMethod("test.Float", "sub"); EXPECT_EQ(float_sub(TFloat(1), 2.0).cast(), -1.0); Function prim_expr_sub = reflection::GetMethod("test.PrimExpr", "sub"); EXPECT_EQ(prim_expr_sub(TPrimExpr("float", 1), 2.0).cast(), -1.0); } TEST(Reflection, InitFunctionBase) { Function int_init = GetInitAttr("test.TestObjA"); Any obj_a = int_init(1, 2); EXPECT_TRUE(obj_a.as() != nullptr); EXPECT_EQ(obj_a.as()->x, 1); EXPECT_EQ(obj_a.as()->y, 2); } TEST(Reflection, InitFunctionDerived) { Function derived_init = GetInitAttr("test.TestObjADerived"); Any obj_derived = derived_init(1, 2, 3); EXPECT_TRUE(obj_derived.as() != nullptr); EXPECT_EQ(obj_derived.as()->x, 1); EXPECT_EQ(obj_derived.as()->y, 2); EXPECT_EQ(obj_derived.as()->z, 3); } TEST(Reflection, ForEachFieldInfo) { const TypeInfo* info = TVMFFIGetTypeInfo(TestObjADerived::RuntimeTypeIndex()); Map field_name_to_offset; reflection::ForEachFieldInfo(info, [&](const TVMFFIFieldInfo* field_info) { field_name_to_offset.Set(String(field_info->name), static_cast(field_info->offset)); }); EXPECT_EQ(field_name_to_offset["x"], sizeof(TVMFFIObject)); EXPECT_EQ(field_name_to_offset["y"], 8 + sizeof(TVMFFIObject)); EXPECT_EQ(field_name_to_offset["z"], 16 + sizeof(TVMFFIObject)); } TEST(Reflection, TypeAttrColumn) { reflection::TypeAttrColumn size_attr("test.size"); EXPECT_EQ(size_attr[TIntObj::RuntimeTypeIndex()].cast(), sizeof(TIntObj)); } TEST(Reflection, TypeAttrColumnBeginIndex) { // Get the column and verify begin_index TVMFFIByteArray attr_name = {"test.size", std::char_traits::length("test.size")}; const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&attr_name); ASSERT_NE(column, nullptr); // begin_index should be >= 0 EXPECT_GE(column->begin_index, 0); // size should cover the range from begin_index EXPECT_GT(column->size, 0); // verify that lookup of a type_index below begin_index returns None reflection::TypeAttrColumn size_attr("test.size"); AnyView result = size_attr[0]; // index 0 is kTVMFFINone, unlikely to have this attr (void)result; // suppress unused variable warning; we only verify no crash occurs // The result may or may not be None depending on begin_index; the key is no crash. // verify the known registered entry still works EXPECT_EQ(size_attr[TIntObj::RuntimeTypeIndex()].cast(), sizeof(TIntObj)); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def_method("testing.Int_GetValue", &TIntObj::GetValue); } TEST(Reflection, FuncRegister) { Function fget_value = Function::GetGlobalRequired("testing.Int_GetValue"); TInt a(12); EXPECT_EQ(fget_value(a).cast(), 12); } TEST(Reflection, ObjectCreator) { namespace refl = tvm::ffi::reflection; refl::ObjectCreator creator("test.Int"); EXPECT_EQ(creator(Map({{"value", 1}})).cast()->value, 1); } TEST(Reflection, AccessPath) { namespace refl = tvm::ffi::reflection; // Test basic path construction and ToSteps() refl::AccessPath path = refl::AccessPath::Root()->Attr("body")->ArrayItem(1); auto steps = path->ToSteps(); EXPECT_EQ(steps.size(), 2); EXPECT_EQ(steps[0]->kind, refl::AccessKind::kAttr); EXPECT_EQ(steps[1]->kind, refl::AccessKind::kArrayItem); EXPECT_EQ(steps[0]->key.cast(), "body"); EXPECT_EQ(steps[1]->key.cast(), 1); // Test PathEqual with identical paths refl::AccessPath path2 = refl::AccessPath::Root()->Attr("body")->ArrayItem(1); EXPECT_TRUE(path->PathEqual(path2)); EXPECT_TRUE(path->IsPrefixOf(path2)); // Test PathEqual with different paths refl::AccessPath path3 = refl::AccessPath::Root()->Attr("body")->ArrayItem(2); EXPECT_FALSE(path->PathEqual(path3)); EXPECT_FALSE(path->IsPrefixOf(path3)); // Test prefix relationship - path4 extends path, so path should be prefix of path4 refl::AccessPath path4 = refl::AccessPath::Root()->Attr("body")->ArrayItem(1)->Attr("body"); EXPECT_FALSE(path->PathEqual(path4)); // Not equal (different lengths) EXPECT_TRUE(path->IsPrefixOf(path4)); // But path is a prefix of path4 // Test completely different paths refl::AccessPath path5 = refl::AccessPath::Root()->ArrayItem(0)->ArrayItem(1)->Attr("body"); EXPECT_FALSE(path->PathEqual(path5)); EXPECT_FALSE(path->IsPrefixOf(path5)); // Test Root path refl::AccessPath root = refl::AccessPath::Root(); auto root_steps = root->ToSteps(); EXPECT_EQ(root_steps.size(), 0); EXPECT_EQ(root->depth, 0); EXPECT_TRUE(root->IsPrefixOf(path)); EXPECT_TRUE(root->IsPrefixOf(root)); EXPECT_TRUE(root->PathEqual(refl::AccessPath::Root())); // Test depth calculations EXPECT_EQ(path->depth, 2); EXPECT_EQ(path4->depth, 3); EXPECT_EQ(root->depth, 0); // Test MapItem access refl::AccessPath map_path = refl::AccessPath::Root()->Attr("data")->MapItem("key1"); auto map_steps = map_path->ToSteps(); EXPECT_EQ(map_steps.size(), 2); EXPECT_EQ(map_steps[0]->kind, refl::AccessKind::kAttr); EXPECT_EQ(map_steps[1]->kind, refl::AccessKind::kMapItem); EXPECT_EQ(map_steps[0]->key.cast(), "data"); EXPECT_EQ(map_steps[1]->key.cast(), "key1"); // Test MapItemMissing access refl::AccessPath map_missing_path = refl::AccessPath::Root()->MapItemMissing(42); auto map_missing_steps = map_missing_path->ToSteps(); EXPECT_EQ(map_missing_steps.size(), 1); EXPECT_EQ(map_missing_steps[0]->kind, refl::AccessKind::kMapItemMissing); EXPECT_EQ(map_missing_steps[0]->key.cast(), 42); // Test ArrayItemMissing access refl::AccessPath array_missing_path = refl::AccessPath::Root()->ArrayItemMissing(5); auto array_missing_steps = array_missing_path->ToSteps(); EXPECT_EQ(array_missing_steps.size(), 1); EXPECT_EQ(array_missing_steps[0]->kind, refl::AccessKind::kArrayItemMissing); EXPECT_EQ(array_missing_steps[0]->key.cast(), 5); // Test FromSteps static method - round trip conversion auto original_steps = path->ToSteps(); refl::AccessPath reconstructed = refl::AccessPath::FromSteps(original_steps); EXPECT_TRUE(path->PathEqual(reconstructed)); EXPECT_EQ(path->depth, reconstructed->depth); // Test complex prefix relationships refl::AccessPath short_path = refl::AccessPath::Root()->Attr("x"); refl::AccessPath medium_path = refl::AccessPath::Root()->Attr("x")->ArrayItem(0); refl::AccessPath long_path = refl::AccessPath::Root()->Attr("x")->ArrayItem(0)->MapItem("z"); EXPECT_TRUE(short_path->IsPrefixOf(medium_path)); EXPECT_TRUE(short_path->IsPrefixOf(long_path)); EXPECT_TRUE(medium_path->IsPrefixOf(long_path)); EXPECT_FALSE(medium_path->IsPrefixOf(short_path)); EXPECT_FALSE(long_path->IsPrefixOf(medium_path)); EXPECT_FALSE(long_path->IsPrefixOf(short_path)); // Test non-prefix relationships refl::AccessPath branch1 = refl::AccessPath::Root()->Attr("x")->ArrayItem(0); refl::AccessPath branch2 = refl::AccessPath::Root()->Attr("x")->ArrayItem(1); EXPECT_FALSE(branch1->IsPrefixOf(branch2)); EXPECT_FALSE(branch2->IsPrefixOf(branch1)); EXPECT_FALSE(branch1->PathEqual(branch2)); // Test GetParent functionality auto parent = path4->GetParent(); EXPECT_TRUE(parent.has_value()); EXPECT_TRUE(parent.value()->PathEqual(path)); auto root_parent = root->GetParent(); EXPECT_FALSE(root_parent.has_value()); } struct TestObjWithFactory : public Object { Array items; int64_t count; explicit TestObjWithFactory(UnsafeInit) {} [[maybe_unused]] static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.TestObjWithFactory", TestObjWithFactory, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("items", &TestObjWithFactory::items, refl::default_factory( Function::FromTyped([]() -> Array { return Array(); }))) .def_ro("count", &TestObjWithFactory::count, refl::default_value(static_cast(0))); } struct TestObjWithAny : public Object { Any value; explicit TestObjWithAny(Any value) : value(std::move(value)) {} [[maybe_unused]] static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.TestObjWithAny", TestObjWithAny, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef().def(refl::init()).def_ro("value", &TestObjWithAny::value); } struct TestObjWithAnyView : public Object { Any value; explicit TestObjWithAnyView(AnyView value) : value(value) {} [[maybe_unused]] static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.TestObjWithAnyView", TestObjWithAnyView, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def(refl::init()) .def_ro("value", &TestObjWithAnyView::value); } TEST(Reflection, InitWithAny) { Function init = GetInitAttr("test.TestObjWithAny"); Any obj1 = init(42); ASSERT_TRUE(obj1.as() != nullptr); EXPECT_EQ(obj1.as()->value.cast(), 42); Any obj2 = init(3.14); ASSERT_TRUE(obj2.as() != nullptr); EXPECT_EQ(obj2.as()->value.cast(), 3.14); Any obj3 = init(String("hello")); ASSERT_TRUE(obj3.as() != nullptr); EXPECT_EQ(obj3.as()->value.cast(), "hello"); } TEST(Reflection, InitWithAnyView) { Function init = GetInitAttr("test.TestObjWithAnyView"); Any obj1 = init(42); ASSERT_TRUE(obj1.as() != nullptr); EXPECT_EQ(obj1.as()->value.cast(), 42); Any obj2 = init(3.14); ASSERT_TRUE(obj2.as() != nullptr); EXPECT_EQ(obj2.as()->value.cast(), 3.14); Any obj3 = init(String("hello")); ASSERT_TRUE(obj3.as() != nullptr); EXPECT_EQ(obj3.as()->value.cast(), "hello"); } TEST(Reflection, DefaultFactoryFlag) { const TVMFFIFieldInfo* info_items = reflection::GetFieldInfo("test.TestObjWithFactory", "items"); EXPECT_TRUE(info_items->flags & kTVMFFIFieldFlagBitMaskHasDefault); EXPECT_TRUE(info_items->flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory); const TVMFFIFieldInfo* info_count = reflection::GetFieldInfo("test.TestObjWithFactory", "count"); EXPECT_TRUE(info_count->flags & kTVMFFIFieldFlagBitMaskHasDefault); EXPECT_FALSE(info_count->flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory); } TEST(Reflection, DefaultFactoryCreation) { namespace refl = tvm::ffi::reflection; refl::ObjectCreator creator("test.TestObjWithFactory"); // Create two objects without providing "items" - each should get a fresh Array Any obj1 = creator(Map({{"count", static_cast(42)}})); Any obj2 = creator(Map({{"count", static_cast(99)}})); auto* p1 = obj1.as(); auto* p2 = obj2.as(); ASSERT_NE(p1, nullptr); ASSERT_NE(p2, nullptr); EXPECT_EQ(p1->count, 42); EXPECT_EQ(p2->count, 99); // Both should have empty arrays EXPECT_EQ(p1->items.size(), 0); EXPECT_EQ(p2->items.size(), 0); // Crucially, the arrays should be distinct objects (not aliased) EXPECT_NE(p1->items.get(), p2->items.get()); } TEST(Reflection, DefaultFactoryNotCalledWhenProvided) { namespace refl = tvm::ffi::reflection; refl::ObjectCreator creator("test.TestObjWithFactory"); Array custom_items; custom_items.push_back(TInt(1)); Any obj = creator(Map({{"items", custom_items}, {"count", static_cast(5)}})); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->items.size(), 1); EXPECT_EQ(p->count, 5); } // --------------------------------------------------------------------------- // Tests for auto-generated __ffi_init__ with init(false) / KwOnly(true) // --------------------------------------------------------------------------- struct TestAutoInitObj : public Object { int64_t a; int64_t b; int64_t c; int64_t d; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("test.AutoInit", TestAutoInitObj, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; // No refl::init<>() — auto-generates __ffi_init__ refl::ObjectDef() .def_rw("a", &TestAutoInitObj::a) .def_rw("b", &TestAutoInitObj::b, refl::init(false), refl::default_value(int64_t{42})) .def_rw("c", &TestAutoInitObj::c, refl::kw_only(true)) .def_rw("d", &TestAutoInitObj::d, refl::default_value(int64_t{99})); } TEST(Reflection, AutoInitPositional) { // Auto-generated init: positional args for non-kw-only init=True fields (a, d) // c is kw_only so it cannot be passed positionally. Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); // Positional: a=1, d=3; keyword: c=2 Any obj = auto_init(int64_t{1}, int64_t{3}, kwargs, String("c"), int64_t{2}); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->a, 1); EXPECT_EQ(p->b, 42); // init=False, gets default EXPECT_EQ(p->c, 2); // kw_only, passed via KWARGS EXPECT_EQ(p->d, 3); // init=True, 2nd positional } TEST(Reflection, AutoInitPartialPositional) { // Provide only a (position 0); c is required but missing → error Function auto_init = GetInitAttr("test.AutoInit"); EXPECT_THROW( { try { auto_init(int64_t{1}); } catch (const std::exception& e) { EXPECT_NE(std::string(e.what()).find("missing required"), std::string::npos); throw; } }, std::exception); } TEST(Reflection, AutoInitWithDefaults) { // Provide a positionally and c via KWARGS; d should use default 99 Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); Any obj = auto_init(int64_t{10}, kwargs, String("c"), int64_t{20}); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->a, 10); EXPECT_EQ(p->b, 42); // default EXPECT_EQ(p->c, 20); // provided via KWARGS EXPECT_EQ(p->d, 99); // default } TEST(Reflection, AutoInitKwargs) { Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); // Positional: a=1, then KWARGS: c=30, d=40 Any obj = auto_init(int64_t{1}, kwargs, String("c"), int64_t{30}, String("d"), int64_t{40}); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->a, 1); EXPECT_EQ(p->b, 42); // default EXPECT_EQ(p->c, 30); EXPECT_EQ(p->d, 40); } TEST(Reflection, AutoInitKwargsOnly) { Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); // No positional args, all via KWARGS Any obj = auto_init(kwargs, String("a"), int64_t{5}, String("c"), int64_t{15}, String("d"), int64_t{25}); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->a, 5); EXPECT_EQ(p->b, 42); EXPECT_EQ(p->c, 15); EXPECT_EQ(p->d, 25); } TEST(Reflection, AutoInitKwargsDuplicate) { Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); // a is provided positionally AND as kwarg → error EXPECT_THROW( { try { auto_init(int64_t{1}, kwargs, String("a"), int64_t{2}, String("c"), int64_t{3}); } catch (const std::exception& e) { EXPECT_NE(std::string(e.what()).find("multiple values"), std::string::npos); throw; } }, std::exception); } TEST(Reflection, AutoInitKwargsUnknown) { Function auto_init = GetInitAttr("test.AutoInit"); ObjectRef kwargs = Function::GetGlobalRequired("ffi.GetKwargsObject")().cast(); EXPECT_THROW( { try { auto_init(kwargs, String("a"), int64_t{1}, String("z"), int64_t{2}, String("c"), int64_t{3}); } catch (const std::exception& e) { EXPECT_NE(std::string(e.what()).find("unexpected keyword"), std::string::npos); throw; } }, std::exception); } TEST(Reflection, AutoInitFlagBits) { // Verify the flag bits are set correctly on the field info. const TVMFFIFieldInfo* fi_a = reflection::GetFieldInfo("test.AutoInit", "a"); EXPECT_FALSE(fi_a->flags & kTVMFFIFieldFlagBitMaskInitOff); EXPECT_FALSE(fi_a->flags & kTVMFFIFieldFlagBitMaskKwOnly); const TVMFFIFieldInfo* fi_b = reflection::GetFieldInfo("test.AutoInit", "b"); EXPECT_TRUE(fi_b->flags & kTVMFFIFieldFlagBitMaskInitOff); EXPECT_FALSE(fi_b->flags & kTVMFFIFieldFlagBitMaskKwOnly); EXPECT_TRUE(fi_b->flags & kTVMFFIFieldFlagBitMaskHasDefault); const TVMFFIFieldInfo* fi_c = reflection::GetFieldInfo("test.AutoInit", "c"); EXPECT_FALSE(fi_c->flags & kTVMFFIFieldFlagBitMaskInitOff); EXPECT_TRUE(fi_c->flags & kTVMFFIFieldFlagBitMaskKwOnly); const TVMFFIFieldInfo* fi_d = reflection::GetFieldInfo("test.AutoInit", "d"); EXPECT_FALSE(fi_d->flags & kTVMFFIFieldFlagBitMaskInitOff); EXPECT_FALSE(fi_d->flags & kTVMFFIFieldFlagBitMaskKwOnly); EXPECT_TRUE(fi_d->flags & kTVMFFIFieldFlagBitMaskHasDefault); } // Simple auto-init test: all fields init=True, no Init/KwOnly traits struct TestAutoInitSimpleObj : public Object { int64_t x; int64_t y; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("test.AutoInitSimple", TestAutoInitSimpleObj, Object); }; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_rw("x", &TestAutoInitSimpleObj::x) .def_rw("y", &TestAutoInitSimpleObj::y); } TEST(Reflection, AutoInitSimple) { Function auto_init = GetInitAttr("test.AutoInitSimple"); Any obj = auto_init(int64_t{10}, int64_t{20}); auto* p = obj.as(); ASSERT_NE(p, nullptr); EXPECT_EQ(p->x, 10); EXPECT_EQ(p->y, 20); } TEST(Reflection, AutoInitSimpleTooManyArgs) { Function auto_init = GetInitAttr("test.AutoInitSimple"); EXPECT_THROW(auto_init(int64_t{1}, int64_t{2}, int64_t{3}), std::exception); } // --------------------------------------------------------------------------- // overload_cast — pick an overload by prefix-matching its parameter types. // --------------------------------------------------------------------------- namespace overload_cast_test { struct Cat {}; struct Dog {}; // Pet: each Feed overload has a unique first arg plus a trailing context // (`int amount`) that the caller doesn't have to spell out. Get is const // vs non-const overloaded with identical params (selected via const_ tag). struct Pet { int Feed(const Cat*, int amount) { return 100 + amount; } int Feed(const Dog*, int amount) { return 200 + amount; } int Get(int x) { return 4000 + x; } int Get(int x) const { return 5000 + x; } }; // Mix: overloads share a leading prefix — spelling more parameters // disambiguates against the longer variant. struct Mix { int Run(int, int) { return 1000; } int Run(int, double) { return 2000; } int Run(int, int, int) { return 3000; } }; int FreeFeed(const Cat*, int x) { return 6000 + x; } int FreeFeed(const Dog*, int x) { return 7000 + x; } template struct CallVia { template static auto Run(Self&& self, Args&&... args) { return (std::forward(self).*Method)(std::forward(args)...); } }; } // namespace overload_cast_test TEST(OverloadCast, PrefixMatch) { using namespace overload_cast_test; namespace refl = tvm::ffi::reflection; Pet p; Cat cat; Dog dog; // (a) Member with unique first arg per overload: spelling only the // disambiguating prefix picks the overload and deduces the // trailing `int amount` from the picked signature. auto p_cat = refl::overload_cast(&Pet::Feed); static_assert(std::is_same_v, "prefix match must deduce trailing arg types"); EXPECT_EQ((p.*p_cat)(&cat, 7), 107); auto p_dog = refl::overload_cast(&Pet::Feed); EXPECT_EQ((p.*p_dog)(&dog, 12), 212); // (b) Free function with the same shape — trailing arg deduced. auto p_free_cat = refl::overload_cast(&FreeFeed); EXPECT_EQ(p_free_cat(&cat, 7), 6007); auto p_free_dog = refl::overload_cast(&FreeFeed); EXPECT_EQ(p_free_dog(&dog, 7), 7007); } TEST(OverloadCast, AmbiguousPrefixRequiresMoreSpelling) { using namespace overload_cast_test; namespace refl = tvm::ffi::reflection; Mix m; // Mix::Run has three overloads: // Run(int, int), Run(int, double), Run(int, int, int) // Spelling only would be ambiguous (all three start with int). // Spelling enough parameters to identify exactly one overload picks it. EXPECT_EQ((m.*refl::overload_cast(&Mix::Run))(0, 0), 1000); EXPECT_EQ((m.*refl::overload_cast(&Mix::Run))(0, 0.0), 2000); EXPECT_EQ((m.*refl::overload_cast(&Mix::Run))(0, 0, 0), 3000); } TEST(OverloadCast, ConstQualifiedMember) { using namespace overload_cast_test; namespace refl = tvm::ffi::reflection; Pet p; const Pet& cp = p; // Non-const overload — no tag. EXPECT_EQ((p.*refl::overload_cast(&Pet::Get))(7), 4007); // Const overload — const_ tag required (even when the const overload // is the only one with that signature, address-of-overload alone // cannot select it from the operator() overload set). EXPECT_EQ((cp.*refl::overload_cast(&Pet::Get, refl::const_))(7), 5007); } TEST(OverloadCast, NonTypeTemplateArgument) { using namespace overload_cast_test; namespace refl = tvm::ffi::reflection; Pet p; Mix m; Cat cat; // Prefix match composed as a non-type template argument. EXPECT_EQ((CallVia(&Pet::Feed)>::Run(p, &cat, 7)), 107); // Disambiguated 3-arg overload as a non-type template argument. EXPECT_EQ((CallVia(&Mix::Run)>::Run(m, 0, 0, 0)), 3000); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_rvalue_ref.cc000066400000000000000000000065101521067262500205770ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(RValueRef, Basic) { auto append = Function::FromTyped([](RValueRef> ref, int val, bool is_unique) -> Array { Array arr = *std::move(ref); EXPECT_EQ(arr.unique(), is_unique); arr.push_back(val); return arr; }); auto a = append(RValueRef(Array({1, 2})), 3, true).cast>(); EXPECT_EQ(a.size(), 3); a = append(RValueRef(std::move(a)), 4, true).cast>(); EXPECT_EQ(a.size(), 4); // pass in lvalue instead, the append still will succeed but array will not be unique a = append(a, 5, false).cast>(); EXPECT_EQ(a.size(), 5); } TEST(RValueRef, ParamChecking) { // try decution // NOLINTNEXTLINE(performance-unnecessary-value-param) Function fadd1 = Function::FromTyped([](TInt a) -> int64_t { return a->value + 1; }); // convert that triggers error EXPECT_THROW( { try { fadd1(RValueRef(TInt(1))); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: test.Int) -> int`. " "Expected `test.Int` but got `ObjectRValueRef`"); throw; } }, ::tvm::ffi::Error); Function fadd2 = Function::FromTyped([](RValueRef> a) -> int { Array arr = *std::move(a); return arr[0] + 1; }); // convert that triggers error EXPECT_THROW( { try { fadd2(RValueRef(Array({1, 2.2}))); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ( error.message(), "Mismatched type on argument #0 when calling: `(0: RValueRef>) -> int`. " "Expected `RValueRef>` but got `RValueRef`"); throw; } }, ::tvm::ffi::Error); // triggered a rvalue based conversion Function func3 = Function::FromTyped([](RValueRef a) -> String { TPrimExpr expr = *std::move(a); return expr->dtype; }); // EXPECT_EQ(func3(RValueRef(String("int32"))).cast(), "int32"); // triggered a lvalue based conversion // EXPECT_EQ(func3(String("int32")).cast(), "int32"); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_shape.cc000066400000000000000000000053731521067262500175530ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include namespace { using namespace tvm::ffi; TEST(Shape, Basic) { Shape shape = Shape({1, 2, 3}); EXPECT_EQ(shape.size(), 3); EXPECT_EQ(shape[0], 1); EXPECT_EQ(shape[1], 2); EXPECT_EQ(shape[2], 3); Shape shape2 = Shape(Array({4, 5, 6, 7})); EXPECT_EQ(shape2.size(), 4); EXPECT_EQ(shape2[0], 4); EXPECT_EQ(shape2[1], 5); EXPECT_EQ(shape2[2], 6); EXPECT_EQ(shape2[3], 7); std::vector vec = {8, 9, 10}; Shape shape3 = Shape(std::move(vec)); EXPECT_EQ(shape3.size(), 3); EXPECT_EQ(shape3[0], 8); EXPECT_EQ(shape3[1], 9); EXPECT_EQ(shape3[2], 10); EXPECT_EQ(shape3.Product(), 8 * 9 * 10); Shape shape4 = Shape(); EXPECT_EQ(shape4.size(), 0); EXPECT_EQ(shape4.Product(), 1); } TEST(Shape, AnyConvert) { Shape shape0 = Shape({1, 2, 3}); Any any0 = shape0; auto shape1 = any0.cast(); EXPECT_EQ(shape1.size(), 3); EXPECT_EQ(shape1[0], 1); EXPECT_EQ(shape1[1], 2); EXPECT_EQ(shape1[2], 3); Array arr({1, 2}); AnyView any_view0 = arr; auto shape2 = any_view0.cast(); EXPECT_EQ(shape2.size(), 2); EXPECT_EQ(shape2[0], 1); EXPECT_EQ(shape2[1], 2); } TEST(Shape, ShapeView) { Shape shape = Shape({1, 2, 3}); ShapeView shape_view = shape; EXPECT_EQ(shape_view.size(), 3); EXPECT_EQ(shape_view[0], 1); EXPECT_EQ(shape_view[1], 2); EXPECT_EQ(shape_view[2], 3); std::vector data = {4, 5, 6}; ShapeView view_from_data(data.data(), data.size()); EXPECT_EQ(view_from_data.size(), 3); EXPECT_EQ(view_from_data[0], 4); EXPECT_EQ(view_from_data[1], 5); EXPECT_EQ(view_from_data[2], 6); std::initializer_list init = {7, 8, 9}; ShapeView view_from_init = init; EXPECT_EQ(view_from_init.size(), 3); EXPECT_EQ(view_from_init[0], 7); EXPECT_EQ(view_from_init[1], 8); EXPECT_EQ(view_from_init[2], 9); EXPECT_EQ(view_from_init.Product(), 7 * 8 * 9); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_string.cc000066400000000000000000000472401521067262500177600ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include namespace { using namespace tvm::ffi; TEST(String, MoveFromStd) { using namespace std; string source = "this is a string"; string expect = source; String s(std::move(source)); string copy = string(s); EXPECT_EQ(copy, expect); EXPECT_EQ(source.size(), 0); // NOLINT(bugprone-use-after-move) } TEST(String, CopyFromStd) { using namespace std; string source = "this is a string"; string expect = source; // NOLINT(performance-unnecessary-copy-initialization) String s{source}; string copy = string(s); EXPECT_EQ(copy, expect); EXPECT_EQ(source.size(), expect.size()); } TEST(String, Assignment) { using namespace std; String s{string{"hello"}}; s = string{"world"}; EXPECT_EQ(s == "world", true); string s2{"world2"}; s = std::move(s2); EXPECT_EQ(s == "world2", true); Any r; r = String("hello"); EXPECT_EQ(r != nullptr, true); } TEST(String, empty) { using namespace std; String s{"hello"}; EXPECT_EQ(s.empty(), false); s = std::string(""); EXPECT_EQ(s.empty(), true); } TEST(String, Comparisons) { using namespace std; string source = "a string"; string mismatch = "a string but longer"; String s{"a string"}; String m{mismatch}; EXPECT_EQ("a str" >= s, false); EXPECT_EQ(s == source, true); EXPECT_EQ(s == mismatch, false); EXPECT_EQ(s == source.data(), true); EXPECT_EQ(s == mismatch.data(), false); EXPECT_EQ(s < m, source < mismatch); EXPECT_EQ(s > m, source > mismatch); EXPECT_EQ(s <= m, source <= mismatch); EXPECT_EQ(s >= m, source >= mismatch); EXPECT_EQ(s == m, source == mismatch); EXPECT_EQ(s != m, source != mismatch); EXPECT_EQ(m < s, mismatch < source); EXPECT_EQ(m > s, mismatch > source); EXPECT_EQ(m <= s, mismatch <= source); EXPECT_EQ(m >= s, mismatch >= source); EXPECT_EQ(m == s, mismatch == source); EXPECT_EQ(m != s, mismatch != source); } TEST(String, Compare) { // string compare const char* String s{"hello"}; EXPECT_EQ(s.compare("hello"), 0); EXPECT_EQ(s.compare(String("hello")), 0); EXPECT_EQ(s.compare("hallo"), 1); EXPECT_EQ(s.compare(String("hallo")), 1); EXPECT_EQ(s.compare("hfllo"), -1); EXPECT_EQ(s.compare(String("hfllo")), -1); // s is longer EXPECT_EQ(s.compare("hell"), 1); EXPECT_EQ(s.compare(String("hell")), 1); // s is shorter EXPECT_EQ(s.compare("hello world"), -1); EXPECT_EQ(s.compare(String("helloworld")), -1); } // Check '\0' handling TEST(String, NullByteHandling) { using namespace std; // Ensure string still compares equal if it contains '\0'. string v1 = "hello world"; size_t v1_size = v1.size(); v1[5] = '\0'; EXPECT_EQ(v1[5], '\0'); EXPECT_EQ(v1.size(), v1_size); String str_v1{v1}; EXPECT_EQ(str_v1.compare(v1), 0); EXPECT_EQ(str_v1.size(), v1_size); // Ensure bytes after '\0' are taken into account for mismatches. string v2 = "aaa one"; string v3 = "aaa two"; v2[3] = '\0'; v3[3] = '\0'; String str_v2{v2}; String str_v3{v3}; EXPECT_EQ(str_v2.compare(str_v3), -1); EXPECT_EQ(str_v2.size(), 7); // strcmp won't be able to detect the mismatch EXPECT_EQ(strcmp(v2.data(), v3.data()), 0); // string::compare can handle \0 since it knows size EXPECT_LT(v2.compare(v3), 0); // If there is mismatch before '\0', should still handle it. string v4 = "acc one"; string v5 = "abb two"; v4[3] = '\0'; v5[3] = '\0'; String str_v4{v4}; String str_v5{v5}; EXPECT_GT(str_v4.compare(str_v5), 0); EXPECT_EQ(str_v4.size(), 7); // strcmp is able to detect the mismatch EXPECT_GT(strcmp(v4.data(), v5.data()), 0); // string::compare can handle \0 since it knows size EXPECT_GT(v4.compare(v5), 0); } TEST(String, CompareSameMemoryRegionDifferentSize) { using namespace std; string source = "a string"; String str_source{source}; char* memory = const_cast(str_source.data()); EXPECT_EQ(str_source.compare(memory), 0); // This changes the string size memory[2] = '\0'; // memory is logically shorter now EXPECT_GT(str_source.compare(memory), 0); } TEST(String, compare) { using namespace std; constexpr auto mismatch1_cstr = "a string but longer"; string source = "a string"; string mismatch1 = mismatch1_cstr; string mismatch2 = "a strin"; string mismatch3 = "a b"; string mismatch4 = "a t"; String str_source{source}; String str_mismatch1{mismatch1_cstr}; String str_mismatch2{mismatch2}; String str_mismatch3{mismatch3}; String str_mismatch4{mismatch4}; // compare with string EXPECT_EQ(str_source.compare(source), 0); EXPECT_TRUE(str_source == source); EXPECT_TRUE(source == str_source); EXPECT_TRUE(str_source <= source); EXPECT_TRUE(source <= str_source); EXPECT_TRUE(str_source >= source); EXPECT_TRUE(source >= str_source); EXPECT_LT(str_source.compare(mismatch1), 0); EXPECT_TRUE(str_source < mismatch1); EXPECT_TRUE(mismatch1 != str_source); EXPECT_GT(str_source.compare(mismatch2), 0); EXPECT_TRUE(str_source > mismatch2); EXPECT_TRUE(mismatch2 < str_source); EXPECT_GT(str_source.compare(mismatch3), 0); EXPECT_TRUE(str_source > mismatch3); EXPECT_LT(str_source.compare(mismatch4), 0); EXPECT_TRUE(str_source < mismatch4); EXPECT_TRUE(mismatch4 > str_source); // compare with char* EXPECT_EQ(str_source.compare(source.data()), 0); EXPECT_TRUE(str_source == source.data()); EXPECT_TRUE(source.data() == str_source); EXPECT_TRUE(str_source <= source.data()); EXPECT_TRUE(source <= str_source.data()); EXPECT_TRUE(str_source >= source.data()); EXPECT_TRUE(source >= str_source.data()); EXPECT_LT(str_source.compare(mismatch1.data()), 0); EXPECT_TRUE(str_source < mismatch1.data()); EXPECT_TRUE(str_source != mismatch1.data()); EXPECT_TRUE(mismatch1.data() != str_source); EXPECT_GT(str_source.compare(mismatch2.data()), 0); EXPECT_TRUE(str_source > mismatch2.data()); EXPECT_TRUE(mismatch2.data() < str_source); EXPECT_GT(str_source.compare(mismatch3.data()), 0); EXPECT_TRUE(str_source > mismatch3.data()); EXPECT_LT(str_source.compare(mismatch4.data()), 0); EXPECT_TRUE(str_source < mismatch4.data()); EXPECT_TRUE(mismatch4.data() > str_source); // compare with String EXPECT_LT(str_source.compare(str_mismatch1), 0); EXPECT_TRUE(str_source < str_mismatch1); EXPECT_GT(str_source.compare(str_mismatch2), 0); EXPECT_TRUE(str_source > str_mismatch2); EXPECT_GT(str_source.compare(str_mismatch3), 0); EXPECT_TRUE(str_source > str_mismatch3); EXPECT_LT(str_source.compare(str_mismatch4), 0); EXPECT_TRUE(str_source < str_mismatch4); } TEST(String, CString) { using namespace std; string source = "this is a string"; string mismatch = "mismatch"; String s{source}; EXPECT_EQ(std::strcmp(s.c_str(), source.data()), 0); EXPECT_NE(std::strcmp(s.c_str(), mismatch.data()), 0); } TEST(String, hash) { using namespace std; string source = "this is a string"; String s{source}; std::hash()(s); std::unordered_map map; String k1{string{"k1"}}; string v1{"v1"}; String k2{string{"k2"}}; string v2{"v2"}; map[k1] = v1; map[k2] = v2; EXPECT_EQ(map[k1], v1); EXPECT_EQ(map[k2], v2); } TEST(String, Cast) { using namespace std; string source = "this is a string"; String s{source}; Any r = s; String s2 = r.cast(); } TEST(String, Concat) { String s1("hello"); String s2("world"); std::string s3("world"); String res1 = s1 + s2; String res2 = s1 + s3; String res3 = s3 + s1; String res4 = s1 + "world"; String res5 = "world" + s1; EXPECT_EQ(res1.compare("helloworld"), 0); EXPECT_EQ(res2.compare("helloworld"), 0); EXPECT_EQ(res3.compare("worldhello"), 0); EXPECT_EQ(res4.compare("helloworld"), 0); EXPECT_EQ(res5.compare("worldhello"), 0); String storage_scope; String res = "The input storage scope \"" + storage_scope + "\" is invalid."; EXPECT_EQ(res.compare("The input storage scope \"\" is invalid."), 0); } TEST(String, Any) { // test anyview promotion to any AnyView view = "hello"; EXPECT_EQ(view.type_index(), TypeIndex::kTVMFFIRawStr); Any b = view; EXPECT_EQ(b.type_index(), TypeIndex::kTVMFFISmallStr); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(b.as().value(), "hello"); EXPECT_TRUE(b.as().has_value()); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(b.try_cast().value(), "hello"); std::string s_world = "world"; view = s_world; // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(view.try_cast().value(), "world"); String s{"hello"}; Any a = s; EXPECT_EQ(a.type_index(), TypeIndex::kTVMFFISmallStr); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(a.as().value(), "hello"); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(a.try_cast().value(), "hello"); Any c = "long string very long"; EXPECT_EQ(c.type_index(), TypeIndex::kTVMFFIStr); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(c.as().value(), "long string very long"); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(c.try_cast().value(), "long string very long"); } TEST(String, Bytes) { Bytes b0; EXPECT_EQ(b0.size(), 0); EXPECT_EQ(b0.operator std::string(), ""); // explicitly test zero element std::string s = {'\0', 'a', 'b', 'c'}; Bytes b = s; EXPECT_EQ(b.size(), 4); EXPECT_EQ(b.operator std::string(), s); TVMFFIByteArray arr{s.data(), static_cast(s.size())}; Bytes b2 = arr; EXPECT_EQ(b2.size(), 4); EXPECT_EQ(b2.operator std::string(), s); } TEST(String, BytesAny) { std::string s = {'\0', 'a', 'b', 'c'}; TVMFFIByteArray arr{s.data(), static_cast(s.size())}; AnyView view = &arr; EXPECT_EQ(view.type_index(), TypeIndex::kTVMFFIByteArrayPtr); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(view.try_cast().value().operator std::string(), s); Any b = view; EXPECT_EQ(b.type_index(), TypeIndex::kTVMFFISmallBytes); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(b.try_cast().value().operator std::string(), s); EXPECT_EQ(b.cast(), s); std::string s2 = "hello long long long string"; s2[0] = '\0'; Any b2 = Bytes(s2); EXPECT_EQ(b2.type_index(), TypeIndex::kTVMFFIBytes); EXPECT_EQ(b2.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s2); EXPECT_EQ(b2.cast(), s2); } TEST(String, StdString) { std::string s1 = "test_string"; AnyView view1 = s1; EXPECT_EQ(view1.type_index(), TypeIndex::kTVMFFIRawStr); EXPECT_EQ(view1.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); TVMFFIByteArray arr1{s1.data(), static_cast(s1.size())}; AnyView view2 = &arr1; EXPECT_EQ(view2.type_index(), TypeIndex::kTVMFFIByteArrayPtr); EXPECT_EQ(view2.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); Bytes bytes1 = s1; AnyView view3 = bytes1; EXPECT_EQ(view3.type_index(), TypeIndex::kTVMFFIBytes); EXPECT_EQ(view3.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); String string1 = s1; AnyView view4 = string1; EXPECT_EQ(view4.type_index(), TypeIndex::kTVMFFIStr); EXPECT_EQ(view4.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); // Test with Any Any any1 = s1; EXPECT_EQ(any1.type_index(), TypeIndex::kTVMFFIStr); EXPECT_EQ(any1.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); Any any2 = &arr1; EXPECT_EQ(any2.type_index(), TypeIndex::kTVMFFIBytes); EXPECT_EQ(any2.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); Any any3 = bytes1; EXPECT_EQ(any3.type_index(), TypeIndex::kTVMFFIBytes); EXPECT_EQ(any3.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); Any any4 = string1; EXPECT_EQ(any4.type_index(), TypeIndex::kTVMFFIStr); EXPECT_EQ(any4.try_cast().value(), // NOLINT(bugprone-unchecked-optional-access) s1); } TEST(String, CAPIAccessor) { using namespace std; String s{"hello"}; TVMFFIByteArray arr{s.data(), s.size()}; EXPECT_EQ(arr.size, 5); EXPECT_EQ(std::string(arr.data, arr.size), "hello"); } TEST(String, BytesHash) { std::vector data1(10); std::vector data2(11); for (size_t i = 0; i < data1.size(); ++i) { data1[i] = static_cast(i); } char* data1_ptr = reinterpret_cast(data1.data()); char* data2_ptr = reinterpret_cast(data2.data()) + 1; std::memcpy(data2_ptr, data1.data(), data1.size() * sizeof(int64_t)); // has of aligned and unaligned data should be the same uint64_t hash1 = details::StableHashBytes(data1_ptr, data1.size() * sizeof(int64_t)); uint64_t hash2 = details::StableHashBytes(data2_ptr, data1.size() * sizeof(int64_t)); EXPECT_EQ(hash1, hash2); } TEST(String, StdHash) { String s1 = "a"; String s2(std::string("a")); EXPECT_EQ(std::hash()(s1), std::hash()(s2)); Bytes s3("a", 1); Bytes s4(std::string("a")); EXPECT_EQ(std::hash()(s3), std::hash()(s4)); } TEST(String, Find) { String s{"hello world"}; EXPECT_EQ(s.find("world"), 6); EXPECT_EQ(s.find("hello"), 0); EXPECT_EQ(s.find("o"), 4); EXPECT_EQ(s.find("o", 5), 7); EXPECT_EQ(s.find("notfound"), String::npos); EXPECT_EQ(s.find(""), 0); EXPECT_EQ(s.find("", 5), 5); EXPECT_EQ(s.find("", 11), 11); EXPECT_EQ(s.find("", 20), String::npos); String pattern{"world"}; EXPECT_EQ(s.find(pattern), 6); String empty{""}; EXPECT_EQ(empty.find("x"), String::npos); EXPECT_EQ(empty.find(""), 0); } TEST(String, Substr) { String s{"hello world"}; EXPECT_EQ(s.substr(0, 5), "hello"); EXPECT_EQ(s.substr(6, 5), "world"); EXPECT_EQ(s.substr(6), "world"); EXPECT_EQ(s.substr(0), "hello world"); EXPECT_EQ(s.substr(11), ""); EXPECT_EQ(s.substr(0, 0), ""); EXPECT_THROW(s.substr(12), std::out_of_range); EXPECT_THROW(s.substr(100), std::out_of_range); String empty{""}; EXPECT_EQ(empty.substr(0), ""); EXPECT_THROW(empty.substr(1), std::out_of_range); } TEST(String, StartsWith) { String s{"hello world"}; EXPECT_TRUE(s.starts_with("hello")); EXPECT_TRUE(s.starts_with("h")); EXPECT_TRUE(s.starts_with("")); EXPECT_TRUE(s.starts_with(String{"hello"})); EXPECT_TRUE(s.starts_with(std::string_view{"hello"})); EXPECT_FALSE(s.starts_with("world")); EXPECT_FALSE(s.starts_with("Hello")); EXPECT_FALSE(s.starts_with("hello world extra")); EXPECT_FALSE(s.starts_with(std::string_view{"world"})); String empty{""}; EXPECT_TRUE(empty.starts_with("")); EXPECT_TRUE(empty.starts_with(std::string_view{""})); EXPECT_FALSE(empty.starts_with("x")); String single{"x"}; EXPECT_TRUE(single.starts_with("x")); EXPECT_TRUE(single.starts_with("")); EXPECT_FALSE(single.starts_with("xy")); } TEST(String, EndsWith) { String s{"hello world"}; EXPECT_TRUE(s.ends_with("world")); EXPECT_TRUE(s.ends_with("d")); EXPECT_TRUE(s.ends_with("")); EXPECT_TRUE(s.ends_with(String{"world"})); EXPECT_TRUE(s.ends_with(std::string_view{"world"})); EXPECT_FALSE(s.ends_with("hello")); EXPECT_FALSE(s.ends_with("World")); EXPECT_FALSE(s.ends_with("extra hello world")); EXPECT_FALSE(s.ends_with(std::string_view{"hello"})); String empty{""}; EXPECT_TRUE(empty.ends_with("")); EXPECT_TRUE(empty.ends_with(std::string_view{""})); EXPECT_FALSE(empty.ends_with("x")); String single{"x"}; EXPECT_TRUE(single.ends_with("x")); EXPECT_TRUE(single.ends_with("")); EXPECT_FALSE(single.ends_with("yx")); } TEST(String, Split) { String s{"a,b,c"}; auto parts = s.Split(','); ASSERT_EQ(parts.size(), 3); EXPECT_EQ(parts[0], "a"); EXPECT_EQ(parts[1], "b"); EXPECT_EQ(parts[2], "c"); // No delimiter present String s2{"hello"}; auto parts2 = s2.Split(','); ASSERT_EQ(parts2.size(), 1); EXPECT_EQ(parts2[0], "hello"); // Empty string String s3{""}; auto parts3 = s3.Split(','); ASSERT_EQ(parts3.size(), 1); EXPECT_EQ(parts3[0], ""); // Delimiter at boundaries String s4{",a,b,"}; auto parts4 = s4.Split(','); ASSERT_EQ(parts4.size(), 4); EXPECT_EQ(parts4[0], ""); EXPECT_EQ(parts4[1], "a"); EXPECT_EQ(parts4[2], "b"); EXPECT_EQ(parts4[3], ""); // Consecutive delimiters String s5{"a,,b"}; auto parts5 = s5.Split(','); ASSERT_EQ(parts5.size(), 3); EXPECT_EQ(parts5[0], "a"); EXPECT_EQ(parts5[1], ""); EXPECT_EQ(parts5[2], "b"); } TEST(String, EscapeStringJSON) { // Basic escaping String s1{"hello"}; EXPECT_EQ(EscapeStringJSON(s1), "\"hello\""); // Special characters String s2{"line1\nline2\ttab"}; EXPECT_EQ(EscapeStringJSON(s2), "\"line1\\nline2\\ttab\""); // Backslash and quote String s3{"a\\b\"c"}; EXPECT_EQ(EscapeStringJSON(s3), "\"a\\\\b\\\"c\""); // Control characters String s4{std::string("a\x01\x1f z", 5)}; EXPECT_EQ(EscapeStringJSON(s4), "\"a\\u0001\\u001f z\""); } TEST(String, EscapedStringPyBasic) { // Plain ASCII String s1{"hello world"}; EXPECT_EQ(EscapedStringPy(s1), "\"hello world\""); // C escape sequences String s2{"a\nb\tc\r"}; EXPECT_EQ(EscapedStringPy(s2), "\"a\\nb\\tc\\r\""); // Backslash and quote String s3{"a\\b\"c"}; EXPECT_EQ(EscapedStringPy(s3), "\"a\\\\b\\\"c\""); } TEST(String, EscapedStringPyControlChars) { // Control characters -> \xNN String s1{std::string("\x01\x02\x7f", 3)}; String result = EscapedStringPy(s1); EXPECT_EQ(result, "\"\\x01\\x02\\x7f\""); } TEST(String, EscapedStringPyANSI) { // ANSI escape: ESC[31m (red) String s1{std::string("\x1b[31mred\x1b[0m", 12)}; String result = EscapedStringPy(s1); EXPECT_EQ(result, "\"\\x1b[31mred\\x1b[0m\""); // ANSI erase line: ESC[K String s2{std::string("\x1b[K", 3)}; EXPECT_EQ(EscapedStringPy(s2), "\"\\x1b[K\""); } TEST(String, EscapedStringPyUTF8) { // 2-byte: U+00E9 (é) = C3 A9 String s1{std::string("\xc3\xa9", 2)}; EXPECT_EQ(EscapedStringPy(s1), "\"\\u00e9\""); // 3-byte: U+4E16 (世) = E4 B8 96 String s2{std::string("\xe4\xb8\x96", 3)}; EXPECT_EQ(EscapedStringPy(s2), "\"\\u4e16\""); // 4-byte: U+1F600 (😀) = F0 9F 98 80 String s3{std::string("\xf0\x9f\x98\x80", 4)}; EXPECT_EQ(EscapedStringPy(s3), "\"\\U0001f600\""); } TEST(String, EscapedStringPyMalformedUTF8) { // Lone continuation byte -> \xNN fallback String s1{std::string("\x80", 1)}; EXPECT_EQ(EscapedStringPy(s1), "\"\\x80\""); // 2-byte leader followed by non-continuation -> fallback for leader String s2{std::string("\xc3\x20", 2)}; String result2 = EscapedStringPy(s2); EXPECT_EQ(result2, "\"\\xc3 \""); // 3-byte leader with bad continuation -> fallback for leader String s3{std::string("\xe4\xb8\x20", 3)}; String result3 = EscapedStringPy(s3); EXPECT_EQ(result3, "\"\\xe4\\xb8 \""); // Truncated 2-byte at end of string String s4{std::string("\xc3", 1)}; EXPECT_EQ(EscapedStringPy(s4), "\"\\xc3\""); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_tensor.cc000066400000000000000000000372621521067262500177670ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include namespace { using namespace tvm::ffi; struct CPUNDAlloc { void AllocData(DLTensor* tensor) { tensor->data = malloc(GetDataSize(*tensor)); } void FreeData(DLTensor* tensor) { free(tensor->data); } }; inline Tensor Empty(const Shape& shape, DLDataType dtype, DLDevice device) { return Tensor::FromNDAlloc(CPUNDAlloc(), shape, dtype, device); } inline Tensor EmptyStrided(const Shape& shape, const Shape& strides, DLDataType dtype, DLDevice device) { return Tensor::FromNDAllocStrided(CPUNDAlloc(), shape, strides, dtype, device); } int TestEnvTensorAllocator(DLTensor* prototype, TVMFFIObjectHandle* out) { Shape shape(prototype->shape, prototype->shape + prototype->ndim); Tensor nd = Empty(shape, prototype->dtype, prototype->device); *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(nd)); return 0; } int TestEnvTensorAllocatorError(DLTensor* prototype, TVMFFIObjectHandle* out) { TVMFFIErrorSetRaisedFromCStr("RuntimeError", "TestEnvTensorAllocatorError"); return -1; } TEST(Tensor, Basic) { Tensor nd = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); Shape shape = nd.shape(); Shape strides = nd.strides(); EXPECT_EQ(shape.size(), 3); EXPECT_EQ(shape[0], 1); EXPECT_EQ(shape[1], 2); EXPECT_EQ(shape[2], 3); EXPECT_EQ(strides.size(), 3); EXPECT_EQ(strides[0], 6); EXPECT_EQ(strides[1], 3); EXPECT_EQ(strides[2], 1); EXPECT_EQ(nd.dtype(), DLDataType({kDLFloat, 32, 1})); for (int64_t i = 0; i < shape.Product(); ++i) { reinterpret_cast(nd.data_ptr())[i] = static_cast(i); } EXPECT_EQ(nd.numel(), 6); EXPECT_EQ(nd.ndim(), 3); EXPECT_EQ(nd.data_ptr(), nd.GetDLTensorPtr()->data); Any any0 = nd; Tensor nd2 = any0.as().value(); // NOLINT(bugprone-unchecked-optional-access) EXPECT_EQ(nd2.dtype(), DLDataType({kDLFloat, 32, 1})); for (int64_t i = 0; i < shape.Product(); ++i) { EXPECT_EQ(reinterpret_cast(nd2.data_ptr())[i], i); } EXPECT_EQ(nd.IsContiguous(), true); EXPECT_EQ(nd2.use_count(), 3); Tensor nd3 = EmptyStrided({2, 3}, {1, 2}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); Shape shape3 = nd3.shape(); Shape strides3 = nd3.strides(); EXPECT_EQ(shape3.size(), 2); EXPECT_EQ(shape3[0], 2); EXPECT_EQ(shape3[1], 3); EXPECT_EQ(strides3.size(), 2); EXPECT_EQ(strides3[0], 1); EXPECT_EQ(strides3[1], 2); } TEST(Tensor, EmptyTensorIsContiguous) { // An empty tensor (any shape dim == 0) is trivially contiguous regardless of // stride values. This matches NumPy / PyTorch semantics. // Use strides that would normally fail the contiguity check to verify the // early-return path in IsContiguous(). Tensor nd = EmptyStrided({4, 0, 4}, {0, 0, 0}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0})); EXPECT_EQ(nd.numel(), 0); EXPECT_EQ(nd.IsContiguous(), true); EXPECT_EQ(nd.is_contiguous(), true); } TEST(Tensor, DLPack) { Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLInt, 16, 1}), DLDevice({kDLCPU, 0})); DLManagedTensor* dlpack = tensor.ToDLPack(); EXPECT_EQ(dlpack->dl_tensor.ndim, 3); EXPECT_EQ(dlpack->dl_tensor.shape[0], 1); EXPECT_EQ(dlpack->dl_tensor.shape[1], 2); EXPECT_EQ(dlpack->dl_tensor.shape[2], 3); EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLInt); EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 16); EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1); EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU); EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0); EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0); EXPECT_EQ(dlpack->dl_tensor.strides[0], 6); EXPECT_EQ(dlpack->dl_tensor.strides[1], 3); EXPECT_EQ(dlpack->dl_tensor.strides[2], 1); EXPECT_EQ(tensor.use_count(), 2); { Tensor tensor2 = Tensor::FromDLPack(dlpack); EXPECT_EQ(tensor2.use_count(), 1); EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr()); EXPECT_EQ(tensor.use_count(), 2); EXPECT_EQ(tensor2.use_count(), 1); } EXPECT_EQ(tensor.use_count(), 1); } TEST(Tensor, DLPackVersioned) { DLDataType dtype = DLDataType({kDLFloat4_e2m1fn, 4, 1}); EXPECT_EQ(GetDataSize(2, dtype), 2 * 4 / 8); Tensor tensor = Empty({2}, dtype, DLDevice({kDLCPU, 0})); DLManagedTensorVersioned* dlpack = tensor.ToDLPackVersioned(); EXPECT_EQ(dlpack->version.major, DLPACK_MAJOR_VERSION); EXPECT_EQ(dlpack->version.minor, DLPACK_MINOR_VERSION); EXPECT_EQ(dlpack->dl_tensor.ndim, 1); EXPECT_EQ(dlpack->dl_tensor.shape[0], 2); EXPECT_EQ(dlpack->dl_tensor.dtype.code, kDLFloat4_e2m1fn); EXPECT_EQ(dlpack->dl_tensor.dtype.bits, 4); EXPECT_EQ(dlpack->dl_tensor.dtype.lanes, 1); EXPECT_EQ(dlpack->dl_tensor.device.device_type, kDLCPU); EXPECT_EQ(dlpack->dl_tensor.device.device_id, 0); EXPECT_EQ(dlpack->dl_tensor.byte_offset, 0); EXPECT_EQ(dlpack->dl_tensor.strides[0], 1); EXPECT_EQ(tensor.use_count(), 2); { Tensor tensor2 = Tensor::FromDLPackVersioned(dlpack); EXPECT_EQ(tensor2.use_count(), 1); EXPECT_EQ(tensor2.data_ptr(), tensor.data_ptr()); EXPECT_EQ(tensor.use_count(), 2); EXPECT_EQ(tensor2.use_count(), 1); } EXPECT_EQ(tensor.use_count(), 1); } TEST(Tensor, EnvAlloc) { // Test successful allocation Tensor tensor = Tensor::FromEnvAlloc(TestEnvTensorAllocator, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); EXPECT_EQ(tensor.use_count(), 1); EXPECT_EQ(tensor.shape().size(), 3); EXPECT_EQ(tensor.size(0), 1); EXPECT_EQ(tensor.size(1), 2); EXPECT_EQ(tensor.size(2), 3); EXPECT_EQ(tensor.size(-3), 1); EXPECT_EQ(tensor.size(-2), 2); EXPECT_EQ(tensor.size(-1), 3); EXPECT_EQ(tensor.stride(0), 6); EXPECT_EQ(tensor.stride(1), 3); EXPECT_EQ(tensor.stride(2), 1); EXPECT_EQ(tensor.stride(-3), 6); EXPECT_EQ(tensor.stride(-2), 3); EXPECT_EQ(tensor.stride(-1), 1); EXPECT_EQ(tensor.dtype().code, kDLFloat); EXPECT_EQ(tensor.dtype().bits, 32); EXPECT_EQ(tensor.dtype().lanes, 1); EXPECT_EQ(tensor.device().device_type, kDLCPU); EXPECT_EQ(tensor.device().device_id, 0); EXPECT_NE(tensor.data_ptr(), nullptr); } TEST(Tensor, EnvAllocError) { // Test error handling in DLPackAlloc EXPECT_THROW( { Tensor::FromEnvAlloc(TestEnvTensorAllocatorError, {1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); }, tvm::ffi::Error); } TEST(Tensor, TensorView) { Tensor tensor = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); TensorView tensor_view = tensor; EXPECT_EQ(tensor_view.shape().size(), 3); EXPECT_EQ(tensor_view.shape()[0], 1); EXPECT_EQ(tensor_view.shape()[1], 2); EXPECT_EQ(tensor_view.shape()[2], 3); EXPECT_EQ(tensor_view.dtype().code, kDLFloat); EXPECT_EQ(tensor_view.dtype().bits, 32); EXPECT_EQ(tensor_view.dtype().lanes, 1); AnyView result = tensor_view; EXPECT_EQ(result.type_index(), TypeIndex::kTVMFFIDLTensorPtr); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) TensorView tensor_view2 = result.as().value(); EXPECT_EQ(tensor_view2.shape().size(), 3); EXPECT_EQ(tensor_view2.shape()[0], 1); EXPECT_EQ(tensor_view2.shape()[1], 2); EXPECT_EQ(tensor_view2.shape()[2], 3); EXPECT_EQ(tensor_view2.dtype().code, kDLFloat); EXPECT_EQ(tensor_view2.dtype().bits, 32); EXPECT_EQ(tensor_view2.dtype().lanes, 1); } TEST(Tensor, TensorViewAsStrided) { // Create a base tensor with shape [2, 3] = 6 elements Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); // Fill with sequential values: [0, 1, 2, 3, 4, 5] float* data = reinterpret_cast(tensor.data_ptr()); size_t element_capacity = GetDataSize(tensor) / sizeof(float); ASSERT_EQ(element_capacity, static_cast(tensor.numel())); for (size_t i = 0; i < element_capacity; ++i) { data[i] = static_cast(i); } TensorView tensor_view = tensor; void* original_data_ptr = tensor_view.data_ptr(); EXPECT_EQ(tensor_view.byte_offset(), 0); // Create a strided view with shape [3, 2] and custom strides // Use local variables to ensure they stay in scope for the TensorView Shape new_shape = {3, 2}; Shape new_strides = {1, 3}; TensorView strided_view = tensor_view.as_strided(new_shape, new_strides); // Verify the view has correct shape and strides EXPECT_EQ(strided_view.shape().size(), 2); EXPECT_EQ(strided_view.shape()[0], 3); EXPECT_EQ(strided_view.shape()[1], 2); EXPECT_EQ(strided_view.strides().size(), 2); EXPECT_EQ(strided_view.strides()[0], 1); EXPECT_EQ(strided_view.strides()[1], 3); // Verify the view shares the same underlying data pointer (no offset) EXPECT_EQ(strided_view.data_ptr(), original_data_ptr); EXPECT_EQ(strided_view.byte_offset(), 0); EXPECT_EQ(strided_view.dtype(), tensor_view.dtype()); // Test with element_offset - for float32, 1 element = 4 bytes Shape offset_shape = {2, 2}; Shape offset_strides = {3, 1}; int64_t element_offset = 1; TensorView offset_view = tensor_view.as_strided(offset_shape, offset_strides, element_offset); EXPECT_EQ(offset_view.shape().size(), 2); EXPECT_EQ(offset_view.shape()[0], 2); EXPECT_EQ(offset_view.shape()[1], 2); EXPECT_EQ(offset_view.strides().size(), 2); EXPECT_EQ(offset_view.strides()[0], 3); EXPECT_EQ(offset_view.strides()[1], 1); // For CPU (direct address device), byte_offset should be added to data pointer // and byte_offset field should be 0 // element_offset=1 for float32 = 4 bytes size_t expected_byte_offset = GetDataSize(static_cast(element_offset), DLDataType({kDLFloat, 32, 1})); EXPECT_EQ(expected_byte_offset, 4); // 1 element * 32 bits / 8 = 4 bytes // The data pointer should be advanced by 4 bytes (1 float element) void* expected_offset_ptr = reinterpret_cast(original_data_ptr) + expected_byte_offset; EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr); EXPECT_EQ(offset_view.byte_offset(), 0); // Should be 0 for direct address devices // Verify data access through the offset view float* offset_data = reinterpret_cast(offset_view.data_ptr()); EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f); // Points to data[1] EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f); // Points to data[4] // Test with larger element_offset int64_t element_offset2 = 2; Shape offset_shape2 = {1, 2}; Shape offset_strides2 = {3, 1}; TensorView offset_view2 = tensor_view.as_strided(offset_shape2, offset_strides2, element_offset2); size_t expected_byte_offset2 = GetDataSize(static_cast(element_offset2), DLDataType({kDLFloat, 32, 1})); EXPECT_EQ(expected_byte_offset2, 8); // 2 elements * 32 bits / 8 = 8 bytes void* expected_offset_ptr2 = reinterpret_cast(original_data_ptr) + expected_byte_offset2; EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2); EXPECT_EQ(offset_view2.byte_offset(), 0); float* offset_data2 = reinterpret_cast(offset_view2.data_ptr()); EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f); // Points to data[2] } TEST(Tensor, AsStrided) { // Create a base tensor with shape [2, 3] = 6 elements Tensor tensor = Empty({2, 3}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); // Fill with sequential values: [0, 1, 2, 3, 4, 5] float* data = reinterpret_cast(tensor.data_ptr()); size_t element_capacity = GetDataSize(tensor) / sizeof(float); ASSERT_EQ(element_capacity, static_cast(tensor.numel())); for (size_t i = 0; i < element_capacity; ++i) { data[i] = static_cast(i); } void* original_data_ptr = tensor.data_ptr(); EXPECT_EQ(tensor.byte_offset(), 0); // Create a strided view with shape [3, 2] and custom strides Shape new_shape = {3, 2}; Shape new_strides = {1, 3}; Tensor strided_view = tensor.as_strided(new_shape, new_strides); // Verify the view has correct shape and strides EXPECT_EQ(strided_view.shape().size(), 2); EXPECT_EQ(strided_view.shape()[0], 3); EXPECT_EQ(strided_view.shape()[1], 2); EXPECT_EQ(strided_view.strides().size(), 2); EXPECT_EQ(strided_view.strides()[0], 1); EXPECT_EQ(strided_view.strides()[1], 3); // Verify the view shares the same underlying data pointer (no offset) EXPECT_EQ(strided_view.data_ptr(), original_data_ptr); EXPECT_EQ(strided_view.byte_offset(), 0); EXPECT_EQ(strided_view.dtype(), tensor.dtype()); // Test with element_offset - for float32, 1 element = 4 bytes Shape offset_shape = {2, 2}; Shape offset_strides = {3, 1}; int64_t element_offset = 1; Tensor offset_view = tensor.as_strided(offset_shape, offset_strides, element_offset); EXPECT_EQ(offset_view.shape().size(), 2); EXPECT_EQ(offset_view.shape()[0], 2); EXPECT_EQ(offset_view.shape()[1], 2); EXPECT_EQ(offset_view.strides().size(), 2); EXPECT_EQ(offset_view.strides()[0], 3); EXPECT_EQ(offset_view.strides()[1], 1); // For CPU (direct address device), byte_offset should be added to data pointer // and byte_offset field should be 0 // element_offset=1 for float32 = 4 bytes size_t expected_byte_offset = GetDataSize(static_cast(element_offset), DLDataType({kDLFloat, 32, 1})); EXPECT_EQ(expected_byte_offset, 4); // 1 element * 32 bits / 8 = 4 bytes // The data pointer should be advanced by 4 bytes (1 float element) void* expected_offset_ptr = reinterpret_cast(original_data_ptr) + expected_byte_offset; EXPECT_EQ(offset_view.data_ptr(), expected_offset_ptr); EXPECT_EQ(offset_view.byte_offset(), 0); // Should be 0 for direct address devices // Verify data access through the offset view float* offset_data = reinterpret_cast(offset_view.data_ptr()); EXPECT_EQ(offset_data[0 * 3 + 0 * 1], 1.0f); // Points to data[1] EXPECT_EQ(offset_data[1 * 3 + 0 * 1], 4.0f); // Points to data[4] // Test with larger element_offset int64_t element_offset2 = 2; Tensor offset_view2 = tensor.as_strided({1, 2}, {3, 1}, element_offset2); size_t expected_byte_offset2 = GetDataSize(static_cast(element_offset2), DLDataType({kDLFloat, 32, 1})); EXPECT_EQ(expected_byte_offset2, 8); // 2 elements * 32 bits / 8 = 8 bytes void* expected_offset_ptr2 = reinterpret_cast(original_data_ptr) + expected_byte_offset2; EXPECT_EQ(offset_view2.data_ptr(), expected_offset_ptr2); EXPECT_EQ(offset_view2.byte_offset(), 0); float* offset_data2 = reinterpret_cast(offset_view2.data_ptr()); EXPECT_EQ(offset_data2[0 * 3 + 0 * 1], 2.0f); // Points to data[2] } TEST(Tensor, SizeStrideOutOfBounds) { Tensor tensor = Empty({2, 3, 4}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0})); EXPECT_THROW({ tensor.size(3); }, tvm::ffi::Error); EXPECT_THROW({ tensor.size(-4); }, tvm::ffi::Error); EXPECT_THROW({ tensor.stride(3); }, tvm::ffi::Error); EXPECT_THROW({ tensor.stride(-4); }, tvm::ffi::Error); TensorView tensor_view = tensor; EXPECT_THROW({ tensor_view.size(3); }, tvm::ffi::Error); EXPECT_THROW({ tensor_view.size(-4); }, tvm::ffi::Error); EXPECT_THROW({ tensor_view.stride(3); }, tvm::ffi::Error); EXPECT_THROW({ tensor_view.stride(-4); }, tvm::ffi::Error); } } // namespace tvm-ffi-0.1.12/tests/cpp/test_tuple.cc000066400000000000000000000157701521067262500176060ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Tuple, Basic) { Tuple tuple0(1, 2.0f); EXPECT_EQ(tuple0.get<0>(), 1); EXPECT_EQ(tuple0.get<1>(), 2.0f); Tuple tuple1 = tuple0; EXPECT_EQ(tuple0.use_count(), 2); // test copy on write tuple1.Set<0>(3); EXPECT_EQ(tuple0.get<0>(), 1); EXPECT_EQ(tuple1.get<0>(), 3); EXPECT_EQ(tuple0.use_count(), 1); EXPECT_EQ(tuple1.use_count(), 1); // copy on write not triggered because // tuple1 is unique. tuple1.Set<1>(4); EXPECT_EQ(tuple1.get<1>(), 4.0f); EXPECT_EQ(tuple1.use_count(), 1); // default state Tuple tuple2; EXPECT_EQ(tuple2.use_count(), 1); tuple2.Set<0>(1); tuple2.Set<1>(2.0f); EXPECT_EQ(tuple2.get<0>(), 1); EXPECT_EQ(tuple2.get<1>(), 2.0f); // tuple of object and primitive Tuple tuple3(1, 2); EXPECT_EQ(tuple3.get<0>()->value, 1); EXPECT_EQ(tuple3.get<1>(), 2); tuple3.Set<0>(4); EXPECT_EQ(tuple3.get<0>()->value, 4); } TEST(Tuple, AnyConvert) { Tuple tuple0(1, 2); AnyView view0 = tuple0; // NOLINTNEXTLINE(bugprone-unchecked-optional-access) Array arr0 = view0.as>().value(); EXPECT_EQ(arr0.size(), 2); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(arr0[0].as().value(), 1); // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(arr0[1].as().value()->value, 2); // directly reuse the underlying storage. auto tuple1 = view0.cast>(); EXPECT_TRUE(tuple0.same_as(tuple1)); Any any0 = view0; // trigger a copy due to implict conversion auto tuple2 = any0.cast>(); EXPECT_TRUE(!tuple0.same_as(tuple2)); EXPECT_EQ(tuple2.get<0>()->value, 1); EXPECT_EQ(tuple2.get<1>()->value, 2); } TEST(Tuple, FromTyped) { // try decution Function fadd1 = Function::FromTyped([](const Tuple& a) -> int { return a.get<0>() + static_cast(a.get<1>()->value); }); int b = fadd1(Tuple(1, 2)).cast(); EXPECT_EQ(b, 3); int c = fadd1(Array({1, 2})).cast(); EXPECT_EQ(c, 3); // convert that triggers error EXPECT_THROW( { try { fadd1(Array({1.1, 2})); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: Tuple) -> int`. " "Expected `Tuple` but got `Array[index 0: float]`"); throw; } }, ::tvm::ffi::Error); EXPECT_THROW( { try { fadd1(Array({1.1})); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: Tuple) -> int`. " "Expected `Tuple` but got `Array[size=1]`"); throw; } }, ::tvm::ffi::Error); } TEST(Tuple, Upcast) { Tuple t0(1, 2.0f); Tuple t1 = t0; EXPECT_EQ(t1.get<0>().cast(), 1); EXPECT_EQ(t1.get<1>().cast(), 2.0f); static_assert(details::type_contains_v, Tuple>); static_assert(details::type_contains_v, Tuple>); static_assert(details::type_contains_v, Tuple>); } TEST(Tuple, ArrayIterForwarding) { Tuple t0(1, 2); Tuple t1(3, 4); Array> arr0 = {t0, t1}; std::vector> vec0 = {t0}; vec0.insert(vec0.end(), arr0.begin(), arr0.end()); EXPECT_EQ(vec0.size(), 3); EXPECT_EQ(vec0[0].get<0>()->value, 1); EXPECT_EQ(vec0[0].get<1>()->value, 2); EXPECT_EQ(vec0[1].get<0>()->value, 1); EXPECT_EQ(vec0[1].get<1>()->value, 2); EXPECT_EQ(vec0[2].get<0>()->value, 3); EXPECT_EQ(vec0[2].get<1>()->value, 4); } TEST(Tuple, ArrayIterForwardSingleElem) { Tuple t0(1); Tuple t1(2); Array> arr0 = {t0, t1}; std::vector> vec0 = {t0}; vec0.insert(vec0.end(), arr0.begin(), arr0.end()); EXPECT_EQ(vec0.size(), 3); EXPECT_EQ(vec0[0].get<0>()->value, 1); EXPECT_EQ(vec0[1].get<0>()->value, 1); EXPECT_EQ(vec0[2].get<0>()->value, 2); } TEST(Tuple, CPPFeatures) { { // deduction guide auto t = Tuple{1, 2.0f, String{"hello"}}; // structural binding (lvalue) auto [a, b, c] = t; EXPECT_EQ(a, 1); EXPECT_EQ(b, 2.0f); EXPECT_EQ(c, "hello"); } { // ADL-friendly get will find the right get function auto p = Tuple{Array{0}}; EXPECT_EQ(p.use_count(), 1); // p<0> and q each hold a reference const auto q = get<0>(p); EXPECT_EQ(q.use_count(), 2); } { auto p = Tuple{Array{0}}; EXPECT_EQ(p.use_count(), 1); // structured binding (rvalue) // move out the only ownership auto [q] = std::move(p); EXPECT_EQ(q.use_count(), 1); } { // ADL-friendly get with move semantics auto p = Tuple{Array{0}}; // normal copy get, won't move out anything { auto& q = p; EXPECT_EQ(q.use_count(), 1); const auto _ = get<0>(q); } EXPECT_TRUE(get<0>(p).defined()); // ref-count of q > 1, so the array obj is not moved out { auto q = p; EXPECT_EQ(q.use_count(), 2); const auto _ = get<0>(std::move(q)); } EXPECT_TRUE(get<0>(p).defined()); // ref-count of q = 1, so the array obj is moved out { auto& q = p; EXPECT_EQ(q.use_count(), 1); // NOTE: we do not use structured binding here, // as it creates a temporary value that will move out the q (i.e. p), // making p in an not defined state (p.data_ == 0). // Our get resolution accepts rvalue reference, // which keeps p in a defined state for testing purposes. // DO NOT rely on this behavior in real code. const auto _ = get<0>(std::move(q)); } EXPECT_FALSE(get<0>(p).defined()); } } } // namespace tvm-ffi-0.1.12/tests/cpp/test_variant.cc000066400000000000000000000115251521067262500201130ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include "./testing_object.h" namespace { using namespace tvm::ffi; using namespace tvm::ffi::testing; TEST(Variant, Basic) { Variant v1 = 1; EXPECT_EQ(v1.get(), 1); Variant v2 = 2.0f; EXPECT_EQ(v2.get(), 2.0f); v2 = v1; EXPECT_EQ(v2.get(), 1); } TEST(Variant, AnyConvert) { Variant v = 1; AnyView view0 = v; // NOLINTNEXTLINE(bugprone-unchecked-optional-access) EXPECT_EQ(view0.as().value(), 1); // implicit convert to variant Any any0 = 1; auto v1 = any0.cast>>(); EXPECT_EQ(v1.get()->value, 1); // move from any to variant Variant v2 = TInt(1); Any any1 = std::move(v2); auto v3 = std::move(any1).cast>(); auto v4 = std::move(v3).get(); EXPECT_EQ(v4->value, 1); EXPECT_EQ(v4.use_count(), 1); } TEST(Variant, ObjectPtrHashEqual) { TInt x = TInt(1); TFloat y = TFloat(1.0f); Variant v0 = x; Variant v1 = y; Variant v2 = v1; // NOLINT(performance-unnecessary-copy-initialization) EXPECT_EQ(ObjectPtrHash()(v0), ObjectPtrHash()(x)); EXPECT_TRUE(!ObjectPtrEqual()(v0, v1)); EXPECT_TRUE(!ObjectPtrEqual()(v0, v2)); } TEST(Variant, FromTyped) { // try decution Function fadd1 = Function::FromTyped([](const Variant& a) -> int64_t { if (auto opt_int = a.as()) { return opt_int.value() + 1; } else { return a.get()->value + 1; } }); int b = fadd1(1).cast(); EXPECT_EQ(b, 2); // convert that triggers error EXPECT_THROW( { try { fadd1(1.1); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ( error.message(), "Mismatched type on argument #0 when calling: `(0: Variant) -> int`. " "Expected `Variant` but got `float`"); throw; } }, ::tvm::ffi::Error); Function fadd2 = Function::FromTyped([](const Array>& a) -> int64_t { if (auto opt_int = a[0].as()) { return opt_int.value() + 1; } else { return a[0].get()->value + 1; } }); int c = fadd2(Array({1, 2})).cast(); EXPECT_EQ(c, 2); // convert that triggers error EXPECT_THROW( { try { fadd2(Array({1, 1.1})); } catch (const Error& error) { EXPECT_EQ(error.kind(), "TypeError"); EXPECT_EQ(error.message(), "Mismatched type on argument #0 when calling: `(0: Array>) -> int`. " "Expected `Array>` but got `Array[index 1: float]`"); throw; } }, ::tvm::ffi::Error); } TEST(Variant, Upcast) { Array a0 = {1, 2, 3}; static_assert(details::type_contains_v>, Array>); Array> a1 = a0; EXPECT_EQ(a1[0].get(), 1); } TEST(Variant, AllObjectRef) { Variant> v0 = TInt(1); EXPECT_EQ(v0.get()->value, 1); static_assert(std::is_base_of_v); Any any0 = v0; EXPECT_EQ(any0.cast()->value, 1); auto v2 = any0.cast>>(); EXPECT_TRUE(v0.same_as(v2)); // assignment operator v0 = Array({TInt(2), TInt(3)}); EXPECT_EQ(v0.get>().size(), 2); EXPECT_EQ(v0.get>()[0]->value, 2); EXPECT_EQ(v0.get>()[1]->value, 3); EXPECT_EQ(sizeof(v0), sizeof(ObjectRef)); } TEST(Variant, PODSameAs) { Variant v0 = 1; Variant v1 = 1; EXPECT_TRUE(v0.same_as(v1)); String s = String("hello long str"); v0 = s; v1 = s; EXPECT_TRUE(v0.same_as(v1)); v1 = String("hello long str"); EXPECT_TRUE(!v0.same_as(v1)); } } // namespace tvm-ffi-0.1.12/tests/cpp/testing_object.h000066400000000000000000000422241521067262500202550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #ifndef TVM_FFI_TESTING_OBJECT_H_ #define TVM_FFI_TESTING_OBJECT_H_ #include #include #include #include #include #include #include #include namespace tvm { namespace ffi { namespace testing { // We deliberately pad extra // in the header to test cases // where the object subclass address // do not align with the base object address // not handling properly will cause buffer overflow class BasePad { public: int64_t extra[4]; }; class TNumberObj : public BasePad, public Object { public: // declare as one slot, with float as overflow static constexpr uint32_t _type_child_slots = 1; static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO("test.Number", TNumberObj, Object); }; class TNumber : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TNumber, ObjectRef, TNumberObj); }; class TIntObj : public TNumberObj { public: int64_t value; TIntObj(int64_t value) : value(value) {} explicit TIntObj(UnsafeInit) {} int64_t GetValue() const { return value; } inline static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.Int", TIntObj, TNumberObj); }; class TInt : public TNumber { public: explicit TInt(int64_t value) { data_ = make_object(value); } static TInt StaticAdd(TInt lhs, TInt rhs) { return TInt(lhs->value + rhs->value); } static int64_t CustomAnyHash(const Any& src) { return static_cast(src.cast()->value + 1024); } static bool CustomAnyEqual(const Any& lhs, const Any& rhs) { return lhs.cast()->value == rhs.cast()->value; } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TInt, TNumber, TIntObj); }; inline void TIntObj::RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("value", &TIntObj::value) .def_static("static_add", &TInt::StaticAdd, "static add method"); // define extra type attributes refl::TypeAttrDef() .def("test.GetValue", &TIntObj::GetValue) .attr("test.size", sizeof(TIntObj)) .attr(refl::type_attr::kAnyHash, reinterpret_cast(&TInt::CustomAnyHash)) .attr(refl::type_attr::kAnyEqual, reinterpret_cast(&TInt::CustomAnyEqual)); // custom json serialization refl::TypeAttrDef() .def(refl::type_attr::kDataToJson, [](const TIntObj* self) -> Map { return Map{{"value", self->value}}; }) .def(refl::type_attr::kDataFromJson, [](Map json_obj) -> TInt { return TInt(json_obj["value"].cast()); }); } class TFloatObj : public TNumberObj { public: double value; TFloatObj(double value) : value(value) {} double Add(double other) const { return value + other; } static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.Float", TFloatObj, TNumberObj); }; class TFloat : public TNumber { public: explicit TFloat(double value) { data_ = make_object(value); } static uint64_t CustomAnyHash(const Any& src) { double value = src.cast()->value; return static_cast(value * 10 + 2048); } static bool CustomAnyEqual(const Any& lhs, const Any& rhs) { return lhs.cast()->value == rhs.cast()->value; } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TFloat, TNumber, TFloatObj); }; inline void TFloatObj::RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("value", &TFloatObj::value, "float value field", refl::default_value(10.0)) .def("sub", [](const TFloatObj* self, double other) -> double { return self->value - other; }) .def("add", &TFloatObj::Add, "add method"); refl::TypeAttrDef() .def(refl::type_attr::kAnyHash, &TFloat::CustomAnyHash) .def(refl::type_attr::kAnyEqual, &TFloat::CustomAnyEqual); } class TPrimExprObj : public Object { public: std::string dtype; double value; TPrimExprObj(std::string dtype, double value) : dtype(dtype), value(value) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_rw("dtype", &TPrimExprObj::dtype, "dtype field", refl::default_value("float")) .def_ro("value", &TPrimExprObj::value, "value field", refl::default_value(0)) .def("sub", [](TPrimExprObj* self, double other) -> double { // this is ok because TPrimExprObj is declared asmutable return self->value - other; }); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.PrimExpr", TPrimExprObj, Object); }; class TPrimExpr : public ObjectRef { public: explicit TPrimExpr(std::string dtype, double value) { data_ = make_object(dtype, value); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TPrimExpr, ObjectRef, TPrimExprObj); }; class TVarObj : public Object { public: std::string name; TVarObj(std::string name) : name(name) {} // need unsafe init constructor for json serialization explicit TVarObj(UnsafeInit) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef().def_ro("name", &TVarObj::name, refl::AttachFieldFlag::SEqHashIgnore()); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.Var", TVarObj, Object); }; class TVar : public ObjectRef { public: explicit TVar(std::string name) { data_ = make_object(name); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TVar, ObjectRef, TVarObj); }; // FreeVar test object that has a sub-field referencing another FreeVar. // This models the "var with nested vars" case (analogous to a relax::Var // whose struct_info contains tir shape vars). It is used to exercise the // difference between SEqHashDefRecursive and SEqHashDefNonRecursive at the // FFI layer: under recursive semantics the nested ``dep`` var rebinds // transitively; under non-recursive semantics it is treated as a use of an // outer-scope binding and equality fails when no such outer binding exists. class TVarWithDepObj : public Object { public: std::string name; // Optional dependency var; when null, this object behaves like a plain // FreeVar with no nested free vars. Optional dep; TVarWithDepObj(std::string name, Optional dep) : name(std::move(name)), dep(std::move(dep)) {} explicit TVarWithDepObj(UnsafeInit) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("name", &TVarWithDepObj::name, refl::AttachFieldFlag::SEqHashIgnore()) // ``dep`` participates in structural equality without any def flag, // so it is a USE position. Whether the FreeVar in ``dep`` may rebind // is decided by the def flag on whichever outer field reaches this // object. .def_ro("dep", &TVarWithDepObj::dep); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.VarWithDep", TVarWithDepObj, Object); }; class TVarWithDep : public ObjectRef { public: explicit TVarWithDep(std::string name, Optional dep = std::nullopt) { data_ = make_object(std::move(name), std::move(dep)); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TVarWithDep, ObjectRef, TVarWithDepObj); }; // Holder with one recursive-def field and one non-recursive-def field. // Used by StructuralEqualHash.NonRecursiveDef tests below. class TDefHolderObj : public Object { public: TVarWithDep def_recursive; TVarWithDep def_non_recursive; TDefHolderObj(TVarWithDep def_recursive, TVarWithDep def_non_recursive) : def_recursive(std::move(def_recursive)), def_non_recursive(std::move(def_non_recursive)) {} explicit TDefHolderObj(UnsafeInit) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("def_recursive", &TDefHolderObj::def_recursive, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("def_non_recursive", &TDefHolderObj::def_non_recursive, refl::AttachFieldFlag::SEqHashDefNonRecursive()); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.DefHolder", TDefHolderObj, Object); }; class TDefHolder : public ObjectRef { public: explicit TDefHolder(TVarWithDep def_recursive, TVarWithDep def_non_recursive) { data_ = make_object(std::move(def_recursive), std::move(def_non_recursive)); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TDefHolder, ObjectRef, TDefHolderObj); }; class TFuncObj : public Object { public: Array params; Array body; Optional comment; // need unsafe init constructor or default constructor for json serialization explicit TFuncObj(UnsafeInit) {} TFuncObj(Array params, Array body, Optional comment) : params(params), body(body), comment(comment) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("params", &TFuncObj::params, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("body", &TFuncObj::body) .def_ro("comment", &TFuncObj::comment, refl::AttachFieldFlag::SEqHashIgnore()); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.Func", TFuncObj, Object); }; class TFunc : public ObjectRef { public: explicit TFunc(Array params, Array body, Optional comment) { data_ = make_object(params, body, comment); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TFunc, ObjectRef, TFuncObj); }; class TCustomFuncObj : public Object { public: Array params; Array body; String comment; TCustomFuncObj(Array params, Array body, String comment) : params(params), body(body), comment(comment) {} bool SEqual(const TCustomFuncObj* other, ffi::TypedFunction cmp) const { if (!cmp(params, other->params, true, "params")) { return false; } if (!cmp(body, other->body, false, "body")) { return false; } return true; } int64_t SHash(int64_t init_hash, ffi::TypedFunction hash) const { int64_t hash_value = init_hash; hash_value = hash(params, hash_value, true); hash_value = hash(body, hash_value, false); return hash_value; } static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("params", &TCustomFuncObj::params) .def_ro("body", &TCustomFuncObj::body) .def_ro("comment", &TCustomFuncObj::comment); refl::TypeAttrDef() .def(refl::type_attr::kSEqual, &TCustomFuncObj::SEqual) .def(refl::type_attr::kSHash, &TCustomFuncObj::SHash); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.CustomFunc", TCustomFuncObj, Object); }; class TCustomFunc : public ObjectRef { public: explicit TCustomFunc(Array params, Array body, String comment) { data_ = make_object(params, body, comment); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TCustomFunc, ObjectRef, TCustomFuncObj); }; // Test object with all POD field types to exercise serialization of every field kind. class TAllFieldsObj : public Object { public: bool v_bool; int64_t v_int; double v_float; DLDataType v_dtype; DLDevice v_device; String v_str; Optional v_opt_str; Array v_array; Map v_map; TAllFieldsObj(bool v_bool, int64_t v_int, double v_float, DLDataType v_dtype, DLDevice v_device, String v_str, Optional v_opt_str, Array v_array, Map v_map) : v_bool(v_bool), v_int(v_int), v_float(v_float), v_dtype(v_dtype), v_device(v_device), v_str(v_str), v_opt_str(v_opt_str), v_array(v_array), v_map(v_map) {} explicit TAllFieldsObj(UnsafeInit) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("v_bool", &TAllFieldsObj::v_bool) .def_ro("v_int", &TAllFieldsObj::v_int) .def_ro("v_float", &TAllFieldsObj::v_float) .def_ro("v_dtype", &TAllFieldsObj::v_dtype) .def_ro("v_device", &TAllFieldsObj::v_device) .def_ro("v_str", &TAllFieldsObj::v_str) .def_ro("v_opt_str", &TAllFieldsObj::v_opt_str) .def_ro("v_array", &TAllFieldsObj::v_array) .def_ro("v_map", &TAllFieldsObj::v_map); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.AllFields", TAllFieldsObj, Object); }; class TAllFields : public ObjectRef { public: explicit TAllFields(bool v_bool, int64_t v_int, double v_float, DLDataType v_dtype, DLDevice v_device, String v_str, Optional v_opt_str, Array v_array, Map v_map) { data_ = make_object(v_bool, v_int, v_float, v_dtype, v_device, v_str, v_opt_str, v_array, v_map); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TAllFields, ObjectRef, TAllFieldsObj); }; // Test object with fields that have default values to test deserialization with missing fields class TWithDefaultsObj : public Object { public: int64_t required_val; int64_t default_int; String default_str; bool default_bool; TWithDefaultsObj(int64_t required_val, int64_t default_int, String default_str, bool default_bool) : required_val(required_val), default_int(default_int), default_str(default_str), default_bool(default_bool) {} explicit TWithDefaultsObj(UnsafeInit) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() .def_ro("required_val", &TWithDefaultsObj::required_val) .def_ro("default_int", &TWithDefaultsObj::default_int, refl::default_value(42)) .def_ro("default_str", &TWithDefaultsObj::default_str, refl::default_value("default")) .def_ro("default_bool", &TWithDefaultsObj::default_bool, refl::default_value(true)); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.WithDefaults", TWithDefaultsObj, Object); }; class TWithDefaults : public ObjectRef { public: explicit TWithDefaults(int64_t required_val, int64_t default_int = 42, String default_str = "default", bool default_bool = true) { data_ = make_object(required_val, default_int, default_str, default_bool); } TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TWithDefaults, ObjectRef, TWithDefaultsObj); }; } // namespace testing template <> inline constexpr bool use_default_type_traits_v = true; template <> struct TypeTraits : public ObjectRefWithFallbackTraitsBase { TVM_FFI_INLINE static testing::TPrimExpr ConvertFallbackValue(StrictBool value) { return testing::TPrimExpr("bool", static_cast(value)); } TVM_FFI_INLINE static testing::TPrimExpr ConvertFallbackValue(int64_t value) { return testing::TPrimExpr("int64", static_cast(value)); } TVM_FFI_INLINE static testing::TPrimExpr ConvertFallbackValue(double value) { return testing::TPrimExpr("float32", static_cast(value)); } // hack into the dtype to store string TVM_FFI_INLINE static testing::TPrimExpr ConvertFallbackValue(String value) { return testing::TPrimExpr(value, 0); } }; } // namespace ffi } // namespace tvm #endif // TVM_FFI_TESTING_OBJECT_H_ tvm-ffi-0.1.12/tests/docker/000077500000000000000000000000001521067262500155625ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/docker/Dockerfile000066400000000000000000000046301521067262500175570ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # Using NVIDIA CUDA base image https://hub.docker.com/r/nvidia/cuda # Based on your host CUDA driver version, you may need to select an older CUDA base image. # See https://gitlab.com/nvidia/container-images/cuda/blob/master/doc/supported-tags.md FROM nvidia/cuda:12.6.3-devel-ubuntu22.04 ARG DEBIAN_FRONTEND=noninteractive # Install prerequisites for external repositories RUN apt-get update \ && apt-get install -y --no-install-recommends \ ca-certificates \ curl \ gnupg \ wget \ lsb-release \ && rm -rf /var/lib/apt/lists/* # Add Kitware APT repository for CMake >= 3.24 RUN wget -O /tmp/kitware-archive.sh https://apt.kitware.com/kitware-archive.sh \ && bash /tmp/kitware-archive.sh \ && rm -f /tmp/kitware-archive.sh # Install build essentials, Git, and Python >= 3.9 RUN apt-get update \ && apt-get install -y --no-install-recommends \ build-essential \ clang \ clang-format \ clang-tidy \ cmake \ doxygen \ git \ graphviz \ ninja-build \ pandoc \ pkg-config \ python-is-python3 \ python3 \ python3-dev \ python3-pip \ python3-setuptools \ python3-venv \ python3-wheel \ unzip \ zip \ && rm -rf /var/lib/apt/lists/* # Upgrade pip to the latest version RUN python3 -m pip install --upgrade pip # Provide a working directory for the project WORKDIR /workspace # Optionally clone a repository during build with --build-arg REPO_URL=... ARG REPO_URL RUN if [ -n "${REPO_URL}" ]; then git clone "${REPO_URL}" repo; fi CMD ["/bin/bash"] tvm-ffi-0.1.12/tests/lint/000077500000000000000000000000001521067262500152615ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/lint/check_asf_header.py000066400000000000000000000347361521067262500210660ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Helper tool to add ASF header to files that cannot be handled by Rat.""" from __future__ import annotations import argparse import fnmatch import subprocess import sys from pathlib import Path header_cstyle = """ /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ """.strip() header_mdstyle = """ """.strip() header_pystyle = """ # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """.strip() header_rststyle = """ .. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """.strip() header_groovystyle = """ // Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. """.strip() header_cmdstyle = """ :: Licensed to the Apache Software Foundation (ASF) under one :: or more contributor license agreements. See the NOTICE file :: distributed with this work for additional information :: regarding copyright ownership. The ASF licenses this file :: to you under the Apache License, Version 2.0 (the :: "License"); you may not use this file except in compliance :: with the License. You may obtain a copy of the License at :: :: http://www.apache.org/licenses/LICENSE-2.0 :: :: Unless required by applicable law or agreed to in writing, :: software distributed under the License is distributed on an :: "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY :: KIND, either express or implied. See the License for the :: specific language governing permissions and limitations :: under the License. """.strip() FMT_MAP = { "sh": header_pystyle, "cc": header_cstyle, "c": header_cstyle, "mm": header_cstyle, "m": header_cstyle, "go": header_cstyle, "java": header_cstyle, "h": header_cstyle, "pyi": header_pystyle, "py": header_pystyle, "toml": header_pystyle, "yml": header_pystyle, "yaml": header_pystyle, "rs": header_cstyle, "md": header_mdstyle, "cmake": header_pystyle, "mk": header_pystyle, "rst": header_rststyle, "gradle": header_groovystyle, "tcl": header_pystyle, "xml": header_mdstyle, "storyboard": header_mdstyle, "pbxproj": header_cstyle, "plist": header_mdstyle, "xcworkspacedata": header_mdstyle, "html": header_mdstyle, "bat": header_cmdstyle, } # Files and patterns to skip during header checking SKIP_LIST: list[str] = [] def should_skip_file(filepath: str) -> bool: """Check if file should be skipped based on SKIP_LIST.""" for pattern in SKIP_LIST: if fnmatch.fnmatch(filepath, pattern): return True return False def get_git_files() -> list[str] | None: """Get list of files tracked by git.""" try: result = subprocess.run( ["git", "ls-files"], check=False, capture_output=True, text=True, cwd=Path.cwd() ) if result.returncode == 0: return [line.strip() for line in result.stdout.split("\n") if line.strip()] else: print("Error: Could not get git files. Make sure you're in a git repository.") print("Git command failed:", result.stderr.strip()) return None except FileNotFoundError: print("Error: Git not found. This tool requires git to be installed.") return None def copyright_line(line: str) -> bool: # Following two items are intentionally break apart # so that the copyright detector won"t detect the file itself. if line.find("Copyright " + "(c)") != -1: return True # break pattern into two lines to avoid false-negative check spattern1 = "Copyright" if line.find(spattern1) != -1 and line.find("by") != -1: return True return False def check_header(fname: str, header: str) -> bool: """Check header status of file without modifying it.""" if not Path(fname).exists(): print(f"ERROR: Cannot find {fname}") return False lines = Path(fname).open().readlines() has_asf_header = False has_copyright = False for i, l in enumerate(lines): if l.find("Licensed to the Apache Software Foundation") != -1: has_asf_header = True elif copyright_line(l): has_copyright = True if has_asf_header and not has_copyright: return True # File is good if not has_asf_header: print(f"ERROR: Missing ASF header in {fname}") print("Run `python tests/lint/check_asf_header.py --fix` to add the header") return False if has_copyright: print(f"ERROR: Has copyright line that should be removed in {fname}") return False return True def collect_files() -> list[str] | None: """Collect all files that need header checking from git.""" files = [] # Get files from git (required) git_files = get_git_files() if git_files is None: # Git is required, so exit if we can't get files return None # Filter git files based on supported types and skip list for git_file in git_files: # Check if file should be skipped if should_skip_file(git_file): continue # Check if this file type is supported suffix = git_file.split(".")[-1] if "." in git_file else "" basename = Path(git_file).name if ( suffix in FMT_MAP or basename == "gradle.properties" or (suffix == "" and basename in ["CMakeLists", "Makefile"]) ): files.append(git_file) return files def add_header(fname: str, header: str) -> None: # noqa: PLR0912 """Add header to file.""" if not Path(fname).exists(): print(f"Cannot find {fname} ...") return lines = Path(fname).open().readlines() has_asf_header = False has_copyright = False for i, l in enumerate(lines): if l.find("Licensed to the Apache Software Foundation") != -1: has_asf_header = True elif copyright_line(l): has_copyright = True lines[i] = "" if has_asf_header and not has_copyright: return with Path(fname).open("w") as outfile: skipline = False if not lines: skipline = False # File is enpty elif lines[0][:2] == "#!": skipline = True elif lines[0][:2] == ""): skipline = True elif lines[0].startswith("// !$"): skipline = True if skipline: outfile.write(lines[0]) if not has_asf_header: outfile.write(header + "\n\n") outfile.write("".join(lines[1:])) else: if not has_asf_header: outfile.write(header + "\n\n") outfile.write("".join(lines)) if not has_asf_header: print(f"Add header to {fname}") if has_copyright: print(f"Removed copyright line from {fname}") def main() -> int: # noqa: PLR0911, PLR0912 parser = argparse.ArgumentParser( description="Check and fix ASF headers in source files tracked by git", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" This tool processes all files tracked by git (using 'git ls-files') and automatically skips files matching patterns in SKIP_LIST (build artifacts, third-party files, etc.). One of --check, --fix, or --dry-run must be specified. Examples: # Show all files that would be processed (dry run) python check_asf_header.py --dry-run # Check all git-tracked files (report errors) python check_asf_header.py --check # Fix all git-tracked files (add/update headers) python check_asf_header.py --fix """, ) parser.add_argument( "--check", action="store_true", help="Check mode: report errors without modifying files", ) parser.add_argument( "--fix", action="store_true", help="Fix mode: fix files with missing or incorrect headers (default)", ) parser.add_argument( "--dry-run", action="store_true", help="Dry run: show files that would be processed without doing anything", ) args = parser.parse_args() # Validate arguments if args.check and args.fix: print("Error: Cannot specify both --check and --fix") return 1 if not args.check and not args.fix and not args.dry_run: print("Error: Must specify one of --check, --fix, or --dry-run") return 1 # Always use git-based approach files = collect_files() if files is None: return 1 # Error already printed by get_git_files() if not files: print("No files found to process") return 0 # Handle dry-run mode if args.dry_run: print(f"Would process {len(files)} files:") for fname in sorted(files): print(f" {fname}") return 0 # Determine mode check_only = args.check error_count = 0 processed_count = 0 for fname in files: processed_count += 1 suffix = fname.split(".")[-1] if "." in fname else "" basename = Path(fname).name # Determine header type if suffix in FMT_MAP: header = FMT_MAP[suffix] elif basename == "gradle.properties": header = FMT_MAP["h"] elif suffix == "" and basename in ["CMakeLists", "Makefile"]: header = FMT_MAP["cmake"] if basename == "CMakeLists" else FMT_MAP["mk"] else: if check_only: print(f"ERROR: Cannot handle file type for {fname}") error_count += 1 else: print(f"Cannot handle {fname} ...") continue if check_only: # Check mode if not check_header(fname, header): error_count += 1 else: # Fix mode add_header(fname, header) if check_only: if error_count > 0: print(f"\nFound {error_count} errors in {processed_count} files") return 1 else: print(f"\nAll {processed_count} files are compliant") return 0 return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/tests/lint/check_file_type.py000066400000000000000000000140051521067262500207500ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Helper tool to check file types that are allowed to checkin.""" import subprocess import sys from pathlib import Path # List of file types we allow ALLOW_EXTENSION = { # source code "cc", "c", "h", "s", "rs", "m", "mm", "g4", "gradle", "js", "cjs", "mjs", "tcl", "scala", "java", "go", "ts", "sh", "py", "pyi", "pxi", "pyd", "pyx", "cu", "cuh", "bat", "ps1", # configurations "mk", "in", "cmake", "xml", "toml", "yml", "yaml", "json", "cfg", # docs "txt", "md", "rst", "css", # sgx "edl", "lds", # ios "pbxproj", "plist", "xcworkspacedata", "storyboard", "xcscheme", # hw/chisel "sbt", "properties", "v", "sdc", # generated parser "interp", "tokens", # interface definition "idl", # opencl file "cl", # zephyr config file "conf", # arduino sketch file "ino", # linker scripts "ld", # Jinja2 templates "j2", # Jenkinsfiles "groovy", # Python-parseable config files "ini", } # List of file names allowed ALLOW_FILE_NAME = { ".gitignore", ".eslintignore", ".gitattributes", "README", "Makefile", "Doxyfile", "pylintrc", ".clang-format", ".clang-tidy", ".gitmodules", "CODEOWNERSHIP", "Dockerfile", "py.typed", } # List of specific files allowed in relpath to ALLOW_SPECIFIC_FILE = { "LICENSE", "NOTICE", "KEYS", "DISCLAIMER", "addons/torch_c_dlpack_ext/LICENSE", "addons/torch_c_dlpack_ext/NOTICE", } def filename_allowed(name: str) -> bool: """Check if name is allowed by the current policy. Paramaters ---------- name : str Input name Returns ------- allowed : bool Whether the filename is allowed. """ arr = name.rsplit(".", 1) if arr[-1] in ALLOW_EXTENSION: return True if Path(name).name in ALLOW_FILE_NAME: return True if name.startswith("3rdparty"): return True if name in ALLOW_SPECIFIC_FILE: return True return False def copyright_line(line: str) -> bool: # Following two items are intentionally break apart # so that the copyright detector won't detect the file itself. if line.find("Copyright " + "(c)") != -1: return True # break pattern into two lines to avoid false-negative check spattern1 = "Copyright" if line.find(spattern1) != -1 and line.find("by") != -1: return True return False def check_asf_copyright(fname: str) -> bool: if fname.endswith(".png"): return True if not Path(fname).is_file(): return True has_asf_header = False has_copyright = False try: for line in Path(fname).open(): if line.find("Licensed to the Apache Software Foundation") != -1: has_asf_header = True if copyright_line(line): has_copyright = True if has_asf_header and has_copyright: return False except UnicodeDecodeError: pass return True def main() -> None: cmd = ["git", "ls-files"] proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) (out, _) = proc.communicate() res = out.decode("utf-8") assert proc.returncode == 0, f"{' '.join(cmd)} errored: {res}" flist = res.split() error_list = [] for fname in flist: if not filename_allowed(fname): error_list.append(fname) if error_list: report = "------File type check report----\n" report += "\n".join(error_list) report += f"\nFound {len(error_list)} files that are not allowed\n" report += ( "We do not check in binary files into the repo.\n" "If necessary, please discuss with committers and" "modify tests/lint/check_file_type.py to enable the file you need.\n" ) sys.stderr.write(report) sys.stderr.flush() sys.exit(-1) asf_copyright_list = [] for fname in res.split(): if not check_asf_copyright(fname): asf_copyright_list.append(fname) if asf_copyright_list: report = "------File type check report----\n" report += "\n".join(asf_copyright_list) + "\n" report += f"------Found {len(asf_copyright_list)} files that has ASF header with copyright message----\n" report += "--- Files with ASF header do not need Copyright lines.\n" report += "--- Contributors retain copyright to their contribution by default.\n" report += "--- If a file comes with a different license, consider put it under the 3rdparty folder instead.\n" report += "---\n" report += "--- You can use the following steps to remove the copyright lines\n" report += "--- Create file_list.txt in your text editor\n" report += "--- Copy paste the above content in file-list into file_list.txt\n" report += "--- python3 tests/lint/add_asf_header.py file_list.txt\n" sys.stderr.write(report) sys.stderr.flush() sys.exit(-1) print("check_file_type.py: all checks passed..") if __name__ == "__main__": main() tvm-ffi-0.1.12/tests/lint/check_version.py000066400000000000000000000146401521067262500204620ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Version consistency linter across Python, C++, and Rust. This script checks that: 1) C++ version macros in headers are internally consistent and match the canonical project version (major/minor/micro) derived from setuptools_scm. 2) Rust crate versions are internally consistent and compatible with the canonical project version. Usage: python tests/lint/check_version.py [--cpp] [--rust] """ from __future__ import annotations import argparse import re import sys from pathlib import Path from typing import Any import setuptools_scm import tomli from packaging.version import Version as packaging_version RE_MAJOR = re.compile(r"^\s*#\s*define\s+TVM_FFI_VERSION_MAJOR\s+(\d+)\b") RE_MINOR = re.compile(r"^\s*#\s*define\s+TVM_FFI_VERSION_MINOR\s+(\d+)\b") RE_PATCH = re.compile(r"^\s*#\s*define\s+TVM_FFI_VERSION_PATCH\s+(\d+)\b") def _version_info() -> dict[str, Any]: """Return project version information using setuptools_scm.""" version = setuptools_scm.get_version() v = packaging_version(version) return { "full": version, "major": v.major, "minor": v.minor, "micro": v.micro, "pre": v.pre, "dev": v.dev, "local": v.local, "post": v.post, "release": v.release, "is_prerelease": v.is_prerelease, "is_postrelease": v.is_postrelease, "is_devrelease": v.is_devrelease, "base_version": v.base_version, "public": v.public, } def _map_pep440_pre_to_semver(pre: tuple[str, int] | None) -> str | None: if pre is None: return None tag, num = pre tag = tag.lower() if tag in {"a", "alpha"}: tag = "alpha" elif tag in {"b", "beta"}: tag = "beta" elif tag in {"rc", "c"}: tag = "rc" else: return None return f"{tag}.{num}" def _check_cpp(version_info: dict) -> list[str]: errors: list[str] = [] c_api_path = Path("include") / "tvm" / "ffi" / "c_api.h" if not c_api_path.exists(): errors.append(f"[C++] Missing expected header file: {c_api_path}") return errors def _scan_cpp_macros() -> tuple[int, int, int]: file = c_api_path major = minor = patch = None for line in file.read_text(encoding="utf-8").splitlines(): if m := RE_MAJOR.match(line): major = int(m.group(1)) if m := RE_MINOR.match(line): minor = int(m.group(1)) if m := RE_PATCH.match(line): patch = int(m.group(1)) return major, minor, patch (major, minor, patch) = _scan_cpp_macros() if major is None or minor is None or patch is None: errors.append(f"[C++] {c_api_path}: No version macros found: {major=}, {minor=}, {patch=}.") return errors exp_major, exp_minor, exp_patch = ( version_info["major"], version_info["minor"], version_info["micro"], ) if (major, minor, patch) != (exp_major, exp_minor, exp_patch): errors.append( f"[C++] {c_api_path}: Macro version mismatch: found {major}.{minor}.{patch}, " f"expected {exp_major}.{exp_minor}.{exp_patch}." ) return errors def _check_rust(version_info: dict) -> list[str]: errors: list[str] = [] rust_dir = Path("rust") found_versions: dict[Path, str] = {} for path in [ rust_dir / "tvm-ffi" / "Cargo.toml", rust_dir / "tvm-ffi-macros" / "Cargo.toml", rust_dir / "tvm-ffi-sys" / "Cargo.toml", ]: found_versions[path] = tomli.loads(path.read_text(encoding="utf-8"))["package"]["version"] if not found_versions: # No crates found, skip silently return errors # 1) All crates must agree on a single version unique_versions = set(found_versions.values()) if len(unique_versions) > 1: errors.append( "[Rust] Crates have inconsistent versions: " + ", ".join(f"{p} -> {v}" for p, v in sorted(found_versions.items())) ) # 2) Optionally enforce compatibility with Python version base = version_info["base_version"] allowed: set[str] = {base} pre = _map_pep440_pre_to_semver(version_info.get("pre")) if pre: allowed.add(f"{base}-{pre}") allowed = sorted(allowed) for path, v in found_versions.items(): if v not in allowed: errors.append( f"[Rust] {path}: version not compatible with project version. Allowed: {allowed}; got: {v}." ) return errors def main() -> int: parser = argparse.ArgumentParser(description="Check version consistency across languages") parser.add_argument("--cpp", action="store_true", help="Check C++ version macros") parser.add_argument("--rust", action="store_true", help="Check Rust crate versions") args = parser.parse_args() info = _version_info() print( f"Project version: {info['full']}\n" f" major: {info['major']}\n" f" minor: {info['minor']}\n" f" micro: {info['micro']}\n" f" pre: {info['pre']}\n" f" dev: {info['dev']}\n" f" local: {info['local']}\n" f" post: {info['post']}\n" f" base_version: {info['base_version']}\n" f" release: {info['release']}\n" f" public: {info['public']}\n" f" is_prerelease: {info['is_prerelease']}\n" f" is_postrelease: {info['is_postrelease']}\n" f" is_devrelease: {info['is_devrelease']}" ) errors: list[str] = [] if args.cpp: errors += _check_cpp(info) if args.rust: errors += _check_rust(info) if errors: print("\n".join(errors)) return 1 return 0 if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/tests/lint/clang_tidy_precommit.py000066400000000000000000000122511521067262500220300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import os import subprocess import sys from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path DEFAULT_SOURCE_EXTS = (".c", ".cc", ".cpp", ".cxx", ".m", ".mm") def parse_args(argv: Sequence[str]) -> Namespace: p = ArgumentParser( prog="clang-tidy-precommit", description=( "Run clang-tidy on changed files using a compile_commands.json " "generated in a dedicated build directory." ), formatter_class=ArgumentDefaultsHelpFormatter, ) p.add_argument( "--build-dir", "-B", default="build-pre-commit", help="CMake build directory containing compile_commands.json.", ) p.add_argument( "--jobs", "-j", type=int, default=max(1, os.cpu_count() or 1), help="Maximum parallel clang-tidy processes.", ) p.add_argument( "--fix", action="store_true", help="Enable clang-tidy in-place fixes (-fix).", ) p.add_argument( "files", nargs="*", help="Files passed by pre-commit (will be filtered by extension).", ) return p.parse_args(list(argv)) def filter_files(files: Sequence[str]) -> list[str]: kept: list[str] = [] for f in files: if not f: continue p = Path(f) if p.is_dir(): for child in p.rglob("*"): if child.is_file() and child.suffix.lower() in DEFAULT_SOURCE_EXTS: kept.append(str(child)) continue if p.suffix.lower() in DEFAULT_SOURCE_EXTS: kept.append(str(p)) return kept def ensure_compile_commands(build_dir: Path) -> Path: cc_path = build_dir / "compile_commands.json" if cc_path.exists(): return cc_path build_dir.mkdir(parents=True, exist_ok=True) cmd = [ "cmake", "-S", ".", "-B", str(build_dir), "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON", "-DCMAKE_BUILD_TYPE=Release", "-DTVM_FFI_BUILD_TESTS=ON", ] ret = subprocess.run(cmd, check=False) if ret.returncode != 0: print( "[clang-tidy-precommit] Failed to generate compile_commands.json via CMake.", file=sys.stderr, ) sys.exit(ret.returncode) if not cc_path.exists(): print( f"[clang-tidy-precommit] Missing {cc_path}. Ensure CMake config succeeded.", file=sys.stderr, ) sys.exit(2) return cc_path def build_base_cmd(p_arg: Path, fix: bool) -> list[str]: cmd: list[str] = ["clang-tidy", f"-p={p_arg!s}", "-quiet"] if fix: cmd.append("-fix") return cmd def run_parallel(cmd: list[str], files: list[str], jobs: int) -> int: def one(f: str) -> tuple[int, str, list[str]]: full_cmd = [*cmd, f] print(f"[RUNNING] {' '.join(full_cmd)}", file=sys.stderr) proc = subprocess.run(full_cmd, check=False, capture_output=True, text=True) out = (proc.stdout or "") + (proc.stderr or "") return proc.returncode, out, " ".join(full_cmd) jobs = max(1, int(jobs or 1)) rc = 0 with ThreadPoolExecutor(max_workers=jobs) as ex: futs = [ex.submit(one, f) for f in files] for fut in as_completed(futs): code, output, full_cmd = fut.result() output = output.strip() if code != 0: print(f"[FAILED] {full_cmd}\n{output}") rc = 1 return rc def main(argv: Sequence[str] | None = None) -> int: args = parse_args(sys.argv[1:] if argv is None else argv) files = filter_files(args.files) if not files: print("[clang-tidy-precommit] No relevant files to lint.") return 0 build_dir = Path(args.build_dir).resolve() cc_path = ensure_compile_commands(build_dir) base_cmd = build_base_cmd(p_arg=cc_path.parent, fix=args.fix) if sys.platform == "darwin": base_cmd = ["xcrun", *base_cmd] rc = run_parallel(base_cmd, files, args.jobs) if rc != 0 and args.fix: print( "[clang-tidy-precommit] clang-tidy reported issues and applied fixes. " "Re-stage your changes if files were modified.", file=sys.stderr, ) return rc if __name__ == "__main__": sys.exit(main()) tvm-ffi-0.1.12/tests/python/000077500000000000000000000000001521067262500156345ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/python/cpp_src/000077500000000000000000000000001521067262500172655ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/python/cpp_src/test_stl.cc000066400000000000000000000060501521067262500214360ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #include #include #include #include #include #include #include #include #include #include #include namespace { auto test_tuple(std::tuple arg) -> std::tuple { return std::make_tuple(std::get<1>(arg), std::get<0>(arg)); } auto test_vector(std::optional>> arg) -> std::optional> { if (arg) { auto result = std::vector{}; result.reserve(arg->size()); for (const auto& row : *arg) { result.push_back(std::accumulate(row.begin(), row.end(), 0)); } return result; } else { return std::nullopt; } } auto test_variant(std::variant>> arg) -> std::variant> { if (std::holds_alternative(arg)) { return "int"; } else if (std::holds_alternative(arg)) { return "float"; } else { auto result = std::vector{}; for (const auto& item : std::get>>(arg)) { if (std::holds_alternative(item)) { result.emplace_back("int"); } else { result.emplace_back("float"); } } return result; } } auto test_map(const std::map& map) -> std::map { auto result = std::map{}; for (const auto& [key, value] : map) { result[value] = key; } return result; } auto test_map_2(const std::unordered_map& map) -> std::unordered_map { auto result = std::unordered_map{}; for (const auto& [key, value] : map) { result[value] = key; } return result; } auto test_function(std::function f) -> std::function { return [fn = std::move(f)] { return fn() + 1; }; } TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_tuple, test_tuple) TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_vector, test_vector) TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_variant, test_variant) TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_map, test_map) TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_map_2, test_map_2) TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_function, test_function) } // namespace tvm-ffi-0.1.12/tests/python/test_access_path.py000066400000000000000000000110201521067262500215140ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from tvm_ffi.access_path import AccessKind, AccessPath def test_root_path() -> None: root = AccessPath.root() assert isinstance(root, AccessPath) steps = root.to_steps() assert len(steps) == 0 assert root == AccessPath.root() def test_path_attr() -> None: path = AccessPath.root().attr("foo") assert isinstance(path, AccessPath) steps = path.to_steps() assert len(steps) == 1 assert steps[0].kind == AccessKind.ATTR assert steps[0].key == "foo" assert path.parent == AccessPath.root() def test_path_array_item() -> None: path = AccessPath.root().array_item(2) assert isinstance(path, AccessPath) steps = path.to_steps() assert len(steps) == 1 assert steps[0].kind == AccessKind.ARRAY_ITEM assert steps[0].key == 2 assert path.parent == AccessPath.root() def test_path_missing_array_element() -> None: path = AccessPath.root().array_item_missing(2) assert isinstance(path, AccessPath) steps = path.to_steps() assert len(steps) == 1 assert steps[0].kind == AccessKind.ARRAY_ITEM_MISSING assert steps[0].key == 2 assert path.parent == AccessPath.root() def test_path_map_item() -> None: path = AccessPath.root().map_item("foo") assert isinstance(path, AccessPath) steps = path.to_steps() assert len(steps) == 1 assert steps[0].kind == AccessKind.MAP_ITEM assert steps[0].key == "foo" assert path.parent == AccessPath.root() def test_path_missing_map_item() -> None: path = AccessPath.root().map_item_missing("foo") assert isinstance(path, AccessPath) steps = path.to_steps() assert len(steps) == 1 assert steps[0].kind == AccessKind.MAP_ITEM_MISSING assert steps[0].key == "foo" assert path.parent == AccessPath.root() def test_path_is_prefix_of() -> None: # Root is prefix of root assert AccessPath.root().is_prefix_of(AccessPath.root()) # Root is prefix of any path assert AccessPath.root().is_prefix_of(AccessPath.root().attr("foo")) # Non-root is not prefix of root assert not AccessPath.root().attr("foo").is_prefix_of(AccessPath.root()) # Path is prefix of itself assert AccessPath.root().attr("foo").is_prefix_of(AccessPath.root().attr("foo")) # Different attrs are not prefixes of each other assert not AccessPath.root().attr("bar").is_prefix_of(AccessPath.root().attr("foo")) # Shorter path is prefix of longer path with same start assert AccessPath.root().attr("foo").is_prefix_of(AccessPath.root().attr("foo").array_item(2)) # Longer path is not prefix of shorter path assert ( not AccessPath.root().attr("foo").array_item(2).is_prefix_of(AccessPath.root().attr("foo")) ) # Different paths are not prefixes assert ( not AccessPath.root().attr("foo").is_prefix_of(AccessPath.root().attr("bar").array_item(2)) ) def test_path_equal() -> None: # Root equals root assert AccessPath.root() == AccessPath.root() # Root does not equal non-root paths assert not (AccessPath.root() == AccessPath.root().attr("foo")) # Non-root does not equal root assert not (AccessPath.root().attr("foo") == AccessPath.root()) # Path equals itself assert AccessPath.root().attr("foo") == AccessPath.root().attr("foo") # Different attrs are not equal assert not (AccessPath.root().attr("bar") == AccessPath.root().attr("foo")) # Shorter path does not equal longer path assert not (AccessPath.root().attr("foo") == AccessPath.root().attr("foo").array_item(2)) # Longer path does not equal shorter path assert not (AccessPath.root().attr("foo").array_item(2) == AccessPath.root().attr("foo")) # Different paths are not equal assert not (AccessPath.root().attr("foo") == AccessPath.root().attr("bar").array_item(2)) tvm-ffi-0.1.12/tests/python/test_build.cc000066400000000000000000000032201521067262500202760ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1 #include "test_build.h" #include #include #include #include #include void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu) tvm-ffi-0.1.12/tests/python/test_build.h000066400000000000000000000017571521067262500201550ustar00rootroot00000000000000/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ #ifndef TVM_FFI_TEST_BUILD_H_ #define TVM_FFI_TEST_BUILD_H_ #include void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y); #endif // TVM_FFI_TEST_BUILD_H_ tvm-ffi-0.1.12/tests/python/test_build.py000066400000000000000000000252031521067262500203460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import pathlib import numpy import pytest import tvm_ffi import tvm_ffi.cpp from tvm_ffi.core import TypeSchema from tvm_ffi.module import Module def test_build_cpp() -> None: cpp_path = pathlib.Path(__file__).parent.resolve() / "test_build.cc" output_lib_path = tvm_ffi.cpp.build( name="hello", sources=[str(cpp_path)], ) mod: Module = tvm_ffi.load_module(output_lib_path) metadata = mod.get_function_metadata("add_one_cpu") assert metadata is not None, "add_one_cpu should have metadata" assert "type_schema" in metadata, f"{'add_one_cpu'}: {metadata}" schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[Tensor, Tensor], None]", f"{'add_one_cpu'}: {schema}" doc = mod.get_function_doc("add_one_cpu") assert doc is None x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) def test_build_inline_with_metadata() -> None: # noqa: PLR0915 """Test functions with various input and output types.""" # Keep module alive until all returned objects are destroyed mod: Module = tvm_ffi.cpp.load_inline( name="test_io_types", cpp_sources=r""" // int input -> int output int square(int x) { return x * x; } // float input -> float output float reciprocal(float x) { return 1.0f / x; } // bool input -> bool output bool negate(bool x) { return !x; } // String input -> String output tvm::ffi::String uppercase_first(tvm::ffi::String s) { std::string result(s.c_str()); if (!result.empty()) { result[0] = std::toupper(result[0]); } return tvm::ffi::String(result); } // Multiple inputs: int, float -> float float weighted_sum(int count, float weight) { return static_cast(count) * weight; } // Multiple inputs: String, int -> String tvm::ffi::String repeat_string(tvm::ffi::String s, int times) { std::string result; for (int i = 0; i < times; ++i) { result += s.c_str(); } return tvm::ffi::String(result); } // Mixed types: bool, int, float, String -> String tvm::ffi::String format_data(bool flag, int count, float value, tvm::ffi::String label) { std::ostringstream oss; oss << label.c_str() << ": flag=" << (flag ? "true" : "false") << ", count=" << count << ", value=" << value; return tvm::ffi::String(oss.str()); } // Tensor input/output void double_tensor(tvm::ffi::TensorView input, tvm::ffi::TensorView output) { TVM_FFI_ICHECK(input.ndim() == 1); TVM_FFI_ICHECK(output.ndim() == 1); TVM_FFI_ICHECK(input.size(0) == output.size(0)); DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(input.dtype() == f32_dtype); TVM_FFI_ICHECK(output.dtype() == f32_dtype); for (int i = 0; i < input.size(0); ++i) { static_cast(output.data_ptr())[i] = static_cast(input.data_ptr())[i] * 2.0f; } } """, functions=[ "square", "reciprocal", "negate", "uppercase_first", "weighted_sum", "repeat_string", "format_data", "double_tensor", ], extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"], ) # Test square: int -> int assert mod.square(5) == 25 metadata = mod.get_function_metadata("square") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[int], int]" # Test reciprocal: float -> float result = mod.reciprocal(2.0) assert abs(result - 0.5) < 0.001 metadata = mod.get_function_metadata("reciprocal") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[float], float]" # Test negate: bool -> bool assert mod.negate(True) is False assert mod.negate(False) is True metadata = mod.get_function_metadata("negate") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[bool], bool]" # Test uppercase_first: String -> String result = mod.uppercase_first("hello") assert result == "Hello" metadata = mod.get_function_metadata("uppercase_first") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[str], str]" # Test weighted_sum: int, float -> float result = mod.weighted_sum(10, 2.5) assert abs(result - 25.0) < 0.001 metadata = mod.get_function_metadata("weighted_sum") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[int, float], float]" # Test repeat_string: String, int -> String result = mod.repeat_string("ab", 3) assert result == "ababab" metadata = mod.get_function_metadata("repeat_string") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[str, int], str]" # Test format_data: bool, int, float, String -> String result = mod.format_data(True, 42, 3.14, "test") assert "test:" in result assert "flag=true" in result assert "count=42" in result assert "value=3.14" in result metadata = mod.get_function_metadata("format_data") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[bool, int, float, str], str]" # Test double_tensor: Tensor, Tensor -> None x = numpy.array([1.0, 2.0, 3.0], dtype=numpy.float32) y = numpy.empty_like(x) mod.double_tensor(x, y) numpy.testing.assert_allclose(y, x * 2.0) metadata = mod.get_function_metadata("double_tensor") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[Tensor, Tensor], None]" def test_build_inline_with_docstrings() -> None: """Test building functions with documentation using the functions dict.""" # Keep module alive until all returned objects are destroyed add_docstring = ( "Add two integers and return the sum.\n" "\n" "Parameters\n" "----------\n" "a : int\n" " First integer\n" "b : int\n" " Second integer\n" "\n" "Returns\n" "-------\n" "result : int\n" " Sum of a and b" ) divide_docstring = "Divides two floats. Returns a/b." mod: Module = tvm_ffi.cpp.load_inline( name="test_docs", cpp_sources=r""" int add(int a, int b) { return a + b; } int subtract(int a, int b) { return a - b; } float divide(float a, float b) { TVM_FFI_ICHECK(b != 0.0f) << "Division by zero"; return a / b; } """, functions={ "add": add_docstring, "subtract": "", # No documentation "divide": divide_docstring, }, extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"], ) # Test add function with full documentation assert mod.add(10, 5) == 15 metadata = mod.get_function_metadata("add") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[int, int], int]" doc = mod.get_function_doc("add") assert doc is not None, "add should have documentation" assert doc == add_docstring # Test subtract function without documentation assert mod.subtract(10, 5) == 5 metadata = mod.get_function_metadata("subtract") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[int, int], int]" doc = mod.get_function_doc("subtract") assert doc is None, "subtract should not have documentation" # Test divide function with short documentation result = mod.divide(10.0, 2.0) assert abs(result - 5.0) < 0.001 metadata = mod.get_function_metadata("divide") assert metadata is not None schema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(schema) == "Callable[[float, float], float]" doc = mod.get_function_doc("divide") assert doc is not None, "divide should have documentation" assert doc == divide_docstring def test_build_without_metadata() -> None: """Test building without metadata export.""" mod: Module = tvm_ffi.cpp.load_inline( name="test_no_meta", cpp_sources=r""" // Note: NOT defining TVM_FFI_DLL_EXPORT_INCLUDE_METADATA int simple_add(int a, int b) { return a + b; } """, functions=["simple_add"], ) # Function should still work result = mod.simple_add(10, 20) assert result == 30 # But metadata should not be available metadata = mod.get_function_metadata("simple_add") assert metadata is None, ( "Metadata should not be available without TVM_FFI_DLL_EXPORT_INCLUDE_METADATA" ) # Doc should also not be available doc = mod.get_function_doc("simple_add") assert doc is None if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/test_build_inline.py000066400000000000000000000105601521067262500217040ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import tempfile from pathlib import Path import numpy import pytest import tvm_ffi.cpp from tvm_ffi.module import Module def test_build_inline_cpp() -> None: output_lib_path = tvm_ffi.cpp.build_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions=["add_one_cpu"], ) mod: Module = tvm_ffi.load_module(output_lib_path) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) def test_build_c_file() -> None: """Test building a C source file to .o and linking into .so.""" c_source = r""" #include TVM_FFI_DLL_EXPORT int __tvm_ffi_c_add( void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { result->type_index = kTVMFFIInt; result->zero_padding = 0; result->v_int64 = args[0].v_int64 + args[1].v_int64; return 0; } """ with tempfile.TemporaryDirectory() as tmpdir: c_file = str(Path(tmpdir) / "test_c.c") Path(c_file).write_text(c_source) obj_path = tvm_ffi.cpp.build( name="test_c", sources=[c_file], output="test_c.o", ) assert obj_path.endswith(".o") lib_path = tvm_ffi.cpp.build( name="test_c_lib", sources=[obj_path], ) mod: Module = tvm_ffi.load_module(lib_path) assert mod.c_add(10, 20) == 30 def test_build_inline_object_output() -> None: """Test building a .o then linking it into a .so and calling the function.""" obj_path = tvm_ffi.cpp.build_inline( name="hello_obj", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions=["add_one_cpu"], output="hello_obj.o", ) assert obj_path.endswith(".o") # Link the object file into a shared library lib_path = tvm_ffi.cpp.build( name="hello_from_obj", sources=[obj_path], ) mod: Module = tvm_ffi.load_module(lib_path) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/test_container.py000066400000000000000000000603161521067262500212350ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import pickle import sys from typing import Any import pytest import tvm_ffi from tvm_ffi import testing from tvm_ffi.container import MISSING as CONTAINER_MISSING from tvm_ffi.core import MISSING if sys.version_info >= (3, 9): # PEP 585 generics from collections.abc import Sequence else: # Python 3.8 # workarounds for python 3.8 # typing-module generics (subscriptable on 3.8) from typing import Sequence def test_array() -> None: a = tvm_ffi.convert([1, 2, 3]) assert isinstance(a, tvm_ffi.Array) assert len(a) == 3 assert a[-1] == 3 a_slice = a[-3:-1] assert isinstance(a_slice, list) # TVM array slicing returns a list[T] instead of Array[T] assert (a_slice[0], a_slice[1]) == (1, 2) def test_bad_constructor_init_state() -> None: """Test when error is raised before __init_handle_by_constructor. This case we need the FFI binding to gracefully handle both repr and dealloc by ensuring the chandle is initialized and there is proper repr code """ with pytest.raises(TypeError): tvm_ffi.Array(1) # ty: ignore[invalid-argument-type] with pytest.raises(TypeError): tvm_ffi.List(1) # ty: ignore[invalid-argument-type] with pytest.raises(AttributeError): tvm_ffi.Map(1) # ty: ignore[invalid-argument-type] def test_array_of_array_map() -> None: a = tvm_ffi.convert([[1, 2, 3], {"A": 5, "B": 6}]) assert isinstance(a, tvm_ffi.Array) assert len(a) == 2 assert isinstance(a[0], tvm_ffi.Array) assert isinstance(a[1], tvm_ffi.Map) assert tuple(a[0]) == (1, 2, 3) assert a[1]["A"] == 5 assert a[1]["B"] == 6 def test_int_map() -> None: amap = tvm_ffi.convert({3: 2, 4: 3}) assert 3 in amap assert len(amap) == 2 dd = dict(amap.items()) assert 3 in dd assert 4 in dd assert 5 not in amap assert tuple(amap.items()) == ((3, 2), (4, 3)) assert tuple(amap.keys()) == (3, 4) assert tuple(amap.values()) == (2, 3) def test_array_map_of_opaque_object() -> None: class MyObject: def __init__(self, value: Any) -> None: self.value = value a = tvm_ffi.convert([MyObject("hello"), MyObject(1)]) assert isinstance(a, tvm_ffi.Array) assert len(a) == 2 assert isinstance(a[0], MyObject) assert a[0].value == "hello" assert isinstance(a[1], MyObject) assert a[1].value == 1 y = tvm_ffi.convert({"a": MyObject(1), "b": MyObject("hello")}) assert isinstance(y, tvm_ffi.Map) assert len(y) == 2 assert isinstance(y["a"], MyObject) assert y["a"].value == 1 assert isinstance(y["b"], MyObject) assert y["b"].value == "hello" def test_str_map() -> None: data = [] for i in reversed(range(10)): data.append((f"a{i}", i)) amap = tvm_ffi.convert({k: v for k, v in data}) assert tuple(amap.items()) == tuple(data) for k, v in data: assert k in amap assert amap[k] == v assert amap.get(k) == v assert tuple(k for k in amap) == tuple(k for k, _ in data) def test_key_not_found() -> None: amap = tvm_ffi.convert({3: 2, 4: 3}) with pytest.raises(KeyError): amap[5] def test_repr() -> None: a = tvm_ffi.convert([1, 2, 3]) assert str(a) == "(1, 2, 3)" amap = tvm_ffi.convert({3: 2, 4: 3}) assert str(amap) == "{3: 2, 4: 3}" smap = tvm_ffi.convert({"a": 1, "b": 2}) assert str(smap) == '{"a": 1, "b": 2}' def test_serialization() -> None: a = tvm_ffi.convert([1, 2, 3]) b = pickle.loads(pickle.dumps(a)) assert str(b) == "(1, 2, 3)" @pytest.mark.parametrize( "a, b, c_expected", [ ( tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([4, 5, 6]), tvm_ffi.Array([1, 2, 3, 4, 5, 6]), ), ( tvm_ffi.Array([1, 2, 3]), [4, 5, 6], tvm_ffi.Array([1, 2, 3, 4, 5, 6]), ), ( [1, 2, 3], tvm_ffi.Array([4, 5, 6]), tvm_ffi.Array([1, 2, 3, 4, 5, 6]), ), ( tvm_ffi.Array([]), tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([1, 2, 3]), ), ( tvm_ffi.Array([1, 2, 3]), [], tvm_ffi.Array([1, 2, 3]), ), ( [], tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([1, 2, 3]), ), ( tvm_ffi.Array([]), [], tvm_ffi.Array([]), ), ( tvm_ffi.Array([]), [], tvm_ffi.Array([]), ), ( tvm_ffi.Array([1, 2, 3]), (4, 5, 6), tvm_ffi.Array([1, 2, 3, 4, 5, 6]), ), ( (1, 2, 3), tvm_ffi.Array([4, 5, 6]), tvm_ffi.Array([1, 2, 3, 4, 5, 6]), ), ], ) def test_array_concat( a: Sequence[int], b: Sequence[int], c_expected: Sequence[int], ) -> None: c_actual = a + b # ty: ignore[unsupported-operator] assert type(c_actual) is type(c_expected) assert len(c_actual) == len(c_expected) assert tuple(c_actual) == tuple(c_expected) def test_large_map_get() -> None: amap = tvm_ffi.convert({k: k**2 for k in range(100)}) assert amap.get(101) is None assert amap.get(3) == 9 @pytest.mark.parametrize( "arr, value, expected", [ ([1, 2, 3, 4, 5], 3, True), ([1, 2, 3, 4, 5], 1, True), ([1, 2, 3, 4, 5], 5, True), ([1, 2, 3, 4, 5], 10, False), ([1, 2, 3, 4, 5], 0, False), ([], 1, False), (["hello", "world"], "hello", True), (["hello", "world"], "world", True), (["hello", "world"], "foo", False), ], ) def test_array_contains(arr: list[Any], value: Any, expected: bool) -> None: a = tvm_ffi.convert(arr) assert (value in a) == expected def test_array_contains_plain_tuple() -> None: a = tvm_ffi.Array([("BLOCK_SIZE", 128)]) assert ("BLOCK_SIZE", 128) in a @pytest.mark.parametrize( "arr, expected", [ ([1, 2, 3], True), ([1], True), ([], False), (["hello"], True), ], ) def test_array_bool(arr: list[Any], expected: bool) -> None: a = tvm_ffi.Array(arr) assert bool(a) is expected @pytest.mark.parametrize( "mapping, expected", [ ({"a": 1, "b": 2}, True), ({"a": 1}, True), ({}, False), ({1: "one"}, True), ], ) def test_map_bool(mapping: dict[Any, Any], expected: bool) -> None: m = tvm_ffi.Map(mapping) assert bool(m) is expected def test_list_basic() -> None: lst = tvm_ffi.List([1, 2, 3]) assert isinstance(lst, tvm_ffi.List) assert len(lst) == 3 assert tuple(lst) == (1, 2, 3) assert lst[0] == 1 assert lst[-1] == 3 assert lst[:] == [1, 2, 3] assert lst[::-1] == [3, 2, 1] def test_list_mutation_methods() -> None: lst = tvm_ffi.List([1, 2, 3]) lst.append(4) assert tuple(lst) == (1, 2, 3, 4) lst.insert(2, 9) assert tuple(lst) == (1, 2, 9, 3, 4) value = lst.pop() assert value == 4 assert tuple(lst) == (1, 2, 9, 3) value = lst.pop(1) assert value == 2 assert tuple(lst) == (1, 9, 3) lst.extend([5, 6]) assert tuple(lst) == (1, 9, 3, 5, 6) lst.clear() assert tuple(lst) == () assert len(lst) == 0 def test_list_setitem_and_delitem() -> None: lst = tvm_ffi.List([0, 1, 2, 3, 4, 5]) lst[1] = 10 lst[-1] = 60 assert tuple(lst) == (0, 10, 2, 3, 4, 60) del lst[2] assert tuple(lst) == (0, 10, 3, 4, 60) del lst[-1] assert tuple(lst) == (0, 10, 3, 4) def test_list_slice_assignment_and_delete() -> None: lst = tvm_ffi.List([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) lst[2:6] = [20, 21] assert tuple(lst) == (0, 1, 20, 21, 6, 7, 8, 9) lst[3:3] = [30, 31] assert tuple(lst) == (0, 1, 20, 30, 31, 21, 6, 7, 8, 9) lst[1:9:2] = [101, 102, 103, 104] assert tuple(lst) == (0, 101, 20, 102, 31, 103, 6, 104, 8, 9) with pytest.raises(ValueError): lst[0:6:2] = [1, 2] del lst[1:8:2] assert tuple(lst) == (0, 20, 31, 6, 8, 9) del lst[2:2] assert tuple(lst) == (0, 20, 31, 6, 8, 9) del lst[1:4] assert tuple(lst) == (0, 8, 9) def test_list_slice_edge_cases() -> None: lst = tvm_ffi.List([0, 1, 2, 3, 4]) lst[3:1] = [9] assert tuple(lst) == (0, 1, 2, 9, 3, 4) lst = tvm_ffi.List([0, 1, 2, 3]) del lst[3:1] assert tuple(lst) == (0, 1, 2, 3) lst = tvm_ffi.List([0, 1, 2, 3, 4, 5]) del lst[5:1:-1] assert tuple(lst) == (0, 1) lst = tvm_ffi.List([0, 1, 2, 3]) del lst[::-1] assert tuple(lst) == () def test_list_contains_bool_repr_and_concat() -> None: lst = tvm_ffi.List([1, 2, 3]) assert 2 in lst assert 5 not in lst assert bool(lst) is True assert str(lst) == "[1, 2, 3]" assert tuple(lst.__add__([4, 5])) == (1, 2, 3, 4, 5) assert tuple(lst.__radd__([0])) == (0, 1, 2, 3) empty = tvm_ffi.List() assert bool(empty) is False assert str(empty) == "[]" def test_list_insert_index_normalization() -> None: lst = tvm_ffi.List([1, 2, 3]) lst.insert(-100, 0) assert tuple(lst) == (0, 1, 2, 3) lst.insert(100, 4) assert tuple(lst) == (0, 1, 2, 3, 4) def test_list_error_cases() -> None: lst = tvm_ffi.List([1, 2, 3]) with pytest.raises(IndexError): _ = lst[3] with pytest.raises(IndexError): _ = lst[-4] with pytest.raises(IndexError): lst[3] = 0 with pytest.raises(IndexError): del lst[3] with pytest.raises(IndexError): tvm_ffi.List().pop() def test_list_pickle_roundtrip() -> None: lst = tvm_ffi.List([1, "a", {"k": 2}]) restored = pickle.loads(pickle.dumps(lst)) assert isinstance(restored, tvm_ffi.List) assert restored[0] == 1 assert restored[1] == "a" assert isinstance(restored[2], tvm_ffi.Map) assert restored[2]["k"] == 2 def test_list_reverse() -> None: lst = tvm_ffi.List([3, 1, 2]) lst.reverse() assert tuple(lst) == (2, 1, 3) empty = tvm_ffi.List() empty.reverse() assert tuple(empty) == () single = tvm_ffi.List([42]) single.reverse() assert tuple(single) == (42,) def test_list_self_aliasing_slice_assignment() -> None: lst = tvm_ffi.List([0, 1, 2, 3, 4]) lst[1:3] = lst assert tuple(lst) == (0, 0, 1, 2, 3, 4, 3, 4) def test_list_is_mutable_and_shared() -> None: lst = tvm_ffi.List([1, 2]) alias = lst alias.append(3) alias[0] = 10 assert tuple(lst) == (10, 2, 3) # --------------------------------------------------------------------------- # Cross-conversion tests: Array <-> List via SeqTypeTraitsBase # --------------------------------------------------------------------------- def test_seq_cross_conv_list_to_array_int() -> None: """List passed to a function expecting Array (new behavior).""" lst = tvm_ffi.List([10, 20, 30]) result = testing.schema_id_arr_int(lst) assert isinstance(result, tvm_ffi.Array) assert list(result) == [10, 20, 30] def test_seq_cross_conv_array_to_list_int() -> None: """Array passed to a function expecting List.""" arr = tvm_ffi.Array([10, 20, 30]) result = testing.schema_id_list_int(arr) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.List) assert list(result) == [10, 20, 30] def test_seq_cross_conv_list_to_array_str() -> None: """List passed to a function expecting Array.""" lst = tvm_ffi.List(["a", "b", "c"]) result = testing.schema_id_arr_str(lst) assert isinstance(result, tvm_ffi.Array) assert list(result) == ["a", "b", "c"] def test_seq_cross_conv_array_to_list_str() -> None: """Array passed to a function expecting List.""" arr = tvm_ffi.Array(["a", "b", "c"]) result = testing.schema_id_list_str(arr) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.List) assert list(result) == ["a", "b", "c"] def test_seq_cross_conv_empty_list_to_array() -> None: """Empty List passed to a function expecting Array.""" lst = tvm_ffi.List([]) result = testing.schema_id_arr_int(lst) assert isinstance(result, tvm_ffi.Array) assert len(result) == 0 def test_seq_cross_conv_empty_array_to_list() -> None: """Empty Array passed to a function expecting List.""" arr = tvm_ffi.Array([]) result = testing.schema_id_list_int(arr) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.List) assert len(result) == 0 def test_seq_cross_conv_python_list_to_array() -> None: """Plain Python list passed to Array function.""" result = testing.schema_id_arr_int([1, 2, 3]) assert isinstance(result, tvm_ffi.Array) assert list(result) == [1, 2, 3] def test_seq_cross_conv_python_list_to_list() -> None: """Plain Python list passed to List function.""" result = testing.schema_id_list_int([1, 2, 3]) assert isinstance(result, tvm_ffi.List) assert list(result) == [1, 2, 3] def test_seq_cross_conv_incompatible_list_to_array() -> None: """List with incompatible element types should fail when cast to Array.""" lst = tvm_ffi.List(["not", "ints"]) with pytest.raises(TypeError): testing.schema_id_arr_int(lst) # type: ignore[invalid-argument-type] def test_seq_cross_conv_incompatible_array_to_list() -> None: """Array with incompatible element types should fail when cast to List.""" arr = tvm_ffi.Array(["not", "ints"]) with pytest.raises(TypeError): testing.schema_id_list_int(arr) # type: ignore[invalid-argument-type] def test_missing_object() -> None: """Test that MISSING is a valid singleton and works with Map.get().""" # MISSING should be an Object instance assert isinstance(MISSING, tvm_ffi.Object) # MISSING should be the same object across imports assert MISSING.same_as(CONTAINER_MISSING) # Map.get() should use MISSING internally m = tvm_ffi.Map({"a": 1}) assert m.get("a") == 1 assert m.get("b") is None assert m.get("b", 42) == 42 # --------------------------------------------------------------------------- # Dict tests # --------------------------------------------------------------------------- def test_dict_basic() -> None: d = tvm_ffi.Dict({"a": 1, "b": 2}) assert len(d) == 2 assert "a" in d assert d["a"] == 1 assert d["b"] == 2 assert "c" not in d def test_dict_setitem_delitem() -> None: d = tvm_ffi.Dict({"a": 1}) d["b"] = 2 assert len(d) == 2 assert d["b"] == 2 del d["a"] assert len(d) == 1 assert "a" not in d with pytest.raises(KeyError): del d["nonexistent"] def test_dict_shared_mutation() -> None: d1 = tvm_ffi.Dict({"a": 1}) d2 = d1 # alias, same underlying object d1["b"] = 2 assert d2["b"] == 2 assert len(d2) == 2 def test_dict_keys_values_items() -> None: d = tvm_ffi.Dict({"x": 10, "y": 20}) keys = set(d.keys()) assert keys == {"x", "y"} values = set(d.values()) assert values == {10, 20} items = set(d.items()) assert items == {("x", 10), ("y", 20)} def test_dict_get_with_default() -> None: d = tvm_ffi.Dict({"a": 1}) assert d.get("a") == 1 assert d.get("b") is None assert d.get("b", 42) == 42 def test_dict_pop() -> None: d = tvm_ffi.Dict({"a": 1, "b": 2}) val = d.pop("a") assert val == 1 assert len(d) == 1 # pop with default val2 = d.pop("nonexistent", 99) assert val2 == 99 # pop missing without default with pytest.raises(KeyError): d.pop("nonexistent") def test_dict_clear() -> None: d = tvm_ffi.Dict({"a": 1, "b": 2}) d.clear() assert len(d) == 0 def test_dict_update() -> None: d = tvm_ffi.Dict({"a": 1}) d.update({"b": 2, "c": 3}) assert len(d) == 3 assert d["b"] == 2 def test_dict_iteration() -> None: d = tvm_ffi.Dict({"a": 1, "b": 2}) keys = list(d) assert set(keys) == {"a", "b"} def test_dict_iteration_many_entries() -> None: """Regression: ensure iterator visits ALL entries, not just the first.""" entries = {f"key_{i}": i for i in range(10)} d = tvm_ffi.Dict(entries) # keys keys_list = list(d.keys()) assert len(keys_list) == 10 assert set(keys_list) == set(entries.keys()) # values values_list = list(d.values()) assert len(values_list) == 10 assert set(values_list) == set(entries.values()) # items items_list = list(d.items()) assert len(items_list) == 10 assert set(items_list) == set(entries.items()) # direct iteration (keys) iter_keys = list(d) assert len(iter_keys) == 10 assert set(iter_keys) == set(entries.keys()) def test_dict_iteration_after_mutation() -> None: """Regression: iteration after insert/delete visits correct elements.""" d = tvm_ffi.Dict({"a": 1, "b": 2, "c": 3}) d["d"] = 4 del d["b"] # should have a, c, d keys = list(d.keys()) assert len(keys) == 3 assert set(keys) == {"a", "c", "d"} values = list(d.values()) assert len(values) == 3 assert set(values) == {1, 3, 4} items = list(d.items()) assert len(items) == 3 assert set(items) == {("a", 1), ("c", 3), ("d", 4)} def test_dict_repr() -> None: d = tvm_ffi.Dict({"a": 1}) r = repr(d) assert isinstance(r, str) # repr should look dict-like assert ":" in r def test_dict_bool() -> None: assert not tvm_ffi.Dict() assert tvm_ffi.Dict({"a": 1}) def test_dict_empty_init() -> None: d = tvm_ffi.Dict() assert len(d) == 0 d["a"] = 1 assert d["a"] == 1 def test_dict_pickle_roundtrip() -> None: d = tvm_ffi.Dict({"a": 1, "b": 2}) data = pickle.dumps(d) d2 = pickle.loads(data) assert isinstance(d2, tvm_ffi.Dict) assert len(d2) == 2 assert d2["a"] == 1 assert d2["b"] == 2 # --------------------------------------------------------------------------- # Map <-> Dict cross-conversion tests # --------------------------------------------------------------------------- def test_map_cross_conv_dict_to_map_str_int() -> None: """Dict passed to a function expecting Map.""" d = tvm_ffi.Dict({"a": 1, "b": 2}) result = testing.schema_id_map_str_int(d) assert isinstance(result, tvm_ffi.Map) assert dict(result.items()) == {"a": 1, "b": 2} def test_map_cross_conv_map_to_dict_str_int() -> None: """Map passed to a function expecting Dict.""" m = tvm_ffi.Map({"a": 1, "b": 2}) result = testing.schema_id_dict_str_int(m) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.Dict) assert result["a"] == 1 assert result["b"] == 2 def test_map_cross_conv_dict_to_map_str_str() -> None: """Dict passed to a function expecting Map.""" d = tvm_ffi.Dict({"x": "hello", "y": "world"}) result = testing.schema_id_map_str_str(d) assert isinstance(result, tvm_ffi.Map) assert dict(result.items()) == {"x": "hello", "y": "world"} def test_map_cross_conv_map_to_dict_str_str() -> None: """Map passed to a function expecting Dict.""" m = tvm_ffi.Map({"x": "hello", "y": "world"}) result = testing.schema_id_dict_str_str(m) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.Dict) assert result["x"] == "hello" assert result["y"] == "world" def test_map_cross_conv_empty_dict_to_map() -> None: """Empty Dict passed to a function expecting Map.""" d = tvm_ffi.Dict({}) result = testing.schema_id_map_str_int(d) assert isinstance(result, tvm_ffi.Map) assert len(result) == 0 def test_map_cross_conv_empty_map_to_dict() -> None: """Empty Map passed to a function expecting Dict.""" m = tvm_ffi.Map({}) result = testing.schema_id_dict_str_int(m) # type: ignore[invalid-argument-type] assert isinstance(result, tvm_ffi.Dict) assert len(result) == 0 def test_map_cross_conv_incompatible_dict_to_map() -> None: """Dict with incompatible value types should fail when cast to Map.""" d = tvm_ffi.Dict({"a": "not_int", "b": "still_not_int"}) with pytest.raises(TypeError): testing.schema_id_map_str_int(d) # type: ignore[invalid-argument-type] def test_map_cross_conv_incompatible_map_to_dict() -> None: """Map with incompatible value types should fail when cast to Dict.""" m = tvm_ffi.Map({"a": "not_int", "b": "still_not_int"}) with pytest.raises(TypeError): testing.schema_id_dict_str_int(m) # type: ignore[invalid-argument-type] # --------------------------------------------------------------------------- # Structural __eq__ / __ne__ / __hash__ tests # --------------------------------------------------------------------------- def test_array_structural_eq() -> None: a = tvm_ffi.Array([1, 2, 3]) b = tvm_ffi.Array([1, 2, 3]) c = tvm_ffi.Array([1, 2, 4]) assert a == b assert a != c assert not (a != b) assert not (a == c) def test_array_eq_empty() -> None: assert tvm_ffi.Array([]) == tvm_ffi.Array([]) def test_array_eq_nested() -> None: a = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3])]) b = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3])]) c = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([4])]) assert a == b assert a != c def test_array_eq_not_implemented_for_unrelated() -> None: a = tvm_ffi.Array([1, 2, 3]) assert a == [1, 2, 3] assert a == (1, 2, 3) assert not (a != [1, 2, 3]) assert a.__eq__("hello") is NotImplemented def test_array_hash() -> None: a = tvm_ffi.Array([1, 2, 3]) b = tvm_ffi.Array([1, 2, 3]) assert hash(a) == hash(b) # Usable in sets and as dict keys s = {a, b} assert len(s) == 1 d = {a: "value"} assert d[b] == "value" def test_list_structural_eq() -> None: a = tvm_ffi.List([1, 2, 3]) b = tvm_ffi.List([1, 2, 3]) c = tvm_ffi.List([1, 2, 4]) assert a == b assert a != c def test_list_eq_empty() -> None: assert tvm_ffi.List([]) == tvm_ffi.List([]) def test_list_eq_not_implemented_for_unrelated() -> None: a = tvm_ffi.List([1, 2, 3]) assert a == [1, 2, 3] assert a == (1, 2, 3) assert not (a != [1, 2, 3]) def test_list_contains_plain_tuple() -> None: a = tvm_ffi.List([("BLOCK_SIZE", 128)]) assert ("BLOCK_SIZE", 128) in a def test_list_hash() -> None: a = tvm_ffi.List([1, 2, 3]) b = tvm_ffi.List([1, 2, 3]) assert hash(a) == hash(b) def test_map_structural_eq() -> None: a = tvm_ffi.Map({"x": 1, "y": 2}) b = tvm_ffi.Map({"x": 1, "y": 2}) c = tvm_ffi.Map({"x": 1, "y": 3}) assert a == b assert a != c def test_map_eq_empty() -> None: assert tvm_ffi.Map({}) == tvm_ffi.Map({}) def test_map_eq_not_implemented_for_unrelated() -> None: a = tvm_ffi.Map({"x": 1}) assert a.__eq__({"x": 1}) is NotImplemented def test_map_hash() -> None: a = tvm_ffi.Map({"x": 1, "y": 2}) b = tvm_ffi.Map({"x": 1, "y": 2}) assert hash(a) == hash(b) s = {a, b} assert len(s) == 1 def test_dict_structural_eq() -> None: a = tvm_ffi.Dict({"x": 1, "y": 2}) b = tvm_ffi.Dict({"x": 1, "y": 2}) c = tvm_ffi.Dict({"x": 1, "y": 3}) assert a == b assert a != c def test_dict_eq_empty() -> None: assert tvm_ffi.Dict({}) == tvm_ffi.Dict({}) def test_dict_eq_not_implemented_for_unrelated() -> None: a = tvm_ffi.Dict({"x": 1}) assert a.__eq__({"x": 1}) is NotImplemented def test_dict_hash() -> None: a = tvm_ffi.Dict({"x": 1, "y": 2}) b = tvm_ffi.Dict({"x": 1, "y": 2}) assert hash(a) == hash(b) tvm-ffi-0.1.12/tests/python/test_container_dlpack_conversion.py000066400000000000000000000170351521067262500250200ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for lazy container DLPack conversion when DLPack exchange API is active.""" from __future__ import annotations import numpy as np import pytest try: import torch import torch.version except ImportError: torch = None # ty: ignore[invalid-assignment] import tvm_ffi pytestmark = pytest.mark.skipif(torch is None, reason="torch is not installed") def test_array_tensor_only() -> None: """Array stays as Array; element access converts to torch.Tensor.""" assert torch is not None x = torch.arange(8, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_array_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Array) assert len(result) == 1 elem = result[0] assert isinstance(elem, torch.Tensor) assert elem.data_ptr() == x.data_ptr() def test_array_mixed() -> None: """Array with Tensor + int + string: lazy conversion on access.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_array_with_mixed") result = f(x, 42) assert isinstance(result, tvm_ffi.Array) assert len(result) == 3 assert isinstance(result[0], torch.Tensor) assert result[0].data_ptr() == x.data_ptr() assert result[1] == 42 assert result[2] == "hello" def test_array_nested() -> None: """Nested Array>: inner arrays also get tagged.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_nested_array_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Array) assert len(result) == 2 # First element is inner array inner = result[0] assert isinstance(inner, tvm_ffi.Array) assert len(inner) == 2 assert isinstance(inner[0], torch.Tensor) assert inner[0].data_ptr() == x.data_ptr() assert inner[1] == 42 # Second element is a tensor assert isinstance(result[1], torch.Tensor) assert result[1].data_ptr() == x.data_ptr() def test_list_with_tensor() -> None: """List with tensor: stays as List, elements convert on access.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_list_with_tensor") result = f(x, 7) assert isinstance(result, tvm_ffi.List) assert len(result) == 2 assert isinstance(result[0], torch.Tensor) assert result[0].data_ptr() == x.data_ptr() assert result[1] == 7 def test_map_with_tensor() -> None: """Map with tensor value: stays as Map, values convert on access.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_map_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Map) assert len(result) == 3 assert isinstance(result["tensor"], torch.Tensor) assert result["tensor"].data_ptr() == x.data_ptr() assert result["value"] == 42 assert result["name"] == "test" def test_dict_with_tensor() -> None: """Dict with tensor value: stays as Dict, values convert on access.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_dict_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Dict) assert len(result) == 2 assert isinstance(result["tensor"], torch.Tensor) assert result["tensor"].data_ptr() == x.data_ptr() assert result["value"] == 42 def test_nested_map_with_array() -> None: """Nested Map with Array values: all containers tagged, lazy conversion on access.""" assert torch is not None x1 = torch.arange(4, dtype=torch.float32) x2 = torch.arange(8, dtype=torch.int32) f = tvm_ffi.get_global_func("testing.make_nested_map_with_tensor") result = f(x1, x2) assert isinstance(result, tvm_ffi.Map) # "array" -> Array with tagged tensors arr = result["array"] assert isinstance(arr, tvm_ffi.Array) assert len(arr) == 2 assert isinstance(arr[0], torch.Tensor) assert isinstance(arr[1], torch.Tensor) # "map" -> nested Map inner_map = result["map"] assert isinstance(inner_map, tvm_ffi.Map) assert isinstance(inner_map["t"], torch.Tensor) # "scalar" -> int assert result["scalar"] == 99 def test_empty_array() -> None: """Empty Array with torch input: stays as empty Array.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_empty_array_with_tensor_input") result = f(x) assert isinstance(result, tvm_ffi.Array) assert len(result) == 0 def test_no_torch_input_no_conversion() -> None: """Without torch tensor input, containers stay as FFI types with no tag.""" x = tvm_ffi.from_dlpack(np.arange(4, dtype="float32")) f = tvm_ffi.get_global_func("testing.make_array_with_tensor") result = f(x) # No torch input, so no dlpack API set -> normal FFI Array return assert isinstance(result, tvm_ffi.Array) assert isinstance(result[0], tvm_ffi.Tensor) def test_data_correctness() -> None: """Verify tensor data is correct after lazy container conversion.""" assert torch is not None x = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_array_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Array) elem = result[0] assert isinstance(elem, torch.Tensor) np.testing.assert_equal(elem.numpy(), x.numpy()) def test_echo_bare_tensor_unchanged() -> None: """Existing behavior: bare tensor return still works.""" assert torch is not None x = torch.arange(128) fecho = tvm_ffi.get_global_func("testing.echo") y = fecho(x) assert isinstance(y, torch.Tensor) assert y.data_ptr() == x.data_ptr() def test_container_preserves_identity() -> None: """Lazy conversion preserves container identity (can be passed back to FFI).""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_array_with_tensor") result = f(x) assert isinstance(result, tvm_ffi.Array) # Pass container back to FFI (echo) fecho = tvm_ffi.get_global_func("testing.echo") echoed = fecho(result) assert isinstance(echoed, tvm_ffi.Array) assert isinstance(echoed[0], torch.Tensor) assert echoed[0].data_ptr() == x.data_ptr() def test_mutable_list_shared_semantics() -> None: """Lazy conversion preserves mutable list shared-reference semantics.""" assert torch is not None x = torch.arange(4, dtype=torch.float32) f = tvm_ffi.get_global_func("testing.make_list_with_tensor") result = f(x, 7) assert isinstance(result, tvm_ffi.List) # The result is the actual FFI List, not a detached copy assert result.same_as(result) tvm-ffi-0.1.12/tests/python/test_cubin_launcher.py000066400000000000000000000334041521067262500222320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Test CUBIN launcher functionality using load_inline.""" from __future__ import annotations import subprocess import sys import tempfile from pathlib import Path import pytest try: import torch import torch.version except ImportError: torch = None # ty: ignore[invalid-assignment] import tvm_ffi.cpp # Check if CUDA is available def _is_cuda_available() -> bool: """Check if CUDA is available for testing.""" if torch is None: return False return torch.cuda.is_available() def _is_cuda_version_greater_than_13() -> bool: """Check if CUDA version is greater than 13.0.""" if torch is None or not torch.cuda.is_available(): return False if torch.version.cuda is None: return False try: # Parse version string into tuple of integers (e.g., "12.1" -> (12, 1)) version_parts = tuple(int(x) for x in torch.version.cuda.split(".")) return version_parts > (13, 0) except (ValueError, TypeError, AttributeError): return False def _compile_kernel_to_cubin() -> bytes: """Compile simple CUDA kernels to CUBIN. Returns the raw CUBIN bytes. """ cuda_code = r""" extern "C" __global__ void add_one_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1.0f; } } extern "C" __global__ void mul_two_cuda(const float* x, float* y, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] * 2.0f; } } """ with tempfile.TemporaryDirectory() as tmpdir: tmppath = Path(tmpdir) cu_file = tmppath / "kernels.cu" cubin_file = tmppath / "kernels.cubin" cu_file.write_text(cuda_code) # Compile to CUBIN using nvcc result = subprocess.run( ["nvcc", "--cubin", "-arch=native", str(cu_file), "-o", str(cubin_file)], capture_output=True, text=True, check=False, ) if result.returncode != 0: pytest.skip(f"nvcc not available or compilation failed: {result.stderr}") # ty: ignore[invalid-argument-type, too-many-positional-arguments] return cubin_file.read_bytes() @pytest.mark.skipif(sys.platform != "linux", reason="CUBIN launcher only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") @pytest.mark.skipif( not _is_cuda_version_greater_than_13(), reason="CUDA version must be greater than 13.0" ) def test_cubin_launcher_add_one() -> None: """Test loading and launching add_one kernel from CUBIN.""" assert torch is not None, "PyTorch is required for this test" cubin_bytes = _compile_kernel_to_cubin() # Define C++ code to load and launch the CUBIN kernel cpp_code = """ #include #include #include #include #include #include #include #include namespace cubin_test { static std::unique_ptr g_module; static std::unique_ptr g_kernel_add_one; static std::unique_ptr g_kernel_mul_two; void LoadCubinData(const tvm::ffi::Bytes& cubin_data) { // Load CUBIN from bytes g_module = std::make_unique(cubin_data); g_kernel_add_one = std::make_unique((*g_module)["add_one_cuda"]); g_kernel_mul_two = std::make_unique((*g_module)["mul_two_cuda"]); } void LaunchAddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_module != nullptr, RuntimeError) << "CUBIN module not loaded"; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Sizes must match"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); void* args[] = { reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n), }; tvm::ffi::dim3 grid((n + 1023) / 1024); tvm::ffi::dim3 block(1024); DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); auto result = g_kernel_add_one->Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } void LaunchMulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_module != nullptr, RuntimeError) << "CUBIN module not loaded"; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Sizes must match"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); void* args[] = { reinterpret_cast(&x_ptr), reinterpret_cast(&y_ptr), reinterpret_cast(&n), }; tvm::ffi::dim3 grid((n + 1023) / 1024); tvm::ffi::dim3 block(1024); DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); auto result = g_kernel_mul_two->Launch(args, grid, block, stream); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(load_cubin_data, cubin_test::LoadCubinData); TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_add_one, cubin_test::LaunchAddOne); TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_mul_two, cubin_test::LaunchMulTwo); } // namespace cubin_test """ # Compile and load the C++ code mod = tvm_ffi.cpp.load_inline( "cubin_test", cuda_sources=cpp_code, extra_ldflags=["-lcudart"], ) # Load CUBIN from bytes load_fn = mod["load_cubin_data"] load_fn(cubin_bytes) # Test add_one kernel launch_add_one = mod["launch_add_one"] n = 256 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") launch_add_one(x, y) expected = x + 1 torch.testing.assert_close(y, expected) # Test mul_two kernel launch_mul_two = mod["launch_mul_two"] x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5 y = torch.empty(n, dtype=torch.float32, device="cuda") launch_mul_two(x, y) expected = x * 2 torch.testing.assert_close(y, expected) @pytest.mark.skipif(sys.platform != "linux", reason="CUBIN launcher only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") @pytest.mark.skipif( not _is_cuda_version_greater_than_13(), reason="CUDA version must be greater than 13.0" ) def test_cubin_launcher_launch_ex() -> None: """Test LaunchEx with ConstructLaunchConfig (no clustering).""" assert torch is not None, "PyTorch is required for this test" cubin_bytes = _compile_kernel_to_cubin() cpp_code = """ #include #include #include #include #include #include namespace cubin_test_launch_ex { static std::unique_ptr g_module; static std::unique_ptr g_kernel_add_one; void LoadCubinData(const tvm::ffi::Bytes& cubin_data) { g_module = std::make_unique(cubin_data); g_kernel_add_one = std::make_unique((*g_module)["add_one_cuda"]); } void LaunchAddOneEx(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_module != nullptr, RuntimeError) << "CUBIN module not loaded"; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; TVM_FFI_CHECK(y.ndim() == 1, ValueError) << "Output must be 1D tensor"; TVM_FFI_CHECK(x.size(0) == y.size(0), ValueError) << "Sizes must match"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); void* args[] = {&x_ptr, &y_ptr, &n}; tvm::ffi::dim3 grid((n + 1023) / 1024); tvm::ffi::dim3 block(1024); DLDevice device = x.device(); auto stream = static_cast( TVMFFIEnvGetStream(device.device_type, device.device_id)); // Use ConstructLaunchConfig + LaunchEx (cluster_dim=1 means no clustering) tvm::ffi::cuda_api::LaunchConfig config; tvm::ffi::cuda_api::LaunchAttrType attr; auto err = tvm::ffi::cuda_api::ConstructLaunchConfig( g_kernel_add_one->GetHandle(), stream, /*smem_size=*/0, grid, block, /*cluster_dim=*/1, config, attr); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(err); auto result = g_kernel_add_one->LaunchEx(args, config); TVM_FFI_CHECK_CUBIN_LAUNCHER_CUDA_ERROR(result); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(load_cubin_data, cubin_test_launch_ex::LoadCubinData); TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_add_one_ex, cubin_test_launch_ex::LaunchAddOneEx); } // namespace cubin_test_launch_ex """ mod = tvm_ffi.cpp.load_inline( "cubin_test_launch_ex", cuda_sources=cpp_code, extra_ldflags=["-lcudart"], ) load_fn = mod["load_cubin_data"] load_fn(cubin_bytes) launch_add_one_ex = mod["launch_add_one_ex"] n = 256 x = torch.arange(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") launch_add_one_ex(x, y) expected = x + 1 torch.testing.assert_close(y, expected) @pytest.mark.skipif(sys.platform != "linux", reason="CUBIN launcher only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") @pytest.mark.skipif( not _is_cuda_version_greater_than_13(), reason="CUDA version must be greater than 13.0" ) def test_cubin_launcher_chained() -> None: """Test chaining multiple kernel launches.""" assert torch is not None, "PyTorch is required for this test" cubin_bytes = _compile_kernel_to_cubin() cpp_code = """ #include #include #include #include #include #include namespace cubin_test_chain { static std::unique_ptr g_module; static std::unique_ptr g_kernel_add_one; static std::unique_ptr g_kernel_mul_two; void LoadCubinData(const tvm::ffi::Bytes& cubin_data) { // Load CUBIN from bytes g_module = std::make_unique(cubin_data); g_kernel_add_one = std::make_unique((*g_module)["add_one_cuda"]); g_kernel_mul_two = std::make_unique((*g_module)["mul_two_cuda"]); } void LaunchAddOne(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_module != nullptr, RuntimeError) << "CUBIN module not loaded"; TVM_FFI_CHECK(x.ndim() == 1, ValueError) << "Input must be 1D tensor"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); void* args[] = {&x_ptr, &y_ptr, &n}; tvm::ffi::dim3 grid((n + 1023) / 1024); tvm::ffi::dim3 block(1024); DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); g_kernel_add_one->Launch(args, grid, block, stream); } void LaunchMulTwo(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { TVM_FFI_CHECK(g_module != nullptr, RuntimeError) << "CUBIN module not loaded"; int64_t n = x.size(0); void* x_ptr = x.data_ptr(); void* y_ptr = y.data_ptr(); void* args[] = {&x_ptr, &y_ptr, &n}; tvm::ffi::dim3 grid((n + 1023) / 1024); tvm::ffi::dim3 block(1024); DLDevice device = x.device(); cudaStream_t stream = static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); g_kernel_mul_two->Launch(args, grid, block, stream); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(load_cubin_data, cubin_test_chain::LoadCubinData); TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_add_one, cubin_test_chain::LaunchAddOne); TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_mul_two, cubin_test_chain::LaunchMulTwo); } // namespace cubin_test_chain """ mod = tvm_ffi.cpp.load_inline("cubin_test_chain", cuda_sources=cpp_code) # Load CUBIN from bytes load_fn = mod["load_cubin_data"] load_fn(cubin_bytes) launch_add_one = mod["launch_add_one"] launch_mul_two = mod["launch_mul_two"] # Test chained execution: (x + 1) * 2 n = 128 x = torch.full((n,), 5.0, dtype=torch.float32, device="cuda") temp = torch.empty(n, dtype=torch.float32, device="cuda") y = torch.empty(n, dtype=torch.float32, device="cuda") launch_add_one(x, temp) # temp = x + 1 = 6 launch_mul_two(temp, y) # y = temp * 2 = 12 expected = torch.full((n,), 12.0, dtype=torch.float32, device="cuda") torch.testing.assert_close(y, expected) tvm-ffi-0.1.12/tests/python/test_current_work_stream_gpu.py000066400000000000000000000073031521067262500242220ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file to # you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import ctypes import pytest try: import torch import tvm_ffi # noqa: F401 from torch.utils import cpp_extension from tvm_ffi import libinfo except ImportError: torch = None # ty: ignore[invalid-assignment] if torch is None: _HAS_TORCH = False _HAS_GPU = False _HAS_DLPACK_EXCHANGE_API = False else: _HAS_TORCH = True _HAS_GPU = bool(torch.cuda.is_available()) _HAS_DLPACK_EXCHANGE_API = bool(hasattr(torch.Tensor, "__dlpack_c_exchange_api__")) @pytest.mark.skipif(not _HAS_TORCH, reason="Requires torch") @pytest.mark.skipif(not _HAS_GPU, reason="Requires GPU runtime") @pytest.mark.skipif(not _HAS_DLPACK_EXCHANGE_API, reason="Requires __dlpack_c_exchange_api__") def test_current_work_stream_matches_torch_stream() -> None: assert torch is not None api_attr = torch.Tensor.__dlpack_c_exchange_api__ # ty: ignore[unresolved-attribute] pythonapi = ctypes.pythonapi pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, b"dlpack_exchange_api") assert api_ptr != 0 source = r""" #include #include void assert_current_work_stream(int64_t api_ptr_int, bool is_hip, int64_t expected_stream) { DLPackExchangeAPI* api = reinterpret_cast(api_ptr_int); TORCH_CHECK(api != nullptr, "API pointer is NULL"); TORCH_CHECK(api->current_work_stream != nullptr, "current_work_stream is NULL"); void* stream_cuda = nullptr; int result_cuda = api->current_work_stream(kDLCUDA, 0, &stream_cuda); TORCH_CHECK(result_cuda == 0, "current_work_stream(kDLCUDA) failed"); TORCH_CHECK(reinterpret_cast(stream_cuda) == expected_stream, "kDLCUDA stream mismatch"); if (is_hip) { void* stream_rocm = nullptr; int result_rocm = api->current_work_stream(kDLROCM, 0, &stream_rocm); TORCH_CHECK(result_rocm == 0, "current_work_stream(kDLROCM) failed"); TORCH_CHECK(reinterpret_cast(stream_rocm) == expected_stream, "kDLROCM stream mismatch"); } } """ include_paths = libinfo.include_paths() include_paths += cpp_extension.include_paths("cuda") mod = cpp_extension.load_inline( name="test_current_work_stream_gpu_ext", cpp_sources=[source], functions=["assert_current_work_stream"], with_cuda=torch.cuda.is_available(), extra_include_paths=include_paths, ) device_id = torch.cuda.current_device() is_hip = torch.version.hip is not None stream = torch.cuda.Stream(device=device_id) with torch.cuda.stream(stream): expected_stream = int(stream.cuda_stream) mod.assert_current_work_stream(api_ptr, is_hip, expected_stream) tvm-ffi-0.1.12/tests/python/test_dataclass_c_class.py000066400000000000000000000654231521067262500227050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for the c_class decorator (register_object + structural dunders).""" from __future__ import annotations import inspect import warnings import pytest import tvm_ffi.testing from tvm_ffi.core import MISSING, TypeInfo from tvm_ffi.dataclasses import Field from tvm_ffi.dataclasses.c_class import _attach_field_objects from tvm_ffi.registry import _warn_missing_field_annotations from tvm_ffi.testing import ( TestCompare, TestHash, _TestCxxClassBase, _TestCxxClassDerived, _TestCxxClassDerivedDerived, _TestCxxInitSubset, _TestCxxKwOnly, ) # --------------------------------------------------------------------------- # 1. Custom __init__ preservation # --------------------------------------------------------------------------- def test_c_class_custom_init() -> None: """c_class preserves user-defined __init__.""" obj = _TestCxxClassBase(v_i64=10, v_i32=20) assert obj.v_i64 == 11 # +1 from custom __init__ assert obj.v_i32 == 22 # +2 from custom __init__ # --------------------------------------------------------------------------- # 2. Auto-generated __init__ with defaults # --------------------------------------------------------------------------- def test_c_class_auto_init_defaults() -> None: """Derived classes use auto-generated __init__ with C++ defaults.""" obj = _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0) assert obj.v_i64 == 1 assert obj.v_i32 == 2 assert obj.v_f64 == 3.0 assert obj.v_f32 == 8.0 # default from C++ def test_c_class_auto_init_all_explicit() -> None: """Auto-generated __init__ accepts all fields explicitly.""" obj = _TestCxxClassDerived(v_i64=123, v_i32=456, v_f64=4.0, v_f32=9.0) assert obj.v_i64 == 123 assert obj.v_i32 == 456 assert obj.v_f64 == 4.0 assert obj.v_f32 == 9.0 # --------------------------------------------------------------------------- # 3. Structural equality (__eq__) # --------------------------------------------------------------------------- def test_c_class_eq() -> None: """c_class installs __eq__ using RecursiveEq.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = _TestCxxClassDerived(1, 2, 3.0, 4.0) assert a == b assert a is not b # different objects c = _TestCxxClassDerived(1, 2, 3.0, 5.0) assert a != c def test_c_class_eq_reflexive() -> None: """Equality is reflexive: an object equals itself.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = a # alias, same object assert a == b def test_c_class_eq_symmetric() -> None: """Equality is symmetric: a == b implies b == a.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = _TestCxxClassDerived(1, 2, 3.0, 4.0) assert a == b assert b == a # --------------------------------------------------------------------------- # 4. Structural hash (__hash__) # --------------------------------------------------------------------------- def test_c_class_hash() -> None: """c_class installs __hash__ using RecursiveHash.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = _TestCxxClassDerived(1, 2, 3.0, 4.0) assert hash(a) == hash(b) def test_c_class_hash_as_dict_key() -> None: """Equal objects can be used interchangeably as dict keys.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = _TestCxxClassDerived(1, 2, 3.0, 4.0) d = {a: "value"} assert d[b] == "value" # --------------------------------------------------------------------------- # 5. Ordering (__lt__, __le__, __gt__, __ge__) # --------------------------------------------------------------------------- def test_c_class_ordering() -> None: """c_class installs ordering operators.""" small = _TestCxxClassDerived(0, 0, 0.0, 0.0) big = _TestCxxClassDerived(100, 100, 100.0, 100.0) assert small < big assert small <= big assert big > small assert big >= small assert not (big < small) assert not (small > big) def test_c_class_ordering_reflexive() -> None: """<= and >= are reflexive.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = a # alias, same object assert a <= b assert a >= b def test_c_class_ordering_antisymmetric() -> None: """If a < b then not b < a.""" a = _TestCxxClassDerived(0, 0, 0.0, 0.0) b = _TestCxxClassDerived(100, 100, 100.0, 100.0) if a < b: assert not (b < a) else: assert not (a < b) # --------------------------------------------------------------------------- # 6. Equality with different types returns NotImplemented # --------------------------------------------------------------------------- def test_c_class_eq_different_type() -> None: """__eq__ returns NotImplemented for unrelated types.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) assert a != "hello" assert a != 42 assert a != 3.14 assert a is not None def test_c_class_ordering_different_type() -> None: """Ordering against unrelated types raises TypeError.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) with pytest.raises(TypeError): a < "hello" # ty: ignore[unsupported-operator] with pytest.raises(TypeError): a <= 42 # ty: ignore[unsupported-operator] with pytest.raises(TypeError): a > 3.14 # ty: ignore[unsupported-operator] with pytest.raises(TypeError): a >= None # ty: ignore[unsupported-operator] # --------------------------------------------------------------------------- # 7. Subclass equality # --------------------------------------------------------------------------- def test_c_class_subclass_eq() -> None: """Subclass instances can be compared to parent instances without crashing.""" derived = _TestCxxClassDerived(1, 2, 3.0, 4.0) derived_derived = _TestCxxClassDerivedDerived( v_i64=1, v_i32=2, v_f64=3.0, v_f32=4.0, v_str="hello", v_bool=True ) # These are different types in the same hierarchy; comparison should # return a bool (the result depends on C++ behavior). result = derived == derived_derived assert isinstance(result, bool) # --------------------------------------------------------------------------- # 8. KwOnly from C++ reflection # --------------------------------------------------------------------------- def test_c_class_kw_only_signature() -> None: """kw_only trait comes from C++ reflection, not Python decorator.""" sig = inspect.signature(_TestCxxKwOnly.__init__) params = sig.parameters for name in ("x", "y", "z", "w"): assert params[name].kind == inspect.Parameter.KEYWORD_ONLY, ( f"Expected {name} to be KEYWORD_ONLY" ) def test_c_class_kw_only_call() -> None: """KwOnly fields can be supplied as keyword arguments.""" obj = _TestCxxKwOnly(x=1, y=2, z=3, w=4) assert obj.x == 1 assert obj.y == 2 assert obj.z == 3 assert obj.w == 4 def test_c_class_kw_only_default() -> None: """KwOnly field with a C++ default can be omitted.""" obj = _TestCxxKwOnly(x=1, y=2, z=3) assert obj.w == 100 def test_c_class_kw_only_rejects_positional() -> None: """KwOnly fields reject positional arguments.""" with pytest.raises(TypeError, match="positional"): _TestCxxKwOnly(1, 2, 3, 4) # ty: ignore[missing-argument, too-many-positional-arguments] # --------------------------------------------------------------------------- # 9. Init subset from C++ reflection # --------------------------------------------------------------------------- def test_c_class_init_subset_signature() -> None: """init=False fields from C++ reflection are excluded from __init__.""" sig = inspect.signature(_TestCxxInitSubset.__init__) params = tuple(sig.parameters) assert "required_field" in params assert "optional_field" not in params assert "note" not in params def test_c_class_init_subset_defaults() -> None: """init=False fields get their default values from C++.""" obj = _TestCxxInitSubset(required_field=42) assert obj.required_field == 42 assert obj.optional_field == -1 # C++ default assert obj.note == "default" # C++ default def test_c_class_init_subset_positional() -> None: """Init-subset fields can be passed positionally.""" obj = _TestCxxInitSubset(7) assert obj.required_field == 7 assert obj.optional_field == -1 def test_c_class_init_subset_field_writable() -> None: """Fields excluded from __init__ can still be assigned after construction.""" obj = _TestCxxInitSubset(required_field=0) obj.optional_field = 11 assert obj.optional_field == 11 # --------------------------------------------------------------------------- # 10. DerivedDerived with defaults # --------------------------------------------------------------------------- def test_c_class_derived_derived_defaults() -> None: """DerivedDerived uses positional args; C++ defaults fill in omitted fields.""" obj = _TestCxxClassDerivedDerived(1, 2, 3.0, True) assert obj.v_i64 == 1 assert obj.v_i32 == 2 assert obj.v_f64 == 3.0 assert obj.v_f32 == 8.0 # C++ default assert obj.v_str == "default" # C++ default assert obj.v_bool is True def test_c_class_derived_derived_all_explicit() -> None: """DerivedDerived with all fields explicitly provided.""" obj = _TestCxxClassDerivedDerived( v_i64=123, v_i32=456, v_f64=4.0, v_f32=9.0, v_str="hello", v_bool=True, ) assert obj.v_i64 == 123 assert obj.v_i32 == 456 assert obj.v_f64 == 4.0 assert obj.v_f32 == 9.0 assert obj.v_str == "hello" assert obj.v_bool is True # --------------------------------------------------------------------------- # 11. Hash / set usage # --------------------------------------------------------------------------- def test_c_class_usable_in_set() -> None: """Equal objects deduplicate in a set.""" a = _TestCxxClassDerived(1, 2, 3.0, 4.0) b = _TestCxxClassDerived(1, 2, 3.0, 4.0) c = _TestCxxClassDerived(5, 6, 7.0, 8.0) s = {a, b, c} assert len(s) == 2 # a and b are equal def test_c_class_unequal_objects_in_set() -> None: """Distinct objects are separate entries in a set.""" objs = {_TestCxxClassDerived(i, i, float(i), float(i)) for i in range(5)} assert len(objs) == 5 # --------------------------------------------------------------------------- # 12. Field annotation warnings # --------------------------------------------------------------------------- def test_c_class_warns_on_missing_field_annotations() -> None: """@c_class warns when reflected fields lack Python annotations.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") field_names = {f.name for f in type_info.fields} assert field_names # sanity: there are reflected fields # A class with no annotations should trigger a warning DummyCls = type("DummyCls", (), {}) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") _warn_missing_field_annotations(DummyCls, type_info, stacklevel=2) assert len(w) == 1 assert "does not annotate" in str(w[0].message) for name in field_names: assert name in str(w[0].message) def test_c_class_no_warning_when_all_fields_annotated() -> None: """@c_class does not warn when all reflected fields are annotated.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") _warn_missing_field_annotations(_TestCxxClassBase, type_info, stacklevel=2) assert len(w) == 0 def test_c_class_warns_only_for_missing_annotations() -> None: """Warning lists only the missing fields, not the annotated ones.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") field_names = sorted(f.name for f in type_info.fields) assert len(field_names) >= 2 # need at least 2 fields for this test # Annotate only the first field, leave the rest unannotated PartialCls = type("PartialCls", (), {"__annotations__": {field_names[0]: int}}) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") _warn_missing_field_annotations(PartialCls, type_info, stacklevel=2) assert len(w) == 1 msg = str(w[0].message) # The annotated field should NOT appear in the warning assert field_names[0] not in msg # The unannotated fields should appear for name in field_names[1:]: assert name in msg def test_c_class_warns_only_own_fields_not_inherited() -> None: """Warning only checks own fields, not parent fields.""" # _TestCxxClassDerived's type_info.fields contains only its own fields # (v_f64, v_f32), not parent fields (v_i64, v_i32). derived_type_info: TypeInfo = getattr(_TestCxxClassDerived, "__tvm_ffi_type_info__") own_field_names = {f.name for f in derived_type_info.fields} assert own_field_names # sanity # A class with no annotations: warning should mention only own fields DummyCls = type("DummyCls", (), {}) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") _warn_missing_field_annotations(DummyCls, derived_type_info, stacklevel=2) assert len(w) == 1 msg = str(w[0].message) for name in own_field_names: assert name in msg # Parent fields should NOT appear in the warning parent_type_info = derived_type_info.parent_type_info if parent_type_info is not None: parent_field_names = {f.name for f in parent_type_info.fields} for name in parent_field_names: assert name not in msg # --------------------------------------------------------------------------- # 13. Field object attachment (_attach_field_objects) # --------------------------------------------------------------------------- def test_c_class_attaches_field_object_per_typefield() -> None: """Every own reflected field gets a ``Field`` instance on ``dataclass_field``.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") assert type_info.fields # sanity for tf in type_info.fields: assert isinstance(tf.dataclass_field, Field) def test_c_class_field_name_matches_typefield() -> None: """``Field.name`` mirrors ``TypeField.name``.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.dataclass_field.name == tf.name def test_c_class_field_private_schema_mirrors_typefield() -> None: """``Field._ty_schema`` is forwarded verbatim from ``TypeField.ty``. For C++-backed fields ``TypeField.ty`` is typically ``None`` (only populated by ``@py_class`` registration), so the helper should just forward the value without fabricating a :class:`TypeSchema`. """ type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.dataclass_field._ty_schema is tf.ty def test_c_class_field_type_resolved_from_annotation() -> None: """``Field.type`` is the resolved Python annotation, not a TypeSchema.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in type_info.fields: # _TestCxxClassBase annotates v_i64 / v_i32 as ``int``. assert tf.dataclass_field.type is int def test_c_class_field_defaults_missing_when_unspecified() -> None: """Fields with no C++ default retain ``MISSING`` on ``Field.default``. ``_TestCxxClassBase.v_i64`` / ``v_i32`` are registered without ``refl::default_value(...)``; the reflection layer should leave ``c_default`` / ``c_default_factory`` as :data:`MISSING`. """ type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.c_default is MISSING assert tf.c_default_factory is MISSING assert tf.dataclass_field.default is MISSING assert tf.dataclass_field.default_factory is MISSING def test_c_class_field_default_value_from_cxx() -> None: """C++ ``refl::default_value(...)`` is exposed on ``Field.default``. ``_TestCxxClassDerived.v_f32`` registers ``refl::default_value(8.0f)``; ``_TestCxxClassDerivedDerived.v_str`` registers ``refl::default_value(String("default"))``. Both should round-trip through ``TypeField.c_default`` / ``Field.default``. """ derived_info: TypeInfo = getattr(_TestCxxClassDerived, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in derived_info.fields} # v_f64: no default → MISSING. assert by_name["v_f64"].c_default is MISSING assert by_name["v_f64"].dataclass_field.default is MISSING # v_f32: C++ default 8.0f. assert by_name["v_f32"].c_default == pytest.approx(8.0) assert by_name["v_f32"].dataclass_field.default == pytest.approx(8.0) assert by_name["v_f32"].c_default_factory is MISSING assert by_name["v_f32"].dataclass_field.default_factory is MISSING dd_info: TypeInfo = getattr(_TestCxxClassDerivedDerived, "__tvm_ffi_type_info__") dd_by_name = {tf.name: tf for tf in dd_info.fields} assert dd_by_name["v_str"].c_default == "default" assert dd_by_name["v_str"].dataclass_field.default == "default" # v_bool has no default. assert dd_by_name["v_bool"].c_default is MISSING def test_c_class_field_default_respects_init_false() -> None: """Defaults are visible even when ``init=False`` (fields filled by C++).""" subset_info: TypeInfo = getattr(_TestCxxInitSubset, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in subset_info.fields} # optional_field / note are init(false) + have C++ defaults. assert by_name["optional_field"].c_default == -1 assert by_name["optional_field"].dataclass_field.default == -1 assert by_name["optional_field"].dataclass_field.init is False assert by_name["note"].c_default == "default" assert by_name["note"].dataclass_field.default == "default" # required_field has no default. assert by_name["required_field"].c_default is MISSING def test_c_class_field_default_respects_kw_only() -> None: """Defaults are visible on kw_only fields too (``_TestCxxKwOnly.w``).""" kw_info: TypeInfo = getattr(_TestCxxKwOnly, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in kw_info.fields} # w is kw_only and has C++ default 100. assert by_name["w"].c_default == 100 assert by_name["w"].dataclass_field.default == 100 assert by_name["w"].dataclass_field.kw_only is True # x / y / z are kw_only without defaults. for name in ("x", "y", "z"): assert by_name[name].c_default is MISSING assert by_name[name].dataclass_field.default is MISSING def test_c_class_field_frozen_matches_typefield() -> None: """``Field.frozen`` tracks ``TypeField.frozen`` (TestIntPair fields are frozen).""" type_info: TypeInfo = getattr(tvm_ffi.testing.TestIntPair, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.dataclass_field.frozen == tf.frozen def test_c_class_field_init_and_kw_only_flags() -> None: """``init`` / ``kw_only`` flags propagate from the reflection layer.""" # _TestCxxKwOnly has *all* fields kw-only. kw_info: TypeInfo = getattr(_TestCxxKwOnly, "__tvm_ffi_type_info__") for tf in kw_info.fields: assert tf.dataclass_field.kw_only is True assert tf.dataclass_field.init is True # _TestCxxInitSubset has some fields with Init(false) (not in __init__). subset_info: TypeInfo = getattr(_TestCxxInitSubset, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in subset_info.fields} assert by_name["required_field"].dataclass_field.init is True # optional_field / note are marked Init(false) in C++. assert by_name["optional_field"].dataclass_field.init is False assert by_name["note"].dataclass_field.init is False def test_c_class_field_doc_matches_typefield() -> None: """``Field.doc`` mirrors ``TypeField.doc``.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.dataclass_field.doc == tf.doc def test_c_class_field_repr_flag_from_cxx() -> None: """``refl::repr(false)`` on a C++ field propagates to ``Field.repr``. ``_TestCxxClassBase`` registers both of its fields with ``refl::repr(false)``; ``_TestCxxClassDerived`` registers its fields without that modifier (default on). """ base_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") for tf in base_info.fields: assert tf.c_repr is False assert tf.dataclass_field.repr is False derived_info: TypeInfo = getattr(_TestCxxClassDerived, "__tvm_ffi_type_info__") for tf in derived_info.fields: assert tf.c_repr is True assert tf.dataclass_field.repr is True def test_c_class_field_compare_flag_from_cxx() -> None: """``refl::compare(false)`` on a C++ field propagates to ``Field.compare``. ``TestCompare.ignored_field`` is registered with ``refl::compare(false)``; ``key`` and ``name`` are registered without it. """ type_info: TypeInfo = getattr(TestCompare, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in type_info.fields} assert by_name["key"].c_compare is True assert by_name["key"].dataclass_field.compare is True assert by_name["name"].c_compare is True assert by_name["name"].dataclass_field.compare is True assert by_name["ignored_field"].c_compare is False assert by_name["ignored_field"].dataclass_field.compare is False def test_c_class_field_hash_flag_from_cxx() -> None: """``refl::hash(false)`` on a C++ field propagates to ``Field.hash``. ``TestHash.hash_ignored`` is registered with ``refl::hash(false)``; ``key`` and ``name`` are registered without it. """ type_info: TypeInfo = getattr(TestHash, "__tvm_ffi_type_info__") by_name = {tf.name: tf for tf in type_info.fields} assert by_name["key"].c_hash is True assert by_name["key"].dataclass_field.hash is True assert by_name["name"].c_hash is True assert by_name["name"].dataclass_field.hash is True assert by_name["hash_ignored"].c_hash is False assert by_name["hash_ignored"].dataclass_field.hash is False def test_c_class_field_structural_eq_default_none() -> None: """Fields registered without ``s_eq_hash_def`` / ``s_eq_hash_ignore`` stay ``None``. None of the current C++ testing fixtures set structural_eq flags on a field, so we verify the default (``None``) is preserved across the reflection boundary. """ for cls in (_TestCxxClassBase, _TestCxxClassDerived, TestCompare, TestHash): type_info: TypeInfo = getattr(cls, "__tvm_ffi_type_info__") for tf in type_info.fields: assert tf.c_structural_eq is None assert tf.dataclass_field.structural_eq is None def test_c_class_attaches_only_own_fields_not_inherited() -> None: """Only own fields get a Field; parent fields are attached on the parent's TypeInfo.""" derived_info: TypeInfo = getattr(_TestCxxClassDerived, "__tvm_ffi_type_info__") own_names = {tf.name for tf in derived_info.fields} # Derived owns v_f64 / v_f32; parent owns v_i64 / v_i32. assert own_names == {"v_f64", "v_f32"} for tf in derived_info.fields: assert isinstance(tf.dataclass_field, Field) # Parent TypeInfo already had Field objects attached when it was decorated. parent_info = derived_info.parent_type_info assert parent_info is not None for tf in parent_info.fields: assert isinstance(tf.dataclass_field, Field) def test_c_class_field_type_none_when_annotation_missing() -> None: """Fields without a Python annotation get ``Field.type == None``.""" # Re-run _attach_field_objects on a bare class (no annotations) — it must # not raise and should leave every Field.type == None. type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") original = [tf.dataclass_field for tf in type_info.fields] bare = type("BareCls", (), {}) try: _attach_field_objects(bare, type_info) for tf in type_info.fields: assert tf.dataclass_field.type is None finally: # Restore so later tests still see the annotated Field.type. for tf, saved in zip(type_info.fields, original): tf.dataclass_field = saved def test_c_class_attaches_field_for_new_decoration() -> None: """A freshly decorated @c_class type has its own Field objects.""" # We don't create a new C++ type here — just re-run the decorator machinery # with a duplicate key-aware workflow isn't necessary. Simply assert that # the existing decorated classes all expose populated dataclass_field. for cls in (_TestCxxClassBase, _TestCxxClassDerived, _TestCxxClassDerivedDerived): ti: TypeInfo = getattr(cls, "__tvm_ffi_type_info__") for tf in ti.fields: assert tf.dataclass_field is not None assert tf.dataclass_field.name == tf.name def test_c_class_attach_is_idempotent() -> None: """Calling ``_attach_field_objects`` twice replaces Fields without error.""" type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") original = [tf.dataclass_field for tf in type_info.fields] try: _attach_field_objects(_TestCxxClassBase, type_info) for tf in type_info.fields: assert isinstance(tf.dataclass_field, Field) assert tf.dataclass_field.type is int finally: for tf, saved in zip(type_info.fields, original): tf.dataclass_field = saved def test_c_class_attach_tolerates_unresolvable_hints() -> None: """``typing.get_type_hints`` exceptions fall back silently; ``Field.type`` stays None.""" # Declaring a type with a string annotation that points at an undefined name # makes get_type_hints raise NameError. The helper must swallow it. class BrokenAnn: x: ThisNameIsNeverDefined # ty: ignore[unresolved-reference] # noqa: F821 type_info: TypeInfo = getattr(_TestCxxClassBase, "__tvm_ffi_type_info__") original = [tf.dataclass_field for tf in type_info.fields] try: _attach_field_objects(BrokenAnn, type_info) for tf in type_info.fields: assert tf.dataclass_field.type is None finally: for tf, saved in zip(type_info.fields, original): tf.dataclass_field = saved tvm-ffi-0.1.12/tests/python/test_dataclass_common.py000066400000000000000000000466271521067262500225730ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: D102, UP006, UP045 """Tests for :func:`tvm_ffi.dataclasses.is_dataclass`, :func:`fields`, :func:`replace`.""" from __future__ import annotations import collections import dataclasses as _dc import itertools import typing from typing import List, Optional import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi.core import MISSING, Object from tvm_ffi.dataclasses import ( KW_ONLY, Field, asdict, astuple, field, fields, is_dataclass, py_class, replace, ) _counter = itertools.count() def _k(base: str) -> str: return f"testing.compat.{base}_{next(_counter)}" # --------------------------------------------------------------------------- # is_dataclass # --------------------------------------------------------------------------- class TestIsDataclass: def test_c_class_instance(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) assert is_dataclass(p) is True def test_c_class_type(self) -> None: assert is_dataclass(tvm_ffi.testing.TestIntPair) is True def test_py_class_instance(self) -> None: @py_class(_k("PC")) class PC(Object): x: int assert is_dataclass(PC(x=1)) is True def test_py_class_type(self) -> None: @py_class(_k("PCT")) class PCT(Object): x: int assert is_dataclass(PCT) is True def test_non_dataclass_types(self) -> None: assert is_dataclass(42) is False assert is_dataclass("hi") is False assert is_dataclass(object()) is False assert is_dataclass(int) is False def test_plain_object_subclass(self) -> None: """Object subclass without @py_class / @c_class shouldn't qualify.""" class Plain(Object): pass assert is_dataclass(Plain) is False def test_stdlib_dataclass_not_recognised(self) -> None: """``common.is_dataclass`` only recognises FFI types.""" @_dc.dataclass class D: x: int assert is_dataclass(D) is False assert is_dataclass(D(1)) is False # --------------------------------------------------------------------------- # fields # --------------------------------------------------------------------------- class TestFields: def test_returns_tuple_of_field(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) result = fields(p) assert isinstance(result, tuple) assert all(isinstance(f, Field) for f in result) def test_c_class_basic(self) -> None: fs = fields(tvm_ffi.testing.TestIntPair) assert [f.name for f in fs] == ["a", "b"] assert all(f.type is int for f in fs) # TestIntPair registers no defaults; the Field descriptors reflect that. assert all(f.default is MISSING for f in fs) def test_c_class_default_value_from_cxx(self) -> None: """``fields()`` exposes ``refl::default_value(...)`` from C++. ``_TestCxxClassDerived.v_f32`` has ``refl::default_value(8.0f)``. """ fs = fields(tvm_ffi.testing._TestCxxClassDerived) by_name = {f.name: f for f in fs} assert by_name["v_f32"].default == pytest.approx(8.0) assert by_name["v_f32"].default_factory is MISSING # v_f64 has no C++ default. assert by_name["v_f64"].default is MISSING def test_c_class_instance_and_type_equivalent(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) assert fields(p) == fields(tvm_ffi.testing.TestIntPair) def test_py_class_with_default(self) -> None: @py_class(_k("PCDef")) class PCDef(Object): x: int y: str = field(default="hi") fs = fields(PCDef) names = [f.name for f in fs] assert names == ["x", "y"] (fx, fy) = fs assert fx.default is MISSING assert fx.default_factory is MISSING assert fy.default == "hi" assert fy.default_factory is MISSING def test_py_class_with_default_factory(self) -> None: @py_class(_k("PCFact")) class PCFact(Object): items: List[int] = field(default_factory=list) (f,) = fields(PCFact) assert f.default is MISSING assert f.default_factory is list def test_py_class_optional_type_resolved(self) -> None: @py_class(_k("PCOpt")) class PCOpt(Object): v: Optional[int] = field(default=None) (f,) = fields(PCOpt) # Depending on the Python version, get_type_hints may preserve # Optional[int] or normalize it to Union[int, None]; inspect the # type's args instead of relying on str() formatting. args = typing.get_args(f.type) assert int in args and type(None) in args assert f.default is None def test_inheritance_parent_first(self) -> None: @py_class(_k("PCP")) class P(Object): x: int y: str = field(default="p") @py_class(_k("PCC")) class C(P): z: float = field(default=1.5) assert [f.name for f in fields(C)] == ["x", "y", "z"] assert fields(C)[2].default == 1.5 def test_c_class_inheritance(self) -> None: """@c_class derived types also walk the parent chain.""" fs = fields(tvm_ffi.testing._TestCxxClassDerived) # Parent v_i64/v_i32 come before child v_f64/v_f32. assert [f.name for f in fs] == ["v_i64", "v_i32", "v_f64", "v_f32"] def test_non_dataclass_raises(self) -> None: with pytest.raises(TypeError, match="c_class or py_class"): fields(42) def test_field_attribute_access(self) -> None: """Returned Field descriptors expose ``name`` / ``type`` / ``default``.""" (fa, _) = fields(tvm_ffi.testing.TestIntPair) assert fa.name == "a" assert fa.type is int assert fa.default is MISSING # --------------------------------------------------------------------------- # replace # --------------------------------------------------------------------------- class TestReplace: def test_c_class_single_field(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) p2 = replace(p, a=10) assert (p2.a, p2.b) == (10, 2) assert (p.a, p.b) == (1, 2) # original unchanged assert not p.same_as(p2) def test_c_class_multiple_fields(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) p2 = replace(p, a=7, b=8) assert (p2.a, p2.b) == (7, 8) def test_c_class_no_changes_is_copy(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) p2 = replace(p) assert (p2.a, p2.b) == (1, 2) assert not p.same_as(p2) def test_py_class_basic(self) -> None: @py_class(_k("PCR")) class PCR(Object): x: int = field(default=5) y: str = field(default="hi") obj = PCR() obj2 = replace(obj, x=99) assert obj2.x == 99 assert obj2.y == "hi" assert obj.x == 5 # original unchanged def test_py_class_inherited_fields(self) -> None: @py_class(_k("PCRParent")) class P(Object): x: int @py_class(_k("PCRChild")) class C(P): y: str obj = C(x=1, y="a") obj2 = replace(obj, x=99) assert obj2.x == 99 assert obj2.y == "a" def test_replace_frozen_field(self) -> None: """replace() goes through __replace__, which uses the frozen-bypass setter.""" # TestIntPair.a and .b are read-only (frozen c_class fields). p = tvm_ffi.testing.TestIntPair(3, 4) p2 = replace(p, a=10) assert p2.a == 10 assert p.a == 3 # still read-only, original untouched # --------------------------------------------------------------------------- # asdict # --------------------------------------------------------------------------- class TestAsdict: def test_c_class_basic(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) assert asdict(p) == {"a": 1, "b": 2} def test_py_class_basic(self) -> None: @py_class(_k("PCAsdict")) class PC(Object): x: int y: str assert asdict(PC(x=1, y="hi")) == {"x": 1, "y": "hi"} def test_py_class_nested(self) -> None: @py_class(_k("PCInner")) class Inner(Object): x: int @py_class(_k("PCOuter")) class Outer(Object): a: Inner b: Inner out = Outer(a=Inner(x=1), b=Inner(x=2)) assert asdict(out) == {"a": {"x": 1}, "b": {"x": 2}} def test_py_class_inheritance(self) -> None: @py_class(_k("PCParent")) class P(Object): x: int @py_class(_k("PCChild")) class C(P): y: str assert asdict(C(x=1, y="a")) == {"x": 1, "y": "a"} def test_ffi_array_recurses_to_list(self) -> None: @py_class(_k("PCWithArray")) class PC(Object): xs: tvm_ffi.Array pc = PC(xs=tvm_ffi.Array([1, 2, 3])) result = asdict(pc) assert result == {"xs": [1, 2, 3]} assert type(result["xs"]) is list def test_ffi_list_recurses_to_list(self) -> None: @py_class(_k("PCWithList")) class PC(Object): xs: tvm_ffi.List pc = PC(xs=tvm_ffi.List([1, 2, 3])) result = asdict(pc) assert result == {"xs": [1, 2, 3]} assert type(result["xs"]) is list def test_ffi_map_recurses_to_dict(self) -> None: @py_class(_k("PCWithMap")) class PC(Object): m: tvm_ffi.Map pc = PC(m=tvm_ffi.Map({"a": 1, "b": 2})) result = asdict(pc) assert result == {"m": {"a": 1, "b": 2}} assert type(result["m"]) is dict def test_ffi_dict_recurses_to_dict(self) -> None: @py_class(_k("PCWithDict")) class PC(Object): d: tvm_ffi.Dict pc = PC(d=tvm_ffi.Dict({"a": 1, "b": 2})) result = asdict(pc) assert result == {"d": {"a": 1, "b": 2}} assert type(result["d"]) is dict def test_ffi_array_of_dataclasses(self) -> None: @py_class(_k("PCItem")) class Item(Object): v: int @py_class(_k("PCBox")) class Box(Object): items: tvm_ffi.Array box = Box(items=tvm_ffi.Array([Item(v=1), Item(v=2)])) assert asdict(box) == {"items": [{"v": 1}, {"v": 2}]} def test_ffi_map_of_dataclasses(self) -> None: @py_class(_k("PCItem2")) class Item(Object): v: int @py_class(_k("PCBox2")) class Box(Object): items: tvm_ffi.Map box = Box(items=tvm_ffi.Map({"a": Item(v=1), "b": Item(v=2)})) assert asdict(box) == {"items": {"a": {"v": 1}, "b": {"v": 2}}} def test_dict_factory(self) -> None: @py_class(_k("PCDF")) class PC(Object): x: int y: int result = asdict(PC(x=1, y=2), dict_factory=collections.OrderedDict) assert isinstance(result, collections.OrderedDict) assert list(result.items()) == [("x", 1), ("y", 2)] def test_dict_factory_recurses(self) -> None: @py_class(_k("PCDFI")) class Inner(Object): v: int @py_class(_k("PCDFO")) class Outer(Object): a: Inner result = asdict(Outer(a=Inner(v=5)), dict_factory=collections.OrderedDict) assert isinstance(result, collections.OrderedDict) assert isinstance(result["a"], collections.OrderedDict) def test_default_factory_list_independent(self) -> None: """Result must be a fresh ``list``, not aliased to any internal state.""" @py_class(_k("PCInd")) class PC(Object): xs: tvm_ffi.Array pc = PC(xs=tvm_ffi.Array([1, 2])) d = asdict(pc) d["xs"].append(99) # Mutating the result must not affect the original. assert list(pc.xs) == [1, 2] def test_type_raises(self) -> None: with pytest.raises(TypeError, match="c_class / py_class instances"): asdict(tvm_ffi.testing.TestIntPair) # passing type, not instance def test_non_dataclass_raises(self) -> None: with pytest.raises(TypeError, match="c_class / py_class instances"): asdict(42) with pytest.raises(TypeError, match="c_class / py_class instances"): asdict([1, 2, 3]) def test_defaultdict_preserved(self) -> None: """``defaultdict`` round-trips with its ``default_factory`` intact. Exercises the ``_asdict_inner`` defaultdict branch directly, since FFI ``Any`` field storage converts a stored ``defaultdict`` into an FFI ``Map`` on readback. Mirrors stdlib ``dataclasses._asdict_inner``'s check on ``type(obj)``. """ from tvm_ffi.dataclasses.common import _asdict_inner # noqa: PLC0415 dd = collections.defaultdict(list) dd["a"].append(1) dd["b"].append(2) result = _asdict_inner(dd, dict) assert type(result) is collections.defaultdict assert result.default_factory is list assert dict(result) == {"a": [1], "b": [2]} # --------------------------------------------------------------------------- # astuple # --------------------------------------------------------------------------- class TestAstuple: def test_c_class_basic(self) -> None: p = tvm_ffi.testing.TestIntPair(1, 2) assert astuple(p) == (1, 2) def test_py_class_basic(self) -> None: @py_class(_k("PCAstuple")) class PC(Object): x: int y: str assert astuple(PC(x=1, y="hi")) == (1, "hi") def test_py_class_nested(self) -> None: @py_class(_k("PCInnerT")) class Inner(Object): x: int @py_class(_k("PCOuterT")) class Outer(Object): a: Inner b: Inner out = Outer(a=Inner(x=1), b=Inner(x=2)) assert astuple(out) == ((1,), (2,)) def test_py_class_inheritance(self) -> None: @py_class(_k("PCParentT")) class P(Object): x: int @py_class(_k("PCChildT")) class C(P): y: str assert astuple(C(x=1, y="a")) == (1, "a") def test_ffi_array_recurses_to_list(self) -> None: @py_class(_k("PCArrT")) class PC(Object): xs: tvm_ffi.Array pc = PC(xs=tvm_ffi.Array([1, 2, 3])) result = astuple(pc) assert result == ([1, 2, 3],) assert type(result[0]) is list def test_ffi_map_recurses_to_dict(self) -> None: @py_class(_k("PCMapT")) class PC(Object): m: tvm_ffi.Map pc = PC(m=tvm_ffi.Map({"a": 1})) result = astuple(pc) assert result == ({"a": 1},) assert type(result[0]) is dict def test_tuple_factory(self) -> None: @py_class(_k("PCTF")) class PC(Object): x: int y: int assert astuple(PC(x=1, y=2), tuple_factory=list) == [1, 2] def test_tuple_factory_recurses(self) -> None: @py_class(_k("PCTFI")) class Inner(Object): v: int @py_class(_k("PCTFO")) class Outer(Object): a: Inner result = astuple(Outer(a=Inner(v=5)), tuple_factory=list) assert result == [[5]] def test_type_raises(self) -> None: with pytest.raises(TypeError, match="c_class / py_class instances"): astuple(tvm_ffi.testing.TestIntPair) def test_non_dataclass_raises(self) -> None: with pytest.raises(TypeError, match="c_class / py_class instances"): astuple(42) with pytest.raises(TypeError, match="c_class / py_class instances"): astuple([1, 2, 3]) def test_defaultdict_preserved(self) -> None: """``defaultdict`` round-trips with its ``default_factory`` intact.""" from tvm_ffi.dataclasses.common import _astuple_inner # noqa: PLC0415 dd = collections.defaultdict(list) dd["a"].append(1) dd["b"].append(2) result = _astuple_inner(dd, tuple) assert type(result) is collections.defaultdict assert result.default_factory is list assert dict(result) == {"a": [1], "b": [2]} # --------------------------------------------------------------------------- # __match_args__ # --------------------------------------------------------------------------- def _match_args(cls: type) -> object: """Read ``__match_args__`` without tripping static attribute checks.""" return getattr(cls, "__match_args__") class TestMatchArgs: def test_py_class_basic(self) -> None: @py_class(_k("MAPy")) class PC(Object): x: int y: str assert _match_args(PC) == ("x", "y") def test_py_class_init_false_excluded(self) -> None: @py_class(_k("MAInitFalse")) class PC(Object): x: int y: int = field(default=0, init=False) assert _match_args(PC) == ("x",) def test_py_class_kw_only_field_excluded(self) -> None: @py_class(_k("MAKwField")) class PC(Object): x: int y: int = field(kw_only=True) assert _match_args(PC) == ("x",) def test_py_class_kw_only_decorator_excludes_all(self) -> None: @py_class(_k("MAKwDeco"), kw_only=True) class PC(Object): x: int y: int assert _match_args(PC) == () def test_py_class_kw_only_sentinel(self) -> None: @py_class(_k("MAKwSent")) class PC(Object): x: int _: KW_ONLY y: int z: int assert _match_args(PC) == ("x",) def test_py_class_inheritance_parent_first(self) -> None: @py_class(_k("MAParent")) class P(Object): x: int y: str @py_class(_k("MAChild")) class C(P): z: float assert _match_args(P) == ("x", "y") assert _match_args(C) == ("x", "y", "z") def test_py_class_opt_out(self) -> None: @py_class(_k("MAOptOut"), match_args=False) class PC(Object): x: int y: int assert "__match_args__" not in PC.__dict__ def test_py_class_user_defined_preserved(self) -> None: @py_class(_k("MAUser")) class PC(Object): x: int y: int __match_args__ = ("y",) assert _match_args(PC) == ("y",) def test_c_class_basic(self) -> None: assert _match_args(tvm_ffi.testing.TestIntPair) == ("a", "b") def test_c_class_inheritance_parent_first(self) -> None: assert _match_args(tvm_ffi.testing._TestCxxClassDerived) == ( "v_i64", "v_i32", "v_f64", "v_f32", ) def test_tuple_type(self) -> None: @py_class(_k("MATupT")) class PC(Object): x: int assert isinstance(_match_args(PC), tuple) tvm-ffi-0.1.12/tests/python/test_dataclass_compare.py000066400000000000000000001161441521067262500227210ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for ffi.RecursiveEq / RecursiveLt / RecursiveLe / RecursiveGt / RecursiveGe.""" from __future__ import annotations import math import struct import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi._ffi_api import RecursiveEq, RecursiveGe, RecursiveGt, RecursiveLe, RecursiveLt from tvm_ffi.testing import ( TestCompare, TestCustomCompare, TestEqWithoutHash, TestIntPair, _TestCxxClassDerived, _TestCxxClassDerivedDerived, create_object, ) def _make_nan_from_payload(payload: int) -> float: """Create a quiet NaN with a deterministic payload.""" bits = 0x7FF8000000000000 | (payload & 0x0007FFFFFFFFFFFF) return struct.unpack(">d", struct.pack(">Q", bits))[0] # --------------------------------------------------------------------------- # Primitives: int # --------------------------------------------------------------------------- def test_int_eq() -> None: assert RecursiveEq(42, 42) assert not RecursiveEq(1, 2) def test_int_ordering() -> None: assert RecursiveLt(1, 2) assert not RecursiveLt(2, 1) assert not RecursiveLt(1, 1) assert RecursiveLe(1, 2) assert RecursiveLe(1, 1) assert not RecursiveLe(2, 1) assert RecursiveGt(2, 1) assert not RecursiveGt(1, 2) assert RecursiveGe(2, 1) assert RecursiveGe(2, 2) def test_int64_extremes_eq() -> None: """Extreme int64 values compare equal to themselves.""" i64_min = -(2**63) i64_max = 2**63 - 1 assert RecursiveEq(i64_max, i64_max) assert RecursiveEq(i64_min, i64_min) def test_int64_extremes_ordering() -> None: """Exercise int64 boundary ordering where naive subtraction would overflow.""" i64_min = -(2**63) i64_max = 2**63 - 1 assert RecursiveLt(i64_min, i64_max) assert RecursiveGt(i64_max, i64_min) assert not RecursiveGt(i64_min, i64_max) assert not RecursiveLe(i64_max, i64_min) # Cases that would give wrong results with naive Sign(a - b) assert RecursiveLt(i64_min, 1) assert RecursiveGt(i64_max, -1) assert not RecursiveLt(i64_max, -1) # --------------------------------------------------------------------------- # Primitives: float # --------------------------------------------------------------------------- def test_float_eq() -> None: assert RecursiveEq(3.14, 3.14) assert not RecursiveEq(1.0, 2.0) def test_float_ordering() -> None: assert RecursiveLt(1.0, 2.0) assert not RecursiveLt(2.0, 1.0) assert RecursiveGe(2.0, 2.0) def test_float_signed_zero() -> None: assert RecursiveEq(-0.0, 0.0) assert not RecursiveLt(-0.0, 0.0) assert RecursiveLe(-0.0, 0.0) assert RecursiveGe(-0.0, 0.0) def test_float_infinity_ordering() -> None: assert RecursiveLt(-math.inf, 0.0) assert RecursiveLt(0.0, math.inf) assert RecursiveLt(-math.inf, math.inf) assert RecursiveGt(math.inf, -math.inf) assert RecursiveEq(math.inf, math.inf) assert RecursiveLt(1.0, math.inf) # --------------------------------------------------------------------------- # NaN handling # --------------------------------------------------------------------------- def test_nan_eq() -> None: """NaN == NaN under RecursiveEq (equality-only mode).""" assert RecursiveEq(math.nan, math.nan) def test_nan_ordering_raises() -> None: """Ordering NaN values raises TypeError.""" with pytest.raises(TypeError): RecursiveLt(math.nan, 1.0) with pytest.raises(TypeError): RecursiveLt(1.0, math.nan) with pytest.raises(TypeError): RecursiveLe(math.nan, math.nan) def test_nan_payloads_eq() -> None: nan1 = _make_nan_from_payload(0x1) nan2 = _make_nan_from_payload(0x2) assert math.isnan(nan1) and math.isnan(nan2) assert RecursiveEq(nan1, nan2) def test_nan_payloads_in_nested_array() -> None: nan1 = _make_nan_from_payload(0xA5) nan2 = _make_nan_from_payload(0x5A) a = tvm_ffi.Array([1.0, nan1, 2.0]) b = tvm_ffi.Array([1.0, nan2, 2.0]) assert RecursiveEq(a, b) with pytest.raises(TypeError): RecursiveLt(a, b) # --------------------------------------------------------------------------- # Primitives: bool # --------------------------------------------------------------------------- def test_bool_eq() -> None: assert RecursiveEq(True, True) assert RecursiveEq(False, False) assert not RecursiveEq(True, False) def test_bool_ordering() -> None: assert RecursiveLt(False, True) assert not RecursiveLt(True, False) # --------------------------------------------------------------------------- # Primitives: string # --------------------------------------------------------------------------- def test_string_eq() -> None: assert RecursiveEq("hello", "hello") assert not RecursiveEq("hello", "world") def test_string_ordering() -> None: assert RecursiveLt("abc", "abd") assert RecursiveLt("abc", "abcd") assert not RecursiveLt("abd", "abc") def test_string_small_boundary_len7_len8() -> None: small = "1234567" # SmallStr large = "12345678" # heap-backed Str assert RecursiveEq(small, "1234567") assert RecursiveLt(small, large) def test_string_embedded_nul() -> None: assert RecursiveEq("a\x00b", "a\x00b") assert RecursiveLt("a\x00b", "a\x00c") # --------------------------------------------------------------------------- # Primitives: bytes # --------------------------------------------------------------------------- def test_bytes_eq() -> None: assert RecursiveEq(b"hello", b"hello") assert not RecursiveEq(b"hello", b"world") def test_bytes_ordering() -> None: assert RecursiveLt(b"abc", b"abd") assert RecursiveLt(b"abc", b"abcd") def test_bytes_small_boundary_len7_len8() -> None: small = b"1234567" # SmallBytes large = b"12345678" # heap-backed Bytes assert RecursiveEq(small, b"1234567") assert RecursiveLt(small, large) def test_bytes_embedded_nul() -> None: assert RecursiveEq(b"a\x00b", b"a\x00b") assert RecursiveLt(b"a\x00b", b"a\x00c") def test_bytes_high_bit() -> None: """Document high-bit byte ordering behavior. Bytes::memncmp uses char comparison; whether char is signed or unsigned is platform-dependent. On platforms where char is signed (most x86/arm64 compilers), 0xff (-1) sorts before 0x00 (0). This is a known pre-existing issue in Bytes::memncmp (see string.h), not specific to recursive_compare. """ # Just verify the two values are distinguishable (not equal) assert not RecursiveEq(b"\x00", b"\xff") # --------------------------------------------------------------------------- # None # --------------------------------------------------------------------------- def test_none_eq() -> None: assert RecursiveEq(None, None) assert not RecursiveEq(None, 42) assert not RecursiveEq(42, None) def test_none_ordering() -> None: """None is less than any non-None value.""" assert RecursiveLt(None, 42) assert RecursiveLt(None, "s") assert not RecursiveGt(None, 42) # --------------------------------------------------------------------------- # Type mismatch # --------------------------------------------------------------------------- def test_type_mismatch_eq() -> None: """RecursiveEq returns False for different types (no throw).""" assert not RecursiveEq(42, "hello") assert not RecursiveEq(1, True) assert not RecursiveEq(1.0, 1) def test_type_mismatch_ordering_raises() -> None: """Ordering different types raises TypeError.""" with pytest.raises(TypeError): RecursiveLt(42, "hello") with pytest.raises(TypeError): RecursiveGt(1.0, 1) # --------------------------------------------------------------------------- # DataType # --------------------------------------------------------------------------- def test_dtype_eq() -> None: assert RecursiveEq(tvm_ffi.dtype("float32"), tvm_ffi.dtype("float32")) assert not RecursiveEq(tvm_ffi.dtype("float32"), tvm_ffi.dtype("float16")) def test_dtype_ordering() -> None: # float32 code=2 bits=32, float16 code=2 bits=16 -> float16 < float32 assert RecursiveLt(tvm_ffi.dtype("float16"), tvm_ffi.dtype("float32")) assert not RecursiveLt(tvm_ffi.dtype("float32"), tvm_ffi.dtype("float16")) # --------------------------------------------------------------------------- # Device # --------------------------------------------------------------------------- def test_device_eq() -> None: assert RecursiveEq(tvm_ffi.Device("cpu", 0), tvm_ffi.Device("cpu", 0)) assert not RecursiveEq(tvm_ffi.Device("cpu", 0), tvm_ffi.Device("cpu", 1)) def test_device_ordering() -> None: assert RecursiveLt(tvm_ffi.Device("cpu", 0), tvm_ffi.Device("cpu", 1)) assert RecursiveLt(tvm_ffi.Device("cpu", 0), tvm_ffi.Device("cuda", 0)) # --------------------------------------------------------------------------- # Containers: Array # --------------------------------------------------------------------------- def test_array_eq() -> None: a = tvm_ffi.Array([1, 2, 3]) b = tvm_ffi.Array([1, 2, 3]) c = tvm_ffi.Array([1, 2, 4]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_array_ordering() -> None: a = tvm_ffi.Array([1, 2, 3]) b = tvm_ffi.Array([1, 2, 4]) c = tvm_ffi.Array([1, 2, 3, 0]) assert RecursiveLt(a, b) assert RecursiveLt(a, c) # shorter < longer when prefix matches assert not RecursiveLt(b, a) def test_array_empty() -> None: empty = tvm_ffi.Array([]) one = tvm_ffi.Array([1]) assert RecursiveEq(empty, empty) assert RecursiveLt(empty, one) # --------------------------------------------------------------------------- # Containers: List # --------------------------------------------------------------------------- def test_list_eq() -> None: a = tvm_ffi.List([1, 2, 3]) b = tvm_ffi.List([1, 2, 3]) c = tvm_ffi.List([1, 2, 4]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_list_ordering() -> None: a = tvm_ffi.List([1, 2]) b = tvm_ffi.List([1, 3]) assert RecursiveLt(a, b) # --------------------------------------------------------------------------- # Shape # --------------------------------------------------------------------------- def test_shape_eq() -> None: a = tvm_ffi.Shape((2, 3, 4)) b = tvm_ffi.Shape((2, 3, 4)) c = tvm_ffi.Shape((2, 3, 5)) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_shape_ordering() -> None: a = tvm_ffi.Shape((2, 3, 4)) b = tvm_ffi.Shape((2, 3, 5)) c = tvm_ffi.Shape((2, 3, 4, 0)) assert RecursiveLt(a, b) assert RecursiveLt(a, c) # --------------------------------------------------------------------------- # Map/Dict equality # --------------------------------------------------------------------------- def test_map_eq() -> None: a = tvm_ffi.Map({"x": 1, "y": 2}) b = tvm_ffi.Map({"x": 1, "y": 2}) c = tvm_ffi.Map({"x": 1, "y": 3}) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_map_different_size() -> None: a = tvm_ffi.Map({"x": 1}) b = tvm_ffi.Map({"x": 1, "y": 2}) assert not RecursiveEq(a, b) def test_map_ordering_raises() -> None: a = tvm_ffi.Map({"x": 1}) b = tvm_ffi.Map({"x": 2}) with pytest.raises(TypeError): RecursiveLt(a, b) def test_map_same_size_different_keys() -> None: a = tvm_ffi.Map({"x": 1}) b = tvm_ffi.Map({"y": 1}) assert not RecursiveEq(a, b) with pytest.raises(TypeError): RecursiveLt(a, b) def test_equal_maps_under_ordering() -> None: """Two separate but equal maps pass Le/Ge without raising TypeError.""" a = tvm_ffi.Map({"x": 1, "y": 2}) b = tvm_ffi.Map({"x": 1, "y": 2}) assert RecursiveLe(a, b) assert RecursiveGe(a, b) assert not RecursiveLt(a, b) assert not RecursiveGt(a, b) def test_dict_eq() -> None: a = tvm_ffi.Dict({"x": 1, "y": 2}) b = tvm_ffi.Dict({"x": 1, "y": 2}) assert RecursiveEq(a, b) def test_dict_ordering_raises() -> None: a = tvm_ffi.Dict({"x": 1}) b = tvm_ffi.Dict({"x": 2}) with pytest.raises(TypeError): RecursiveLt(a, b) def test_equal_dicts_under_ordering() -> None: """Two separate but equal dicts pass Le/Ge without raising TypeError.""" a = tvm_ffi.Dict({"x": 1, "y": 2}) b = tvm_ffi.Dict({"x": 1, "y": 2}) assert RecursiveLe(a, b) assert RecursiveGe(a, b) # --------------------------------------------------------------------------- # Reflected objects: TestIntPair # --------------------------------------------------------------------------- def test_reflected_obj_eq() -> None: a = TestIntPair(1, 2) b = TestIntPair(1, 2) c = TestIntPair(1, 3) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_reflected_obj_ordering() -> None: a = TestIntPair(1, 2) b = TestIntPair(1, 3) c = TestIntPair(2, 0) assert RecursiveLt(a, b) # first field equal, second: 2 < 3 assert RecursiveLt(a, c) # first field: 1 < 2 # --------------------------------------------------------------------------- # CompareOff flag: TestCompare # --------------------------------------------------------------------------- def test_compare_off_ignored_field() -> None: """ignored_field is excluded from comparison via Compare(false).""" a = TestCompare(1, "x", 100) b = TestCompare(1, "x", 999) assert RecursiveEq(a, b) def test_compare_off_key_differs() -> None: a = TestCompare(1, "x", 100) b = TestCompare(2, "x", 100) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) def test_compare_off_name_differs() -> None: a = TestCompare(1, "a", 100) b = TestCompare(1, "b", 100) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) # --------------------------------------------------------------------------- # Same pointer fast path # --------------------------------------------------------------------------- def test_same_pointer() -> None: x = TestIntPair(42, 99) assert RecursiveEq(x, x) # --------------------------------------------------------------------------- # Different object types # --------------------------------------------------------------------------- def test_different_obj_types_eq() -> None: """RecursiveEq returns False for different object types.""" a = TestIntPair(1, 2) b = TestCompare(1, "x", 0) assert not RecursiveEq(a, b) def test_different_obj_types_ordering_raises() -> None: a = TestIntPair(1, 2) b = TestCompare(1, "x", 0) with pytest.raises(TypeError): RecursiveLt(a, b) # --------------------------------------------------------------------------- # Nested objects # --------------------------------------------------------------------------- def test_nested_objects_in_array() -> None: a1 = TestIntPair(1, 2) a2 = TestIntPair(3, 4) b1 = TestIntPair(1, 2) b2 = TestIntPair(3, 4) arr_a = tvm_ffi.Array([a1, a2]) arr_b = tvm_ffi.Array([b1, b2]) assert RecursiveEq(arr_a, arr_b) def test_nested_objects_in_array_differ() -> None: a1 = TestIntPair(1, 2) a2 = TestIntPair(3, 4) b1 = TestIntPair(1, 2) b2 = TestIntPair(3, 5) arr_a = tvm_ffi.Array([a1, a2]) arr_b = tvm_ffi.Array([b1, b2]) assert not RecursiveEq(arr_a, arr_b) assert RecursiveLt(arr_a, arr_b) # --------------------------------------------------------------------------- # Le / Ge / Gt derived operators # --------------------------------------------------------------------------- def test_le_ge_gt() -> None: assert RecursiveLe(1, 2) assert RecursiveLe(2, 2) assert not RecursiveLe(3, 2) assert RecursiveGe(2, 1) assert RecursiveGe(2, 2) assert not RecursiveGe(1, 2) assert RecursiveGt(2, 1) assert not RecursiveGt(1, 1) assert not RecursiveGt(0, 1) # --------------------------------------------------------------------------- # Nested containers # --------------------------------------------------------------------------- def test_array_of_arrays_eq() -> None: a = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 4])]) b = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 4])]) c = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 5])]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_array_of_arrays_ordering() -> None: a = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 4])]) b = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 5])]) assert RecursiveLt(a, b) # differ at depth-2: 4 < 5 assert not RecursiveLt(b, a) def test_list_of_lists_eq() -> None: a = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([3])]) b = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([3])]) c = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([4])]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_array_of_shapes() -> None: a = tvm_ffi.Array([tvm_ffi.Shape((1, 2)), tvm_ffi.Shape((3, 4))]) b = tvm_ffi.Array([tvm_ffi.Shape((1, 2)), tvm_ffi.Shape((3, 4))]) c = tvm_ffi.Array([tvm_ffi.Shape((1, 2)), tvm_ffi.Shape((3, 5))]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) assert RecursiveLt(a, c) def test_array_of_arrays_different_inner_lengths() -> None: a = tvm_ffi.Array([tvm_ffi.Array([1, 2])]) b = tvm_ffi.Array([tvm_ffi.Array([1, 2, 3])]) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) # inner [1,2] < [1,2,3] def test_three_level_nested_containers() -> None: a = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([1])])]) b = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([1])])]) c = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([2])])]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) assert RecursiveLt(a, c) # --------------------------------------------------------------------------- # Objects with container fields (TestObjectDerived) # --------------------------------------------------------------------------- def test_object_with_array_field_eq() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20, 30]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20, 30]), ) assert RecursiveEq(a, b) def test_object_with_array_field_differ() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 21]), ) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) def test_object_with_map_field_eq() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"a": 1, "b": 2}), v_array=tvm_ffi.Array([]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"a": 1, "b": 2}), v_array=tvm_ffi.Array([]), ) assert RecursiveEq(a, b) def test_object_with_map_field_differ() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"a": 1}), v_array=tvm_ffi.Array([]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"a": 2}), v_array=tvm_ffi.Array([]), ) assert not RecursiveEq(a, b) def test_object_primitive_field_differ_short_circuits() -> None: """First field (v_i64) differs; container fields are the same.""" a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"k": 1}), v_array=tvm_ffi.Array([1]), ) b = create_object( "testing.TestObjectDerived", v_i64=2, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({"k": 1}), v_array=tvm_ffi.Array([1]), ) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) # ordering uses v_i64: 1 < 2 # --------------------------------------------------------------------------- # Objects nested inside object fields # --------------------------------------------------------------------------- def test_object_array_of_objects() -> None: a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestIntPair(1, 2), TestIntPair(3, 4), ] ), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestIntPair(1, 2), TestIntPair(3, 4), ] ), ) assert RecursiveEq(a, b) def test_object_array_of_objects_differ() -> None: a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestIntPair(1, 2), TestIntPair(3, 4), ] ), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestIntPair(1, 2), TestIntPair(3, 5), ] ), ) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) def test_object_map_with_object_values() -> None: a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map( { "x": TestIntPair(1, 2), } ), v_array=tvm_ffi.Array([]), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map( { "x": TestIntPair(1, 2), } ), v_array=tvm_ffi.Array([]), ) assert RecursiveEq(a, b) def test_deep_object_in_object() -> None: """TestObjectDerived.v_array contains another TestObjectDerived -> 3-level nesting.""" inner_a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=1.0, v_str="inner", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([42]), ) a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="outer", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([inner_a]), ) inner_b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=1.0, v_str="inner", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([42]), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="outer", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([inner_b]), ) assert RecursiveEq(a, b) # --------------------------------------------------------------------------- # Inherited field comparison # --------------------------------------------------------------------------- def test_inherited_fields_eq() -> None: a = _TestCxxClassDerived(10, 20, 1.5, 2.5) b = _TestCxxClassDerived(10, 20, 1.5, 2.5) assert RecursiveEq(a, b) def test_inherited_fields_differ_in_base() -> None: a = _TestCxxClassDerived(10, 20, 1.5, 2.5) b = _TestCxxClassDerived(99, 20, 1.5, 2.5) assert not RecursiveEq(a, b) assert RecursiveLt(a, b) def test_three_level_inheritance_eq_and_differ() -> None: # Positional order: required (v_i64, v_i32, v_f64, v_bool), then optional (v_f32, v_str) a = _TestCxxClassDerivedDerived(1, 2, 3.0, True, 4.0, "hi") b = _TestCxxClassDerivedDerived(1, 2, 3.0, True, 4.0, "hi") assert RecursiveEq(a, b) c = _TestCxxClassDerivedDerived(1, 2, 3.0, False, 4.0, "hi") assert not RecursiveEq(a, c) # --------------------------------------------------------------------------- # CompareOff with nesting # --------------------------------------------------------------------------- def test_compare_off_inside_array() -> None: a = tvm_ffi.Array( [ TestCompare(1, "x", 100), TestCompare(2, "y", 200), ] ) b = tvm_ffi.Array( [ TestCompare(1, "x", 999), TestCompare(2, "y", 888), ] ) assert RecursiveEq(a, b) def test_compare_off_inside_nested_object() -> None: """TestObjectDerived.v_array contains TestCompare with different ignored_field.""" a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestCompare(1, "n", 100), ] ), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array( [ TestCompare(1, "n", 999), ] ), ) assert RecursiveEq(a, b) # --------------------------------------------------------------------------- # SmallStr / Str cross-variant in nested context # --------------------------------------------------------------------------- def test_object_with_short_vs_long_string_field() -> None: """SmallStr (<=7 bytes) vs Str (>7 bytes) stored in TestObjectDerived.v_str.""" a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="hi", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="a_very_long_string", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) assert not RecursiveEq(a, b) def test_array_of_mixed_length_strings() -> None: """Array mixing SmallStr (<=7 bytes) and Str (>7 bytes).""" a = tvm_ffi.Array(["hi", "a_very_long_string", "ok"]) b = tvm_ffi.Array(["hi", "a_very_long_string", "ok"]) c = tvm_ffi.Array(["hi", "a_very_long_string", "zz"]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) assert RecursiveLt(a, c) # --------------------------------------------------------------------------- # Map / Dict with nested values # --------------------------------------------------------------------------- def test_map_with_array_values_eq() -> None: a = tvm_ffi.Map({"k": tvm_ffi.Array([1, 2])}) b = tvm_ffi.Map({"k": tvm_ffi.Array([1, 2])}) c = tvm_ffi.Map({"k": tvm_ffi.Array([1, 3])}) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) def test_dict_with_object_values_eq() -> None: a = tvm_ffi.Dict( { "k": TestIntPair(1, 2), } ) b = tvm_ffi.Dict( { "k": TestIntPair(1, 2), } ) c = tvm_ffi.Dict( { "k": TestIntPair(1, 3), } ) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) # --------------------------------------------------------------------------- # None in nested contexts # --------------------------------------------------------------------------- def test_array_with_none_elements() -> None: a = tvm_ffi.Array([None, 1, None]) b = tvm_ffi.Array([None, 1, None]) c = tvm_ffi.Array([None, 2, None]) assert RecursiveEq(a, b) assert not RecursiveEq(a, c) assert RecursiveLt(a, c) # element 1: 1 < 2 def test_object_with_none_in_array_field() -> None: a = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([None, 1]), ) b = create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([None, 1]), ) assert RecursiveEq(a, b) # --------------------------------------------------------------------------- # Cross-container and function-object edge cases # --------------------------------------------------------------------------- def test_array_list_type_mismatch() -> None: arr = tvm_ffi.Array([1, 2]) lst = tvm_ffi.List([1, 2]) assert not RecursiveEq(arr, lst) with pytest.raises(TypeError): RecursiveLt(arr, lst) def test_map_dict_type_mismatch() -> None: m = tvm_ffi.Map({"k": 1}) d = tvm_ffi.Dict({"k": 1}) assert not RecursiveEq(m, d) with pytest.raises(TypeError): RecursiveLt(m, d) def test_function_objects_compare_equal() -> None: """Function objects have no reflected fields, so distinct functions compare equal. This is by design: reflection-based comparison only considers reflected fields, and Function has none. """ f_add_one = tvm_ffi.get_global_func("testing.add_one") f_nop = tvm_ffi.get_global_func("testing.nop") assert RecursiveEq(f_add_one, f_nop) def test_function_same_pointer_eq() -> None: """Same function object compared with itself returns True via pointer identity.""" f = tvm_ffi.get_global_func("testing.add_one") assert RecursiveEq(f, f) # --------------------------------------------------------------------------- # Cycle safety: run in subprocess to avoid hanging/crashing pytest worker # --------------------------------------------------------------------------- def test_cyclic_list_same_pointer_eq() -> None: """Same cyclic list compared with itself returns True via pointer identity.""" lst = tvm_ffi.List() lst.append(lst) assert RecursiveEq(lst, lst) def test_cyclic_list_eq_returns_true() -> None: """Two distinct cyclic lists are considered equal by RecursiveEq. In eq-only mode, the cycle detector treats a re-encountered (lhs, rhs) pair as equal, allowing the recursion to terminate. """ a = tvm_ffi.List() b = tvm_ffi.List() a.append(a) b.append(b) assert RecursiveEq(a, b) def test_cyclic_list_ordering_raises() -> None: """Ordering two distinct cyclic lists raises ValueError.""" a = tvm_ffi.List() b = tvm_ffi.List() a.append(a) b.append(b) with pytest.raises(ValueError, match="cyclic reference"): RecursiveLt(a, b) def test_cyclic_dict_eq_returns_true() -> None: """Two distinct cyclic dicts are considered equal by RecursiveEq. In eq-only mode, the cycle detector treats a re-encountered (lhs, rhs) pair as equal, allowing the recursion to terminate. """ a = tvm_ffi.Dict() b = tvm_ffi.Dict() a["self"] = a b["self"] = b assert RecursiveEq(a, b) def test_cyclic_dict_ordering_raises() -> None: """Ordering two distinct cyclic dicts raises ValueError.""" a = tvm_ffi.Dict() b = tvm_ffi.Dict() a["self"] = a b["self"] = b with pytest.raises(ValueError, match="cyclic reference"): RecursiveLt(a, b) # --------------------------------------------------------------------------- # Comparator laws # --------------------------------------------------------------------------- def test_ordering_laws_on_int_pairs() -> None: """Verify ordering laws (trichotomy, antisymmetry, transitivity) on TestIntPair.""" values = [ TestIntPair(0, 0), TestIntPair(0, 1), TestIntPair(1, 0), TestIntPair(1, 1), ] for a in values: for b in values: lt = RecursiveLt(a, b) eq = RecursiveEq(a, b) gt = RecursiveGt(a, b) le = RecursiveLe(a, b) ge = RecursiveGe(a, b) # Trichotomy: exactly one of a < b, a == b, a > b assert (lt + eq + gt) == 1, f"Trichotomy violated for ({a}, {b})" # Consistency of derived operators assert le == (lt or eq) assert ge == (gt or eq) assert lt == (not ge) assert gt == (not le) # Antisymmetry: (a <= b and b <= a) implies a == b assert eq == (le and ge), f"Antisymmetry violated for ({a}, {b})" # Transitivity for triplets for a in values: for b in values: for c in values: if RecursiveLt(a, b) and RecursiveLt(b, c): assert RecursiveLt(a, c), f"Lt transitivity violated on ({a}, {b}, {c})" if RecursiveLe(a, b) and RecursiveLe(b, c): assert RecursiveLe(a, c), f"Le transitivity violated on ({a}, {b}, {c})" # --------------------------------------------------------------------------- # Deep nesting (iterative stack handles depth > 128) # --------------------------------------------------------------------------- def _make_nested_singleton_array(depth: int) -> object: value: object = 0 for _ in range(depth): value = tvm_ffi.Array([value]) return value def test_depth_1000_nested_eq() -> None: """Deep nested arrays compare correctly with iterative stack.""" a = _make_nested_singleton_array(1000) b = _make_nested_singleton_array(1000) assert RecursiveEq(a, b) # --------------------------------------------------------------------------- # Custom __ffi_eq__ / __ffi_compare__ hooks: TestCustomCompare # --------------------------------------------------------------------------- def test_custom_eq_ignores_label() -> None: """TestCustomCompare.__ffi_eq__ compares only `key`, ignoring `label`.""" a = TestCustomCompare(42, "alpha") b = TestCustomCompare(42, "beta") assert RecursiveEq(a, b) def test_custom_eq_different_key() -> None: a = TestCustomCompare(1, "same") b = TestCustomCompare(2, "same") assert not RecursiveEq(a, b) def test_custom_compare_ordering() -> None: """Ordering uses __ffi_compare__ hook (key only).""" a = TestCustomCompare(1, "zzz") b = TestCustomCompare(2, "aaa") assert RecursiveLt(a, b) assert not RecursiveLt(b, a) def test_custom_eq_in_container() -> None: """Custom-hooked objects inside an Array.""" a = tvm_ffi.Array( [ TestCustomCompare(1, "x"), TestCustomCompare(2, "y"), ] ) b = tvm_ffi.Array( [ TestCustomCompare(1, "different"), TestCustomCompare(2, "labels"), ] ) assert RecursiveEq(a, b) # --------------------------------------------------------------------------- # __ffi_eq__-only types: eq and ordering may diverge (no __ffi_compare__) # --------------------------------------------------------------------------- def test_eq_only_type_eq_uses_hook() -> None: """__ffi_eq__-only type: RecursiveEq uses the hook (compares only key).""" a = TestEqWithoutHash(42, "alpha") b = TestEqWithoutHash(42, "beta") assert RecursiveEq(a, b) def test_eq_only_type_ordering_uses_reflection() -> None: """__ffi_eq__-only type: ordering falls back to field-by-field reflection. Without __ffi_compare__, ordering sees the differing `label` field even though __ffi_eq__ ignores it. This is expected — register __ffi_compare__ for consistent ordering semantics. """ a = TestEqWithoutHash(42, "alpha") b = TestEqWithoutHash(42, "beta") # Eq says equal (hook), but ordering sees label difference (reflection) assert RecursiveEq(a, b) assert RecursiveLt(a, b) # "alpha" < "beta" # --------------------------------------------------------------------------- # __ffi_compare__-equipped types: ordering consistency guaranteed # --------------------------------------------------------------------------- def test_custom_compare_ordering_consistency() -> None: """TestCustomCompare has __ffi_compare__: Eq(a,b) implies not Lt/Gt and both Le/Ge.""" a = TestCustomCompare(42, "alpha") b = TestCustomCompare(42, "beta") assert RecursiveEq(a, b) assert not RecursiveLt(a, b) assert not RecursiveGt(a, b) assert RecursiveLe(a, b) assert RecursiveGe(a, b) # --------------------------------------------------------------------------- # Custom __ffi_eq__ / __ffi_compare__ hooks via @py_class # --------------------------------------------------------------------------- import itertools as _itertools_cmp from typing import Any as _Any_cmp from typing import Callable as _Callable_cmp from tvm_ffi._ffi_api import RecursiveHash as _RecursiveHash_cmp from tvm_ffi.core import Object as _Object_cmp from tvm_ffi.dataclasses import py_class as _py_class_cmp _counter_cmp = _itertools_cmp.count() def _unique_key_cmp(base: str) -> str: return f"testing.cmp_pc.{base}_{next(_counter_cmp)}" @_py_class_cmp(_unique_key_cmp("PyEqHash")) class _PyEqHash(_Object_cmp): key: int label: str def __ffi_hash__(self, fn_hash: _Callable_cmp[..., _Any_cmp]) -> int: return fn_hash(self.key) def __ffi_eq__(self, other: _PyEqHash, fn_eq: _Callable_cmp[..., _Any_cmp]) -> bool: return fn_eq(self.key, other.key) @_py_class_cmp(_unique_key_cmp("PyCmp")) class _PyCmp(_Object_cmp): key: int label: str def __ffi_hash__(self, fn_hash: _Callable_cmp[..., _Any_cmp]) -> int: return fn_hash(self.key) def __ffi_eq__(self, other: _PyCmp, fn_eq: _Callable_cmp[..., _Any_cmp]) -> bool: return fn_eq(self.key, other.key) def __ffi_compare__(self, other: _PyCmp, fn_cmp: _Callable_cmp[..., _Any_cmp]) -> int: return fn_cmp(self.key, other.key) def test_py_class_custom_eq_ignores_label() -> None: assert RecursiveEq(_PyEqHash(42, "alpha"), _PyEqHash(42, "beta")) def test_py_class_custom_eq_different_key() -> None: assert not RecursiveEq(_PyEqHash(1, "same"), _PyEqHash(2, "same")) def test_py_class_custom_eq_hash_consistency() -> None: a, b = _PyEqHash(42, "alpha"), _PyEqHash(42, "beta") assert RecursiveEq(a, b) assert _RecursiveHash_cmp(a) == _RecursiveHash_cmp(b) def test_py_class_custom_compare_ordering() -> None: a = _PyCmp(1, "zzz") b = _PyCmp(2, "aaa") assert RecursiveLt(a, b) assert RecursiveLe(a, b) assert not RecursiveGt(a, b) assert not RecursiveGe(a, b) def test_py_class_custom_compare_equal_keys() -> None: a = _PyCmp(42, "alpha") b = _PyCmp(42, "beta") assert RecursiveEq(a, b) assert RecursiveLe(a, b) assert RecursiveGe(a, b) assert not RecursiveLt(a, b) assert not RecursiveGt(a, b) tvm-ffi-0.1.12/tests/python/test_dataclass_copy.py000066400000000000000000001360171521067262500222460ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: D102, UP006, UP045 """Tests for __copy__, __deepcopy__, and __replace__ on FFI objects.""" from __future__ import annotations import copy import itertools import pickle from typing import Dict, List, Optional import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi._ffi_api import DeepCopy from tvm_ffi.core import Object from tvm_ffi.dataclasses import py_class _counter_pc = itertools.count() def _unique_key_pc(base: str) -> str: return f"testing.copy_pc.{base}_{next(_counter_pc)}" # --------------------------------------------------------------------------- # # __copy__ # --------------------------------------------------------------------------- # class TestShallowCopy: """Tests for copy.copy() / __copy__.""" def test_basic_fields(self) -> None: pair = tvm_ffi.testing.TestIntPair(1, 2) pair_copy = copy.copy(pair) assert pair_copy.a == 1 assert pair_copy.b == 2 def test_creates_new_object(self) -> None: pair = tvm_ffi.testing.TestIntPair(3, 7) pair_copy = copy.copy(pair) assert not pair.same_as(pair_copy) def test_mutable_fields(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=42, v_str="hello") obj_copy = copy.copy(obj) assert obj_copy.v_i64 == 42 # ty: ignore[unresolved-attribute] assert obj_copy.v_str == "hello" # ty: ignore[unresolved-attribute] assert obj_copy.v_f64 == 10.0 # ty: ignore[unresolved-attribute] assert not obj.same_as(obj_copy) def test_mutating_copy_does_not_affect_original(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=1, v_str="a") obj_copy = copy.copy(obj) obj_copy.v_i64 = 99 # ty: ignore[unresolved-attribute] obj_copy.v_str = "z" # ty: ignore[unresolved-attribute] assert obj.v_i64 == 1 # ty: ignore[unresolved-attribute] assert obj.v_str == "a" # ty: ignore[unresolved-attribute] def test_derived_type_preserves_type(self) -> None: v_map = tvm_ffi.convert({"k": 1}) v_array = tvm_ffi.convert([1, 2]) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=5, v_map=v_map, v_array=v_array, ) obj_copy = copy.copy(obj) assert not obj.same_as(obj_copy) assert isinstance(obj_copy, tvm_ffi.testing.TestObjectDerived) assert obj_copy.v_i64 == 5 # shallow copy shares sub-objects assert obj_copy.v_map.same_as(obj.v_map) # ty: ignore[unresolved-attribute] assert obj_copy.v_array.same_as(obj.v_array) # ty: ignore[unresolved-attribute] def test_auto_copy_for_cxx_class(self) -> None: # _TestCxxClassBase is copy-constructible, so copy is auto-enabled # Note: _TestCxxClassBase.__init__ adds 1 to v_i64 and 2 to v_i32 obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) obj_copy = copy.copy(obj) assert obj_copy.v_i64 == 2 assert obj_copy.v_i32 == 4 assert not obj.same_as(obj_copy) def test_non_copyable_type_raises(self) -> None: obj = tvm_ffi.testing.TestNonCopyable(42) with pytest.raises(TypeError, match="does not support copy"): copy.copy(obj) # --------------------------------------------------------------------------- # # __deepcopy__ # --------------------------------------------------------------------------- # class TestDeepCopy: """Tests for copy.deepcopy() / __deepcopy__.""" def test_basic_fields(self) -> None: pair = tvm_ffi.testing.TestIntPair(5, 10) pair_deep = copy.deepcopy(pair) assert pair_deep.a == 5 assert pair_deep.b == 10 assert not pair.same_as(pair_deep) def test_nested_objects_are_copied(self) -> None: inner = tvm_ffi.testing.TestIntPair(1, 2) v_array = tvm_ffi.convert([inner]) v_map = tvm_ffi.convert({"x": "y"}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=10, v_map=v_map, v_array=v_array, ) obj_deep = copy.deepcopy(obj) # top-level is a new object assert not obj.same_as(obj_deep) # nested TestIntPair should also be a new object assert not obj.v_array[0].same_as(obj_deep.v_array[0]) # ty: ignore[unresolved-attribute] # but values are preserved assert obj_deep.v_array[0].a == 1 # ty: ignore[unresolved-attribute] assert obj_deep.v_array[0].b == 2 # ty: ignore[unresolved-attribute] def test_shared_references_preserved(self) -> None: """Two array slots pointing to the same object should still share after deepcopy.""" shared = tvm_ffi.testing.TestIntPair(7, 8) v_array = tvm_ffi.convert([shared, shared]) v_map = tvm_ffi.convert({"a": "b"}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) assert obj.v_array[0].same_as(obj.v_array[1]) # ty: ignore[unresolved-attribute] obj_deep = copy.deepcopy(obj) # the copies should still share assert obj_deep.v_array[0].same_as(obj_deep.v_array[1]) # ty: ignore[unresolved-attribute] # but they must be distinct from the originals assert not obj.v_array[0].same_as(obj_deep.v_array[0]) # ty: ignore[unresolved-attribute] def test_shared_containers_preserved(self) -> None: """Two array slots pointing to the same container should still share after deepcopy.""" inner = tvm_ffi.convert([1, 2, 3]) v_array = tvm_ffi.convert([inner, inner]) v_map = tvm_ffi.convert({"a": "b"}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) assert obj.v_array[0].same_as(obj.v_array[1]) # ty: ignore[unresolved-attribute] obj_deep = copy.deepcopy(obj) # the copied containers should still share identity assert obj_deep.v_array[0].same_as(obj_deep.v_array[1]) # ty: ignore[unresolved-attribute] # but they must be distinct from the originals assert not obj.v_array[0].same_as(obj_deep.v_array[0]) # ty: ignore[unresolved-attribute] def test_original_untouched(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=42, v_str="original") obj_deep = copy.deepcopy(obj) obj_deep.v_i64 = 0 # ty: ignore[unresolved-attribute] obj_deep.v_str = "modified" # ty: ignore[unresolved-attribute] assert obj.v_i64 == 42 # ty: ignore[unresolved-attribute] assert obj.v_str == "original" # ty: ignore[unresolved-attribute] def test_self_referencing_cycle(self) -> None: """An object whose array field contains itself should deepcopy correctly.""" v_map = tvm_ffi.convert({"a": "b"}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=tvm_ffi.convert([]), ) # Create self-reference: obj.v_array = [obj] obj.v_array = tvm_ffi.convert([obj]) # ty: ignore[unresolved-attribute] obj_deep = copy.deepcopy(obj) assert not obj.same_as(obj_deep) # The cycle should be preserved: copy -> copy.v_array[0] -> same copy assert obj_deep.v_array[0].same_as(obj_deep) # ty: ignore[unresolved-attribute] def test_mutual_reference_cycle(self) -> None: """Two objects referencing each other should deepcopy with cycle preserved.""" v_map = tvm_ffi.convert({"a": "b"}) obj_a = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=tvm_ffi.convert([]), ) obj_b = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=2, v_map=v_map, v_array=tvm_ffi.convert([obj_a]), ) # Close the cycle: A -> B and B -> A obj_a.v_array = tvm_ffi.convert([obj_b]) # ty: ignore[unresolved-attribute] deep_a = copy.deepcopy(obj_a) assert not obj_a.same_as(deep_a) # deep_a.v_array[0] is the copy of obj_b deep_b = deep_a.v_array[0] # ty: ignore[unresolved-attribute] assert not obj_b.same_as(deep_b) # The cycle should be preserved: deep_a -> deep_b -> deep_a assert deep_b.v_array[0].same_as(deep_a) # Values are preserved assert deep_a.v_i64 == 1 # ty: ignore[unresolved-attribute] assert deep_b.v_i64 == 2 def test_array_root(self) -> None: """Deepcopy with a bare Array as root should create a new array.""" inner = tvm_ffi.testing.TestIntPair(1, 2) arr = tvm_ffi.convert([inner, "hello", 42]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) # inner object is deep-copied assert not arr[0].same_as(arr_deep[0]) assert arr_deep[0].a == 1 # primitives and strings preserved assert arr_deep[1] == "hello" assert arr_deep[2] == 42 def test_map_root(self) -> None: """Deepcopy with a bare Map as root should create a new map.""" inner = tvm_ffi.testing.TestIntPair(3, 4) m = tvm_ffi.convert({"key": inner}) m_deep = copy.deepcopy(m) assert not m.same_as(m_deep) # inner object is deep-copied assert not m["key"].same_as(m_deep["key"]) assert m_deep["key"].a == 3 def test_dict_root(self) -> None: """Deepcopy with a bare Dict as root should create a new dict.""" inner = tvm_ffi.testing.TestIntPair(3, 4) d = tvm_ffi.Dict({"key": inner}) d_deep = copy.deepcopy(d) assert not d.same_as(d_deep) # inner object is deep-copied assert not d["key"].same_as(d_deep["key"]) assert d_deep["key"].a == 3 def test_auto_deepcopy_for_cxx_class(self) -> None: # _TestCxxClassBase is copy-constructible, so deepcopy is auto-enabled # Note: _TestCxxClassBase.__init__ adds 1 to v_i64 and 2 to v_i32 obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) obj_deep = copy.deepcopy(obj) assert obj_deep.v_i64 == 2 assert obj_deep.v_i32 == 4 assert not obj.same_as(obj_deep) def test_non_copyable_type_raises(self) -> None: obj = tvm_ffi.testing.TestNonCopyable(42) with pytest.raises((TypeError, RuntimeError), match="not copy-constructible"): copy.deepcopy(obj) def test_long_string_in_array(self) -> None: """Strings exceeding inline threshold are heap-allocated objects. deepcopy must treat them as immutable terminals, not call CopyObject. """ long_str = "a" * 100 arr = tvm_ffi.convert([long_str]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert arr_deep[0] == long_str def test_long_string_in_object_field(self) -> None: """Heap-allocated string as a field value should survive deepcopy.""" long_str = "x" * 200 obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=1, v_str=long_str) obj_deep = copy.deepcopy(obj) assert obj_deep.v_str == long_str # ty: ignore[unresolved-attribute] def test_any_field_with_object(self) -> None: """Any-typed field containing an object must be recursively copied.""" inner = tvm_ffi.testing.TestIntPair(3, 4) obj = tvm_ffi.testing.create_object("testing.TestDeepCopyEdges", v_any=inner, v_obj=inner) obj_deep = copy.deepcopy(obj) assert not obj.same_as(obj_deep) assert not inner.same_as(obj_deep.v_any) # ty: ignore[unresolved-attribute] assert obj_deep.v_any.a == 3 # ty: ignore[unresolved-attribute] def test_any_field_with_array(self) -> None: """Any-typed field containing an Array must be recursively copied.""" inner_arr = tvm_ffi.convert([1, 2, 3]) obj = tvm_ffi.testing.create_object( "testing.TestDeepCopyEdges", v_any=inner_arr, v_obj=inner_arr ) obj_deep = copy.deepcopy(obj) assert not inner_arr.same_as(obj_deep.v_any) # ty: ignore[unresolved-attribute] assert list(obj_deep.v_any) == [1, 2, 3] # ty: ignore[unresolved-attribute] def test_any_field_with_map(self) -> None: """Any-typed field containing a Map must be recursively copied.""" inner_map = tvm_ffi.convert({"k": "v"}) obj = tvm_ffi.testing.create_object( "testing.TestDeepCopyEdges", v_any=inner_map, v_obj=inner_map ) obj_deep = copy.deepcopy(obj) assert not inner_map.same_as(obj_deep.v_any) # ty: ignore[unresolved-attribute] assert obj_deep.v_any["k"] == "v" # ty: ignore[unresolved-attribute] def test_objectref_field_with_array(self) -> None: """ObjectRef field holding runtime Array must go through Resolve.""" inner_arr = tvm_ffi.convert([10, 20]) obj = tvm_ffi.testing.create_object( "testing.TestDeepCopyEdges", v_any=None, v_obj=inner_arr ) obj_deep = copy.deepcopy(obj) assert not inner_arr.same_as(obj_deep.v_obj) # ty: ignore[unresolved-attribute] assert list(obj_deep.v_obj) == [10, 20] # ty: ignore[unresolved-attribute] def test_objectref_field_with_map(self) -> None: """ObjectRef field holding runtime Map must go through Resolve.""" inner_map = tvm_ffi.convert({"a": 1}) obj = tvm_ffi.testing.create_object( "testing.TestDeepCopyEdges", v_any=None, v_obj=inner_map ) obj_deep = copy.deepcopy(obj) assert not inner_map.same_as(obj_deep.v_obj) # ty: ignore[unresolved-attribute] assert obj_deep.v_obj["a"] == 1 # ty: ignore[unresolved-attribute] def test_any_field_sharing_preserved(self) -> None: """Shared references through Any and ObjectRef fields are preserved.""" shared = tvm_ffi.testing.TestIntPair(5, 6) obj = tvm_ffi.testing.create_object("testing.TestDeepCopyEdges", v_any=shared, v_obj=shared) obj_deep = copy.deepcopy(obj) # Both fields should point to the same copied object assert obj_deep.v_any.same_as(obj_deep.v_obj) # ty: ignore[unresolved-attribute] assert not shared.same_as(obj_deep.v_any) # ty: ignore[unresolved-attribute] # --------------------------------------------------------------------------- # # Deep copy branch coverage (C++ dataclass.cc) # --------------------------------------------------------------------------- # _deep_copy = tvm_ffi.get_global_func("ffi.DeepCopy") class TestDeepCopyBranches: """Branch-coverage tests targeting every code path in deep_copy.cc.""" # --- Run(): primitive passthrough (type_index < kStaticObjectBegin) --- def test_primitive_int(self) -> None: assert _deep_copy(42) == 42 def test_primitive_float(self) -> None: result = _deep_copy(3.14) assert abs(result - 3.14) < 1e-9 def test_primitive_str(self) -> None: assert _deep_copy("hello") == "hello" def test_primitive_none(self) -> None: assert _deep_copy(None) is None def test_primitive_bool(self) -> None: assert _deep_copy(True) is True assert _deep_copy(False) is False # --- Resolve(): array with various element types --- def test_array_all_ints(self) -> None: """All elements are primitives — Resolve() returns each as-is.""" arr = tvm_ffi.convert([1, 2, 3]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert list(arr_deep) == [1, 2, 3] def test_array_all_strings(self) -> None: arr = tvm_ffi.convert(["a", "bb", "ccc"]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert list(arr_deep) == ["a", "bb", "ccc"] def test_array_with_none_elements(self) -> None: arr = tvm_ffi.convert([None, 1, None]) arr_deep = copy.deepcopy(arr) assert arr_deep[0] is None assert arr_deep[1] == 1 assert arr_deep[2] is None def test_array_mixed_primitive_types(self) -> None: """Array with int, float, str, bool, None — all primitives.""" arr = tvm_ffi.convert([42, 3.14, "hi", True, None]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert arr_deep[0] == 42 assert abs(arr_deep[1] - 3.14) < 1e-9 assert arr_deep[2] == "hi" assert arr_deep[3] is True assert arr_deep[4] is None def test_array_mixed_with_objects_and_containers(self) -> None: """Array with int, str, None, object, nested array, nested map.""" inner_obj = tvm_ffi.testing.TestIntPair(1, 2) inner_arr = tvm_ffi.convert([10, 20]) inner_map = tvm_ffi.convert({"k": "v"}) arr = tvm_ffi.convert([42, "hello", None, inner_obj, inner_arr, inner_map]) arr_deep = copy.deepcopy(arr) # primitives pass through assert arr_deep[0] == 42 assert arr_deep[1] == "hello" assert arr_deep[2] is None # object is deep-copied assert not arr[3].same_as(arr_deep[3]) assert arr_deep[3].a == 1 # nested array is deep-copied assert not arr[4].same_as(arr_deep[4]) assert list(arr_deep[4]) == [10, 20] # nested map is deep-copied assert not arr[5].same_as(arr_deep[5]) assert arr_deep[5]["k"] == "v" def test_array_empty(self) -> None: arr = tvm_ffi.convert([]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert len(arr_deep) == 0 def test_array_nested_arrays(self) -> None: """Array of arrays — Resolve() recurses into each nested array.""" a = tvm_ffi.convert([1, 2]) b = tvm_ffi.convert([3, 4]) outer = tvm_ffi.convert([a, b]) outer_deep = copy.deepcopy(outer) assert not outer.same_as(outer_deep) assert not outer[0].same_as(outer_deep[0]) assert not outer[1].same_as(outer_deep[1]) assert list(outer_deep[0]) == [1, 2] assert list(outer_deep[1]) == [3, 4] def test_array_nested_maps(self) -> None: """Array of maps.""" m = tvm_ffi.convert({"x": 1}) arr = tvm_ffi.convert([m]) arr_deep = copy.deepcopy(arr) assert not arr[0].same_as(arr_deep[0]) assert arr_deep[0]["x"] == 1 # --- Resolve(): map with various key/value types --- def test_map_primitive_keys_and_values(self) -> None: m = tvm_ffi.convert({"a": 1, "b": 2, "c": 3}) m_deep = copy.deepcopy(m) assert not m.same_as(m_deep) assert m_deep["a"] == 1 assert m_deep["b"] == 2 assert m_deep["c"] == 3 def test_map_with_container_values(self) -> None: inner_arr = tvm_ffi.convert([1, 2]) m = tvm_ffi.convert({"arr": inner_arr}) m_deep = copy.deepcopy(m) assert not m["arr"].same_as(m_deep["arr"]) assert list(m_deep["arr"]) == [1, 2] def test_map_with_none_values(self) -> None: m = tvm_ffi.convert({"a": None, "b": 1}) m_deep = copy.deepcopy(m) assert not m.same_as(m_deep) assert m_deep["a"] is None assert m_deep["b"] == 1 def test_map_empty(self) -> None: m = tvm_ffi.convert({}) m_deep = copy.deepcopy(m) assert not m.same_as(m_deep) assert len(m_deep) == 0 # --- Resolve(): dict with various key/value types --- def test_dict_primitive_keys_and_values(self) -> None: d = tvm_ffi.Dict({"a": 1, "b": 2, "c": 3}) d_deep = copy.deepcopy(d) assert not d.same_as(d_deep) assert d_deep["a"] == 1 assert d_deep["b"] == 2 assert d_deep["c"] == 3 def test_dict_with_container_values(self) -> None: inner_arr = tvm_ffi.convert([1, 2]) d = tvm_ffi.Dict({"arr": inner_arr}) d_deep = copy.deepcopy(d) assert not d["arr"].same_as(d_deep["arr"]) assert list(d_deep["arr"]) == [1, 2] def test_dict_with_none_values(self) -> None: d = tvm_ffi.Dict({"a": None, "b": 1}) d_deep = copy.deepcopy(d) assert not d.same_as(d_deep) assert d_deep["a"] is None assert d_deep["b"] == 1 def test_dict_empty(self) -> None: d = tvm_ffi.Dict({}) d_deep = copy.deepcopy(d) assert not d.same_as(d_deep) assert len(d_deep) == 0 # --- Resolve(): copy_map_ hit (shared references across containers) --- def test_shared_array_identity_in_outer_array(self) -> None: """Same array appears 3 times — all copies share identity.""" shared = tvm_ffi.convert([1, 2]) outer = tvm_ffi.convert([shared, shared, shared]) outer_deep = copy.deepcopy(outer) assert outer_deep[0].same_as(outer_deep[1]) assert outer_deep[1].same_as(outer_deep[2]) assert not outer[0].same_as(outer_deep[0]) def test_shared_map_identity_in_outer_array(self) -> None: shared = tvm_ffi.convert({"x": 1}) outer = tvm_ffi.convert([shared, shared]) outer_deep = copy.deepcopy(outer) assert outer_deep[0].same_as(outer_deep[1]) assert not outer[0].same_as(outer_deep[0]) def test_shared_dict_identity_in_outer_array(self) -> None: shared = tvm_ffi.Dict({"x": 1}) outer = tvm_ffi.convert([shared, shared]) outer_deep = copy.deepcopy(outer) assert outer_deep[0].same_as(outer_deep[1]) assert not outer[0].same_as(outer_deep[0]) def test_shared_object_across_array_and_map(self) -> None: """Same object referenced from both v_array and v_map.""" pair = tvm_ffi.testing.TestIntPair(7, 8) v_array = tvm_ffi.convert([pair]) v_map = tvm_ffi.convert({"p": pair}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) obj_deep = copy.deepcopy(obj) # both fields should refer to the same copied object deep_from_arr = obj_deep.v_array[0] # ty: ignore[unresolved-attribute] deep_from_map = obj_deep.v_map["p"] # ty: ignore[unresolved-attribute] assert deep_from_arr.same_as(deep_from_map) assert not pair.same_as(deep_from_arr) # --- ResolveFields: container field with only primitives --- # Resolve() always rebuilds the container, so setter is always called. def test_field_array_only_primitives(self) -> None: v_array = tvm_ffi.convert([1, 2, 3]) v_map = tvm_ffi.convert({"k": "v"}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) obj_deep = copy.deepcopy(obj) assert not obj.v_array.same_as(obj_deep.v_array) # ty: ignore[unresolved-attribute] assert list(obj_deep.v_array) == [1, 2, 3] # ty: ignore[unresolved-attribute] def test_field_map_only_primitives(self) -> None: v_array = tvm_ffi.convert([]) v_map = tvm_ffi.convert({"x": 1, "y": 2}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) obj_deep = copy.deepcopy(obj) assert not obj.v_map.same_as(obj_deep.v_map) # ty: ignore[unresolved-attribute] assert obj_deep.v_map["x"] == 1 # ty: ignore[unresolved-attribute] assert obj_deep.v_map["y"] == 2 # ty: ignore[unresolved-attribute] # --- ResolveFields: empty container fields --- def test_field_empty_containers(self) -> None: v_array = tvm_ffi.convert([]) v_map = tvm_ffi.convert({}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=v_map, v_array=v_array, ) obj_deep = copy.deepcopy(obj) assert not obj.v_array.same_as(obj_deep.v_array) # ty: ignore[unresolved-attribute] assert not obj.v_map.same_as(obj_deep.v_map) # ty: ignore[unresolved-attribute] assert len(obj_deep.v_array) == 0 # ty: ignore[unresolved-attribute] assert len(obj_deep.v_map) == 0 # ty: ignore[unresolved-attribute] # --- ResolveFields: shared container across multiple objects --- def test_shared_container_field_across_objects(self) -> None: """Two objects share the same v_array — copy_map_ deduplicates.""" shared_arr = tvm_ffi.convert([1, 2, 3]) obj_a = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=tvm_ffi.convert({}), v_array=shared_arr, ) obj_b = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=2, v_map=tvm_ffi.convert({}), v_array=shared_arr, ) outer = tvm_ffi.convert([obj_a, obj_b]) outer_deep = copy.deepcopy(outer) deep_a = outer_deep[0] deep_b = outer_deep[1] # both should share the same deep-copied array assert deep_a.v_array.same_as(deep_b.v_array) assert not shared_arr.same_as(deep_a.v_array) # --- CopyObject: unsupported type nested in container --- def test_cxx_class_in_array(self) -> None: # Note: _TestCxxClassBase.__init__ adds 1 to v_i64 and 2 to v_i32 obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) arr = tvm_ffi.convert([obj]) arr_deep = copy.deepcopy(arr) assert not arr.same_as(arr_deep) assert not arr[0].same_as(arr_deep[0]) assert arr_deep[0].v_i64 == 2 assert arr_deep[0].v_i32 == 4 def test_cxx_class_in_map_value(self) -> None: # Note: _TestCxxClassBase.__init__ adds 1 to v_i64 and 2 to v_i32 obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) m = tvm_ffi.convert({"k": obj}) m_deep = copy.deepcopy(m) assert not m.same_as(m_deep) assert not m["k"].same_as(m_deep["k"]) assert m_deep["k"].v_i64 == 2 assert m_deep["k"].v_i32 == 4 def test_non_copyable_type_in_array(self) -> None: obj = tvm_ffi.testing.TestNonCopyable(1) arr = tvm_ffi.convert([obj]) with pytest.raises(RuntimeError, match="not copy-constructible"): copy.deepcopy(arr) def test_non_copyable_type_in_map_value(self) -> None: obj = tvm_ffi.testing.TestNonCopyable(1) m = tvm_ffi.convert({"k": obj}) with pytest.raises(RuntimeError, match="not copy-constructible"): copy.deepcopy(m) # --- Deeply nested structures --- def test_deeply_nested_containers(self) -> None: """Array > Map > Array > object — all levels resolved.""" pair = tvm_ffi.testing.TestIntPair(9, 10) inner_arr = tvm_ffi.convert([pair]) inner_map = tvm_ffi.convert({"items": inner_arr}) outer = tvm_ffi.convert([inner_map]) outer_deep = copy.deepcopy(outer) deep_pair = outer_deep[0]["items"][0] assert not pair.same_as(deep_pair) assert deep_pair.a == 9 assert deep_pair.b == 10 def test_object_with_deeply_nested_field(self) -> None: """Object whose array field contains a map containing an object.""" pair = tvm_ffi.testing.TestIntPair(5, 6) inner_map = tvm_ffi.convert({"pair": pair}) v_array = tvm_ffi.convert([inner_map]) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=tvm_ffi.convert({}), v_array=v_array, ) obj_deep = copy.deepcopy(obj) deep_pair = obj_deep.v_array[0]["pair"] # ty: ignore[unresolved-attribute] assert not pair.same_as(deep_pair) assert deep_pair.a == 5 # --- Cycle preservation with immutable root containers --- def test_cycle_list_root_map_backref_preserved(self) -> None: """Control case: List root with Map back-reference should preserve cycle.""" root_list = tvm_ffi.List() m = tvm_ffi.Map({"list": root_list}) root_list.append(m) deep_list = copy.deepcopy(root_list) assert not root_list.same_as(deep_list) assert deep_list[0]["list"].same_as(deep_list) def test_cycle_map_root_list_backref_preserved(self) -> None: """Map root with List child pointing back should preserve cycle to root copy.""" l = tvm_ffi.List() m = tvm_ffi.Map({"list": l}) l.append(m) deep_map = copy.deepcopy(m) assert not m.same_as(deep_map) assert not l.same_as(deep_map["list"]) assert deep_map["list"][0].same_as(deep_map) def test_cycle_array_root_list_backref_preserved(self) -> None: """Array root with List child pointing back should preserve cycle to root copy.""" l = tvm_ffi.List() a = tvm_ffi.Array([l]) l.append(a) deep_arr = copy.deepcopy(a) assert not a.same_as(deep_arr) assert not l.same_as(deep_arr[0]) assert deep_arr[0][0].same_as(deep_arr) def test_cycle_array_root_dict_backref_preserved(self) -> None: """Array root with Dict child pointing back should preserve cycle to root copy.""" d = tvm_ffi.Dict() a = tvm_ffi.Array([d]) d["self"] = a deep_arr = copy.deepcopy(a) assert not a.same_as(deep_arr) assert not d.same_as(deep_arr[0]) assert deep_arr[0]["self"].same_as(deep_arr) def test_cycle_map_root_dict_backref_preserved(self) -> None: """Map root with Dict child pointing back should preserve cycle to root copy.""" d = tvm_ffi.Dict() m = tvm_ffi.Map({"dict": d}) d["self"] = m deep_map = copy.deepcopy(m) assert not m.same_as(deep_map) assert not d.same_as(deep_map["dict"]) assert deep_map["dict"]["self"].same_as(deep_map) def test_cycle_map_root_backref_identity_not_duplicated(self) -> None: """Back-references in a map-root cycle should point to the root copied map.""" shared_list = tvm_ffi.List() m = tvm_ffi.Map({"l1": shared_list, "l2": shared_list}) shared_list.append(m) deep_map = copy.deepcopy(m) assert deep_map["l1"].same_as(deep_map["l2"]) assert deep_map["l1"][0].same_as(deep_map) def test_cycle_map_root_list_key_backref_preserved(self) -> None: """Map-root cycles through keys should preserve back-reference to copied root.""" key_list = tvm_ffi.List() m = tvm_ffi.Map({key_list: 1}) key_list.append(m) deep_map = copy.deepcopy(m) deep_key = next(iter(deep_map.keys())) assert isinstance(deep_key, tvm_ffi.List) assert deep_key[0].same_as(deep_map) def test_cycle_map_root_dict_key_backref_preserved(self) -> None: """Map-root cycles through Dict keys should preserve back-reference to copied root.""" key_dict = tvm_ffi.Dict() m = tvm_ffi.Map({key_dict: 1}) key_dict["self"] = m deep_map = copy.deepcopy(m) deep_key = next(iter(deep_map.keys())) assert isinstance(deep_key, tvm_ffi.Dict) assert deep_key["self"].same_as(deep_map) def test_cycle_array_root_dict_contains_root_as_key(self) -> None: """Array root with Dict child using the root as key should fix key to copied root.""" d = tvm_ffi.Dict() root = tvm_ffi.Array([d]) d[root] = 1 deep_root = copy.deepcopy(root) deep_dict = deep_root[0] deep_key = next(iter(deep_dict.keys())) assert not root.same_as(deep_root) assert deep_key.same_as(deep_root) assert not deep_key.same_as(root) def test_cycle_map_root_dict_contains_root_as_key(self) -> None: """Map root with Dict child using the root as key should fix key to copied root.""" d = tvm_ffi.Dict() root = tvm_ffi.Map({"d": d}) d[root] = 1 deep_root = copy.deepcopy(root) deep_dict = deep_root["d"] deep_key = next(iter(deep_dict.keys())) assert not root.same_as(deep_root) assert deep_key.same_as(deep_root) assert not deep_key.same_as(root) # --- Python deepcopy protocol consistency for immutable Shape --- def test_shape_root_python_deepcopy_matches_ffi_deepcopy(self) -> None: """copy.deepcopy(Shape) should be consistent with ffi.DeepCopy.""" deep_copy_fn = tvm_ffi.get_global_func("ffi.DeepCopy") s = tvm_ffi.Shape((2, 3, 4)) ffi_copied = deep_copy_fn(s) py_copied = copy.deepcopy(s) assert py_copied == ffi_copied assert isinstance(py_copied, type(s)) def test_shape_inside_python_container_deepcopy(self) -> None: """Python container deepcopy should handle Shape payloads.""" s = tvm_ffi.Shape((1, 2)) payload = [s, {"shape": s}] copied = copy.deepcopy(payload) assert copied[0] == s assert copied[1]["shape"] == s # ty: ignore[invalid-argument-type] # --- Cycle fixup: immutable container → reflected object back-reference --- def test_cycle_array_root_object_backreference(self) -> None: """Array A → Object X, X.v_array = A. Deep copy from A.""" obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=42, v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) arr = tvm_ffi.Array([obj]) obj.v_array = arr # ty: ignore[unresolved-attribute] arr_deep = _deep_copy(arr) assert not arr.same_as(arr_deep) obj_deep = arr_deep[0] assert not obj.same_as(obj_deep) assert obj_deep.v_i64 == 42 assert not obj_deep.v_array.same_as(arr) assert obj_deep.v_array.same_as(arr_deep) def test_cycle_map_root_object_backreference(self) -> None: """Map M → Object X, X.v_map = M. Deep copy from M.""" obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=7, v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) m = tvm_ffi.Map({"key": obj}) obj.v_map = m # ty: ignore[unresolved-attribute] m_deep = _deep_copy(m) assert not m.same_as(m_deep) obj_deep = m_deep["key"] assert not obj.same_as(obj_deep) assert obj_deep.v_i64 == 7 assert not obj_deep.v_map.same_as(m) assert obj_deep.v_map.same_as(m_deep) def test_cycle_nested_array_object_array(self) -> None: """Array → Object → Array → Object → back to root Array.""" inner = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) outer = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=2, v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([inner]), ) root_arr = tvm_ffi.Array([outer]) inner.v_array = root_arr # ty: ignore[unresolved-attribute] root_deep = _deep_copy(root_arr) assert not root_arr.same_as(root_deep) outer_deep = root_deep[0] inner_deep = outer_deep.v_array[0] assert not inner_deep.v_array.same_as(root_arr) assert inner_deep.v_array.same_as(root_deep) # --------------------------------------------------------------------------- # # __replace__ # --------------------------------------------------------------------------- # class TestReplace: """Tests for __replace__.""" def test_replace_writable_fields(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=1, v_str="a") obj2 = obj.__replace__(v_i64=99) # ty: ignore[unresolved-attribute] assert obj2.v_i64 == 99 assert obj2.v_str == "a" assert not obj.same_as(obj2) def test_replace_multiple_fields(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=1, v_str="a") obj2 = obj.__replace__(v_i64=42, v_str="world") # ty: ignore[unresolved-attribute] assert obj2.v_i64 == 42 assert obj2.v_str == "world" def test_replace_no_kwargs_is_copy(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=7, v_str="hi") obj2 = obj.__replace__() # ty: ignore[unresolved-attribute] assert obj2.v_i64 == 7 assert obj2.v_str == "hi" assert not obj.same_as(obj2) def test_original_unchanged(self) -> None: obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=5, v_str="x") obj.__replace__(v_i64=100) # ty: ignore[unresolved-attribute] assert obj.v_i64 == 5 # ty: ignore[unresolved-attribute] def test_replace_readonly_field(self) -> None: # __replace__ uses the FFIProperty.set() escape hatch, # so it works even on frozen / read-only fields. pair = tvm_ffi.testing.TestIntPair(3, 4) pair2 = pair.__replace__(a=10) # ty: ignore[unresolved-attribute] assert pair2.a == 10 assert pair2.b == 4 assert pair.a == 3 # original unchanged def test_auto_replace_for_cxx_class(self) -> None: # _TestCxxClassBase is copy-constructible, so replace is auto-enabled # Note: _TestCxxClassBase.__init__ adds 1 to v_i64 and 2 to v_i32 obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) obj2 = obj.__replace__(v_i64=99) # ty: ignore[unresolved-attribute] assert obj2.v_i64 == 99 assert obj2.v_i32 == 4 assert not obj.same_as(obj2) def test_non_copyable_type_raises(self) -> None: obj = tvm_ffi.testing.TestNonCopyable(42) with pytest.raises(TypeError, match="does not support copy"): obj.__replace__() # ty: ignore[unresolved-attribute] # --------------------------------------------------------------------------- # # @py_class copy/deepcopy with rich field types # --------------------------------------------------------------------------- # class TestPyClassCopyRichFields: """copy.copy / copy.deepcopy with container and optional fields on @py_class.""" def test_shallow_copy_containers(self) -> None: @py_class(_unique_key_pc("SCCont")) class SCCont(Object): x: int items: List[int] data: Dict[str, int] obj = SCCont(x=42, items=[1, 2, 3], data={"a": 1}) obj2 = copy.copy(obj) assert obj2.x == 42 assert len(obj2.items) == 3 assert obj2.data["a"] == 1 assert obj.items.same_as(obj2.items) # ty:ignore[unresolved-attribute] assert obj.data.same_as(obj2.data) # ty:ignore[unresolved-attribute] def test_deep_copy_containers(self) -> None: @py_class(_unique_key_pc("DCCont")) class DCCont(Object): x: int items: List[int] data: Dict[str, int] obj = DCCont(x=42, items=[1, 2, 3], data={"a": 1}) obj2 = copy.deepcopy(obj) assert obj2.x == 42 assert len(obj2.items) == 3 assert obj2.data["a"] == 1 assert not obj.items.same_as(obj2.items) # ty:ignore[unresolved-attribute] assert not obj.data.same_as(obj2.data) # ty:ignore[unresolved-attribute] def test_shallow_copy_nested_containers(self) -> None: @py_class(_unique_key_pc("SCNest")) class SCNest(Object): matrix: List[List[int]] obj = SCNest(matrix=[[1, 2], [3, 4]]) obj2 = copy.copy(obj) assert obj.matrix.same_as(obj2.matrix) # ty:ignore[unresolved-attribute] def test_deep_copy_nested_containers(self) -> None: @py_class(_unique_key_pc("DCNest")) class DCNest(Object): matrix: List[List[int]] obj = DCNest(matrix=[[1, 2], [3, 4]]) obj2 = copy.deepcopy(obj) assert obj2.matrix[0][0] == 1 assert obj2.matrix[1][1] == 4 assert not obj.matrix.same_as(obj2.matrix) # ty:ignore[unresolved-attribute] def test_deep_copy_mutation_independent(self) -> None: @py_class(_unique_key_pc("DCMutInd")) class DCMutInd(Object): x: int items: List[int] obj = DCMutInd(x=1, items=[10, 20]) obj2 = copy.deepcopy(obj) obj2.x = 99 assert obj.x == 1 obj2.items[0] = 999 assert obj.items[0] == 10 def test_shallow_copy_optional_fields(self) -> None: @py_class(_unique_key_pc("SCOpt")) class SCOpt(Object): x: Optional[int] items: Optional[List[int]] obj = SCOpt(x=42, items=[1, 2]) obj2 = copy.copy(obj) assert obj2.x == 42 assert len(obj2.items) == 2 # ty:ignore[invalid-argument-type] def test_deep_copy_with_none_optional(self) -> None: @py_class(_unique_key_pc("DCOptNone")) class DCOptNone(Object): x: Optional[int] items: Optional[List[int]] obj = DCOptNone(x=None, items=None) obj2 = copy.deepcopy(obj) assert obj2.x is None assert obj2.items is None def test_replace_with_containers(self) -> None: @py_class(_unique_key_pc("ReplCont")) class ReplCont(Object): x: int items: List[int] obj = ReplCont(x=1, items=[1, 2, 3]) obj2 = obj.__replace__(x=99) # ty:ignore[unresolved-attribute] assert obj2.x == 99 assert tuple(obj2.items) == (1, 2, 3) assert obj.x == 1 # --------------------------------------------------------------------------- # # DeepCopy FFI with @py_class containers # --------------------------------------------------------------------------- # class TestPyClassDeepCopyContainers: """DeepCopy FFI function with @py_class container fields.""" def test_deep_copy_list_field(self) -> None: @py_class(_unique_key_pc("DCList")) class DCList(Object): items: List[int] obj = DCList(items=[1, 2, 3]) obj2 = DeepCopy(obj) assert tuple(obj2.items) == (1, 2, 3) assert not obj.items.same_as(obj2.items) # ty:ignore[unresolved-attribute] def test_deep_copy_dict_field(self) -> None: @py_class(_unique_key_pc("DCDict")) class DCDict(Object): data: Dict[str, int] obj = DCDict(data={"a": 1, "b": 2}) obj2 = DeepCopy(obj) assert obj2.data["a"] == 1 assert not obj.data.same_as(obj2.data) # ty:ignore[unresolved-attribute] def test_deep_copy_nested(self) -> None: @py_class(_unique_key_pc("DCNested")) class DCNested(Object): matrix: List[List[int]] obj = DCNested(matrix=[[1, 2], [3, 4]]) obj2 = DeepCopy(obj) assert obj2.matrix[0][0] == 1 assert not obj.matrix.same_as(obj2.matrix) # ty:ignore[unresolved-attribute] def test_deep_copy_optional_none(self) -> None: @py_class(_unique_key_pc("DCOptN")) class DCOptN(Object): items: Optional[List[int]] obj = DCOptN(items=None) assert DeepCopy(obj).items is None def test_deep_copy_optional_value(self) -> None: @py_class(_unique_key_pc("DCOptV")) class DCOptV(Object): items: Optional[List[int]] obj = DCOptV(items=[1, 2, 3]) obj2 = DeepCopy(obj) assert tuple(obj2.items) == (1, 2, 3) assert not obj.items.same_as(obj2.items) # ty:ignore[unresolved-attribute] # --------------------------------------------------------------------------- # # Copy of @py_class with custom __init__ # --------------------------------------------------------------------------- # class TestPyClassCopyCustomInit: """Copy of @py_class with init=False and custom __init__.""" def _make_cls(self) -> type: @py_class(_unique_key_pc("CopyCI"), init=False) class CopyCI(Object): a: int b: str def __init__(self, *, b: str, a: int) -> None: self.a = a self.b = b return CopyCI def test_shallow_copy_custom_init(self) -> None: CopyCI = self._make_cls() src = CopyCI(a=1, b="hello") dst = copy.copy(src) assert not src.same_as(dst) assert dst.a == 1 assert dst.b == "hello" def test_deep_copy_custom_init(self) -> None: CopyCI = self._make_cls() src = CopyCI(a=1, b="hello") dst = copy.deepcopy(src) assert not src.same_as(dst) assert dst.a == 1 assert dst.b == "hello" # --------------------------------------------------------------------------- # # Pickle roundtrip for @py_class # --------------------------------------------------------------------------- # # Pickle requires classes to be importable at module level. @py_class(_unique_key_pc("PickleBasic")) class _PickleBasic(Object): a: int b: float c: str d: bool @py_class(_unique_key_pc("PickleOptV")) class _PickleOptV(Object): a: Optional[int] b: Optional[str] @py_class(_unique_key_pc("PickleCont")) class _PickleCont(Object): items: List[int] data: Dict[str, int] @py_class(_unique_key_pc("PickleCI"), init=False) class _PickleCI(Object): a: int b: str def __init__(self, *, b: str, a: int) -> None: self.a = a self.b = b class TestPyClassPickleRoundtrip: """Pickle serialization/deserialization for @py_class objects.""" def test_pickle_basic_fields(self) -> None: obj = _PickleBasic(a=1, b=2.0, c="hello", d=True) obj2 = pickle.loads(pickle.dumps(obj)) assert obj2.a == 1 assert obj2.b == 2.0 assert obj2.c == "hello" assert obj2.d is True def test_pickle_optional_with_values(self) -> None: obj = _PickleOptV(a=42, b="world") obj2 = pickle.loads(pickle.dumps(obj)) assert obj2.a == 42 assert obj2.b == "world" def test_pickle_optional_with_none(self) -> None: obj = _PickleOptV(a=None, b=None) obj2 = pickle.loads(pickle.dumps(obj)) assert obj2.a is None assert obj2.b is None def test_pickle_container_fields(self) -> None: obj = _PickleCont(items=[1, 2, 3], data={"a": 1, "b": 2}) obj2 = pickle.loads(pickle.dumps(obj)) assert tuple(obj2.items) == (1, 2, 3) assert obj2.data["a"] == 1 def test_pickle_custom_init(self) -> None: obj = _PickleCI(a=1, b="hello") obj2 = pickle.loads(pickle.dumps(obj)) assert obj2.a == 1 assert obj2.b == "hello" tvm-ffi-0.1.12/tests/python/test_dataclass_enum.py000066400000000000000000000741171521067262500222420ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for ``Enum`` subclasses with ``type_key``-parameterized inheritance.""" from __future__ import annotations import itertools from typing import ClassVar import pytest import tvm_ffi from tvm_ffi import core from tvm_ffi.core import Object from tvm_ffi.dataclasses import Enum, EnumAttrMap, IntEnum, StrEnum, auto, entry from tvm_ffi.dataclasses.enum import ( ENUM_ATTRS_ATTR, ENUM_ENTRIES_ATTR, ENUM_VALUE_ENTRIES_ATTR, _EnumEntry, ) # --------------------------------------------------------------------------- # Unique type-key generator — ensures no collisions across tests. # --------------------------------------------------------------------------- _counter = itertools.count() def _unique_key(base: str) -> str: return f"testing.py_enum_sub.{base}_{next(_counter)}" # --------------------------------------------------------------------------- # Attribute-carrying form — entry(...) with extra field kwargs # --------------------------------------------------------------------------- def test_attribute_carrying_basic() -> None: class Activation(Enum, type_key=_unique_key("Activation")): output_zero: bool is_monotonic: bool relu: ClassVar[Activation] = entry(output_zero=True, is_monotonic=True) gelu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=False) silu: ClassVar[Activation] = entry(output_zero=False, is_monotonic=True) assert isinstance(Activation.relu, Activation) assert isinstance(Activation.gelu, Activation) assert isinstance(Activation.silu, Activation) assert Activation.relu.output_zero is True # ty: ignore[unresolved-attribute] assert Activation.relu.is_monotonic is True # ty: ignore[unresolved-attribute] assert Activation.gelu.output_zero is False # ty: ignore[unresolved-attribute] assert Activation.silu.is_monotonic is True # ty: ignore[unresolved-attribute] # Ordinals auto-assigned in declaration order. assert Activation.relu._value == 0 assert Activation.gelu._value == 1 assert Activation.silu._value == 2 assert Activation.relu._name == "relu" assert Activation.gelu._name == "gelu" assert Activation.get("relu").same_as(Activation.relu) assert Activation.get("gelu").same_as(Activation.gelu) assert Activation.get("silu").same_as(Activation.silu) def test_entry_rejects__value_kwarg() -> None: """``entry(_value=...)`` conflicts with the auto-assigned ordinal.""" with pytest.raises(TypeError): class _Bad(Enum, type_key=_unique_key("BadValue")): flag: bool a: ClassVar[_Bad] = entry(flag=True, _value=5) def test_entry_rejects__name_kwarg() -> None: """``entry(_name=...)`` conflicts with the auto-assigned declaration key.""" with pytest.raises(TypeError): class _Bad(Enum, type_key=_unique_key("BadName")): flag: bool a: ClassVar[_Bad] = entry(flag=True, _name="other") def test_get_missing_raises() -> None: class Missing(Enum, type_key=_unique_key("Missing")): flag: bool yes: ClassVar[Missing] = entry(flag=True) with pytest.raises(KeyError): Missing.get("no-such-entry") def test_public_value_and_name_are_hidden() -> None: class Hidden(Enum, type_key=_unique_key("Hidden")): yes = auto() assert Hidden.yes._value == 0 assert Hidden.yes._name == "yes" with pytest.raises(AttributeError): _ = Hidden.yes.value with pytest.raises(AttributeError): _ = Hidden.yes.name def test_int_enum_payload_value() -> None: class Colors(IntEnum, type_key=_unique_key("IntEnum")): red = entry(value=10) blue = entry(value=20) assert Colors.red.value == 10 assert Colors.blue.value == 20 assert Colors.red._value == 0 assert Colors.blue._value == 1 assert Colors.red._name == "red" assert Colors.get("red").same_as(Colors.red) assert list(Colors.all_entries()) == [Colors.red, Colors.blue] def test_int_enum_payload_literal_sugar() -> None: class Priority(IntEnum, type_key=_unique_key("IntEnumLiteral")): low = 10 high = 20 assert isinstance(Priority.low, Priority) assert isinstance(Priority.high, Priority) assert not isinstance(Priority.low, int) assert Priority.low.value == 10 assert Priority.high.value == 20 assert Priority.low._name == "low" assert Priority.high._name == "high" assert list(Priority.all_entries()) == [Priority.low, Priority.high] def test_str_enum_payload_value() -> None: class Tokens(StrEnum, type_key=_unique_key("StrEnum")): add = entry(value="+") mul = entry(value="*") assert Tokens.add.value == "+" assert Tokens.mul.value == "*" assert Tokens.add._value == 0 assert Tokens.mul._value == 1 assert Tokens.add._name == "add" assert Tokens.get("mul").same_as(Tokens.mul) def test_str_enum_payload_literal_sugar() -> None: class Opcode(StrEnum, type_key=_unique_key("StrEnumLiteral")): add = "+" mul = "*" assert isinstance(Opcode.add, Opcode) assert isinstance(Opcode.mul, Opcode) assert not isinstance(Opcode.add, str) assert Opcode.add.value == "+" assert Opcode.mul.value == "*" assert Opcode.add._name == "add" assert Opcode.mul._name == "mul" assert list(Opcode.all_entries()) == [Opcode.add, Opcode.mul] def test_payload_enum_compat_behaviors() -> None: class Priority(IntEnum, type_key=_unique_key("IntEnumCompat")): low = 10 high = 20 class Opcode(StrEnum, type_key=_unique_key("StrEnumCompat")): add = "+" mul = "*" assert Priority.low.name == "low" # ty: ignore[possibly-missing-attribute] assert repr(Priority.low) == "Priority.low" assert Priority.low == 10 assert Priority.low != 20 assert str(Priority.low) == "10" assert hash(Priority.low) == hash(10) assert Priority(10).same_as(Priority.low) # ty: ignore[missing-argument] assert Priority("low").same_as(Priority.low) # ty: ignore[missing-argument, invalid-argument-type] assert Opcode.add.name == "add" # ty: ignore[possibly-missing-attribute] assert repr(Opcode.add) == "Opcode.add" assert Opcode.add == "+" assert Opcode.add != "*" assert str(Opcode.add) == "+" assert hash(Opcode.add) == hash("+") assert Opcode("+").same_as(Opcode.add) # ty: ignore[missing-argument, invalid-argument-type] assert Opcode("add").same_as(Opcode.add) # ty: ignore[missing-argument, invalid-argument-type] def test_payload_literal_sugar_preserves_annotated_field_defaults() -> None: class Opcode(StrEnum, type_key=_unique_key("StrEnumLiteralDefault")): arity: int = 0 add = "+" mul = "*" assert isinstance(Opcode.add, Opcode) assert isinstance(Opcode.mul, Opcode) assert Opcode.add.arity == 0 # ty: ignore[unresolved-attribute] assert Opcode.mul.arity == 0 # ty: ignore[unresolved-attribute] assert Opcode.add.value == "+" assert Opcode.mul.value == "*" def test_int_enum_payload_literal_sugar_rejects_invalid_payload() -> None: with pytest.raises(TypeError): class _Bad(IntEnum, type_key=_unique_key("IntEnumLiteralBad")): nope = "x" def test_str_enum_payload_literal_sugar_rejects_invalid_payload() -> None: with pytest.raises(TypeError): class _Bad(StrEnum, type_key=_unique_key("StrEnumLiteralBad")): nope = 42 def test_int_enum_populates_value_entries_typeattr() -> None: class Priority(IntEnum, type_key=_unique_key("IntEnumValueEntries")): low = entry(value=10) high = 20 # literal sugar type_info = Priority.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] value_entries = core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) assert value_entries is not None assert value_entries[10].same_as(Priority.low) assert value_entries[20].same_as(Priority.high) def test_str_enum_populates_value_entries_typeattr() -> None: class Opcode(StrEnum, type_key=_unique_key("StrEnumValueEntries")): add = entry(value="+") mul = "*" # literal sugar type_info = Opcode.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] value_entries = core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) assert value_entries is not None assert value_entries["+"].same_as(Opcode.add) assert value_entries["*"].same_as(Opcode.mul) def test_plain_enum_does_not_create_value_entries_typeattr() -> None: class Status(Enum, type_key=_unique_key("PlainEnumNoValue")): ok: ClassVar[Status] = auto() type_info = Status.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] value_entries = core._lookup_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR) assert value_entries is None def test_payload_literal_sugar_preserves_methods_and_properties() -> None: class Priority(IntEnum, type_key=_unique_key("IntEnumWithMethods")): low = 1 high = 10 def is_high(self) -> bool: return self.value >= 5 @classmethod def default(cls) -> Priority: return cls.low # ty: ignore[invalid-return-type] @property def doubled(self) -> int: return self.value * 2 assert Priority.low.is_high() is False # ty: ignore[possibly-missing-attribute] assert Priority.high.is_high() is True # ty: ignore[possibly-missing-attribute] assert Priority.default().same_as(Priority.low) assert Priority.high.doubled == 20 # ty: ignore[possibly-missing-attribute] assert list(Priority.all_entries()) == [Priority.low, Priority.high] def test_all_entries_iteration_order() -> None: class Ordered(Enum, type_key=_unique_key("Ordered")): tag: str a: ClassVar[Ordered] = entry(tag="first") b: ClassVar[Ordered] = entry(tag="second") c: ClassVar[Ordered] = entry(tag="third") values = list(Ordered.all_entries()) assert len(values) == 3 assert values[0].same_as(Ordered.a) assert values[1].same_as(Ordered.b) assert values[2].same_as(Ordered.c) def test_class_iteration_and_len() -> None: class Ordered(Enum, type_key=_unique_key("OrderedIter")): a = auto() b = auto() c = auto() assert list(Ordered) == [Ordered.a, Ordered.b, Ordered.c] assert len(Ordered) == 3 def test_frozen_variants() -> None: class Frozen(Enum, type_key=_unique_key("Frozen")): flag: bool yes: ClassVar[Frozen] = entry(flag=True) with pytest.raises(AttributeError): Frozen.yes.flag = False # ty: ignore[invalid-assignment] # --------------------------------------------------------------------------- # Bare ClassVar[Cls] (no assignment) — Python-side blank entries # --------------------------------------------------------------------------- def test_bare_classvar_without_cxx_entries() -> None: """``ClassVar[Cls]`` with no value registers a new blank Python entry when the type key has no C++ backing. """ class Status(Enum, type_key=_unique_key("Status")): ok: ClassVar[Status] err: ClassVar[Status] retry: ClassVar[Status] assert isinstance(Status.ok, Status) assert Status.ok._value == 0 assert Status.err._value == 1 assert Status.retry._value == 2 assert Status.ok._name == "ok" assert Status.err._name == "err" assert list(Status.all_entries()) == [Status.ok, Status.err, Status.retry] assert Status.get("ok").same_as(Status.ok) def test_bare_classvar_mixed_with_entry() -> None: """Bare ``ClassVar`` and ``ClassVar = entry(...)`` may mix in one class.""" class Kind(Enum, type_key=_unique_key("Kind")): blank: ClassVar[Kind] tag: str = "" # ordinary field; extra fields follow named: ClassVar[Kind] = entry(tag="hi") # ``ClassVar`` binders are processed before ``entry(...)`` assignments. assert Kind.blank._value == 0 assert Kind.named._value == 1 assert Kind.blank._name == "blank" assert Kind.named._name == "named" assert Kind.named.tag == "hi" # ty: ignore[unresolved-attribute] # --------------------------------------------------------------------------- # Bare ``name = entry(...)`` sugar (no ClassVar annotation) # --------------------------------------------------------------------------- def test_bare_entry_sugar_form() -> None: """``name = entry(...)`` without a ``ClassVar`` annotation is picked up.""" class Activation(Enum, type_key=_unique_key("ActivationBare")): output_zero: bool is_monotonic: bool relu = entry(output_zero=True, is_monotonic=True) gelu = entry(output_zero=False, is_monotonic=False) assert isinstance(Activation.relu, Activation) assert Activation.relu.output_zero is True # ty: ignore[unresolved-attribute] assert Activation.gelu.output_zero is False # ty: ignore[unresolved-attribute] assert list(Activation.all_entries()) == [Activation.relu, Activation.gelu] # --------------------------------------------------------------------------- # ``auto()`` — simple Python-side entries without ClassVar annotation # --------------------------------------------------------------------------- def test_auto_basic_no_annotation() -> None: """``name = auto()`` registers a py-side entry with dense auto-ordinals.""" class Priority(Enum, type_key=_unique_key("Priority")): low = auto() medium = auto() high = auto() assert isinstance(Priority.low, Priority) assert Priority.low._value == 0 assert Priority.medium._value == 1 assert Priority.high._value == 2 assert Priority.low._name == "low" assert Priority.high._name == "high" assert list(Priority.all_entries()) == [Priority.low, Priority.medium, Priority.high] def test_auto_with_classvar_annotation() -> None: """``name: ClassVar[Cls] = auto()`` is equivalent to the annotation-less form.""" class Stage(Enum, type_key=_unique_key("Stage")): init: ClassVar[Stage] = auto() run: ClassVar[Stage] = auto() done: ClassVar[Stage] = auto() assert Stage.init._value == 0 assert Stage.run._value == 1 assert Stage.done._value == 2 def test_auto_mixed_with_bare_classvar() -> None: """``auto()`` may coexist with bare ``ClassVar`` binders in one class. Bare ``ClassVar`` binders are processed first (in annotation order), then ``auto()`` / ``entry(...)`` sentinels in class-body order — so ordinals reflect that deterministic two-phase order. """ class Mixed(Enum, type_key=_unique_key("Mixed")): alpha: ClassVar[Mixed] beta = auto() gamma: ClassVar[Mixed] # Binders (alpha, gamma) come first in annotation order, then sentinels. assert Mixed.alpha._value == 0 assert Mixed.gamma._value == 1 assert Mixed.beta._value == 2 assert {v._name for v in Mixed.all_entries()} == {"alpha", "beta", "gamma"} def test_auto_mixed_with_entry() -> None: """``auto()`` and ``entry(...)`` compose on an attribute-carrying enum.""" class Op(Enum, type_key=_unique_key("OpMixedAuto")): arity: int = 0 noop = auto() add = entry(arity=2) neg = entry(arity=1) assert Op.noop._value == 0 assert Op.add._value == 1 assert Op.neg._value == 2 assert Op.noop.arity == 0 # ty: ignore[unresolved-attribute] assert Op.add.arity == 2 # ty: ignore[unresolved-attribute] def test_auto_rejects_already_registered_name() -> None: """``auto()`` on a name already in the entries dict is rejected. ``testing.TestEnumVariant`` pre-registers ``Alpha`` / ``Beta`` from C++, so attempting to *register* (rather than bind) ``Alpha`` via ``auto()`` must fail — bare ``ClassVar[Cls]`` is the way to bind to an existing entry. """ with pytest.raises(RuntimeError): class _Shadow(Enum, type_key="testing.TestEnumVariant"): Alpha = auto() def test_auto_returns_fresh_sentinels() -> None: """Each ``auto()`` call returns a distinct sentinel instance.""" a, b = auto(), auto() assert isinstance(a, _EnumEntry) assert isinstance(b, _EnumEntry) assert a is not b assert a.args == () assert a.kwargs == {} # --------------------------------------------------------------------------- # all_entries / attr_dict # --------------------------------------------------------------------------- def test_all_entries_indexed_by_ordinal() -> None: class K(Enum, type_key=_unique_key("AllEntries")): a: ClassVar[K] b: ClassVar[K] c: ClassVar[K] entries = list(K.all_entries()) assert len(entries) == 3 assert entries[0].same_as(K.a) assert entries[1].same_as(K.b) assert entries[2].same_as(K.c) def test_attr_dict_direct_access() -> None: """The ``attr_dict`` class property returns the live per-variant attrs map.""" class Op(Enum, type_key=_unique_key("OpDirect")): arity: int add: ClassVar[Op] = entry(arity=2) neg: ClassVar[Op] = entry(arity=1) has_side_effects = Op.def_attr("has_side_effects", default=False) has_side_effects[Op.add] = False has_side_effects[Op.neg] = True # Direct read via class-level property. column = Op.attr_dict["has_side_effects"] assert column[Op.add._value] is False assert column[Op.neg._value] is True # --------------------------------------------------------------------------- # EnumAttrMap / def_attr # --------------------------------------------------------------------------- def test_def_attr_basic_get_set() -> None: class Op(Enum, type_key=_unique_key("Op")): arity: int add: ClassVar[Op] = entry(arity=2) neg: ClassVar[Op] = entry(arity=1) cost = Op.def_attr("cost", default=0) assert isinstance(cost, EnumAttrMap) assert cost[Op.add] == 0 assert cost[Op.neg] == 0 cost[Op.add] = 5 cost[Op.neg] = 2 assert cost[Op.add] == 5 assert cost[Op.neg] == 2 def test_def_attr_missing_raises_without_default() -> None: class OpStrict(Enum, type_key=_unique_key("OpStrict")): arity: int add: ClassVar[OpStrict] = entry(arity=2) attr = OpStrict.def_attr("strict_attr") with pytest.raises(KeyError): _ = attr[OpStrict.add] def test_def_attr_get_method_default() -> None: class Op(Enum, type_key=_unique_key("OpGetDefault")): arity: int add: ClassVar[Op] = entry(arity=2) attr = Op.def_attr("cost") assert attr.get(Op.add, -1) == -1 attr[Op.add] = 9 assert attr.get(Op.add, -1) == 9 def test_def_attr_rejects_foreign_variant() -> None: class Left(Enum, type_key=_unique_key("LeftEnum")): v: int one: ClassVar[Left] = entry(v=1) class Right(Enum, type_key=_unique_key("RightEnum")): v: int one: ClassVar[Right] = entry(v=1) attr = Left.def_attr("x", default=0) with pytest.raises(TypeError): attr[Right.one] = 1 def test_def_attr_contains() -> None: class Op(Enum, type_key=_unique_key("OpC")): arity: int add: ClassVar[Op] = entry(arity=2) neg: ClassVar[Op] = entry(arity=1) cost = Op.def_attr("cost", default=0) assert Op.add not in cost cost[Op.add] = 5 assert Op.add in cost assert Op.neg not in cost def test_def_attr_rejects_none_write() -> None: """``None`` is reserved as the column's "unset" sentinel.""" class Op(Enum, type_key=_unique_key("OpNone")): arity: int add: ClassVar[Op] = entry(arity=2) cost = Op.def_attr("cost", default=0) with pytest.raises(TypeError, match="reserved as the 'unset' sentinel"): cost[Op.add] = None assert Op.add not in cost def test_def_attr_accepts_fresh_wrapper_from_get() -> None: """Variants returned by ``Cls.get(...)`` may be fresh Python wrappers whose ``id`` differs from the cached class attribute. Ordinal-indexed lookup must still resolve correctly. """ class Op(Enum, type_key=_unique_key("OpFresh")): arity: int add: ClassVar[Op] = entry(arity=2) cost = Op.def_attr("cost", default=0) cost[Op.add] = 7 fresh = Op.get("add") assert fresh.same_as(Op.add) assert cost[fresh] == 7 assert fresh in cost # --------------------------------------------------------------------------- # `entry` sentinel sanity checks # --------------------------------------------------------------------------- def test_entry_sentinel_reprs() -> None: e = entry(1, 2, name_key="x") assert isinstance(e, _EnumEntry) assert e.args == (1, 2) assert e.kwargs == {"name_key": "x"} assert "entry(" in repr(e) def test_entry_attribute_access_outside_class_body() -> None: """A naked ``entry()`` call returns the sentinel — never a real instance.""" e = entry(output_zero=True) assert not isinstance(e, Object) # --------------------------------------------------------------------------- # TypeAttr-level verification # --------------------------------------------------------------------------- def test_enum_entries_typeattr_is_mapping() -> None: class WithAttr(Enum, type_key=_unique_key("WithAttr")): v: int one: ClassVar[WithAttr] = entry(v=1) tinfo = WithAttr.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] stored = core._lookup_type_attr(tinfo.type_index, ENUM_ENTRIES_ATTR) assert stored is not None assert "one" in stored def test_enum_attrs_typeattr_stored_under_unified_column() -> None: """``def_attr`` writes into the ``__ffi_enum_attrs__`` Dict column.""" class WithAttr(Enum, type_key=_unique_key("UnifiedAttrs")): v: int one: ClassVar[WithAttr] = entry(v=1) attr = WithAttr.def_attr("color", default="?") attr[WithAttr.one] = "red" tinfo = WithAttr.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] stored = core._lookup_type_attr(tinfo.type_index, ENUM_ATTRS_ATTR) assert stored is not None assert "color" in stored assert stored["color"][WithAttr.one._value] == "red" # --------------------------------------------------------------------------- # C++-backed enum — auto-detected when type_key is already registered # --------------------------------------------------------------------------- def test_cxx_enum_obj_get_returns_singleton() -> None: """``EnumObj::Get`` (wired as ``testing.enum_variant_get``) returns the same singleton as ``Enum.get`` and the Python-side binder. """ cxx_get = tvm_ffi.get_global_func("testing.enum_variant_get") class Variant(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Variant] Beta: ClassVar[Variant] assert cxx_get("Alpha").same_as(Variant.Alpha) assert cxx_get("Beta").same_as(Variant.Beta) with pytest.raises(RuntimeError, match="no instance named"): cxx_get("Nope") def test_cxx_backed_classvar_binds_to_existing_entries() -> None: """``ClassVar[Cls]`` (no assignment) binds to entries registered on the C++ side via ``refl::EnumDef``. ``testing.TestEnumVariant`` is registered in C++ with two entries, ``Alpha`` and ``Beta``, each with a ``code`` attr attached via ``set_attr``. The Python subclass picks these up without declaring any new entries of its own. """ class Variant(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Variant] Beta: ClassVar[Variant] assert isinstance(Variant.Alpha, Variant) assert Variant.Alpha._name == "Alpha" assert Variant.Beta._name == "Beta" # Ordinals come from C++ (registered in Alpha, Beta declaration order). assert Variant.Alpha._value == 0 assert Variant.Beta._value == 1 assert Variant.get("Alpha").same_as(Variant.Alpha) # C++-stored `code` attr is visible through attr_dict. code_col = Variant.attr_dict["code"] assert code_col[Variant.Alpha._value] == 10 assert code_col[Variant.Beta._value] == 20 def test_cxx_backed_reads_entries_typeattr() -> None: class Variant2(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Variant2] tinfo = Variant2.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] stored = core._lookup_type_attr(tinfo.type_index, ENUM_ENTRIES_ATTR) assert stored is not None assert "Alpha" in stored def test_cxx_backed_binder_typo_raises_descriptive_error() -> None: """A ``ClassVar[Cls]`` binder naming an entry that isn't in the C++ registry raises a ``RuntimeError`` that names both the typo and the known C++-registered entries, rather than falling through to the ``Enum``-base ``init=False`` guard. """ with pytest.raises(RuntimeError) as excinfo: class _Typo(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[_Typo] Neta: ClassVar[_Typo] # intended to be "Beta" msg = str(excinfo.value) assert "'Neta'" in msg assert "'testing.TestEnumVariant'" in msg assert "'Alpha'" in msg assert "'Beta'" in msg assert "C++" in msg assert "ClassVar" in msg def test_cxx_backed_mixed_entries_via_auto() -> None: """A Python subclass of a cxx-backed enum may add new Python-side entries via ``auto()`` alongside bare ``ClassVar`` binders for existing C++ entries. Ordinals for the new Python entries continue from the count of existing entries, preserving the dense-ordinal invariant across the mixed set. The ``__ffi_enum_entries__`` dict lives at the type-index level and is shared with every other Python subclass of the same ``type_key`` — so we pick unique variant names (``Mixed*``) to avoid collisions with other tests that also bind to ``testing.TestEnumVariant``. """ class Mixed(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Mixed] Beta: ClassVar[Mixed] MixedOne = auto() MixedTwo = auto() # Alpha/Beta come from C++ with ordinals 0 and 1. assert Mixed.Alpha._value == 0 assert Mixed.Beta._value == 1 assert Mixed.Alpha._name == "Alpha" assert Mixed.Beta._name == "Beta" # Python-side entries extend the dense ordinal sequence from the C++ count. assert Mixed.MixedOne._name == "MixedOne" assert Mixed.MixedTwo._name == "MixedTwo" assert Mixed.MixedOne._value == Mixed.Beta._value + 1 assert Mixed.MixedTwo._value == Mixed.Beta._value + 2 # Round-trip through ``get`` / ``all_entries``. assert Mixed.get("MixedOne").same_as(Mixed.MixedOne) assert Mixed.get("MixedTwo").same_as(Mixed.MixedTwo) assert {"Alpha", "Beta", "MixedOne", "MixedTwo"}.issubset( {entry._name for entry in Mixed.all_entries()} ) # Python-side variants are real subclass instances of the cxx-backed type. assert isinstance(Mixed.MixedOne, Mixed) assert isinstance(Mixed.MixedTwo, Mixed) # Existing C++ attrs remain unaffected; new Python variants have no attrs yet. code = Mixed.attr_dict["code"] assert code[Mixed.Alpha._value] == 10 assert code[Mixed.Beta._value] == 20 def test_cxx_backed_python_entry_accepts_def_attr() -> None: """``def_attr`` writes still work for Python-side variants on a cxx-backed enum.""" class WithPy(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[WithPy] AttrOne = auto() tag = WithPy.def_attr("tag", default=None) tag[WithPy.AttrOne] = "py-side" assert tag[WithPy.AttrOne] == "py-side" # Column was widened to the new ordinal; C++-registered entries retain default. assert tag.get(WithPy.Alpha) is None # --------------------------------------------------------------------------- # Default ReprPrint for EnumObj subclasses + MISSING/KWARGS sentinels # --------------------------------------------------------------------------- def test_default_repr_python_backed() -> None: """Python-only enum subclasses format each variant as ``.``.""" key = _unique_key("ReprPy") class Priority(Enum, type_key=key): low = auto() medium = auto() high = auto() assert repr(Priority.low) == f"{key}.low" assert repr(Priority.medium) == f"{key}.medium" assert repr(Priority.high) == f"{key}.high" def test_default_repr_cxx_backed() -> None: """C++-registered enum subclasses format with the C++ type_key.""" class Variant(Enum, type_key="testing.TestEnumVariant"): Alpha: ClassVar[Variant] Beta: ClassVar[Variant] assert repr(Variant.Alpha) == "testing.TestEnumVariant.Alpha" assert repr(Variant.Beta) == "testing.TestEnumVariant.Beta" def test_default_repr_in_nested_container() -> None: """Enum repr applies recursively when a variant is nested inside a Dict/List.""" key = _unique_key("ReprNested") class Color(Enum, type_key=key): red = auto() green = auto() all_entries_repr = repr(list(Color.all_entries())) assert f"{key}.red" in all_entries_repr assert f"{key}.green" in all_entries_repr all_entries = [repr(v) for v in Color.all_entries()] assert all_entries == [f"{key}.red", f"{key}.green"] def test_default_repr_with_attribute_carrying_variant() -> None: """Attribute-carrying entries still render with the ``.`` form.""" key = _unique_key("ReprWithAttrs") class Op(Enum, type_key=key): arity: int add: ClassVar[Op] = entry(arity=2) neg: ClassVar[Op] = entry(arity=1) assert repr(Op.add) == f"{key}.add" assert repr(Op.neg) == f"{key}.neg" def test_missing_and_kwargs_sentinel_repr() -> None: """The built-in MISSING and KWARGS singletons render with angle-bracket tags.""" assert repr(core.MISSING) == "" assert repr(core.KWARGS) == "" tvm-ffi-0.1.12/tests/python/test_dataclass_frozen.py000066400000000000000000000512461521067262500225770ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for frozen support in ``@py_class``.""" # ruff: noqa: D102 from __future__ import annotations import copy import itertools from typing import Dict, List, Optional import pytest from tvm_ffi.core import FFIProperty, Object # ty: ignore[unresolved-import] from tvm_ffi.dataclasses import field, py_class _counter = itertools.count() def _unique_key(base: str) -> str: return f"testing.frozen.{base}_{next(_counter)}" # --------------------------------------------------------------------------- # Basic frozen class # --------------------------------------------------------------------------- class TestFrozenClassBasic: def test_init_works_normally(self) -> None: @py_class(_unique_key("basic_init"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) assert p.x == 1 assert p.y == 2 def test_assignment_blocked_after_init(self) -> None: @py_class(_unique_key("basic_blocked"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) with pytest.raises(AttributeError): p.x = 10 # ty: ignore[invalid-assignment] def test_all_fields_blocked(self) -> None: @py_class(_unique_key("all_blocked"), frozen=True) class Rec(Object): a: int b: str c: float r = Rec(1, "hi", 3.0) for attr in ("a", "b", "c"): with pytest.raises(AttributeError): setattr(r, attr, None) def test_reading_fields_works(self) -> None: @py_class(_unique_key("read_ok"), frozen=True) class Pt(Object): x: int p = Pt(42) assert p.x == 42 def test_del_blocked(self) -> None: @py_class(_unique_key("del_blocked"), frozen=True) class Pt(Object): x: int p = Pt(1) with pytest.raises(AttributeError): del p.x # --------------------------------------------------------------------------- # Frozen with various field types # --------------------------------------------------------------------------- class TestFrozenFieldTypes: def test_frozen_optional_field(self) -> None: @py_class(_unique_key("opt_field"), frozen=True) class Opt(Object): v: Optional[int] # noqa: UP045 o1 = Opt(None) assert o1.v is None o2 = Opt(42) assert o2.v == 42 with pytest.raises(AttributeError): o1.v = 5 # ty: ignore[invalid-assignment] def test_frozen_object_field(self) -> None: @py_class(_unique_key("obj_inner"), frozen=True) class Inner(Object): val: int @py_class(_unique_key("obj_outer"), frozen=True) class Outer(Object): child: Inner inner = Inner(10) outer = Outer(inner) assert outer.child.val == 10 with pytest.raises(AttributeError): outer.child = Inner(20) # ty: ignore[invalid-assignment] def test_frozen_list_field(self) -> None: @py_class(_unique_key("list_field"), frozen=True) class HasList(Object): items: List[int] # noqa: UP006 h = HasList([1, 2, 3]) assert len(h.items) == 3 with pytest.raises(AttributeError): h.items = [4, 5] # ty: ignore[invalid-assignment] def test_frozen_dict_field(self) -> None: @py_class(_unique_key("dict_field"), frozen=True) class HasDict(Object): mapping: Dict[str, int] # noqa: UP006 h = HasDict({"a": 1}) assert h.mapping["a"] == 1 with pytest.raises(AttributeError): h.mapping = {"b": 2} # ty: ignore[invalid-assignment] # --------------------------------------------------------------------------- # Frozen with defaults # --------------------------------------------------------------------------- class TestFrozenDefaults: def test_frozen_field_with_default(self) -> None: @py_class(_unique_key("def_val"), frozen=True) class Cfg(Object): x: int = 10 c = Cfg() assert c.x == 10 with pytest.raises(AttributeError): c.x = 20 # ty: ignore[invalid-assignment] def test_frozen_field_with_default_factory(self) -> None: @py_class(_unique_key("def_factory"), frozen=True) class Cfg(Object): items: List[int] = field(default_factory=list) # noqa: UP006 c = Cfg() assert len(c.items) == 0 with pytest.raises(AttributeError): c.items = [1] # ty: ignore[invalid-assignment] def test_frozen_init_false_with_default(self) -> None: @py_class(_unique_key("init_false_def"), frozen=True) class Cfg(Object): x: int tag: str = field(default="default", init=False) c = Cfg(5) assert c.tag == "default" with pytest.raises(AttributeError): c.tag = "other" # ty: ignore[invalid-assignment] # --------------------------------------------------------------------------- # Per-field frozen on a non-frozen class # --------------------------------------------------------------------------- class TestPerFieldFrozen: def test_field_frozen_true_on_mutable_class(self) -> None: @py_class(_unique_key("per_field")) class Rec(Object): mutable: int immutable: int = field(frozen=True) r = Rec(1, 2) r.mutable = 10 # OK assert r.mutable == 10 with pytest.raises(AttributeError): r.immutable = 20 def test_field_frozen_true_init_works(self) -> None: @py_class(_unique_key("per_field_init")) class Rec(Object): x: int = field(frozen=True) r = Rec(42) assert r.x == 42 def test_multiple_mixed_fields(self) -> None: @py_class(_unique_key("mixed_fields")) class Rec(Object): a: int = field(frozen=True) b: int c: int = field(frozen=True) r = Rec(1, 2, 3) r.b = 20 # OK with pytest.raises(AttributeError): r.a = 10 with pytest.raises(AttributeError): r.c = 30 # --------------------------------------------------------------------------- # Escape hatch: type(obj).field.set(obj, val) # --------------------------------------------------------------------------- class TestEscapeHatch: def test_escape_hatch_sets_frozen_field(self) -> None: @py_class(_unique_key("esc_basic"), frozen=True) class Pt(Object): x: int p = Pt(1) type(p).x.set(p, 99) # ty: ignore[unresolved-attribute] assert p.x == 99 def test_escape_hatch_on_field_level_frozen(self) -> None: @py_class(_unique_key("esc_field")) class Rec(Object): val: int = field(frozen=True) r = Rec(5) type(r).val.set(r, 50) # ty: ignore[unresolved-attribute] assert r.val == 50 def test_escape_hatch_multiple_fields(self) -> None: @py_class(_unique_key("esc_multi"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) type(p).x.set(p, 10) # ty: ignore[unresolved-attribute] type(p).y.set(p, 20) # ty: ignore[unresolved-attribute] assert p.x == 10 assert p.y == 20 def test_escape_hatch_preserves_type_coercion(self) -> None: @py_class(_unique_key("esc_coerce"), frozen=True) class HasList(Object): items: List[int] # noqa: UP006 h = HasList([1]) type(h).items.set(h, [10, 20]) # ty: ignore[unresolved-attribute] assert len(h.items) == 2 def test_regular_setattr_still_blocked_after_escape_hatch(self) -> None: @py_class(_unique_key("esc_still_frozen"), frozen=True) class Pt(Object): x: int p = Pt(1) type(p).x.set(p, 99) # ty: ignore[unresolved-attribute] with pytest.raises(AttributeError): p.x = 100 # ty: ignore[invalid-assignment] def test_escape_hatch_on_mutable_field_also_works(self) -> None: @py_class(_unique_key("esc_mutable")) class Pt(Object): x: int p = Pt(1) type(p).x.set(p, 99) # ty: ignore[unresolved-attribute] assert p.x == 99 # --------------------------------------------------------------------------- # copy / deepcopy / __replace__ # --------------------------------------------------------------------------- class TestFrozenCopy: def test_copy_copy_frozen_class(self) -> None: @py_class(_unique_key("copy_basic"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) p2 = copy.copy(p) assert p2.x == 1 and p2.y == 2 assert not p.same_as(p2) def test_copy_copy_result_is_also_frozen(self) -> None: @py_class(_unique_key("copy_frozen"), frozen=True) class Pt(Object): x: int p2 = copy.copy(Pt(1)) with pytest.raises(AttributeError): p2.x = 10 # ty: ignore[invalid-assignment] def test_deepcopy_frozen_class(self) -> None: @py_class(_unique_key("deepcopy"), frozen=True) class Pt(Object): x: int p = Pt(42) p2 = copy.deepcopy(p) assert p2.x == 42 assert not p.same_as(p2) def test_deepcopy_result_is_also_frozen(self) -> None: @py_class(_unique_key("deepcopy_frozen"), frozen=True) class Pt(Object): x: int p2 = copy.deepcopy(Pt(1)) with pytest.raises(AttributeError): p2.x = 10 # ty: ignore[invalid-assignment] def test_replace_on_frozen_class(self) -> None: @py_class(_unique_key("replace"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) p2 = p.__replace__(x=10) # ty: ignore[unresolved-attribute] assert p2.x == 10 and p2.y == 2 assert p.x == 1 # original unchanged def test_replace_multiple_fields(self) -> None: @py_class(_unique_key("replace_multi"), frozen=True) class Pt(Object): x: int y: int p = Pt(1, 2) p2 = p.__replace__(x=10, y=20) # ty: ignore[unresolved-attribute] assert p2.x == 10 and p2.y == 20 def test_replace_result_is_frozen(self) -> None: @py_class(_unique_key("replace_frozen"), frozen=True) class Pt(Object): x: int p2 = Pt(1).__replace__(x=10) # ty: ignore[unresolved-attribute] with pytest.raises(AttributeError): p2.x = 99 # --------------------------------------------------------------------------- # Inheritance # --------------------------------------------------------------------------- class TestFrozenInheritance: def test_frozen_parent_mutable_child(self) -> None: @py_class(_unique_key("inh_parent_frozen"), frozen=True) class Parent(Object): a: int @py_class(_unique_key("inh_child_mutable")) class Child(Parent): # ty: ignore[invalid-frozen-dataclass-subclass] b: int c = Child(1, 2) assert c.a == 1 and c.b == 2 # Parent field stays frozen (property installed by Parent class) with pytest.raises(AttributeError): c.a = 10 # ty: ignore[invalid-assignment] # Child's own field is mutable c.b = 20 # ty: ignore[invalid-assignment] assert c.b == 20 def test_frozen_parent_frozen_child(self) -> None: @py_class(_unique_key("inh_both_frozen_p"), frozen=True) class Parent(Object): a: int @py_class(_unique_key("inh_both_frozen_c"), frozen=True) class Child(Parent): b: int c = Child(1, 2) with pytest.raises(AttributeError): c.a = 10 # ty: ignore[invalid-assignment] with pytest.raises(AttributeError): c.b = 20 # ty: ignore[invalid-assignment] def test_mutable_parent_frozen_child(self) -> None: @py_class(_unique_key("inh_parent_mutable")) class Parent(Object): a: int @py_class(_unique_key("inh_child_frozen"), frozen=True) class Child(Parent): # ty: ignore[invalid-frozen-dataclass-subclass] b: int c = Child(1, 2) # Parent field is mutable (property installed by Parent class) c.a = 10 # ty: ignore[invalid-assignment] assert c.a == 10 # Child's own field is frozen with pytest.raises(AttributeError): c.b = 20 # ty: ignore[invalid-assignment] def test_three_level_frozen_inheritance(self) -> None: @py_class(_unique_key("inh_l1"), frozen=True) class L1(Object): a: int @py_class(_unique_key("inh_l2"), frozen=True) class L2(L1): b: int @py_class(_unique_key("inh_l3"), frozen=True) class L3(L2): c: int obj = L3(1, 2, 3) assert obj.a == 1 and obj.b == 2 and obj.c == 3 for attr in ("a", "b", "c"): with pytest.raises(AttributeError): setattr(obj, attr, 99) def test_escape_hatch_on_inherited_frozen_field(self) -> None: @py_class(_unique_key("inh_esc_p"), frozen=True) class Parent(Object): a: int @py_class(_unique_key("inh_esc_c"), frozen=True) class Child(Parent): b: int c = Child(1, 2) # Escape hatch for parent field must go through Parent class descriptor Parent.a.set(c, 10) # ty: ignore[unresolved-attribute] assert c.a == 10 def test_replace_on_inherited_frozen(self) -> None: @py_class(_unique_key("inh_replace_p"), frozen=True) class Parent(Object): a: int @py_class(_unique_key("inh_replace_c"), frozen=True) class Child(Parent): b: int c = Child(1, 2) c2 = c.__replace__(a=10, b=20) # ty: ignore[unresolved-attribute] assert c2.a == 10 and c2.b == 20 # --------------------------------------------------------------------------- # kw_only + frozen # --------------------------------------------------------------------------- class TestFrozenKwOnly: def test_frozen_with_kw_only(self) -> None: @py_class(_unique_key("kw_frozen"), frozen=True, kw_only=True) class Cfg(Object): x: int y: int c = Cfg(x=1, y=2) assert c.x == 1 and c.y == 2 with pytest.raises(AttributeError): c.x = 10 # ty: ignore[invalid-assignment] # --------------------------------------------------------------------------- # __post_init__ interaction # --------------------------------------------------------------------------- class TestFrozenPostInit: def test_post_init_called(self) -> None: log: list[bool] = [] @py_class(_unique_key("post_init_called"), frozen=True) class Pt(Object): x: int def __post_init__(self) -> None: log.append(True) Pt(1) assert log == [True] def test_post_init_can_read_fields(self) -> None: captured: list[int] = [] @py_class(_unique_key("post_init_read"), frozen=True) class Pt(Object): x: int def __post_init__(self) -> None: captured.append(self.x) Pt(42) assert captured == [42] def test_post_init_cannot_set_frozen_fields(self) -> None: @py_class(_unique_key("post_init_set"), frozen=True) class Pt(Object): x: int def __post_init__(self) -> None: self.x = 999 # should fail # ty: ignore[invalid-assignment] with pytest.raises(AttributeError): Pt(1) def test_post_init_can_use_escape_hatch(self) -> None: @py_class(_unique_key("post_init_esc"), frozen=True) class Pt(Object): x: int def __post_init__(self) -> None: type(self).x.set(self, self.x * 2) # ty: ignore[unresolved-attribute] p = Pt(5) assert p.x == 10 # --------------------------------------------------------------------------- # eq / hash + frozen # --------------------------------------------------------------------------- class TestFrozenEqHash: def test_frozen_with_eq(self) -> None: @py_class(_unique_key("eq"), frozen=True, eq=True) class Pt(Object): x: int assert Pt(1) == Pt(1) assert Pt(1) != Pt(2) def test_frozen_with_hash(self) -> None: @py_class(_unique_key("hash"), frozen=True, eq=True, unsafe_hash=True) class Pt(Object): x: int assert hash(Pt(1)) == hash(Pt(1)) s = {Pt(1), Pt(2)} assert len(s) == 2 # --------------------------------------------------------------------------- # FFIProperty descriptor checks # --------------------------------------------------------------------------- class TestFFIPropertyDescriptor: def test_frozen_field_is_ffi_property(self) -> None: @py_class(_unique_key("desc_type"), frozen=True) class Pt(Object): x: int assert isinstance(Pt.__dict__["x"], FFIProperty) def test_frozen_field_fset_is_none(self) -> None: @py_class(_unique_key("desc_fset"), frozen=True) class Pt(Object): x: int assert Pt.__dict__["x"].fset is None def test_mutable_field_is_ffi_property(self) -> None: @py_class(_unique_key("desc_mutable")) class Pt(Object): x: int assert isinstance(Pt.__dict__["x"], FFIProperty) def test_mutable_field_fset_is_not_none(self) -> None: @py_class(_unique_key("desc_mutable_fset")) class Pt(Object): x: int assert Pt.__dict__["x"].fset is not None def test_mutable_field_set_method_also_works(self) -> None: @py_class(_unique_key("desc_mutable_set")) class Pt(Object): x: int p = Pt(1) type(p).x.set(p, 99) # ty: ignore[unresolved-attribute] assert p.x == 99 # --------------------------------------------------------------------------- # Edge cases # --------------------------------------------------------------------------- class TestFrozenEdgeCases: def test_frozen_class_no_own_fields(self) -> None: @py_class(_unique_key("no_fields"), frozen=True) class Empty(Object): pass Empty() # should not raise def test_frozen_single_field(self) -> None: @py_class(_unique_key("single"), frozen=True) class Single(Object): x: int s = Single(1) assert s.x == 1 with pytest.raises(AttributeError): s.x = 2 # ty: ignore[invalid-assignment] def test_frozen_field_with_none_default(self) -> None: @py_class(_unique_key("none_def"), frozen=True) class Opt(Object): v: Optional[int] = None # noqa: UP045 o = Opt() assert o.v is None with pytest.raises(AttributeError): o.v = 1 # ty: ignore[invalid-assignment] def test_multiple_instances_independent(self) -> None: @py_class(_unique_key("multi_inst"), frozen=True) class Pt(Object): x: int a = Pt(1) b = Pt(2) assert a.x == 1 assert b.x == 2 with pytest.raises(AttributeError): a.x = 99 # ty: ignore[invalid-assignment] with pytest.raises(AttributeError): b.x = 99 # ty: ignore[invalid-assignment] def test_frozen_instance_as_field_value(self) -> None: @py_class(_unique_key("inner_frozen"), frozen=True) class Inner(Object): val: int @py_class(_unique_key("outer_mut")) class Outer(Object): child: Inner inner = Inner(10) outer = Outer(inner) # Inner is still frozen even when held in a mutable outer with pytest.raises(AttributeError): outer.child.val = 99 # ty: ignore[invalid-assignment] # But the outer field itself is mutable outer.child = Inner(20) assert outer.child.val == 20 tvm-ffi-0.1.12/tests/python/test_dataclass_hash.py000066400000000000000000001002061521067262500222060ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for ffi.RecursiveHash.""" from __future__ import annotations import math import struct import time from collections.abc import Callable import numpy as np import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi._ffi_api import RecursiveEq, RecursiveHash from tvm_ffi.testing import ( TestCompare, TestCustomCompare, TestCustomHash, TestEqWithoutHash, TestHash, TestIntPair, _TestCxxClassDerived, _TestCxxClassDerivedDerived, create_object, ) def _make_nan_from_payload(payload: int) -> float: """Create a quiet NaN with a deterministic payload.""" bits = 0x7FF8000000000000 | (payload & 0x0007FFFFFFFFFFFF) return struct.unpack(">d", struct.pack(">Q", bits))[0] # --------------------------------------------------------------------------- # Primitives: int # --------------------------------------------------------------------------- def test_int_hash_equal_values() -> None: assert RecursiveHash(42) == RecursiveHash(42) def test_int_hash_different_values() -> None: assert RecursiveHash(1) != RecursiveHash(2) def test_int64_extremes_hash() -> None: i64_min = -(2**63) i64_max = 2**63 - 1 assert RecursiveHash(i64_min) == RecursiveHash(i64_min) assert RecursiveHash(i64_max) == RecursiveHash(i64_max) assert RecursiveHash(i64_min) != RecursiveHash(i64_max) # --------------------------------------------------------------------------- # Primitives: float # --------------------------------------------------------------------------- def test_float_hash_equal_values() -> None: assert RecursiveHash(3.14) == RecursiveHash(3.14) def test_float_hash_different_values() -> None: assert RecursiveHash(1.0) != RecursiveHash(2.0) def test_float_signed_zero_hash() -> None: """Both +0.0 and -0.0 hash the same (consistent with RecursiveEq).""" assert RecursiveHash(-0.0) == RecursiveHash(0.0) def test_float_infinity_hash() -> None: assert RecursiveHash(math.inf) == RecursiveHash(math.inf) assert RecursiveHash(-math.inf) == RecursiveHash(-math.inf) assert RecursiveHash(math.inf) != RecursiveHash(-math.inf) # --------------------------------------------------------------------------- # NaN handling # --------------------------------------------------------------------------- def test_nan_hash() -> None: """All NaN values hash the same (consistent with RecursiveEq).""" assert RecursiveHash(math.nan) == RecursiveHash(math.nan) def test_nan_payloads_hash_equal() -> None: nan1 = _make_nan_from_payload(0x1) nan2 = _make_nan_from_payload(0x2) assert math.isnan(nan1) and math.isnan(nan2) assert RecursiveHash(nan1) == RecursiveHash(nan2) def test_nan_payloads_in_nested_array_hash() -> None: nan1 = _make_nan_from_payload(0xA5) nan2 = _make_nan_from_payload(0x5A) a = tvm_ffi.Array([1.0, nan1, 2.0]) b = tvm_ffi.Array([1.0, nan2, 2.0]) assert RecursiveHash(a) == RecursiveHash(b) # --------------------------------------------------------------------------- # Primitives: bool # --------------------------------------------------------------------------- def test_bool_hash_equal() -> None: assert RecursiveHash(True) == RecursiveHash(True) assert RecursiveHash(False) == RecursiveHash(False) def test_bool_hash_different() -> None: assert RecursiveHash(True) != RecursiveHash(False) # --------------------------------------------------------------------------- # Primitives: string # --------------------------------------------------------------------------- def test_string_hash_equal() -> None: assert RecursiveHash("hello") == RecursiveHash("hello") def test_string_hash_different() -> None: assert RecursiveHash("hello") != RecursiveHash("world") def test_string_small_boundary_hash() -> None: small = "1234567" # SmallStr large = "12345678" # heap-backed Str assert RecursiveHash(small) == RecursiveHash("1234567") assert RecursiveHash(small) != RecursiveHash(large) def test_string_embedded_nul_hash() -> None: assert RecursiveHash("a\x00b") == RecursiveHash("a\x00b") assert RecursiveHash("a\x00b") != RecursiveHash("a\x00c") # --------------------------------------------------------------------------- # Primitives: bytes # --------------------------------------------------------------------------- def test_bytes_hash_equal() -> None: assert RecursiveHash(b"hello") == RecursiveHash(b"hello") def test_bytes_hash_different() -> None: assert RecursiveHash(b"hello") != RecursiveHash(b"world") def test_bytes_small_boundary_hash() -> None: small = b"1234567" # SmallBytes large = b"12345678" # heap-backed Bytes assert RecursiveHash(small) == RecursiveHash(b"1234567") assert RecursiveHash(small) != RecursiveHash(large) # --------------------------------------------------------------------------- # None # --------------------------------------------------------------------------- def test_none_hash() -> None: assert RecursiveHash(None) == RecursiveHash(None) def test_none_vs_other_hash() -> None: assert RecursiveHash(None) != RecursiveHash(0) # --------------------------------------------------------------------------- # DataType # --------------------------------------------------------------------------- def test_dtype_hash_equal() -> None: assert RecursiveHash(tvm_ffi.dtype("float32")) == RecursiveHash(tvm_ffi.dtype("float32")) def test_dtype_hash_different() -> None: assert RecursiveHash(tvm_ffi.dtype("float32")) != RecursiveHash(tvm_ffi.dtype("float16")) # --------------------------------------------------------------------------- # Device # --------------------------------------------------------------------------- def test_device_hash_equal() -> None: assert RecursiveHash(tvm_ffi.Device("cpu", 0)) == RecursiveHash(tvm_ffi.Device("cpu", 0)) def test_device_hash_different() -> None: assert RecursiveHash(tvm_ffi.Device("cpu", 0)) != RecursiveHash(tvm_ffi.Device("cpu", 1)) # --------------------------------------------------------------------------- # Containers: Array # --------------------------------------------------------------------------- def test_array_hash_equal() -> None: a = tvm_ffi.Array([1, 2, 3]) b = tvm_ffi.Array([1, 2, 3]) assert RecursiveHash(a) == RecursiveHash(b) def test_array_hash_different() -> None: a = tvm_ffi.Array([1, 2, 3]) c = tvm_ffi.Array([1, 2, 4]) assert RecursiveHash(a) != RecursiveHash(c) def test_array_empty_hash() -> None: empty1 = tvm_ffi.Array([]) empty2 = tvm_ffi.Array([]) assert RecursiveHash(empty1) == RecursiveHash(empty2) def test_array_different_length_hash() -> None: a = tvm_ffi.Array([1, 2]) b = tvm_ffi.Array([1, 2, 3]) assert RecursiveHash(a) != RecursiveHash(b) # --------------------------------------------------------------------------- # Containers: List # --------------------------------------------------------------------------- def test_list_hash_equal() -> None: a = tvm_ffi.List([1, 2, 3]) b = tvm_ffi.List([1, 2, 3]) assert RecursiveHash(a) == RecursiveHash(b) def test_list_hash_different() -> None: a = tvm_ffi.List([1, 2, 3]) c = tvm_ffi.List([1, 2, 4]) assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # Shape # --------------------------------------------------------------------------- def test_shape_hash_equal() -> None: a = tvm_ffi.Shape((2, 3, 4)) b = tvm_ffi.Shape((2, 3, 4)) assert RecursiveHash(a) == RecursiveHash(b) def test_shape_hash_different() -> None: a = tvm_ffi.Shape((2, 3, 4)) c = tvm_ffi.Shape((2, 3, 5)) assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # Map/Dict # --------------------------------------------------------------------------- def test_map_hash_equal() -> None: a = tvm_ffi.Map({"x": 1, "y": 2}) b = tvm_ffi.Map({"x": 1, "y": 2}) assert RecursiveHash(a) == RecursiveHash(b) def test_map_hash_different_values() -> None: a = tvm_ffi.Map({"x": 1, "y": 2}) c = tvm_ffi.Map({"x": 1, "y": 3}) assert RecursiveHash(a) != RecursiveHash(c) def test_map_hash_different_size() -> None: a = tvm_ffi.Map({"x": 1}) b = tvm_ffi.Map({"x": 1, "y": 2}) assert RecursiveHash(a) != RecursiveHash(b) def test_map_hash_order_independent() -> None: """Map hash should be the same regardless of insertion order.""" a = tvm_ffi.Map({"x": 1, "y": 2, "z": 3}) b = tvm_ffi.Map({"z": 3, "x": 1, "y": 2}) assert RecursiveHash(a) == RecursiveHash(b) def test_dict_hash_equal() -> None: a = tvm_ffi.Dict({"x": 1, "y": 2}) b = tvm_ffi.Dict({"x": 1, "y": 2}) assert RecursiveHash(a) == RecursiveHash(b) def test_dict_hash_different() -> None: a = tvm_ffi.Dict({"x": 1}) b = tvm_ffi.Dict({"x": 2}) assert RecursiveHash(a) != RecursiveHash(b) # --------------------------------------------------------------------------- # Reflected objects: TestIntPair # --------------------------------------------------------------------------- def test_reflected_obj_hash_equal() -> None: a = TestIntPair(1, 2) b = TestIntPair(1, 2) assert RecursiveHash(a) == RecursiveHash(b) def test_reflected_obj_hash_different() -> None: a = TestIntPair(1, 2) c = TestIntPair(1, 3) assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # HashOff flag: TestHash # --------------------------------------------------------------------------- def test_hash_off_ignored_field() -> None: """hash_ignored is excluded from hashing via Hash(false).""" a = TestHash(1, "x", 100) b = TestHash(1, "x", 999) assert RecursiveHash(a) == RecursiveHash(b) def test_hash_off_key_differs() -> None: a = TestHash(1, "x", 100) b = TestHash(2, "x", 100) assert RecursiveHash(a) != RecursiveHash(b) def test_hash_off_name_differs() -> None: a = TestHash(1, "a", 100) b = TestHash(1, "b", 100) assert RecursiveHash(a) != RecursiveHash(b) # --------------------------------------------------------------------------- # CompareOff implies hash-off: TestCompare # --------------------------------------------------------------------------- def test_compare_off_implies_hash_off() -> None: """Fields with Compare(false) are also excluded from hashing.""" a = TestCompare(1, "x", 100) b = TestCompare(1, "x", 999) assert RecursiveHash(a) == RecursiveHash(b) # --------------------------------------------------------------------------- # Same pointer fast path # --------------------------------------------------------------------------- def test_same_pointer_hash() -> None: x = TestIntPair(42, 99) assert RecursiveHash(x) == RecursiveHash(x) # --------------------------------------------------------------------------- # Different types produce different hashes # --------------------------------------------------------------------------- def test_different_types_hash() -> None: """Different type indices should generally produce different hashes.""" assert RecursiveHash(1) != RecursiveHash(1.0) assert RecursiveHash(1) != RecursiveHash(True) # --------------------------------------------------------------------------- # Nested containers # --------------------------------------------------------------------------- def test_array_of_arrays_hash() -> None: a = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 4])]) b = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 4])]) c = tvm_ffi.Array([tvm_ffi.Array([1, 2]), tvm_ffi.Array([3, 5])]) assert RecursiveHash(a) == RecursiveHash(b) assert RecursiveHash(a) != RecursiveHash(c) def test_list_of_lists_hash() -> None: a = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([3])]) b = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([3])]) c = tvm_ffi.List([tvm_ffi.List([1, 2]), tvm_ffi.List([4])]) assert RecursiveHash(a) == RecursiveHash(b) assert RecursiveHash(a) != RecursiveHash(c) def test_three_level_nested_hash() -> None: a = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([1])])]) b = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([1])])]) c = tvm_ffi.Array([tvm_ffi.Array([tvm_ffi.Array([2])])]) assert RecursiveHash(a) == RecursiveHash(b) assert RecursiveHash(a) != RecursiveHash(c) def test_map_with_array_values_hash() -> None: a = tvm_ffi.Map({"k": tvm_ffi.Array([1, 2])}) b = tvm_ffi.Map({"k": tvm_ffi.Array([1, 2])}) c = tvm_ffi.Map({"k": tvm_ffi.Array([1, 3])}) assert RecursiveHash(a) == RecursiveHash(b) assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # Inherited field hashing # --------------------------------------------------------------------------- def test_inherited_fields_hash_equal() -> None: a = _TestCxxClassDerived(10, 20, 1.5, 2.5) b = _TestCxxClassDerived(10, 20, 1.5, 2.5) assert RecursiveHash(a) == RecursiveHash(b) def test_inherited_fields_differ_in_base_hash() -> None: a = _TestCxxClassDerived(10, 20, 1.5, 2.5) b = _TestCxxClassDerived(99, 20, 1.5, 2.5) assert RecursiveHash(a) != RecursiveHash(b) def test_three_level_inheritance_hash() -> None: a = _TestCxxClassDerivedDerived(1, 2, 3.0, True, 4.0, "hi") b = _TestCxxClassDerivedDerived(1, 2, 3.0, True, 4.0, "hi") assert RecursiveHash(a) == RecursiveHash(b) c = _TestCxxClassDerivedDerived(1, 2, 3.0, False, 4.0, "hi") assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # Objects with container fields # --------------------------------------------------------------------------- def test_object_with_array_field_hash() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20, 30]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20, 30]), ) assert RecursiveHash(a) == RecursiveHash(b) def test_object_with_array_field_differ_hash() -> None: a = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 20]), ) b = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="s", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([10, 21]), ) assert RecursiveHash(a) != RecursiveHash(b) # --------------------------------------------------------------------------- # Array/List type mismatch: same content, different types # --------------------------------------------------------------------------- def test_array_vs_list_different_hash() -> None: arr = tvm_ffi.Array([1, 2]) lst = tvm_ffi.List([1, 2]) assert RecursiveHash(arr) != RecursiveHash(lst) def test_map_vs_dict_different_hash() -> None: m = tvm_ffi.Map({"k": 1}) d = tvm_ffi.Dict({"k": 1}) assert RecursiveHash(m) != RecursiveHash(d) # --------------------------------------------------------------------------- # None in nested contexts # --------------------------------------------------------------------------- def test_array_with_none_elements_hash() -> None: a = tvm_ffi.Array([None, 1, None]) b = tvm_ffi.Array([None, 1, None]) c = tvm_ffi.Array([None, 2, None]) assert RecursiveHash(a) == RecursiveHash(b) assert RecursiveHash(a) != RecursiveHash(c) # --------------------------------------------------------------------------- # Cycle safety: run in subprocess # --------------------------------------------------------------------------- def test_cyclic_list_same_pointer_hash() -> None: """Same cyclic list hashed with itself succeeds (pointer identity short-circuit).""" lst = tvm_ffi.List() lst.append(lst) # Should not raise; same pointer returns a deterministic hash h = RecursiveHash(lst) assert h == RecursiveHash(lst) def test_cyclic_list_handled() -> None: """Distinct cyclic lists are handled gracefully via on-stack cycle detection.""" a = tvm_ffi.List() a.append(a) b = tvm_ffi.List() b.append(b) h_a = RecursiveHash(a) h_b = RecursiveHash(b) assert h_a == h_b def test_cyclic_dict_handled() -> None: """Cyclic dict is handled gracefully via on-stack cycle detection.""" d = tvm_ffi.Dict() d["self"] = d h = RecursiveHash(d) assert h == RecursiveHash(d) # --------------------------------------------------------------------------- # Consistency law: RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b) # --------------------------------------------------------------------------- def test_consistency_primitives() -> None: """Verify hash consistency for various primitive pairs.""" pairs = [ (42, 42), (3.14, 3.14), (True, True), (False, False), ("hello", "hello"), (b"hello", b"hello"), (None, None), ] for a, b in pairs: assert RecursiveEq(a, b), f"Expected RecursiveEq({a!r}, {b!r})" assert RecursiveHash(a) == RecursiveHash(b), ( f"Hash mismatch for equal values: {a!r} and {b!r}" ) def test_consistency_nan() -> None: nan1 = _make_nan_from_payload(0x1) nan2 = _make_nan_from_payload(0x2) assert RecursiveEq(nan1, nan2) assert RecursiveHash(nan1) == RecursiveHash(nan2) def test_consistency_signed_zero() -> None: assert RecursiveEq(-0.0, 0.0) assert RecursiveHash(-0.0) == RecursiveHash(0.0) def test_consistency_containers() -> None: """Verify hash consistency for containers.""" pairs = [ (tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([1, 2, 3])), (tvm_ffi.List([1, 2, 3]), tvm_ffi.List([1, 2, 3])), (tvm_ffi.Shape((2, 3, 4)), tvm_ffi.Shape((2, 3, 4))), (tvm_ffi.Map({"x": 1, "y": 2}), tvm_ffi.Map({"x": 1, "y": 2})), (tvm_ffi.Dict({"x": 1, "y": 2}), tvm_ffi.Dict({"x": 1, "y": 2})), ] for a, b in pairs: assert RecursiveEq(a, b), f"Expected RecursiveEq for {a!r} and {b!r}" assert RecursiveHash(a) == RecursiveHash(b), "Hash mismatch for equal containers" def test_consistency_reflected_objects() -> None: """Verify hash consistency for reflected objects.""" a = TestIntPair(1, 2) b = TestIntPair(1, 2) assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) def test_consistency_compare_off() -> None: """Fields excluded from comparison are also excluded from hash.""" a = TestCompare(1, "x", 100) b = TestCompare(1, "x", 999) assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) def test_consistency_hash_off() -> None: """Fields excluded from hashing produce same hash when they differ.""" a = TestHash(1, "x", 100) b = TestHash(1, "x", 999) assert RecursiveHash(a) == RecursiveHash(b) def test_consistency_law_on_int_pairs() -> None: """Verify: RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b).""" values = [ TestIntPair(0, 0), TestIntPair(0, 1), TestIntPair(1, 0), TestIntPair(1, 1), ] for a in values: for b in values: if RecursiveEq(a, b): assert RecursiveHash(a) == RecursiveHash(b), ( f"Hash consistency violated: RecursiveEq({a},{b}) but hashes differ" ) def _make_nested_singleton_array(depth: int) -> object: value: object = 0 for _ in range(depth): value = tvm_ffi.Array([value]) return value # --------------------------------------------------------------------------- # Shared-reference aliasing invariants # --------------------------------------------------------------------------- def test_aliasing_consistency_array_of_reflected_objects() -> None: shared = TestIntPair(11, 22) aliased = tvm_ffi.Array([shared, shared]) duplicated = tvm_ffi.Array( [ TestIntPair(11, 22), TestIntPair(11, 22), ] ) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_list_of_reflected_objects() -> None: shared = TestIntPair(13, 26) aliased = tvm_ffi.List([shared, shared]) duplicated = tvm_ffi.List( [ TestIntPair(13, 26), TestIntPair(13, 26), ] ) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_array_of_arrays() -> None: shared = tvm_ffi.Array([1, 2, 3]) aliased = tvm_ffi.Array([shared, shared]) duplicated = tvm_ffi.Array([tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([1, 2, 3])]) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_list_of_lists() -> None: shared = tvm_ffi.List([1, 2, 3]) aliased = tvm_ffi.List([shared, shared]) duplicated = tvm_ffi.List([tvm_ffi.List([1, 2, 3]), tvm_ffi.List([1, 2, 3])]) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_shape_objects() -> None: shared = tvm_ffi.Shape((3, 4)) aliased = tvm_ffi.Array([shared, shared]) duplicated = tvm_ffi.Array([tvm_ffi.Shape((3, 4)), tvm_ffi.Shape((3, 4))]) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_map_shared_values() -> None: shared = TestIntPair(31, 41) aliased = tvm_ffi.Map({"x": shared, "y": shared}) duplicated = tvm_ffi.Map( { "x": TestIntPair(31, 41), "y": TestIntPair(31, 41), } ) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_dict_shared_values() -> None: shared = tvm_ffi.Array([7, 8, 9]) aliased = tvm_ffi.Dict({"x": shared, "y": shared}) duplicated = tvm_ffi.Dict({"x": tvm_ffi.Array([7, 8, 9]), "y": tvm_ffi.Array([7, 8, 9])}) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) def test_aliasing_consistency_reflected_object_fields() -> None: shared = TestIntPair(5, 6) aliased = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="shared", v_map=tvm_ffi.Map({"k": shared}), v_array=tvm_ffi.Array([shared]), ) duplicated = create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="shared", v_map=tvm_ffi.Map({"k": TestIntPair(5, 6)}), v_array=tvm_ffi.Array([TestIntPair(5, 6)]), ) assert RecursiveEq(aliased, duplicated) assert RecursiveHash(aliased) == RecursiveHash(duplicated) # --------------------------------------------------------------------------- # Map/Dict order-independence with shared references # --------------------------------------------------------------------------- def test_map_hash_order_independent_with_shared_values() -> None: shared = TestIntPair(1, 2) a = tvm_ffi.Map({"a": shared, "b": shared, "c": shared}) b = tvm_ffi.Map({"b": shared, "a": shared, "c": shared}) assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) def test_dict_hash_order_independent_with_shared_values() -> None: shared = tvm_ffi.Array([1, 2]) a = tvm_ffi.Dict({"a": shared, "b": shared, "c": shared}) b = tvm_ffi.Dict({"b": shared, "a": shared, "c": shared}) assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) # --------------------------------------------------------------------------- # Recursion-depth boundary # --------------------------------------------------------------------------- def test_depth_127_nested_arrays_allowed() -> None: RecursiveHash(_make_nested_singleton_array(127)) def test_depth_1000_nested_arrays_works() -> None: """Deep graphs now succeed thanks to iterative (heap-based) stack.""" RecursiveHash(_make_nested_singleton_array(1000)) # --------------------------------------------------------------------------- # Additional adversarial quality checks # --------------------------------------------------------------------------- def test_map_hash_collision_swap_values() -> None: """Swapping values across two keys should not trivially collide.""" a = tvm_ffi.Map({"a": 0, "b": 1}) b = tvm_ffi.Map({"a": 1, "b": 0}) assert not RecursiveEq(a, b) assert RecursiveHash(a) != RecursiveHash(b) def test_array_hash_collision_small_int_pairs() -> None: """Distinct short arrays should not have obvious low-domain collisions.""" a = tvm_ffi.Array([1, 2]) b = tvm_ffi.Array([2, 61]) assert not RecursiveEq(a, b) assert RecursiveHash(a) != RecursiveHash(b) def test_function_hash_consistent_with_eq() -> None: """Functions have no reflected fields, so RecursiveEq treats all as equal. Hash must be consistent: RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b). """ f_add_one = tvm_ffi.get_global_func("testing.add_one") f_nop = tvm_ffi.get_global_func("testing.nop") assert RecursiveEq(f_add_one, f_nop) assert RecursiveHash(f_add_one) == RecursiveHash(f_nop) def test_tensor_hash_consistent_with_eq() -> None: """Tensors have no reflected fields, so RecursiveEq treats all as equal. Hash must be consistent: RecursiveEq(a, b) => RecursiveHash(a) == RecursiveHash(b). """ t1 = tvm_ffi.from_dlpack(np.array([1, 2], dtype="int32")) t2 = tvm_ffi.from_dlpack(np.array([[9, 8, 7]], dtype="int64")) assert RecursiveEq(t1, t2) assert RecursiveHash(t1) == RecursiveHash(t2) def _make_shared_binary_dag(depth: int) -> object: node: object = 1 for _ in range(depth): # Two edges point to the same child object (DAG with heavy sharing). node = tvm_ffi.Array([node, node]) return node def test_shared_dag_hash_scaling_not_exponential() -> None: """Hashing shared DAGs should scale roughly linearly in depth.""" d18 = _make_shared_binary_dag(18) d19 = _make_shared_binary_dag(19) # Warm-up run to mitigate one-time setup costs RecursiveHash(_make_shared_binary_dag(10)) repeats = 10 t0 = time.perf_counter() for _ in range(repeats): RecursiveHash(d18) t18 = (time.perf_counter() - t0) / repeats t0 = time.perf_counter() for _ in range(repeats): RecursiveHash(d19) t19 = (time.perf_counter() - t0) / repeats # With memoization this ratio should stay close to 1x; 1.6x leaves buffer for noise. assert t19 <= t18 * 2.0, f"Unexpected super-linear scaling: d18={t18:.6f}s d19={t19:.6f}s" # --------------------------------------------------------------------------- # Custom __ffi_hash__ hook: TestCustomHash # --------------------------------------------------------------------------- def test_custom_hash_ignores_label() -> None: """TestCustomHash hashes only `key`, ignoring `label`.""" a = TestCustomHash(42, "alpha") b = TestCustomHash(42, "beta") assert RecursiveHash(a) == RecursiveHash(b) def test_custom_hash_different_key() -> None: a = TestCustomHash(1, "same") b = TestCustomHash(2, "same") assert RecursiveHash(a) != RecursiveHash(b) def test_custom_hash_in_container() -> None: """Custom-hooked objects inside an Array.""" a = tvm_ffi.Array( [ TestCustomHash(1, "x"), TestCustomHash(2, "y"), ] ) b = tvm_ffi.Array( [ TestCustomHash(1, "different"), TestCustomHash(2, "labels"), ] ) assert RecursiveHash(a) == RecursiveHash(b) def test_custom_hash_consistency_with_eq() -> None: """RecursiveEq(a,b) => RecursiveHash(a)==RecursiveHash(b) for TestCustomCompare.""" a = TestCustomCompare(42, "alpha") b = TestCustomCompare(42, "beta") assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) # --------------------------------------------------------------------------- # Failing regression tests for Eq=>Hash invariant # --------------------------------------------------------------------------- @pytest.mark.parametrize("key", [-7, -1, 0, 1, 7, 1024]) @pytest.mark.parametrize("lhs_label,rhs_label", [("alpha", "beta"), ("x", "y"), ("foo", "bar")]) def test_custom_compare_eq_implies_hash_same_direct( key: int, lhs_label: str, rhs_label: str ) -> None: lhs = TestCustomCompare(key, lhs_label) rhs = TestCustomCompare(key, rhs_label) assert RecursiveEq(lhs, rhs) assert RecursiveHash(lhs) == RecursiveHash(rhs) @pytest.mark.parametrize("key", [-3, -1, 0, 2, 11]) @pytest.mark.parametrize( "wrap", [ lambda obj: tvm_ffi.Array([obj]), lambda obj: tvm_ffi.List([obj]), lambda obj: tvm_ffi.Array([0, obj, 1]), lambda obj: tvm_ffi.List([0, obj, 1]), lambda obj: tvm_ffi.Map({"k": obj}), lambda obj: tvm_ffi.Dict({"k": obj}), lambda obj: tvm_ffi.Array([tvm_ffi.Array([obj])]), lambda obj: tvm_ffi.List([tvm_ffi.List([obj])]), ], ) def test_custom_compare_eq_implies_hash_same_in_wrappers( key: int, wrap: Callable[[object], object] ) -> None: lhs_obj = TestCustomCompare(key, "left") rhs_obj = TestCustomCompare(key, "right") lhs = wrap(lhs_obj) rhs = wrap(rhs_obj) assert RecursiveEq(lhs, rhs) assert RecursiveHash(lhs) == RecursiveHash(rhs) # --------------------------------------------------------------------------- # Guard: __ffi_eq__ without __ffi_hash__ must raise in RecursiveHash # --------------------------------------------------------------------------- def test_eq_without_hash_raises() -> None: """RecursiveHash rejects types that define __ffi_eq__ but not __ffi_hash__.""" obj = TestEqWithoutHash(1, "hello") with pytest.raises(ValueError, match="__ffi_eq__ or __ffi_compare__ but not __ffi_hash__"): RecursiveHash(obj) def test_eq_without_hash_inside_container_raises() -> None: """The guard also triggers when the object is nested inside a container.""" obj = TestEqWithoutHash(1, "hello") arr = tvm_ffi.Array([obj]) with pytest.raises(ValueError, match="__ffi_eq__ or __ffi_compare__ but not __ffi_hash__"): RecursiveHash(arr) # --------------------------------------------------------------------------- # Custom __ffi_hash__ hook via @py_class # --------------------------------------------------------------------------- import itertools as _itertools_hash from typing import Any as _Any_hash from typing import Callable as _Callable_hash from tvm_ffi.core import Object as _Object_hash from tvm_ffi.dataclasses import py_class as _py_class_hash _counter_hash = _itertools_hash.count() def _unique_key_hash(base: str) -> str: return f"testing.hash_pc.{base}_{next(_counter_hash)}" @_py_class_hash(_unique_key_hash("PyCustomHash")) class _PyCustomHash(_Object_hash): key: int label: str def __ffi_hash__(self, fn_hash: _Callable_hash[..., _Any_hash]) -> int: return fn_hash(self.key) def test_py_class_custom_hash_ignores_label() -> None: a = _PyCustomHash(42, "alpha") b = _PyCustomHash(42, "beta") assert RecursiveHash(a) == RecursiveHash(b) def test_py_class_custom_hash_different_key() -> None: a = _PyCustomHash(1, "same") b = _PyCustomHash(2, "same") assert RecursiveHash(a) != RecursiveHash(b) tvm-ffi-0.1.12/tests/python/test_dataclass_init.py000066400000000000000000001347001521067262500222340ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Comprehensive tests for reflection-driven auto-generated ``__ffi_init__``. This file exercises: 1. metadata emitted by C++ for auto-init traits 2. Python ``__init__`` signature generation 3. constructor behavior across positional/kw-only/default/init=False combinations 4. low-level KWARGS protocol via TypeAttrColumn __ffi_init__ lookup 5. inheritance behavior for auto-generated init 6. copy/deepcopy/replace interplay with auto-init objects 7. re-initialization, isinstance checks, and instance isolation """ # ruff: noqa: D102 from __future__ import annotations import copy import inspect from typing import Any import pytest from tvm_ffi import core from tvm_ffi.dataclasses import py_class from tvm_ffi.testing import ( TestCompare, TestHash, TestIntPair, _TestCxxAutoInit, _TestCxxAutoInitAllInitOff, _TestCxxAutoInitChild, _TestCxxAutoInitKwOnlyDefaults, _TestCxxAutoInitParent, _TestCxxAutoInitSimple, _TestCxxClassDerived, _TestCxxClassDerivedDerived, _TestCxxNoAutoInit, ) from tvm_ffi.testing.testing import requires_py313 def _ffi_init(obj: Any, *args: Any) -> None: """Look up __ffi_init__ from the TypeAttrColumn and call it via __init_handle_by_constructor__.""" type_index = type(obj).__tvm_ffi_type_info__.type_index ffi_init = core._lookup_type_attr(type_index, "__ffi_init__") assert ffi_init is not None, f"No __ffi_init__ TypeAttrColumn entry for {type(obj)}" obj.__init_handle_by_constructor__(ffi_init, *args) def _field_map(type_cls: type) -> dict[str, Any]: return {field.name: field for field in getattr(type_cls, "__tvm_ffi_type_info__").fields} def _method_metadata(type_cls: type, method_name: str) -> dict[str, Any]: type_info = getattr(type_cls, "__tvm_ffi_type_info__") for method in type_info.methods: if method.name == method_name: return method.metadata raise AssertionError(f"Cannot find method metadata: {type_cls.__name__}.{method_name}") class TestAutoInitMetadata: def test_auto_init_method_marked(self) -> None: type_info = getattr(_TestCxxAutoInit, "__tvm_ffi_type_info__") ffi_init = core._lookup_type_attr(type_info.type_index, "__ffi_init__") assert ffi_init is not None # Auto-generated init is only in TypeAttrColumn, not in the method list assert not any(m.name == "__ffi_init__" for m in type_info.methods) def test_field_bitmask_init_kw_only_and_defaults(self) -> None: fields = _field_map(_TestCxxAutoInit) assert fields["a"].c_init is True assert fields["b"].c_init is False assert fields["b"].c_has_default is True assert fields["c"].c_kw_only is True assert fields["c"].c_init is True assert fields["d"].c_has_default is True def test_all_init_off_field_bitmask(self) -> None: fields = _field_map(_TestCxxAutoInitAllInitOff) assert fields["x"].c_init is False assert fields["x"].c_has_default is True assert fields["y"].c_init is False assert fields["y"].c_has_default is True assert fields["z"].c_init is False assert fields["z"].c_has_default is False def test_kw_only_defaults_field_bitmask(self) -> None: fields = _field_map(_TestCxxAutoInitKwOnlyDefaults) assert fields["p_default"].c_has_default is True assert fields["k_required"].c_kw_only is True assert fields["k_default"].c_kw_only is True assert fields["k_default"].c_has_default is True assert fields["hidden"].c_init is False assert fields["hidden"].c_has_default is True class TestAutoInitSignature: def test_auto_init_signature_layout(self) -> None: sig = inspect.signature(_TestCxxAutoInit.__init__) assert tuple(sig.parameters) == ("self", "a", "d", "c") assert sig.parameters["a"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD assert sig.parameters["d"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD assert sig.parameters["c"].kind == inspect.Parameter.KEYWORD_ONLY def test_auto_init_signature_required_vs_default(self) -> None: sig = inspect.signature(_TestCxxAutoInit.__init__) assert sig.parameters["a"].default is inspect.Parameter.empty assert sig.parameters["c"].default is inspect.Parameter.empty assert sig.parameters["d"].default is not inspect.Parameter.empty def test_simple_signature(self) -> None: sig = inspect.signature(_TestCxxAutoInitSimple.__init__) assert tuple(sig.parameters) == ("self", "x", "y") assert all( p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD for p in (sig.parameters["x"], sig.parameters["y"]) ) def test_all_init_off_signature_is_no_arg(self) -> None: sig = inspect.signature(_TestCxxAutoInitAllInitOff.__init__) assert tuple(sig.parameters) == ("self",) def test_kw_only_defaults_signature_layout(self) -> None: sig = inspect.signature(_TestCxxAutoInitKwOnlyDefaults.__init__) assert tuple(sig.parameters) == ( "self", "p_required", "p_default", "k_required", "k_default", ) assert sig.parameters["p_required"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD assert sig.parameters["p_default"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD assert sig.parameters["k_required"].kind == inspect.Parameter.KEYWORD_ONLY assert sig.parameters["k_default"].kind == inspect.Parameter.KEYWORD_ONLY assert sig.parameters["p_required"].default is inspect.Parameter.empty assert sig.parameters["k_required"].default is inspect.Parameter.empty assert sig.parameters["p_default"].default is not inspect.Parameter.empty assert sig.parameters["k_default"].default is not inspect.Parameter.empty def test_inheritance_signature_includes_parent_fields(self) -> None: sig = inspect.signature(_TestCxxAutoInitChild.__init__) # Required positional params must precede default ones in Python, # so child_required comes before parent_default. assert tuple(sig.parameters) == ( "self", "parent_required", "child_required", "parent_default", "child_kw_only", ) assert sig.parameters["child_kw_only"].kind == inspect.Parameter.KEYWORD_ONLY assert sig.parameters["parent_default"].default is not inspect.Parameter.empty def test_init_false_field_not_in_signature(self) -> None: """B is init=False and should not appear in signature.""" sig = inspect.signature(_TestCxxAutoInit.__init__) assert "b" not in sig.parameters def test_hidden_field_not_in_signature(self) -> None: """Hidden is init=False and should not appear in signature.""" sig = inspect.signature(_TestCxxAutoInitKwOnlyDefaults.__init__) assert "hidden" not in sig.parameters def test_kw_only_params_after_positional(self) -> None: """Keyword-only params should come after positional in the signature.""" sig = inspect.signature(_TestCxxAutoInit.__init__) params = list(sig.parameters.values()) saw_kw_only = False for p in params[1:]: # skip self if p.kind == inspect.Parameter.KEYWORD_ONLY: saw_kw_only = True elif saw_kw_only: assert p.kind == inspect.Parameter.KEYWORD_ONLY, ( f"Parameter {p.name} is {p.kind} after a KEYWORD_ONLY param" ) def test_child_signature_required_before_optional(self) -> None: """The child signature should have all required positional before optional.""" sig = inspect.signature(_TestCxxAutoInitChild.__init__) params = list(sig.parameters.values())[1:] # skip self positional = [p for p in params if p.kind != inspect.Parameter.KEYWORD_ONLY] saw_default = False for p in positional: if p.default is not inspect.Parameter.empty: saw_default = True elif saw_default: pytest.fail(f"Required param '{p.name}' appears after a default param") class TestAutoInitConstruction: def test_auto_init_minimal_required(self) -> None: obj = _TestCxxAutoInit(1, c=3) assert obj.a == 1 assert obj.b == 42 assert obj.c == 3 assert obj.d == 99 def test_auto_init_all_keyword_arguments(self) -> None: obj = _TestCxxAutoInit(a=10, c=30, d=20) assert obj.a == 10 assert obj.b == 42 assert obj.c == 30 assert obj.d == 20 def test_auto_init_keyword_order_irrelevant(self) -> None: obj = _TestCxxAutoInit(c=7, d=8, a=9) assert obj.a == 9 assert obj.c == 7 assert obj.d == 8 assert obj.b == 42 def test_auto_init_second_positional_maps_to_d(self) -> None: obj = _TestCxxAutoInit(1, 2, c=3) assert obj.a == 1 assert obj.d == 2 assert obj.c == 3 assert obj.b == 42 def test_mutate_fields_after_construction(self) -> None: obj = _TestCxxAutoInit(1, c=2) obj.b = 100 obj.c = 999 assert obj.b == 100 assert obj.c == 999 class TestAutoInitErrors: def test_missing_required_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(c=3) # ty: ignore[missing-argument] def test_missing_required_kw_only(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1) # ty: ignore[missing-argument] def test_kw_only_rejects_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, 2, 3) # ty: ignore[missing-argument, too-many-positional-arguments] def test_duplicate_argument_detection(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, 2, c=3, d=4) # ty: ignore[parameter-already-assigned] def test_unexpected_keyword(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, c=2, nope=3) # ty: ignore[unknown-argument] def test_init_false_field_rejected_in_python_init(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, b=2, c=3) # ty: ignore[unknown-argument] def test_type_mismatch_for_required_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit("x", c=2) # ty: ignore[invalid-argument-type] def test_type_mismatch_for_kw_only(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, c="x") # ty: ignore[invalid-argument-type] def test_type_mismatch_for_defaultable_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInit(1, d="x", c=3) # ty: ignore[invalid-argument-type] def test_init_false_field_rejected_via_dict_unpacking(self) -> None: """B is init=False, rejected even when passed via **dict.""" with pytest.raises(TypeError): _TestCxxAutoInit(**{"a": 1, "c": 2, "b": 3}) def test_positional_and_keyword_same_field(self) -> None: """Providing a field both positionally and as keyword should error.""" with pytest.raises(TypeError): _TestCxxAutoInit(1, a=2, c=3) # ty: ignore[parameter-already-assigned] def test_none_for_required_integer_field(self) -> None: """Passing None where an int64_t is expected should raise TypeError.""" with pytest.raises(TypeError): _TestCxxAutoInitSimple(None, 1) # ty: ignore[invalid-argument-type] def test_none_for_keyword_integer_field(self) -> None: """Passing None as keyword where int64_t is expected.""" with pytest.raises(TypeError): _TestCxxAutoInitSimple(x=None, y=1) # ty: ignore[invalid-argument-type] class TestAutoInitLowLevelFfiInit: def test_low_level_kwargs_protocol(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) _ffi_init(obj, core.KWARGS, "a", 1, "c", 3) assert obj.a == 1 assert obj.b == 42 assert obj.c == 3 assert obj.d == 99 def test_low_level_kwargs_even_pairs_required(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, "a", 1, "c") def test_low_level_kwargs_duplicate_name(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, "a", 1, "a", 2, "c", 3) def test_low_level_kwargs_unknown_name(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, "a", 1, "unknown", 2, "c", 3) def test_low_level_kwargs_key_must_be_string(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, 1, 2, "a", 3, "c", 4) def test_low_level_positional_too_many(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, 1, 2, 3, 4) def test_low_level_missing_required(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, 1) def test_low_level_kw_only_field_rejected_positionally(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, 1, 2) def test_low_level_init_false_field_rejected(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, "a", 1, "b", 2, "c", 3) def test_low_level_kwargs_all_init_fields_explicit(self) -> None: """Providing all init=True fields via KWARGS.""" obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) _ffi_init(obj, core.KWARGS, "a", 10, "c", 30, "d", 40) assert obj.a == 10 assert obj.b == 42 # init=False, default assert obj.c == 30 assert obj.d == 40 def test_low_level_kwargs_empty_string_key(self) -> None: """Empty string as keyword name should be rejected as unknown.""" obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError, match="unexpected keyword"): _ffi_init(obj, core.KWARGS, "", 1, "a", 2, "c", 3) def test_low_level_kwargs_odd_kv_count(self) -> None: """Odd number of key-value args after KWARGS sentinel.""" obj = _TestCxxAutoInitSimple.__new__(_TestCxxAutoInitSimple) with pytest.raises(TypeError): _ffi_init(obj, core.KWARGS, "x") def test_low_level_positional_only_simple(self) -> None: """Positional-only mode (no KWARGS sentinel).""" obj = _TestCxxAutoInitSimple.__new__(_TestCxxAutoInitSimple) _ffi_init(obj, 10, 20) assert obj.x == 10 assert obj.y == 20 def test_low_level_zero_args_missing_required(self) -> None: """Zero args for a type that requires them.""" obj = _TestCxxAutoInitSimple.__new__(_TestCxxAutoInitSimple) with pytest.raises(TypeError, match="missing required"): _ffi_init(obj) def test_low_level_kwargs_sentinel_only_no_kv_pairs(self) -> None: """KWARGS sentinel with zero key-value pairs; required fields still missing.""" obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError, match="missing required"): _ffi_init(obj, core.KWARGS) def test_low_level_kwargs_positional_then_sentinel_no_kv(self) -> None: """Positional args followed by KWARGS sentinel but no key-value pairs.""" obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) # a=1 positionally, sentinel, no KV pairs → c is still missing with pytest.raises(TypeError, match="missing required"): _ffi_init(obj, 1, core.KWARGS) def test_low_level_child_positional_routing(self) -> None: """Verify the inheritance positional fix at the raw __ffi_init__ level. After stable_partition, pos_indices should be: [parent_required, child_required, parent_default] So 2 positional args map to parent_required=1, child_required=2. """ obj = _TestCxxAutoInitChild.__new__(_TestCxxAutoInitChild) _ffi_init(obj, 1, 2, core.KWARGS, "child_kw_only", 3) assert obj.parent_required == 1 assert obj.child_required == 2 assert obj.parent_default == 5 # default assert obj.child_kw_only == 3 def test_low_level_child_all_three_positional(self) -> None: """Three positional args for child at the raw protocol level.""" obj = _TestCxxAutoInitChild.__new__(_TestCxxAutoInitChild) _ffi_init(obj, 1, 2, 3, core.KWARGS, "child_kw_only", 4) assert obj.parent_required == 1 assert obj.child_required == 2 assert obj.parent_default == 3 assert obj.child_kw_only == 4 def test_low_level_child_too_many_positional(self) -> None: """Four positional args exceed the 3 positional slots for child.""" obj = _TestCxxAutoInitChild.__new__(_TestCxxAutoInitChild) with pytest.raises(TypeError, match="positional"): _ffi_init(obj, 1, 2, 3, 4, core.KWARGS, "child_kw_only", 5) class TestAutoInitSimple: def test_simple_positional(self) -> None: obj = _TestCxxAutoInitSimple(10, 20) assert obj.x == 10 assert obj.y == 20 def test_simple_keyword(self) -> None: obj = _TestCxxAutoInitSimple(x=10, y=20) assert obj.x == 10 assert obj.y == 20 def test_simple_mixed(self) -> None: obj = _TestCxxAutoInitSimple(10, y=20) assert obj.x == 10 assert obj.y == 20 def test_simple_missing_required(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitSimple(x=10) # ty: ignore[missing-argument] def test_simple_too_many_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitSimple(1, 2, 3) # ty: ignore[too-many-positional-arguments] def test_simple_unexpected_keyword(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitSimple(x=1, y=2, z=3) # ty: ignore[unknown-argument] def test_simple_low_level_kwargs(self) -> None: obj = _TestCxxAutoInitSimple.__new__(_TestCxxAutoInitSimple) _ffi_init(obj, core.KWARGS, "x", 10, "y", 20) assert obj.x == 10 assert obj.y == 20 def test_zero_values(self) -> None: obj = _TestCxxAutoInitSimple(0, 0) assert obj.x == 0 assert obj.y == 0 def test_negative_values(self) -> None: obj = _TestCxxAutoInitSimple(-1, -9999999) assert obj.x == -1 assert obj.y == -9999999 def test_large_values(self) -> None: large = 2**62 obj = _TestCxxAutoInitSimple(large, -large) assert obj.x == large assert obj.y == -large def test_float_for_integer_field(self) -> None: """Passing float where int64_t is expected - should truncate or error.""" try: obj = _TestCxxAutoInitSimple(1.5, 2) # ty: ignore[invalid-argument-type] assert isinstance(obj.x, int) except TypeError: pass # Also acceptable def test_bool_for_integer_field(self) -> None: """Booleans are valid ints in Python but should work correctly.""" obj = _TestCxxAutoInitSimple(True, False) assert obj.x == 1 assert obj.y == 0 class TestAutoInitAllInitOff: def test_no_arg_constructor(self) -> None: obj = _TestCxxAutoInitAllInitOff() assert obj.x == 7 assert obj.y == 9 assert obj.z == 1234 def test_rejects_positional_args(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitAllInitOff(1) # ty: ignore[too-many-positional-arguments] def test_rejects_keyword_args(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitAllInitOff(x=1) # ty: ignore[unknown-argument] def test_low_level_empty_init(self) -> None: obj = _TestCxxAutoInitAllInitOff.__new__(_TestCxxAutoInitAllInitOff) _ffi_init(obj) assert obj.x == 7 assert obj.y == 9 assert obj.z == 1234 def test_mutate_fields(self) -> None: obj = _TestCxxAutoInitAllInitOff() obj.x = 101 obj.y = 202 obj.z = 303 assert (obj.x, obj.y, obj.z) == (101, 202, 303) def test_empty_kwargs_dict_star(self) -> None: """Passing **{} should be fine (empty kwargs).""" obj = _TestCxxAutoInitAllInitOff(**{}) assert obj.x == 7 def test_low_level_rejects_positional(self) -> None: obj = _TestCxxAutoInitAllInitOff.__new__(_TestCxxAutoInitAllInitOff) with pytest.raises(TypeError): _ffi_init(obj, 1) class TestAutoInitKwOnlyDefaults: def test_minimal_required(self) -> None: obj = _TestCxxAutoInitKwOnlyDefaults(1, k_required=2) assert obj.p_required == 1 assert obj.p_default == 11 assert obj.k_required == 2 assert obj.k_default == 22 assert obj.hidden == 33 def test_override_defaults(self) -> None: obj = _TestCxxAutoInitKwOnlyDefaults(p_required=1, p_default=4, k_required=5, k_default=6) assert obj.p_required == 1 assert obj.p_default == 4 assert obj.k_required == 5 assert obj.k_default == 6 assert obj.hidden == 33 def test_missing_required_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitKwOnlyDefaults(k_required=2) # ty: ignore[missing-argument] def test_missing_required_kw_only(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitKwOnlyDefaults(1) # ty: ignore[missing-argument] def test_kw_only_rejects_positional(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitKwOnlyDefaults(1, 2, 3) # ty: ignore[missing-argument, too-many-positional-arguments] def test_hidden_init_false_field_not_accepted(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitKwOnlyDefaults(1, k_required=2, hidden=4) # ty: ignore[unknown-argument] def test_type_mismatch(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitKwOnlyDefaults("x", k_required=2) # ty: ignore[invalid-argument-type] def test_low_level_kwargs_call(self) -> None: obj = _TestCxxAutoInitKwOnlyDefaults.__new__(_TestCxxAutoInitKwOnlyDefaults) _ffi_init(obj, core.KWARGS, "p_required", 1, "k_required", 2) assert obj.p_required == 1 assert obj.p_default == 11 assert obj.k_required == 2 assert obj.k_default == 22 assert obj.hidden == 33 def test_kw_only_via_dict_unpacking(self) -> None: """Verify kw_only fields work via **dict.""" kwargs = {"k_required": 100, "k_default": 200} obj = _TestCxxAutoInitKwOnlyDefaults(1, **kwargs) assert obj.p_required == 1 assert obj.p_default == 11 # default assert obj.k_required == 100 assert obj.k_default == 200 assert obj.hidden == 33 # init=False default class TestAutoInitInheritance: def test_parent_constructor(self) -> None: obj = _TestCxxAutoInitParent(10) assert obj.parent_required == 10 assert obj.parent_default == 5 def test_parent_all_keyword(self) -> None: obj = _TestCxxAutoInitParent(parent_required=10, parent_default=20) assert obj.parent_required == 10 assert obj.parent_default == 20 def test_parent_positional_then_keyword(self) -> None: obj = _TestCxxAutoInitParent(10, parent_default=20) assert obj.parent_required == 10 assert obj.parent_default == 20 def test_child_constructor_uses_parent_and_child_fields(self) -> None: obj = _TestCxxAutoInitChild(parent_required=1, child_required=2, child_kw_only=3) assert obj.parent_required == 1 assert obj.parent_default == 5 assert obj.child_required == 2 assert obj.child_kw_only == 3 def test_child_all_keyword_with_parent_default_override(self) -> None: # fmt: off obj = _TestCxxAutoInitChild(parent_required=1, child_required=2, parent_default=99, child_kw_only=3) # fmt: on assert obj.parent_required == 1 assert obj.child_required == 2 assert obj.parent_default == 99 assert obj.child_kw_only == 3 def test_child_missing_parent_required(self) -> None: with pytest.raises(TypeError): _TestCxxAutoInitChild(child_required=2, child_kw_only=3) # ty: ignore[missing-argument] def test_child_two_positional_args_routes_correctly(self) -> None: """Calling Child(1, 2, child_kw_only=3) should set parent_required=1, child_required=2. The Python signature is: (self, parent_required, child_required, parent_default=..., *, child_kw_only) So positional arg 0 = parent_required, positional arg 1 = child_required. """ obj = _TestCxxAutoInitChild(1, 2, child_kw_only=3) assert obj.parent_required == 1 assert obj.child_required == 2 assert obj.parent_default == 5 # should use the default assert obj.child_kw_only == 3 def test_child_three_positional_args_no_silent_swap(self) -> None: """Calling Child(1, 2, 3, child_kw_only=4) should map correctly. Python signature order: parent_required=1, child_required=2, parent_default=3 """ obj = _TestCxxAutoInitChild(1, 2, 3, child_kw_only=4) assert obj.parent_required == 1 assert obj.child_required == 2 assert obj.parent_default == 3 assert obj.child_kw_only == 4 def test_child_one_positional_rest_keyword(self) -> None: """Mix of positional and keyword to verify correct mapping.""" obj = _TestCxxAutoInitChild(10, child_required=20, parent_default=30, child_kw_only=40) assert obj.parent_required == 10 assert obj.child_required == 20 assert obj.parent_default == 30 assert obj.child_kw_only == 40 class TestAutoInitCopyBehavior: """Test copy/deepcopy/replace interplay with auto-init objects.""" def test_shallow_copy(self) -> None: obj = _TestCxxAutoInitSimple(10, 20) obj_copy = copy.copy(obj) assert obj_copy.x == 10 assert obj_copy.y == 20 assert not obj.same_as(obj_copy) def test_deepcopy(self) -> None: obj = _TestCxxAutoInitSimple(10, 20) obj_copy = copy.deepcopy(obj) assert obj_copy.x == 10 assert obj_copy.y == 20 assert not obj.same_as(obj_copy) @requires_py313 def test_replace(self) -> None: obj = _TestCxxAutoInit(1, c=3) replaced = copy.replace(obj, a=100, c=300) # type: ignore[attr-defined] assert replaced.a == 100 assert replaced.b == 42 assert replaced.c == 300 assert replaced.d == 99 def test_copy_preserves_init_false_field(self) -> None: """After construction, mutating the init=False field and copying.""" obj = _TestCxxAutoInit(1, c=3) assert obj.b == 42 obj.b = 999 assert obj.b == 999 obj_copy = copy.copy(obj) assert obj_copy.b == 999 def test_copy_preserves_default_override(self) -> None: """Override a default field, then copy should preserve the override.""" obj = _TestCxxAutoInit(1, c=3, d=55) obj_copy = copy.copy(obj) assert obj_copy.d == 55 def test_deepcopy_all_init_off(self) -> None: """Deepcopy of an object with all fields init=False.""" obj = _TestCxxAutoInitAllInitOff() obj.x = 111 obj.y = 222 obj.z = 333 obj_copy = copy.deepcopy(obj) assert obj_copy.x == 111 assert obj_copy.y == 222 assert obj_copy.z == 333 assert not obj.same_as(obj_copy) @requires_py313 def test_replace_kw_only_defaults(self) -> None: obj = _TestCxxAutoInitKwOnlyDefaults(1, k_required=2) replaced = copy.replace(obj, k_required=99, p_default=88) # type: ignore[attr-defined] assert replaced.p_required == 1 assert replaced.p_default == 88 assert replaced.k_required == 99 assert replaced.k_default == 22 assert replaced.hidden == 33 class TestAutoInitReinitialization: """Test what happens when __ffi_init__ is called multiple times.""" def test_reinit_changes_handle(self) -> None: """Calling __ffi_init__ again should create a new underlying object.""" obj = _TestCxxAutoInit(1, c=3) original_handle = obj.__chandle__() assert obj.a == 1 _ffi_init(obj, core.KWARGS, "a", 100, "c", 300) assert obj.a == 100 assert obj.c == 300 assert obj.__chandle__() != original_handle def test_reinit_resets_init_false_field(self) -> None: """Re-initialization should reset init=False fields to defaults.""" obj = _TestCxxAutoInit(1, c=3) obj.b = 999 assert obj.b == 999 _ffi_init(obj, core.KWARGS, "a", 2, "c", 4) assert obj.b == 42 # reset to default class TestAutoInitTypeChecks: """Verify isinstance relationships for auto-init objects.""" def test_parent_isinstance(self) -> None: obj = _TestCxxAutoInitParent(1) assert isinstance(obj, _TestCxxAutoInitParent) assert isinstance(obj, core.Object) def test_child_isinstance_parent(self) -> None: obj = _TestCxxAutoInitChild(parent_required=1, child_required=2, child_kw_only=3) assert isinstance(obj, _TestCxxAutoInitChild) assert isinstance(obj, _TestCxxAutoInitParent) assert isinstance(obj, core.Object) def test_parent_not_isinstance_child(self) -> None: """A parent object should not pass isinstance check for a child class.""" obj = _TestCxxAutoInitParent(1) assert not isinstance(obj, _TestCxxAutoInitChild) class TestAutoInitInstanceIsolation: """Verify that multiple instances don't share mutable state.""" def test_separate_instances_are_independent(self) -> None: a = _TestCxxAutoInitSimple(1, 2) b = _TestCxxAutoInitSimple(3, 4) assert a.x == 1 assert b.x == 3 assert not a.same_as(b) def test_mutating_one_instance_doesnt_affect_another(self) -> None: a = _TestCxxAutoInit(1, c=3) b = _TestCxxAutoInit(1, c=3) assert a.b == 42 assert b.b == 42 a.b = 999 assert a.b == 999 assert b.b == 42 def test_all_init_off_instances_are_independent(self) -> None: a = _TestCxxAutoInitAllInitOff() b = _TestCxxAutoInitAllInitOff() a.x = 100 assert b.x == 7 class TestAutoInitReinitInitOffNoDefault: """Test reinit behavior for init=False fields with and without reflection defaults.""" def test_reinit_init_false_with_default_resets(self) -> None: """Fields b (init=False, default=42) should reset to default on reinit.""" obj = _TestCxxAutoInit(1, c=3) obj.b = 999 _ffi_init(obj, core.KWARGS, "a", 2, "c", 4) assert obj.b == 42 def test_reinit_init_false_without_reflection_default(self) -> None: """Field z has init=False AND no reflection default (c_has_default=False). On reinit, z should get whatever the C++ creator sets (1234). """ obj = _TestCxxAutoInitAllInitOff() assert obj.z == 1234 obj.z = 9999 assert obj.z == 9999 # Reinit via low-level call _ffi_init(obj) # z has no reflection default, so creator's C++ default (1234) is used assert obj.z == 1234 class TestAutoInitErrorMessages: """Verify that error messages name the correct field after pos_indices reordering.""" def test_missing_child_required_names_correct_field(self) -> None: """When child_required is missing, error should mention 'child_required'.""" with pytest.raises(TypeError, match="child_required"): _TestCxxAutoInitChild(parent_required=1, child_kw_only=3) # ty: ignore[missing-argument] def test_missing_parent_required_names_correct_field(self) -> None: """When parent_required is missing, error should mention 'parent_required'.""" with pytest.raises(TypeError, match="parent_required"): _TestCxxAutoInitChild(child_required=2, child_kw_only=3) # ty: ignore[missing-argument] def test_missing_kw_only_names_correct_field(self) -> None: """When child_kw_only is missing, error should mention 'child_kw_only'.""" with pytest.raises(TypeError, match="child_kw_only"): _TestCxxAutoInitChild(parent_required=1, child_required=2) # ty: ignore[missing-argument] def test_too_many_positional_error_message(self) -> None: """Error for too many positional args should mention the correct count.""" with pytest.raises(TypeError, match="3 positional"): _TestCxxAutoInitChild(1, 2, 3, 4, child_kw_only=5) # ty: ignore[too-many-positional-arguments] def test_unexpected_keyword_error_message(self) -> None: """Error for unknown keyword should mention the keyword name.""" with pytest.raises(TypeError, match="bogus"): _TestCxxAutoInit(1, c=2, bogus=3) # ty: ignore[unknown-argument] def test_duplicate_arg_error_message(self) -> None: """Error for duplicate argument should mention the field name.""" with pytest.raises(TypeError, match=r"multiple values.*a"): _TestCxxAutoInit(1, c=2, a=3) # ty: ignore[parameter-already-assigned] class TestAutoInitLowLevelKwOnlyDefaults: """Low-level KWARGS protocol tests for the KwOnlyDefaults type.""" def test_low_level_positional_plus_kwargs_kw_only(self) -> None: """Positional arg for p_required, then KWARGS for kw_only fields.""" obj = _TestCxxAutoInitKwOnlyDefaults.__new__(_TestCxxAutoInitKwOnlyDefaults) _ffi_init(obj, 1, core.KWARGS, "k_required", 2) assert obj.p_required == 1 assert obj.p_default == 11 assert obj.k_required == 2 assert obj.k_default == 22 assert obj.hidden == 33 def test_low_level_two_positional_plus_kwargs(self) -> None: """Two positional args (p_required, p_default) then KWARGS for kw_only.""" obj = _TestCxxAutoInitKwOnlyDefaults.__new__(_TestCxxAutoInitKwOnlyDefaults) _ffi_init(obj, 1, 2, core.KWARGS, "k_required", 3, "k_default", 4) assert obj.p_required == 1 assert obj.p_default == 2 assert obj.k_required == 3 assert obj.k_default == 4 assert obj.hidden == 33 def test_low_level_all_via_kwargs(self) -> None: """All init=True fields via KWARGS, no positional.""" obj = _TestCxxAutoInitKwOnlyDefaults.__new__(_TestCxxAutoInitKwOnlyDefaults) _ffi_init( obj, core.KWARGS, "p_required", 10, "p_default", 20, "k_required", 30, "k_default", 40 ) assert obj.p_required == 10 assert obj.p_default == 20 assert obj.k_required == 30 assert obj.k_default == 40 assert obj.hidden == 33 class TestClassLevelInitFalse: """init(false) passed to ObjectDef constructor suppresses __ffi_init__.""" def test_no_ffi_init_method(self) -> None: type_index = getattr(_TestCxxNoAutoInit, "__tvm_ffi_type_info__").type_index ffi_init = core._lookup_type_attr(type_index, "__ffi_init__") assert ffi_init is None def test_has_fields(self) -> None: type_info = getattr(_TestCxxNoAutoInit, "__tvm_ffi_type_info__") field_names = [f.name for f in type_info.fields] assert field_names == ["x", "y"] def test_direct_construction_raises(self) -> None: with pytest.raises(TypeError, match="cannot be constructed directly"): _TestCxxNoAutoInit(1, 2) # ty: ignore[too-many-positional-arguments] def test_has_shallow_copy(self) -> None: type_info = getattr(_TestCxxNoAutoInit, "__tvm_ffi_type_info__") # __ffi_shallow_copy__ is now TypeAttrColumn only, not a method fn = core._lookup_type_attr(type_info.type_index, "__ffi_shallow_copy__") assert fn is not None # ########################################################################### # kw_only error message regression tests # ########################################################################### class TestKwOnlyErrorMessages: """Verify that error messages distinguish keyword-only from positional.""" def test_missing_kw_only_error_says_keyword_only(self) -> None: """Missing required kw_only field should say 'keyword-only'.""" with pytest.raises(TypeError, match="keyword-only"): _TestCxxAutoInit(1) # ty: ignore[missing-argument] def test_missing_positional_error_does_not_say_keyword_only(self) -> None: """Missing required positional field should NOT say 'keyword-only'.""" with pytest.raises(TypeError, match="missing required argument") as exc_info: _TestCxxAutoInit(c=3) # ty: ignore[missing-argument] assert "keyword-only" not in str(exc_info.value) def test_child_missing_kw_only_error_says_keyword_only(self) -> None: """Inherited kw_only field should produce 'keyword-only' error.""" with pytest.raises(TypeError, match="keyword-only"): _TestCxxAutoInitChild(parent_required=1, child_required=2) # ty: ignore[missing-argument] def test_kw_only_defaults_missing_error(self) -> None: """KwOnlyDefaults type: missing k_required should say 'keyword-only'.""" with pytest.raises(TypeError, match="keyword-only"): _TestCxxAutoInitKwOnlyDefaults(1) # ty: ignore[missing-argument] def test_low_level_kw_only_error_says_keyword_only(self) -> None: """Low-level __ffi_init__ also gives 'keyword-only' error.""" obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) with pytest.raises(TypeError, match="keyword-only"): _ffi_init(obj, 1) # ########################################################################### # __ffi_init__ TypeMethod registration regression tests # # These tests verify that def(init<>) in C++ registers __ffi_init__ as a # TypeMethod (not just TypeAttrColumn), preserving type_schema metadata. # They also verify the __ffi_init__ instance method wrapper and the MISSING # sentinel used for default parameter values. # See commit 2987899a for the original fix. # ########################################################################### class TestFfiInitAsTypeMethod: """Verify that classes using ``def(init<>)`` expose __ffi_init__ as a TypeMethod.""" @pytest.mark.parametrize("cls", [TestIntPair, TestCompare, TestHash]) def test_explicit_init_has_ffi_init_type_method(self, cls: type) -> None: type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] method_names = {m.name for m in type_info.methods} assert "__ffi_init__" in method_names def test_auto_init_does_not_have_ffi_init_type_method(self) -> None: type_info = _TestCxxAutoInit.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] method_names = {m.name for m in type_info.methods} assert "__ffi_init__" not in method_names @pytest.mark.parametrize("cls", [TestIntPair, TestCompare, TestHash]) def test_ffi_init_type_method_has_func(self, cls: type) -> None: type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] for method in type_info.methods: if method.name == "__ffi_init__": assert method.func is not None break else: pytest.fail(f"__ffi_init__ not found in {cls.__name__} methods") # ty: ignore[invalid-argument-type] class TestFfiInitAsInstanceMethod: """Verify that __ffi_init__ is available as a callable method on registered classes.""" def test_explicit_init_class_has_ffi_init_attr(self) -> None: assert hasattr(TestIntPair, "__ffi_init__") def test_auto_init_class_has_ffi_init_attr(self) -> None: assert hasattr(_TestCxxAutoInit, "__ffi_init__") def test_ffi_init_constructs_explicit_init_object(self) -> None: obj = TestIntPair.__new__(TestIntPair) obj.__ffi_init__(10, 20) # ty: ignore[unresolved-attribute] assert obj.a == 10 assert obj.b == 20 assert obj.sum() == 30 def test_ffi_init_constructs_auto_init_object(self) -> None: obj = _TestCxxAutoInitSimple.__new__(_TestCxxAutoInitSimple) obj.__ffi_init__(7, 8) # ty: ignore[unresolved-attribute] assert obj.x == 7 assert obj.y == 8 def test_ffi_init_with_kwargs_protocol(self) -> None: obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit) obj.__ffi_init__(a=1, c=3) # ty: ignore[unresolved-attribute] assert obj.a == 1 assert obj.b == 42 assert obj.c == 3 assert obj.d == 99 class TestFfiInitTypeOwnerGuard: """Verify the type-owner guard prevents subclass misuse of __ffi_init__.""" def test_same_type_ffi_init_succeeds(self) -> None: obj = TestIntPair.__new__(TestIntPair) TestIntPair.__ffi_init__(obj, 1, 2) # ty: ignore[unresolved-attribute] assert obj.a == 1 assert obj.b == 2 def test_subclass_ffi_init_raises_type_error(self) -> None: obj = _TestCxxClassDerivedDerived.__new__(_TestCxxClassDerivedDerived) with pytest.raises(TypeError, match="not supported"): _TestCxxClassDerived.__ffi_init__(obj, 1, 2, 3.0) # ty: ignore[unresolved-attribute] class TestMissingSentinelDefaults: """Verify that auto-generated __init__ signatures use core.MISSING for defaults.""" def test_default_params_use_core_missing(self) -> None: sig = inspect.signature(_TestCxxAutoInit.__init__) d_param = sig.parameters["d"] assert d_param.default is not inspect.Parameter.empty assert d_param.default is core.MISSING def test_kw_only_default_uses_core_missing(self) -> None: sig = inspect.signature(_TestCxxAutoInitKwOnlyDefaults.__init__) assert sig.parameters["k_default"].default is core.MISSING assert sig.parameters["p_default"].default is core.MISSING def test_required_params_have_no_default(self) -> None: sig = inspect.signature(_TestCxxAutoInit.__init__) assert sig.parameters["a"].default is inspect.Parameter.empty assert sig.parameters["c"].default is inspect.Parameter.empty def test_missing_not_sent_to_ffi(self) -> None: """MISSING sentinel values should be stripped, not forwarded to C++.""" obj = _TestCxxAutoInitKwOnlyDefaults(p_required=1, k_required=2) assert obj.p_default == 11 assert obj.k_default == 22 assert obj.hidden == 33 def test_derived_class_defaults(self) -> None: obj = _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0) assert obj.v_f32 == 8.0 class TestFfiInitDualRegistration: """Verify __ffi_init__ is available from both TypeMethod and TypeAttrColumn.""" @pytest.mark.parametrize("cls", [TestIntPair, TestCompare]) def test_explicit_init_in_type_attr_column(self, cls: type) -> None: type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] ffi_init = core._lookup_type_attr(type_info.type_index, "__ffi_init__") assert ffi_init is not None def test_auto_init_in_type_attr_column(self) -> None: type_info = _TestCxxAutoInit.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] ffi_init = core._lookup_type_attr(type_info.type_index, "__ffi_init__") assert ffi_init is not None def test_explicit_init_method_and_attr_produce_same_result(self) -> None: type_info = TestIntPair.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] method_func = None for method in type_info.methods: if method.name == "__ffi_init__": method_func = method.func break assert method_func is not None attr_func = core._lookup_type_attr(type_info.type_index, "__ffi_init__") assert attr_func is not None obj1 = TestIntPair.__new__(TestIntPair) obj1.__init_handle_by_constructor__(method_func, 10, 20) obj2 = TestIntPair.__new__(TestIntPair) obj2.__init_handle_by_constructor__(attr_func, 10, 20) assert obj1.a == obj2.a == 10 assert obj1.b == obj2.b == 20 # ########################################################################### # Python 3.14 annotation regression (PEP 749) # # Python 3.14 stores annotations lazily via __annotate_func__ instead of # directly in cls.__dict__["__annotations__"]. This broke @py_class field # discovery when it used cls.__dict__.get("__annotations__", {}). # ########################################################################### @py_class("testing.PyClassSimple") class _PyClassSimple(core.Object): x: int y: int @py_class("testing.PyClassWithDefault") class _PyClassWithDefault(core.Object): a: int b: int = 42 class TestPyClassAnnotationDiscovery: """Regression: @py_class must discover fields on Python 3.14+ (PEP 749).""" def test_fields_registered(self) -> None: ti: core.TypeInfo = _PyClassSimple.__tvm_ffi_type_info__ # type: ignore[unresolved-attribute] names = [f.name for f in ti.fields] assert names == ["x", "y"] def test_construct_kwargs(self) -> None: obj = _PyClassSimple(x=1, y=2) assert obj.x == 1 assert obj.y == 2 def test_construct_positional(self) -> None: obj = _PyClassSimple(10, 20) assert obj.x == 10 assert obj.y == 20 def test_default_field(self) -> None: obj = _PyClassWithDefault(a=7) assert obj.a == 7 assert obj.b == 42 def test_override_default(self) -> None: obj = _PyClassWithDefault(a=1, b=2) assert obj.a == 1 assert obj.b == 2 tvm-ffi-0.1.12/tests/python/test_dataclass_py_class.py000066400000000000000000005474401521067262500231170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for Python-defined TVM-FFI types: ``@py_class`` decorator and low-level Field API.""" # ruff: noqa: D102, PLR0124, PLW1641, UP006, UP045 from __future__ import annotations import copy import gc import inspect import itertools import math from typing import Any, ClassVar, Dict, List, Optional import pytest import tvm_ffi from tvm_ffi import core from tvm_ffi._dunder import _install_dataclass_dunders from tvm_ffi._ffi_api import DeepCopy, RecursiveEq, RecursiveHash, ReprPrint from tvm_ffi.core import MISSING, Object, TypeInfo, TypeSchema, _to_py_class_value from tvm_ffi.dataclasses import KW_ONLY, Field, IntEnum, StrEnum, entry, field, fields, py_class from tvm_ffi.registry import _add_class_attrs from tvm_ffi.testing import TestObjectBase as _TestObjectBase from tvm_ffi.testing.testing import requires_py310 # --------------------------------------------------------------------------- # Unique type key generator (avoids collisions across tests) # --------------------------------------------------------------------------- _counter = itertools.count() def _unique_key(base: str) -> str: return f"testing.py_class_dec.{base}_{next(_counter)}" def _get_type_info(cls: type) -> TypeInfo: ret = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] assert isinstance(ret, TypeInfo), f"Expected TypeInfo, got {type(ret)}" return ret # --------------------------------------------------------------------------- # Low-level helpers for _make_type-based tests # --------------------------------------------------------------------------- _counter_ff = itertools.count() def _unique_key_ff(base: str) -> str: """Return a globally unique type key for low-level field tests.""" return f"testing.py_class.{base}_{next(_counter_ff)}" def _make_type( name: str, fields: List[Field], *, parent: type = core.Object, eq: bool = False, unsafe_hash: bool = False, repr: bool = True, ) -> type: """Create, register, and fully set up a Python-defined TVM-FFI type. Returns the ready-to-use Python class. """ type_key = _unique_key_ff(name) parent_info = core._type_cls_to_type_info(parent) assert parent_info is not None cls = type(name, (parent,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) info._register_fields(fields) setattr(cls, "__tvm_ffi_type_info__", info) _add_class_attrs(cls, info) _install_dataclass_dunders( cls, init=True, repr=repr, eq=eq, order=False, unsafe_hash=unsafe_hash, ) return cls # ########################################################################### # 1. Basic registration # ########################################################################### class TestBasicRegistration: """@py_class decorator with different calling conventions.""" def test_bare_decorator(self) -> None: @py_class(_unique_key("Bare")) class Bare(Object): x: int info = _get_type_info(Bare) assert info is not None assert len(info.fields) == 1 assert info.fields[0].name == "x" def test_decorator_with_options(self) -> None: @py_class(_unique_key("Opts"), eq=True) class Opts(Object): x: int assert hasattr(Opts, "__eq__") assert Opts(x=1) == Opts(x=1) def test_auto_type_key(self) -> None: @py_class(_unique_key("AutoKey")) class AutoKey(Object): x: int info = _get_type_info(AutoKey) assert info.type_key.startswith("testing.") def test_explicit_type_key(self) -> None: key = _unique_key("ExplicitKey") @py_class(key) class ExplicitKey(Object): x: int assert _get_type_info(ExplicitKey).type_key == key def test_empty_class(self) -> None: @py_class(_unique_key("Empty")) class Empty(Object): pass obj = Empty() assert obj is not None def test_isinstance_check(self) -> None: @py_class(_unique_key("InstCheck")) class InstCheck(Object): x: int obj = InstCheck(x=42) assert isinstance(obj, InstCheck) assert isinstance(obj, Object) # ########################################################################### # 2. Field parsing # ########################################################################### class TestFieldParsing: """Annotation-to-Field conversion.""" def test_int_field(self) -> None: @py_class(_unique_key("IntFld")) class IntFld(Object): x: int obj = IntFld(x=42) assert obj.x == 42 def test_float_field(self) -> None: @py_class(_unique_key("FltFld")) class FltFld(Object): x: float obj = FltFld(x=3.14) assert abs(obj.x - 3.14) < 1e-10 def test_str_field(self) -> None: @py_class(_unique_key("StrFld")) class StrFld(Object): x: str obj = StrFld(x="hello") assert obj.x == "hello" def test_bool_field(self) -> None: @py_class(_unique_key("BoolFld")) class BoolFld(Object): x: bool obj = BoolFld(x=True) assert obj.x is True def test_payload_enum_fields_end_to_end(self) -> None: """IntEnum/StrEnum fields on @py_class accept raw payloads and expose enum objects.""" class Priority(IntEnum, type_key=_unique_key("Priority")): low = entry(value=10) high = entry(value=20) class Opcode(StrEnum, type_key=_unique_key("Opcode")): add = entry(value="+") mul = entry(value="*") @py_class(_unique_key("EnumFields")) class Instruction(Object): priority: Priority opcode: Opcode from_payload = Instruction(priority=20, opcode="*") # ty: ignore[invalid-argument-type] assert isinstance(from_payload.priority, Priority) assert isinstance(from_payload.opcode, Opcode) assert from_payload.priority.same_as(Priority.high) assert from_payload.opcode.same_as(Opcode.mul) assert from_payload.priority.value == 20 assert from_payload.opcode.value == "*" from_enum = Instruction(priority=Priority.low, opcode=Opcode.add) assert from_enum.priority.same_as(Priority.low) assert from_enum.opcode.same_as(Opcode.add) from_enum.priority = 20 # ty: ignore[invalid-assignment] from_enum.opcode = "*" # ty: ignore[invalid-assignment] assert from_enum.priority.same_as(Priority.high) assert from_enum.opcode.same_as(Opcode.mul) with pytest.raises((TypeError, RuntimeError), match="expected"): from_enum.priority = 99 # ty: ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError), match="expected"): from_enum.opcode = "/" # ty: ignore[invalid-assignment] @requires_py310 def test_optional_field(self) -> None: @py_class(_unique_key("OptFld")) class OptFld(Object): x: Optional[int] obj = OptFld(x=42) assert obj.x == 42 obj2 = OptFld(x=None) assert obj2.x is None def test_multiple_fields(self) -> None: @py_class(_unique_key("Multi")) class Multi(Object): a: int b: float c: str obj = Multi(a=1, b=2.0, c="three") assert obj.a == 1 assert obj.b == 2.0 assert obj.c == "three" # ########################################################################### # 3. Defaults # ########################################################################### class TestDefaults: """Default values and default_factory.""" def test_bare_default(self) -> None: @py_class(_unique_key("BareDef")) class BareDef(Object): x: int y: int = 10 obj = BareDef(x=1) assert obj.y == 10 def test_field_default(self) -> None: @py_class(_unique_key("FldDef")) class FldDef(Object): x: int = field(default=42) obj = FldDef() assert obj.x == 42 def test_field_default_factory(self) -> None: call_count = 0 def make_default() -> int: nonlocal call_count call_count += 1 return 99 @py_class(_unique_key("FldFact")) class FldFact(Object): x: int = field(default_factory=make_default) obj1 = FldFact() assert obj1.x == 99 obj2 = FldFact() assert obj2.x == 99 assert call_count == 2 def test_default_and_factory_mutually_exclusive(self) -> None: with pytest.raises(ValueError, match="cannot specify both"): field(default=1, default_factory=int) def test_non_callable_factory_rejected(self) -> None: with pytest.raises(TypeError, match="default_factory must be a callable"): field(default_factory=42) # ty: ignore[invalid-argument-type] def test_required_before_optional(self) -> None: @py_class(_unique_key("ReqOpt")) class ReqOpt(Object): a: int b: int = 10 obj = ReqOpt(1) assert obj.a == 1 assert obj.b == 10 # ########################################################################### # 4. KW_ONLY # ########################################################################### class TestKwOnly: """Keyword-only field support.""" def test_kw_only_sentinel(self) -> None: @py_class(_unique_key("KWSent")) class KWSent(Object): a: int _: KW_ONLY b: int = 10 obj = KWSent(1, b=20) # ty: ignore[missing-argument] assert obj.a == 1 assert obj.b == 20 with pytest.raises(TypeError): KWSent(1, 2) # ty: ignore[invalid-argument-type] def test_decorator_level_kw_only(self) -> None: @py_class(_unique_key("DecKW"), kw_only=True) class DecKW(Object): a: int b: int = 10 obj = DecKW(a=1) assert obj.a == 1 assert obj.b == 10 with pytest.raises(TypeError): DecKW(1) # ty: ignore[missing-argument,too-many-positional-arguments] def test_field_level_kw_only_override(self) -> None: @py_class(_unique_key("FldKW")) class FldKW(Object): a: int b: int = field(default=10, kw_only=True) obj = FldKW(1) assert obj.a == 1 assert obj.b == 10 with pytest.raises(TypeError): FldKW(1, 2) # b is keyword-only # ########################################################################### # 5. ClassVar # ########################################################################### class TestClassVar: """ClassVar annotations are skipped.""" def test_classvar_skipped(self) -> None: @py_class(_unique_key("CV")) class CV(Object): x: int count: ClassVar[int] = 0 info = _get_type_info(CV) field_names = [f.name for f in info.fields] assert "x" in field_names assert "count" not in field_names def test_classvar_preserved_on_class(self) -> None: @py_class(_unique_key("CVPres")) class CVPres(Object): x: int tag: ClassVar[str] = "hello" assert CVPres.tag == "hello" # ########################################################################### # 6. Init generation # ########################################################################### class TestInit: """Auto-generated __init__.""" def test_positional_args(self) -> None: @py_class(_unique_key("Pos")) class Pos(Object): a: int b: str obj = Pos(1, "hello") assert obj.a == 1 assert obj.b == "hello" def test_keyword_args(self) -> None: @py_class(_unique_key("Kw")) class Kw(Object): a: int b: str obj = Kw(a=1, b="hello") assert obj.a == 1 assert obj.b == "hello" def test_init_false_field(self) -> None: @py_class(_unique_key("NoInit")) class NoInit(Object): a: int b: int = field(default=99, init=False) obj = NoInit(a=1) assert obj.a == 1 assert obj.b == 99 def test_user_defined_init_preserved(self) -> None: @py_class(_unique_key("UserInit"), init=False) class UserInit(Object): a: int def __init__(self, val: int) -> None: self.a = val obj = UserInit(42) assert obj.a == 42 def test_required_after_optional_reordered(self) -> None: """Required positional fields are reordered before optional ones in __init__.""" @py_class(_unique_key("ReorderOwn")) class ReorderOwn(Object): x: int = 0 y: int # ty: ignore[dataclass-field-order] sig = inspect.signature(ReorderOwn.__init__) param_names = [n for n in sig.parameters if n != "self"] assert param_names[0] == "y" # required comes first assert param_names[1] == "x" # optional comes second obj = ReorderOwn(y=1) # ty: ignore[missing-argument] assert obj.x == 0 assert obj.y == 1 def test_required_after_optional_in_parent(self) -> None: """Child required fields are reordered before parent optional fields.""" @py_class(_unique_key("OptParent")) class OptParent(Object): x: int y: int = 0 @py_class(_unique_key("ReqChild")) class ReqChild(OptParent): z: int sig = inspect.signature(ReqChild.__init__) param_names = [n for n in sig.parameters if n != "self"] # required (x, z) before optional (y) assert param_names == ["x", "z", "y"] obj = ReqChild(x=1, z=3) assert obj.x == 1 assert obj.y == 0 assert obj.z == 3 def test_kw_only_exempt_from_reorder(self) -> None: """kw_only fields are not reordered with positional fields.""" @py_class(_unique_key("KwReorder")) class KwReorder(Object): x: int = 0 _: KW_ONLY # ty: ignore[dataclass-field-order] y: int # ty: ignore[dataclass-field-order] sig = inspect.signature(KwReorder.__init__) params = sig.parameters assert params["x"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD assert params["y"].kind == inspect.Parameter.KEYWORD_ONLY obj = KwReorder(y=1) # ty: ignore[missing-argument] assert obj.x == 0 assert obj.y == 1 def test_mixed_positional_and_kw_only_with_defaults(self) -> None: """Mixed positional/kw_only fields with defaults produce correct signature.""" @py_class(_unique_key("MixedSig")) class MixedSig(Object): a: int = 0 b: int # ty: ignore[dataclass-field-order] _: KW_ONLY # ty: ignore[dataclass-field-order] c: int = 10 d: int # ty: ignore[dataclass-field-order] sig = inspect.signature(MixedSig.__init__) param_names = [n for n in sig.parameters if n != "self"] # positional: b (required) before a (optional); kw_only: d (required) before c (optional) assert param_names == ["b", "a", "d", "c"] obj = MixedSig(b=2, d=4) # ty: ignore[missing-argument] assert obj.a == 0 assert obj.b == 2 assert obj.c == 10 assert obj.d == 4 def test_init_false_excluded_from_signature(self) -> None: """init=False fields do not appear in __init__ signature.""" @py_class(_unique_key("InitFalseSig")) class InitFalseSig(Object): a: int b: int = field(default=99, init=False) c: str sig = inspect.signature(InitFalseSig.__init__) param_names = [n for n in sig.parameters if n != "self"] assert "b" not in param_names assert "a" in param_names assert "c" in param_names # ########################################################################### # 7. __post_init__ # ########################################################################### class TestPostInit: """__post_init__ support.""" def test_post_init_called(self) -> None: post_init_called = False @py_class(_unique_key("PostInit")) class PostInit(Object): x: int def __post_init__(self) -> None: nonlocal post_init_called post_init_called = True PostInit(x=1) assert post_init_called def test_post_init_sees_field_values(self) -> None: @py_class(_unique_key("PostInitVal")) class PostInitVal(Object): x: int y: int = 10 def __post_init__(self) -> None: # Fields should be set before __post_init__ is called assert self.x is not None assert self.y == 10 PostInitVal(x=5) # ########################################################################### # 8. Repr # ########################################################################### class TestRepr: """__repr__ generation.""" def test_repr_generated(self) -> None: @py_class(_unique_key("Repr")) class Repr(Object): x: int y: str obj = Repr(x=1, y="hello") r = repr(obj) assert "1" in r assert "hello" in r def test_repr_disabled(self) -> None: @py_class(_unique_key("NoRepr"), repr=False) class NoRepr(Object): x: int obj = NoRepr(x=1) # Should use default object repr r = repr(obj) assert "NoRepr" in r or "object at" in r # ########################################################################### # 9. Equality # ########################################################################### class TestEquality: """__eq__ and __ne__ generation.""" def test_eq_enabled(self) -> None: @py_class(_unique_key("Eq"), eq=True) class Eq(Object): x: int y: str assert Eq(x=1, y="a") == Eq(x=1, y="a") assert Eq(x=1, y="a") != Eq(x=2, y="a") def test_eq_disabled_by_default(self) -> None: @py_class(_unique_key("NoEq")) class NoEq(Object): x: int a = NoEq(x=1) b = NoEq(x=1) # Without eq, identity comparison assert a != b assert a == a # ########################################################################### # 10. Order # ########################################################################### class TestOrder: """Comparison methods.""" def test_order_enabled(self) -> None: @py_class(_unique_key("Ord"), eq=True, order=True) class Ord(Object): x: int assert Ord(x=1) < Ord(x=2) assert Ord(x=2) > Ord(x=1) assert Ord(x=1) <= Ord(x=1) assert Ord(x=1) >= Ord(x=1) # ########################################################################### # 11. Hash # ########################################################################### class TestHash: """__hash__ generation.""" def test_unsafe_hash(self) -> None: @py_class(_unique_key("Hash"), eq=True, unsafe_hash=True) class Hash(Object): x: int a = Hash(x=1) b = Hash(x=1) assert hash(a) == hash(b) # Can be used in sets s = {a, b} assert len(s) == 1 # ########################################################################### # 12. Copy # ########################################################################### class TestCopy: """__copy__, __deepcopy__, __replace__.""" def test_shallow_copy(self) -> None: @py_class(_unique_key("SCopy")) class SCopy(Object): x: int obj = SCopy(x=42) obj2 = copy.copy(obj) assert obj2.x == 42 def test_deep_copy(self) -> None: @py_class(_unique_key("DCopy")) class DCopy(Object): x: int obj = DCopy(x=42) obj2 = copy.deepcopy(obj) assert obj2.x == 42 def test_replace(self) -> None: @py_class(_unique_key("Repl")) class Repl(Object): x: int y: str obj = Repl(x=1, y="a") obj2 = obj.__replace__(x=2) # ty: ignore[unresolved-attribute] assert obj2.x == 2 assert obj2.y == "a" # ########################################################################### # 13. Inheritance # ########################################################################### class TestInheritance: """Inheritance between py_class types.""" def test_child_adds_fields(self) -> None: @py_class(_unique_key("Parent")) class Parent(Object): x: int @py_class(_unique_key("Child")) class Child(Parent): y: str obj = Child(x=1, y="hello") assert obj.x == 1 assert obj.y == "hello" def test_child_isinstance(self) -> None: @py_class(_unique_key("P2")) class P2(Object): x: int @py_class(_unique_key("C2")) class C2(P2): y: str obj = C2(x=1, y="hello") assert isinstance(obj, C2) assert isinstance(obj, P2) assert isinstance(obj, Object) def test_three_level_inheritance(self) -> None: @py_class(_unique_key("L1")) class L1(Object): a: int @py_class(_unique_key("L2")) class L2(L1): b: int @py_class(_unique_key("L3")) class L3(L2): c: int obj = L3(a=1, b=2, c=3) assert obj.a == 1 assert obj.b == 2 assert obj.c == 3 def test_collects_fields_from_non_py_class_parent(self) -> None: @py_class(_unique_key("NPCNode")) class Node(Object): x: int class BaseBinOp(Node): lhs: int rhs: int @py_class(_unique_key("NPCAdd")) class Add(BaseBinOp): pass obj = Add(lhs=1, rhs=2, x=0) # ty: ignore[unknown-argument] assert obj.x == 0 assert obj.lhs == 1 assert obj.rhs == 2 assert [f.name for f in fields(Add)] == ["x", "lhs", "rhs"] assert [f.name for f in _get_type_info(Add).fields] == ["lhs", "rhs"] def test_collects_non_py_class_parent_field_options(self) -> None: @py_class(_unique_key("NPCOptNode")) class Node(Object): x: int class BaseOp(Node): kind: ClassVar[str] = "binop" lhs: int hidden: int = field(default=99, init=False) _: KW_ONLY rhs: int @py_class(_unique_key("NPCOptAdd")) class Add(BaseOp): scale: int obj = Add(0, 1, rhs=3, scale=2) # ty: ignore[parameter-already-assigned,unknown-argument] assert obj.x == 0 assert obj.lhs == 1 assert obj.scale == 2 assert obj.rhs == 3 assert obj.hidden == 99 assert [f.name for f in fields(Add)] == ["x", "lhs", "hidden", "rhs", "scale"] assert [f.name for f in _get_type_info(Add).fields] == [ "lhs", "hidden", "rhs", "scale", ] with pytest.raises(TypeError): Add(0, 1, 2, rhs=3) # ty: ignore[too-many-positional-arguments,unknown-argument] def test_registered_parent_non_py_class_fields_not_duplicated(self) -> None: @py_class(_unique_key("NPCDedupNode")) class Node(Object): x: int class BaseBinOp(Node): lhs: int rhs: int @py_class(_unique_key("NPCDedupAdd")) class Add(BaseBinOp): op_id: int @py_class(_unique_key("NPCDedupWeightedAdd")) class WeightedAdd(Add): weight: int obj = WeightedAdd(x=0, lhs=1, rhs=2, op_id=3, weight=4) # ty: ignore[unknown-argument] assert obj.x == 0 assert obj.lhs == 1 assert obj.rhs == 2 assert obj.op_id == 3 assert obj.weight == 4 assert [f.name for f in fields(WeightedAdd)] == ["x", "lhs", "rhs", "op_id", "weight"] assert [f.name for f in _get_type_info(WeightedAdd).fields] == ["weight"] def test_multiple_non_py_class_parents_single_ffi_lineage(self) -> None: @py_class(_unique_key("MROBase")) class Base(Object): x: int class NonFFIClassC(Base): c: int class NonFFIClassA: a: int class NonFFIClassB: b: int @py_class(_unique_key("MROChild")) class Child(NonFFIClassA, NonFFIClassB, NonFFIClassC): y: int obj = Child(x=0, c=3, b=2, a=1, y=4) # ty: ignore[unknown-argument] assert obj.x == 0 assert obj.a == 1 assert obj.b == 2 assert obj.c == 3 assert obj.y == 4 assert [f.name for f in fields(Child)] == ["x", "c", "b", "a", "y"] assert [f.name for f in _get_type_info(Child).fields] == ["c", "b", "a", "y"] def test_non_py_class_parent_override_preserves_field_position(self) -> None: @py_class(_unique_key("OverrideBase")) class Base(Object): x: int class Parent(Base): a: int b: int c: int @py_class(_unique_key("OverrideChild")) class Child(Parent): b: str d: int obj = Child(x=0, a=1, b="two", c=3, d=4) # ty: ignore[unknown-argument] assert obj.b == "two" assert [f.name for f in fields(Child)] == ["x", "a", "b", "c", "d"] own_fields = _get_type_info(Child).fields assert [f.name for f in own_fields] == ["a", "b", "c", "d"] assert own_fields[1].dataclass_field is not None assert own_fields[1].dataclass_field.type is str # ########################################################################### # 14. Forward references / deferred resolution # ########################################################################### class TestForwardReferences: """Deferred annotation resolution for mutual and self-references.""" @requires_py310 def test_self_reference(self) -> None: @py_class(_unique_key("SelfRef")) class SelfRef(Object): value: int next_node: SelfRef | None leaf = SelfRef(value=2, next_node=None) head = SelfRef(value=1, next_node=leaf) assert head.next_node is not None assert head.next_node.value == 2 @requires_py310 def test_mutual_reference(self) -> None: """Two classes that reference each other.""" @py_class(_unique_key("Foo")) class Foo(Object): value: int bar: Bar | None @py_class(_unique_key("Bar")) class Bar(Object): value: int foo: Foo | None bar = Bar(value=2, foo=None) foo = Foo(value=1, bar=bar) assert foo.bar is not None assert foo.bar.value == 2 @requires_py310 def test_deferred_resolution_on_instantiation(self) -> None: """Forward ref resolved on first instantiation.""" @py_class(_unique_key("Early")) class Early(Object): value: int ref: Late | None # At this point, Early's fields are deferred because Late doesn't exist @py_class(_unique_key("Late")) class Late(Object): value: int # Now Early should resolve (either via flush or on instantiation) obj = Early(value=1, ref=Late(value=2)) assert obj.ref is not None assert obj.ref.value == 2 @requires_py310 def test_cross_module_forward_ref_via_circular_import(self) -> None: """Forward ref to a class in another module that isn't in the declaring module's globals (circular import / TYPE_CHECKING-gated) still resolves. Mirrors the loom codegen case where ``weave_ir.TaskSpec`` has a field ``body: tuple[Op, ...]`` and ``Op`` lives in ``ops`` — ``ops`` already imports from ``weave_ir``, so ``Op`` cannot appear in ``weave_ir``'s globals without introducing a circular import. """ # Module B: the target of the forward ref. Define it first so it's # registered in ``_PY_CLASS_BY_MODULE`` under its own module key when # module A's phase-2 runs. ns_b: Dict[str, Any] = { "__name__": "testing.cross_mod_b", "py_class": py_class, "Object": Object, "_unique_key": _unique_key, } exec( "from __future__ import annotations\n" "@py_class(_unique_key('CrossChild'))\n" "class CrossChild(Object):\n" " x: int\n", ns_b, ) child_cls = ns_b["CrossChild"] assert child_cls.__module__ == "testing.cross_mod_b" # Module A: references ``CrossChild`` by name but does NOT have it # in its globals (simulating a ``TYPE_CHECKING``-gated import). ns_a: Dict[str, Any] = { "__name__": "testing.cross_mod_a", "py_class": py_class, "Object": Object, "_unique_key": _unique_key, } exec( "from __future__ import annotations\n" "@py_class(_unique_key('CrossHolder'))\n" "class Holder(Object):\n" " value: int\n" " child: CrossChild | None = None\n", ns_a, ) holder_cls = ns_a["Holder"] assert holder_cls.__module__ == "testing.cross_mod_a" # Instantiation forces phase-2 to run; the cross-module localns # fallback should pick up CrossChild from module B even though it's # neither in module A's globals nor in module A's per-module registry. child = child_cls(x=7) holder = holder_cls(value=1, child=child) assert holder.child is not None assert holder.child.x == 7 # ########################################################################### # 15. User-defined dunder preservation # ########################################################################### class TestDunderPreservation: """User-defined dunders are not overwritten.""" def test_user_repr_preserved(self) -> None: @py_class(_unique_key("UserRepr")) class UserRepr(Object): x: int def __repr__(self) -> str: return f"Custom({self.x})" obj = UserRepr(x=42) assert repr(obj) == "Custom(42)" def test_user_eq_preserved(self) -> None: @py_class(_unique_key("UserEq"), eq=True) class UserEq(Object): x: int def __eq__(self, other: object) -> bool: return False assert not (UserEq(x=1) == UserEq(x=1)) # ########################################################################### # 16. field() API # ########################################################################### class TestFieldAPI: """field() function returns a Field.""" def test_field_returns_field(self) -> None: f = field(default=42) assert isinstance(f, Field) assert f.default == 42 def test_field_defaults(self) -> None: f = field() assert f.init is True assert f.repr is True assert f.hash is None # None = follow compare assert f.compare is True def test_field_kw_only_missing_by_default(self) -> None: f = field() assert f.kw_only is None def test_field_repr_false(self) -> None: @py_class(_unique_key("FldRepr")) class FldRepr(Object): x: int y: int = field(default=0, repr=False) obj = FldRepr(x=1) r = repr(obj) assert "1" in r # y with repr=False should not appear in repr # (depends on C++ ReprPrint implementation respecting the flag) # ########################################################################### # 17. Edge cases # ########################################################################### class TestEdgeCases: """Edge cases and error conditions.""" def test_no_ffi_parent_raises(self) -> None: with pytest.raises(TypeError, match="must inherit from"): @py_class(_unique_key("NoPar")) class NoPar: # no Object parent! x: int def test_only_classvar(self) -> None: @py_class(_unique_key("OnlyCV")) class OnlyCV(Object): count: ClassVar[int] = 0 obj = OnlyCV() assert obj is not None def test_mutation_after_creation(self) -> None: @py_class(_unique_key("Mut")) class Mut(Object): x: int obj = Mut(x=1) obj.x = 42 assert obj.x == 42 # ########################################################################### # 18. hash=None tri-state # ########################################################################### class TestHashTriState: """field(hash=None) means 'follow compare' (native dataclass semantics).""" def test_hash_none_follows_compare_true(self) -> None: """hash=None + compare=True → field participates in hash.""" @py_class(_unique_key("HNT"), eq=True, unsafe_hash=True) class HNT(Object): x: int # default: compare=True, hash=None → hash=True a = HNT(x=1) b = HNT(x=1) assert hash(a) == hash(b) def test_hash_none_follows_compare_false(self) -> None: """hash=None + compare=False → field excluded from hash.""" @py_class(_unique_key("HNF"), eq=True, unsafe_hash=True) class HNF(Object): x: int y: int = field(compare=False) # hash=None → follows compare=False # y doesn't participate in hash, so different y values → same hash a = HNF(x=1, y=10) b = HNF(x=1, y=20) assert hash(a) == hash(b) def test_hash_explicit_true_with_compare_true(self) -> None: """hash=True + compare=True → field participates in hash.""" @py_class(_unique_key("HET"), eq=True, unsafe_hash=True) class HET(Object): x: int = field(hash=True) # compare=True (default) a = HET(x=1) b = HET(x=2) assert hash(a) != hash(b) def test_hash_explicit_false(self) -> None: """hash=False excludes field from hashing even with compare=True.""" @py_class(_unique_key("HEF"), eq=True, unsafe_hash=True) class HEF(Object): x: int y: int = field(hash=False) # compare=True but hash=False a = HEF(x=1, y=10) b = HEF(x=1, y=20) assert hash(a) == hash(b) # ########################################################################### # 19. Deferred resolution + user __init__ / init=False # ########################################################################### class TestDeferredInitPreservation: """Deferred resolution preserves user-defined __init__ and init=False.""" @requires_py310 def test_deferred_with_user_init(self) -> None: """User-defined __init__ is preserved after deferred resolution.""" @py_class(_unique_key("DefUI")) class DefUI(Object): value: int ref: DefUILate | None def __init__(self, value: int) -> None: self.value = value self.ref = None @py_class(_unique_key("DefUILate")) class DefUILate(Object): x: int # DefUI should use the user-defined __init__ (one positional arg) obj = DefUI(42) assert obj.value == 42 assert obj.ref is None @requires_py310 def test_deferred_with_init_false(self) -> None: """init=False is respected after deferred resolution.""" @py_class(_unique_key("DefNoInit"), init=False) class DefNoInit(Object): value: int ref: DefNoInitLate | None def __init__(self, v: int) -> None: self.value = v self.ref = None @py_class(_unique_key("DefNoInitLate")) class DefNoInitLate(Object): x: int obj = DefNoInit(10) assert obj.value == 10 # ########################################################################### # 21. order=True requires eq=True # ########################################################################### class TestOrderEqValidation: """order=True without eq=True is rejected.""" def test_order_without_eq_raises(self) -> None: with pytest.raises(ValueError, match="order=True requires eq=True"): @py_class(_unique_key("OrdNoEq"), order=True) class OrdNoEq(Object): x: int # ########################################################################### # 23. Registration rollback on failure # ########################################################################### class TestRegistrationRollback: """Failed decorations don't permanently poison the type registry.""" def test_failed_decoration_allows_retry(self) -> None: key = _unique_key("Rollback") with pytest.raises(Exception): @py_class(key) class Bad(Object): x: object # unsupported annotation type # The type key should be available for reuse @py_class(key) class Good(Object): x: int y: int = 0 assert Good(x=1).y == 0 # ########################################################################### # 24. User-defined __replace__ preserved # ########################################################################### class TestUserReplace: """User-defined __replace__ is not overwritten by py_class.""" def test_user_replace_preserved(self) -> None: @py_class(_unique_key("UserRepl")) class UserRepl(Object): x: int def __replace__(self, **changes: object) -> str: return "custom" obj = UserRepl(x=1) assert obj.__replace__(x=2) == "custom" # ########################################################################### # 25. default_factory=None raises # ########################################################################### class TestDefaultFactoryNone: """Explicit default_factory=None matches stdlib semantics (raises).""" def test_explicit_none_raises(self) -> None: with pytest.raises(TypeError, match="default_factory must be a callable"): field(default_factory=None) # ########################################################################### # 26. Adversarial edge cases for init reordering # ########################################################################### class TestInitReorderingAdversarial: """Tricky scenarios that catch bugs in naive init-signature generation.""" def test_positional_call_maps_to_required_not_declared_order(self) -> None: """Positional arg 1 maps to the first *required* field, not the first declared.""" @py_class(_unique_key("PosMap")) class PosMap(Object): x: int = 0 # optional, declared first y: int # ty: ignore[dataclass-field-order] # required, declared second # Positional call: first arg is y (required), not x (optional) obj = PosMap(42) # ty: ignore[missing-argument] assert obj.y == 42 assert obj.x == 0 def test_relative_order_preserved_within_groups(self) -> None: """Within required and optional groups, declaration order is preserved.""" @py_class(_unique_key("RelOrder")) class RelOrder(Object): a: int = 0 b: int # ty: ignore[dataclass-field-order] c: int = 1 d: int # ty: ignore[dataclass-field-order] sig = inspect.signature(RelOrder.__init__) param_names = [n for n in sig.parameters if n != "self"] # required: b, d (declaration order); optional: a, c (declaration order) assert param_names == ["b", "d", "a", "c"] obj = RelOrder(10, 20) # ty: ignore[missing-argument] assert obj.b == 10 assert obj.d == 20 assert obj.a == 0 assert obj.c == 1 def test_default_factory_counts_as_optional(self) -> None: """default_factory makes a field optional for reordering purposes.""" @py_class(_unique_key("DFReorder")) class DFReorder(Object): items: str = field(default_factory=lambda: "hello") count: int # ty: ignore[dataclass-field-order] sig = inspect.signature(DFReorder.__init__) param_names = [n for n in sig.parameters if n != "self"] assert param_names[0] == "count" # required first assert param_names[1] == "items" # optional (factory) second obj = DFReorder(count=5) assert obj.count == 5 assert obj.items == "hello" def test_three_level_hierarchy_reorder(self) -> None: """Required fields from all levels come before optional fields from all levels.""" @py_class(_unique_key("G1")) class G1(Object): a: int # required @py_class(_unique_key("P1")) class P1(G1): b: int = 0 # optional @py_class(_unique_key("C1")) class C1(P1): c: int # required sig = inspect.signature(C1.__init__) param_names = [n for n in sig.parameters if n != "self"] # required (a, c) before optional (b) assert param_names == ["a", "c", "b"] obj = C1(a=1, c=3) assert obj.a == 1 assert obj.b == 0 assert obj.c == 3 def test_kw_only_false_overrides_sentinel(self) -> None: """kw_only=False on a field after KW_ONLY sentinel makes it positional.""" @py_class(_unique_key("KwOverride")) class KwOverride(Object): _: KW_ONLY a: int # kw_only (inherits sentinel) b: int = field(kw_only=False) # positional (explicit override) sig = inspect.signature(KwOverride.__init__) assert sig.parameters["a"].kind == inspect.Parameter.KEYWORD_ONLY assert sig.parameters["b"].kind == inspect.Parameter.POSITIONAL_OR_KEYWORD obj = KwOverride(42, a=1) # ty: ignore[missing-argument,invalid-argument-type] assert obj.b == 42 assert obj.a == 1 def test_init_false_field_gets_default(self) -> None: """init=False field with default is set to default, not left uninitialized.""" @py_class(_unique_key("InitFalseDef")) class InitFalseDef(Object): visible: int hidden: str = field(default="secret", init=False) obj = InitFalseDef(visible=1) assert obj.hidden == "secret" def test_post_init_sees_reordered_fields(self) -> None: """__post_init__ sees correct values even when __init__ reorders fields.""" seen: Dict[str, int] = {} @py_class(_unique_key("PostReorder")) class PostReorder(Object): x: int = 0 y: int # ty: ignore[dataclass-field-order] def __post_init__(self) -> None: seen["x"] = self.x seen["y"] = self.y PostReorder(y=10, x=20) assert seen == {"x": 20, "y": 10} @requires_py310 def test_deferred_forward_ref_with_reordering(self) -> None: """Deferred forward-reference resolution still produces correct reordering.""" @py_class(_unique_key("DeferReorder")) class DeferReorder(Object): opt: DeferLate | None = None req: int # ty: ignore[dataclass-field-order] @py_class(_unique_key("DeferLate")) class DeferLate(Object): x: int sig = inspect.signature(DeferReorder.__init__) param_names = [n for n in sig.parameters if n != "self"] assert param_names[0] == "req" assert param_names[1] == "opt" obj = DeferReorder(req=1) assert obj.req == 1 assert obj.opt is None def test_all_optional_preserves_declaration_order(self) -> None: """When all fields are optional, declaration order is preserved.""" @py_class(_unique_key("AllOpt")) class AllOpt(Object): c: int = 3 a: int = 1 b: int = 2 sig = inspect.signature(AllOpt.__init__) param_names = [n for n in sig.parameters if n != "self"] assert param_names == ["c", "a", "b"] obj = AllOpt() assert obj.c == 3 assert obj.a == 1 assert obj.b == 2 def test_all_required_preserves_declaration_order(self) -> None: """When all fields are required, declaration order is preserved.""" @py_class(_unique_key("AllReq")) class AllReq(Object): c: int a: int b: int sig = inspect.signature(AllReq.__init__) param_names = [n for n in sig.parameters if n != "self"] assert param_names == ["c", "a", "b"] obj = AllReq(10, 20, 30) assert obj.c == 10 assert obj.a == 20 assert obj.b == 30 # ########################################################################### # 1. Registration # ########################################################################### class TestRegisterPyClass: """Low-level _register_py_class: type allocation, ancestors, field lifecycle.""" def test_basic_registration(self) -> None: type_key = _unique_key_ff("RegBasic") parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls = type("RegBasic", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) assert info is not None assert info.type_key == type_key def test_type_index_allocated(self) -> None: type_key = _unique_key_ff("RegIndex") parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls = type("RegIndex", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) assert isinstance(info.type_index, int) assert info.type_index > 0 def test_ancestors_include_parent(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None type_key = _unique_key_ff("RegAncestors") cls = type("RegAncestors", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) assert parent_info.type_index in info.type_ancestors def test_parent_type_info_set(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None type_key = _unique_key_ff("RegParent") cls = type("RegParent", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) assert info.parent_type_info is parent_info def test_initial_fields_none_and_methods_empty(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None type_key = _unique_key_ff("RegEmpty") cls = type("RegEmpty", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, type_key, cls) assert info.fields is None assert len(info.methods) == 0 def test_two_registrations_different_indices(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls1 = type("RegDiff1", (core.Object,), {"__slots__": ()}) cls2 = type("RegDiff2", (core.Object,), {"__slots__": ()}) info1 = core._register_py_class(parent_info, _unique_key_ff("RegDiff1"), cls1) info2 = core._register_py_class(parent_info, _unique_key_ff("RegDiff2"), cls2) assert info1.type_index != info2.type_index def test_fields_none_before_registration(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls = type("Pending", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, _unique_key_ff("Pending"), cls) assert info.fields is None def test_register_fields_is_instance_method(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls = type("PendingM", (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, _unique_key_ff("PendingM"), cls) assert hasattr(info, "_register_fields") def test_duplicate_type_key_raises(self) -> None: parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None type_key = _unique_key_ff("Dup") cls1 = type("Dup1", (core.Object,), {"__slots__": ()}) core._register_py_class(parent_info, type_key, cls1) cls2 = type("Dup2", (core.Object,), {"__slots__": ()}) with pytest.raises((RuntimeError, ValueError)): core._register_py_class(parent_info, type_key, cls2) def test_duplicate_type_key_preserves_original(self) -> None: """After rejected duplicate, original entry is intact.""" parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None type_key = _unique_key_ff("DupPreserve") cls1 = type("DupPreserve1", (core.Object,), {"__slots__": ()}) info1 = core._register_py_class(parent_info, type_key, cls1) info1._register_fields([Field(name="x", _ty_schema=TypeSchema("int"))]) setattr(cls1, "__tvm_ffi_type_info__", info1) _add_class_attrs(cls1, info1) cls2 = type("DupPreserve2", (core.Object,), {"__slots__": ()}) with pytest.raises((RuntimeError, ValueError)): core._register_py_class(parent_info, type_key, cls2) reloaded = core._lookup_or_register_type_info_from_type_key(type_key) assert reloaded.type_cls is cls1 assert [f.name for f in reloaded.fields] == ["x"] # ########################################################################### # 2. Field Registration # ########################################################################### class TestFieldRegistration: """Low-level _register_fields: field types, metadata, offsets.""" def test_int_field_registered(self) -> None: cls = _make_type( "FldInt", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) info = getattr(cls, "__tvm_ffi_type_info__") assert len(info.fields) == 1 assert info.fields[0].name == "x" def test_float_field_registered(self) -> None: cls = _make_type( "FldFloat", [Field(name="val", _ty_schema=TypeSchema("float"), default=0.0)], ) info = getattr(cls, "__tvm_ffi_type_info__") assert info.fields[0].name == "val" def test_str_field_registered(self) -> None: cls = _make_type( "FldStr", [Field(name="s", _ty_schema=TypeSchema("str"), default="hello")], ) info = getattr(cls, "__tvm_ffi_type_info__") assert info.fields[0].name == "s" def test_bool_field_registered(self) -> None: cls = _make_type( "FldBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=False)], ) info = getattr(cls, "__tvm_ffi_type_info__") assert info.fields[0].name == "flag" def test_multiple_fields_count(self) -> None: cls = _make_type( "FldMulti", [ Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("float"), default=0.0), Field(name="c", _ty_schema=TypeSchema("str"), default="x"), ], ) info = getattr(cls, "__tvm_ffi_type_info__") assert len(info.fields) == 3 assert [f.name for f in info.fields] == ["a", "b", "c"] def test_field_offsets_increasing(self) -> None: cls = _make_type( "FldOff", [ Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("float"), default=MISSING), Field(name="c", _ty_schema=TypeSchema("str"), default=MISSING), ], ) info = getattr(cls, "__tvm_ffi_type_info__") offsets = [f.offset for f in info.fields] for i in range(1, len(offsets)): assert offsets[i] > offsets[i - 1], f"Field offsets not increasing: {offsets}" def test_ffi_init_method_registered(self) -> None: cls = _make_type( "FldInit", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) info = getattr(cls, "__tvm_ffi_type_info__") ffi_init = core._lookup_type_attr(info.type_index, "__ffi_init__") assert ffi_init is not None def test_field_metadata_repr_flag(self) -> None: cls = _make_type( "FldReprMeta", [ Field( name="visible", _ty_schema=TypeSchema("int"), default=MISSING, repr=True, ), Field( name="hidden", _ty_schema=TypeSchema("int"), default=0, repr=False, ), ], ) info = getattr(cls, "__tvm_ffi_type_info__") assert len(info.fields) == 2 # ########################################################################### # 3. Field Descriptor # ########################################################################### class TestFieldDescriptor: """Field class: validation, defaults, default_factory checks.""" def test_compare_default_is_false(self) -> None: f = Field(name="x", _ty_schema=TypeSchema("int")) assert f.compare is False def test_default_and_factory_mutually_exclusive(self) -> None: with pytest.raises(ValueError): Field(name="x", _ty_schema=TypeSchema("int"), default=0, default_factory=lambda: 0) def test_factory_must_be_callable(self) -> None: with pytest.raises(TypeError, match="callable"): Field(name="x", _ty_schema=TypeSchema("int"), default_factory=0) # ty: ignore[invalid-argument-type] def test_non_callable_factory_rejected(self) -> None: with pytest.raises(TypeError, match="callable"): Field(name="x", _ty_schema=TypeSchema("int"), default_factory="not_callable") # ty: ignore[invalid-argument-type] # ########################################################################### # 4. Construction # ########################################################################### class TestConstruction: """Low-level __init__ via _make_type: positional/keyword args, defaults, errors.""" def test_keyword_args(self) -> None: Cls = _make_type( "ConKw", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("float"), default=MISSING), ], ) obj = Cls(x=42, y=3.14) assert obj.x == 42 assert obj.y == pytest.approx(3.14) def test_positional_args(self) -> None: Cls = _make_type( "ConPos", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("float"), default=MISSING), ], ) obj = Cls(10, 2.5) assert obj.x == 10 assert obj.y == pytest.approx(2.5) def test_mixed_positional_and_keyword(self) -> None: Cls = _make_type( "ConMixed", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("float"), default=MISSING), ], ) obj = Cls(7, y=1.5) assert obj.x == 7 assert obj.y == pytest.approx(1.5) def test_default_value_int(self) -> None: Cls = _make_type( "ConDefInt", [Field(name="x", _ty_schema=TypeSchema("int"), default=99)], ) assert Cls().x == 99 def test_default_value_float(self) -> None: Cls = _make_type( "ConDefFloat", [Field(name="x", _ty_schema=TypeSchema("float"), default=1.5)], ) assert Cls().x == pytest.approx(1.5) def test_default_value_str(self) -> None: Cls = _make_type( "ConDefStr", [Field(name="s", _ty_schema=TypeSchema("str"), default="hello")], ) assert Cls().s == "hello" def test_override_default(self) -> None: Cls = _make_type( "ConOverride", [Field(name="x", _ty_schema=TypeSchema("int"), default=0)], ) assert Cls(x=42).x == 42 def test_required_and_optional_together(self) -> None: Cls = _make_type( "ConReqOpt", [ Field(name="required", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="optional", _ty_schema=TypeSchema("float"), default=0.0), ], ) obj = Cls(required=5) assert obj.required == 5 assert obj.optional == pytest.approx(0.0) def test_missing_required_raises(self) -> None: Cls = _make_type( "ConMissing", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) with pytest.raises(TypeError): Cls() def test_extra_kwarg_raises(self) -> None: Cls = _make_type( "ConExtra", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) with pytest.raises(TypeError): Cls(x=1, bogus=2) def test_str_field_construction(self) -> None: Cls = _make_type( "ConStr", [Field(name="name", _ty_schema=TypeSchema("str"), default=MISSING)], ) assert Cls(name="world").name == "world" def test_bool_field_construction(self) -> None: Cls = _make_type( "ConBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING)], ) assert Cls(flag=True).flag is True assert Cls(flag=False).flag is False def test_kw_only_field(self) -> None: Cls = _make_type( "ConKwOnly", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="y", _ty_schema=TypeSchema("int"), default=MISSING, kw_only=True, ), ], ) obj = Cls(1, y=2) assert obj.x == 1 assert obj.y == 2 def test_kw_only_rejects_positional(self) -> None: Cls = _make_type( "ConKwOnlyReject", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="y", _ty_schema=TypeSchema("int"), default=MISSING, kw_only=True, ), ], ) with pytest.raises(TypeError): Cls(1, 2) def test_isinstance_check(self) -> None: Cls = _make_type( "ConIsInstance", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=1) assert isinstance(obj, Cls) assert isinstance(obj, core.Object) # ########################################################################### # 5. Getter / Setter # ########################################################################### class TestGetterSetter: """Field access: get/set POD, str, bool, mutation isolation.""" def test_get_int(self) -> None: Cls = _make_type( "GSInt", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) assert Cls(x=42).x == 42 def test_set_int(self) -> None: Cls = _make_type( "GSSetInt", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=1) obj.x = 100 assert obj.x == 100 def test_get_float(self) -> None: Cls = _make_type( "GSFloat", [Field(name="val", _ty_schema=TypeSchema("float"), default=MISSING)], ) assert Cls(val=3.14).val == pytest.approx(3.14) def test_set_float(self) -> None: Cls = _make_type( "GSSetFloat", [Field(name="val", _ty_schema=TypeSchema("float"), default=MISSING)], ) obj = Cls(val=1.0) obj.val = 2.718 assert obj.val == pytest.approx(2.718) def test_get_str(self) -> None: Cls = _make_type( "GSStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) assert Cls(s="hello").s == "hello" def test_set_str(self) -> None: Cls = _make_type( "GSSetStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) obj = Cls(s="hello") obj.s = "world" assert obj.s == "world" def test_get_bool(self) -> None: Cls = _make_type( "GSBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING)], ) assert Cls(flag=True).flag is True def test_set_bool(self) -> None: Cls = _make_type( "GSSetBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING)], ) obj = Cls(flag=True) obj.flag = False assert obj.flag is False def test_mutation_isolated(self) -> None: Cls = _make_type( "GSIsolate", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) a = Cls(x=1) b = Cls(x=1) a.x = 99 assert a.x == 99 assert b.x == 1 def test_multiple_fields_mutation(self) -> None: Cls = _make_type( "GSMultiMut", [ Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("float"), default=MISSING), Field(name="c", _ty_schema=TypeSchema("str"), default=MISSING), ], ) obj = Cls(a=1, b=2.0, c="x") obj.a = 10 obj.b = 20.0 obj.c = "y" assert obj.a == 10 assert obj.b == pytest.approx(20.0) assert obj.c == "y" def test_set_array_field(self) -> None: Cls = _make_type( "GSSetArr", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(arr=[1]) obj.arr = tvm_ffi.Array([4, 5, 6]) assert len(obj.arr) == 3 assert obj.arr[0] == 4 # ########################################################################### # 6. ObjectRef Fields # ########################################################################### class TestObjectRefFields: """Fields holding ObjectRef types: Array, custom objects.""" def test_array_field(self) -> None: Cls = _make_type( "ObjArr", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(arr=tvm_ffi.Array([1, 2, 3])) assert len(obj.arr) == 3 assert obj.arr[0] == 1 assert obj.arr[2] == 3 def test_array_field_from_list(self) -> None: Cls = _make_type( "ObjArrList", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(arr=[10, 20, 30]) assert len(obj.arr) == 3 assert obj.arr[1] == 20 def test_nested_object_field(self) -> None: Inner = _make_type( "ObjInner", [Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING)], ) inner_info = getattr(Inner, "__tvm_ffi_type_info__") inner_schema = TypeSchema(inner_info.type_key, origin_type_index=inner_info.type_index) Outer = _make_type( "ObjOuter", [Field(name="child", _ty_schema=inner_schema, default=MISSING)], ) assert Outer(child=Inner(val=42)).child.val == 42 # ########################################################################### # 7. Optional Fields # ########################################################################### class TestOptionalFields: """Optional/Union fields: None and non-None values.""" def test_optional_int_with_value(self) -> None: Cls = _make_type( "OptIntV", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=MISSING, ), ], ) assert Cls(x=42).x == 42 def test_optional_int_with_none(self) -> None: Cls = _make_type( "OptIntN", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), ], ) assert Cls().x is None def test_optional_str_with_value(self) -> None: Cls = _make_type( "OptStrV", [ Field( name="s", _ty_schema=TypeSchema("Optional", (TypeSchema("str"),)), default=MISSING, ), ], ) assert Cls(s="hello").s == "hello" def test_optional_str_with_none(self) -> None: Cls = _make_type( "OptStrN", [ Field( name="s", _ty_schema=TypeSchema("Optional", (TypeSchema("str"),)), default=None, ), ], ) assert Cls().s is None def test_optional_set_to_none(self) -> None: Cls = _make_type( "OptSet", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(x=42) obj.x = None assert obj.x is None def test_optional_set_back_to_value(self) -> None: Cls = _make_type( "OptBack", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), ], ) obj = Cls() obj.x = 99 assert obj.x == 99 def test_all_optional_fields_default_none(self) -> None: Cls = _make_type( "AllOpt", [ Field( name="a", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), Field( name="b", _ty_schema=TypeSchema("Optional", (TypeSchema("str"),)), default=None, ), Field( name="c", _ty_schema=TypeSchema("Optional", (TypeSchema("float"),)), default=None, ), ], ) obj = Cls() assert obj.a is None assert obj.b is None assert obj.c is None obj.a = 42 assert obj.a == 42 obj.b = "hello" assert obj.b == "hello" def test_optional_object_none_and_back(self) -> None: Cls = _make_type( "OptObjRound", [ Field( name="ref", _ty_schema=TypeSchema("Optional", (TypeSchema("Object"),)), default=None, ), ], ) obj = Cls() assert obj.ref is None obj.ref = tvm_ffi.Array([1]) assert len(obj.ref) == 1 obj.ref = None assert obj.ref is None def test_union_int_str(self) -> None: """Union[int, str] field should accept both types.""" Cls = _make_type( "UnionIntStr", [ Field( name="val", _ty_schema=TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))), default=MISSING, ), ], ) obj = Cls(val=42) assert obj.val == 42 obj.val = "hello" assert obj.val == "hello" def test_union_int_str_rejects_float(self) -> None: """Union[int, str] should reject float (not in union).""" Cls = _make_type( "UnionReject", [ Field( name="val", _ty_schema=TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))), default=MISSING, ), ], ) obj = Cls(val=1) with pytest.raises((TypeError, RuntimeError)): obj.val = 3.14 def test_optional_union(self) -> None: """Optional[Union[int, str]] should accept None, int, and str.""" Cls = _make_type( "OptUnion", [ Field( name="val", _ty_schema=TypeSchema( "Optional", (TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))),), ), default=None, ), ], ) obj = Cls() assert obj.val is None obj.val = 42 assert obj.val == 42 obj.val = "hi" assert obj.val == "hi" obj.val = None assert obj.val is None # ########################################################################### # 8. Any Fields # ########################################################################### class TestAnyField: """Fields with TypeSchema('Any'): hold any value type.""" def test_any_holds_int(self) -> None: Cls = _make_type( "AnyI", [Field(name="val", _ty_schema=TypeSchema("Any"), default=None)], ) assert Cls(val=42).val == 42 def test_any_holds_str(self) -> None: Cls = _make_type( "AnyS", [Field(name="val", _ty_schema=TypeSchema("Any"), default=None)], ) assert Cls(val="hello").val == "hello" def test_any_holds_none(self) -> None: Cls = _make_type( "AnyN", [Field(name="val", _ty_schema=TypeSchema("Any"), default=None)], ) assert Cls().val is None def test_any_holds_object(self) -> None: Cls = _make_type( "AnyObj", [Field(name="val", _ty_schema=TypeSchema("Any"), default=None)], ) arr = tvm_ffi.Array([1, 2]) assert len(Cls(val=arr).val) == 2 def test_any_type_change(self) -> None: Cls = _make_type( "AnyChg", [Field(name="val", _ty_schema=TypeSchema("Any"), default=None)], ) obj = Cls() obj.val = 42 assert obj.val == 42 obj.val = "hello" assert obj.val == "hello" obj.val = None assert obj.val is None obj.val = tvm_ffi.Array([1]) assert len(obj.val) == 1 # ########################################################################### # 9. Default Factory # ########################################################################### class TestDefaultFactory: """default_factory support: fresh instances, override, various types.""" def test_factory_produces_fresh_instances(self) -> None: Cls = _make_type( "DFFresh", [ Field( name="data", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default_factory=lambda: tvm_ffi.Array([]), ), ], ) a = Cls() b = Cls() assert not a.data.same_as(b.data) def test_factory_with_content(self) -> None: Cls = _make_type( "DFContent", [ Field( name="items", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default_factory=lambda: tvm_ffi.Array([1, 2, 3]), ), ], ) obj = Cls() assert len(obj.items) == 3 assert obj.items[0] == 1 def test_factory_override(self) -> None: Cls = _make_type( "DFOverride", [Field(name="x", _ty_schema=TypeSchema("int"), default_factory=lambda: 42)], ) assert Cls(x=99).x == 99 def test_factory_str(self) -> None: Cls = _make_type( "DFStr", [Field(name="s", _ty_schema=TypeSchema("str"), default_factory=lambda: "generated")], ) assert Cls().s == "generated" # ########################################################################### # 10. Repr # ########################################################################### class TestFieldRepr: """Low-level repr via _make_type: field values, repr=False exclusion.""" def test_repr_includes_fields(self) -> None: Cls = _make_type( "ReprBasic", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("float"), default=0.0), ], ) r = ReprPrint(Cls(x=42, y=3.14)) assert "x=42" in r assert "y=3.14" in r def test_repr_str_field(self) -> None: Cls = _make_type( "ReprStr", [Field(name="name", _ty_schema=TypeSchema("str"), default=MISSING)], ) assert '"hello"' in ReprPrint(Cls(name="hello")) def test_repr_bool_field(self) -> None: Cls = _make_type( "ReprBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING)], ) assert "flag=True" in ReprPrint(Cls(flag=True)) def test_repr_false_excluded(self) -> None: Cls = _make_type( "ReprExcl", [ Field(name="visible", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="hidden", _ty_schema=TypeSchema("int"), default=0, repr=False, ), ], ) r = ReprPrint(Cls(visible=42)) assert "visible=42" in r assert "hidden" not in r def test_python_repr_delegates(self) -> None: Cls = _make_type( "ReprDeleg", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) assert "x=7" in repr(Cls(x=7)) def test_repr_contains_type_key(self) -> None: Cls = _make_type( "ReprKey", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) info = getattr(Cls, "__tvm_ffi_type_info__") assert info.type_key in ReprPrint(Cls(x=1)) def test_repr_optional_none(self) -> None: Cls = _make_type( "ReprOptNone", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), ], ) r = ReprPrint(Cls()) assert isinstance(r, str) def test_repr_array_field(self) -> None: Cls = _make_type( "ReprArr", [ Field( name="items", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) r = ReprPrint(Cls(items=[1, 2, 3])) assert isinstance(r, str) # ########################################################################### # 11. Hash # ########################################################################### class TestFieldHash: """Low-level hash via _make_type: equal objects same hash, hash=False exclusion.""" def test_equal_objects_same_hash(self) -> None: Cls = _make_type( "HashEq", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="y", _ty_schema=TypeSchema("float"), default=MISSING, compare=True, ), ], eq=True, unsafe_hash=True, ) assert RecursiveHash(Cls(x=1, y=2.0)) == RecursiveHash(Cls(x=1, y=2.0)) def test_different_objects_different_hash(self) -> None: Cls = _make_type( "HashDiff", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, unsafe_hash=True, ) assert RecursiveHash(Cls(x=1)) != RecursiveHash(Cls(x=2)) def test_hash_false_field_ignored(self) -> None: Cls = _make_type( "HashIgn", [ Field( name="key", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="ignored", _ty_schema=TypeSchema("int"), default=0, hash=False, ), ], eq=True, unsafe_hash=True, ) assert RecursiveHash(Cls(key=42, ignored=100)) == RecursiveHash(Cls(key=42, ignored=999)) def test_hash_dunder_installed(self) -> None: Cls = _make_type( "HashDunder", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], eq=True, unsafe_hash=True, ) assert isinstance(hash(Cls(x=42)), int) def test_usable_as_dict_key(self) -> None: Cls = _make_type( "HashDict", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, unsafe_hash=True, ) assert {Cls(x=1): "value"}[Cls(x=1)] == "value" def test_usable_in_set(self) -> None: Cls = _make_type( "HashSet", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, unsafe_hash=True, ) assert len({Cls(x=1), Cls(x=1), Cls(x=2)}) == 2 # ########################################################################### # 12. Equality # ########################################################################### class TestFieldEquality: """Low-level equality via _make_type: structural compare, compare=False exclusion.""" def test_equal_objects(self) -> None: Cls = _make_type( "EqEqual", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="y", _ty_schema=TypeSchema("float"), default=MISSING, compare=True, ), ], eq=True, ) assert Cls(x=1, y=2.0) == Cls(x=1, y=2.0) def test_different_objects(self) -> None: Cls = _make_type( "EqDiff", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, ) assert Cls(x=1) != Cls(x=2) def test_compare_false_field_ignored(self) -> None: Cls = _make_type( "EqIgn", [ Field( name="key", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="ignored", _ty_schema=TypeSchema("int"), default=0, compare=False, ), ], eq=True, ) assert RecursiveEq(Cls(key=42, ignored=100), Cls(key=42, ignored=999)) def test_compare_off_excludes_from_eq(self) -> None: """Fields with compare=False (default) are ignored by RecursiveEq.""" Cls = _make_type( "CmpOff", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], eq=True, ) assert RecursiveEq(Cls(x=1), Cls(x=2)) def test_compare_true_includes_in_eq(self) -> None: Cls = _make_type( "CmpOn", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, ) assert not RecursiveEq(Cls(x=1), Cls(x=2)) def test_eq_reflexive(self) -> None: Cls = _make_type( "EqRefl", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, ) a = Cls(x=42) assert a == a def test_eq_symmetric(self) -> None: Cls = _make_type( "EqSym", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), ], eq=True, ) a, b = Cls(x=1), Cls(x=1) assert a == b assert b == a def test_eq_with_str_field(self) -> None: Cls = _make_type( "EqStr", [ Field( name="s", _ty_schema=TypeSchema("str"), default=MISSING, compare=True, ), ], eq=True, ) assert RecursiveEq(Cls(s="hello"), Cls(s="hello")) assert not RecursiveEq(Cls(s="hello"), Cls(s="world")) def test_eq_hash_consistency(self) -> None: Cls = _make_type( "EqHashCon", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="y", _ty_schema=TypeSchema("float"), default=MISSING, compare=True, ), ], eq=True, unsafe_hash=True, ) a, b = Cls(x=1, y=2.0), Cls(x=1, y=2.0) assert RecursiveEq(a, b) assert RecursiveHash(a) == RecursiveHash(b) # ########################################################################### # 13. Edge Cases # ########################################################################### class TestFieldEdgeCases: """Low-level edge cases via _make_type: empty class, extreme values, init=False.""" def test_empty_class_no_fields(self) -> None: Cls = _make_type("EdgeEmpty", []) obj = Cls() assert isinstance(obj, core.Object) assert isinstance(obj, Cls) def test_empty_class_repr(self) -> None: Cls = _make_type("EdgeEmptyRepr", []) info = getattr(Cls, "__tvm_ffi_type_info__") assert info.type_key in ReprPrint(Cls()) def test_bool_true_and_false(self) -> None: Cls = _make_type( "EdgeBool", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING)], ) assert Cls(flag=True).flag is True assert Cls(flag=False).flag is False def test_bool_default_false(self) -> None: Cls = _make_type( "EdgeBoolDef", [Field(name="flag", _ty_schema=TypeSchema("bool"), default=False)], ) assert Cls().flag is False def test_multiple_types_together(self) -> None: Cls = _make_type( "EdgeMulti", [ Field( name="i", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING), Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("bool"), default=MISSING), ], ) obj = Cls(i=42, f=3.14, s="test", b=True) assert obj.i == 42 assert obj.f == pytest.approx(3.14) assert obj.s == "test" assert obj.b is True def test_pod_and_objectref_mixed(self) -> None: Cls = _make_type( "EdgeMixed", [ Field(name="count", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="items", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), Field(name="label", _ty_schema=TypeSchema("str"), default=""), ], ) obj = Cls(count=3, items=[1, 2, 3]) assert obj.count == 3 assert len(obj.items) == 3 assert obj.label == "" def test_multiple_types_with_defaults(self) -> None: Cls = _make_type( "EdgeMultiDef", [ Field(name="i", _ty_schema=TypeSchema("int"), default=0), Field(name="f", _ty_schema=TypeSchema("float"), default=1.0), Field(name="s", _ty_schema=TypeSchema("str"), default="default"), Field(name="b", _ty_schema=TypeSchema("bool"), default=True), ], ) obj = Cls() assert obj.i == 0 assert obj.f == pytest.approx(1.0) assert obj.s == "default" assert obj.b is True def test_zero_values(self) -> None: Cls = _make_type( "EdgeZero", [ Field(name="i", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING), ], ) obj = Cls(i=0, f=0.0) assert obj.i == 0 assert obj.f == 0.0 def test_negative_values(self) -> None: Cls = _make_type( "EdgeNeg", [ Field(name="i", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING), ], ) obj = Cls(i=-42, f=-3.14) assert obj.i == -42 assert obj.f == pytest.approx(-3.14) def test_large_int(self) -> None: Cls = _make_type( "EdgeLargeInt", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) large = 2**62 assert Cls(x=large).x == large def test_empty_string_field(self) -> None: Cls = _make_type( "EdgeEmptyStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) assert Cls(s="").s == "" def test_long_string_field(self) -> None: Cls = _make_type( "EdgeLongStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) long_str = "a" * 1000 assert Cls(s=long_str).s == long_str def test_equality_empty_class(self) -> None: Cls = _make_type("EdgeEmptyEq", [], eq=True, unsafe_hash=True) assert RecursiveEq(Cls(), Cls()) assert RecursiveHash(Cls()) == RecursiveHash(Cls()) def test_init_false_field_excluded_from_init(self) -> None: Cls = _make_type( "EdgeInitFalse", [ Field(name="visible", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="internal", _ty_schema=TypeSchema("int"), default=0, init=False, ), ], ) obj = Cls(visible=42) assert obj.visible == 42 assert obj.internal == 0 def test_init_false_field_rejected_as_kwarg(self) -> None: Cls = _make_type( "EdgeInitFalseReject", [ Field(name="visible", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="internal", _ty_schema=TypeSchema("int"), default=0, init=False, ), ], ) with pytest.raises(TypeError): Cls(visible=1, internal=2) def test_init_false_field_writable(self) -> None: Cls = _make_type( "EdgeInitFalseWrite", [ Field(name="visible", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="internal", _ty_schema=TypeSchema("int"), default=0, init=False, ), ], ) obj = Cls(visible=1) obj.internal = 99 assert obj.internal == 99 # ########################################################################### # 14. Inheritance (Python-defined parent) # ########################################################################### class TestFieldInheritance: """Low-level inheritance via _make_type: field offsets, parent-child layout.""" def test_child_fields_after_parent(self) -> None: Parent = _make_type( "InhParent", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) Child = _make_type( "InhChild", [Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING)], parent=Parent, ) obj = Child(1, 2) assert obj.x == 1 assert obj.y == 2 def test_child_field_offsets_non_overlapping(self) -> None: Parent = _make_type( "InhParentOff", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) Child = _make_type( "InhChildOff", [Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING)], parent=Parent, ) p_info = getattr(Parent, "__tvm_ffi_type_info__") c_info = getattr(Child, "__tvm_ffi_type_info__") parent_end = p_info.fields[0].offset + p_info.fields[0].size assert c_info.fields[0].offset >= parent_end def test_mutation_no_aliasing(self) -> None: Parent = _make_type( "InhParentAlias", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) Child = _make_type( "InhChildAlias", [Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING)], parent=Parent, ) obj = Child(1, 2) obj.y = 9 assert obj.x == 1 assert obj.y == 9 def test_three_level_inheritance(self) -> None: """Object → A → B → C: all fields accessible and non-overlapping.""" A = _make_type( "InhA", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], ) B = _make_type( "InhB", [Field(name="b", _ty_schema=TypeSchema("str"), default=MISSING)], parent=A, ) C = _make_type( "InhC", [Field(name="c", _ty_schema=TypeSchema("float"), default=MISSING)], parent=B, ) obj = C(a=1, b="two", c=3.0) assert obj.a == 1 assert obj.b == "two" assert obj.c == pytest.approx(3.0) def test_three_level_offsets_non_overlapping(self) -> None: A = _make_type( "InhAOff", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], ) B = _make_type( "InhBOff", [Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING)], parent=A, ) C = _make_type( "InhCOff", [Field(name="c", _ty_schema=TypeSchema("int"), default=MISSING)], parent=B, ) a_info = getattr(A, "__tvm_ffi_type_info__") b_info = getattr(B, "__tvm_ffi_type_info__") c_info = getattr(C, "__tvm_ffi_type_info__") a_end = a_info.fields[0].offset + a_info.fields[0].size b_end = b_info.fields[0].offset + b_info.fields[0].size assert b_info.fields[0].offset >= a_end assert c_info.fields[0].offset >= b_end def test_three_level_mutation_no_aliasing(self) -> None: A = _make_type( "InhAMut", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], ) B = _make_type( "InhBMut", [Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING)], parent=A, ) C = _make_type( "InhCMut", [Field(name="c", _ty_schema=TypeSchema("int"), default=MISSING)], parent=B, ) obj = C(a=1, b=2, c=3) obj.c = 99 assert obj.a == 1 assert obj.b == 2 assert obj.c == 99 obj.a = 77 assert obj.a == 77 assert obj.b == 2 assert obj.c == 99 def test_three_level_isinstance(self) -> None: A = _make_type( "InhAIs", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], ) B = _make_type( "InhBIs", [Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING)], parent=A, ) C = _make_type( "InhCIs", [Field(name="c", _ty_schema=TypeSchema("int"), default=MISSING)], parent=B, ) obj = C(a=1, b=2, c=3) assert isinstance(obj, C) assert isinstance(obj, B) assert isinstance(obj, A) assert isinstance(obj, core.Object) def test_three_level_deep_copy(self) -> None: A = _make_type( "InhACopy", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], ) B = _make_type( "InhBCopy", [Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING)], parent=A, ) C = _make_type( "InhCCopy", [Field(name="c", _ty_schema=TypeSchema("int"), default=MISSING)], parent=B, ) obj = C(a=1, b=2, c=3) obj_copy = DeepCopy(obj) assert not obj.same_as(obj_copy) assert obj_copy.a == 1 assert obj_copy.b == 2 assert obj_copy.c == 3 obj_copy.c = 99 assert obj.c == 3 # -- Regression: empty intermediate parent must propagate total_size ----- def test_empty_intermediate_parent_offsets(self) -> None: """GrandParent(x,y) -> Parent(no fields) -> Child(a,b): no overlap.""" GP = _make_type( "OffGP", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING), ], ) P = _make_type("OffP", [], parent=GP) C = _make_type( "OffC", [ Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING), ], parent=P, ) obj = C(x=10, y=20, a=30, b=40) assert obj.x == 10 assert obj.y == 20 assert obj.a == 30 assert obj.b == 40 def test_empty_intermediate_parent_total_size(self) -> None: """Parent with no own fields inherits total_size from its parent.""" GP = _make_type( "OffGP2", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING), ], ) P = _make_type("OffP2", [], parent=GP) gp_info = getattr(GP, "__tvm_ffi_type_info__") p_info = getattr(P, "__tvm_ffi_type_info__") assert p_info.total_size >= gp_info.total_size def test_empty_intermediate_child_offsets_non_overlapping(self) -> None: """Child field offsets must start at or after Parent.total_size.""" GP = _make_type( "OffGP3", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING), ], ) P = _make_type("OffP3", [], parent=GP) C = _make_type( "OffC3", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], parent=P, ) p_info = getattr(P, "__tvm_ffi_type_info__") c_info = getattr(C, "__tvm_ffi_type_info__") assert c_info.fields[0].offset >= p_info.total_size def test_two_empty_intermediate_parents(self) -> None: """GP(x) -> Mid1() -> Mid2() -> Leaf(a): offsets propagated through two empty layers.""" GP = _make_type( "OffGP4", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) Mid1 = _make_type("OffMid1", [], parent=GP) Mid2 = _make_type("OffMid2", [], parent=Mid1) Leaf = _make_type( "OffLeaf", [Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING)], parent=Mid2, ) obj = Leaf(x=1, a=2) assert obj.x == 1 assert obj.a == 2 gp_info = getattr(GP, "__tvm_ffi_type_info__") leaf_info = getattr(Leaf, "__tvm_ffi_type_info__") assert leaf_info.fields[0].offset >= gp_info.total_size def test_empty_intermediate_mutation_no_aliasing(self) -> None: """Mutating Child.a must not corrupt GrandParent.x through an empty parent.""" GP = _make_type( "OffGP5", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("int"), default=MISSING), ], ) P = _make_type("OffP5", [], parent=GP) C = _make_type( "OffC5", [ Field(name="a", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("int"), default=MISSING), ], parent=P, ) obj = C(x=10, y=20, a=30, b=40) obj.a = 99 assert obj.x == 10 assert obj.y == 20 obj.x = 77 assert obj.a == 99 assert obj.b == 40 def test_empty_intermediate_py_class_decorator(self) -> None: """Same bug via @py_class (the high-level API).""" @py_class(_unique_key("OffDecGP")) class GP(Object): x: int = 0 y: int = 0 @py_class(_unique_key("OffDecP")) class P(GP): pass @py_class(_unique_key("OffDecC")) class C(P): a: int = 0 b: int = 0 obj = C(x=10, y=20, a=30, b=40) assert obj.x == 10 assert obj.y == 20 assert obj.a == 30 assert obj.b == 40 # ########################################################################### # 15. Mutual / Self References # ########################################################################### class TestMutualReferences: """Low-level mutual and self-referential type fields via two-phase registration.""" def _register_bare(self, name: str) -> tuple[type, core.TypeInfo]: """Register a type with no fields (phase 1 of two-phase).""" parent_info = core._type_cls_to_type_info(core.Object) assert parent_info is not None cls = type(name, (core.Object,), {"__slots__": ()}) info = core._register_py_class(parent_info, _unique_key_ff(name), cls) return cls, info def _finalize(self, cls: type, info: core.TypeInfo, fields: List[Field]) -> None: """Register fields and install class attrs (phase 2 of two-phase).""" info._register_fields(fields) setattr(cls, "__tvm_ffi_type_info__", info) _add_class_attrs(cls, info) _install_dataclass_dunders( cls, init=True, repr=True, eq=False, order=False, unsafe_hash=False, ) def test_mutual_references(self) -> None: """Foo has Optional[Bar], Bar has Optional[Foo].""" Foo, foo_info = self._register_bare("MutFoo") Bar, bar_info = self._register_bare("MutBar") foo_schema = TypeSchema(foo_info.type_key, origin_type_index=foo_info.type_index) bar_schema = TypeSchema(bar_info.type_key, origin_type_index=bar_info.type_index) self._finalize( Foo, foo_info, [ Field(name="a", _ty_schema=TypeSchema("str"), default=MISSING), Field( name="bar", _ty_schema=TypeSchema("Optional", (bar_schema,)), default=None, ), ], ) self._finalize( Bar, bar_info, [ Field( name="foo", _ty_schema=TypeSchema("Optional", (foo_schema,)), default=None, ), ], ) foo = Foo(a="hello") bar = Bar() bar.foo = foo foo.bar = bar assert foo.bar.foo.a == "hello" def test_self_referential_field(self) -> None: """Bar has Optional[Bar] (self-reference).""" Bar, bar_info = self._register_bare("SelfRef") bar_schema = TypeSchema(bar_info.type_key, origin_type_index=bar_info.type_index) self._finalize( Bar, bar_info, [ Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="next", _ty_schema=TypeSchema("Optional", (bar_schema,)), default=None, ), ], ) a = Bar(val=1) b = Bar(val=2, next=a) assert b.next.val == 1 assert a.next is None # Circular: a → b → a a.next = b assert a.next.next.val == 1 def test_typed_mutual_ref_rejects_wrong_type(self) -> None: """Optional[Foo] field should reject Bar objects.""" Foo, foo_info = self._register_bare("TypedFoo") Bar, bar_info = self._register_bare("TypedBar") foo_schema = TypeSchema(foo_info.type_key, origin_type_index=foo_info.type_index) self._finalize( Foo, foo_info, [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), ], ) self._finalize( Bar, bar_info, [ Field( name="foo", _ty_schema=TypeSchema("Optional", (foo_schema,)), default=None, ), ], ) bar = Bar() bar.foo = Foo(x=1) # OK assert bar.foo.x == 1 with pytest.raises((TypeError, RuntimeError)): bar.foo = bar # Bar is not Foo # ########################################################################### # 16. Inheritance (native C++ parent) # ########################################################################### class TestNativeParentInheritance: """Low-level Python child of C++ TestObjectBase: offsets, fields, methods, copy.""" def test_non_overlapping_offsets(self) -> None: parent_info = core._type_cls_to_type_info(_TestObjectBase) assert parent_info is not None Child = _make_type( "InhNativeChild", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) child_info = getattr(Child, "__tvm_ffi_type_info__") parent_end = max(f.offset + f.size for f in parent_info.fields) assert child_info.fields[0].offset >= parent_end def test_preserves_parent_fields(self) -> None: Child = _make_type( "InhNativePreserve", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) obj = Child(extra=7, v_i64=1, v_f64=2.0, v_str="x") assert obj.extra == 7 assert obj.v_i64 == 1 assert obj.v_f64 == 2.0 assert obj.v_str == "x" def test_mutation_no_aliasing(self) -> None: Child = _make_type( "InhNativeMut", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) obj = Child(extra=7, v_i64=1, v_f64=2.0, v_str="x") obj.extra = 33 assert obj.extra == 33 assert obj.v_i64 == 1 assert obj.v_f64 == 2.0 assert obj.v_str == "x" def test_parent_method_uses_parent_state(self) -> None: Child = _make_type( "InhNativeMethod", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) obj = Child(extra=7, v_i64=1, v_f64=2.0, v_str="x") assert obj.add_i64(5) == 6 def test_copy_preserves_all_fields(self) -> None: Child = _make_type( "InhNativeCopy", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) obj = Child(extra=7, v_i64=1, v_f64=2.0, v_str="x") obj_copy = copy.copy(obj) assert obj_copy.extra == 7 assert obj_copy.v_i64 == 1 assert obj_copy.v_f64 == 2.0 assert obj_copy.v_str == "x" def test_deepcopy_preserves_all_fields(self) -> None: Child = _make_type( "InhNativeDeepCopy", [Field(name="extra", _ty_schema=TypeSchema("int"), default=MISSING)], parent=_TestObjectBase, ) obj = Child(extra=7, v_i64=1, v_f64=2.0, v_str="x") obj_copy = copy.deepcopy(obj) assert obj_copy.extra == 7 assert obj_copy.v_i64 == 1 assert obj_copy.v_f64 == 2.0 assert obj_copy.v_str == "x" # ########################################################################### # 16. Deep Copy # ########################################################################### class TestDeepCopy: """Low-level DeepCopy via _make_type: nested ObjectRef, mutation independence.""" def test_deep_copy_basic(self) -> None: Cls = _make_type( "DCBasic", [ Field( name="x", _ty_schema=TypeSchema("int"), default=MISSING, compare=True, ), Field( name="s", _ty_schema=TypeSchema("str"), default=MISSING, compare=True, ), ], eq=True, ) obj = Cls(x=42, s="hello") obj_copy = DeepCopy(obj) assert not obj.same_as(obj_copy) assert RecursiveEq(obj, obj_copy) def test_deep_copy_nested_objectref(self) -> None: Cls = _make_type( "DCNested", [ Field( name="items", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(items=tvm_ffi.Array([1, 2, 3])) obj_copy = DeepCopy(obj) assert not obj.items.same_as(obj_copy.items) assert len(obj_copy.items) == 3 def test_deep_copy_mutate_independent(self) -> None: Cls = _make_type( "DCMut", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=1) obj_copy = DeepCopy(obj) obj_copy.x = 99 assert obj.x == 1 def test_python_deepcopy_dunder(self) -> None: Cls = _make_type( "DCPython", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=42) obj_copy = copy.deepcopy(obj) assert obj_copy.x == 42 assert not obj.same_as(obj_copy) # ########################################################################### # 17. Memory / Lifetime # ########################################################################### class TestMemoryLifetime: """Low-level ref-counting: ObjectRef/Any fields are properly ref-counted.""" def test_objectref_field_kept_alive(self) -> None: Cls = _make_type( "MemAlive", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) arr = tvm_ffi.Array([1, 2, 3]) obj = Cls(arr=arr) del arr gc.collect() assert len(obj.arr) == 3 def test_multiple_objects_independent_lifetime(self) -> None: Cls = _make_type( "MemIndep", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) shared = tvm_ffi.Array([10, 20]) a = Cls(arr=shared) b = Cls(arr=shared) del a gc.collect() assert len(b.arr) == 2 assert b.arr[0] == 10 def test_str_field_any_storage(self) -> None: Cls = _make_type( "MemStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) assert Cls(s="hi").s == "hi" long_str = "a" * 500 assert Cls(s=long_str).s == long_str # ########################################################################### # 18. Bool Alignment # ########################################################################### class TestBoolAlignment: """Low-level bool field layout: 1-byte packing, padding, alternating layouts.""" def test_bool_then_int_alignment(self) -> None: Cls = _make_type( "BoolAlign", [ Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING), ], ) info = getattr(Cls, "__tvm_ffi_type_info__") assert info.fields[0].offset == 24 assert info.fields[1].offset % 8 == 0 assert info.fields[1].offset >= info.fields[0].offset + 1 def test_bool_then_int_values(self) -> None: Cls = _make_type( "BoolAlignVal", [ Field(name="flag", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING), ], ) obj = Cls(flag=True, val=42) assert obj.flag is True assert obj.val == 42 obj.flag = False assert obj.flag is False assert obj.val == 42 def test_multiple_bools_packed(self) -> None: Cls = _make_type( "MultiBool", [ Field(name="a", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="c", _ty_schema=TypeSchema("bool"), default=MISSING), ], ) info = getattr(Cls, "__tvm_ffi_type_info__") assert [f.offset for f in info.fields] == [24, 25, 26] obj = Cls(a=True, b=False, c=True) assert obj.a is True assert obj.b is False assert obj.c is True def test_bool_int_bool_int_alternating(self) -> None: Cls = _make_type( "BoolIntBoolInt", [ Field(name="b1", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="i1", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="b2", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="i2", _ty_schema=TypeSchema("int"), default=MISSING), ], ) obj = Cls(b1=True, i1=100, b2=False, i2=200) assert obj.b1 is True assert obj.i1 == 100 assert obj.b2 is False assert obj.i2 == 200 # ########################################################################### # 19. Type Conversion Errors # ########################################################################### class TestTypeConversionErrors: """Low-level type conversion: wrong-type setter/construction raises.""" def test_set_int_field_to_str_raises(self) -> None: Cls = _make_type( "ErrIntStr", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=1) with pytest.raises((TypeError, RuntimeError)): obj.x = "not_an_int" def test_set_str_field_to_int_raises(self) -> None: Cls = _make_type( "ErrStrInt", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) obj = Cls(s="hello") with pytest.raises((TypeError, RuntimeError)): obj.s = 42 def test_construct_with_wrong_type_raises(self) -> None: Cls = _make_type( "ErrInit", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) with pytest.raises((TypeError, RuntimeError)): Cls(x="bad") def test_set_wrong_type_preserves_old_value(self) -> None: """Failed type-checked mutation preserves old value.""" Cls = _make_type( "ErrPreserve", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=42) with pytest.raises((TypeError, RuntimeError)): obj.x = "bad_value" assert obj.x == 42 def test_type_schema_convert_raises_directly(self) -> None: """TypeSchema.convert raises TypeError for incompatible values.""" ts = TypeSchema("int") assert _to_py_class_value(ts.convert(42)) == 42 with pytest.raises(TypeError): ts.convert("not_an_int") def test_set_non_optional_object_field_to_none_raises(self) -> None: """A non-Optional Object field must reject None.""" Cls = _make_type( "ErrObjNone", [ Field( name="child", _ty_schema=TypeSchema("Object"), default=MISSING, ), ], ) obj = Cls(child=tvm_ffi.Array([1])) with pytest.raises((TypeError, RuntimeError)): obj.child = None def test_construct_non_optional_object_field_with_none_raises(self) -> None: """Constructing with None for a non-Optional Object field must fail.""" Cls = _make_type( "ErrObjNoneInit", [ Field( name="child", _ty_schema=TypeSchema("Object"), default=MISSING, ), ], ) with pytest.raises((TypeError, RuntimeError)): Cls(child=None) def test_optional_object_field_accepts_none(self) -> None: """An Optional[Object] field should accept None.""" Cls = _make_type( "OptObjNone", [ Field( name="child", _ty_schema=TypeSchema("Optional", (TypeSchema("Object"),)), default=None, ), ], ) obj = Cls() assert obj.child is None obj.child = tvm_ffi.Array([1, 2]) assert len(obj.child) == 2 obj.child = None assert obj.child is None def test_set_bool_field_to_str_raises(self) -> None: Cls = _make_type( "ErrBoolStr", [Field(name="b", _ty_schema=TypeSchema("bool"), default=MISSING)], ) obj = Cls(b=True) with pytest.raises((TypeError, RuntimeError)): obj.b = "not_a_bool" def test_set_array_field_to_int_raises(self) -> None: Cls = _make_type( "ErrArrInt", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(arr=[1]) with pytest.raises((TypeError, RuntimeError)): obj.arr = 42 def test_construct_multiple_wrong_types_first_caught(self) -> None: """When the first field has a wrong type, the error is caught.""" Cls = _make_type( "ErrMulti", [ Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="y", _ty_schema=TypeSchema("str"), default=MISSING), ], ) with pytest.raises((TypeError, RuntimeError)): Cls(x="bad", y="ok") def test_set_optional_to_wrong_inner_type_raises(self) -> None: Cls = _make_type( "ErrOptWrong", [ Field( name="x", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), ], ) obj = Cls() with pytest.raises((TypeError, RuntimeError)): obj.x = "not_an_int" # ########################################################################### # 20. Setter / Getter Corner Cases # ########################################################################### class TestSetterGetterCornerCases: """Low-level setter/getter corner cases: conversions, nesting, edge values.""" # --- Bool / int coercion --- def test_bool_field_accepts_true_false(self) -> None: Cls = _make_type( "SGBool", [Field(name="b", _ty_schema=TypeSchema("bool"), default=MISSING)], ) obj = Cls(b=True) assert obj.b is True obj.b = False assert obj.b is False def test_int_field_accepts_bool(self) -> None: """Python bool is a subclass of int — FFI should accept it.""" Cls = _make_type( "SGIntBool", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=True) assert obj.x == 1 obj.x = False assert obj.x == 0 # --- Float edge values --- def test_float_field_inf_nan(self) -> None: Cls = _make_type( "SGFloatEdge", [Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING)], ) obj = Cls(f=float("inf")) assert math.isinf(obj.f) obj.f = float("-inf") assert math.isinf(obj.f) and obj.f < 0 obj.f = float("nan") assert math.isnan(obj.f) def test_float_field_accepts_int(self) -> None: Cls = _make_type( "SGFloatInt", [Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING)], ) obj = Cls(f=42) assert obj.f == pytest.approx(42.0) # --- String edge values --- def test_str_field_unicode(self) -> None: Cls = _make_type( "SGStrUni", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) obj = Cls(s="日本語テスト 🎉") assert obj.s == "日本語テスト 🎉" def test_str_field_null_bytes(self) -> None: Cls = _make_type( "SGStrNull", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) s = "hello\x00world" obj = Cls(s=s) assert obj.s == s # --- Multiple mutations --- def test_repeated_mutation_same_field(self) -> None: Cls = _make_type( "SGRepeat", [Field(name="x", _ty_schema=TypeSchema("int"), default=MISSING)], ) obj = Cls(x=0) for i in range(100): obj.x = i assert obj.x == 99 def test_repeated_str_mutation(self) -> None: """Stress: repeated str assignment should not leak.""" Cls = _make_type( "SGRepeatStr", [Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING)], ) obj = Cls(s="init") for i in range(100): obj.s = f"value_{i}" assert obj.s == "value_99" def test_repeated_objectref_mutation(self) -> None: """Stress: repeated ObjectRef assignment should properly DecRef old values.""" Cls = _make_type( "SGRepeatArr", [ Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), ], ) obj = Cls(arr=[0]) for i in range(50): obj.arr = tvm_ffi.Array([i]) assert obj.arr[0] == 49 # --- Nested object fields --- def test_nested_two_levels(self) -> None: Inner = _make_type( "SGInner", [Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING)], ) inner_info = getattr(Inner, "__tvm_ffi_type_info__") inner_schema = TypeSchema(inner_info.type_key, origin_type_index=inner_info.type_index) Outer = _make_type( "SGOuter", [Field(name="child", _ty_schema=inner_schema, default=MISSING)], ) obj = Outer(child=Inner(val=42)) assert obj.child.val == 42 # Mutate inner through outer new_inner = Inner(val=99) obj.child = new_inner assert obj.child.val == 99 def test_self_referential_optional_field(self) -> None: """A type with an Optional[Self] field (stored as Any).""" Cls = _make_type( "SGSelfRef", [ Field(name="val", _ty_schema=TypeSchema("int"), default=MISSING), Field( name="next", _ty_schema=TypeSchema("Optional", (TypeSchema("Object"),)), default=None, ), ], ) a = Cls(val=1) b = Cls(val=2, next=a) assert b.val == 2 assert b.next.val == 1 assert a.next is None # --- Default factory edge cases --- def test_default_factory_called_each_time(self) -> None: call_count = [0] def factory() -> int: call_count[0] += 1 return call_count[0] Cls = _make_type( "SGFactoryCount", [Field(name="x", _ty_schema=TypeSchema("int"), default_factory=factory)], ) a = Cls() b = Cls() c = Cls() assert a.x == 1 assert b.x == 2 assert c.x == 3 # --- Mixed types in one class --- def test_all_pod_plus_objectref_plus_optional(self) -> None: Cls = _make_type( "SGKitchenSink", [ Field(name="i", _ty_schema=TypeSchema("int"), default=MISSING), Field(name="f", _ty_schema=TypeSchema("float"), default=MISSING), Field(name="b", _ty_schema=TypeSchema("bool"), default=MISSING), Field(name="s", _ty_schema=TypeSchema("str"), default=MISSING), Field( name="arr", _ty_schema=TypeSchema("Array", (TypeSchema("int"),)), default=MISSING, ), Field( name="opt", _ty_schema=TypeSchema("Optional", (TypeSchema("int"),)), default=None, ), ], ) obj = Cls(i=1, f=2.0, b=True, s="hi", arr=[10, 20]) assert obj.i == 1 assert obj.f == pytest.approx(2.0) assert obj.b is True assert obj.s == "hi" assert len(obj.arr) == 2 assert obj.opt is None # Mutate all fields obj.i = -1 obj.f = -2.0 obj.b = False obj.s = "bye" obj.arr = tvm_ffi.Array([30]) obj.opt = 42 assert obj.i == -1 assert obj.f == pytest.approx(-2.0) assert obj.b is False assert obj.s == "bye" assert len(obj.arr) == 1 assert obj.opt == 42 # ########################################################################### # 21. FFI Global Function Existence # ########################################################################### class TestFFIGlobalFunctions: """Low-level FFI global function registration checks.""" def test_py_class_register_type_attr_columns_exists(self) -> None: assert ( tvm_ffi.get_global_func("ffi._PyClassRegisterTypeAttrColumns", allow_missing=True) is not None ) def test_get_field_getter_exists(self) -> None: assert tvm_ffi.get_global_func("ffi.MakeFieldGetter", allow_missing=True) is not None def test_make_field_setter_exists(self) -> None: assert tvm_ffi.get_global_func("ffi.MakeFieldSetter", allow_missing=True) is not None def test_make_new_removed(self) -> None: assert tvm_ffi.get_global_func("ffi.MakeNew", allow_missing=True) is None # ########################################################################### # 22. Container field annotations # ########################################################################### class TestContainerFieldAnnotations: """Container field annotations: List[T], Dict[K,V], nested.""" def test_list_int_field(self) -> None: @py_class(_unique_key("ListInt")) class ListInt(Object): items: List[int] obj = ListInt(items=[1, 2, 3]) assert len(obj.items) == 3 assert obj.items[0] == 1 assert obj.items[2] == 3 def test_list_int_from_tuple(self) -> None: @py_class(_unique_key("ListIntTup")) class ListIntTup(Object): items: List[int] obj = ListIntTup(items=(10, 20, 30)) # ty:ignore[invalid-argument-type] assert len(obj.items) == 3 assert obj.items[1] == 20 def test_dict_str_int_field(self) -> None: @py_class(_unique_key("DictStrInt")) class DictStrInt(Object): mapping: Dict[str, int] obj = DictStrInt(mapping={"a": 1, "b": 2}) assert len(obj.mapping) == 2 assert obj.mapping["a"] == 1 assert obj.mapping["b"] == 2 def test_list_list_int_field(self) -> None: @py_class(_unique_key("ListListInt")) class ListListInt(Object): matrix: List[List[int]] obj = ListListInt(matrix=[[1, 2, 3], [4, 5, 6]]) assert len(obj.matrix) == 2 assert len(obj.matrix[0]) == 3 assert obj.matrix[0][0] == 1 assert obj.matrix[1][2] == 6 def test_dict_str_list_int_field(self) -> None: @py_class(_unique_key("DictStrListInt")) class DictStrListInt(Object): data: Dict[str, List[int]] obj = DictStrListInt(data={"x": [1, 2, 3], "y": [4, 5, 6]}) assert len(obj.data) == 2 assert tuple(obj.data["x"]) == (1, 2, 3) assert tuple(obj.data["y"]) == (4, 5, 6) def test_container_field_set(self) -> None: @py_class(_unique_key("ContSet")) class ContSet(Object): items: List[int] obj = ContSet(items=[1, 2]) assert tuple(obj.items) == (1, 2) obj.items = [3, 4, 5] assert len(obj.items) == 3 assert obj.items[0] == 3 def test_dict_field_set(self) -> None: @py_class(_unique_key("DictSet")) class DictSet(Object): mapping: Dict[str, int] obj = DictSet(mapping={"a": 1}) obj.mapping = {"b": 2, "c": 3} assert len(obj.mapping) == 2 assert obj.mapping["b"] == 2 def test_container_shared_reference(self) -> None: @py_class(_unique_key("ContShare")) class ContShare(Object): a: List[int] b: List[int] obj = ContShare(a=[1, 2], b=[1, 2]) assert tuple(obj.a) == tuple(obj.b) def test_untyped_list_field(self) -> None: @py_class(_unique_key("UList")) class UList(Object): items: list obj = UList(items=[1, "two", 3.0]) assert len(obj.items) == 3 assert obj.items[0] == 1 assert obj.items[1] == "two" def test_untyped_dict_field(self) -> None: @py_class(_unique_key("UDict")) class UDict(Object): data: dict obj = UDict(data={"a": 1, "b": "two"}) assert len(obj.data) == 2 # ########################################################################### # 23. Optional container fields # ########################################################################### class TestOptionalContainerFields: """Optional[List[T]], Optional[Dict[K,V]] via @py_class.""" @requires_py310 def test_optional_list_int(self) -> None: @py_class(_unique_key("OptListInt")) class OptListInt(Object): items: Optional[List[int]] obj = OptListInt(items=[1, 2, 3]) assert len(obj.items) == 3 # ty:ignore[invalid-argument-type] obj.items = None assert obj.items is None obj.items = [4, 5] assert len(obj.items) == 2 @requires_py310 def test_optional_dict_str_int(self) -> None: @py_class(_unique_key("OptDictStrInt")) class OptDictStrInt(Object): data: Optional[Dict[str, int]] obj = OptDictStrInt(data={"a": 1}) assert obj.data["a"] == 1 # ty:ignore[not-subscriptable] obj.data = None assert obj.data is None obj.data = {"b": 2} assert obj.data["b"] == 2 @requires_py310 def test_optional_list_list_int(self) -> None: @py_class(_unique_key("OptLLI")) class OptLLI(Object): matrix: Optional[List[List[int]]] obj = OptLLI(matrix=[[1, 2], [3, 4]]) assert obj.matrix[0][0] == 1 # ty:ignore[not-subscriptable] obj.matrix = None assert obj.matrix is None @requires_py310 def test_optional_dict_str_list_int(self) -> None: @py_class(_unique_key("OptDSLI")) class OptDSLI(Object): data: Optional[Dict[str, List[int]]] obj = OptDSLI(data={"x": [1, 2, 3]}) assert tuple(obj.data["x"]) == (1, 2, 3) # ty:ignore[not-subscriptable] obj.data = None assert obj.data is None def test_optional_list_with_typing_optional(self) -> None: @py_class(_unique_key("OptListTyping")) class OptListTyping(Object): items: Optional[List[int]] obj = OptListTyping(items=[1, 2, 3]) assert len(obj.items) == 3 # ty:ignore[invalid-argument-type] obj.items = None assert obj.items is None def test_optional_dict_with_typing_optional(self) -> None: @py_class(_unique_key("OptDictTyping")) class OptDictTyping(Object): data: Optional[Dict[str, int]] obj = OptDictTyping(data={"a": 1}) assert obj.data["a"] == 1 # ty:ignore[not-subscriptable] obj.data = None assert obj.data is None # ########################################################################### # 24. Callable / Function fields # ########################################################################### class TestFunctionField: """Function/Callable field via @py_class decorator.""" def test_function_field(self) -> None: @py_class(_unique_key("FuncFld")) class FuncFld(Object): func: tvm_ffi.Function fn = tvm_ffi.convert(lambda x: x + 1) obj = FuncFld(func=fn) assert obj.func(1) == 2 def test_function_field_set(self) -> None: @py_class(_unique_key("FuncSet")) class FuncSet(Object): func: tvm_ffi.Function fn1 = tvm_ffi.convert(lambda x: x + 1) fn2 = tvm_ffi.convert(lambda x: x + 2) obj = FuncSet(func=fn1) assert obj.func(1) == 2 obj.func = fn2 assert obj.func(1) == 3 @requires_py310 def test_optional_function_field(self) -> None: @py_class(_unique_key("OptFunc")) class OptFunc(Object): func: Optional[tvm_ffi.Function] obj = OptFunc(func=None) assert obj.func is None obj.func = tvm_ffi.convert(lambda x: x * 2) assert obj.func(3) == 6 obj.func = None assert obj.func is None # ########################################################################### # 25. Any-typed fields (decorator level) # ########################################################################### class TestAnyFieldDecorator: """Any-typed field via @py_class decorator.""" def test_any_holds_int(self) -> None: @py_class(_unique_key("AnyI")) class AnyI(Object): val: Any assert AnyI(val=42).val == 42 def test_any_holds_str(self) -> None: @py_class(_unique_key("AnyS")) class AnyS(Object): val: Any assert AnyS(val="hello").val == "hello" def test_any_holds_none(self) -> None: @py_class(_unique_key("AnyN")) class AnyN(Object): val: Any = None assert AnyN().val is None def test_any_holds_list(self) -> None: @py_class(_unique_key("AnyL")) class AnyL(Object): val: Any assert len(AnyL(val=[1, 2, 3]).val) == 3 def test_any_type_change(self) -> None: @py_class(_unique_key("AnyChg")) class AnyChg(Object): val: Any = None obj = AnyChg() assert obj.val is None obj.val = 42 assert obj.val == 42 obj.val = "hello" assert obj.val == "hello" obj.val = tvm_ffi.Array([1, 2]) assert len(obj.val) == 2 obj.val = None assert obj.val is None # ########################################################################### # 26. Post-init field mutation # ########################################################################### class TestPostInitMutation: """__post_init__ that mutates field values.""" def test_post_init_mutates_str(self) -> None: @py_class(_unique_key("PostMut")) class PostMut(Object): a: int b: str def __post_init__(self) -> None: self.b = self.b.upper() obj = PostMut(a=1, b="hello") assert obj.a == 1 assert obj.b == "HELLO" def test_post_init_computes_derived(self) -> None: @py_class(_unique_key("PostDeriv")) class PostDeriv(Object): x: int doubled: int = 0 def __post_init__(self) -> None: self.doubled = self.x * 2 assert PostDeriv(x=5).doubled == 10 # ########################################################################### # 27. Custom __init__ with init=False # ########################################################################### class TestCustomInitFalse: """Custom __init__ with init=False and reordered parameters.""" def test_custom_init_reordered_params(self) -> None: @py_class(_unique_key("CustomOrd"), init=False) class CustomOrd(Object): a: int b: float c: str d: bool def __init__(self, b: float, c: str, a: int, d: bool) -> None: self.a = a self.b = b self.c = c self.d = d obj = CustomOrd(b=2.0, c="3", a=1, d=True) assert obj.a == 1 assert obj.b == 2.0 assert obj.c == "3" assert obj.d is True def test_custom_init_keyword_only(self) -> None: @py_class(_unique_key("CustomKW"), init=False) class CustomKW(Object): a: int b: str def __init__(self, *, b: str, a: int) -> None: self.a = a self.b = b obj = CustomKW(a=1, b="hello") assert obj.a == 1 assert obj.b == "hello" # ########################################################################### # 28. Inheritance with defaults and containers # ########################################################################### class TestInheritanceWithDefaults: """Inheritance with default values and container fields.""" def test_base_with_default_factory(self) -> None: @py_class(_unique_key("BaseDef")) class BaseDef(Object): a: int b: List[int] = field(default_factory=list) obj = BaseDef(a=42) assert obj.a == 42 assert len(obj.b) == 0 def test_derived_adds_optional_fields(self) -> None: @py_class(_unique_key("BaseD")) class BaseD(Object): a: int b: List[int] = field(default_factory=list) @py_class(_unique_key("DerivedD")) class DerivedD(BaseD): c: Optional[int] = None d: Optional[str] = "default" obj = DerivedD(a=12) assert obj.a == 12 assert len(obj.b) == 0 assert obj.c is None assert obj.d == "default" def test_derived_interleaved_required_optional(self) -> None: @py_class(_unique_key("BaseIL")) class BaseIL(Object): a: int b: List[int] = field(default_factory=list) @py_class(_unique_key("DerivedIL")) class DerivedIL(BaseIL): c: int d: Optional[str] = "default" obj = DerivedIL(a=1, c=2) assert obj.a == 1 assert len(obj.b) == 0 assert obj.c == 2 assert obj.d == "default" def test_three_level_with_defaults(self) -> None: @py_class(_unique_key("L1D")) class L1D(Object): a: int @py_class(_unique_key("L2D")) class L2D(L1D): b: Optional[int] = None c: Optional[str] = "hello" @py_class(_unique_key("L3D")) class L3D(L2D): d: str obj = L3D(a=1, d="world") assert obj.a == 1 assert obj.b is None assert obj.c == "hello" assert obj.d == "world" def test_derived_with_container_init(self) -> None: @py_class(_unique_key("BaseC")) class BaseC(Object): items: List[int] @py_class(_unique_key("DerivedC")) class DerivedC(BaseC): name: str obj = DerivedC(items=[1, 2, 3], name="test") assert len(obj.items) == 3 assert obj.name == "test" # ########################################################################### # 29. Decorator-level type validation # ########################################################################### class TestFieldTypeValidation: """Type validation on set for @py_class fields.""" def test_set_int_to_str_raises(self) -> None: @py_class(_unique_key("ValInt")) class ValInt(Object): x: int obj = ValInt(x=1) with pytest.raises((TypeError, RuntimeError)): obj.x = "not_an_int" # ty:ignore[invalid-assignment] def test_set_str_to_int_raises(self) -> None: @py_class(_unique_key("ValStr")) class ValStr(Object): x: str obj = ValStr(x="hello") with pytest.raises((TypeError, RuntimeError)): obj.x = 42 # ty:ignore[invalid-assignment] def test_set_bool_to_str_raises(self) -> None: @py_class(_unique_key("ValBool")) class ValBool(Object): x: bool obj = ValBool(x=True) with pytest.raises((TypeError, RuntimeError)): obj.x = "not_a_bool" # ty:ignore[invalid-assignment] def test_set_list_to_wrong_type_raises(self) -> None: @py_class(_unique_key("ValList")) class ValList(Object): items: List[int] obj = ValList(items=[1, 2]) with pytest.raises((TypeError, RuntimeError)): obj.items = "not_a_list" # ty:ignore[invalid-assignment] def test_set_dict_to_wrong_type_raises(self) -> None: @py_class(_unique_key("ValDict")) class ValDict(Object): data: Dict[str, int] obj = ValDict(data={"a": 1}) with pytest.raises((TypeError, RuntimeError)): obj.data = "not_a_dict" # ty:ignore[invalid-assignment] # ########################################################################### # 30. Three-level inheritance with containers # ########################################################################### class TestDerivedDerivedContainers: """Three-level inheritance with container fields and init reordering.""" def test_three_level_with_container_and_defaults(self) -> None: @py_class(_unique_key("DD_L1")) class L1(Object): a: int b: List[int] = field(default_factory=list) @py_class(_unique_key("DD_L2")) class L2(L1): c: Optional[int] = None d: Optional[str] = "hello" @py_class(_unique_key("DD_L3")) class L3(L2): e: str obj = L3(a=1, e="world", b=[1, 2]) assert obj.a == 1 assert tuple(obj.b) == (1, 2) assert obj.c is None assert obj.d == "hello" assert obj.e == "world" def test_three_level_positional_call(self) -> None: @py_class(_unique_key("DD2_L1")) class L1(Object): a: int @py_class(_unique_key("DD2_L2")) class L2(L1): b: List[int] = field(default_factory=list) @py_class(_unique_key("DD2_L3")) class L3(L2): c: str obj = L3(a=1, c="x") assert obj.a == 1 assert len(obj.b) == 0 assert obj.c == "x" # ########################################################################### # 31. Container field mutation and type rejection # ########################################################################### class TestContainerFieldMutation: """Container field set, mutation, and type rejection.""" def test_untyped_list_mutation(self) -> None: obj = _make_multi_type_obj() assert len(obj.list_any) == 3 assert obj.list_any[0] == 1 obj.list_any = [4, 3.0, "two"] assert len(obj.list_any) == 3 assert obj.list_any[0] == 4 assert obj.list_any[2] == "two" def test_untyped_dict_mutation(self) -> None: obj = _make_multi_type_obj() assert len(obj.dict_any) == 2 obj.dict_any = {"4": 4, "3": "two", "2": 3.0} assert len(obj.dict_any) == 3 assert obj.dict_any["4"] == 4 assert obj.dict_any["3"] == "two" def test_list_any_type_rejection(self) -> None: @py_class(_unique_key("LAReject")) class LAReject(Object): items: List[Any] obj = LAReject(items=[1, 2]) with pytest.raises((TypeError, RuntimeError)): obj.items = "wrong" # ty:ignore[invalid-assignment] def test_list_list_int_type_rejection(self) -> None: @py_class(_unique_key("LLReject")) class LLReject(Object): matrix: List[List[int]] obj = LLReject(matrix=[[1, 2]]) with pytest.raises((TypeError, RuntimeError)): obj.matrix = [4, 3, 2, 1] # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.matrix = None # ty:ignore[invalid-assignment] assert len(obj.matrix) == 1 def test_dict_any_any_type_rejection(self) -> None: @py_class(_unique_key("DAAReject")) class DAAReject(Object): data: Dict[Any, Any] obj = DAAReject(data={1: 2}) with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = None # ty:ignore[invalid-assignment] def test_dict_str_any_type_rejection(self) -> None: @py_class(_unique_key("DSAReject")) class DSAReject(Object): data: Dict[str, Any] obj = DSAReject(data={"a": 1}) with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = None # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = {4: 4, 3.0: 3} # ty:ignore[invalid-assignment] def test_dict_str_list_int_type_rejection(self) -> None: @py_class(_unique_key("DSLReject")) class DSLReject(Object): data: Dict[str, List[int]] obj = DSLReject(data={"a": [1, 2]}) with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = None # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = {"a": 1, "b": [2]} # ty:ignore[invalid-assignment] # ########################################################################### # 32. Optional field set/unset cycles with type rejection # ########################################################################### class TestOptionalFieldCycles: """Optional field set → None → set-back cycles with type rejection.""" def test_opt_func_type_rejection(self) -> None: @py_class(_unique_key("OptFuncR")) class OptFuncR(Object): func: Optional[tvm_ffi.Function] obj = OptFuncR(func=None) assert obj.func is None obj.func = tvm_ffi.convert(lambda x: x + 2) assert obj.func(1) == 3 obj.func = None assert obj.func is None with pytest.raises((TypeError, RuntimeError)): obj.func = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.func = 42 # ty:ignore[invalid-assignment] def test_opt_ulist_cycle(self) -> None: @py_class(_unique_key("OptUListC")) class OptUListC(Object): items: Optional[list] obj = OptUListC(items=None) assert obj.items is None obj.items = [4, 3.0, "two"] assert len(obj.items) == 3 assert obj.items[0] == 4 with pytest.raises((TypeError, RuntimeError)): obj.items = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.items = 42 # ty:ignore[invalid-assignment] obj.items = None assert obj.items is None def test_opt_udict_cycle(self) -> None: @py_class(_unique_key("OptUDictC")) class OptUDictC(Object): data: Optional[dict] obj = OptUDictC(data=None) assert obj.data is None obj.data = {"4": 4, "3": "two"} assert len(obj.data) == 2 with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] obj.data = None assert obj.data is None def test_opt_list_any_cycle(self) -> None: @py_class(_unique_key("OptLAC")) class OptLAC(Object): items: Optional[List[Any]] obj = OptLAC(items=[1, 2.0, "three"]) assert len(obj.items) == 3 # ty:ignore[invalid-argument-type] obj.items = [4, 3.0, "two"] assert obj.items[0] == 4 with pytest.raises((TypeError, RuntimeError)): obj.items = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.items = 42 # ty:ignore[invalid-assignment] obj.items = None assert obj.items is None def test_opt_list_list_int_cycle(self) -> None: @py_class(_unique_key("OptLLIC")) class OptLLIC(Object): matrix: Optional[List[List[int]]] obj = OptLLIC(matrix=[[1, 2, 3], [4, 5, 6]]) assert tuple(obj.matrix[0]) == (1, 2, 3) # ty:ignore[not-subscriptable] obj.matrix = [[4, 3, 2]] assert len(obj.matrix) == 1 with pytest.raises((TypeError, RuntimeError)): obj.matrix = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.matrix = [1, 2, 3] # ty:ignore[invalid-assignment] obj.matrix = None assert obj.matrix is None def test_opt_dict_any_any_cycle(self) -> None: @py_class(_unique_key("OptDAAC")) class OptDAAC(Object): data: Optional[Dict[Any, Any]] obj = OptDAAC(data=None) assert obj.data is None obj.data = {4: 4, "three": "two"} assert len(obj.data) == 2 with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] obj.data = None assert obj.data is None def test_opt_dict_str_any_cycle(self) -> None: @py_class(_unique_key("OptDSAC")) class OptDSAC(Object): data: Optional[Dict[str, Any]] obj = OptDSAC(data={"a": 1}) assert obj.data["a"] == 1 # ty:ignore[not-subscriptable] obj.data = {} assert len(obj.data) == 0 with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] obj.data = None assert obj.data is None def test_opt_dict_any_str_cycle(self) -> None: @py_class(_unique_key("OptDASC")) class OptDASC(Object): data: Optional[Dict[Any, str]] obj = OptDASC(data={1: "a", "two": "b"}) assert obj.data[1] == "a" # ty:ignore[not-subscriptable] obj.data = {4: "4", "three": "two"} assert len(obj.data) == 2 with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] obj.data = None assert obj.data is None def test_opt_dict_str_list_int_cycle(self) -> None: @py_class(_unique_key("OptDSLIC")) class OptDSLIC(Object): data: Optional[Dict[str, List[int]]] obj = OptDSLIC(data={"1": [1, 2, 3], "2": [4, 5, 6]}) assert tuple(obj.data["1"]) == (1, 2, 3) # ty:ignore[not-subscriptable] obj.data = {"a": [7, 8]} assert tuple(obj.data["a"]) == (7, 8) with pytest.raises((TypeError, RuntimeError)): obj.data = "wrong" # ty:ignore[invalid-assignment] with pytest.raises((TypeError, RuntimeError)): obj.data = 42 # ty:ignore[invalid-assignment] obj.data = None assert obj.data is None # ########################################################################### # 33. Multi-type field class # ########################################################################### @py_class(_unique_key("MultiType")) class _PyClassMultiType(Object): """@py_class with many field types for cross-cutting field tests.""" bool_: bool i64: int f64: float str_: str any_val: Any list_int: List[int] list_any: list dict_str_int: Dict[str, int] dict_any: dict list_list_int: List[List[int]] dict_str_list_int: Dict[str, List[int]] opt_bool: Optional[bool] opt_int: Optional[int] opt_float: Optional[float] opt_str: Optional[str] opt_list_int: Optional[List[int]] opt_dict_str_int: Optional[Dict[str, int]] def _make_multi_type_obj() -> _PyClassMultiType: return _PyClassMultiType( bool_=False, i64=64, f64=2.5, str_="world", any_val="hello", list_int=[1, 2, 3], list_any=[1, "two", 3.0], dict_str_int={"a": 1, "b": 2}, dict_any={"x": 1, "y": "two"}, list_list_int=[[1, 2, 3], [4, 5, 6]], dict_str_list_int={"p": [1, 2], "q": [3, 4]}, opt_bool=True, opt_int=-64, opt_float=None, opt_str=None, opt_list_int=[10, 20], opt_dict_str_int=None, ) class TestMultiTypeFieldOps: """Per-field get/set/validation on a many-field @py_class type.""" def test_bool_field(self) -> None: obj = _make_multi_type_obj() assert obj.bool_ is False obj.bool_ = True assert obj.bool_ is True with pytest.raises((TypeError, RuntimeError)): obj.bool_ = "not_a_bool" # ty:ignore[invalid-assignment] def test_int_field(self) -> None: obj = _make_multi_type_obj() assert obj.i64 == 64 obj.i64 = -128 assert obj.i64 == -128 with pytest.raises((TypeError, RuntimeError)): obj.i64 = "wrong" # ty:ignore[invalid-assignment] def test_float_field(self) -> None: obj = _make_multi_type_obj() assert abs(obj.f64 - 2.5) < 1e-10 obj.f64 = 5.0 assert abs(obj.f64 - 5.0) < 1e-10 with pytest.raises((TypeError, RuntimeError)): obj.f64 = "wrong" # ty:ignore[invalid-assignment] def test_str_field(self) -> None: obj = _make_multi_type_obj() assert obj.str_ == "world" obj.str_ = "hello" assert obj.str_ == "hello" def test_any_field(self) -> None: obj = _make_multi_type_obj() assert obj.any_val == "hello" obj.any_val = 42 assert obj.any_val == 42 obj.any_val = [1, 2] assert len(obj.any_val) == 2 def test_list_int_field(self) -> None: obj = _make_multi_type_obj() assert tuple(obj.list_int) == (1, 2, 3) obj.list_int = [4, 5] assert len(obj.list_int) == 2 with pytest.raises((TypeError, RuntimeError)): obj.list_int = "wrong" # ty:ignore[invalid-assignment] def test_untyped_list_field(self) -> None: obj = _make_multi_type_obj() assert len(obj.list_any) == 3 assert obj.list_any[0] == 1 assert obj.list_any[1] == "two" obj.list_any = [4, 3.0, "new"] assert len(obj.list_any) == 3 def test_dict_str_int_field(self) -> None: obj = _make_multi_type_obj() assert obj.dict_str_int["a"] == 1 obj.dict_str_int = {"c": 3} assert obj.dict_str_int["c"] == 3 with pytest.raises((TypeError, RuntimeError)): obj.dict_str_int = "wrong" # ty:ignore[invalid-assignment] def test_untyped_dict_field(self) -> None: obj = _make_multi_type_obj() assert len(obj.dict_any) == 2 obj.dict_any = {"new": 42} assert obj.dict_any["new"] == 42 def test_nested_list_field(self) -> None: obj = _make_multi_type_obj() assert tuple(obj.list_list_int[0]) == (1, 2, 3) assert tuple(obj.list_list_int[1]) == (4, 5, 6) def test_nested_dict_field(self) -> None: obj = _make_multi_type_obj() assert tuple(obj.dict_str_list_int["p"]) == (1, 2) assert tuple(obj.dict_str_list_int["q"]) == (3, 4) def test_optional_bool(self) -> None: obj = _make_multi_type_obj() assert obj.opt_bool is True obj.opt_bool = False assert obj.opt_bool is False obj.opt_bool = None assert obj.opt_bool is None def test_optional_int(self) -> None: obj = _make_multi_type_obj() assert obj.opt_int == -64 obj.opt_int = None assert obj.opt_int is None obj.opt_int = 128 assert obj.opt_int == 128 def test_optional_float(self) -> None: obj = _make_multi_type_obj() assert obj.opt_float is None obj.opt_float = 1.5 assert abs(obj.opt_float - 1.5) < 1e-10 obj.opt_float = None assert obj.opt_float is None def test_optional_str(self) -> None: obj = _make_multi_type_obj() assert obj.opt_str is None obj.opt_str = "hello" assert obj.opt_str == "hello" obj.opt_str = None assert obj.opt_str is None def test_optional_list_int(self) -> None: obj = _make_multi_type_obj() assert tuple(obj.opt_list_int) == (10, 20) # ty:ignore[invalid-argument-type] obj.opt_list_int = None assert obj.opt_list_int is None obj.opt_list_int = [30] assert len(obj.opt_list_int) == 1 def test_optional_dict_str_int(self) -> None: obj = _make_multi_type_obj() assert obj.opt_dict_str_int is None obj.opt_dict_str_int = {"z": 99} assert obj.opt_dict_str_int["z"] == 99 obj.opt_dict_str_int = None assert obj.opt_dict_str_int is None class TestMultiTypeCopy: """Copy with the comprehensive multi-type class.""" def test_shallow_copy_comprehensive(self) -> None: obj = _make_multi_type_obj() obj2 = copy.copy(obj) assert obj2.bool_ == obj.bool_ assert obj2.i64 == obj.i64 assert obj2.f64 == obj.f64 assert obj2.str_ == obj.str_ assert obj2.any_val == obj.any_val assert obj.list_int.same_as(obj2.list_int) # ty:ignore[unresolved-attribute] assert obj.dict_str_int.same_as(obj2.dict_str_int) # ty:ignore[unresolved-attribute] def test_deep_copy_comprehensive(self) -> None: obj = _make_multi_type_obj() obj2 = copy.deepcopy(obj) assert obj2.bool_ == obj.bool_ assert obj2.i64 == obj.i64 assert not obj.list_int.same_as(obj2.list_int) # ty:ignore[unresolved-attribute] assert not obj.dict_str_int.same_as(obj2.dict_str_int) # ty:ignore[unresolved-attribute] assert tuple(obj2.list_int) == (1, 2, 3) assert obj2.dict_str_int["a"] == 1 # --------------------------------------------------------------------------- # _collect_py_methods allowlist and method introspection # --------------------------------------------------------------------------- class TestPyMethodAllowlist: """Only names in ``_FFI_RECOGNIZED_METHODS`` are collected by ``_collect_py_methods``.""" def test_system_methods_not_in_allowlist(self) -> None: from tvm_ffi.dataclasses.py_class import _collect_py_methods # noqa: PLC0415 @py_class(_unique_key("Allow")) class Allow(core.Object): x: int def __ffi_init__(self, x: int) -> None: # ty: ignore[invalid-method-override] pass def __ffi_shallow_copy__(self) -> None: pass def __ffi_repr__(self, fn_repr: Any) -> str: return "repr" collected = _collect_py_methods(Allow) assert collected is not None names = {name for name, _, _ in collected} assert "__ffi_repr__" in names assert "__ffi_init__" not in names assert "__ffi_shallow_copy__" not in names def test_arbitrary_ffi_dunder_not_collected(self) -> None: from tvm_ffi.dataclasses.py_class import _collect_py_methods # noqa: PLC0415 @py_class(_unique_key("Arb")) class Arb(core.Object): x: int def __ffi_custom_op__(self, y: int) -> int: return self.x + y collected = _collect_py_methods(Arb) assert collected is None class TestPyMethodIntrospection: """Registered __ffi_* hooks appear in TypeAttrColumn (not TypeMethod).""" def test_ffi_repr_in_type_attr(self) -> None: @py_class(_unique_key("IntrRepr")) class IntrRepr(core.Object): x: int def __ffi_repr__(self, fn_repr: Any) -> str: return "repr" info = getattr(IntrRepr, "__tvm_ffi_type_info__") # User-defined hooks are registered as TypeAttrColumn, not TypeMethod repr_attr = core._lookup_type_attr(info.type_index, "__ffi_repr__") assert repr_attr is not None # System-generated hooks are also TypeAttrColumn only ffi_init = core._lookup_type_attr(info.type_index, "__ffi_init__") assert ffi_init is not None ffi_copy = core._lookup_type_attr(info.type_index, "__ffi_shallow_copy__") assert ffi_copy is not None # --------------------------------------------------------------------------- # super().__init__() support for @py_class(init=False) subclasses # --------------------------------------------------------------------------- class TestSuperInitPattern: """Regression tests for using ``super().__init__()`` + field assignment in ``@py_class(init=False)`` custom ``__init__`` methods. Previously, ``super().__init__()`` would dispatch to the parent's auto-generated ``__init__`` which called ``self.__ffi_init__()`` with the **child** type's C++ constructor (requiring field arguments that weren't provided), causing a crash. """ def test_basic_super_init_with_field_setters(self) -> None: """The original crash scenario: init=False with super().__init__() and field setters.""" @py_class(_unique_key("SIBase")) class SIBase(Object): pass @py_class(_unique_key("SIChild"), init=False) class SIChild(SIBase): x: int y: str def __init__(self, x: int, y: str) -> None: super().__init__() self.x = x self.y = y obj = SIChild(42, "hello") assert obj.x == 42 assert obj.y == "hello" def test_super_init_deep_hierarchy(self) -> None: """super().__init__() through multiple levels of py_class inheritance.""" @py_class(_unique_key("SIDH_Node")) class Node(Object): pass @py_class(_unique_key("SIDH_BaseType")) class BaseType(Node): pass @py_class(_unique_key("SIDH_PtrType"), init=False) class PtrType(BaseType): base_type: Any specifiers: list use_bracket: bool def __init__( self, base_type: Any, specifiers: Optional[list] = None, use_bracket: bool = False, ) -> None: super().__init__() self.base_type = base_type self.specifiers = list(specifiers) if specifiers else [] self.use_bracket = use_bracket void_p = PtrType("void") assert void_p.base_type == "void" assert list(void_p.specifiers) == [] assert void_p.use_bracket is False int_p = PtrType("int", ["const", "volatile"], True) assert int_p.base_type == "int" assert list(int_p.specifiers) == ["const", "volatile"] assert int_p.use_bracket is True def test_super_init_with_defaults_from_calloc(self) -> None: """Fields not set after super().__init__() should be zero-initialized.""" @py_class(_unique_key("SIDef")) class SIDefBase(Object): pass @py_class(_unique_key("SIDef_Child"), init=False) class SIDefChild(SIDefBase): a: int b: float c: bool d: Any def __init__(self) -> None: super().__init__() # Don't set any fields — they should be zero/None from calloc obj = SIDefChild() assert obj.a == 0 assert obj.b == 0.0 assert obj.c is False assert obj.d is None def test_super_init_intermediate_custom_init(self) -> None: """Chained super().__init__() through an intermediate init=False class.""" @py_class(_unique_key("SIChain_A")) class A(Object): pass @py_class(_unique_key("SIChain_B"), init=False) class B(A): x: int def __init__(self, x: int) -> None: super().__init__() self.x = x @py_class(_unique_key("SIChain_C"), init=False) class C(B): y: str def __init__(self, x: int, y: str) -> None: super().__init__(x) self.y = y obj = C(10, "world") assert obj.x == 10 assert obj.y == "world" def test_super_init_intermediate_auto_init(self) -> None: """py_class init=False: no need to call super().__init__(), just set fields directly.""" @py_class(_unique_key("SIAI_Mid")) class Mid(Object): a: int @py_class(_unique_key("SIAI_Leaf"), init=False) class Leaf(Mid): b: str def __init__(self, a: int, b: str) -> None: # For py_class, __new__ already allocated via __ffi_new__. # No need to call super().__init__() — just set fields directly. self.a = a self.b = b obj = Leaf(5, "hi") assert obj.a == 5 assert obj.b == "hi" def test_normal_init_unaffected(self) -> None: """Normal init=True construction must still work correctly.""" @py_class(_unique_key("SINorm")) class SINorm(Object): x: int y: str obj = SINorm(1, "a") assert obj.x == 1 assert obj.y == "a" def test_non_pyclass_subclass_no_args_errors(self) -> None: """A non-py_class subclass calling parent init with no args should still error for required fields (not silently create an empty object). """ @py_class(_unique_key("SINonPC")) class SINonPC(Object): x: int class Plain(SINonPC): pass with pytest.raises(TypeError): Plain() # type: ignore[missing-argument] def test_super_init_isinstance(self) -> None: """Objects created via super().__init__() pattern have correct isinstance.""" @py_class(_unique_key("SIInst_B")) class Base(Object): pass @py_class(_unique_key("SIInst_C"), init=False) class Child(Base): val: int def __init__(self, val: int) -> None: super().__init__() self.val = val obj = Child(99) assert isinstance(obj, Child) assert isinstance(obj, Base) assert isinstance(obj, Object) def test_super_init_field_overwrite(self) -> None: """Fields can be overwritten multiple times after super().__init__().""" @py_class(_unique_key("SIOverwrite_B")) class Base(Object): pass @py_class(_unique_key("SIOverwrite_C"), init=False) class Child(Base): x: int def __init__(self, x: int) -> None: super().__init__() self.x = 0 self.x = x obj = Child(42) assert obj.x == 42 def test_super_init_copy_deepcopy(self) -> None: """copy/deepcopy work on objects created via the super().__init__() pattern.""" @py_class(_unique_key("SICopyBase")) class SICopyBase(Object): pass @py_class(_unique_key("SICopy"), init=False) class SICopy(SICopyBase): x: int y: str def __init__(self, x: int, y: str) -> None: super().__init__() self.x = x self.y = y obj = SICopy(42, "hello") obj2 = copy.copy(obj) assert obj2.x == 42 assert obj2.y == "hello" assert not obj.same_as(obj2) obj3 = copy.deepcopy(obj) assert obj3.x == 42 assert obj3.y == "hello" assert not obj.same_as(obj3) def test_super_init_direct_from_object(self) -> None: """super().__init__() works when inheriting directly from Object (no intermediate).""" @py_class(_unique_key("SIDirect"), init=False) class SIDirect(Object): x: int y: str def __init__(self, x: int, y: str) -> None: super().__init__() self.x = x self.y = y obj = SIDirect(10, "direct") assert obj.x == 10 assert obj.y == "direct" assert isinstance(obj, Object) def test_super_init_direct_from_object_copy(self) -> None: """copy/deepcopy work for init=False classes inheriting directly from Object.""" @py_class(_unique_key("SIDirectCopy"), init=False) class SIDirectCopy(Object): x: int y: str def __init__(self, x: int, y: str) -> None: super().__init__() self.x = x self.y = y obj = SIDirectCopy(42, "hello") obj2 = copy.copy(obj) assert obj2.x == 42 assert obj2.y == "hello" assert not obj.same_as(obj2) obj3 = copy.deepcopy(obj) assert obj3.x == 42 assert obj3.y == "hello" assert not obj.same_as(obj3) class TestDtypeDeviceFields: """Regression: @py_class should accept tvm_ffi.dtype and tvm_ffi.Device as field types.""" def test_dtype_field(self) -> None: @py_class(_unique_key("DtypeField")) class DtypeHolder(Object): dt: tvm_ffi.dtype obj = DtypeHolder(dt=tvm_ffi.dtype("float32")) assert obj.dt == "float32" assert isinstance(obj.dt, tvm_ffi.dtype) def test_dtype_field_setter(self) -> None: @py_class(_unique_key("DtypeFieldSet")) class DtypeHolder2(Object): dt: tvm_ffi.dtype obj = DtypeHolder2(dt=tvm_ffi.dtype("float32")) obj.dt = tvm_ffi.dtype("int8") assert obj.dt == "int8" def test_device_field(self) -> None: @py_class(_unique_key("DeviceField")) class DeviceHolder(Object): dev: tvm_ffi.Device dev = tvm_ffi.device("cpu", 0) obj = DeviceHolder(dev=dev) assert obj.dev == dev def test_dtype_device_together(self) -> None: @py_class(_unique_key("DtypeDeviceTogether")) class DtypeDeviceHolder(Object): dt: tvm_ffi.dtype dev: tvm_ffi.Device name: str dev = tvm_ffi.device("cpu", 0) obj = DtypeDeviceHolder(dt=tvm_ffi.dtype("float16"), dev=dev, name="test") assert obj.dt == "float16" assert obj.dev == dev assert obj.name == "test" def test_optional_dtype_field(self) -> None: @py_class(_unique_key("OptDtype")) class OptDtype(Object): dt: Optional[tvm_ffi.dtype] = None obj_none = OptDtype() assert obj_none.dt is None obj_val = OptDtype(dt=tvm_ffi.dtype("bfloat16")) assert obj_val.dt == "bfloat16" def test_optional_device_field(self) -> None: @py_class(_unique_key("OptDevice")) class OptDevice(Object): dev: Optional[tvm_ffi.Device] = None obj_none = OptDevice() assert obj_none.dev is None obj_val = OptDevice(dev=tvm_ffi.device("cpu", 0)) assert obj_val.dev == tvm_ffi.device("cpu", 0) # ########################################################################### # kw_only regression tests (py_class via __ffi_init__) # ########################################################################### class TestPyClassKwOnlyRegression: """Regression tests ensuring kw_only enforcement via C++ __ffi_init__.""" def test_missing_kw_only_error_says_keyword_only(self) -> None: """Missing required kw_only field produces 'keyword-only' in the error.""" @py_class(_unique_key("KwOnlyErr")) class KwOnlyErr(Object): x: int _: KW_ONLY y: int with pytest.raises(TypeError, match="keyword-only"): KwOnlyErr(1) # ty: ignore[missing-argument] def test_missing_positional_error_not_keyword_only(self) -> None: """Missing required positional field does NOT say 'keyword-only'.""" @py_class(_unique_key("PosErr")) class PosErr(Object): x: int _: KW_ONLY y: int with pytest.raises(TypeError, match="missing required argument") as exc_info: PosErr(y=1) # ty: ignore[missing-argument] assert "keyword-only" not in str(exc_info.value) def test_kw_only_with_default(self) -> None: """kw_only field with default can be omitted.""" @py_class(_unique_key("KwOnlyDef")) class KwOnlyDef(Object): x: int _: KW_ONLY y: int = 42 obj = KwOnlyDef(1) # ty: ignore[missing-argument] assert obj.x == 1 assert obj.y == 42 def test_kw_only_rejects_positional(self) -> None: """kw_only fields cannot be passed positionally.""" @py_class(_unique_key("KwOnlyReject")) class KwOnlyReject(Object): x: int _: KW_ONLY y: int with pytest.raises(TypeError): KwOnlyReject(1, 2) # ty: ignore[missing-argument, invalid-argument-type] def test_kw_only_all_fields(self) -> None: """All fields kw_only via decorator kwarg.""" @py_class(_unique_key("AllKw"), kw_only=True) class AllKw(Object): a: int b: int = 10 obj = AllKw(a=1) assert obj.a == 1 assert obj.b == 10 with pytest.raises(TypeError): AllKw(1) # ty: ignore[missing-argument, too-many-positional-arguments] def test_kw_only_field_override_false(self) -> None: """kw_only=False on a field after KW_ONLY sentinel makes it positional.""" @py_class(_unique_key("KwOverride2")) class KwOverride2(Object): _: KW_ONLY a: int b: int = field(kw_only=False) obj = KwOverride2(42, a=1) # ty: ignore[missing-argument, invalid-argument-type] assert obj.b == 42 assert obj.a == 1 def test_decorator_kw_only_with_override(self) -> None: """Decorator-level kw_only with field-level override.""" @py_class(_unique_key("DecKwOverride"), kw_only=True) class DecKwOverride(Object): a: int = field(kw_only=False) b: int obj = DecKwOverride(1, b=2) # ty: ignore[missing-argument, too-many-positional-arguments] assert obj.a == 1 assert obj.b == 2 def test_kw_only_inheritance(self) -> None: """kw_only enforcement works through inheritance.""" @py_class(_unique_key("KwParent")) class KwParent(Object): x: int @py_class(_unique_key("KwChild")) class KwChild(KwParent): _: KW_ONLY y: int obj = KwChild(1, y=2) # ty: ignore[missing-argument] assert obj.x == 1 assert obj.y == 2 with pytest.raises(TypeError): KwChild(1, 2) # ty: ignore[missing-argument, invalid-argument-type] with pytest.raises(TypeError, match="keyword-only"): KwChild(1) # ty: ignore[missing-argument] # ########################################################################### # No-leak / no-spurious-allocation tests for py_class # ########################################################################### class TestPyClassNoLeak: """Verify that the removal of custom __new__ eliminates the memory leak. The original bug: ``@py_class`` installed a custom ``__new__`` that called ``__ffi_new__`` to pre-allocate a C++ object. When a py_class object was returned from C++ through ``make_ret_object``, ``cls.__new__(cls)`` triggered that custom ``__new__``, allocating a spurious C++ object whose refcount was never decremented. """ def test_no_custom_new_on_py_class(self) -> None: """py_class must NOT install a custom __new__.""" @py_class(_unique_key("NoNew")) class NoNew(Object): x: int assert "__new__" not in NoNew.__dict__ def test_no_custom_new_with_user_init(self) -> None: """py_class with user-defined __init__ must NOT install a custom __new__.""" @py_class(_unique_key("NoNewUI"), init=False) class NoNewUI(Object): x: int def __init__(self, x: int) -> None: self.x = x assert "__new__" not in NoNewUI.__dict__ def test_make_ret_no_spurious_alloc(self) -> None: """Objects returned from C++ (via DeepCopy) must not trigger spurious allocation.""" @py_class(_unique_key("RetTest")) class RetTest(Object): x: int obj = RetTest(42) # DeepCopy returns through make_ret_object obj2 = DeepCopy(obj) assert obj2.x == 42 assert not obj.same_as(obj2) def test_repeated_roundtrip_no_leak(self) -> None: """Repeated construct + DeepCopy cycles must not leak.""" @py_class(_unique_key("RoundTrip")) class RoundTrip(Object): x: int y: str = "default" gc.collect() for i in range(1000): o = RoundTrip(i, str(i)) o2 = DeepCopy(o) del o, o2 gc.collect() def test_user_init_wraps_metadata(self) -> None: """Wrapped user __init__ preserves docstring and __wrapped__.""" @py_class(_unique_key("WrapMeta"), init=False) class WrapMeta(Object): x: int def __init__(self, x: int) -> None: """Initialize WrapMeta.""" self.x = x assert WrapMeta.__init__.__doc__ == "Initialize WrapMeta." assert hasattr(WrapMeta.__init__, "__wrapped__") def test_chandle_guard_skips_on_make_ret(self) -> None: """Auto-generated __init__ with chandle guard: make_ret never calls __init__.""" @py_class(_unique_key("ChandleGuard")) class ChandleGuard(Object): x: int obj = ChandleGuard(7) obj2 = DeepCopy(obj) # If __init__ were called by make_ret, it would fail (no args) # or overwrite fields. Verify the value survived intact. assert obj2.x == 7 def test_super_init_noop_after_ffi_init(self) -> None: """super().__init__() is a no-op when chandle is already set.""" @py_class(_unique_key("SuperBase")) class SuperBase(Object): pass call_log: list[str] = [] @py_class(_unique_key("SuperChild"), init=False) class SuperChild(SuperBase): x: int def __init__(self, x: int) -> None: call_log.append("before_super") super().__init__() call_log.append("after_super") self.x = x obj = SuperChild(42) assert obj.x == 42 assert call_log == ["before_super", "after_super"] def test_user_init_mismatched_signature(self) -> None: """User __init__ whose args don't match field layout still works.""" @py_class(_unique_key("Mismatch"), init=False) class Mismatch(Object): value: int ref: Optional[Object] def __init__(self, val: int) -> None: self.value = val self.ref = None obj = Mismatch(99) assert obj.value == 99 assert obj.ref is None tvm-ffi-0.1.12/tests/python/test_dataclass_repr.py000066400000000000000000000602621521067262500222420ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for __ffi_repr__ / ffi.ReprPrint.""" from __future__ import annotations import ast import re import numpy as np import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi._ffi_api import ReprPrint # Regex building blocks A = r"0x[0-9a-f]+" # hex address def _check(result: str, pattern: str) -> None: """Assert result matches pattern with re.fullmatch, with a clear error message.""" assert re.fullmatch(pattern, result), ( f"fullmatch failed:\n result: {result!r}\n pattern: {pattern!r}" ) def test_repr_primitives() -> None: """Test repr of primitive types.""" assert ReprPrint(42) == "42" assert ReprPrint(0) == "0" assert ReprPrint(-1) == "-1" assert ReprPrint(True) == "True" assert ReprPrint(False) == "False" assert ReprPrint(None) == "None" def test_repr_float() -> None: """Test repr of floating point.""" assert ReprPrint(3.14) == "3.14" assert ReprPrint(0.0) == "0" assert ReprPrint(1e10) == "1e+10" def test_repr_string() -> None: """Test repr of FFI String (both SmallStr and StringObj).""" # SmallStr (<=7 bytes) assert ReprPrint("hello") == '"hello"' # StringObj (>7 bytes) assert ReprPrint("hello world") == '"hello world"' # ---------- Array (tuple format) ---------- def test_repr_array() -> None: """Test repr of FFI Array uses tuple format.""" assert ReprPrint(tvm_ffi.Array([1, 2, 3])) == "(1, 2, 3)" def test_repr_array_single() -> None: """Test repr of single-element Array has trailing comma.""" assert ReprPrint(tvm_ffi.Array([42])) == "(42,)" def test_repr_array_empty() -> None: """Test repr of empty Array.""" assert ReprPrint(tvm_ffi.Array([])) == "()" def test_repr_array_nested_strings() -> None: """Test repr of Array containing strings.""" assert ReprPrint(tvm_ffi.Array(["a", "b"])) == '("a", "b")' def test_repr_array_python_repr() -> None: """Test that Array.__repr__ uses the centralized repr.""" assert repr(tvm_ffi.Array([1, 2])) == "(1, 2)" # ---------- List ---------- def test_repr_list() -> None: """Test repr of FFI List uses list format.""" assert ReprPrint(tvm_ffi.List([10, 20])) == "[10, 20]" def test_repr_list_single() -> None: """Test repr of single-element List (no trailing comma).""" assert ReprPrint(tvm_ffi.List([99])) == "[99]" def test_repr_list_empty() -> None: """Test repr of empty List.""" assert ReprPrint(tvm_ffi.List([])) == "[]" def test_repr_list_nested_strings() -> None: """Test repr of List containing strings.""" assert ReprPrint(tvm_ffi.List(["x", "y"])) == '["x", "y"]' # ---------- Map ---------- def test_repr_map() -> None: """Test repr of FFI Map.""" assert ReprPrint(tvm_ffi.Map({"key": "value"})) == '{"key": "value"}' def test_repr_map_empty() -> None: """Test repr of empty Map.""" assert ReprPrint(tvm_ffi.Map({})) == "{}" # ---------- Dict ---------- def test_repr_dict() -> None: """Test repr of FFI Dict.""" assert ReprPrint(tvm_ffi.Dict({"key": "value"})) == '{"key": "value"}' def test_repr_dict_empty() -> None: """Test repr of empty Dict.""" assert ReprPrint(tvm_ffi.Dict({})) == "{}" def test_repr_dict_int_keys() -> None: """Test repr of Dict with integer keys.""" d = tvm_ffi.Dict({1: 2, 3: 4}) result = ReprPrint(d) # Dict iteration order is hash-dependent; match either ordering. _check(result, r"(?:\{1: 2, 3: 4\}|\{3: 4, 1: 2\})") def test_repr_dict_with_array_values() -> None: """Test repr of Dict with Array values.""" assert ReprPrint(tvm_ffi.Dict({1: tvm_ffi.Array([10, 20])})) == "{1: (10, 20)}" def test_repr_dict_with_object_values() -> None: """Test repr of Dict with object values.""" pair = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) d = tvm_ffi.Dict({"obj": pair}) assert ReprPrint(d) == '{"obj": testing.TestIntPair(a=1, b=2)}' # ---------- Tensor ---------- def test_repr_tensor() -> None: """Test repr of Tensor shows dtype, shape, device (no address by default).""" x = tvm_ffi.from_dlpack(np.zeros((3, 4), dtype="float32")) assert ReprPrint(x) == "float32[3, 4]@cpu:0" def test_repr_tensor_int8() -> None: """Test repr of Tensor with int8 dtype.""" x = tvm_ffi.from_dlpack(np.zeros((2,), dtype="int8")) assert ReprPrint(x) == "int8[2]@cpu:0" # ---------- Shape ---------- def test_repr_shape() -> None: """Test repr of Shape.""" assert ReprPrint(tvm_ffi.Shape((5, 6))) == "Shape(5, 6)" # ---------- User-defined objects ---------- def test_repr_user_object_all_fields() -> None: """Test repr of user-defined object with all fields shown (no address by default).""" obj = tvm_ffi.testing.create_object("testing.TestIntPair", a=10, b=20) assert ReprPrint(obj) == "testing.TestIntPair(a=10, b=20)" def test_repr_user_object_repr_off() -> None: """Test repr of object with Repr(false) fields excluded.""" # Positional order: required first (v_i64, v_i32, v_f64), then optional (v_f32) obj = tvm_ffi.testing._TestCxxClassDerived(1, 2, 3.5, 4.5) assert ReprPrint(obj) == "testing.TestCxxClassDerived(v_f64=3.5, v_f32=4.5)" def test_repr_python_repr() -> None: """Test that Python __repr__ delegates to ReprPrint.""" obj = tvm_ffi.testing.create_object("testing.TestIntPair", a=5, b=6) assert repr(obj) == "testing.TestIntPair(a=5, b=6)" # ---------- DAG / shared references (full form on every occurrence) ---------- def test_repr_duplicate_reference() -> None: """Test that duplicate object references use full form on every occurrence.""" inner = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) arr = tvm_ffi.Array([inner, inner]) result = ReprPrint(arr) assert result == "(testing.TestIntPair(a=1, b=2), testing.TestIntPair(a=1, b=2))" def test_repr_shared_in_map_values() -> None: """Test that the same Array shared in Map values uses full form on both.""" shared = tvm_ffi.Array([1, 2]) m = tvm_ffi.Map({"a": shared, "b": shared}) result = ReprPrint(m) # Map iteration order is hash-dependent; match either ordering. pat_ab = r'\{"a": \(1, 2\), "b": \(1, 2\)\}' pat_ba = r'\{"b": \(1, 2\), "a": \(1, 2\)\}' _check(result, rf"(?:{pat_ab}|{pat_ba})") def test_repr_shared_across_nesting_levels() -> None: """Test shared object across different nesting levels uses full form everywhere.""" leaf = tvm_ffi.testing.create_object("testing.TestIntPair", a=7, b=8) arr = tvm_ffi.Array([leaf, tvm_ffi.Array([leaf])]) result = ReprPrint(arr) assert result == "(testing.TestIntPair(a=7, b=8), (testing.TestIntPair(a=7, b=8),))" def test_repr_triple_shared_reference() -> None: """Test object appearing three times -- full form every time.""" inner = tvm_ffi.testing.create_object("testing.TestIntPair", a=0, b=0) arr = tvm_ffi.Array([inner, inner, inner]) result = ReprPrint(arr) assert result == ( "(testing.TestIntPair(a=0, b=0), " "testing.TestIntPair(a=0, b=0), " "testing.TestIntPair(a=0, b=0))" ) # ---------- Nested containers ---------- def test_repr_array_of_arrays() -> None: """Test repr of Array containing Arrays.""" inner = tvm_ffi.Array([1, 2]) outer = tvm_ffi.Array([inner, tvm_ffi.Array([3])]) assert ReprPrint(outer) == "((1, 2), (3,))" def test_repr_map_of_containers() -> None: """Test repr of Map containing Array values.""" m = tvm_ffi.Map({"a": tvm_ffi.Array([1, 2])}) assert ReprPrint(m) == '{"a": (1, 2)}' def test_repr_list_of_lists() -> None: """Test repr of List containing Lists.""" inner = tvm_ffi.List([1, 2]) outer = tvm_ffi.List([inner, tvm_ffi.List([3])]) assert ReprPrint(outer) == "[[1, 2], [3]]" # ---------- Nested dataclasses ---------- def test_repr_nested_dataclass() -> None: """Test repr of object with object-typed fields.""" inner = tvm_ffi.testing.create_object("testing.TestIntPair", a=10, b=20) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.5, v_str="hi", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([inner]), ) assert ReprPrint(obj) == ( 'testing.TestObjectDerived(v_i64=1, v_f64=2.5, v_str="hi", ' "v_map={}, " "v_array=(testing.TestIntPair(a=10, b=20),))" ) def test_repr_object_with_none_field() -> None: """Test repr of object where container fields are empty.""" obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) assert ( ReprPrint(obj) == 'testing.TestObjectDerived(v_i64=0, v_f64=0, v_str="", v_map={}, v_array=())' ) # ---------- Deep nesting ---------- def test_repr_deeply_nested_arrays() -> None: """Test repr of deeply nested Arrays (4 levels).""" a = tvm_ffi.Array([1]) for _ in range(3): a = tvm_ffi.Array([a]) assert ReprPrint(a) == "((((1,),),),)" def test_repr_deeply_nested_lists() -> None: """Test repr of deeply nested Lists (4 levels).""" lst = tvm_ffi.List([1]) for _ in range(3): lst = tvm_ffi.List([lst]) assert ReprPrint(lst) == "[[[[1]]]]" def test_repr_mixed_container_nesting() -> None: """Test repr of mixed Array/List/Map nesting.""" inner_list = tvm_ffi.List([1, 2]) inner_arr = tvm_ffi.Array([inner_list]) m = tvm_ffi.Map({"nested": inner_arr}) assert ReprPrint(m) == '{"nested": ([1, 2],)}' # ---------- Shared reference patterns ---------- def test_repr_dataclass_shared_subobject() -> None: """Test repr of two dataclasses sharing the same sub-object (full form in both).""" shared = tvm_ffi.testing.create_object("testing.TestIntPair", a=5, b=5) obj1 = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([shared]), ) obj2 = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=2, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([shared]), ) arr = tvm_ffi.Array([obj1, obj2]) result = ReprPrint(arr) assert result == ( "(" 'testing.TestObjectDerived(v_i64=1, v_f64=0, v_str="", v_map={}, ' "v_array=(testing.TestIntPair(a=5, b=5),)), " 'testing.TestObjectDerived(v_i64=2, v_f64=0, v_str="", v_map={}, ' "v_array=(testing.TestIntPair(a=5, b=5),))" ")" ) # ---------- Container with dataclass nesting ---------- def test_repr_array_of_dataclasses() -> None: """Test repr of Array of user-defined objects.""" objs = [tvm_ffi.testing.create_object("testing.TestIntPair", a=i, b=i * 10) for i in range(3)] arr = tvm_ffi.Array(objs) assert ReprPrint(arr) == ( "(testing.TestIntPair(a=0, b=0), " "testing.TestIntPair(a=1, b=10), " "testing.TestIntPair(a=2, b=20))" ) def test_repr_map_with_object_values() -> None: """Test repr of Map with object values.""" pair = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) m = tvm_ffi.Map({"obj": pair}) assert ReprPrint(m) == '{"obj": testing.TestIntPair(a=1, b=2)}' # ---------- Repr(false) inheritance ---------- def test_repr_derived_derived_shows_all_own_fields() -> None: """TestCxxClassDerivedDerived should show v_f64, v_f32, v_str, v_bool (not v_i64, v_i32).""" # Positional order: required (v_i64, v_i32, v_f64, v_bool), then optional (v_f32, v_str) obj = tvm_ffi.testing._TestCxxClassDerivedDerived(1, 2, 3.0, True, 4.0, "test") assert ( ReprPrint(obj) == 'testing.TestCxxClassDerivedDerived(v_f64=3, v_f32=4, v_str="test", v_bool=True)' ) # ---------- Edge cases: special values ---------- def test_repr_large_integer() -> None: """Test repr of large integers.""" assert ReprPrint(2**62) == str(2**62) assert ReprPrint(-(2**62)) == str(-(2**62)) def test_repr_negative_float() -> None: """Test repr of negative floats.""" assert ReprPrint(-1.5) == "-1.5" def test_repr_empty_string() -> None: """Test repr of empty string (SmallStr).""" assert ReprPrint("") == '""' def test_repr_string_with_spaces() -> None: """Test repr of string with spaces.""" assert ReprPrint("a b c") == '"a b c"' def test_repr_array_of_none() -> None: """Test repr of Array containing None values.""" assert ReprPrint(tvm_ffi.Array([None, None])) == "(None, None)" def test_repr_array_of_booleans() -> None: """Test repr of Array containing boolean values.""" assert ReprPrint(tvm_ffi.Array([True, False])) == "(True, False)" def test_repr_array_of_mixed_types() -> None: """Test repr of Array containing mixed primitive types.""" assert ReprPrint(tvm_ffi.Array([1, "hello", True, None])) == '(1, "hello", True, None)' def test_repr_map_int_keys() -> None: """Test repr of Map with integer keys.""" m = tvm_ffi.Map({1: 2, 3: 4}) result = ReprPrint(m) # Map iteration order is hash-dependent; match either ordering. _check(result, r"(?:\{1: 2, 3: 4\}|\{3: 4, 1: 2\})") def test_repr_map_with_array_values() -> None: """Test repr of Map with Array values.""" assert ReprPrint(tvm_ffi.Map({1: tvm_ffi.Array([10, 20])})) == "{1: (10, 20)}" # ---------- Nested dataclass edge cases ---------- def test_repr_dataclass_with_array_field() -> None: """Test repr of dataclass whose field is an Array of objects.""" pair1 = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) pair2 = tvm_ffi.testing.create_object("testing.TestIntPair", a=3, b=4) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=0, v_f64=0.0, v_str="test", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([pair1, pair2]), ) assert ReprPrint(obj) == ( 'testing.TestObjectDerived(v_i64=0, v_f64=0, v_str="test", ' "v_map={}, " "v_array=(testing.TestIntPair(a=1, b=2), testing.TestIntPair(a=3, b=4)))" ) def test_repr_dataclass_with_map_field() -> None: """Test repr of dataclass whose field is a Map.""" obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=42, v_f64=1.0, v_str="s", v_map=tvm_ffi.Map({"x": 10}), v_array=tvm_ffi.Array([]), ) assert ReprPrint(obj) == ( 'testing.TestObjectDerived(v_i64=42, v_f64=1, v_str="s", v_map={"x": 10}, v_array=())' ) # ---------- Cycle detection ---------- def test_repr_self_reference_cycle() -> None: """Test that self-referencing cycles show '...' marker.""" obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=2.0, v_str="hi", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) obj.v_array = tvm_ffi.Array([obj]) # type: ignore[unresolved-attribute] result = ReprPrint(obj) assert result == ( 'testing.TestObjectDerived(v_i64=1, v_f64=2, v_str="hi", v_map={}, v_array=(...,))' ) def test_repr_mutual_reference_cycle() -> None: """Test that mutual reference cycles show '...' marker.""" v_map = tvm_ffi.Map({}) obj_a = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=0.0, v_str="a", v_map=v_map, v_array=tvm_ffi.Array([]), ) obj_b = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=2, v_f64=0.0, v_str="b", v_map=v_map, v_array=tvm_ffi.Array([obj_a]), ) obj_a.v_array = tvm_ffi.Array([obj_b]) # type: ignore[unresolved-attribute] result = ReprPrint(obj_a) assert result == ( 'testing.TestObjectDerived(v_i64=1, v_f64=0, v_str="a", v_map={}, ' "v_array=(testing.TestObjectDerived(v_i64=2, v_f64=0, " 'v_str="b", v_map={}, v_array=(...,)),))' ) # ---------- TVM_FFI_REPR_WITH_ADDR ---------- def test_repr_with_addr_user_object(monkeypatch: pytest.MonkeyPatch) -> None: """Test that user objects show address when TVM_FFI_REPR_WITH_ADDR is set.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") obj = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) _check(ReprPrint(obj), rf"testing\.TestIntPair@{A}\(a=1, b=2\)") def test_repr_with_addr_array(monkeypatch: pytest.MonkeyPatch) -> None: """Test that Array shows address suffix when TVM_FFI_REPR_WITH_ADDR is set.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") arr = tvm_ffi.Array([1, 2, 3]) _check(ReprPrint(arr), rf"\(1, 2, 3\)@{A}") def test_repr_with_addr_list(monkeypatch: pytest.MonkeyPatch) -> None: """Test that List shows address suffix when TVM_FFI_REPR_WITH_ADDR is set.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") lst = tvm_ffi.List([10, 20]) _check(ReprPrint(lst), rf"\[10, 20\]@{A}") def test_repr_with_addr_dict(monkeypatch: pytest.MonkeyPatch) -> None: """Test that Dict shows address suffix when TVM_FFI_REPR_WITH_ADDR is set.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") d = tvm_ffi.Dict({"a": 1}) _check(ReprPrint(d), rf'\{{"a": 1\}}@{A}') def test_repr_with_addr_dag(monkeypatch: pytest.MonkeyPatch) -> None: """Test DAG with addresses: both occurrences show full form with same address.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") inner = tvm_ffi.testing.create_object("testing.TestIntPair", a=1, b=2) arr = tvm_ffi.Array([inner, inner]) result = ReprPrint(arr) _check( result, rf"\(testing\.TestIntPair@(?P{A})\(a=1, b=2\), " rf"testing\.TestIntPair@(?P=a)\(a=1, b=2\)\)@{A}", ) def test_repr_with_addr_cycle(monkeypatch: pytest.MonkeyPatch) -> None: """Test cycle with addresses: '...@ADDR' points back to the cyclic object.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=0.0, v_str="", v_map=tvm_ffi.Map({}), v_array=tvm_ffi.Array([]), ) obj.v_array = tvm_ffi.Array([obj]) # type: ignore[unresolved-attribute] result = ReprPrint(obj) _check( result, rf"testing\.TestObjectDerived@(?P{A})\(" rf'v_i64=1, v_f64=0, v_str="", v_map=\{{\}}@{A}, ' rf"v_array=\(\.\.\.@(?P=obj),\)@{A}" rf"\)", ) def test_repr_with_addr_tensor(monkeypatch: pytest.MonkeyPatch) -> None: """Test that Tensor shows address suffix when TVM_FFI_REPR_WITH_ADDR is set.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") x = tvm_ffi.from_dlpack(np.zeros((3, 4), dtype="float32")) _check(ReprPrint(x), rf"float32\[3, 4\]@cpu:0@{A}") def test_repr_with_addr_no_fields(monkeypatch: pytest.MonkeyPatch) -> None: """Test that object with no visible fields shows TypeKey@ADDR with env var.""" monkeypatch.setenv("TVM_FFI_REPR_WITH_ADDR", "1") # TestCxxClassBase has v_i64 and v_i32, both with Repr(false) obj = tvm_ffi.testing._TestCxxClassBase(v_i64=1, v_i32=2) _check(ReprPrint(obj), rf"testing\.TestCxxClassBase@{A}") # ---------- Additional corner cases (fail-first) ---------- @pytest.mark.parametrize( "value", [ 'a"b', "a\\b", "\\", '"', "a\nb", "a\rb", "\x1b", "a\x00b", ], ) def test_repr_string_literal_roundtrip_special_chars(value: str) -> None: """String repr should be parseable and round-trip through ast.literal_eval.""" assert ast.literal_eval(ReprPrint(value)) == value @pytest.mark.parametrize( ("value", "expected"), [ (tvm_ffi.Array(['a"b']), ('a"b',)), (tvm_ffi.List(["a\\b"]), ["a\\b"]), (tvm_ffi.Array(["\\"]), ("\\",)), (tvm_ffi.Dict({"k": "a\nb"}), {"k": "a\nb"}), (tvm_ffi.Map({"k": "a\x00b"}), {"k": "a\x00b"}), ], ) def test_repr_container_literal_roundtrip_special_strings(value: object, expected: object) -> None: """Container repr with string payloads should remain parseable literals.""" assert ast.literal_eval(ReprPrint(value)) == expected def test_repr_device_trn_name() -> None: """DLDeviceType.kDLTrn should print as trn:, not unknown:.""" assert ReprPrint(tvm_ffi.Device("trn", 0)) == "trn:0" def test_repr_unregistered_object_no_duplicate_field_names() -> None: """Inherited fields should not appear twice in generic repr.""" obj = tvm_ffi.testing.make_unregistered_object() result = ReprPrint(obj) assert result.count("v1=") == 1 # --------------------------------------------------------------------------- # # @py_class repr # --------------------------------------------------------------------------- # import itertools as _itertools_repr from typing import Optional as _Optional_repr from tvm_ffi.core import Object as _Object_repr from tvm_ffi.dataclasses import py_class as _py_class_repr _counter_repr = _itertools_repr.count() def _unique_key_repr(base: str) -> str: return f"testing.repr_pc.{base}_{next(_counter_repr)}" def test_repr_py_class_base() -> None: """Repr of a simple @py_class contains field names and values.""" @_py_class_repr(_unique_key_repr("ReprBase")) class ReprBase(_Object_repr): a: int b: str r = repr(ReprBase(a=1, b="hello")) assert "a=1" in r or "a: 1" in r assert "hello" in r def test_repr_py_class_derived() -> None: """Repr of a derived @py_class shows all fields including parent.""" @_py_class_repr(_unique_key_repr("ReprP")) class ReprP(_Object_repr): base_a: int base_b: str @_py_class_repr(_unique_key_repr("ReprD")) class ReprD(ReprP): derived_a: float derived_b: _Optional_repr[str] # noqa: UP045 r = repr(ReprD(base_a=1, base_b="b", derived_a=2.0, derived_b="c")) assert "1" in r assert "2" in r def test_repr_py_class_in_array() -> None: """@py_class objects inside Array have proper repr.""" @_py_class_repr(_unique_key_repr("ReprInArr")) class ReprInArr(_Object_repr): x: int r = repr(tvm_ffi.Array([ReprInArr(x=1), ReprInArr(x=2)])) assert "1" in r assert "2" in r # --------------------------------------------------------------------------- # Custom __ffi_repr__ hook via @py_class # --------------------------------------------------------------------------- from typing import Any as _Any_repr from typing import Callable as _Callable_repr def test_py_class_custom_ffi_repr() -> None: """ReprPrint dispatches the user-defined __ffi_repr__ hook.""" @_py_class_repr(_unique_key_repr("CRepr")) class CRepr(_Object_repr): value: int def __ffi_repr__(self, fn_repr: _Callable_repr[..., _Any_repr]) -> str: return f"" assert ReprPrint(CRepr(42)) == "" assert ReprPrint(CRepr(999)) == "" def test_py_class_ffi_repr_with_fields_and_copy() -> None: """Fields work normally and copy preserves __ffi_repr__ behaviour.""" import copy as _copy_repr # noqa: PLC0415 @_py_class_repr(_unique_key_repr("FnR")) class FnR(_Object_repr): a: int b: str def __ffi_repr__(self, fn_repr: _Callable_repr[..., _Any_repr]) -> str: return f"FnR({self.a}, {self.b!r})" obj = FnR(10, "hi") assert obj.a == 10 assert obj.b == "hi" assert ReprPrint(obj) == "FnR(10, 'hi')" obj2 = _copy_repr.copy(obj) assert ReprPrint(obj2) == "FnR(10, 'hi')" if __name__ == "__main__": pytest.main([__file__, "-v"]) tvm-ffi-0.1.12/tests/python/test_device.py000066400000000000000000000103511521067262500205040ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import ctypes import pickle import numpy import pytest import tvm_ffi from tvm_ffi import DLDeviceType def test_device() -> None: device = tvm_ffi.Device("cuda", 0) assert device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCUDA assert device.index == 0 assert str(device) == "cuda:0" assert device.__repr__() == "device(type='cuda', index=0)" def test_device_from_str() -> None: device = tvm_ffi.device("ext_dev:0") assert device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLExtDev assert device.index == 0 assert str(device) == "ext_dev:0" assert device.__repr__() == "device(type='ext_dev', index=0)" @pytest.mark.parametrize( "dev_str, expected_device_type, expect_device_id", [ ("cpu", DLDeviceType.kDLCPU, 0), ("cuda", DLDeviceType.kDLCUDA, 0), ("cuda:0", DLDeviceType.kDLCUDA, 0), ("cuda:3", DLDeviceType.kDLCUDA, 3), ("metal:2", DLDeviceType.kDLMetal, 2), ], ) def test_device_dlpack_device_type( dev_str: str, expected_device_type: DLDeviceType, expect_device_id: int, ) -> None: dev = tvm_ffi.device(dev_str) assert dev.dlpack_device_type() == expected_device_type assert dev.index == expect_device_id @pytest.mark.parametrize( "dev_type, dev_id, expected_device_type, expect_device_id", [ ("cpu", 0, DLDeviceType.kDLCPU, 0), ("cuda", 0, DLDeviceType.kDLCUDA, 0), (DLDeviceType.kDLCUDA, 0, DLDeviceType.kDLCUDA, 0), ("cuda", 3, DLDeviceType.kDLCUDA, 3), (DLDeviceType.kDLMetal, 2, DLDeviceType.kDLMetal, 2), # id from numpy ("cpu", numpy.int32(1), DLDeviceType.kDLCPU, 1), # id from torch (py dependency not ready in environment) # ("cpu", torch.tensor(1, dtype=torch.int32), DLDeviceType.kDLCPU, 1), ], ) def test_device_with_dev_id( dev_type: str | DLDeviceType, dev_id: int, expected_device_type: DLDeviceType, expect_device_id: int, ) -> None: dev = tvm_ffi.device(dev_type, dev_id) assert dev.dlpack_device_type() == expected_device_type assert dev.index == expect_device_id @pytest.mark.parametrize("dev_type, dev_id", [("cpu:0:0", None), ("cpu:?", None), ("cpu:", None)]) def test_deive_type_error(dev_type: str, dev_id: int | None) -> None: with pytest.raises(ValueError): tvm_ffi.device(dev_type, dev_id) def test_deive_id_error() -> None: with pytest.raises(TypeError): tvm_ffi.device("cpu", "?") # ty: ignore[invalid-argument-type] def test_device_pickle() -> None: device = tvm_ffi.device("cuda", 0) device_pickled = pickle.loads(pickle.dumps(device)) assert device_pickled.dlpack_device_type() == device.dlpack_device_type() assert device_pickled.index == device.index def test_device_class_override() -> None: class MyDevice(tvm_ffi.Device): pass old_device = tvm_ffi.core._CLASS_DEVICE tvm_ffi.core._set_class_device(MyDevice) device = tvm_ffi.device("cuda", 0) assert isinstance(device, MyDevice) tvm_ffi.core._set_class_device(old_device) def test_cuda_stream_handling() -> None: class MyDummyStream: def __init__(self, stream: int) -> None: self.stream = stream def __cuda_stream__(self) -> tuple[str, int]: return ("cuda", self.stream) stream = MyDummyStream(1) echo = tvm_ffi.get_global_func("testing.echo") y = echo(stream) assert isinstance(y, ctypes.c_void_p) assert y.value == 1 tvm-ffi-0.1.12/tests/python/test_dlpack_exchange_api.py000066400000000000000000000307161521067262500232050ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file to # you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import ctypes import sys import pytest try: import torch # Import tvm_ffi to load the DLPack exchange API extension # This sets torch.Tensor.__dlpack_c_exchange_api__ import tvm_ffi from torch.utils import cpp_extension from tvm_ffi import libinfo except ImportError: torch = None # ty: ignore[invalid-assignment] # Check if DLPack Exchange API is available _has_dlpack_api = torch is not None and hasattr(torch.Tensor, "__dlpack_c_exchange_api__") _has_gpu = torch is not None and torch.cuda.is_available() @pytest.mark.skipif(not _has_dlpack_api, reason="PyTorch DLPack Exchange API not available") def test_dlpack_exchange_api() -> None: # xfail the test on windows platform, it seems to be a bug in torch extension building on windows if sys.platform.startswith("win"): pytest.xfail("DLPack Exchange API test is known to fail on Windows platform") assert torch is not None assert hasattr(torch.Tensor, "__dlpack_c_exchange_api__") api_attr = torch.Tensor.__dlpack_c_exchange_api__ # PyCapsule - extract the pointer as integer pythonapi = ctypes.pythonapi # Set restype to c_size_t to get integer directly (avoids c_void_p quirks) pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] capsule_name = b"dlpack_exchange_api" api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, capsule_name) assert api_ptr != 0, "API pointer from PyCapsule should not be NULL" tensor = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) source = """ #include #include #include void test_dlpack_api(at::Tensor tensor, int64_t api_ptr_int, bool cuda_available) { DLPackExchangeAPI* api = reinterpret_cast(api_ptr_int); // Test 1: API structure and version { TORCH_CHECK(api != nullptr, "API pointer is NULL"); TORCH_CHECK(api->header.version.major == DLPACK_MAJOR_VERSION, "Expected major version ", DLPACK_MAJOR_VERSION, ", got ", api->header.version.major); TORCH_CHECK(api->header.version.minor == DLPACK_MINOR_VERSION, "Expected minor version ", DLPACK_MINOR_VERSION, ", got ", api->header.version.minor); TORCH_CHECK(api->managed_tensor_allocator != nullptr, "managed_tensor_allocator is NULL"); TORCH_CHECK(api->managed_tensor_from_py_object_no_sync != nullptr, "managed_tensor_from_py_object_no_sync is NULL"); TORCH_CHECK(api->managed_tensor_to_py_object_no_sync != nullptr, "managed_tensor_to_py_object_no_sync is NULL"); TORCH_CHECK(api->dltensor_from_py_object_no_sync != nullptr, "dltensor_from_py_object_no_sync is NULL"); TORCH_CHECK(api->current_work_stream != nullptr, "current_work_stream is NULL"); } // Test 2: managed_tensor_allocator { DLTensor prototype; prototype.device.device_type = kDLCPU; prototype.device.device_id = 0; prototype.ndim = 3; int64_t shape[3] = {3, 4, 5}; prototype.shape = shape; prototype.strides = nullptr; DLDataType dtype; dtype.code = kDLFloat; dtype.bits = 32; dtype.lanes = 1; prototype.dtype = dtype; prototype.data = nullptr; prototype.byte_offset = 0; DLManagedTensorVersioned* out_tensor = nullptr; int result = api->managed_tensor_allocator(&prototype, &out_tensor, nullptr, nullptr); TORCH_CHECK(result == 0, "Allocator failed with code ", result); TORCH_CHECK(out_tensor != nullptr, "Allocator returned NULL"); TORCH_CHECK(out_tensor->dl_tensor.ndim == 3, "Expected ndim 3, got ", out_tensor->dl_tensor.ndim); TORCH_CHECK(out_tensor->dl_tensor.shape[0] == 3, "Expected shape[0] = 3, got ", out_tensor->dl_tensor.shape[0]); TORCH_CHECK(out_tensor->dl_tensor.shape[1] == 4, "Expected shape[1] = 4, got ", out_tensor->dl_tensor.shape[1]); TORCH_CHECK(out_tensor->dl_tensor.shape[2] == 5, "Expected shape[2] = 5, got ", out_tensor->dl_tensor.shape[2]); TORCH_CHECK(out_tensor->dl_tensor.dtype.code == kDLFloat, "Expected dtype code kDLFloat, got ", out_tensor->dl_tensor.dtype.code); TORCH_CHECK(out_tensor->dl_tensor.dtype.bits == 32, "Expected dtype bits 32, got ", out_tensor->dl_tensor.dtype.bits); TORCH_CHECK(out_tensor->dl_tensor.device.device_type == kDLCPU, "Expected device type kDLCPU, got ", out_tensor->dl_tensor.device.device_type); if (out_tensor->deleter) { out_tensor->deleter(out_tensor); } } // Test 3: managed_tensor_from_py_object_no_sync { std::unique_ptr py_obj(THPVariable_Wrap(tensor), &Py_DecRef); TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); DLManagedTensorVersioned* out_tensor = nullptr; int result = api->managed_tensor_from_py_object_no_sync(py_obj.get(), &out_tensor); TORCH_CHECK(result == 0, "from_py_object_no_sync failed with code ", result); TORCH_CHECK(out_tensor != nullptr, "from_py_object_no_sync returned NULL"); TORCH_CHECK(out_tensor->version.major == DLPACK_MAJOR_VERSION, "Expected major version ", DLPACK_MAJOR_VERSION, ", got ", out_tensor->version.major); TORCH_CHECK(out_tensor->version.minor == DLPACK_MINOR_VERSION, "Expected minor version ", DLPACK_MINOR_VERSION, ", got ", out_tensor->version.minor); TORCH_CHECK(out_tensor->dl_tensor.ndim == 3, "Expected ndim 3, got ", out_tensor->dl_tensor.ndim); TORCH_CHECK(out_tensor->dl_tensor.shape[0] == 2, "Expected shape[0] = 2, got ", out_tensor->dl_tensor.shape[0]); TORCH_CHECK(out_tensor->dl_tensor.shape[1] == 3, "Expected shape[1] = 3, got ", out_tensor->dl_tensor.shape[1]); TORCH_CHECK(out_tensor->dl_tensor.shape[2] == 4, "Expected shape[2] = 4, got ", out_tensor->dl_tensor.shape[2]); TORCH_CHECK(out_tensor->dl_tensor.dtype.code == kDLFloat, "Expected dtype code kDLFloat, got ", out_tensor->dl_tensor.dtype.code); TORCH_CHECK(out_tensor->dl_tensor.dtype.bits == 32, "Expected dtype bits 32, got ", out_tensor->dl_tensor.dtype.bits); TORCH_CHECK(out_tensor->dl_tensor.data != nullptr, "Data pointer is NULL"); if (out_tensor->deleter) { out_tensor->deleter(out_tensor); } } // Test 4: managed_tensor_to_py_object_no_sync { std::unique_ptr py_obj(THPVariable_Wrap(tensor), &Py_DecRef); TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); DLManagedTensorVersioned* managed_tensor = nullptr; int result = api->managed_tensor_from_py_object_no_sync(py_obj.get(), &managed_tensor); TORCH_CHECK(result == 0, "from_py_object_no_sync failed"); TORCH_CHECK(managed_tensor != nullptr, "from_py_object_no_sync returned NULL"); std::unique_ptr py_obj_out(nullptr, &Py_DecRef); PyObject* py_obj_out_raw = nullptr; result = api->managed_tensor_to_py_object_no_sync(managed_tensor, reinterpret_cast(&py_obj_out_raw)); py_obj_out.reset(py_obj_out_raw); TORCH_CHECK(result == 0, "to_py_object_no_sync failed with code ", result); TORCH_CHECK(py_obj_out.get() != nullptr, "to_py_object_no_sync returned NULL"); TORCH_CHECK(THPVariable_Check(py_obj_out.get()), "Returned PyObject is not a Tensor"); at::Tensor result_tensor = THPVariable_Unpack(py_obj_out.get()); TORCH_CHECK(result_tensor.dim() == 3, "Expected 3 dimensions, got ", result_tensor.dim()); TORCH_CHECK(result_tensor.size(0) == 2, "Expected size(0) = 2, got ", result_tensor.size(0)); TORCH_CHECK(result_tensor.size(1) == 3, "Expected size(1) = 3, got ", result_tensor.size(1)); TORCH_CHECK(result_tensor.size(2) == 4, "Expected size(2) = 4, got ", result_tensor.size(2)); TORCH_CHECK(result_tensor.scalar_type() == at::kFloat, "Expected dtype kFloat, got ", result_tensor.scalar_type()); } // Test 5: dltensor_from_py_object_no_sync { std::unique_ptr py_obj(THPVariable_Wrap(tensor), &Py_DecRef); TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); DLTensor dltensor; int result = api->dltensor_from_py_object_no_sync(py_obj.get(), &dltensor); TORCH_CHECK(result == 0, "dltensor_from_py_object_no_sync failed with code ", result); TORCH_CHECK(dltensor.ndim == 3, "Expected ndim 3, got ", dltensor.ndim); TORCH_CHECK(dltensor.shape[0] == 2, "Expected shape[0] = 2, got ", dltensor.shape[0]); TORCH_CHECK(dltensor.shape[1] == 3, "Expected shape[1] = 3, got ", dltensor.shape[1]); TORCH_CHECK(dltensor.shape[2] == 4, "Expected shape[2] = 4, got ", dltensor.shape[2]); TORCH_CHECK(dltensor.dtype.code == kDLFloat, "Expected dtype code kDLFloat, got ", dltensor.dtype.code); TORCH_CHECK(dltensor.dtype.bits == 32, "Expected dtype bits 32, got ", dltensor.dtype.bits); TORCH_CHECK(dltensor.data != nullptr, "Data pointer is NULL"); } // Test 6: current_work_stream (CUDA if available, otherwise CPU) { void* stream_out = nullptr; DLDeviceType device_type = cuda_available ? kDLCUDA : kDLCPU; int result = api->current_work_stream(device_type, 0, &stream_out); TORCH_CHECK(result == 0, "current_work_stream failed with code ", result); } } """ include_paths = libinfo.include_paths() if torch.cuda.is_available(): include_paths += cpp_extension.include_paths("cuda") mod = cpp_extension.load_inline( name="dlpack_test", cpp_sources=[source], functions=["test_dlpack_api"], extra_include_paths=include_paths, ) # Run the comprehensive test mod.test_dlpack_api(tensor, api_ptr, torch.cuda.is_available()) @pytest.mark.skipif( not (_has_dlpack_api and _has_gpu), reason="PyTorch DLPack Exchange API with GPU is not available", ) def test_dlpack_exchange_api_gpu_tensor_metadata() -> None: assert torch is not None echo = tvm_ffi.get_global_func("testing.echo") for shape in [(512,), (512, 512), (2, 3, 4)]: source = torch.empty(shape, device="cuda", dtype=torch.float16) tvm_tensor = tvm_ffi.from_dlpack(source) assert tvm_tensor.shape == shape assert tvm_tensor.dtype == tvm_ffi.dtype("float16") echoed = echo(source) assert tuple(echoed.shape) == shape assert echoed.dtype == source.dtype assert echoed.device == source.device @pytest.mark.skipif(not _has_dlpack_api, reason="PyTorch DLPack Exchange API not available") def test_from_dlpack_torch() -> None: # Covers from_dlpack to use fallback fastpath assert torch is not None tensor = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) tensor_from_dlpack = tvm_ffi.from_dlpack(tensor) assert tensor_from_dlpack.shape == tensor.shape assert tensor_from_dlpack.dtype == tvm_ffi.float32 if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) tvm-ffi-0.1.12/tests/python/test_dtype.py000066400000000000000000000166501521067262500204020ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import pickle from typing import Any import numpy as np import pytest import tvm_ffi from packaging.version import Version def test_dtype() -> None: float32 = tvm_ffi.dtype("float32") assert float32.__repr__() == "dtype('float32')" assert type(float32) == tvm_ffi.dtype x = np.array([1, 2, 3], dtype=float32) assert x.dtype == float32 @pytest.mark.parametrize( "dtype_str, expected_size", [ ("float32", 4), ("float32x4", 16), ("float8_e5m2x4", 4), ("float6_e2m3fnx4", 3), ("float4_e2m1fnx4", 2), ("uint8", 1), ("bool", 1), ], ) def test_dtype_itemsize(dtype_str: str, expected_size: int) -> None: dtype = tvm_ffi.dtype(dtype_str) assert dtype.itemsize == expected_size @pytest.mark.parametrize("dtype_str", ["int32xvscalex4"]) def test_dtype_itemmize_error(dtype_str: str) -> None: with pytest.raises(ValueError): tvm_ffi.dtype(dtype_str).itemsize @pytest.mark.parametrize( "dtype_str", [ "float32", "float32x4", "float8_e5m2x4", "float6_e2m3fnx4", "float4_e2m1fnx4", "uint8", "bool", ], ) def test_dtype_pickle(dtype_str: str) -> None: dtype = tvm_ffi.dtype(dtype_str) dtype_pickled = pickle.loads(pickle.dumps(dtype)) assert dtype_pickled.type_code == dtype.type_code assert dtype_pickled.bits == dtype.bits assert dtype_pickled.lanes == dtype.lanes @pytest.mark.parametrize("dtype_str", ["float32", "bool"]) def test_dtype_with_lanes(dtype_str: str) -> None: dtype = tvm_ffi.dtype(dtype_str) dtype_with_lanes = dtype.with_lanes(4) assert dtype_with_lanes.type_code == dtype.type_code assert dtype_with_lanes.bits == dtype.bits assert dtype_with_lanes.lanes == 4 _fecho = tvm_ffi.get_global_func("testing.echo") def _check_dtype(dtype: Any, code: int, bits: int, lanes: int) -> None: echo_dtype = _fecho(dtype) assert isinstance(echo_dtype, tvm_ffi.dtype) assert echo_dtype.type_code == code assert echo_dtype.bits == bits assert echo_dtype.lanes == lanes converted_dtype = tvm_ffi.convert(dtype) assert isinstance(converted_dtype, tvm_ffi.dtype) assert converted_dtype.type_code == code assert converted_dtype.bits == bits assert converted_dtype.lanes == lanes def test_torch_dtype_conversion() -> None: torch = pytest.importorskip("torch") _check_dtype(torch.int8, 0, 8, 1) _check_dtype(torch.short, 0, 16, 1) _check_dtype(torch.int16, 0, 16, 1) _check_dtype(torch.int32, 0, 32, 1) _check_dtype(torch.int, 0, 32, 1) _check_dtype(torch.int64, 0, 64, 1) _check_dtype(torch.long, 0, 64, 1) _check_dtype(torch.uint8, 1, 8, 1) _check_dtype(torch.uint16, 1, 16, 1) _check_dtype(torch.uint32, 1, 32, 1) _check_dtype(torch.uint64, 1, 64, 1) _check_dtype(torch.float16, 2, 16, 1) _check_dtype(torch.half, 2, 16, 1) _check_dtype(torch.float32, 2, 32, 1) _check_dtype(torch.float, 2, 32, 1) _check_dtype(torch.float64, 2, 64, 1) _check_dtype(torch.double, 2, 64, 1) _check_dtype(torch.bfloat16, 4, 16, 1) _check_dtype(torch.bool, 6, 8, 1) _check_dtype(torch.float8_e4m3fn, 10, 8, 1) _check_dtype(torch.float8_e4m3fnuz, 11, 8, 1) _check_dtype(torch.float8_e5m2, 12, 8, 1) _check_dtype(torch.float8_e5m2fnuz, 13, 8, 1) if hasattr(torch, "float8_e8m0fnu"): _check_dtype(torch.float8_e8m0fnu, 14, 8, 1) def test_numpy_dtype_conversion() -> None: np = pytest.importorskip("numpy") _check_dtype(np.dtype(np.int8), 0, 8, 1) _check_dtype(np.dtype(np.int16), 0, 16, 1) _check_dtype(np.dtype(np.int32), 0, 32, 1) _check_dtype(np.dtype(np.int64), 0, 64, 1) _check_dtype(np.dtype(np.uint8), 1, 8, 1) _check_dtype(np.dtype(np.uint16), 1, 16, 1) _check_dtype(np.dtype(np.uint32), 1, 32, 1) _check_dtype(np.dtype(np.uint64), 1, 64, 1) _check_dtype(np.dtype(np.float16), 2, 16, 1) _check_dtype(np.dtype(np.float32), 2, 32, 1) _check_dtype(np.dtype(np.float64), 2, 64, 1) def test_ml_dtypes_dtype_conversion() -> None: np = pytest.importorskip("numpy") ml_dtypes = pytest.importorskip("ml_dtypes") if Version(ml_dtypes.__version__) < Version("0.4.0"): pytest.skip("ml_dtypes < 0.4.0") # ty: ignore[invalid-argument-type, too-many-positional-arguments] return _check_dtype(np.dtype(ml_dtypes.int2), 0, 2, 1) _check_dtype(np.dtype(ml_dtypes.int4), 0, 4, 1) _check_dtype(np.dtype(ml_dtypes.uint2), 1, 2, 1) _check_dtype(np.dtype(ml_dtypes.uint4), 1, 4, 1) _check_dtype(np.dtype(ml_dtypes.bfloat16), 4, 16, 1) _check_dtype(np.dtype(ml_dtypes.float8_e3m4), 7, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e4m3), 8, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e4m3b11fnuz), 9, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e4m3fn), 10, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e4m3fnuz), 11, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e5m2), 12, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e5m2fnuz), 13, 8, 1) _check_dtype(np.dtype(ml_dtypes.float8_e8m0fnu), 14, 8, 1) _check_dtype(np.dtype(ml_dtypes.float6_e2m3fn), 15, 6, 1) _check_dtype(np.dtype(ml_dtypes.float6_e3m2fn), 16, 6, 1) _check_dtype(np.dtype(ml_dtypes.float4_e2m1fn), 17, 4, 1) def test_builtin_dtype_conversion() -> None: _check_dtype(tvm_ffi.bool, 6, 8, 1) _check_dtype(tvm_ffi.int8, 0, 8, 1) _check_dtype(tvm_ffi.int16, 0, 16, 1) _check_dtype(tvm_ffi.int32, 0, 32, 1) _check_dtype(tvm_ffi.int64, 0, 64, 1) _check_dtype(tvm_ffi.uint8, 1, 8, 1) _check_dtype(tvm_ffi.uint16, 1, 16, 1) _check_dtype(tvm_ffi.uint32, 1, 32, 1) _check_dtype(tvm_ffi.uint64, 1, 64, 1) _check_dtype(tvm_ffi.float16, 2, 16, 1) _check_dtype(tvm_ffi.float32, 2, 32, 1) _check_dtype(tvm_ffi.float64, 2, 64, 1) _check_dtype(tvm_ffi.bfloat16, 4, 16, 1) _check_dtype(tvm_ffi.float8_e4m3fn, 10, 8, 1) _check_dtype(tvm_ffi.float8_e4m3fnuz, 11, 8, 1) _check_dtype(tvm_ffi.float8_e5m2, 12, 8, 1) _check_dtype(tvm_ffi.float8_e5m2fnuz, 13, 8, 1) _check_dtype(tvm_ffi.float8_e8m0fnu, 14, 8, 1) _check_dtype(tvm_ffi.float4_e2m1fnx2, 17, 4, 2) def test_dtype_from_dlpack_data_type() -> None: dtype = tvm_ffi.dtype.from_dlpack_data_type((0, 8, 1)) assert dtype.type_code == 0 assert dtype.bits == 8 assert dtype.lanes == 1 def test_dtype_bool() -> None: dtype = tvm_ffi.dtype("bool") assert dtype.type_code == 6 assert dtype.bits == 8 assert dtype.lanes == 1 dtype_with_lanes = dtype.with_lanes(4) assert dtype_with_lanes.type_code == 6 assert dtype_with_lanes.bits == 8 assert dtype_with_lanes.lanes == 4 assert dtype_with_lanes == "boolx4" tvm-ffi-0.1.12/tests/python/test_error.py000066400000000000000000000132671521067262500204070ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import gc import weakref from typing import NoReturn import pytest import tvm_ffi def test_parse_backtrace() -> None: backtrace = """ File "test.py", line 1, in File "test.py", line 3, in run_test """ parsed = tvm_ffi.error._parse_backtrace(backtrace) assert len(parsed) == 2 assert parsed[0] == ("test.py", 1, "") assert parsed[1] == ("test.py", 3, "run_test") def test_error_from_cxx() -> None: test_raise_error = tvm_ffi.get_global_func("testing.test_raise_error") try: test_raise_error("ValueError", "error XYZ") except ValueError as e: assert e.__tvm_ffi_error__.kind == "ValueError" # ty: ignore[unresolved-attribute] assert e.__tvm_ffi_error__.message == "error XYZ" # ty: ignore[unresolved-attribute] assert e.__tvm_ffi_error__.backtrace.find("TestRaiseError") != -1 # ty: ignore[unresolved-attribute] fapply = tvm_ffi.convert(lambda f, *args: f(*args)) with pytest.raises(TypeError): fapply(test_raise_error, "TypeError", "error XYZ") # wrong number of arguments with pytest.raises(TypeError): tvm_ffi.convert(lambda x: x)() def test_error_from_nested_pyfunc() -> None: fapply = tvm_ffi.convert(lambda f, *args: f(*args)) cxx_test_raise_error = tvm_ffi.get_global_func("testing.test_raise_error") cxx_test_apply = tvm_ffi.get_global_func("testing.apply") record_object = [] def raise_error() -> None: try: fapply(cxx_test_raise_error, "ValueError", "error XYZ") except ValueError as e: assert e.__tvm_ffi_error__.kind == "ValueError" # ty: ignore[unresolved-attribute] assert e.__tvm_ffi_error__.message == "error XYZ" # ty: ignore[unresolved-attribute] assert e.__tvm_ffi_error__.backtrace.find("TestRaiseError") != -1 # ty: ignore[unresolved-attribute] record_object.append(e.__tvm_ffi_error__) # ty: ignore[unresolved-attribute] raise e try: cxx_test_apply(raise_error) except ValueError as e: backtrace = e.__tvm_ffi_error__.backtrace # ty: ignore[unresolved-attribute] assert e.__tvm_ffi_error__.same_as(record_object[0]) # ty: ignore[unresolved-attribute] assert backtrace.count("TestRaiseError") == 1 # The following lines may fail if debug symbols are missing try: assert backtrace.count("TestApply") == 1 assert backtrace.count("") == 1 pos_cxx_raise = backtrace.find("TestRaiseError") pos_cxx_apply = backtrace.find("TestApply") pos_lambda = backtrace.find("") assert pos_cxx_raise < pos_lambda assert pos_lambda < pos_cxx_apply except Exception as e: pytest.xfail("May fail if debug symbols are missing") # ty: ignore[invalid-argument-type, too-many-positional-arguments] def test_error_traceback_update() -> None: fecho = tvm_ffi.get_global_func("testing.echo") def raise_error() -> NoReturn: raise ValueError("error XYZ") try: raise_error() except ValueError as e: ffi_error = tvm_ffi.convert(e) assert ffi_error.backtrace.find("raise_error") != -1 def raise_cxx_error() -> None: cxx_test_raise_error = tvm_ffi.get_global_func("testing.test_raise_error") cxx_test_raise_error("ValueError", "error XYZ") try: raise_cxx_error() except ValueError as e: assert e.__tvm_ffi_error__.backtrace.find("raise_cxx_error") == -1 # ty: ignore[unresolved-attribute] ffi_error1 = tvm_ffi.convert(e) ffi_error2 = fecho(e) assert ffi_error1.backtrace.find("raise_cxx_error") != -1 assert ffi_error2.backtrace.find("raise_cxx_error") != -1 def test_error_no_cyclic_reference() -> None: # This test case ensures that when an error is raised from C++ side, # there is no cyclic reference that slows down the garbage collection. # Please see `_with_append_backtrace` in error.py # temporarily disable gc gc.disable() try: # We should create a class as a probe to detect gc activity # beacuse weakref doesn't support list, dict or other trivial types class SampleObject: ... # trigger a C++ side KeyError by accessing a non-existent key def trigger_cpp_side_error() -> None: try: tmp_map: tvm_ffi.Map = tvm_ffi.Map(dict()) tmp_map["a"] except KeyError: pass def may_create_cyclic_reference() -> weakref.ReferenceType: obj = SampleObject() trigger_cpp_side_error() return weakref.ref(obj) wref = may_create_cyclic_reference() # if the object is not collected, wref() will return the object assert wref() is None, "Cyclic reference occurs inside error handling pipeline" finally: # re-enable gc whenever exception occurs gc.enable() tvm-ffi-0.1.12/tests/python/test_examples.py000066400000000000000000000032341521067262500210650ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # testcases appearing in example docstrings from typing import Any import tvm_ffi def test_register_global_func() -> None: # we can use decorator to register a function @tvm_ffi.register_global_func("example.echo") def echo(x: Any) -> Any: return x # After registering, we can get the function by its name f = tvm_ffi.get_global_func("example.echo") assert f(1) == 1 # we can also directly register a function tvm_ffi.register_global_func("example.add_one", lambda x: x + 1) f = tvm_ffi.get_global_func("example.add_one") assert f(1) == 2 def test_array() -> None: a = tvm_ffi.convert([1, 2, 3]) assert isinstance(a, tvm_ffi.Array) assert len(a) == 3 def test_map() -> None: amap = tvm_ffi.convert({"a": 1, "b": 2}) assert isinstance(amap, tvm_ffi.Map) assert len(amap) == 2 assert amap["a"] == 1 assert amap["b"] == 2 tvm-ffi-0.1.12/tests/python/test_free_threaded_python_helpers.py000066400000000000000000000023221521067262500251500ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import sys import pytest import tvm_ffi def _is_free_threaded_python() -> bool: return hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled() @pytest.mark.skipif(not _is_free_threaded_python(), reason="requires free-threaded Python") def test_pyobject_deleter_handles_last_ref() -> None: drop_last_ref = getattr(tvm_ffi.core, "_testing_drop_last_ref_without_thread_state") drop_last_ref() tvm-ffi-0.1.12/tests/python/test_function.py000066400000000000000000000415101521067262500210730ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import ctypes import gc import sys from typing import Any import numpy as np import pytest import tvm_ffi import tvm_ffi.cpp try: import torch except ImportError: torch = None # ty: ignore[invalid-assignment] _HAS_TORCH_DLPACK_API = torch is not None and hasattr(torch.Tensor, "__dlpack_c_exchange_api__") def test_echo() -> None: fecho = tvm_ffi.get_global_func("testing.echo") assert isinstance(fecho, tvm_ffi.Function) # test each type assert fecho(None) is None # test bool bool_result = fecho(True) assert isinstance(bool_result, bool) assert bool_result is True bool_result = fecho(False) assert isinstance(bool_result, bool) assert bool_result is False # test int/float assert fecho(1) == 1 assert fecho(1.2) == 1.2 # test str str_result = fecho("hello") assert isinstance(str_result, str) assert str_result == "hello" # test bytes bytes_result = fecho(b"abc") assert isinstance(bytes_result, bytes) assert bytes_result == b"abc" # test dtype dtype_result = fecho(tvm_ffi.dtype("float32")) assert isinstance(dtype_result, tvm_ffi.dtype) assert dtype_result == tvm_ffi.dtype("float32") # test device device_result = fecho(tvm_ffi.device("cuda:1")) assert isinstance(device_result, tvm_ffi.Device) assert device_result.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCUDA assert device_result.index == 1 assert str(device_result) == "cuda:1" assert device_result.__repr__() == "device(type='cuda', index=1)" # test c_void_p c_void_p_result = fecho(ctypes.c_void_p(0x12345678)) assert isinstance(c_void_p_result, ctypes.c_void_p) assert c_void_p_result.value == 0x12345678 # test c_void_p for nullptr c_void_p_nullptr_result = fecho(ctypes.c_void_p(0)) c_void_p_nullptr_result is None # test function: aka object fadd = tvm_ffi.convert(lambda a, b: a + b) fadd1 = fecho(fadd) assert fadd1(1, 2) == 3 assert fadd1.same_as(fadd) def check_tensor() -> None: np_data = np.arange(10, dtype="int32") if not hasattr(np_data, "__dlpack__"): return # test Tensor x = tvm_ffi.from_dlpack(np_data) assert isinstance(x, tvm_ffi.Tensor) tensor_result = fecho(x) assert isinstance(tensor_result, tvm_ffi.Tensor) assert tensor_result.shape == (10,) assert tensor_result.dtype == tvm_ffi.dtype("int32") assert tensor_result.device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCPU assert tensor_result.device.index == 0 check_tensor() def test_return_raw_str_bytes() -> None: assert tvm_ffi.convert(lambda: "hello")() == "hello" assert tvm_ffi.convert(lambda: b"hello")() == b"hello" assert tvm_ffi.convert(lambda: bytearray(b"hello"))() == b"hello" def test_string_bytes_passing() -> None: fecho = tvm_ffi.get_global_func("testing.echo") use_count = tvm_ffi.get_global_func("testing.object_use_count") # small string assert fecho("hello") == "hello" # large string x = "hello" * 100 y = fecho(x) assert y == x assert y._tvm_ffi_cached_object is not None use_count(y) == 1 # small bytes assert fecho(b"hello") == b"hello" # large bytes x2 = b"hello" * 100 y2 = fecho(x2) assert y2 == x2 assert y2._tvm_ffi_cached_object is not None fecho(y2) == 1 def test_nested_container_passing() -> None: # test and make sure our ref counting is correct fecho = tvm_ffi.get_global_func("testing.echo") use_count = tvm_ffi.get_global_func("testing.object_use_count") obj = tvm_ffi.convert((1, 2, 3)) assert use_count(obj) == 1 y = fecho([obj, {"a": 1, "b": obj}]) assert use_count(y) == 1 assert use_count(obj) == 3 assert use_count(y[1]) == 2 def test_pyfunc_convert() -> None: def add(a: int, b: int) -> int: return a + b fadd = tvm_ffi.convert(add) assert isinstance(fadd, tvm_ffi.Function) assert fadd(1, 2) == 3 def fapply(f: Any, *args: Any) -> Any: return f(*args) fapply = tvm_ffi.convert(fapply) assert fapply(add, 1, 3.3) == 4.3 def test_global_func() -> None: @tvm_ffi.register_global_func("mytest.echo") def echo(x: Any) -> Any: return x f = tvm_ffi.get_global_func("mytest.echo") assert f.same_as(echo) assert f(1) == 1 assert "mytest.echo" in tvm_ffi.registry.list_global_func_names() tvm_ffi.registry.remove_global_func("mytest.echo") assert "mytest.echo" not in tvm_ffi.registry.list_global_func_names() assert tvm_ffi.get_global_func("mytest.echo", allow_missing=True) is None def test_rvalue_ref() -> None: use_count = tvm_ffi.get_global_func("testing.object_use_count") def callback(x: Any, expected_count: int) -> Any: # The use count of TVM FFI objects is decremented as part of # `ObjectRef.__del__`, which runs when the Python object is # destructed. However, Python object destruction is not # deterministic, and even CPython's reference-counting is # considered an implementation detail. Therefore, to ensure # correct results from this test, `gc.collect()` must be # explicitly called. gc.collect() assert expected_count == use_count(x) return x._move() f = tvm_ffi.convert(callback) def check0() -> None: x = tvm_ffi.convert([1, 2]) assert use_count(x) == 1 f(x, 2) f(x._move(), 1) assert x.__ctypes_handle__().value is None def check1() -> None: x = tvm_ffi.convert([1, 2]) assert use_count(x) == 1 y = f(x, 2) f(x._move(), 2) assert x.__ctypes_handle__().value is None assert y.__ctypes_handle__().value is not None check0() check1() def test_echo_with_opaque_object() -> None: class MyObject: def __init__(self, value: Any) -> None: self.value = value fecho = tvm_ffi.get_global_func("testing.echo") x = MyObject("hello") assert sys.getrefcount(x) == 2 y = fecho(x) assert isinstance(y, MyObject) assert y is x assert sys.getrefcount(x) == 3 def py_callback(z: Any) -> Any: """Python callback with opaque object.""" assert z is x return z fcallback = tvm_ffi.convert(py_callback) z = fcallback(x) assert z is x assert sys.getrefcount(x) == 4 def test_function_from_c_symbol() -> None: add_one_c_symbol = tvm_ffi.get_global_func("testing.get_add_one_c_symbol")() fadd_one = tvm_ffi.Function.__from_extern_c__(add_one_c_symbol) assert fadd_one(1) == 2 assert fadd_one(2) == 3 with pytest.raises(TypeError): fadd_one(None) keep_alive = [1, 2, 3] base_ref_count = sys.getrefcount(keep_alive) fadd_one = tvm_ffi.Function.__from_extern_c__(add_one_c_symbol, keep_alive_object=keep_alive) assert fadd_one(1) == 2 assert fadd_one(2) == 3 assert sys.getrefcount(keep_alive) == base_ref_count + 1 fadd_one = None assert sys.getrefcount(keep_alive) == base_ref_count def test_function_from_mlir_packed_safe_call() -> None: add_one_c_symbol = tvm_ffi.get_global_func("testing.get_mlir_add_one_c_symbol")() fadd_one = tvm_ffi.Function.__from_mlir_packed_safe_call__(add_one_c_symbol) assert fadd_one(1) == 2 assert fadd_one(2) == 3 keep_alive = [1, 2, 3] base_ref_count = sys.getrefcount(keep_alive) fadd_one = tvm_ffi.Function.__from_mlir_packed_safe_call__( add_one_c_symbol, keep_alive_object=keep_alive ) with pytest.raises(TypeError): fadd_one(None) assert fadd_one(1) == 2 assert fadd_one(2) == 3 assert sys.getrefcount(keep_alive) == base_ref_count + 1 fadd_one = None assert sys.getrefcount(keep_alive) == base_ref_count def test_function_subclass() -> None: class JitFunction: def __init__(self, metadata: Any) -> None: self.metadata = metadata class MyFunction(tvm_ffi.Function, JitFunction): def __init__(self, metadata: Any) -> None: # Explicitly initialize the mixin. `super()` is not used because `tvm_ffi.Function` # is an extension type without a standard `__init__`. JitFunction.__init__(self, metadata) # When subclassing a Cython cdef class and overriding `__init__`, # special methods like `__call__` may not be inherited automatically. # This explicit assignment ensures the subclass remains callable. __call__ = tvm_ffi.Function.__call__ f = tvm_ffi.convert(lambda x: x) assert isinstance(f, tvm_ffi.Function) f_sub = MyFunction(128) # move handle from f to f_sub an existing function f_sub.__move_handle_from__(f) assert isinstance(f_sub, MyFunction) assert isinstance(f_sub, JitFunction) assert f_sub.metadata == 128 y: int = f_sub(2) assert y == 2 echo = tvm_ffi.get_global_func("testing.echo") fechoed = echo(f_sub) assert isinstance(fechoed, tvm_ffi.Function) assert fechoed.__chandle__() == f_sub.__chandle__() assert fechoed(10) == 10 def test_function_with_opaque_ptr_protocol() -> None: class MyObject: def __init__(self, value: Any) -> None: self.value = value def __tvm_ffi_opaque_ptr__(self) -> Any: return self.value fecho = tvm_ffi.get_global_func("testing.echo") x = MyObject(10) y = fecho(x) assert isinstance(y, ctypes.c_void_p) assert y.value == 10 def test_function_with_dlpack_data_type_protocol() -> None: class DLPackDataTypeProtocol: def __init__(self, dlpack_data_type: tuple[int, int, int]) -> None: self.dlpack_data_type = dlpack_data_type def __dlpack_data_type__(self) -> tuple[int, int, int]: return self.dlpack_data_type dtype = tvm_ffi.dtype("float32") fecho = tvm_ffi.get_global_func("testing.echo") x = DLPackDataTypeProtocol((dtype.type_code, dtype.bits, dtype.lanes)) y = fecho(x) assert y == dtype converted_y = tvm_ffi.convert(x) assert converted_y == dtype def test_function_with_dlpack_device_protocol() -> None: device = tvm_ffi.device("cuda:1") class DLPackDeviceProtocol: def __init__(self, device: tvm_ffi.Device) -> None: self.device = device def __dlpack_device__(self) -> tuple[int, int]: return (self.device.dlpack_device_type(), self.device.index) fecho = tvm_ffi.get_global_func("testing.echo") x = DLPackDeviceProtocol(device) y = fecho(x) assert y == device def test_integral_float_variants_passing() -> None: fecho = tvm_ffi.get_global_func("testing.echo") y = fecho(np.int32(1)) assert isinstance(y, int) assert y == 1 y = fecho(np.float64(2.0)) assert isinstance(y, float) assert y == 2.0 class IntProtocol: def __init__(self, value: int) -> None: self.value = value def __tvm_ffi_int__(self) -> int: return self.value y = fecho(IntProtocol(10)) assert isinstance(y, int) assert y == 10 class FloatProtocol: def __init__(self, value: float) -> None: self.value = value def __tvm_ffi_float__(self) -> float: return self.value y = fecho(FloatProtocol(10)) assert isinstance(y, float) assert y == 10 def test_function_with_value_protocol() -> None: class ValueProtocol: def __init__(self, value: Any) -> None: self.value = value def __tvm_ffi_value__(self) -> Any: return self.value fecho = tvm_ffi.get_global_func("testing.echo") assert fecho(ValueProtocol(10)) == 10 assert tuple(fecho(ValueProtocol([1, 2, 3]))) == (1, 2, 3) assert tuple(fecho(ValueProtocol([1, 2, ValueProtocol(3)]))) == (1, 2, 3) nested_value_protocol = ValueProtocol(ValueProtocol(ValueProtocol(10))) assert fecho(nested_value_protocol) == 10 nested_value_protocol = ValueProtocol([ValueProtocol(1), ValueProtocol(2), ValueProtocol(3)]) assert tuple(fecho(nested_value_protocol)) == (1, 2, 3) def test_convert_func_tensor_cls_missing_attribute() -> None: """Passing a class without __dlpack_c_exchange_api__ raises TypeError.""" class DummyTensor: pass with pytest.raises(TypeError, match="__dlpack_c_exchange_api__"): tvm_ffi.convert_func(lambda x: x, tensor_cls=DummyTensor) with pytest.raises(TypeError, match="__dlpack_c_exchange_api__"): tvm_ffi.convert_func(lambda x: x, tensor_cls=object) def test_convert_func_raises_propagates() -> None: """An exception raised inside the callback propagates out to the caller.""" def raises(x: int) -> None: raise ValueError(f"boom {x}") f = tvm_ffi.convert_func(raises) with pytest.raises(ValueError, match="boom 42"): f(42) @pytest.mark.skipif( not _HAS_TORCH_DLPACK_API, reason="torch.Tensor.__dlpack_c_exchange_api__ not available", ) def test_convert_func_with_torch_tensor_cls() -> None: """tensor_cls=torch.Tensor delivers torch.Tensor instances to the callback. Asserts the type *inside* the callback (which runs on the C++ -> Python side of the FFI boundary) — the return value's Python type depends on the outer caller's conversion path, so we verify shape survives the round-trip rather than isinstance on the return. """ calls = 0 def callback(a: Any, b: Any, c: Any) -> Any: nonlocal calls calls += 1 assert isinstance(a, torch.Tensor) assert isinstance(b, torch.Tensor) assert isinstance(c, torch.Tensor) assert list(a.shape) == [2] assert list(b.shape) == [3] assert list(c.shape) == [4] return b f = tvm_ffi.convert_func(callback, tensor_cls=torch.Tensor) a = torch.zeros(2) b = torch.ones(3) c = torch.full((4,), 2.0) out = f(a, b, c) assert calls == 1 assert tuple(out.shape) == (3,) def test_callback_rawstr_and_bytearrayptr_args() -> None: """Regression: C++ -> Python callback with kTVMFFIRawStr / kTVMFFIByteArrayPtr args. When C++ invokes a Python callback with non-owning RawStr or ByteArrayPtr arg shapes, the callback arg setter must materialise a Python str / bytes directly rather than hitting the ``raise ValueError`` guard that formerly existed in ``TVMFFICyCallbackArgSetterFactory``. Two trampolines are compiled via cpp.load_inline: - ``invoke_with_raw_str(callback)`` — calls ``callback("hello rawstr")`` using a C-string literal, which the TypeTraits pack as kTVMFFIRawStr. - ``invoke_with_byte_array_ptr(callback)`` — calls ``callback(&byte_arr)`` where ``byte_arr`` is a ``TVMFFIByteArray`` on the stack, packed as kTVMFFIByteArrayPtr. """ mod = tvm_ffi.cpp.load_inline( name="test_callback_rawstr_bytearrayptr", cpp_sources=r""" void invoke_with_raw_str(tvm::ffi::Function callback) { // Passing a string literal packs as kTVMFFIRawStr (const char* TypeTraits). callback("hello rawstr"); } void invoke_with_byte_array_ptr(tvm::ffi::Function callback) { // Passing a TVMFFIByteArray* packs as kTVMFFIByteArrayPtr. static const char kData[] = "hello bytearrayptr"; TVMFFIByteArray byte_arr{kData, sizeof(kData) - 1}; callback(&byte_arr); } """, functions=["invoke_with_raw_str", "invoke_with_byte_array_ptr"], ) # --- kTVMFFIRawStr path --- str_received: list[Any] = [] str_cb = tvm_ffi.convert(lambda x: str_received.append(x)) mod.invoke_with_raw_str(str_cb) assert len(str_received) == 1 assert isinstance(str_received[0], str), f"expected str, got {type(str_received[0])}" assert str_received[0] == "hello rawstr" # --- kTVMFFIByteArrayPtr path --- bytes_received: list[Any] = [] bytes_cb = tvm_ffi.convert(lambda x: bytes_received.append(x)) mod.invoke_with_byte_array_ptr(bytes_cb) assert len(bytes_received) == 1 assert isinstance(bytes_received[0], bytes), f"expected bytes, got {type(bytes_received[0])}" assert bytes_received[0] == b"hello bytearrayptr" tvm-ffi-0.1.12/tests/python/test_libinfo.py000066400000000000000000000123521521067262500206720ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for the tvm-ffi-config command line utility.""" import subprocess import sys from pathlib import Path from typing import Callable import pytest from tvm_ffi import libinfo def _stdout_for(*args: str) -> str: """Invoke tvm-ffi-config with the provided arguments and return stdout with trailing whitespace removed.""" result = subprocess.run( [ sys.executable, "-m", "tvm_ffi.config", *args, ], check=True, capture_output=True, text=True, ) assert result.stderr == "" return result.stdout.strip() @pytest.mark.parametrize( ("flag", "expected_fn", "is_dir"), [ ("--includedir", libinfo.find_include_path, True), ("--dlpack-includedir", libinfo.find_dlpack_include_path, True), ("--cmakedir", libinfo.find_cmake_path, True), ("--sourcedir", libinfo.find_source_path, True), ("--cython-lib-path", libinfo.find_cython_lib, False), ], ) def test_basic_path_flags(flag: str, expected_fn: Callable[[], str], is_dir: bool) -> None: output = _stdout_for(flag) assert output == expected_fn() path = Path(output) assert path.exists() assert path.is_dir() if is_dir else path.is_file() def test_libdir_matches_library_parent() -> None: expected_dir = Path(libinfo.find_libtvm_ffi()).parent output = _stdout_for("--libdir") assert output == str(expected_dir) assert Path(output).is_dir() assert Path(libinfo.find_libtvm_ffi()).is_file() def test_libfiles_reports_platform_library() -> None: output = _stdout_for("--libfiles") if sys.platform.startswith("win32"): expected = libinfo.find_windows_implib() else: expected = libinfo.find_libtvm_ffi() assert output == expected assert Path(output).is_file() def test_libs_reports_link_target() -> None: output = _stdout_for("--libs") if sys.platform.startswith("win32"): assert output == libinfo.find_windows_implib() else: assert output == "-ltvm_ffi" def test_cxxflags_include_paths_and_standard() -> None: include_dir = libinfo.find_include_path() dlpack_dir = libinfo.find_dlpack_include_path() assert _stdout_for("--cxxflags") == f"-I{include_dir} -I{dlpack_dir} -std=c++17" def test_cflags_include_paths() -> None: include_dir = libinfo.find_include_path() dlpack_dir = libinfo.find_dlpack_include_path() assert _stdout_for("--cflags") == f"-I{include_dir} -I{dlpack_dir}" def test_ldflags_only_on_unix() -> None: output = _stdout_for("--ldflags") if sys.platform.startswith("win32"): assert output == "" else: libdir = Path(libinfo.find_libtvm_ffi()).parent assert output == f"-L{libdir}" assert libdir.is_dir() def test_cmakedir_contains_config_file() -> None: cmake_dir = Path(_stdout_for("--cmakedir")) assert (cmake_dir / "tvm_ffi-config.cmake").is_file() def test_find_python_helper_include_path() -> None: path = libinfo.find_python_helper_include_path() assert Path(path).is_dir() assert (Path(path) / "tvm_ffi_python_helpers.h").is_file() def _platform_lib_filename(target_name: str) -> str: """Return the platform-appropriate shared library filename for ``target_name``.""" if sys.platform.startswith("win32"): return f"{target_name}.dll" if sys.platform.startswith("darwin"): return f"lib{target_name}.dylib" return f"lib{target_name}.so" def test_find_library_by_basename_extra_lib_paths(tmp_path: Path) -> None: """A directory passed via ``extra_lib_paths`` is searched in the dev fallback.""" lib_file = tmp_path / _platform_lib_filename("dummy") lib_file.write_bytes(b"") result = libinfo._find_library_by_basename( "nonexistent-dist-name", "dummy", extra_lib_paths=[tmp_path] ) assert isinstance(result, Path) assert result.resolve() == lib_file.resolve() # Sanity: the result must live under the temp tree, not tvm_ffi's own tree. assert tmp_path.resolve() in result.resolve().parents def test_find_library_by_basename_extra_lib_paths_type_error(tmp_path: Path) -> None: """Non-Path elements in ``extra_lib_paths`` raise a TypeError naming the offender.""" with pytest.raises(TypeError, match=r"extra_lib_paths\[1\] must be a pathlib\.Path"): libinfo._find_library_by_basename( "apache-tvm-ffi", "tvm_ffi", extra_lib_paths=[tmp_path, str(tmp_path)], # ty: ignore[invalid-argument-type] ) tvm-ffi-0.1.12/tests/python/test_load_inline.py000066400000000000000000000366611521067262500215360ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import numpy import pytest try: import torch except ImportError: torch = None # ty: ignore[invalid-assignment] import tvm_ffi.cpp from tvm_ffi.module import Module def test_load_inline_cpp() -> None: mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions=["add_one_cpu"], ) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) def test_load_inline_cpp_with_docstrings() -> None: mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions={"add_one_cpu": "add two float32 1D tensors element-wise"}, ) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) def test_load_inline_cpp_multiple_sources() -> None: mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=[ r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, r""" void add_two_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 2; } } """, ], functions=["add_one_cpu", "add_two_cpu"], ) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) def test_load_inline_cpp_build_dir() -> None: mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } """, functions=["add_one_cpu"], build_directory="./build/build_add_one", ) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_load_inline_cuda() -> None: mod: Module = tvm_ffi.cpp.load_inline( name="hello", cuda_sources=r""" __global__ void AddOneKernel(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1; } } void add_one_cuda(tvm::ffi::TensorView x, tvm::ffi::TensorView y, int64_t raw_stream) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; int64_t n = x.size(0); int64_t nthread_per_block = 256; int64_t nblock = (n + nthread_per_block - 1) / nthread_per_block; // Obtain the current stream from the environment // it will be set to torch.cuda.current_stream() when calling the function // with torch.Tensors cudaStream_t stream = static_cast( TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); TVM_FFI_ICHECK_EQ(reinterpret_cast(stream), raw_stream) << "stream must be the same as raw_stream"; // launch the kernel AddOneKernel<<>>(static_cast(x.data_ptr()), static_cast(y.data_ptr()), n); } """, functions=["add_one_cuda"], ) if torch is not None: # test with raw stream x_cuda = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32, device="cuda") y_cuda = torch.empty_like(x_cuda) mod.add_one_cuda(x_cuda, y_cuda, 0) torch.testing.assert_close(x_cuda + 1, y_cuda) # test with torch stream y_cuda = torch.empty_like(x_cuda) stream = torch.cuda.Stream() with torch.cuda.stream(stream): mod.add_one_cuda(x_cuda, y_cuda, stream.cuda_stream) stream.synchronize() torch.testing.assert_close(x_cuda + 1, y_cuda) @pytest.mark.skipif(torch is None, reason="Requires torch") def test_load_inline_with_env_tensor_allocator() -> None: assert torch is not None if not hasattr(torch.Tensor, "__dlpack_c_exchange_api__"): pytest.skip("Torch does not support __dlpack_c_exchange_api__") # ty: ignore[invalid-argument-type, too-many-positional-arguments] mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" #include #include #include namespace ffi = tvm::ffi; ffi::Tensor return_add_one(ffi::Map> kwargs) { ffi::Tensor x = kwargs["x"].get<0>(); // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; // allocate a new tensor with the env tensor allocator // it will be redirected to torch.empty when calling the function ffi::Tensor y = ffi::Tensor::FromEnvAlloc( TVMFFIEnvTensorAlloc, ffi::Shape({x.size(0)}), f32_dtype, x.device()); int64_t n = x.size(0); for (int i = 0; i < n; ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } return y; } """, functions=["return_add_one"], ) assert torch is not None def run_check() -> None: """Must run in a separate function to ensure deletion happens before mod unloads. When a module returns an object, the object deleter address is part of the loaded library. We need to keep the module loaded until the object is deleted. """ assert torch is not None x_cpu = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32, device="cpu") # test support for nested container passing y_cpu = mod.return_add_one({"x": [x_cpu]}) assert isinstance(y_cpu, torch.Tensor) assert y_cpu.shape == (5,) assert y_cpu.dtype == torch.float32 torch.testing.assert_close(x_cpu + 1, y_cpu) run_check() @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_load_inline_both() -> None: assert torch is not None mod: Module = tvm_ffi.cpp.load_inline( name="hello", cpp_sources=r""" void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; for (int i = 0; i < x.size(0); ++i) { static_cast(y.data_ptr())[i] = static_cast(x.data_ptr())[i] + 1; } } void add_one_cuda(tvm::ffi::TensorView x, tvm::ffi::TensorView y); """, cuda_sources=r""" __global__ void AddOneKernel(float* x, float* y, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { y[idx] = x[idx] + 1; } } void add_one_cuda(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { // implementation of a library function TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; DLDataType f32_dtype{kDLFloat, 32, 1}; TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; int64_t n = x.size(0); int64_t nthread_per_block = 256; int64_t nblock = (n + nthread_per_block - 1) / nthread_per_block; // Obtain the current stream from the environment // it will be set to torch.cuda.current_stream() when calling the function // with torch.Tensors cudaStream_t stream = static_cast( TVMFFIEnvGetStream(x.device().device_type, x.device().device_id)); // launch the kernel AddOneKernel<<>>(static_cast(x.data_ptr()), static_cast(y.data_ptr()), n); } """, functions=["add_one_cpu", "add_one_cuda"], ) x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32) y = numpy.empty_like(x) mod.add_one_cpu(x, y) numpy.testing.assert_equal(x + 1, y) x_cuda = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32, device="cuda") y_cuda = torch.empty_like(x_cuda) mod.add_one_cuda(x_cuda, y_cuda) torch.testing.assert_close(x_cuda + 1, y_cuda) @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_cuda_memory_alloc_noleak() -> None: assert torch is not None mod = tvm_ffi.cpp.load_inline( name="hello", cuda_sources=r""" #include #include namespace ffi = tvm::ffi; ffi::Tensor return_tensor(tvm::ffi::TensorView x) { ffi::Tensor y = ffi::Tensor::FromEnvAlloc( TVMFFIEnvTensorAlloc, x.shape(), x.dtype(), x.device()); return y; } """, functions=["return_tensor"], ) def run_check() -> None: """Must run in a separate function to ensure deletion happens before mod unloads.""" assert torch is not None x = torch.arange(1024 * 1024, dtype=torch.float32, device="cuda") current_allocated = torch.cuda.memory_allocated() repeat = 8 for i in range(repeat): mod.return_tensor(x) diff = torch.cuda.memory_allocated() - current_allocated # memory should not grow as we loop over assert diff <= 1024**2 * 8 run_check() tvm-ffi-0.1.12/tests/python/test_metadata.py000066400000000000000000000221511521067262500210260ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from typing import Any import pytest from tvm_ffi import get_global_func_metadata, register_global_func, remove_global_func from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr from tvm_ffi.testing import _SchemaAllTypes def _replace_container_types(ty: str) -> str: return { "Array": "Sequence", "List": "MutableSequence", "Map": "Mapping", "Dict": "MutableMapping", }.get(ty, ty) @pytest.mark.parametrize( "func_name,expected", [ ("testing.schema_id_int", "Callable[[int], int]"), ("testing.schema_id_float", "Callable[[float], float]"), ("testing.schema_id_bool", "Callable[[bool], bool]"), ("testing.schema_id_device", "Callable[[Device], Device]"), ("testing.schema_id_dtype", "Callable[[dtype], dtype]"), ("testing.schema_id_string", "Callable[[str], str]"), ("testing.schema_id_bytes", "Callable[[bytes], bytes]"), ("testing.schema_id_func", "Callable[[Callable[..., Any]], Callable[..., Any]]"), ( "testing.schema_id_func_typed", "Callable[[Callable[[int, float, Callable[..., Any]], None]], Callable[[int, float, Callable[..., Any]], None]]", ), ("testing.schema_id_any", "Callable[[Any], Any]"), ("testing.schema_id_object", "Callable[[Object], Object]"), ("testing.schema_id_dltensor", "Callable[[Tensor], Tensor]"), ("testing.schema_id_tensor", "Callable[[Tensor], Tensor]"), ("testing.schema_tensor_view_input", "Callable[[Tensor], None]"), ("testing.schema_id_opt_int", "Callable[[int | None], int | None]"), ("testing.schema_id_opt_str", "Callable[[str | None], str | None]"), ("testing.schema_id_opt_obj", "Callable[[Object | None], Object | None]"), ("testing.schema_id_arr_int", "Callable[[Array[int]], Array[int]]"), ("testing.schema_id_arr_str", "Callable[[Array[str]], Array[str]]"), ("testing.schema_id_arr_obj", "Callable[[Array[Object]], Array[Object]]"), ("testing.schema_id_arr", "Callable[[Array[Any]], Array[Any]]"), ("testing.schema_id_map_str_int", "Callable[[Map[str, int]], Map[str, int]]"), ("testing.schema_id_map_str_str", "Callable[[Map[str, str]], Map[str, str]]"), ("testing.schema_id_map_str_obj", "Callable[[Map[str, Object]], Map[str, Object]]"), ("testing.schema_id_map", "Callable[[Map[Any, Any]], Map[Any, Any]]"), ("testing.schema_id_dict_str_int", "Callable[[Dict[str, int]], Dict[str, int]]"), ("testing.schema_id_dict_str_str", "Callable[[Dict[str, str]], Dict[str, str]]"), ("testing.schema_id_variant_int_str", "Callable[[int | str], int | str]"), ("testing.schema_packed", "Callable[..., Any]"), ( "testing.schema_arr_map_opt", "Callable[[Array[int | None], Map[str, Array[int]], str | None], Map[str, Array[int]]]", ), ( "testing.schema_variant_mix", "Callable[[int | str | Array[int]], int | str | Array[int]]", ), ("testing.schema_no_args", "Callable[[], int]"), ("testing.schema_no_return", "Callable[[int], None]"), ("testing.schema_no_args_no_return", "Callable[[], None]"), ], ) def test_schema_global_func(func_name: str, expected: str) -> None: metadata: dict[str, Any] = get_global_func_metadata(func_name) actual: TypeSchema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(actual) == expected, f"{func_name}: {actual}" assert actual.repr(_replace_container_types) == ( expected.replace("Array", "Sequence") .replace("List", "MutableSequence") .replace("Map", "Mapping") .replace("Dict", "MutableMapping") ) @pytest.mark.parametrize( "field_name,expected", [ ("v_bool", "bool"), ("v_int", "int"), ("v_float", "float"), ("v_device", "Device"), ("v_dtype", "dtype"), ("v_string", "str"), ("v_bytes", "bytes"), ("v_opt_int", "int | None"), ("v_opt_str", "str | None"), ("v_arr_int", "Array[int]"), ("v_arr_str", "Array[str]"), ("v_map_str_int", "Map[str, int]"), ("v_map_str_arr_int", "Map[str, Array[int]]"), ("v_variant", "str | Array[int] | Map[str, int]"), ("v_opt_arr_variant", "Array[int | str] | None"), ], ) def test_schema_field(field_name: str, expected: str) -> None: type_info: TypeInfo = getattr(_SchemaAllTypes, "__tvm_ffi_type_info__") for field in type_info.fields: if field.name == field_name: actual: TypeSchema = TypeSchema.from_json_str(field.metadata["type_schema"]) assert str(actual) == expected, f"{field_name}: {actual}" assert actual.repr(_replace_container_types) == ( expected.replace("Array", "Sequence") .replace("List", "MutableSequence") .replace("Map", "Mapping") .replace("Dict", "MutableMapping") ) break else: raise ValueError(f"Field not found: {field_name}") @pytest.mark.parametrize( "method_name,expected", [ ("add_int", "Callable[[testing.SchemaAllTypes, int], int]"), ("append_int", "Callable[[testing.SchemaAllTypes, Array[int], int], Array[int]]"), ("maybe_concat", "Callable[[testing.SchemaAllTypes, str | None, str | None], str | None]"), ( "merge_map", "Callable[[testing.SchemaAllTypes, Map[str, Array[int]], Map[str, Array[int]]], Map[str, Array[int]]]", ), ("make_with", "Callable[[int, float, str], testing.SchemaAllTypes]"), ], ) def test_schema_member_method(method_name: str, expected: str) -> None: type_info: TypeInfo = getattr(_SchemaAllTypes, "__tvm_ffi_type_info__") for method in type_info.methods: if method.name == method_name: actual: TypeSchema = TypeSchema.from_json_str(method.metadata["type_schema"]) assert str(actual) == expected, f"{method_name}: {actual}" assert actual.repr(_replace_container_types) == ( expected.replace("Array", "Sequence") .replace("List", "MutableSequence") .replace("Map", "Mapping") .replace("Dict", "MutableMapping") ) break else: raise ValueError(f"Method not found: {method_name}") def test_metadata_global_func() -> None: metadata: dict[str, Any] = get_global_func_metadata("testing.schema_id_int") assert len(metadata) == 4 assert "type_schema" in metadata assert metadata["bool_attr"] is True assert metadata["int_attr"] == 1 assert metadata["str_attr"] == "hello" def test_metadata_empty_for_python_func() -> None: @register_global_func("test.python_func_no_metadata") def simple_func(x: int) -> int: return x + 1 metadata = get_global_func_metadata("test.python_func_no_metadata") assert metadata == {} remove_global_func("test.python_func_no_metadata") def test_metadata_field() -> None: type_info: TypeInfo = getattr(_SchemaAllTypes, "__tvm_ffi_type_info__") for field in type_info.fields: if field.name == "v_bool": assert len(field.metadata) == 4 assert "type_schema" in field.metadata assert field.metadata["bool_attr"] is True assert field.metadata["int_attr"] == 1 assert field.metadata["str_attr"] == "hello" break else: raise ValueError("Field not found: v_bool") def test_metadata_member_method() -> None: type_info: TypeInfo = getattr(_SchemaAllTypes, "__tvm_ffi_type_info__") for method in type_info.methods: if method.name == "add_int": assert len(method.metadata) == 4 assert "type_schema" in method.metadata assert method.metadata["bool_attr"] is True assert method.metadata["int_attr"] == 1 assert method.metadata["str_attr"] == "hello" break else: raise ValueError("Method not found: add_int") def test_mem_fn_as_global_func() -> None: metadata: dict[str, Any] = get_global_func_metadata("testing.TestIntPairSum") type_schema: TypeSchema = TypeSchema.from_json_str(metadata["type_schema"]) assert str(type_schema) == "Callable[[testing.TestIntPair], int]" def test_class_metadata_none() -> None: assert _lookup_type_attr(64, "__metadata__") is None tvm-ffi-0.1.12/tests/python/test_object.py000066400000000000000000000436101521067262500205170ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import sys from typing import Any, Sequence import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi.core import TypeInfo def test_make_object() -> None: # with default values obj0 = tvm_ffi.testing.create_object("testing.TestObjectBase") assert isinstance(obj0, tvm_ffi.testing.TestObjectBase) assert obj0.v_i64 == 10 assert obj0.v_f64 == 10.0 assert obj0.v_str == "hello" def test_make_object_via_init() -> None: obj0 = tvm_ffi.testing.TestIntPair(1, 2) assert obj0.a == 1 assert obj0.b == 2 def test_method() -> None: obj0 = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=12) assert isinstance(obj0, tvm_ffi.testing.TestObjectBase) assert obj0.add_i64(1) == 13 assert type(obj0).add_i64.__doc__ == "add_i64 method" assert type(obj0).v_i64.__doc__ == "i64 field" def test_attribute() -> None: obj = tvm_ffi.testing.TestIntPair(3, 4) assert obj.a == 3 assert obj.b == 4 assert type(obj).a.__doc__ == "Field `a`" assert type(obj).b.__doc__ == "Field `b`" def test_attribute_no_doc() -> None: from tvm_ffi.testing import TestObjectBase # noqa: PLC0415 # The docs for v_f64 and v_str are None because they are not explicitly set. assert TestObjectBase.v_i64.__doc__ == "i64 field" assert TestObjectBase.v_f64.__doc__ is None assert TestObjectBase.v_str.__doc__ is None def test_setter() -> None: # test setter obj0 = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=10, v_str="hello") assert isinstance(obj0, tvm_ffi.testing.TestObjectBase) assert obj0.v_i64 == 10 obj0.v_i64 = 11 assert obj0.v_i64 == 11 obj0.v_str = "world" assert obj0.v_str == "world" with pytest.raises(TypeError): obj0.v_str = 1 # ty: ignore[invalid-assignment] with pytest.raises(TypeError): obj0.v_i64 = "hello" # ty: ignore[invalid-assignment] def test_derived_object() -> None: with pytest.raises(TypeError): obj0 = tvm_ffi.testing.create_object("testing.TestObjectDerived") v_map = tvm_ffi.convert({"a": 1}) v_array = tvm_ffi.convert([1, 2, 3]) obj0 = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=20, v_map=v_map, v_array=v_array ) assert isinstance(obj0, tvm_ffi.testing.TestObjectDerived) assert obj0.v_map.same_as(v_map) # ty: ignore[unresolved-attribute] assert obj0.v_array.same_as(v_array) # ty: ignore[unresolved-attribute] assert obj0.v_i64 == 20 assert obj0.v_f64 == 10.0 assert obj0.v_str == "hello" obj0.v_i64 = 21 assert obj0.v_i64 == 21 class MyObject: def __init__(self, value: Any) -> None: self.value = value def test_opaque_object() -> None: obj0 = MyObject("hello") base_count = sys.getrefcount(obj0) ref_count = tvm_ffi.get_global_func("testing.object_use_count") assert sys.getrefcount(obj0) == base_count obj0_converted = tvm_ffi.convert(obj0) assert ref_count(obj0_converted) == 1 assert sys.getrefcount(obj0) == base_count + 1 assert isinstance(obj0_converted, tvm_ffi.core.OpaquePyObject) obj0_cpy = obj0_converted.pyobject() assert obj0_cpy is obj0 assert sys.getrefcount(obj0) == base_count + 2 obj0_converted = None assert sys.getrefcount(obj0) == base_count + 1 obj0_cpy = None assert sys.getrefcount(obj0) == base_count def test_opaque_type_error() -> None: obj0 = MyObject("hello") with pytest.raises(TypeError) as e: tvm_ffi.testing.add_one(obj0) # ty: ignore[invalid-argument-type] assert ( "Mismatched type on argument #0 when calling: `testing.add_one(0: int) -> int`. Expected `int` but got `ffi.OpaquePyObject`" in str(e.value) ) def test_object_init() -> None: # Registered class with auto-generated __ffi_init__ (all fields have defaults) obj = tvm_ffi.testing.TestObjectBase() assert obj.v_i64 == 10 assert obj.v_f64 == 10.0 assert obj.v_str == "hello" # Registered class with __c_ffi_init__ should work fine pair = tvm_ffi.testing.TestIntPair(3, 4) assert pair.a == 3 and pair.b == 4 # FFI-returned objects should work fine obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=7) assert obj.__chandle__() != 0 assert obj.v_i64 == 7 # ty: ignore[unresolved-attribute] def test_register_object_wires_init() -> None: """Regression: register_object must wire __init__ from __ffi_init__. Without _install_init in register_object, calling TestIntPair(1, 2) falls through to object.__init__ which silently ignores args on Cython extension types (CObject has custom tp_new). The result is chandle = NULL. Any subsequent field access dereferences chandle + offset → segfault. """ # Construction must produce a live handle, not NULL. pair = tvm_ffi.testing.TestIntPair(3, 7) assert pair.__chandle__() != 0 assert pair.a == 3 assert pair.b == 7 # __new__ alone must NOT produce a usable object (chandle stays NULL). bare = tvm_ffi.testing.TestIntPair.__new__(tvm_ffi.testing.TestIntPair) assert bare.__chandle__() == 0 def test_id_and_is() -> None: x = tvm_ffi.testing.TestIntPair(1, 2) y = tvm_ffi.testing.TestIntPair(1, 2) # id_ is the integer handle, aliasing __chandle__. assert x.id_ == x.__chandle__() assert x.id_ != 0 assert x.id_ != y.id_ # is_ tests handle identity, aliasing same_as. assert x.is_(x) assert not x.is_(y) assert x.is_(x) == x.same_as(x) assert not x.is_(object()) assert not x.is_(None) def test_object_protocol() -> None: class CompactObject: def __init__(self, backend_obj: Any) -> None: self.backend_obj = backend_obj def __tvm_ffi_object__(self) -> Any: return self.backend_obj x = tvm_ffi.convert([]) assert isinstance(x, tvm_ffi.Object) x_compact = CompactObject(x) test_echo = tvm_ffi.get_global_func("testing.echo") y = test_echo(x_compact) assert y.__chandle__() == x.__chandle__() def test_unregistered_object_fallback() -> None: def _check_type(x: Any) -> None: type_info: TypeInfo = type(x).__tvm_ffi_type_info__ assert type_info.type_key == "testing.TestUnregisteredObject" assert x.v1 == 41 assert x.v2 == 42 assert x.get_v1_plus_one() == 42 assert x.get_v2_plus_two() == 44 assert type(x).__name__ == "TestUnregisteredObject" assert type(x).__module__ == "testing" assert type(x).__qualname__ == "testing.TestUnregisteredObject" assert "Auto-generated fallback class" in type(x).__doc__ assert "Get (v1 + 1) from TestUnregisteredBaseObject" in type(x).get_v1_plus_one.__doc__ assert "Get (v2 + 2) from TestUnregisteredObject" in type(x).get_v2_plus_two.__doc__ obj = tvm_ffi.testing.make_unregistered_object() _check_type(obj) for _ in range(5): obj = tvm_ffi.testing.make_unregistered_object() _check_type(obj) @pytest.mark.parametrize( ("test_cls", "make_instance"), [ ( tvm_ffi.testing.TestObjectBase, lambda: tvm_ffi.testing.create_object("testing.TestObjectBase"), ), ( tvm_ffi.testing.TestIntPair, lambda: tvm_ffi.testing.TestIntPair(1, 2), ), ( tvm_ffi.testing.TestObjectDerived, lambda: tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=20, v_map=tvm_ffi.convert({"a": 1}), v_array=tvm_ffi.convert([1, 2]), ), ), ], ) def test_object_subclass_slots(test_cls: type, make_instance: Any) -> None: slots = test_cls.__dict__.get("__slots__") assert slots == () assert "__dict__" not in test_cls.__dict__ assert "__weakref__" not in test_cls.__dict__ obj = make_instance() with pytest.raises(AttributeError): obj._tvm_ffi_test_attr = "nope" @pytest.mark.parametrize( "test_cls, type_key, parent_cls", [ (tvm_ffi.Object, "ffi.Object", None), (tvm_ffi.Tensor, "ffi.Tensor", tvm_ffi.Object), (tvm_ffi.core.DataType, "DataType", None), (tvm_ffi.Device, "Device", None), (tvm_ffi.Shape, "ffi.Shape", tvm_ffi.Object), (tvm_ffi.Module, "ffi.Module", tvm_ffi.Object), (tvm_ffi.Function, "ffi.Function", tvm_ffi.Object), (tvm_ffi.core.Error, "ffi.Error", tvm_ffi.Object), (tvm_ffi.core.String, "ffi.String", tvm_ffi.Object), (tvm_ffi.core.Bytes, "ffi.Bytes", tvm_ffi.Object), (tvm_ffi.Tensor, "ffi.Tensor", tvm_ffi.Object), (tvm_ffi.Array, "ffi.Array", tvm_ffi.Object), (tvm_ffi.List, "ffi.List", tvm_ffi.Object), (tvm_ffi.Map, "ffi.Map", tvm_ffi.Object), (tvm_ffi.access_path.AccessStep, "ffi.reflection.AccessStep", tvm_ffi.Object), (tvm_ffi.access_path.AccessPath, "ffi.reflection.AccessPath", tvm_ffi.Object), (tvm_ffi.testing.TestIntPair, "testing.TestIntPair", tvm_ffi.Object), (tvm_ffi.testing.TestObjectBase, "testing.TestObjectBase", tvm_ffi.Object), ( tvm_ffi.testing.TestObjectDerived, "testing.TestObjectDerived", tvm_ffi.testing.TestObjectBase, ), (tvm_ffi.testing._TestCxxClassBase, "testing.TestCxxClassBase", tvm_ffi.Object), ( tvm_ffi.testing._TestCxxClassDerived, "testing.TestCxxClassDerived", tvm_ffi.testing._TestCxxClassBase, ), ( tvm_ffi.testing._TestCxxClassDerivedDerived, "testing.TestCxxClassDerivedDerived", tvm_ffi.testing._TestCxxClassDerived, ), (tvm_ffi.testing._TestCxxInitSubset, "testing.TestCxxInitSubset", tvm_ffi.Object), (tvm_ffi.testing._SchemaAllTypes, "testing.SchemaAllTypes", tvm_ffi.Object), ], ) def test_type_info_attachment(test_cls: type, type_key: str, parent_cls: type | None) -> None: type_info = tvm_ffi.core._type_cls_to_type_info(test_cls) assert type_info is not None assert type_info.type_cls is test_cls assert type_info.type_key == type_key, ( f"Expected type key `{type_key}` but got `{type_info.type_key}`" ) parent_type_info = type_info.parent_type_info if parent_type_info is None: assert parent_cls is None else: assert parent_type_info.type_cls is parent_cls, ( f"Expected parent type {parent_cls}, but got {parent_type_info.type_cls}" ) def test_get_registered_type_keys() -> None: keys = tvm_ffi.registry.get_registered_type_keys() assert isinstance(keys, Sequence) assert all(isinstance(k, str) for k in keys) keys = set(keys) assert "ffi.Object" in keys assert "ffi.String" in keys for ty in [ "None", "int", "bool", "float", "void*", "DataType", "Device", "DLTensor*", "const char*", "TVMFFIByteArray*", "ObjectRValueRef", ]: assert ty in keys, f"Expected to find `{ty}` in registered type keys, but it was not found." keys.remove(ty) for ty in keys: assert ty.startswith("ffi.") or ty.startswith("testing."), ( f"Expected type key `{ty}` to start with `ffi.` or `testing.`" ) # --------------------------------------------------------------------------- # isinstance / issubclass correctness for the Object type hierarchy # --------------------------------------------------------------------------- class TestIsinstanceIssubclass: """Verify that isinstance/issubclass respect the actual type hierarchy. Regression tests for a bug where _ObjectSlotsMeta.__instancecheck__ and __subclasscheck__ returned True for *any* CObject against *any* Object subclass, making e.g. isinstance(Map(...), Array) == True. """ # -- containers: sibling types must not match each other ---------------- def test_map_not_isinstance_array(self) -> None: """Map instance should not pass isinstance check for Array.""" m = tvm_ffi.Map({"a": 1}) assert not isinstance(m, tvm_ffi.Array) def test_array_not_isinstance_map(self) -> None: """Array instance should not pass isinstance check for Map.""" a = tvm_ffi.Array([1, 2]) assert not isinstance(a, tvm_ffi.Map) def test_list_not_isinstance_dict(self) -> None: """List instance should not pass isinstance check for Dict.""" lst = tvm_ffi.List([1, 2]) assert not isinstance(lst, tvm_ffi.Dict) def test_dict_not_isinstance_list(self) -> None: """Dict instance should not pass isinstance check for List.""" d = tvm_ffi.Dict({"k": 1}) assert not isinstance(d, tvm_ffi.List) def test_array_not_isinstance_list(self) -> None: """Array instance should not pass isinstance check for List.""" a = tvm_ffi.Array([1]) assert not isinstance(a, tvm_ffi.List) def test_map_not_isinstance_dict(self) -> None: """Map instance should not pass isinstance check for Dict.""" m = tvm_ffi.Map({"a": 1}) assert not isinstance(m, tvm_ffi.Dict) # -- containers: positive isinstance ------------------------------------ def test_array_isinstance_object(self) -> None: """Array instance should pass isinstance check for Object.""" a = tvm_ffi.Array([1]) assert isinstance(a, tvm_ffi.Object) def test_map_isinstance_object(self) -> None: """Map instance should pass isinstance check for Object.""" m = tvm_ffi.Map({"a": 1}) assert isinstance(m, tvm_ffi.Object) def test_list_isinstance_object(self) -> None: """List instance should pass isinstance check for Object.""" lst = tvm_ffi.List([1]) assert isinstance(lst, tvm_ffi.Object) def test_dict_isinstance_object(self) -> None: """Dict instance should pass isinstance check for Object.""" d = tvm_ffi.Dict({"k": 1}) assert isinstance(d, tvm_ffi.Object) # -- registered user types: parent / child ------------------------------ def test_derived_isinstance_base(self) -> None: """Derived object should pass isinstance check for its base class.""" obj = tvm_ffi.testing.TestObjectDerived( v_map={"a": 1}, v_array=[1], ) assert isinstance(obj, tvm_ffi.testing.TestObjectBase) assert isinstance(obj, tvm_ffi.Object) def test_base_not_isinstance_derived(self) -> None: """Base object should not pass isinstance check for a derived class.""" obj = tvm_ffi.testing.TestObjectBase() assert not isinstance(obj, tvm_ffi.testing.TestObjectDerived) def test_derived_isinstance_own_class(self) -> None: """Derived object should pass isinstance check for its own class.""" obj = tvm_ffi.testing.TestObjectDerived( v_map={"a": 1}, v_array=[1], ) assert isinstance(obj, tvm_ffi.testing.TestObjectDerived) # -- cross-hierarchy: user object vs container -------------------------- def test_user_object_not_isinstance_array(self) -> None: """User-defined object should not pass isinstance check for Array.""" obj = tvm_ffi.testing.TestObjectBase() assert not isinstance(obj, tvm_ffi.Array) def test_array_not_isinstance_user_object(self) -> None: """Array should not pass isinstance check for a user-defined type.""" a = tvm_ffi.Array([1]) assert not isinstance(a, tvm_ffi.testing.TestObjectBase) # -- issubclass mirrors isinstance -------------------------------------- def test_issubclass_derived_base(self) -> None: """Derived class should be a subclass of its base.""" assert issubclass(tvm_ffi.testing.TestObjectDerived, tvm_ffi.testing.TestObjectBase) def test_issubclass_base_not_derived(self) -> None: """Base class should not be a subclass of its derived class.""" assert not issubclass(tvm_ffi.testing.TestObjectBase, tvm_ffi.testing.TestObjectDerived) def test_issubclass_array_not_map(self) -> None: """Array should not be a subclass of Map.""" assert not issubclass(tvm_ffi.Array, tvm_ffi.Map) def test_issubclass_map_not_array(self) -> None: """Map should not be a subclass of Array.""" assert not issubclass(tvm_ffi.Map, tvm_ffi.Array) def test_issubclass_all_containers_are_object(self) -> None: """All container types should be subclasses of Object.""" for cls in (tvm_ffi.Array, tvm_ffi.List, tvm_ffi.Map, tvm_ffi.Dict): assert issubclass(cls, tvm_ffi.Object) def test_issubclass_user_type_is_object(self) -> None: """User-defined types should be subclasses of Object.""" assert issubclass(tvm_ffi.testing.TestObjectBase, tvm_ffi.Object) assert issubclass(tvm_ffi.testing.TestObjectDerived, tvm_ffi.Object) def test_issubclass_sibling_containers(self) -> None: """Sibling container types should not be subclasses of each other.""" assert not issubclass(tvm_ffi.List, tvm_ffi.Array) assert not issubclass(tvm_ffi.Dict, tvm_ffi.Map) assert not issubclass(tvm_ffi.Array, tvm_ffi.List) assert not issubclass(tvm_ffi.Map, tvm_ffi.Dict) tvm-ffi-0.1.12/tests/python/test_optional_torch_c_dlpack.py000066400000000000000000000125501521067262500241140ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import builtins import ctypes import subprocess import sys from pathlib import Path from types import SimpleNamespace from typing import Any import pytest try: import torch import torch.version except ImportError: torch = None # ty: ignore[invalid-assignment] import tvm_ffi from tvm_ffi import _optional_torch_c_dlpack IS_WINDOWS = sys.platform.startswith("win") def _fake_torch_module( *, cuda_available: bool, cuda_version: str | None = None, hip_version: str | None = None, include_cuda_attr: bool = True, include_hip_attr: bool = True, ) -> Any: version = SimpleNamespace() if include_cuda_attr: version.cuda = cuda_version if include_hip_attr: version.hip = hip_version return SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: cuda_available), version=version, ) def test_torch_extension_device() -> None: assert ( _optional_torch_c_dlpack._torch_extension_device( _fake_torch_module(cuda_available=False, cuda_version=None, hip_version=None) ) == "cpu" ) assert ( _optional_torch_c_dlpack._torch_extension_device( _fake_torch_module(cuda_available=True, cuda_version="12.8", hip_version=None) ) == "cuda" ) assert ( _optional_torch_c_dlpack._torch_extension_device( _fake_torch_module(cuda_available=True, cuda_version=None, hip_version="7.2") ) == "rocm" ) assert ( _optional_torch_c_dlpack._torch_extension_device( _fake_torch_module( cuda_available=True, include_cuda_attr=False, include_hip_attr=False, ) ) == "cuda" ) def test_existing_torch_dlpack_api_is_preferred_on_rocm(monkeypatch: pytest.MonkeyPatch) -> None: torch_module = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: True), version=SimpleNamespace(cuda=None, hip="7.2"), Tensor=SimpleNamespace(__dlpack_c_exchange_api__=object()), ) original_import = builtins.__import__ def guarded_import(name: str, *args: Any, **kwargs: Any) -> Any: if name == "torch_c_dlpack_ext": raise AssertionError("torch_c_dlpack_ext should not be imported") return original_import(name, *args, **kwargs) monkeypatch.setitem(sys.modules, "torch", torch_module) monkeypatch.setitem(sys.modules, "torch.version", torch_module.version) monkeypatch.setattr(builtins, "__import__", guarded_import) assert _optional_torch_c_dlpack.load_torch_c_dlpack_extension() is None @pytest.mark.skipif(torch is None, reason="torch is not installed") def test_build_torch_c_dlpack_extension() -> None: assert torch is not None build_script = Path(tvm_ffi.__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py" args = [ sys.executable, str(build_script), "--output-dir", "./output-dir", "--libname", "libtorch_c_dlpack_addon_test.so", ] # First use "torch.cuda.is_available()" to check whether GPU environment # is available. Then determine the GPU type. if torch.cuda.is_available(): device = _optional_torch_c_dlpack._torch_extension_device(torch) if device == "cuda": args.append("--build-with-cuda") elif device == "rocm": args.append("--build-with-rocm") else: raise ValueError("Cannot determine whether to build with CUDA or ROCm.") subprocess.run(args, check=True) lib_path = str(Path("./output-dir/libtorch_c_dlpack_addon_test.so").resolve()) assert Path(lib_path).exists() lib = ctypes.CDLL(lib_path) func = lib.TorchDLPackExchangeAPIPtr func.restype = ctypes.c_int64 ptr = func() assert ptr != 0 @pytest.mark.skipif(torch is None, reason="torch is not installed") def test_parallel_build() -> None: build_script = Path(tvm_ffi.__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py" num_processes = 4 output_dir = "./output-dir-parallel" libname = "libtorch_c_dlpack_addon_test.so" processes = [] for i in range(num_processes): p = subprocess.Popen( [sys.executable, str(build_script), "--output-dir", output_dir, "--libname", libname] ) processes.append((p, output_dir)) for p, output_dir in processes: p.wait() assert p.returncode == 0 lib_path = str(Path(f"{output_dir}/{libname}").resolve()) assert Path(lib_path).exists() if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/test_serialization.py000066400000000000000000000501371521067262500221300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for JSON graph serialization/deserialization roundtrips.""" from __future__ import annotations import json from typing import Any, Callable import pytest import tvm_ffi import tvm_ffi.testing from tvm_ffi.serialization import from_json_graph_str, to_json_graph_str # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _roundtrip(obj: Any) -> Any: """Serialize then deserialize and return the result.""" return from_json_graph_str(to_json_graph_str(obj)) def _assert_roundtrip_eq(obj: Any, cmp: Callable[..., Any] | None = None) -> None: """Assert that roundtrip preserves the value.""" result = _roundtrip(obj) if cmp is not None: cmp(result) elif isinstance(obj, float): assert isinstance(result, float) assert result == pytest.approx(obj) elif obj is None: assert result is None else: _assert_any_equal(result, obj) def _assert_any_equal(a: Any, b: Any) -> None: """Recursively compare two tvm_ffi values for equality.""" if isinstance(b, tvm_ffi.Array): assert len(a) == len(b) for x, y in zip(a, b): _assert_any_equal(x, y) elif isinstance(b, (tvm_ffi.Map, tvm_ffi.Dict)): assert len(a) == len(b) for k in b: _assert_any_equal(a[k], b[k]) elif isinstance(b, tvm_ffi.Shape): assert list(a) == list(b) elif isinstance(b, str): # tvm_ffi String inherits from str assert str(a) == str(b) else: assert a == b # --------------------------------------------------------------------------- # Primitive types # --------------------------------------------------------------------------- class TestNone: """Roundtrip tests for None.""" def test_none(self) -> None: """None roundtrips to None.""" assert _roundtrip(None) is None class TestBool: """Roundtrip tests for bool values.""" def test_true(self) -> None: """True roundtrips to True.""" assert _roundtrip(True) is True def test_false(self) -> None: """False roundtrips to False.""" assert _roundtrip(False) is False class TestInt: """Roundtrip tests for integer values.""" def test_zero(self) -> None: """Zero roundtrips correctly.""" result = _roundtrip(0) assert result == 0 def test_positive(self) -> None: """Positive int roundtrips correctly.""" result = _roundtrip(42) assert result == 42 def test_negative(self) -> None: """Negative int roundtrips correctly.""" result = _roundtrip(-1) assert result == -1 def test_large_positive(self) -> None: """Large positive int roundtrips correctly.""" result = _roundtrip(10**15) assert result == 10**15 def test_large_negative(self) -> None: """Large negative int roundtrips correctly.""" result = _roundtrip(-(10**15)) assert result == -(10**15) class TestFloat: """Roundtrip tests for float values.""" def test_zero(self) -> None: """Float zero roundtrips correctly.""" result = _roundtrip(0.0) assert result == 0.0 def test_positive(self) -> None: """Positive float roundtrips correctly.""" result = _roundtrip(3.14159) assert result == pytest.approx(3.14159) def test_negative(self) -> None: """Negative float roundtrips correctly.""" result = _roundtrip(-2.718) assert result == pytest.approx(-2.718) def test_very_small(self) -> None: """Very small float roundtrips correctly.""" result = _roundtrip(1e-300) assert result == pytest.approx(1e-300) def test_very_large(self) -> None: """Very large float roundtrips correctly.""" result = _roundtrip(1e300) assert result == pytest.approx(1e300) # --------------------------------------------------------------------------- # String types # --------------------------------------------------------------------------- class TestString: """Roundtrip tests for ffi.String values.""" def test_empty(self) -> None: """Empty string roundtrips correctly.""" _assert_roundtrip_eq(tvm_ffi.convert("")) def test_short(self) -> None: """Short string roundtrips correctly.""" _assert_roundtrip_eq(tvm_ffi.convert("hello")) def test_long(self) -> None: """Long string roundtrips correctly.""" _assert_roundtrip_eq(tvm_ffi.convert("x" * 1000)) def test_special_chars(self) -> None: """String with special characters roundtrips correctly.""" _assert_roundtrip_eq(tvm_ffi.convert('hello\nworld\t"quotes"')) def test_unicode(self) -> None: """Unicode string roundtrips correctly.""" _assert_roundtrip_eq(tvm_ffi.convert("hello 世界")) # --------------------------------------------------------------------------- # DataType # --------------------------------------------------------------------------- class TestDataType: """Roundtrip tests for DLDataType values.""" def test_int32(self) -> None: """int32 dtype roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.dtype("int32")) result = from_json_graph_str(s) assert str(result) == "int32" def test_float64(self) -> None: """float64 dtype roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.dtype("float64")) result = from_json_graph_str(s) assert str(result) == "float64" def test_bool(self) -> None: """Bool dtype roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.dtype("bool")) result = from_json_graph_str(s) assert str(result) == "bool" def test_vector(self) -> None: """Vector dtype roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.dtype("float32x4")) result = from_json_graph_str(s) assert str(result) == "float32x4" # --------------------------------------------------------------------------- # Device # --------------------------------------------------------------------------- class TestDevice: """Roundtrip tests for DLDevice values.""" def test_cpu(self) -> None: """CPU device roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.Device("cpu", 0)) result = from_json_graph_str(s) assert result.dlpack_device_type() == tvm_ffi.Device("cpu", 0).dlpack_device_type() assert result.index == 0 def test_cuda(self) -> None: """CUDA device roundtrips correctly.""" s = to_json_graph_str(tvm_ffi.Device("cuda", 1)) result = from_json_graph_str(s) assert result.dlpack_device_type() == tvm_ffi.Device("cuda", 1).dlpack_device_type() assert result.index == 1 # --------------------------------------------------------------------------- # Containers # --------------------------------------------------------------------------- class TestArray: """Roundtrip tests for ffi.Array containers.""" def test_empty(self) -> None: """Empty array roundtrips correctly.""" arr = tvm_ffi.convert([]) _assert_roundtrip_eq(arr) def test_single_element(self) -> None: """Single-element array roundtrips correctly.""" arr = tvm_ffi.convert([42]) result = _roundtrip(arr) assert len(result) == 1 assert result[0] == 42 def test_multiple_elements(self) -> None: """Multi-element array roundtrips correctly.""" arr = tvm_ffi.convert([1, 2, 3]) result = _roundtrip(arr) assert len(result) == 3 assert list(result) == [1, 2, 3] def test_mixed_types(self) -> None: """Array with mixed types roundtrips correctly.""" arr = tvm_ffi.convert([42, "hello", True, None]) result = _roundtrip(arr) assert len(result) == 4 assert result[0] == 42 assert result[1] == "hello" assert result[2] is True assert result[3] is None def test_nested_arrays(self) -> None: """Nested arrays roundtrip correctly.""" inner1 = tvm_ffi.convert([1, 2]) inner2 = tvm_ffi.convert([3]) outer = tvm_ffi.convert([inner1, inner2]) result = _roundtrip(outer) assert len(result) == 2 assert list(result[0]) == [1, 2] assert list(result[1]) == [3] def test_duplicated_elements(self) -> None: """Array with duplicated elements roundtrips correctly.""" arr = tvm_ffi.convert([42, 42, 42]) result = _roundtrip(arr) assert len(result) == 3 assert all(x == 42 for x in result) class TestMap: """Roundtrip tests for ffi.Map containers.""" def test_empty(self) -> None: """Empty map roundtrips correctly.""" m = tvm_ffi.convert({}) _assert_roundtrip_eq(m) def test_single_entry(self) -> None: """Single-entry map roundtrips correctly.""" m = tvm_ffi.convert({"key": 42}) result = _roundtrip(m) assert len(result) == 1 assert result["key"] == 42 def test_multiple_entries(self) -> None: """Multi-entry map roundtrips correctly.""" m = tvm_ffi.convert({"a": 1, "b": 2, "c": 3}) result = _roundtrip(m) assert len(result) == 3 assert result["a"] == 1 assert result["b"] == 2 assert result["c"] == 3 def test_mixed_value_types(self) -> None: """Map with mixed value types roundtrips correctly.""" m = tvm_ffi.convert({"int": 42, "str": "hello", "bool": True, "none": None}) result = _roundtrip(m) assert result["int"] == 42 assert result["str"] == "hello" assert result["bool"] is True assert result["none"] is None def test_nested_map(self) -> None: """Nested maps roundtrip correctly.""" inner = tvm_ffi.convert({"x": 1}) outer = tvm_ffi.convert({"inner": inner}) result = _roundtrip(outer) assert result["inner"]["x"] == 1 def test_map_with_array_value(self) -> None: """Map with array values roundtrips correctly.""" arr = tvm_ffi.convert([10, 20]) m = tvm_ffi.convert({"nums": arr}) result = _roundtrip(m) assert list(result["nums"]) == [10, 20] class TestDict: """Roundtrip tests for ffi.Dict containers.""" def test_empty(self) -> None: """Empty dict roundtrips correctly.""" d = tvm_ffi.Dict({}) _assert_roundtrip_eq(d) def test_single_entry(self) -> None: """Single-entry dict roundtrips correctly.""" d = tvm_ffi.Dict({"key": 42}) result = _roundtrip(d) assert len(result) == 1 assert result["key"] == 42 def test_multiple_entries(self) -> None: """Multi-entry dict roundtrips correctly.""" d = tvm_ffi.Dict({"a": 1, "b": 2, "c": 3}) result = _roundtrip(d) assert len(result) == 3 assert result["a"] == 1 assert result["b"] == 2 assert result["c"] == 3 def test_mixed_value_types(self) -> None: """Dict with mixed value types roundtrips correctly.""" d = tvm_ffi.Dict({"int": 42, "str": "hello", "bool": True, "none": None}) result = _roundtrip(d) assert result["int"] == 42 assert result["str"] == "hello" assert result["bool"] is True assert result["none"] is None def test_nested_dict(self) -> None: """Nested dicts roundtrip correctly.""" inner = tvm_ffi.Dict({"x": 1}) outer = tvm_ffi.Dict({"inner": inner}) result = _roundtrip(outer) assert result["inner"]["x"] == 1 def test_dict_with_array_value(self) -> None: """Dict with array values roundtrips correctly.""" arr = tvm_ffi.convert([10, 20]) d = tvm_ffi.Dict({"nums": arr}) result = _roundtrip(d) assert list(result["nums"]) == [10, 20] class TestShape: """Roundtrip tests for ffi.Shape containers.""" def test_empty(self) -> None: """Empty shape roundtrips correctly.""" shape = tvm_ffi.Shape(()) _assert_roundtrip_eq(shape) def test_1d(self) -> None: """1D shape roundtrips correctly.""" shape = tvm_ffi.Shape((10,)) result = _roundtrip(shape) assert list(result) == [10] def test_nd(self) -> None: """N-D shape roundtrips correctly.""" shape = tvm_ffi.Shape((1, 2, 3, 4)) result = _roundtrip(shape) assert list(result) == [1, 2, 3, 4] # --------------------------------------------------------------------------- # Objects with reflection # --------------------------------------------------------------------------- class TestObjectSerialization: """Roundtrip tests for objects with reflection metadata.""" def test_int_pair_roundtrip(self) -> None: """TestIntPair has refl::init and POD int64 fields.""" pair = tvm_ffi.testing.TestIntPair(3, 7) s = to_json_graph_str(pair) result = from_json_graph_str(s) assert result.a == 3 assert result.b == 7 def test_int_pair_zero_values(self) -> None: """TestIntPair with zero values roundtrips correctly.""" pair = tvm_ffi.testing.TestIntPair(0, 0) result = _roundtrip(pair) assert result.a == 0 assert result.b == 0 def test_int_pair_negative_values(self) -> None: """TestIntPair with negative values roundtrips correctly.""" pair = tvm_ffi.testing.TestIntPair(-100, -200) result = _roundtrip(pair) assert result.a == -100 assert result.b == -200 def test_int_pair_large_values(self) -> None: """TestIntPair with large values roundtrips correctly.""" pair = tvm_ffi.testing.TestIntPair(10**15, -(10**15)) result = _roundtrip(pair) assert result.a == 10**15 assert result.b == -(10**15) # --------------------------------------------------------------------------- # JSON structure verification # --------------------------------------------------------------------------- class TestJSONStructure: """Tests verifying the internal JSON graph structure.""" def test_null_json_structure(self) -> None: """None produces a single node with type 'None'.""" s = to_json_graph_str(None) parsed = json.loads(s) assert parsed["root_index"] == 0 assert len(parsed["nodes"]) == 1 assert parsed["nodes"][0]["type"] == "None" def test_bool_json_structure(self) -> None: """Bool produces a node with type 'bool'.""" s = to_json_graph_str(True) parsed = json.loads(s) assert parsed["nodes"][0]["type"] == "bool" assert parsed["nodes"][0]["data"] is True def test_int_json_structure(self) -> None: """Int produces a node with type 'int'.""" s = to_json_graph_str(42) parsed = json.loads(s) assert parsed["nodes"][0]["type"] == "int" assert parsed["nodes"][0]["data"] == 42 def test_float_json_structure(self) -> None: """Float produces a node with type 'float'.""" s = to_json_graph_str(3.14) parsed = json.loads(s) assert parsed["nodes"][0]["type"] == "float" assert parsed["nodes"][0]["data"] == pytest.approx(3.14) def test_string_json_structure(self) -> None: """String produces a node with type 'ffi.String'.""" s = to_json_graph_str(tvm_ffi.convert("hello")) parsed = json.loads(s) assert parsed["nodes"][parsed["root_index"]]["type"] == "ffi.String" assert parsed["nodes"][parsed["root_index"]]["data"] == "hello" def test_array_json_structure(self) -> None: """Array data contains node index references.""" s = to_json_graph_str(tvm_ffi.convert([1, 2])) parsed = json.loads(s) root = parsed["nodes"][parsed["root_index"]] assert root["type"] == "ffi.Array" # data should be list of node indices assert isinstance(root["data"], list) assert len(root["data"]) == 2 def test_map_json_structure(self) -> None: """Map data contains flattened key-value node index pairs.""" s = to_json_graph_str(tvm_ffi.convert({"a": 1})) parsed = json.loads(s) root = parsed["nodes"][parsed["root_index"]] assert root["type"] == "ffi.Map" assert isinstance(root["data"], list) # key-value pairs flattened: [key_idx, val_idx] assert len(root["data"]) == 2 def test_dict_json_structure(self) -> None: """Dict data contains flattened key-value node index pairs.""" s = to_json_graph_str(tvm_ffi.Dict({"a": 1})) parsed = json.loads(s) root = parsed["nodes"][parsed["root_index"]] assert root["type"] == "ffi.Dict" assert isinstance(root["data"], list) # key-value pairs flattened: [key_idx, val_idx] assert len(root["data"]) == 2 def test_node_dedup(self) -> None: """Duplicate values should share the same node index.""" s = to_json_graph_str(tvm_ffi.convert([42, 42, 42])) parsed = json.loads(s) root = parsed["nodes"][parsed["root_index"]] # all three elements should reference the same node assert root["data"][0] == root["data"][1] == root["data"][2] def test_object_pod_fields_are_inlined(self) -> None: """POD fields (int, bool, float) are inlined directly via field_static_type_index.""" pair = tvm_ffi.testing.TestIntPair(3, 7) s = to_json_graph_str(pair) parsed = json.loads(s) root = parsed["nodes"][parsed["root_index"]] assert root["type"] == "testing.TestIntPair" # fields 'a' and 'b' should be inlined as direct int values (not node indices) assert root["data"]["a"] == 3 assert root["data"]["b"] == 7 # --------------------------------------------------------------------------- # Metadata # --------------------------------------------------------------------------- class TestMetadata: """Tests for optional metadata in serialized output.""" def test_with_metadata(self) -> None: """Metadata dict appears in serialized JSON when provided.""" s = to_json_graph_str(42, {"version": "1.0"}) parsed = json.loads(s) assert "metadata" in parsed assert parsed["metadata"]["version"] == "1.0" def test_without_metadata(self) -> None: """Metadata key is absent when not provided.""" s = to_json_graph_str(42) parsed = json.loads(s) assert "metadata" not in parsed # --------------------------------------------------------------------------- # Error cases # --------------------------------------------------------------------------- class TestErrors: """Tests for error handling in deserialization.""" def test_invalid_json_string(self) -> None: """Invalid JSON string raises an error.""" with pytest.raises(Exception): from_json_graph_str("not valid json") def test_empty_json_string(self) -> None: """Empty string raises an error.""" with pytest.raises(Exception): from_json_graph_str("") def test_missing_root_index(self) -> None: """JSON without root_index raises an error.""" with pytest.raises(Exception): from_json_graph_str('{"nodes": []}') def test_missing_nodes(self) -> None: """JSON without nodes raises an error.""" with pytest.raises(Exception): from_json_graph_str('{"root_index": 0}') tvm-ffi-0.1.12/tests/python/test_stl.py000066400000000000000000000035321521067262500200520ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import pathlib import pytest import tvm_ffi.cpp from tvm_ffi.module import Module def test_stl() -> None: cpp_path = pathlib.Path(__file__).parent.resolve() / "cpp_src" / "test_stl.cc" output_lib_path = tvm_ffi.cpp.build( name="test_stl", sources=[str(cpp_path)], ) mod: Module = tvm_ffi.load_module(output_lib_path) assert list(mod.test_tuple([1, 2.5])) == [2.5, 1] assert mod.test_vector(None) is None assert list(mod.test_vector([[1, 2], [3, 4]])) == [3, 7] assert mod.test_variant(1) == "int" assert mod.test_variant(1.0) == "float" assert list(mod.test_variant([1, 1.0])) == ["int", "float"] assert dict(mod.test_map({"a": 1, "b": 2})) == {1: "a", 2: "b"} assert dict(mod.test_map_2({"a": 1, "b": 2})) == {1: "a", 2: "b"} assert mod.test_function(lambda: 0)() == 1 assert mod.test_function(lambda: 10)() == 11 with pytest.raises(TypeError): mod.test_tuple([1.5, 2.5]) with pytest.raises(TypeError): mod.test_function(lambda: 0)(100) if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/test_stream.py000066400000000000000000000122511521067262500205410ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import ctypes import pytest import tvm_ffi import tvm_ffi.cpp try: import torch except ImportError: torch = None # ty: ignore[invalid-assignment] try: from cuda.bindings import driver as cuda_driver except ImportError: cuda_driver = None @pytest.mark.skipif(cuda_driver is None, reason="Requires cuda-python") def test_cuda_driver_stream() -> None: assert cuda_driver is not None echo = tvm_ffi.get_global_func("testing.echo") stream = cuda_driver.CUstream(0) y = echo(stream) assert y is not None z = echo(cuda_driver.CUstream(1)) assert isinstance(z, ctypes.c_void_p) assert z.value == 1 def gen_check_stream_mod() -> tvm_ffi.Module: return tvm_ffi.cpp.load_inline( name="check_stream", cpp_sources=""" void check_stream(int device_type, int device_id, uint64_t stream) { uint64_t cur_stream = reinterpret_cast(TVMFFIEnvGetStream(device_type, device_id)); TVM_FFI_ICHECK_EQ(cur_stream, stream); } """, functions=["check_stream"], ) def test_raw_stream() -> None: mod = gen_check_stream_mod() device = tvm_ffi.device("cuda:0") stream_1 = 123456789 stream_2 = 987654321 with tvm_ffi.use_raw_stream(device, stream_1): mod.check_stream(device.dlpack_device_type(), device.index, stream_1) assert tvm_ffi.get_raw_stream(device) == stream_1 with tvm_ffi.use_raw_stream(device, stream_2): mod.check_stream(device.dlpack_device_type(), device.index, stream_2) assert tvm_ffi.get_raw_stream(device) == stream_2 mod.check_stream(device.dlpack_device_type(), device.index, stream_1) assert tvm_ffi.get_raw_stream(device) == stream_1 @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_torch_stream() -> None: assert torch is not None mod = gen_check_stream_mod() device_id = torch.cuda.current_device() device = tvm_ffi.device("cuda", device_id) device_type = device.dlpack_device_type() stream_1 = torch.cuda.Stream(device_id) stream_2 = torch.cuda.Stream(device_id) with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_1)): assert torch.cuda.current_stream() == stream_1 mod.check_stream(device_type, device_id, stream_1.cuda_stream) with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_2)): assert torch.cuda.current_stream() == stream_2 mod.check_stream(device_type, device_id, stream_2.cuda_stream) assert torch.cuda.current_stream() == stream_1 mod.check_stream(device_type, device_id, stream_1.cuda_stream) @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_torch_current_stream() -> None: assert torch is not None mod = gen_check_stream_mod() device_id = torch.cuda.current_device() device = tvm_ffi.device("cuda", device_id) device_type = device.dlpack_device_type() stream_1 = torch.cuda.Stream(device_id) stream_2 = torch.cuda.Stream(device_id) with torch.cuda.stream(stream_1): assert torch.cuda.current_stream() == stream_1 with tvm_ffi.use_torch_stream(): mod.check_stream(device_type, device_id, stream_1.cuda_stream) with torch.cuda.stream(stream_2): assert torch.cuda.current_stream() == stream_2 with tvm_ffi.use_torch_stream(): mod.check_stream(device_type, device_id, stream_2.cuda_stream) assert torch.cuda.current_stream() == stream_1 with tvm_ffi.use_torch_stream(): mod.check_stream(device_type, device_id, stream_1.cuda_stream) @pytest.mark.skipif( torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" ) def test_torch_graph() -> None: assert torch is not None mod = gen_check_stream_mod() device_id = torch.cuda.current_device() device = tvm_ffi.device("cuda", device_id) device_type = device.dlpack_device_type() graph = torch.cuda.CUDAGraph() stream = torch.cuda.Stream(device_id) x = torch.zeros(1, device="cuda") with tvm_ffi.use_torch_stream(torch.cuda.graph(graph, stream=stream)): assert torch.cuda.current_stream() == stream mod.check_stream(device_type, device_id, stream.cuda_stream) # avoid cuda graph no capture warning x = x + 1 tvm-ffi-0.1.12/tests/python/test_string.py000066400000000000000000000042301521067262500205520ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import pickle import tvm_ffi def test_string() -> None: fecho = tvm_ffi.get_global_func("testing.echo") s = tvm_ffi.core.String("hello") s2 = fecho(s) assert s2 == "hello" s3 = tvm_ffi.convert("hello") assert isinstance(s3, str) x = "hello long string" assert fecho(x) == x s4 = pickle.loads(pickle.dumps(s)) assert s4 == "hello" def test_bytes() -> None: fecho = tvm_ffi.get_global_func("testing.echo") b = tvm_ffi.core.Bytes(b"hello") assert isinstance(b, tvm_ffi.core.Bytes) b2 = fecho(b) assert b2 == b"hello" b3 = tvm_ffi.convert(b"hello") assert isinstance(b3, tvm_ffi.core.Bytes) assert isinstance(b3, bytes) b4 = tvm_ffi.convert(bytearray(b"hello")) assert isinstance(b4, tvm_ffi.core.Bytes) assert isinstance(b4, bytes) b5 = pickle.loads(pickle.dumps(b)) assert b5 == b"hello" assert isinstance(b5, tvm_ffi.core.Bytes) def test_string_find_substr() -> None: s = tvm_ffi.core.String("hello world") assert s.find("world") == 6 assert s.find("hello") == 0 assert s.find("o") == 4 assert s.find("o", 5) == 7 assert s.find("notfound") == -1 assert s.find("") == 0 assert s.find("", 5) == 5 assert s.find("", 11) == 11 assert s.find("", 20) == -1 assert s[6:11] == "world" assert s[0:5] == "hello" assert s[6:] == "world" assert s[:5] == "hello" tvm-ffi-0.1.12/tests/python/test_structural.py000066400000000000000000000131251521067262500214570ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import numpy as np import tvm_ffi import tvm_ffi.testing _recursive_eq = tvm_ffi.get_global_func("ffi.RecursiveEq") def test_structural_key_basic() -> None: k1 = tvm_ffi.StructuralKey({"a": [1, 2], "b": [3, {"c": 4}]}) k2 = tvm_ffi.StructuralKey({"b": [3, {"c": 4}], "a": [1, 2]}) k3 = tvm_ffi.StructuralKey({"a": [1, 2], "b": [3, {"c": 5}]}) assert tvm_ffi.structural_hash(k1.key) == k1.__hash__() assert tvm_ffi.structural_hash(k2.key) == k2.__hash__() assert k1 == k2 assert k1 != k3 assert hash(k1) == hash(k2) assert tvm_ffi.structural_equal(k1.key, k2.key) assert not tvm_ffi.structural_equal(k1.key, k3.key) def test_structural_helpers() -> None: lhs = {"items": [1, 2, {"k": 3}], "meta": {"tag": "x"}} rhs = {"meta": {"tag": "x"}, "items": [1, 2, {"k": 3}]} other = {"items": [1, 2, {"k": 4}], "meta": {"tag": "x"}} assert tvm_ffi.structural_equal(lhs, rhs) assert not tvm_ffi.structural_equal(lhs, other) assert tvm_ffi.structural_hash(lhs) == tvm_ffi.structural_hash(rhs) assert tvm_ffi.structural_hash(lhs) != tvm_ffi.structural_hash(other) assert tvm_ffi.get_first_structural_mismatch(lhs, rhs) is None assert tvm_ffi.get_first_structural_mismatch(lhs, other) is not None def test_structural_key_in_map() -> None: k1 = tvm_ffi.StructuralKey({"x": [1, 2], "y": [3]}) k2 = tvm_ffi.StructuralKey({"y": [3], "x": [1, 2]}) k3 = tvm_ffi.StructuralKey({"x": [1, 2], "y": [5]}) m = tvm_ffi.Map({k1: 1, k2: 2, k3: 3}) assert len(m) == 2 assert m[k1] == 2 assert m[k2] == 2 assert m[k3] == 3 def test_structural_equal_dict() -> None: d1 = tvm_ffi.Dict({"a": 1, "b": 2, "c": 3}) d2 = tvm_ffi.Dict({"c": 3, "b": 2, "a": 1}) d3 = tvm_ffi.Dict({"a": 1, "b": 2, "c": 4}) assert tvm_ffi.structural_equal(d1, d2) assert tvm_ffi.structural_hash(d1) == tvm_ffi.structural_hash(d2) assert not tvm_ffi.structural_equal(d1, d3) assert tvm_ffi.structural_hash(d1) != tvm_ffi.structural_hash(d3) assert tvm_ffi.get_first_structural_mismatch(d1, d2) is None assert tvm_ffi.get_first_structural_mismatch(d1, d3) is not None def test_structural_dict_vs_map_different_type() -> None: m = tvm_ffi.Map({"a": 1, "b": 2}) d = tvm_ffi.Dict({"a": 1, "b": 2}) # Different type_index => not structurally equal assert not tvm_ffi.structural_equal(m, d) assert tvm_ffi.structural_hash(m) != tvm_ffi.structural_hash(d) def test_structural_key_in_python_dict() -> None: k1 = tvm_ffi.StructuralKey({"name": ["a", "b"], "ver": [1]}) k2 = tvm_ffi.StructuralKey({"ver": [1], "name": ["a", "b"]}) k3 = tvm_ffi.StructuralKey({"name": ["a", "c"], "ver": [1]}) data = {k1: "a", k3: "b"} assert data[k2] == "a" assert data[k3] == "b" def test_structural_key_tensor_content_policy() -> None: t1_np = np.array([1.0, 2.0, 3.0], dtype="float32") t2_np = np.array([1.0, 2.0, 4.0], dtype="float32") if not hasattr(t1_np, "__dlpack__"): return t1 = tvm_ffi.from_dlpack(t1_np) t2 = tvm_ffi.from_dlpack(t2_np) # Default policy compares tensor content. assert not tvm_ffi.structural_equal(t1, t2) # Optional policy can ignore tensor content. assert tvm_ffi.structural_equal(t1, t2, skip_tensor_content=True) # StructuralKey should follow default structural policy. k1 = tvm_ffi.StructuralKey(t1) k2 = tvm_ffi.StructuralKey(t2) assert k1 != k2 data = {k1: "a", k2: "b"} assert len(data) == 2 # ---------- RecursiveEq cycle tests ---------- def test_recursive_eq_self_referencing_cycle() -> None: """RecursiveEq should return True for structurally equivalent cycles.""" v_map = tvm_ffi.Map({}) obj = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=1, v_f64=0.0, v_str="", v_map=v_map, v_array=tvm_ffi.Array([]), ) obj.v_array = tvm_ffi.Array([obj]) # type: ignore[unresolved-attribute] # Self-referencing object compared to itself — identity short-circuits. assert _recursive_eq(obj, obj) def test_recursive_eq_mutual_cycle() -> None: """RecursiveEq should return True for two distinct but structurally equivalent cyclic graphs.""" v_map = tvm_ffi.Map({}) def make_cyclic(v_i64: int) -> object: o = tvm_ffi.testing.create_object( "testing.TestObjectDerived", v_i64=v_i64, v_f64=0.0, v_str="x", v_map=v_map, v_array=tvm_ffi.Array([]), ) o.v_array = tvm_ffi.Array([o]) # type: ignore[unresolved-attribute] return o a = make_cyclic(42) b = make_cyclic(42) # Two distinct objects with identical structure and self-referencing cycles. assert _recursive_eq(a, b) # Different content should not be equal. c = make_cyclic(99) assert not _recursive_eq(a, c) tvm-ffi-0.1.12/tests/python/test_structural_py_class.py000066400000000000000000000324671521067262500233660ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for structural equality/hashing on py_class-defined types. Mirrors the C++ tests in tests/cpp/extra/test_structural_equal_hash.cc, porting the object-level tests (FreeVar, FuncDefAndIgnoreField, etc.) to Python using ``@py_class(structural_eq=...)`` and ``field(structural_eq=...)``. """ from __future__ import annotations import pytest import tvm_ffi from tvm_ffi import get_first_structural_mismatch, structural_equal, structural_hash from tvm_ffi.dataclasses import field, py_class # --------------------------------------------------------------------------- # Type definitions (mirror testing_object.h) # --------------------------------------------------------------------------- @py_class("testing.py.Var", structural_eq="var") class TVar(tvm_ffi.Object): """Variable node — compared by binding position, not by name. Mirrors C++ TVarObj with: _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar name field has SEqHashIgnore """ name: str = field(structural_eq="ignore") @py_class("testing.py.Int", structural_eq="tree") class TInt(tvm_ffi.Object): """Simple integer literal node.""" value: int @py_class("testing.py.Func", structural_eq="tree") class TFunc(tvm_ffi.Object): """Function node with definition region and ignored comment. Mirrors C++ TFuncObj with: params has SEqHashDef comment has SEqHashIgnore """ params: list = field(structural_eq="def") body: list comment: str = field(structural_eq="ignore", default="") @py_class("testing.py.Expr", structural_eq="tree") class TExpr(tvm_ffi.Object): """A simple expression node for tree-comparison tests.""" value: int @py_class("testing.py.Metadata", structural_eq="const-tree") class TMetadata(tvm_ffi.Object): """Immutable metadata node — pointer shortcut is safe (no var children).""" tag: str version: int @py_class("testing.py.Binding", structural_eq="dag") class TBinding(tvm_ffi.Object): """Binding node — sharing structure is semantically meaningful.""" name: str value: int # --------------------------------------------------------------------------- # Tests: FreeVar (mirrors C++ FreeVar test) # --------------------------------------------------------------------------- class TestFreeVar: """Test structural_eq="var" kind (C++ kTVMFFISEqHashKindFreeVar).""" def test_free_var_equal_with_mapping(self) -> None: """Two different vars are equal when map_free_vars=True.""" a = TVar("a") b = TVar("b") assert structural_equal(a, b, map_free_vars=True) def test_free_var_not_equal_without_mapping(self) -> None: """Two different vars are NOT equal by default (no mapping).""" a = TVar("a") b = TVar("b") assert not structural_equal(a, b) def test_free_var_hash_differs_without_mapping(self) -> None: """Without mapping, different vars produce different hashes.""" a = TVar("a") b = TVar("b") assert structural_hash(a) != structural_hash(b) def test_free_var_hash_equal_with_mapping(self) -> None: """With map_free_vars, positional hashing makes them equal.""" a = TVar("a") b = TVar("b") assert structural_hash(a, map_free_vars=True) == structural_hash(b, map_free_vars=True) def test_free_var_same_pointer(self) -> None: """Same variable is always equal to itself.""" x = TVar("x") assert structural_equal(x, x) def test_free_var_name_ignored(self) -> None: """The name field is structural_eq="ignore", so it doesn't affect comparison.""" a = TVar("different_name_a") b = TVar("different_name_b") # Names differ, but with mapping they are still equal assert structural_equal(a, b, map_free_vars=True) # --------------------------------------------------------------------------- # Tests: FuncDefAndIgnoreField (mirrors C++ FuncDefAndIgnoreField test) # --------------------------------------------------------------------------- class TestFuncDefAndIgnore: """Test structural_eq="def" and structural_eq="ignore" on fields.""" def test_alpha_equivalent_functions(self) -> None: """fun(x){1, x} with comment_a == fun(y){1, y} with comment_b.""" x = TVar("x") y = TVar("y") fa = TFunc(params=[x], body=[TInt(1), x], comment="comment a") fb = TFunc(params=[y], body=[TInt(1), y], comment="comment b") assert structural_equal(fa, fb) assert structural_hash(fa) == structural_hash(fb) def test_different_body(self) -> None: """fun(x){1, x} != fun(x){1, 2} — body differs at index 1.""" x = TVar("x") fa = TFunc(params=[x], body=[TInt(1), x], comment="comment a") fc = TFunc(params=[x], body=[TInt(1), TInt(2)], comment="comment c") assert not structural_equal(fa, fc) def test_mismatch_path(self) -> None: """GetFirstMismatch reports the correct access path.""" x = TVar("x") fa = TFunc(params=[x], body=[TInt(1), x]) fc = TFunc(params=[x], body=[TInt(1), TInt(2)]) mismatch = get_first_structural_mismatch(fa, fc) assert mismatch is not None def test_comment_ignored(self) -> None: """Identical structure with different comments are equal.""" x = TVar("x") f1 = TFunc(params=[x], body=[TInt(1)], comment="first") f2 = TFunc(params=[x], body=[TInt(1)], comment="second") assert structural_equal(f1, f2) assert structural_hash(f1) == structural_hash(f2) def test_inconsistent_var_usage(self) -> None: """fun(x,y){x+x} != fun(a,b){a+b} — inconsistent variable mapping.""" x, y = TVar("x"), TVar("y") a, b = TVar("a"), TVar("b") f1 = TFunc(params=[x, y], body=[x, x]) f2 = TFunc(params=[a, b], body=[a, b]) assert not structural_equal(f1, f2) def test_multi_param_alpha_equiv(self) -> None: """fun(x,y){x, y} == fun(a,b){a, b} — consistent variable renaming.""" x, y = TVar("x"), TVar("y") a, b = TVar("a"), TVar("b") f1 = TFunc(params=[x, y], body=[x, y]) f2 = TFunc(params=[a, b], body=[a, b]) assert structural_equal(f1, f2) assert structural_hash(f1) == structural_hash(f2) def test_swapped_params_not_equal(self) -> None: """fun(x,y){x, y} != fun(a,b){b, a} — reversed usage.""" x, y = TVar("x"), TVar("y") a, b = TVar("a"), TVar("b") f1 = TFunc(params=[x, y], body=[x, y]) f2 = TFunc(params=[a, b], body=[b, a]) assert not structural_equal(f1, f2) def test_nested_functions(self) -> None: """Nested function with inner binding — alpha-equivalence is scoped.""" x = TVar("x") y = TVar("y") inner_x = TFunc(params=[x], body=[x]) inner_y = TFunc(params=[y], body=[y]) outer_a = TFunc(params=[], body=[inner_x]) outer_b = TFunc(params=[], body=[inner_y]) assert structural_equal(outer_a, outer_b) assert structural_hash(outer_a) == structural_hash(outer_b) # --------------------------------------------------------------------------- # Tests: tree kind basics # --------------------------------------------------------------------------- class TestTreeNode: """Test structural_eq="tree" kind.""" def test_equal_content(self) -> None: """Two tree nodes with identical content are structurally equal.""" a = TExpr(value=42) b = TExpr(value=42) assert structural_equal(a, b) assert structural_hash(a) == structural_hash(b) def test_different_content(self) -> None: """Two tree nodes with different content are not equal.""" a = TExpr(value=1) b = TExpr(value=2) assert not structural_equal(a, b) assert structural_hash(a) != structural_hash(b) def test_sharing_invisible(self) -> None: """Under "tree", sharing doesn't affect equality.""" s = TExpr(value=10) # Two arrays referencing the same object vs two copies shared = tvm_ffi.Array([s, s]) copies = tvm_ffi.Array([TExpr(value=10), TExpr(value=10)]) assert structural_equal(shared, copies) assert structural_hash(shared) == structural_hash(copies) # --------------------------------------------------------------------------- # Tests: const-tree kind # --------------------------------------------------------------------------- class TestConstTreeNode: """Test structural_eq="const-tree" kind.""" def test_equal_content(self) -> None: """Two const-tree nodes with identical content are structurally equal.""" a = TMetadata(tag="v1", version=1) b = TMetadata(tag="v1", version=1) assert structural_equal(a, b) assert structural_hash(a) == structural_hash(b) def test_different_content(self) -> None: """Two const-tree nodes with different content are not equal.""" a = TMetadata(tag="v1", version=1) b = TMetadata(tag="v1", version=2) assert not structural_equal(a, b) def test_same_pointer_shortcircuits(self) -> None: """Same pointer should be equal (the const-tree optimization).""" a = TMetadata(tag="test", version=1) assert structural_equal(a, a) # --------------------------------------------------------------------------- # Tests: dag kind # --------------------------------------------------------------------------- class TestDAGNode: """Test structural_eq="dag" kind.""" def test_same_dag_shape(self) -> None: """Two DAGs with the same sharing shape are equal.""" s1 = TBinding(name="s", value=1) s2 = TBinding(name="s", value=1) dag1 = tvm_ffi.Array([s1, s1]) # shared dag2 = tvm_ffi.Array([s2, s2]) # shared (same shape) assert structural_equal(dag1, dag2) def test_dag_vs_tree_not_equal(self) -> None: """A DAG (shared) vs tree (independent copies) are NOT equal.""" shared = TBinding(name="s", value=1) copy1 = TBinding(name="s", value=1) copy2 = TBinding(name="s", value=1) dag = tvm_ffi.Array([shared, shared]) tree = tvm_ffi.Array([copy1, copy2]) assert not structural_equal(dag, tree) def test_dag_vs_tree_hash_differs(self) -> None: """DAG and tree with same content should hash differently.""" shared = TBinding(name="s", value=1) copy1 = TBinding(name="s", value=1) copy2 = TBinding(name="s", value=1) dag = tvm_ffi.Array([shared, shared]) tree = tvm_ffi.Array([copy1, copy2]) assert structural_hash(dag) != structural_hash(tree) def test_reverse_bijection(self) -> None: """(a, b) vs (a, a) where a ≅ b — reverse map detects inconsistency.""" a = TBinding(name="a", value=1) b = TBinding(name="b", value=1) # same content as a lhs = tvm_ffi.Array([a, b]) rhs = tvm_ffi.Array([a, a]) # note: same object twice assert not structural_equal(lhs, rhs) # --------------------------------------------------------------------------- # Tests: unsupported kind (default) # --------------------------------------------------------------------------- class TestUnsupported: """Test that types without structural_eq= raise on structural comparison.""" def test_unsupported_raises_on_hash(self) -> None: """Structural hashing raises TypeError for types without structural_eq=.""" @py_class("testing.py.Plain") class Plain(tvm_ffi.Object): x: int with pytest.raises(TypeError): structural_hash(Plain(x=1)) def test_unsupported_raises_on_equal(self) -> None: """Structural equality raises TypeError for types without structural_eq=.""" @py_class("testing.py.Plain2") class Plain2(tvm_ffi.Object): x: int with pytest.raises(TypeError): structural_equal(Plain2(x=1), Plain2(x=1)) # --------------------------------------------------------------------------- # Tests: validation # --------------------------------------------------------------------------- class TestValidation: """Test that invalid structural_eq= values are rejected.""" def test_invalid_type_structure(self) -> None: """Invalid type-level structural_eq= value raises ValueError.""" with pytest.raises(ValueError, match="structural_eq"): py_class(structural_eq="invalid") def test_invalid_field_structure(self) -> None: """Invalid field-level structural_eq= value raises ValueError.""" with pytest.raises(ValueError, match="structural_eq"): field(structural_eq="bad_value") tvm-ffi-0.1.12/tests/python/test_stubgen.py000066400000000000000000000572531521067262500207300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations from pathlib import Path import pytest import tvm_ffi.stub.cli as stub_cli from tvm_ffi.core import TypeSchema from tvm_ffi.stub import consts as C from tvm_ffi.stub.cli import _stage_2, _stage_3 from tvm_ffi.stub.codegen import ( generate_all, generate_export, generate_ffi_api, generate_global_funcs, generate_import_section, generate_init, generate_object, ) from tvm_ffi.stub.file_utils import CodeBlock, FileInfo from tvm_ffi.stub.utils import ( FuncInfo, ImportItem, InitConfig, NamedTypeSchema, ObjectInfo, Options, ) def _identity_ty_map(name: str) -> str: return name def _default_ty_map() -> dict[str, str]: return C.TY_MAP_DEFAULTS.copy() def _type_suffix(name: str) -> str: return C.TY_MAP_DEFAULTS.get(name, name).rsplit(".", 1)[-1] def test_codeblock_from_begin_line_variants() -> None: cases = [ (f"{C.STUB_BEGIN} global/demo", "global", ("demo", "")), (f"{C.STUB_BEGIN} global/demo@.registry", "global", ("demo", ".registry")), (f"{C.STUB_BEGIN} object/demo.TypeBase", "object", "demo.TypeBase"), (f"{C.STUB_BEGIN} ty-map/custom", "ty-map", "custom"), (f"{C.STUB_BEGIN} import-section", "import-section", ""), ] for lineno, (line, kind, param) in enumerate(cases, start=1): block = CodeBlock.from_begin_line(lineno, line) assert block.kind == kind assert block.param == param assert block.lineno_start == lineno assert block.lineno_end is None assert block.lines == [] def test_codeblock_from_begin_line_ty_map_and_unknown() -> None: line = f"{C.STUB_TY_MAP} custom -> mapped" block = CodeBlock.from_begin_line(5, line) assert block.kind == "ty-map" assert block.param == "custom -> mapped" assert block.lineno_start == 5 assert block.lineno_end == 5 with pytest.raises(ValueError): CodeBlock.from_begin_line(1, f"{C.STUB_BEGIN} unsupported/kind") def test_fileinfo_from_file_skip_and_missing_markers(tmp_path: Path) -> None: skip = tmp_path / "skip.py" skip.write_text(f"print('hi')\n{C.STUB_SKIP_FILE}\n", encoding="utf-8") assert FileInfo.from_file(skip) is None plain = tmp_path / "plain.py" plain.write_text("print('plain')\n", encoding="utf-8") assert FileInfo.from_file(plain) is None def test_fileinfo_from_file_parses_blocks(tmp_path: Path) -> None: content = "\n".join( [ "first = 1", f"{C.STUB_BEGIN} global/demo.func", "in_stub = True", C.STUB_END, f"{C.STUB_TY_MAP} x -> y", ] ) path = tmp_path / "demo.py" path.write_text(content, encoding="utf-8") info = FileInfo.from_file(path) assert info is not None assert info.path == path.resolve() assert len(info.code_blocks) == 3 first, stub, ty_map = info.code_blocks assert first.kind is None and first.lines == ["first = 1"] assert stub.kind == "global" assert stub.param == ("demo.func", "") assert stub.lineno_start == 2 assert stub.lineno_end == 4 assert stub.lines == [ f"{C.STUB_BEGIN} global/demo.func", "in_stub = True", C.STUB_END, ] assert ty_map.kind == "ty-map" assert ty_map.param == "x -> y" assert ty_map.lineno_start == ty_map.lineno_end == 5 assert ty_map.lines == [f"{C.STUB_TY_MAP} x -> y"] def test_fileinfo_from_file_error_paths(tmp_path: Path) -> None: nested = tmp_path / "nested.py" nested.write_text( "\n".join( [ f"{C.STUB_BEGIN} global/outer", f"{C.STUB_BEGIN} global/inner", ] ), encoding="utf-8", ) with pytest.raises(ValueError, match="Nested stub not permitted"): FileInfo.from_file(nested) unmatched_end = tmp_path / "unmatched.py" unmatched_end.write_text(C.STUB_END + "\n", encoding="utf-8") with pytest.raises(ValueError, match="Unmatched"): FileInfo.from_file(unmatched_end) unclosed = tmp_path / "unclosed.py" unclosed.write_text(f"{C.STUB_BEGIN} global/method\n", encoding="utf-8") with pytest.raises(ValueError, match="Unclosed stub block"): FileInfo.from_file(unclosed) def test_funcinfo_gen_variants() -> None: called: list[str] = [] def ty_map(name: str) -> str: called.append(name) return name schema_no_args = NamedTypeSchema("demo.no_args", TypeSchema("Callable", ())) func = FuncInfo(schema=schema_no_args, is_member=False) assert func.gen(ty_map, indent=2) == " def no_args(*args: Any) -> Any: ..." assert called == ["Any"] schema_member = NamedTypeSchema( "pkg.Class.method", TypeSchema( "Callable", ( TypeSchema("str"), TypeSchema("int"), TypeSchema("float"), ), ), ) member_func = FuncInfo(schema=schema_member, is_member=True) assert ( member_func.gen(_identity_ty_map, indent=0) == "def method(self, _1: float, /) -> str: ..." ) schema_bad = NamedTypeSchema("bad", TypeSchema("int")) with pytest.raises(ValueError): FuncInfo(schema=schema_bad, is_member=False).gen(_identity_ty_map, indent=0) def test_objectinfo_gen_fields_and_methods() -> None: ty_calls: list[str] = [] def ty_map(name: str) -> str: ty_calls.append(name) return {"list": "Sequence", "dict": "Mapping"}.get(name, name) info = ObjectInfo( fields=[ NamedTypeSchema("field_a", TypeSchema("list", (TypeSchema("int"),))), NamedTypeSchema( "field_b", TypeSchema("dict", (TypeSchema("str"), TypeSchema("float"))) ), ], methods=[ FuncInfo( schema=NamedTypeSchema("demo.static", TypeSchema("Callable", (TypeSchema("int"),))), is_member=False, ), FuncInfo( schema=NamedTypeSchema( "demo.member", TypeSchema("Callable", (TypeSchema("str"), TypeSchema("bytes"))), ), is_member=True, ), ], ) assert info.gen_fields(ty_map, indent=2) == [ " field_a: Sequence[int]", " field_b: Mapping[str, float]", ] assert ty_calls.count("list") == 1 and ty_calls.count("dict") == 1 methods = info.gen_methods(_identity_ty_map, indent=2) assert methods == [ " @staticmethod", " def static() -> int: ...", " def member(self, /) -> str: ...", ] def test_type_schema_container_origins() -> None: """Test that Array/List/Map/Dict origins are distinct and validated correctly.""" # Array and List: 0 or 1 arg, default to (Any,) for origin in ("Array", "List"): s = TypeSchema(origin) assert s.args == (TypeSchema("Any"),), f"{origin} should default to (Any,)" s = TypeSchema(origin, (TypeSchema("int"),)) assert s.repr() == f"{origin}[int]" # Map and Dict: 0 or 2 args, default to (Any, Any) for origin in ("Map", "Dict"): s = TypeSchema(origin) assert s.args == (TypeSchema("Any"), TypeSchema("Any")), ( f"{origin} should default to (Any, Any)" ) s = TypeSchema(origin, (TypeSchema("str"), TypeSchema("float"))) assert s.repr() == f"{origin}[str, float]" # from_json_str round-trip through _TYPE_SCHEMA_ORIGIN_CONVERTER s = TypeSchema.from_json_str('{"type":"ffi.Array","args":[{"type":"int"}]}') assert s.origin == "Array" assert s.repr() == "Array[int]" s = TypeSchema.from_json_str('{"type":"ffi.List","args":[{"type":"str"}]}') assert s.origin == "List" assert s.repr() == "List[str]" s = TypeSchema.from_json_str('{"type":"ffi.Map","args":[{"type":"str"},{"type":"int"}]}') assert s.origin == "Map" assert s.repr() == "Map[str, int]" s = TypeSchema.from_json_str('{"type":"ffi.Dict","args":[{"type":"str"},{"type":"float"}]}') assert s.origin == "Dict" assert s.repr() == "Dict[str, float]" # Backward compat: "list" and "dict" origins still work s = TypeSchema("list", (TypeSchema("int"),)) assert s.repr() == "list[int]" s = TypeSchema("dict", (TypeSchema("str"), TypeSchema("int"))) assert s.repr() == "dict[str, int]" def test_objectinfo_gen_fields_container_types() -> None: """Test that ObjectInfo fields render distinct container annotations.""" info = ObjectInfo( fields=[ NamedTypeSchema("arr", TypeSchema("Array", (TypeSchema("int"),))), NamedTypeSchema("lst", TypeSchema("List", (TypeSchema("str"),))), NamedTypeSchema("mp", TypeSchema("Map", (TypeSchema("str"), TypeSchema("int")))), NamedTypeSchema("dt", TypeSchema("Dict", (TypeSchema("str"), TypeSchema("float")))), ], methods=[], ) assert info.gen_fields(_type_suffix, indent=0) == [ "arr: Sequence[int]", "lst: MutableSequence[str]", "mp: Mapping[str, int]", "dt: MutableMapping[str, float]", ] def test_generate_global_funcs_updates_block() -> None: code = CodeBlock( kind="global", param=("demo", "mockpkg"), lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} global/demo@mockpkg", C.STUB_END], ) funcs = [ FuncInfo( schema=NamedTypeSchema( "demo.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ), is_member=False, ) ] opts = Options(indent=2) imports: list[ImportItem] = [] generate_global_funcs(code, funcs, _default_ty_map(), imports, opts) assert imports == [ ImportItem("mockpkg.init_ffi_api", alias="_FFI_INIT_FUNC"), ImportItem("typing.TYPE_CHECKING"), ] assert code.lines == [ f"{C.STUB_BEGIN} global/demo@mockpkg", "# fmt: off", '_FFI_INIT_FUNC("demo", __name__)', "if TYPE_CHECKING:", " def add_one(_0: int, /) -> int: ...", "# fmt: on", C.STUB_END, ] def test_generate_global_funcs_noop_on_empty_list() -> None: code = CodeBlock( kind="global", param=("empty", ""), lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} global/empty", C.STUB_END], ) imports: list[ImportItem] = [] generate_global_funcs(code, [], _default_ty_map(), imports, Options()) assert code.lines == [f"{C.STUB_BEGIN} global/empty", C.STUB_END] assert imports == [] def test_generate_global_funcs_respects_custom_import_from() -> None: code = CodeBlock( kind="global", param=("demo", "custom.mod"), lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} global/demo@custom.mod", C.STUB_END], ) funcs = [ FuncInfo( schema=NamedTypeSchema( "demo.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ), is_member=False, ) ] imports: list[ImportItem] = [] generate_global_funcs(code, funcs, _default_ty_map(), imports, Options(indent=0)) assert ImportItem("custom.mod.init_ffi_api", alias="_FFI_INIT_FUNC") in imports def test_generate_global_funcs_aliases_colliding_type() -> None: """When a function name matches a type name, the type import gets an alias.""" code = CodeBlock( kind="global", param=("demo", "mockpkg"), lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} global/demo@mockpkg", C.STUB_END], ) # Function "demo.Foo" returns type "demo.Foo" — name collision funcs = [ FuncInfo( schema=NamedTypeSchema( "demo.Foo", TypeSchema("Callable", (TypeSchema("demo.Foo"), TypeSchema("Any"))), ), is_member=False, ) ] ty_map = _default_ty_map() ty_map["demo.Foo"] = "somepkg.Foo" imports: list[ImportItem] = [] generate_global_funcs(code, funcs, ty_map, imports, Options(indent=4)) # The type import should use an alias to avoid shadowing the function assert ImportItem("somepkg.Foo", type_checking_only=True, alias="_Foo") in imports # The function annotation should use the alias assert any("-> _Foo:" in line for line in code.lines) def test_generate_object_fields_only_block() -> None: code = CodeBlock( kind="object", param="demo.TypeDerived", lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} object/demo.TypeDerived", C.STUB_END], ) opts = Options(indent=4) imports: list[ImportItem] = [] info = ObjectInfo( fields=[ NamedTypeSchema("field_a", TypeSchema("int")), NamedTypeSchema("field_b", TypeSchema("float")), ], methods=[], type_key="demo.TypeDerived", parent_type_key="demo.Parent", ) generate_object( code, _default_ty_map(), imports, opts, info, ) assert imports == [] expected = [ f"{C.STUB_BEGIN} object/demo.TypeDerived", " " * code.indent + "# fmt: off", *[(" " * code.indent) + line for line in info.gen_fields(_type_suffix, indent=0)], " " * code.indent + "# fmt: on", C.STUB_END, ] assert code.lines == expected def test_generate_object_with_methods() -> None: code = CodeBlock( kind="object", param="demo.IntPair", lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} object/demo.IntPair", C.STUB_END], ) opts = Options(indent=4) imports: list[ImportItem] = [] info = ObjectInfo( fields=[], methods=[ FuncInfo.from_schema( "demo.IntPair.__ffi_init__", TypeSchema("Callable", (TypeSchema("None"), TypeSchema("int"), TypeSchema("int"))), is_member=True, ), FuncInfo.from_schema( "demo.IntPair.sum", TypeSchema("Callable", (TypeSchema("int"),)), is_member=True, ), ], type_key="demo.IntPair", parent_type_key="demo.Parent", ) generate_object(code, _default_ty_map(), imports, opts, info) assert set(imports) == {ImportItem("typing.TYPE_CHECKING")} assert code.lines[0] == f"{C.STUB_BEGIN} object/demo.IntPair" assert code.lines[-1] == C.STUB_END assert "# fmt: off" in code.lines[1] assert any("if TYPE_CHECKING:" in line for line in code.lines) method_lines = [line for line in code.lines if "def __ffi_init__" in line or "def sum" in line] # __ffi_init__ from TypeMethod is rendered as an instance method (self, ...) -> None assert any(line.strip().startswith("def __ffi_init__(self") for line in method_lines) assert any(line.strip().startswith("def sum") for line in method_lines) def test_generate_import_section_groups_modules() -> None: code = CodeBlock( kind="import-section", param="", lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} import", C.STUB_END], ) imports = [ ImportItem("typing.Any", type_checking_only=True), ImportItem("demo_pkg.Tensor", type_checking_only=True), ImportItem("demo.TestObjectBase", type_checking_only=True), ImportItem("custom.mod.Type", type_checking_only=True), ] opts = Options(indent=4) generate_import_section(code, imports, opts) expected_prefix = [ f"{C.STUB_BEGIN} import", "# fmt: off", "# isort: off", "from __future__ import annotations", "from typing import TYPE_CHECKING", "if TYPE_CHECKING:", ] assert code.lines[: len(expected_prefix)] == expected_prefix assert " from demo import TestObjectBase" in code.lines assert " from demo_pkg import Tensor" in code.lines assert " from custom.mod import Type" in code.lines assert " from typing import Any" in code.lines assert code.lines[-2:] == ["# fmt: on", C.STUB_END] def test_generate_import_section_no_imports_noop() -> None: code = CodeBlock( kind="import-section", param="", lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} import", C.STUB_END], ) before = list(code.lines) generate_import_section(code, [], Options()) assert code.lines == before def test_generate_all_builds_sorted_and_deduped_list() -> None: code = CodeBlock( kind="global", param="all", lineno_start=1, lineno_end=2, lines=[" " + C.STUB_BEGIN + " global/all", C.STUB_END], ) generate_all( code, names={"tvm_ffi.foo", "bar", "pkg.baz", "bar"}, # duplicates stripped opt=Options(indent=2), ) assert code.lines == [ " " + C.STUB_BEGIN + " global/all", ' "bar",', ' "baz",', ' "foo",', C.STUB_END, ] def test_generate_all_noop_on_empty_names() -> None: code = CodeBlock( kind="global", param="all-empty", lineno_start=1, lineno_end=2, lines=[C.STUB_BEGIN + " global/all-empty", C.STUB_END], ) before = list(code.lines) generate_all(code, names=set(), opt=Options()) assert code.lines == before def test_generate_all_uses_isort_style_ordering() -> None: code = CodeBlock( kind="global", param="all-mixed", lineno_start=1, lineno_end=2, lines=[C.STUB_BEGIN + " global/all-mixed", C.STUB_END], ) names = {"foo", "Bar", "LIB", "baz", "Alpha", "CONST"} generate_all(code, names=names, opt=Options(indent=0)) assert code.lines == [ C.STUB_BEGIN + " global/all-mixed", '"CONST",', '"LIB",', '"Alpha",', '"Bar",', '"baz",', '"foo",', C.STUB_END, ] def test_stage_3_adds_LIB_when_load_lib_imported(tmp_path: Path) -> None: path = tmp_path / "demo.py" global_block = CodeBlock( kind="global", param=("testing", ""), lineno_start=2, lineno_end=3, lines=[f"{C.STUB_BEGIN} global/testing", C.STUB_END], ) import_obj_block = CodeBlock( kind="import-object", param=("tvm_ffi.libinfo.load_lib_module", "False", "_FFI_LOAD_LIB"), lineno_start=1, lineno_end=1, lines=[f"{C.STUB_IMPORT_OBJECT} tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB"], ) all_block = CodeBlock( kind="__all__", param="", lineno_start=4, lineno_end=5, lines=[f"{C.STUB_BEGIN} __all__", C.STUB_END], ) file_info = FileInfo( path=path, lines=tuple( line for block in (import_obj_block, global_block, all_block) for line in block.lines ), code_blocks=[import_obj_block, global_block, all_block], ) funcs = [ FuncInfo.from_schema( "testing.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ) ] _stage_3( file_info, Options(dry_run=True), _default_ty_map(), {"testing": funcs}, ) lib_lines = [line for line in all_block.lines if "LIB" in line] assert any("LIB" in line for line in lib_lines) def test_generate_export_builds_all_extension() -> None: code = CodeBlock( kind="export", param="ffi_api", lineno_start=1, lineno_end=2, lines=[f"{C.STUB_BEGIN} export/ffi_api", C.STUB_END], ) generate_export(code) full_text = "\n".join(code.lines) assert "from .ffi_api import *" in full_text assert "ffi_api__all__" in full_text def test_generate_init_with_and_without_existing_export_block() -> None: code_no_blocks = generate_init([], "demo") assert "Package demo." in code_no_blocks assert f"{C.STUB_BEGIN} export/_ffi_api" in code_no_blocks code_with_export = generate_init( [ CodeBlock( kind="export", param="_ffi_api", lineno_start=1, lineno_end=2, lines=["", ""], ) ], "demo", ) assert code_with_export == "" def test_generate_ffi_api_without_objects_includes_sections() -> None: init_cfg = InitConfig(pkg="pkg", shared_target="pkg_shared", prefix="pkg.") code = generate_ffi_api( [], _default_ty_map(), "demo.mod", [], init_cfg, is_root=False, ) assert f"{C.STUB_BEGIN} import-section" in code assert f"{C.STUB_BEGIN} global/demo.mod" in code assert C.STUB_BEGIN + " __all__" in code assert "LIB =" not in code def test_generate_ffi_api_with_objects_imports_parents() -> None: init_cfg = InitConfig(pkg="pkg", shared_target="pkg_shared", prefix="pkg.") obj_info = ObjectInfo( fields=[], methods=[], type_key="demo.TypeDerived", parent_type_key="demo.Parent", ) parent_key = obj_info.parent_type_key code = generate_ffi_api( [], _default_ty_map(), "demo", [obj_info], init_cfg, is_root=False, ) assert C.STUB_IMPORT_OBJECT in code # register_object prompt assert f"{C.STUB_BEGIN} object/{obj_info.type_key}" in code assert parent_key is not None parent_import_prompt = ( f"{C.STUB_IMPORT_OBJECT} {parent_key};False;_{parent_key.replace('.', '_')}" ) assert parent_import_prompt in code def test_stage_2_filters_prefix_and_marks_root( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: prefixes: dict[str, list[FuncInfo]] = {"demo.sub": [], "demo": [], "other": []} monkeypatch.setattr(stub_cli, "collect_type_keys", lambda: prefixes) monkeypatch.setattr(stub_cli, "toposort_objects", lambda objs: []) global_funcs = { "demo.sub": [ FuncInfo.from_schema( "demo.sub.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ) ], "demo": [ FuncInfo.from_schema( "demo.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ) ], "other": [ FuncInfo.from_schema( "other.add_one", TypeSchema("Callable", (TypeSchema("int"), TypeSchema("int"))), ) ], } _stage_2( files=[], ty_map=_default_ty_map(), init_cfg=InitConfig(pkg="demo-pkg", shared_target="demo_shared", prefix="demo."), init_path=tmp_path, global_funcs=global_funcs, ) root_api = tmp_path / "demo" / "_ffi_api.py" sub_api = tmp_path / "demo" / "sub" / "_ffi_api.py" other_api = tmp_path / "other" / "_ffi_api.py" assert root_api.exists() assert sub_api.exists() assert not other_api.exists() root_text = root_api.read_text(encoding="utf-8") sub_text = sub_api.read_text(encoding="utf-8") assert 'LIB = _FFI_LOAD_LIB("demo-pkg", "demo_shared")' in root_text assert "LIB =" not in sub_text tvm-ffi-0.1.12/tests/python/test_tensor.py000066400000000000000000000152701521067262500205640ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations from typing import Any, NamedTuple, NoReturn import numpy.typing as npt import pytest try: import torch import torch.version except ImportError: torch = None # ty: ignore[invalid-assignment] import numpy as np import tvm_ffi def test_tensor_attributes() -> None: data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int16") if not hasattr(data, "__dlpack__"): return x = tvm_ffi.from_dlpack(data) assert isinstance(x, tvm_ffi.Tensor) assert x.shape == (10, 8, 4, 2) assert x.ndim == 4 assert x.numel() == 640 assert x.size(0) == 10 assert x.size(-1) == 2 assert x.is_contiguous() assert x.strides == (64, 8, 2, 1) assert x.dtype == tvm_ffi.dtype("int16") assert x.device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCPU assert x.device.index == 0 x2 = np.from_dlpack(x) np.testing.assert_equal(x2, data) def test_empty_tensor_is_contiguous() -> None: # Empty tensors are trivially contiguous regardless of what # strides the producer reports (numpy 2.3+ via __dlpack__ now # reports (0, 0, 0) for shape (4, 0, 4)). See PR #607 review # comment for context. data: npt.NDArray[Any] = np.zeros((4, 0, 4), dtype="int16") if not hasattr(data, "__dlpack__"): return x = tvm_ffi.from_dlpack(data) assert x.is_contiguous() def test_non_contiguous_tensor_attributes() -> None: data: npt.NDArray[Any] = np.zeros((4, 4, 4), dtype="int16") slice = data[1:3, :, 1:3] if not hasattr(slice, "__dlpack__"): return x = tvm_ffi.from_dlpack(slice) assert isinstance(x, tvm_ffi.Tensor) assert x.shape == (2, 4, 2) assert x.numel() == 16 assert x.size(0) == 2 assert x.size(-1) == 2 assert not x.is_contiguous() assert x.strides == (16, 4, 1) def test_shape_object() -> None: shape = tvm_ffi.Shape((10, 8, 4, 2)) assert isinstance(shape, tvm_ffi.Shape) assert shape == (10, 8, 4, 2) fecho = tvm_ffi.convert(lambda x: x) shape2: tvm_ffi.Shape = fecho(shape) assert shape2._tvm_ffi_cached_object.same_as(shape._tvm_ffi_cached_object) assert isinstance(shape2, tvm_ffi.Shape) assert isinstance(shape2, tuple) shape3: tvm_ffi.Shape = tvm_ffi.convert(shape) assert shape3._tvm_ffi_cached_object.same_as(shape._tvm_ffi_cached_object) assert isinstance(shape3, tvm_ffi.Shape) @pytest.mark.skipif(torch is None, reason="Fast torch dlpack importer is not enabled") def test_tensor_auto_dlpack() -> None: assert torch is not None x = torch.arange(128) fecho = tvm_ffi.get_global_func("testing.echo") y = fecho(x) assert isinstance(y, torch.Tensor) assert y.data_ptr() == x.data_ptr() assert y.dtype == x.dtype assert y.shape == x.shape assert y.device == x.device np.testing.assert_equal(y.numpy(), x.numpy()) @pytest.mark.skipif(torch is None, reason="Fast torch dlpack importer is not enabled") def test_tensor_auto_dlpack_with_error() -> None: assert torch is not None x = torch.arange(128) def raise_torch_error(x: Any) -> NoReturn: raise ValueError("error XYZ") f = tvm_ffi.convert(raise_torch_error) with pytest.raises(ValueError): # pass in torch argment to trigger the error in set allocator path f(x) def test_tensor_class_override() -> None: class MyTensor(tvm_ffi.Tensor): pass old_tensor = tvm_ffi.core._CLASS_TENSOR tvm_ffi.core._set_class_tensor(MyTensor) data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int16") if not hasattr(data, "__dlpack__"): return x = tvm_ffi.from_dlpack(data) fecho = tvm_ffi.get_global_func("testing.echo") y = fecho(x) assert isinstance(y, MyTensor) tvm_ffi.core._set_class_tensor(old_tensor) def test_tvm_ffi_tensor_compatible() -> None: class MyTensor: def __init__(self, tensor: tvm_ffi.Tensor) -> None: """Initialize the MyTensor.""" self._tensor = tensor def __tvm_ffi_object__(self) -> tvm_ffi.Tensor: """Implement __tvm_ffi_object__ protocol.""" return self._tensor data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int32") if not hasattr(data, "__dlpack__"): return x = tvm_ffi.from_dlpack(data) y = MyTensor(x) fecho = tvm_ffi.get_global_func("testing.echo") z = fecho(y) assert z.__chandle__() == x.__chandle__() class MyNamedTuple(NamedTuple): a: MyTensor b: int args = MyNamedTuple(a=y, b=1) z = fecho(args) assert z[0].__chandle__() == x.__chandle__() assert z[1] == 1 class MyCustom: def __init__(self, a: MyTensor, b: int) -> None: self.a = a self.b = b def __tvm_ffi_value__(self) -> Any: """Implement __tvm_ffi_value__ protocol.""" return (self.a, self.b) z = fecho(MyCustom(a=y, b=2)) assert z[0].__chandle__() == x.__chandle__() assert z[1] == 2 @pytest.mark.skipif( torch is None or not torch.cuda.is_available() or torch.version.hip is None, reason="ROCm is not enabled in PyTorch", ) def test_tensor_from_pytorch_rocm() -> None: assert torch is not None @tvm_ffi.register_global_func("testing.check_device", override=True) def _check_device(x: tvm_ffi.Tensor) -> str: return x.device.type # PyTorch uses device name "cuda" to represent ROCm device x = torch.randn(128, device="cuda") device_type = tvm_ffi.get_global_func("testing.check_device")(x) assert device_type == "rocm" def test_optional_tensor_view() -> None: optional_tensor_view_has_value = tvm_ffi.get_global_func( "testing.optional_tensor_view_has_value" ) assert not optional_tensor_view_has_value(None) x: npt.NDArray[Any] = np.zeros((128,), dtype="float32") if not hasattr(x, "__dlpack__"): return assert optional_tensor_view_has_value(x) if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/test_type_converter.py000066400000000000000000005552221521067262500223300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for TypeSchema type conversion to CAny.""" from __future__ import annotations import collections.abc import ctypes import itertools import os import sys import typing from numbers import Integral from typing import Callable, Iterator, Optional, Union import pytest import tvm_ffi from tvm_ffi.core import ( CAny, ObjectConvertible, TypeSchema, _lookup_type_attr, _object_type_key_to_index, _to_py_class_value, ) from tvm_ffi.dataclasses import IntEnum, StrEnum, entry # Python 3.9+ supports list[int], dict[str, int], tuple[int, ...] at runtime. # On 3.8, these raise TypeError("'type' object is not subscriptable"). _PY39 = sys.version_info >= (3, 9) from tvm_ffi.testing import ( TestIntPair, TestObjectBase, TestObjectDerived, _TestCxxClassBase, _TestCxxClassDerived, _TestCxxClassDerivedDerived, ) from tvm_ffi.testing.testing import requires_py39, requires_py310 # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- _TYPE_KEY_COUNTER = itertools.count() def S(origin: str, *args: TypeSchema) -> TypeSchema: """Shorthand constructor for TypeSchema (string-based).""" return TypeSchema(origin, tuple(args)) def _unique_type_key(base: str) -> str: return f"testing.type_converter.{base}_{next(_TYPE_KEY_COUNTER)}" def _make_int_enum_type() -> typing.Any: class Colors(IntEnum, type_key=_unique_type_key("IntEnum")): red = entry(value=10) blue = entry(value=20) return Colors def _make_str_enum_type() -> typing.Any: class Tokens(StrEnum, type_key=_unique_type_key("StrEnum")): add = entry(value="+") mul = entry(value="*") return Tokens # Annotation-based constructor — the main subject under test. A = TypeSchema.from_annotation # --------------------------------------------------------------------------- # Category 1: POD type exact match (check_value) # --------------------------------------------------------------------------- class TestPODExactMatch: def test_int(self) -> None: """Test int.""" A(int).check_value(42) def test_float(self) -> None: """Test float.""" A(float).check_value(3.14) def test_bool_true(self) -> None: """Test bool true.""" A(bool).check_value(True) def test_bool_false(self) -> None: """Test bool false.""" A(bool).check_value(False) def test_str(self) -> None: """Test str.""" A(str).check_value("hello") def test_bytes(self) -> None: """Test bytes.""" A(bytes).check_value(b"data") def test_none(self) -> None: """Test none.""" A(type(None)).check_value(None) # --------------------------------------------------------------------------- # Category 2: Implicit conversions (mirrors TryCastFromAnyView) # --------------------------------------------------------------------------- class TestImplicitConversions: def test_bool_to_int(self) -> None: """Bool -> int is OK (C++: int accepts bool).""" A(int).check_value(True) def test_int_to_float(self) -> None: """Int -> float is OK (C++: float accepts int).""" A(float).check_value(42) def test_bool_to_float(self) -> None: """Bool -> float is OK (C++: float accepts bool).""" A(float).check_value(True) def test_int_to_bool(self) -> None: """Int -> bool is OK (C++: bool accepts int).""" A(bool).check_value(1) # --------------------------------------------------------------------------- # Category 3: Rejection cases # --------------------------------------------------------------------------- class TestRejections: def test_str_not_int(self) -> None: """Test str not int.""" with pytest.raises(TypeError, match="expected int"): A(int).check_value("hello") def test_float_not_int(self) -> None: """Test float not int.""" with pytest.raises(TypeError): A(int).check_value(3.14) def test_none_not_int(self) -> None: """Test none not int.""" with pytest.raises(TypeError): A(int).check_value(None) def test_int_not_str(self) -> None: """Test int not str.""" with pytest.raises(TypeError): A(str).check_value(42) def test_str_not_bool(self) -> None: """Test str not bool.""" with pytest.raises(TypeError): A(bool).check_value("hello") def test_none_not_str(self) -> None: """Test none not str.""" with pytest.raises(TypeError): A(str).check_value(None) def test_int_not_bytes(self) -> None: """Test int not bytes.""" with pytest.raises(TypeError): A(bytes).check_value(42) def test_int_not_none(self) -> None: """Test int not none.""" with pytest.raises(TypeError): A(type(None)).check_value(42) # --------------------------------------------------------------------------- # Category 4: Special types # --------------------------------------------------------------------------- class TestSpecialTypes: def test_device_pass(self) -> None: """Test device pass.""" dev = tvm_ffi.Device("cpu", 0) A(tvm_ffi.Device).check_value(dev) def test_device_fail(self) -> None: """Test device fail.""" with pytest.raises(TypeError): A(tvm_ffi.Device).check_value(42) def test_dtype_pass(self) -> None: """Test dtype pass.""" dt = tvm_ffi.core.DataType("float32") A(tvm_ffi.core.DataType).check_value(dt) def test_dtype_str_pass(self) -> None: """Str accepted as dtype (will be parsed).""" A(tvm_ffi.core.DataType).check_value("float32") def test_dtype_fail(self) -> None: """Test dtype fail.""" with pytest.raises(TypeError): A(tvm_ffi.core.DataType).check_value(42) def test_opaque_ptr_pass(self) -> None: """Test opaque ptr pass.""" A(ctypes.c_void_p).check_value(ctypes.c_void_p(0)) def test_opaque_ptr_none_pass(self) -> None: """Test opaque ptr none pass.""" A(ctypes.c_void_p).check_value(None) def test_opaque_ptr_fail(self) -> None: """Test opaque ptr fail.""" with pytest.raises(TypeError): A(ctypes.c_void_p).check_value(42) def test_callable_pass_function(self) -> None: """Test callable pass function.""" A(Callable).check_value(lambda x: x) def test_callable_pass_builtin(self) -> None: """Test callable pass builtin.""" A(Callable).check_value(len) def test_callable_fail(self) -> None: """Test callable fail.""" with pytest.raises(TypeError): A(Callable).check_value(42) def test_collections_abc_callable_pass_function(self) -> None: """collections.abc.Callable accepts Python functions.""" A(collections.abc.Callable).check_value(lambda x: x) def test_collections_abc_callable_pass_builtin(self) -> None: """collections.abc.Callable accepts builtins.""" A(collections.abc.Callable).check_value(len) def test_collections_abc_callable_fail(self) -> None: """collections.abc.Callable rejects non-callables.""" with pytest.raises(TypeError, match="expected Callable"): A(collections.abc.Callable).check_value(42) def test_callable_cobject_wraps_to_function(self) -> None: """Callable CObjects are wrapped instead of asserting.""" class CallableObj(TestObjectBase): def __call__(self, x: int) -> int: return x + 1 obj = CallableObj(v_i64=1, v_f64=2.0, v_str="s") with pytest.raises(TypeError, match=r"expected Callable, got .*TestObjectBase"): A(Callable).check_value(obj) # --------------------------------------------------------------------------- # Category 5: Object types # --------------------------------------------------------------------------- class TestObjectTypes: def test_object_pass(self) -> None: """Any CObject passes TypeSchema('Object').""" f = tvm_ffi.get_global_func("testing.echo") A(tvm_ffi.core.Object).check_value(f) def test_object_fail(self) -> None: """Test object fail.""" with pytest.raises(TypeError): A(tvm_ffi.core.Object).check_value(42) def test_specific_object_pass(self) -> None: """A Function object should pass its own type schema.""" f = tvm_ffi.get_global_func("testing.echo") A(Callable).check_value(f) def test_function_from_extern_c_exists(self) -> None: """ffi.FunctionFromExternC should be registered.""" fn = tvm_ffi.get_global_func("ffi.FunctionFromExternC", allow_missing=True) assert fn is not None, "ffi.FunctionFromExternC not registered" # --------------------------------------------------------------------------- # Category 6: Payload enums # --------------------------------------------------------------------------- class TestPayloadEnums: def test_int_enum_convert_from_int(self) -> None: """IntEnum accepts its user-visible integer payload.""" Colors = _make_int_enum_type() result = _to_py_class_value(A(Colors).convert(20)) assert result.same_as(Colors.blue) def test_int_enum_passthrough_existing_object(self) -> None: """IntEnum keeps the object passthrough path for existing enum objects.""" Colors = _make_int_enum_type() result = _to_py_class_value(A(Colors).convert(Colors.red)) assert result.same_as(Colors.red) def test_int_enum_rejects_unknown_payload(self) -> None: """IntEnum still rejects unmatched integer payloads.""" Colors = _make_int_enum_type() with pytest.raises(TypeError, match="expected"): A(Colors).check_value(99) def test_str_enum_convert_from_str(self) -> None: """StrEnum accepts its user-visible string payload.""" Tokens = _make_str_enum_type() result = _to_py_class_value(A(Tokens).convert("*")) assert result.same_as(Tokens.mul) def test_str_enum_rejects_unknown_payload(self) -> None: """StrEnum still rejects unmatched string payloads.""" Tokens = _make_str_enum_type() with pytest.raises(TypeError, match="expected"): A(Tokens).check_value("/") # --------------------------------------------------------------------------- # Category 7: Optional # --------------------------------------------------------------------------- class TestOptional: def test_none_passes(self) -> None: """Test none passes.""" A(Optional[int]).check_value(None) def test_inner_type_passes(self) -> None: """Test inner type passes.""" A(Optional[int]).check_value(42) def test_wrong_type_fails(self) -> None: """Test wrong type fails.""" with pytest.raises(TypeError, match="expected int"): A(Optional[int]).check_value("hello") def test_nested_optional(self) -> None: """Test nested optional.""" schema = A(Optional[Optional[int]]) schema.check_value(None) schema.check_value(42) # --------------------------------------------------------------------------- # Category 8: Union / Variant # --------------------------------------------------------------------------- class TestUnion: def test_first_alt_passes(self) -> None: """Test first alt passes.""" A(Union[int, str]).check_value(42) def test_second_alt_passes(self) -> None: """Test second alt passes.""" A(Union[int, str]).check_value("hello") def test_no_alt_matches(self) -> None: """Test no alt matches.""" with pytest.raises(TypeError, match="got float"): A(Union[int, str]).check_value(3.14) def test_bool_matches_int_alt(self) -> None: """Bool is accepted by the int alternative.""" A(Union[int, str]).check_value(True) # --------------------------------------------------------------------------- # Category 9: Containers # --------------------------------------------------------------------------- class TestContainers: @requires_py39 def test_array_list_pass(self) -> None: """Test array list pass.""" A(tuple[int, ...]).check_value([1, 2, 3]) @requires_py39 def test_array_tuple_pass(self) -> None: """Test array tuple pass.""" A(tuple[int, ...]).check_value((1, 2, 3)) @requires_py39 def test_array_wrong_element(self) -> None: """Test array wrong element.""" with pytest.raises(TypeError, match=r"element \[1\].*expected int"): A(tuple[int, ...]).check_value([1, "x"]) @requires_py39 def test_array_empty_pass(self) -> None: """Test array empty pass.""" A(tuple[int, ...]).check_value([]) @requires_py39 def test_array_any_pass(self) -> None: """Test array any pass.""" A(tuple[typing.Any, ...]).check_value([1, "x", None]) @requires_py39 def test_array_wrong_container_type(self) -> None: """Test array wrong container type.""" with pytest.raises(TypeError, match="expected Array"): A(tuple[int, ...]).check_value(42) @requires_py39 def test_array_rejects_generator(self) -> None: """Generators are not accepted by Array schemas.""" def gen() -> Iterator[int]: yield 1 yield 2 with pytest.raises(TypeError, match="expected Array"): A(tuple[int, ...]).check_value(gen()) @requires_py39 def test_array_rejects_string(self) -> None: """Strings are not accepted by Array schemas.""" with pytest.raises(TypeError, match="expected Array"): A(tuple[int, ...]).check_value("hello") @requires_py39 def test_list_pass(self) -> None: """Test list pass.""" A(list[str]).check_value(["a", "b"]) @requires_py39 def test_map_pass(self) -> None: """Test map pass.""" A(tvm_ffi.Map[str, int]).check_value({"a": 1, "b": 2}) @requires_py39 def test_map_wrong_key(self) -> None: """Test map wrong key.""" with pytest.raises(TypeError, match="expected str"): A(tvm_ffi.Map[str, int]).check_value({1: 2}) @requires_py39 def test_map_wrong_value(self) -> None: """Test map wrong value.""" with pytest.raises(TypeError, match="expected int"): A(tvm_ffi.Map[str, int]).check_value({"a": "b"}) @requires_py39 def test_map_empty_pass(self) -> None: """Test map empty pass.""" A(tvm_ffi.Map[str, int]).check_value({}) @requires_py39 def test_dict_pass(self) -> None: """Test dict pass.""" A(dict[str, int]).check_value({"a": 1}) @requires_py39 def test_map_wrong_container(self) -> None: """Test map wrong container.""" with pytest.raises(TypeError, match="expected Map"): A(tvm_ffi.Map[str, int]).check_value([1, 2]) @requires_py39 def test_map_rejects_non_mapping_pairs(self) -> None: """Lists of pairs are not accepted by Map schemas.""" with pytest.raises(TypeError, match="expected Map"): A(tvm_ffi.Map[str, int]).check_value([("a", 1)]) # --------------------------------------------------------------------------- # Category 9: Nested types # --------------------------------------------------------------------------- class TestNestedTypes: @requires_py39 def test_array_optional_int(self) -> None: """Test array optional int.""" A(tuple[Optional[int], ...]).check_value([1, None, 2]) @requires_py39 def test_map_str_array_int(self) -> None: """Test map str array int.""" A(tvm_ffi.Map[str, tuple[int, ...]]).check_value({"a": [1, 2]}) @requires_py39 def test_map_str_array_int_nested_fail(self) -> None: """Test map str array int nested fail.""" with pytest.raises(TypeError, match="expected int"): A(tvm_ffi.Map[str, tuple[int, ...]]).check_value({"a": [1, "x"]}) @requires_py39 def test_union_with_containers(self) -> None: """Test union with containers.""" schema = A(Union[int, tuple[str, ...]]) schema.check_value(42) schema.check_value(["a", "b"]) with pytest.raises(TypeError): schema.check_value(3.14) # --------------------------------------------------------------------------- # Category 10: Any # --------------------------------------------------------------------------- class TestAny: def test_int(self) -> None: """Test int.""" A(typing.Any).check_value(42) def test_none(self) -> None: """Test none.""" A(typing.Any).check_value(None) def test_str(self) -> None: """Test str.""" A(typing.Any).check_value("hello") def test_list(self) -> None: """Test list.""" A(typing.Any).check_value([1, 2, 3]) def test_object(self) -> None: """Test object.""" A(typing.Any).check_value(object()) def test_object_convertible_convert(self) -> None: """Any eagerly unwraps ObjectConvertible via asobject().""" inner = TestIntPair(1, 2) class Convertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner result = _to_py_class_value(A(typing.Any).convert(Convertible())) assert result.same_as(inner) def test_object_convertible_error(self) -> None: """Any surfaces asobject() failures during eager normalization.""" class BadConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: raise RuntimeError("broken") with pytest.raises(TypeError, match=r"asobject\(\) failed"): A(typing.Any).check_value(BadConvertible()) def test_object_protocol_convert(self) -> None: """Any eagerly unwraps __tvm_ffi_object__ before dispatch.""" inner = TestIntPair(1, 2) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner result = _to_py_class_value(A(typing.Any).convert(ObjProto())) assert result.same_as(inner) def test_object_protocol_error(self) -> None: """Any surfaces __tvm_ffi_object__ failures during eager normalization.""" class BadProto: def __tvm_ffi_object__(self) -> object: raise RuntimeError("broken") with pytest.raises(TypeError, match=r"__tvm_ffi_object__\(\) failed"): A(typing.Any).check_value(BadProto()) # --------------------------------------------------------------------------- # Category 11: Error message quality # --------------------------------------------------------------------------- class TestErrorMessages: def test_basic_type_mismatch(self) -> None: """Test basic type mismatch.""" with pytest.raises(TypeError, match=r"expected int, got str"): A(int).check_value("hello") @requires_py39 def test_nested_array_error(self) -> None: """Test nested array error.""" with pytest.raises(TypeError, match=r"element \[2\].*expected int, got str"): A(tuple[int, ...]).check_value([1, 2, "x"]) @requires_py39 def test_nested_map_error(self) -> None: """Test nested map error.""" with pytest.raises(TypeError, match=r"value for key 'b'.*expected int, got str"): A(tvm_ffi.Map[str, int]).check_value({"a": 1, "b": "x"}) def test_union_error_lists_alternatives(self) -> None: """Test union error lists alternatives.""" with pytest.raises(TypeError, match="got float") as exc_info: A(Union[int, str]).check_value(3.14) err = str(exc_info.value) assert "int" in err assert "str" in err def test_schema_in_error_message(self) -> None: """check_value includes the schema repr in the TypeError.""" with pytest.raises(TypeError, match=r"type check failed for"): A(int).check_value("hello") def test_convert_error_message(self) -> None: """Convert includes the schema repr in the TypeError.""" with pytest.raises(TypeError, match=r"type conversion failed for"): A(int).convert("hello") # --------------------------------------------------------------------------- # Category 12: from_type_index factory # --------------------------------------------------------------------------- class TestFromTypeIndex: def test_int(self) -> None: """Test int.""" schema = TypeSchema.from_type_index(1) # kTVMFFIInt assert schema.origin == "int" assert schema.origin_type_index == 1 def test_float(self) -> None: """Test float.""" schema = TypeSchema.from_type_index(3) # kTVMFFIFloat assert schema.origin == "float" def test_bool(self) -> None: """Test bool.""" schema = TypeSchema.from_type_index(2) # kTVMFFIBool assert schema.origin == "bool" def test_array_with_args(self) -> None: """Test array with args.""" schema = TypeSchema.from_type_index(71, (A(int),)) # kTVMFFIArray assert schema.origin == "Array" assert len(schema.args) == 1 assert schema.args[0].origin == "int" def test_roundtrip_check(self) -> None: """from_type_index then check_value works correctly.""" schema = TypeSchema.from_type_index(1) # int schema.check_value(42) with pytest.raises(TypeError): schema.check_value("hello") def test_none(self) -> None: """Test none.""" schema = TypeSchema.from_type_index(0) # kTVMFFINone assert schema.origin == "None" schema.check_value(None) def test_any(self) -> None: """Test any.""" schema = TypeSchema.from_type_index(-1) # kTVMFFIAny assert schema.origin == "Any" schema.check_value("anything") def test_str(self) -> None: """Test str.""" schema = TypeSchema.from_type_index(65) # kTVMFFIStr assert schema.origin == "str" schema.check_value("hello") def test_map_with_args(self) -> None: """Test map with args.""" schema = TypeSchema.from_type_index(72, (A(str), A(int))) # kTVMFFIMap assert schema.origin == "Map" schema.check_value({"a": 1}) # --------------------------------------------------------------------------- # Category 13: Edge cases # --------------------------------------------------------------------------- class TestEdgeCases: def test_bytearray_passes_bytes(self) -> None: """Test bytearray passes bytes.""" A(bytes).check_value(bytearray(b"data")) @requires_py39 def test_tuple_passes_array(self) -> None: """Tuple is accepted as a sequence type for Array.""" A(tuple[int, ...]).check_value((1, 2, 3)) def test_empty_union_is_rejected(self) -> None: """Union requires at least 2 args.""" with pytest.raises(ValueError, match="at least two"): TypeSchema("Union", ()) def test_origin_type_index_auto_computed(self) -> None: """origin_type_index is automatically computed from origin string.""" schema = A(int) assert schema.origin_type_index == 1 # kTVMFFIInt schema = A(float) assert schema.origin_type_index == 3 # kTVMFFIFloat schema = A(Optional[int]) assert schema.origin_type_index == -2 # structural def test_check_value_succeeds_on_valid(self) -> None: """Test check value succeeds on valid input.""" A(int).check_value(42) def test_check_value_raises_on_failure(self) -> None: """Test check value raises TypeError on failure.""" with pytest.raises(TypeError, match="expected int"): A(int).check_value("hello") @requires_py39 def test_tuple_type_schema(self) -> None: """Test tuple type schema.""" schema = A(tuple[int, str]) schema.check_value((1, "a")) with pytest.raises(TypeError): schema.check_value((1, 2)) with pytest.raises(TypeError): schema.check_value((1,)) def test_numpy_int_passes_int(self) -> None: """Numpy integer types should pass int check via Integral.""" np = pytest.importorskip("numpy") A(int).check_value(np.int64(42)) A(float).check_value(np.float64(3.14)) # =========================================================================== # Type Converter Tests (convert) # =========================================================================== # --------------------------------------------------------------------------- # Category 14: POD conversion results # --------------------------------------------------------------------------- class TestConvertPOD: def test_int_passthrough(self) -> None: """Int -> int returns the same value.""" result = _to_py_class_value(A(int).convert(42)) assert result == 42 assert type(result) is int def test_bool_to_int(self) -> None: """Bool -> int actually converts to int.""" result = _to_py_class_value(A(int).convert(True)) assert result == 1 assert type(result) is int def test_bool_false_to_int(self) -> None: """Test bool false to int.""" result = _to_py_class_value(A(int).convert(False)) assert result == 0 assert type(result) is int def test_float_passthrough(self) -> None: """Test float passthrough.""" result = _to_py_class_value(A(float).convert(3.14)) assert result == 3.14 assert type(result) is float def test_int_to_float(self) -> None: """Int -> float actually converts.""" result = _to_py_class_value(A(float).convert(42)) assert result == 42.0 assert type(result) is float def test_bool_to_float(self) -> None: """Bool -> float actually converts.""" result = _to_py_class_value(A(float).convert(True)) assert result == 1.0 assert type(result) is float def test_bool_passthrough(self) -> None: """Test bool passthrough.""" result = _to_py_class_value(A(bool).convert(True)) assert result is True assert type(result) is bool def test_int_to_bool(self) -> None: """Int -> bool actually converts.""" result = _to_py_class_value(A(bool).convert(1)) assert result is True assert type(result) is bool def test_int_zero_to_bool(self) -> None: """Test int zero to bool.""" result = _to_py_class_value(A(bool).convert(0)) assert result is False assert type(result) is bool def test_str_passthrough(self) -> None: """Test str passthrough — returns tvm_ffi.String (subclass of str).""" result = _to_py_class_value(A(str).convert("hello")) assert result == "hello" assert isinstance(result, str) assert isinstance(result, tvm_ffi.core.String) def test_bytes_passthrough(self) -> None: """Test bytes passthrough — returns tvm_ffi.Bytes (subclass of bytes).""" result = _to_py_class_value(A(bytes).convert(b"data")) assert result == b"data" assert isinstance(result, bytes) assert isinstance(result, tvm_ffi.core.Bytes) def test_bytearray_to_bytes(self) -> None: """Bytearray -> bytes converts to tvm_ffi.Bytes.""" result = _to_py_class_value(A(bytes).convert(bytearray(b"data"))) assert result == b"data" assert isinstance(result, bytes) assert isinstance(result, tvm_ffi.core.Bytes) # --------------------------------------------------------------------------- # Category 15: None disambiguation (critical design point) # --------------------------------------------------------------------------- class TestNoneDisambiguation: def test_none_converts_successfully_for_none_schema(self) -> None: """TypeSchema('None').convert(None) returns None as a valid result.""" result = _to_py_class_value(A(type(None)).convert(None)) assert result is None def test_none_converts_successfully_for_optional(self) -> None: """Optional[int].convert(None) returns None as a valid result.""" result = _to_py_class_value(A(Optional[int]).convert(None)) assert result is None def test_none_fails_for_int(self) -> None: """TypeSchema('int').convert(None) raises TypeError.""" with pytest.raises(TypeError, match="expected int, got None"): A(int).convert(None) def test_convert_none_success(self) -> None: """Convert returns None for Optional[int] with None input.""" result = _to_py_class_value(A(Optional[int]).convert(None)) assert result is None def test_convert_none_failure(self) -> None: """Convert raises TypeError for failed conversion.""" with pytest.raises(TypeError, match="expected int"): A(int).convert(None) def test_convert_success_with_value(self) -> None: """Convert returns converted value on success.""" result = _to_py_class_value(A(int).convert(True)) assert result == 1 assert type(result) is int def test_opaque_ptr_none_converts(self) -> None: """ctypes.c_void_p accepts None and converts it to a null opaque pointer.""" result = _to_py_class_value(A(ctypes.c_void_p).convert(None)) assert isinstance(result, ctypes.c_void_p) assert result.value is None def test_convert_opaque_ptr_none(self) -> None: """Test convert opaque ptr none.""" result = _to_py_class_value(A(ctypes.c_void_p).convert(None)) assert isinstance(result, ctypes.c_void_p) assert result.value is None # --------------------------------------------------------------------------- # Category 16: Special type conversions # --------------------------------------------------------------------------- class TestConvertSpecialTypes: def test_dtype_str_converts(self) -> None: """Str -> dtype actually creates a DataType object.""" result = _to_py_class_value(A(tvm_ffi.core.DataType).convert("float32")) assert isinstance(result, tvm_ffi.core.DataType) assert str(result) == "float32" def test_dtype_passthrough(self) -> None: """Test dtype passthrough.""" dt = tvm_ffi.core.DataType("int32") result = _to_py_class_value(A(tvm_ffi.core.DataType).convert(dt)) assert str(result) == str(dt) def test_device_passthrough(self) -> None: """Test device passthrough.""" dev = tvm_ffi.Device("cpu", 0) result = _to_py_class_value(A(tvm_ffi.Device).convert(dev)) assert str(result) == str(dev) def test_callable_passthrough(self) -> None: """Test callable passthrough.""" fn = lambda x: x result = _to_py_class_value(A(Callable).convert(fn)) assert callable(result) def test_opaque_ptr_passthrough(self) -> None: """Test opaque ptr passthrough.""" ptr = ctypes.c_void_p(42) result = _to_py_class_value(A(ctypes.c_void_p).convert(ptr)) assert result is not None # --------------------------------------------------------------------------- # Category 17: Container conversion results # --------------------------------------------------------------------------- class TestConvertContainers: @requires_py39 def test_array_converts_bool_elements_to_int(self) -> None: """Array[int] with bool elements converts them to int.""" result = _to_py_class_value(A(tuple[int, ...]).convert([True, False, 1])) assert list(result) == [1, 0, 1] assert all(type(x) is int for x in result) @requires_py39 def test_array_int_passthrough(self) -> None: """Array[int] with int elements returns ffi.Array.""" result = _to_py_class_value(A(tuple[int, ...]).convert([1, 2, 3])) assert list(result) == [1, 2, 3] @requires_py39 def test_array_any_passthrough(self) -> None: """Array[Any] wraps into ffi.Array.""" original = [1, "x", None] result = _to_py_class_value(A(tuple[typing.Any, ...]).convert(original)) assert isinstance(result, tvm_ffi.Array) @requires_py39 def test_map_converts_values(self) -> None: """Map[str, float] converts int values to float.""" result = _to_py_class_value(A(tvm_ffi.Map[str, float]).convert({"a": 1, "b": 2})) assert isinstance(result, tvm_ffi.Map) assert type(result["a"]) is float assert type(result["b"]) is float assert result["a"] == 1.0 assert result["b"] == 2.0 @requires_py39 def test_map_any_float_converts_values(self) -> None: """Map[Any, float] still converts values when keys are Any.""" result = _to_py_class_value(A(tvm_ffi.Map[typing.Any, float]).convert({"a": 1, "b": 2})) assert isinstance(result, tvm_ffi.Map) assert type(result["a"]) is float @requires_py39 def test_map_any_any_passthrough(self) -> None: """Map[Any, Any] wraps into ffi.Map.""" original = {"a": 1} result = _to_py_class_value(A(tvm_ffi.Map[typing.Any, typing.Any]).convert(original)) assert isinstance(result, tvm_ffi.Map) @requires_py39 def test_map_empty_dict_convert(self) -> None: """Empty dict converts to Map[str, int].""" result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert({})) assert len(result) == 0 @requires_py39 def test_dict_empty_dict_convert(self) -> None: """Empty dict converts to Dict[str, int].""" result = _to_py_class_value(A(dict[str, int]).convert({})) assert len(result) == 0 @requires_py39 def test_tuple_converts_elements(self) -> None: """tuple[int, float] converts elements positionally.""" result = _to_py_class_value(A(tuple[int, float]).convert((True, 42))) assert list(result) == [1, 42.0] assert type(result[0]) is int assert type(result[1]) is float @requires_py39 def test_nested_array_in_map(self) -> None: """Map[str, Array[int]] recursively converts elements.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[int, ...]]).convert({"a": [True, False]}) ) assert isinstance(result, tvm_ffi.Map) assert list(result["a"]) == [1, 0] assert all(type(x) is int for x in result["a"]) @requires_py39 def test_array_optional_int_all_none(self) -> None: """Array[Optional[int]] accepts an all-None payload.""" result = _to_py_class_value(A(tuple[Optional[int], ...]).convert([None, None, None])) assert list(result) == [None, None, None] # --------------------------------------------------------------------------- # Category 18: Optional/Union conversion results # --------------------------------------------------------------------------- class TestConvertComposite: def test_optional_converts_inner(self) -> None: """Optional[float].convert(42) converts int -> float.""" result = _to_py_class_value(A(Optional[float]).convert(42)) assert result == 42.0 assert type(result) is float def test_optional_none(self) -> None: """Test optional none.""" result = _to_py_class_value(A(Optional[float]).convert(None)) assert result is None def test_union_picks_first_match(self) -> None: """Union[int, str] converts bool via int alternative.""" result = _to_py_class_value(A(Union[int, str]).convert(True)) assert result == 1 assert type(result) is int def test_union_second_match(self) -> None: """Test union second match.""" result = _to_py_class_value(A(Union[int, str]).convert("hello")) assert result == "hello" def test_any_passthrough(self) -> None: """Any returns value as-is.""" result = _to_py_class_value(A(typing.Any).convert(42)) assert result == 42 result = _to_py_class_value(A(typing.Any).convert(None)) assert result is None # --------------------------------------------------------------------------- # Category 19: Convert rejection cases # --------------------------------------------------------------------------- class TestConvertRejections: def test_int_rejects_str(self) -> None: """Test int rejects str.""" with pytest.raises(TypeError, match="expected int, got str"): A(int).convert("hello") def test_int_rejects_float(self) -> None: """Test int rejects float.""" with pytest.raises(TypeError, match="expected int, got float"): A(int).convert(3.14) def test_str_rejects_int(self) -> None: """Test str rejects int.""" with pytest.raises(TypeError, match="expected str, got int"): A(str).convert(42) @requires_py39 def test_array_rejects_wrong_element(self) -> None: """Test array rejects wrong element.""" with pytest.raises(TypeError, match=r"element \[1\].*expected int, got str"): A(tuple[int, ...]).convert([1, "x"]) @requires_py39 def test_map_rejects_wrong_value(self) -> None: """Test map rejects wrong value.""" with pytest.raises(TypeError, match=r"value for key 'a'.*expected int, got str"): A(tvm_ffi.Map[str, int]).convert({"a": "x"}) @requires_py39 def test_tuple_rejects_wrong_length(self) -> None: """Test tuple rejects wrong length.""" with pytest.raises(TypeError, match=r"expected tuple of length 2"): A(tuple[int, str]).convert((1,)) def test_convert_failure_raises(self) -> None: """Test convert failure raises TypeError.""" with pytest.raises(TypeError, match="expected int"): A(int).convert("hello") # --------------------------------------------------------------------------- # Category 20: Numpy conversion # --------------------------------------------------------------------------- class TestConvertNumpy: def test_numpy_int_to_int(self) -> None: """Test numpy int to int.""" np = pytest.importorskip("numpy") result = _to_py_class_value(A(int).convert(np.int64(42))) assert result == 42 assert type(result) is int def test_numpy_float_to_float(self) -> None: """Test numpy float to float.""" np = pytest.importorskip("numpy") result = _to_py_class_value(A(float).convert(np.float64(3.14))) assert result == pytest.approx(3.14) # np.float64 is a subclass of float, so isinstance check passes # and the value is returned as-is (no forced conversion to plain float) assert isinstance(result, float) # =========================================================================== # Nested Conversion Tests (with inner-level conversions) # =========================================================================== # --------------------------------------------------------------------------- # Category 21: Array nested with Optional/Union (inner conversion) # --------------------------------------------------------------------------- class TestNestedArrayComposite: @requires_py39 def test_array_optional_float_with_bool(self) -> None: """Array[Optional[float]] converts bool elements to float.""" result = _to_py_class_value(A(tuple[Optional[float], ...]).convert([True, None, 3])) assert list(result) == [1.0, None, 3.0] assert type(result[0]) is float assert result[1] is None assert type(result[2]) is float @requires_py39 def test_array_optional_int_with_bool(self) -> None: """Array[Optional[int]] converts bool elements to int.""" result = _to_py_class_value(A(tuple[Optional[int], ...]).convert([True, None, 2])) assert list(result) == [1, None, 2] assert type(result[0]) is int assert result[1] is None @requires_py39 def test_array_union_int_str_with_bool(self) -> None: """Array[Union[int, str]] converts bool via int alternative.""" result = _to_py_class_value(A(tuple[Union[int, str], ...]).convert([True, "hello", False])) assert list(result) == [1, "hello", 0] assert type(result[0]) is int assert type(result[1]) is str assert type(result[2]) is int @requires_py39 def test_array_union_float_str_with_int(self) -> None: """Array[Union[float, str]] converts int via float alternative.""" result = _to_py_class_value(A(tuple[Union[float, str], ...]).convert([42, "hi", True])) assert list(result) == [42.0, "hi", 1.0] assert type(result[0]) is float assert type(result[2]) is float @requires_py39 def test_array_optional_float_all_none(self) -> None: """Array[Optional[float]] with all None elements.""" result = _to_py_class_value(A(tuple[Optional[float], ...]).convert([None, None])) assert list(result) == [None, None] @requires_py39 def test_array_optional_float_empty(self) -> None: """Array[Optional[float]] with empty list.""" result = _to_py_class_value(A(tuple[Optional[float], ...]).convert([])) assert list(result) == [] @requires_py39 def test_array_union_failure_in_element(self) -> None: """Array[Union[int, str]] fails when element matches no alternative.""" with pytest.raises(TypeError, match=r"element \[1\].*got float"): A(tuple[Union[int, str], ...]).check_value([1, 3.14]) # --------------------------------------------------------------------------- # Category 22: Map/Dict nested with Optional/Union (inner conversion) # --------------------------------------------------------------------------- class TestNestedMapComposite: @requires_py39 def test_map_str_optional_float_with_int(self) -> None: """Map[str, Optional[float]] converts int values to float.""" result = _to_py_class_value( A(tvm_ffi.Map[str, Optional[float]]).convert({"a": 1, "b": None}) ) assert type(result["a"]) is float assert result["a"] == 1.0 assert result["b"] is None @requires_py39 def test_map_str_union_int_str(self) -> None: """Map[str, Union[int, str]] converts bool values via int.""" result = _to_py_class_value( A(tvm_ffi.Map[str, Union[int, str]]).convert({"x": True, "y": "hello"}) ) assert result["x"] == 1 assert result["y"] == "hello" assert type(result["x"]) is int @requires_py39 def test_dict_str_optional_int(self) -> None: """Dict[str, Optional[int]] with bool conversion.""" result = _to_py_class_value( A(dict[str, Optional[int]]).convert({"a": True, "b": None, "c": 42}) ) assert result["a"] == 1 assert result["b"] is None assert result["c"] == 42 assert type(result["a"]) is int @requires_py39 def test_map_str_optional_float_failure(self) -> None: """Map[str, Optional[float]] fails for non-float non-None value.""" with pytest.raises(TypeError, match="expected float"): A(tvm_ffi.Map[str, Optional[float]]).check_value({"a": "bad"}) # --------------------------------------------------------------------------- # Category 23: Nested containers (container inside container) # --------------------------------------------------------------------------- class TestNestedContainerInContainer: @requires_py39 def test_array_of_array_int(self) -> None: """Array[Array[int]] with inner bool->int conversion.""" result = _to_py_class_value(A(tuple[tuple[int, ...], ...]).convert([[True, False], [1, 2]])) assert [list(row) for row in result] == [[1, 0], [1, 2]] assert all(type(x) is int for row in result for x in row) @requires_py39 def test_array_of_array_float(self) -> None: """Array[Array[float]] with inner int->float conversion.""" result = _to_py_class_value(A(tuple[tuple[float, ...], ...]).convert([[1, 2], [True, 3]])) assert [list(row) for row in result] == [[1.0, 2.0], [1.0, 3.0]] assert all(type(x) is float for row in result for x in row) @requires_py39 def test_map_str_array_float(self) -> None: """Map[str, Array[float]] with int->float conversion in arrays.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[float, ...]]).convert({"a": [1, 2], "b": [True, 3]}) ) assert list(result["a"]) == [1.0, 2.0] assert list(result["b"]) == [1.0, 3.0] assert all(type(x) is float for x in result["a"]) assert all(type(x) is float for x in result["b"]) @requires_py39 def test_dict_str_array_int(self) -> None: """Dict[str, Array[int]] with bool->int conversion.""" result = _to_py_class_value(A(dict[str, tuple[int, ...]]).convert({"a": [True, False]})) assert list(result["a"]) == [1, 0] assert all(type(x) is int for x in result["a"]) @requires_py39 def test_array_of_map_str_int(self) -> None: """Array[Map[str, int]] with bool->int value conversion.""" result = _to_py_class_value( A(tuple[tvm_ffi.Map[str, int], ...]).convert([{"x": True}, {"y": 2}]) ) assert result[0]["x"] == 1 assert result[1]["y"] == 2 assert type(result[0]["x"]) is int @requires_py39 def test_map_str_map_str_float(self) -> None: """Map[str, Map[str, float]] double nested with int->float.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tvm_ffi.Map[str, float]]).convert({"outer": {"inner": 42}}) ) assert result["outer"]["inner"] == 42.0 assert type(result["outer"]["inner"]) is float @requires_py39 def test_list_of_list_int(self) -> None: """List[List[int]] with bool->int conversion.""" result = _to_py_class_value(A(list[list[int]]).convert([[True, 1], [False, 2]])) assert [list(row) for row in result] == [[1, 1], [0, 2]] assert all(type(x) is int for row in result for x in row) @requires_py39 def test_nested_failure_array_of_array(self) -> None: """Array[Array[int]] error propagation through nested arrays.""" with pytest.raises(TypeError, match="expected int"): A(tuple[tuple[int, ...], ...]).check_value([[1, 2], [3, "bad"]]) @requires_py39 def test_empty_inner_containers(self) -> None: """Map[str, Array[int]] with empty inner arrays.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[int, ...]]).convert({"a": [], "b": []}) ) assert list(result["a"]) == [] assert list(result["b"]) == [] @requires_py39 def test_array_of_array_of_array_int(self) -> None: """Three-level nested Array[int] conversion still works.""" schema = A(tuple[tuple[tuple[int, ...], ...], ...]) data = [[[1, 2], [True, False]], [[3], [4, 5, 6]]] result = _to_py_class_value(schema.convert(data)) assert list(result[0][0]) == [1, 2] assert list(result[0][1]) == [1, 0] assert type(result[0][1][0]) is int @requires_py39 def test_map_of_map_of_array_float(self) -> None: """Nested map-to-array conversion still coerces inner values.""" schema = A(tvm_ffi.Map[str, tvm_ffi.Map[str, tuple[float, ...]]]) data = {"outer": {"inner": [1, 2, True]}} result = _to_py_class_value(schema.convert(data)) assert list(result["outer"]["inner"]) == [1.0, 2.0, 1.0] assert type(result["outer"]["inner"][0]) is float # --------------------------------------------------------------------------- # Category 24: Optional/Union wrapping containers # --------------------------------------------------------------------------- class TestOptionalUnionWrappingContainers: @requires_py39 def test_optional_array_int_with_conversion(self) -> None: """Optional[Array[int]] converts inner bool elements.""" schema = A(Optional[tuple[int, ...]]) result = _to_py_class_value(schema.convert([True, 2])) assert list(result) == [1, 2] assert type(result[0]) is int @requires_py39 def test_optional_array_int_none(self) -> None: """Optional[Array[int]] accepts None.""" result = _to_py_class_value(A(Optional[tuple[int, ...]]).convert(None)) assert result is None @requires_py39 def test_optional_map_str_float(self) -> None: """Optional[Map[str, float]] converts inner int values.""" result = _to_py_class_value(A(Optional[tvm_ffi.Map[str, float]]).convert({"a": 1})) assert result["a"] == 1.0 assert type(result["a"]) is float @requires_py39 def test_optional_map_str_float_none(self) -> None: """Optional[Map[str, float]] accepts None.""" result = _to_py_class_value(A(Optional[tvm_ffi.Map[str, float]]).convert(None)) assert result is None @requires_py39 def test_union_array_int_or_map_str_int(self) -> None: """Union[Array[int], Map[str, int]] matches first with conversion.""" schema = A(Union[tuple[int, ...], tvm_ffi.Map[str, int]]) # list matches Array alternative result = _to_py_class_value(schema.convert([True, 2])) assert list(result) == [1, 2] assert type(result[0]) is int @requires_py39 def test_union_array_int_or_map_str_int_dict(self) -> None: """Union[Array[int], Map[str, int]] matches Map for dict input.""" schema = A(Union[tuple[int, ...], tvm_ffi.Map[str, int]]) result = _to_py_class_value(schema.convert({"a": True})) assert result["a"] == 1 assert type(result["a"]) is int @requires_py39 def test_union_int_or_array_optional_float(self) -> None: """Union[int, Array[Optional[float]]] matches array with nested conversions.""" schema = A(Union[int, tuple[Optional[float], ...]]) result = _to_py_class_value(schema.convert([True, None, 1])) assert list(result) == [1.0, None, 1.0] assert type(result[0]) is float assert result[1] is None @requires_py39 def test_optional_optional_array_int(self) -> None: """Optional[Optional[Array[int]]] with inner conversion.""" schema = A(Optional[Optional[tuple[int, ...]]]) assert _to_py_class_value(schema.convert(None)) is None result = _to_py_class_value(schema.convert([True, 2])) assert list(result) == [1, 2] assert type(result[0]) is int # --------------------------------------------------------------------------- # Category 25: Tuple nested with other types # --------------------------------------------------------------------------- class TestNestedTuple: @requires_py39 def test_array_of_tuple_int_float(self) -> None: """Array[tuple[int, float]] with element-wise conversion.""" result = _to_py_class_value( A(tuple[tuple[int, float], ...]).convert([(True, 1), (2, True)]) ) # Check element values; FFI storage may normalize float 1.0 to int 1 # when stored inside an ffi.Array, so we only check values not types. assert result[0][0] == 1 assert result[0][1] == 1.0 assert result[1][0] == 2 assert result[1][1] == 1.0 @requires_py39 def test_map_str_tuple_int_str(self) -> None: """Map[str, tuple[int, str]] with inner bool->int conversion.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[int, str]]).convert({"a": (True, "hello")}) ) assert result["a"][0] == 1 assert str(result["a"][1]) == "hello" assert type(result["a"][0]) is int @requires_py39 def test_tuple_of_array_int_and_map(self) -> None: """tuple[Array[int], Map[str, float]] nested conversion.""" schema = A(tuple[tuple[int, ...], tvm_ffi.Map[str, float]]) result = _to_py_class_value(schema.convert(([True, 2], {"k": 3}))) assert list(result[0]) == [1, 2] assert result[1]["k"] == 3.0 assert type(result[0][0]) is int assert type(result[1]["k"]) is float @requires_py39 def test_tuple_of_optional_int_and_optional_float(self) -> None: """tuple[Optional[int], Optional[float]] with conversions.""" schema = A(tuple[Optional[int], Optional[float]]) result = _to_py_class_value(schema.convert((True, None))) assert list(result) == [1, None] assert type(result[0]) is int assert result[1] is None @requires_py39 def test_tuple_nested_failure(self) -> None: """tuple[Array[int], str] error propagation from inner array.""" with pytest.raises(TypeError, match=r"element .0..*element .1..*expected int"): A(tuple[tuple[int, ...], str]).check_value(([1, "bad"], "ok")) # --------------------------------------------------------------------------- # Category 26: Deep nesting (3+ levels) # --------------------------------------------------------------------------- class TestDeepNesting: @requires_py39 def test_map_str_array_optional_int(self) -> None: """Map[str, Array[Optional[int]]] with 3-level nesting and conversion.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[Optional[int], ...]]).convert({"a": [1, None, True]}) ) assert list(result["a"]) == [1, None, 1] assert type(result["a"][0]) is int assert result["a"][1] is None assert type(result["a"][2]) is int @requires_py39 def test_array_map_str_optional_float(self) -> None: """Array[Map[str, Optional[float]]] with 3-level nesting.""" result = _to_py_class_value( A(tuple[tvm_ffi.Map[str, Optional[float]], ...]).convert( [{"x": 1, "y": None}, {"z": True}] ) ) assert result[0]["x"] == 1.0 assert result[0]["y"] is None assert result[1]["z"] == 1.0 assert type(result[0]["x"]) is float assert type(result[1]["z"]) is float @requires_py39 def test_optional_array_map_str_int(self) -> None: """Optional[Array[Map[str, int]]] 3 levels deep.""" schema = A(Optional[tuple[tvm_ffi.Map[str, int], ...]]) result = _to_py_class_value(schema.convert([{"a": True}, {"b": 2}])) assert result[0]["a"] == 1 assert result[1]["b"] == 2 assert type(result[0]["a"]) is int assert _to_py_class_value(schema.convert(None)) is None @requires_py39 def test_map_str_array_array_int(self) -> None: """Map[str, Array[Array[int]]] 3-level container nesting.""" result = _to_py_class_value( A(tvm_ffi.Map[str, tuple[tuple[int, ...], ...]]).convert({"m": [[True, 1], [False, 2]]}) ) assert [list(row) for row in result["m"]] == [[1, 1], [0, 2]] assert all(type(x) is int for row in result["m"] for x in row) @requires_py39 def test_array_array_optional_float(self) -> None: """Array[Array[Optional[float]]] deep nesting with None and conversion.""" result = _to_py_class_value( A(tuple[tuple[Optional[float], ...], ...]).convert([[1, None], [True, 3.14]]) ) assert list(result[0]) == [1.0, None] assert list(result[1]) == [1.0, 3.14] assert type(result[0][0]) is float assert result[0][1] is None assert type(result[1][0]) is float @requires_py39 def test_deep_nesting_failure_propagation(self) -> None: """Error from deepest level propagates with full path info.""" with pytest.raises(TypeError, match=r"value for key 'key'.*element .1..*expected int"): A(tvm_ffi.Map[str, tuple[Optional[int], ...]]).check_value({"key": [1, "bad"]}) # --------------------------------------------------------------------------- # Category 27: FFI container inputs (tvm_ffi.Array/List/Map/Dict) # --------------------------------------------------------------------------- class TestFFIContainerInputs: @requires_py39 def test_ffi_array_with_element_conversion(self) -> None: """tvm_ffi.Array([True, 2]) passes Array[int] with bool->int conversion.""" arr = tvm_ffi.Array([True, 2, 3]) result = _to_py_class_value(A(tuple[int, ...]).convert(arr)) assert list(result) == [1, 2, 3] assert type(result[0]) is int @requires_py39 def test_ffi_array_any_passthrough(self) -> None: """tvm_ffi.Array passes Array[Any] as-is.""" arr = tvm_ffi.Array([1, "x", None]) result = _to_py_class_value(A(tuple[typing.Any, ...]).convert(arr)) assert result.same_as(arr) @requires_py39 def test_ffi_list_with_list_schema(self) -> None: """tvm_ffi.List passes List[int] with conversion.""" lst = tvm_ffi.List([True, 2]) result = _to_py_class_value(A(list[int]).convert(lst)) assert list(result) == [1, 2] assert type(result[0]) is int @requires_py39 def test_ffi_list_accepted_by_array_schema(self) -> None: """tvm_ffi.List passes Array schema (C++ allows cross-type via kOtherTypeIndex).""" lst = tvm_ffi.List([1, 2]) A(tuple[int, ...]).check_value(lst) @requires_py39 def test_ffi_array_accepted_by_list_schema(self) -> None: """tvm_ffi.Array passes List schema (C++ allows cross-type via kOtherTypeIndex).""" arr = tvm_ffi.Array([1, 2]) A(list[int]).check_value(arr) @requires_py39 def test_ffi_map_with_value_conversion(self) -> None: """tvm_ffi.Map passes Map[str, int] with bool->int conversion.""" m = tvm_ffi.Map({"a": True, "b": 2}) result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert(m)) assert result["a"] == 1 assert result["b"] == 2 assert type(result["a"]) is int @requires_py39 def test_ffi_map_any_any_passthrough(self) -> None: """tvm_ffi.Map passes Map[Any, Any] as-is.""" m = tvm_ffi.Map({"a": 1}) result = _to_py_class_value(A(tvm_ffi.Map[typing.Any, typing.Any]).convert(m)) assert result.same_as(m) @requires_py39 def test_ffi_dict_with_dict_schema(self) -> None: """tvm_ffi.Dict passes Dict[str, float] with int->float conversion.""" d = tvm_ffi.Dict({"x": 1, "y": 2}) result = _to_py_class_value(A(dict[str, float]).convert(d)) assert result["x"] == 1.0 assert result["y"] == 2.0 assert type(result["x"]) is float @requires_py39 def test_ffi_dict_accepted_by_map_schema(self) -> None: """tvm_ffi.Dict passes Map schema (C++ allows cross-type via kOtherTypeIndex).""" d = tvm_ffi.Dict({"a": 1}) A(tvm_ffi.Map[str, int]).check_value(d) @requires_py39 def test_ffi_map_accepted_by_dict_schema(self) -> None: """tvm_ffi.Map passes Dict schema (C++ allows cross-type via kOtherTypeIndex).""" m = tvm_ffi.Map({"a": 1}) A(dict[str, int]).check_value(m) @requires_py39 def test_ffi_array_nested_optional_float(self) -> None: """tvm_ffi.Array with nested Optional[float] conversion.""" arr = tvm_ffi.Array([1, None, True]) result = _to_py_class_value(A(tuple[Optional[float], ...]).convert(arr)) assert list(result) == [1.0, None, 1.0] assert type(result[0]) is float assert result[1] is None @requires_py39 def test_ffi_map_nested_array_int(self) -> None: """tvm_ffi.Map with value being a Python list, converted as Array[int].""" # Map values are already stored; create a map with array values m = tvm_ffi.Map({"k": tvm_ffi.Array([True, 2])}) result = _to_py_class_value(A(tvm_ffi.Map[str, tuple[int, ...]]).convert(m)) assert list(result["k"]) == [1, 2] assert type(result["k"][0]) is int @requires_py39 def test_ffi_array_wrong_element_type(self) -> None: """tvm_ffi.Array with wrong element type gives clear error.""" arr = tvm_ffi.Array([1, "bad", 3]) with pytest.raises(TypeError, match=r"element \[1\].*expected int"): A(tuple[int, ...]).check_value(arr) @requires_py39 def test_ffi_map_wrong_value_type(self) -> None: """tvm_ffi.Map with wrong value type gives clear error.""" m = tvm_ffi.Map({"a": 1, "b": "bad"}) with pytest.raises(TypeError, match=r"value for key.*expected int"): A(tvm_ffi.Map[str, int]).check_value(m) def test_ffi_array_object_schema(self) -> None: """tvm_ffi.Array passes Object schema (it is a CObject).""" arr = tvm_ffi.Array([1, 2]) A(tvm_ffi.core.Object).check_value(arr) def test_ffi_map_object_schema(self) -> None: """tvm_ffi.Map passes Object schema (it is a CObject).""" m = tvm_ffi.Map({"a": 1}) A(tvm_ffi.core.Object).check_value(m) # --------------------------------------------------------------------------- # Category 28: Mixed Python and FFI containers in nesting # --------------------------------------------------------------------------- class TestMixedPythonFFIContainers: @requires_py39 def test_python_list_of_ffi_arrays(self) -> None: """Python list containing tvm_ffi.Array elements, Array[Array[int]].""" inner1 = tvm_ffi.Array([True, 2]) inner2 = tvm_ffi.Array([3, False]) result = _to_py_class_value(A(tuple[tuple[int, ...], ...]).convert([inner1, inner2])) assert [list(row) for row in result] == [[1, 2], [3, 0]] @requires_py39 def test_python_dict_with_ffi_array_values(self) -> None: """Python dict with tvm_ffi.Array values, Map[str, Array[float]].""" val = tvm_ffi.Array([1, True]) result = _to_py_class_value(A(tvm_ffi.Map[str, tuple[float, ...]]).convert({"k": val})) assert list(result["k"]) == [1.0, 1.0] assert all(type(x) is float for x in result["k"]) @requires_py39 def test_ffi_map_with_python_list_in_union(self) -> None: """Union[Map[str, int], Array[int]] with tvm_ffi.Map input.""" schema = A(Union[tvm_ffi.Map[str, int], tuple[int, ...]]) m = tvm_ffi.Map({"a": True}) result = _to_py_class_value(schema.convert(m)) assert result["a"] == 1 assert type(result["a"]) is int @requires_py39 def test_ffi_array_in_optional(self) -> None: """Optional[Array[int]] with tvm_ffi.Array input.""" arr = tvm_ffi.Array([True, 2]) result = _to_py_class_value(A(Optional[tuple[int, ...]]).convert(arr)) assert list(result) == [1, 2] assert type(result[0]) is int # --------------------------------------------------------------------------- # Category 29: Error propagation through deeply nested FFI containers # --------------------------------------------------------------------------- class TestNestedErrorPropagation: @requires_py39 def test_array_array_int_inner_failure(self) -> None: """Error path: Array[Array[int]] -> element [1] -> element [0].""" with pytest.raises(TypeError, match=r"element \[1\].*element \[0\].*expected int, got str"): A(tuple[tuple[int, ...], ...]).convert([[1], ["bad"]]) @requires_py39 def test_map_array_int_inner_failure(self) -> None: """Error path: Map -> value for key 'k' -> element [2].""" with pytest.raises( TypeError, match=r"value for key 'k'.*element \[2\].*expected int, got str", ): A(tvm_ffi.Map[str, tuple[int, ...]]).convert({"k": [1, 2, "bad"]}) @requires_py39 def test_array_map_int_inner_failure(self) -> None: """Error path: Array -> element [0] -> value for key 'x'.""" with pytest.raises( TypeError, match=r"element \[0\].*value for key 'x'.*expected int, got str", ): A(tuple[tvm_ffi.Map[str, int], ...]).convert([{"x": "bad"}]) @requires_py39 def test_optional_array_int_inner_failure(self) -> None: """Error path through Optional -> Array -> element.""" with pytest.raises(TypeError, match=r"element \[1\].*expected int, got str"): A(Optional[tuple[int, ...]]).convert([1, "bad"]) @requires_py39 def test_tuple_array_int_inner_failure(self) -> None: """Error path: tuple -> element [0] -> element [1].""" with pytest.raises(TypeError, match=r"element \[0\].*element \[1\].*expected int, got str"): A(tuple[tuple[int, ...], str]).convert(([1, "bad"], "ok")) @requires_py39 def test_deep_3_level_error(self) -> None: """Error at 3 levels deep: Map -> Array -> Optional -> type mismatch.""" with pytest.raises(TypeError, match=r"value for key 'key'.*element .1..*expected int"): A(tvm_ffi.Map[str, tuple[Optional[int], ...]]).check_value({"key": [1, "bad"]}) @requires_py39 def test_ffi_array_nested_error(self) -> None: """Error from tvm_ffi.Array in nested context.""" arr = tvm_ffi.Array([1, "bad", 3]) with pytest.raises(TypeError, match=r"element \[1\].*expected int"): A(tuple[int, ...]).convert(arr) # --------------------------------------------------------------------------- # Category 30: Custom object type exact match # --------------------------------------------------------------------------- class TestCustomObjectExactMatch: def test_test_int_pair_pass(self) -> None: """TestIntPair passes TypeSchema('testing.TestIntPair').""" obj = TestIntPair(1, 2) A(TestIntPair).check_value(obj) def test_test_object_base_pass(self) -> None: """TestObjectBase passes its own schema.""" obj = TestObjectBase(v_i64=10, v_f64=1.5, v_str="hi") A(TestObjectBase).check_value(obj) def test_test_object_derived_pass(self) -> None: """TestObjectDerived passes its own schema.""" obj = TestObjectDerived(v_map={"a": 1}, v_array=[1], v_i64=0, v_f64=0.0, v_str="") A(TestObjectDerived).check_value(obj) def test_cxx_class_base_pass(self) -> None: """_TestCxxClassBase passes its own schema.""" obj = _TestCxxClassBase(v_i64=1, v_i32=2) A(_TestCxxClassBase).check_value(obj) def test_cxx_class_derived_pass(self) -> None: """_TestCxxClassDerived passes its own schema.""" obj = _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0) A(_TestCxxClassDerived).check_value(obj) def test_cxx_class_derived_derived_pass(self) -> None: """_TestCxxClassDerivedDerived passes its own schema.""" obj = _TestCxxClassDerivedDerived(v_i64=1, v_i32=2, v_f64=3.0, v_bool=True) A(_TestCxxClassDerivedDerived).check_value(obj) # --------------------------------------------------------------------------- # Category 31: Custom object type hierarchy (subclass passes parent schema) # --------------------------------------------------------------------------- class TestCustomObjectHierarchy: def test_derived_passes_base_schema(self) -> None: """TestObjectDerived passes TypeSchema('testing.TestObjectBase').""" obj = TestObjectDerived(v_map={"a": 1}, v_array=[1], v_i64=0, v_f64=0.0, v_str="") A(TestObjectBase).check_value(obj) def test_derived_passes_object_schema(self) -> None: """TestObjectDerived passes TypeSchema('Object').""" obj = TestObjectDerived(v_map={"a": 1}, v_array=[1], v_i64=0, v_f64=0.0, v_str="") A(tvm_ffi.core.Object).check_value(obj) def test_cxx_derived_passes_base(self) -> None: """_TestCxxClassDerived passes TestCxxClassBase schema.""" obj = _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0) A(_TestCxxClassBase).check_value(obj) def test_cxx_derived_derived_passes_base(self) -> None: """_TestCxxClassDerivedDerived passes TestCxxClassBase schema (2-level up).""" obj = _TestCxxClassDerivedDerived(v_i64=1, v_i32=2, v_f64=3.0, v_bool=True) A(_TestCxxClassBase).check_value(obj) def test_cxx_derived_derived_passes_derived(self) -> None: """_TestCxxClassDerivedDerived passes TestCxxClassDerived schema (1-level up).""" obj = _TestCxxClassDerivedDerived(v_i64=1, v_i32=2, v_f64=3.0, v_bool=True) A(_TestCxxClassDerived).check_value(obj) def test_all_custom_objects_pass_object_schema(self) -> None: """Every custom object passes the generic Object schema.""" objs = [ TestIntPair(1, 2), TestObjectBase(v_i64=10, v_f64=1.5, v_str="hi"), _TestCxxClassBase(v_i64=1, v_i32=2), _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0), _TestCxxClassDerivedDerived(v_i64=1, v_i32=2, v_f64=3.0, v_bool=True), ] schema = A(tvm_ffi.core.Object) for obj in objs: schema.check_value(obj) # --------------------------------------------------------------------------- # Category 32: Custom object type rejection # --------------------------------------------------------------------------- class TestCustomObjectRejection: def test_wrong_object_type(self) -> None: """TestIntPair fails TypeSchema('testing.TestObjectBase').""" obj = TestIntPair(1, 2) with pytest.raises(TypeError, match=r"testing.TestIntPair"): A(TestObjectBase).check_value(obj) def test_base_fails_derived_schema(self) -> None: """Parent object fails child schema (TestObjectBase fails TestObjectDerived).""" obj = TestObjectBase(v_i64=10, v_f64=1.5, v_str="hi") with pytest.raises(TypeError, match=r"testing.TestObjectBase"): A(TestObjectDerived).check_value(obj) def test_non_object_fails_custom_schema(self) -> None: """Plain int fails custom object schema.""" with pytest.raises(TypeError, match=r"expected testing\.TestIntPair.*got int"): A(TestIntPair).check_value(42) def test_none_fails_custom_schema(self) -> None: """None fails custom object schema.""" with pytest.raises(TypeError, match="got None"): A(TestIntPair).check_value(None) def test_string_fails_custom_schema(self) -> None: """String fails custom object schema.""" with pytest.raises(TypeError, match="got str"): A(TestIntPair).check_value("hello") def test_cxx_base_fails_derived_schema(self) -> None: """_TestCxxClassBase fails _TestCxxClassDerived schema.""" obj = _TestCxxClassBase(v_i64=1, v_i32=2) with pytest.raises(TypeError): A(_TestCxxClassDerived).check_value(obj) def test_sibling_types_reject_each_other(self) -> None: """TestIntPair and TestCxxClassBase are unrelated -- reject each other.""" pair = TestIntPair(1, 2) base = _TestCxxClassBase(v_i64=1, v_i32=2) with pytest.raises(TypeError): A(_TestCxxClassBase).check_value(pair) with pytest.raises(TypeError): A(TestIntPair).check_value(base) # --------------------------------------------------------------------------- # Category 33: Custom objects in containers # --------------------------------------------------------------------------- class TestCustomObjectInContainers: @requires_py39 def test_array_of_custom_objects(self) -> None: """Array[testing.TestIntPair] with matching elements.""" objs = [TestIntPair(1, 2), TestIntPair(3, 4)] A(tuple[TestIntPair, ...]).check_value(objs) @requires_py39 def test_array_of_custom_objects_wrong_type(self) -> None: """Array[testing.TestIntPair] with wrong element type fails.""" objs = [TestIntPair(1, 2), _TestCxxClassBase(v_i64=1, v_i32=2)] with pytest.raises(TypeError, match=r"element \[1\]"): A(tuple[TestIntPair, ...]).check_value(objs) @requires_py39 def test_array_of_base_with_derived_elements(self) -> None: """Array[testing.TestObjectBase] accepts derived elements via hierarchy.""" base = TestObjectBase(v_i64=1, v_f64=1.0, v_str="a") derived = TestObjectDerived(v_map={"a": 1}, v_array=[1], v_i64=0, v_f64=0.0, v_str="") A(tuple[TestObjectBase, ...]).check_value([base, derived]) @requires_py39 def test_map_str_to_custom_object(self) -> None: """Map[str, testing.TestIntPair] pass.""" objs = {"a": TestIntPair(1, 2), "b": TestIntPair(3, 4)} A(tvm_ffi.Map[str, TestIntPair]).check_value(objs) @requires_py39 def test_map_str_to_custom_object_wrong_value(self) -> None: """Map[str, testing.TestIntPair] with int value fails.""" data = {"a": TestIntPair(1, 2), "b": 42} with pytest.raises(TypeError, match="value for key 'b'"): A(tvm_ffi.Map[str, TestIntPair]).check_value(data) @requires_py39 def test_ffi_array_of_custom_objects(self) -> None: """tvm_ffi.Array of custom objects passes Array[Object].""" arr = tvm_ffi.Array([TestIntPair(1, 2), TestObjectBase(v_i64=1, v_f64=2.0, v_str="s")]) A(tuple[tvm_ffi.core.Object, ...]).check_value(arr) @requires_py39 def test_ffi_array_of_custom_objects_specific_type(self) -> None: """tvm_ffi.Array of TestIntPair passes Array[testing.TestIntPair].""" arr = tvm_ffi.Array([TestIntPair(1, 2), TestIntPair(3, 4)]) A(tuple[TestIntPair, ...]).check_value(arr) @requires_py39 def test_ffi_map_with_custom_object_values(self) -> None: """tvm_ffi.Map with custom object values passes.""" m = tvm_ffi.Map({"x": TestIntPair(1, 2), "y": TestIntPair(3, 4)}) A(tvm_ffi.Map[str, TestIntPair]).check_value(m) # --------------------------------------------------------------------------- # Category 34: Optional/Union with custom objects # --------------------------------------------------------------------------- class TestCustomObjectOptionalUnion: def test_optional_custom_object_with_value(self) -> None: """Optional[testing.TestIntPair] with actual object.""" obj = TestIntPair(1, 2) A(Optional[TestIntPair]).check_value(obj) def test_optional_custom_object_with_none(self) -> None: """Optional[testing.TestIntPair] with None.""" A(Optional[TestIntPair]).check_value(None) def test_optional_custom_object_wrong_type(self) -> None: """Optional[testing.TestIntPair] with wrong object type.""" obj = _TestCxxClassBase(v_i64=1, v_i32=2) with pytest.raises(TypeError): A(Optional[TestIntPair]).check_value(obj) def test_union_custom_object_and_int(self) -> None: """Union[testing.TestIntPair, int] with object.""" obj = TestIntPair(1, 2) A(Union[TestIntPair, int]).check_value(obj) def test_union_custom_object_and_int_with_int(self) -> None: """Union[testing.TestIntPair, int] with int.""" A(Union[TestIntPair, int]).check_value(42) def test_union_custom_object_and_int_with_wrong(self) -> None: """Union[testing.TestIntPair, int] with str fails.""" with pytest.raises(TypeError): A(Union[TestIntPair, int]).check_value("bad") def test_union_two_custom_objects(self) -> None: """Union of two custom types accepts both.""" pair = TestIntPair(1, 2) base = _TestCxxClassBase(v_i64=1, v_i32=2) schema = A(Union[TestIntPair, _TestCxxClassBase]) schema.check_value(pair) schema.check_value(base) def test_union_two_custom_objects_rejects_third(self) -> None: """Union of two custom types rejects a third.""" obj = TestObjectBase(v_i64=1, v_f64=2.0, v_str="s") with pytest.raises(TypeError): A(Union[TestIntPair, _TestCxxClassBase]).check_value(obj) # --------------------------------------------------------------------------- # Category 35: Custom objects with from_type_index # --------------------------------------------------------------------------- class TestCustomObjectFromTypeIndex: def test_from_type_index_custom_object(self) -> None: """from_type_index resolves a custom object type and validates.""" obj = TestIntPair(1, 2) tindex = tvm_ffi.core._object_type_key_to_index("testing.TestIntPair") assert tindex is not None schema = TypeSchema.from_type_index(tindex) assert schema.origin == "testing.TestIntPair" schema.check_value(obj) def test_from_type_index_rejects_wrong_object(self) -> None: """from_type_index schema rejects wrong object type.""" tindex = tvm_ffi.core._object_type_key_to_index("testing.TestIntPair") assert tindex is not None schema = TypeSchema.from_type_index(tindex) with pytest.raises(TypeError): schema.check_value(_TestCxxClassBase(v_i64=1, v_i32=2)) def test_from_type_index_hierarchy(self) -> None: """from_type_index for base type accepts derived objects.""" tindex = tvm_ffi.core._object_type_key_to_index("testing.TestObjectBase") assert tindex is not None schema = TypeSchema.from_type_index(tindex) derived = TestObjectDerived(v_map={"a": 1}, v_array=[1], v_i64=0, v_f64=0.0, v_str="") schema.check_value(derived) # --------------------------------------------------------------------------- # Category 36: Custom objects in nested containers # --------------------------------------------------------------------------- class TestCustomObjectNestedContainers: @requires_py39 def test_array_of_optional_custom_object(self) -> None: """Array[Optional[testing.TestIntPair]] with mix of objects and None.""" data = [TestIntPair(1, 2), None, TestIntPair(3, 4)] A(tuple[Optional[TestIntPair], ...]).check_value(data) @requires_py39 def test_map_str_to_array_of_custom_objects(self) -> None: """Map[str, Array[testing.TestIntPair]] with nested objects.""" data = { "group1": [TestIntPair(1, 2), TestIntPair(3, 4)], "group2": [TestIntPair(5, 6)], } A(tvm_ffi.Map[str, tuple[TestIntPair, ...]]).check_value(data) @requires_py39 def test_array_of_union_custom_objects(self) -> None: """Array[Union[testing.TestIntPair, testing.TestCxxClassBase]].""" data = [TestIntPair(1, 2), _TestCxxClassBase(v_i64=1, v_i32=2), TestIntPair(5, 6)] A(tuple[Union[TestIntPair, _TestCxxClassBase], ...]).check_value(data) @requires_py39 def test_optional_array_of_custom_objects(self) -> None: """Optional[Array[testing.TestIntPair]] with array.""" data = [TestIntPair(1, 2)] A(Optional[tuple[TestIntPair, ...]]).check_value(data) @requires_py39 def test_optional_array_of_custom_objects_none(self) -> None: """Optional[Array[testing.TestIntPair]] with None.""" A(Optional[tuple[TestIntPair, ...]]).check_value(None) @requires_py39 def test_nested_error_with_custom_object(self) -> None: """Array[testing.TestIntPair] error message includes type keys.""" data = [TestIntPair(1, 2), _TestCxxClassBase(v_i64=1, v_i32=2)] with pytest.raises( TypeError, match=r"element \[1\].*testing.TestIntPair.*testing.TestCxxClassBase" ): A(tuple[TestIntPair, ...]).check_value(data) @requires_py39 def test_map_nested_error_with_custom_object(self) -> None: """Map value error for custom object includes key and type info.""" data = {"ok": TestIntPair(1, 2), "bad": 42} with pytest.raises( TypeError, match=r"value for key 'bad'.*expected testing\.TestIntPair.*got int" ): A(tvm_ffi.Map[str, TestIntPair]).check_value(data) @requires_py39 def test_deep_nested_custom_objects(self) -> None: """Map[str, Array[Optional[testing.TestIntPair]]] deep nesting.""" data = { "a": [TestIntPair(1, 2), None], "b": [None, TestIntPair(3, 4), TestIntPair(5, 6)], } A(tvm_ffi.Map[str, tuple[Optional[TestIntPair], ...]]).check_value(data) @requires_py39 def test_deep_nested_custom_objects_error(self) -> None: """Map[str, Array[testing.TestIntPair]] error at 3 levels.""" data = {"k": [TestIntPair(1, 2), "bad"]} with pytest.raises(TypeError, match=r"value for key 'k'.*element .1."): A(tvm_ffi.Map[str, tuple[TestIntPair, ...]]).check_value(data) @requires_py39 def test_tuple_with_custom_object(self) -> None: """tuple[testing.TestIntPair, int, str] with custom object.""" data = (TestIntPair(1, 2), 42, "hello") A(tuple[TestIntPair, int, str]).check_value(data) @requires_py39 def test_tuple_with_custom_object_wrong(self) -> None: """tuple[testing.TestIntPair, int] with wrong object in first position.""" data = (_TestCxxClassBase(v_i64=1, v_i32=2), 42) with pytest.raises(TypeError, match=r"element \[0\]"): A(tuple[TestIntPair, int]).check_value(data) # --------------------------------------------------------------------------- # Category 37: Lowercase Python-native origins ("list", "dict") # --------------------------------------------------------------------------- class TestLowercaseOrigins: def test_list_origin_accepts_python_list(self) -> None: """TypeSchema("list", ...) should validate elements, not passthrough.""" S("list", S("int")).check_value([1, 2, 3]) # S: lowercase "list" is an internal origin def test_list_origin_rejects_bad_elements(self) -> None: """TypeSchema("list", (int,)).check_value(["x"]) should fail.""" with pytest.raises(TypeError, match=r"element \[0\]"): S("list", S("int")).check_value(["x"]) # S: lowercase "list" is an internal origin def test_list_origin_converts_elements(self) -> None: """TypeSchema("list", (float,)).convert([1, True]) does int->float.""" # S: lowercase "list" is an internal origin result = _to_py_class_value(S("list", S("float")).convert([1, True])) assert isinstance(result, tvm_ffi.List) assert list(result) == [1.0, 1.0] assert all(type(x) is float for x in result) def test_dict_origin_accepts_python_dict(self) -> None: """TypeSchema("dict", ...) should validate key/value types.""" S("dict", S("str"), S("int")).check_value( {"a": 1} ) # S: lowercase "dict" is an internal origin def test_dict_origin_rejects_bad_values(self) -> None: """TypeSchema("dict", (str, int)).check_value({"a": "x"}) should fail.""" with pytest.raises(TypeError, match="value for key 'a'"): S("dict", S("str"), S("int")).check_value( {"a": "x"} ) # S: lowercase "dict" is an internal origin def test_dict_origin_converts_values(self) -> None: """TypeSchema("dict", (str, float)).convert({"a": 1}) does int->float.""" # S: lowercase "dict" is an internal origin result = _to_py_class_value(S("dict", S("str"), S("float")).convert({"a": 1, "b": True})) assert isinstance(result, tvm_ffi.Dict) assert result["a"] == 1.0 assert result["b"] == 1.0 assert all(type(v) is float for v in result.values()) def test_list_origin_no_args_accepts_anything(self) -> None: """TypeSchema("list") with no args accepts any list (element type is Any).""" S("list").check_value([1, "a", None]) # S: lowercase "list" is an internal origin def test_dict_origin_no_args_accepts_anything(self) -> None: """TypeSchema("dict") with no args accepts any dict.""" S("dict").check_value({"a": 1, 2: "b"}) # S: lowercase "dict" is an internal origin def test_list_origin_rejects_non_list(self) -> None: """TypeSchema("list") rejects non-sequence types.""" with pytest.raises(TypeError, match="got int"): S("list").check_value(42) # S: lowercase "list" is an internal origin def test_dict_origin_rejects_non_dict(self) -> None: """TypeSchema("dict") rejects non-dict types.""" with pytest.raises(TypeError): S("dict").check_value([1, 2]) # S: lowercase "dict" is an internal origin # --------------------------------------------------------------------------- # Category 38: Cross-type container conversions (Array<->List, Map<->Dict) # --------------------------------------------------------------------------- class TestCrossTypeContainers: @requires_py39 def test_array_schema_accepts_ffi_list(self) -> None: """Array[int] schema accepts tvm_ffi.List (C++ kOtherTypeIndex).""" lst = tvm_ffi.List([1, 2, 3]) A(tuple[int, ...]).check_value(lst) @requires_py39 def test_list_schema_accepts_ffi_array(self) -> None: """List[int] schema accepts tvm_ffi.Array (C++ kOtherTypeIndex).""" arr = tvm_ffi.Array([1, 2, 3]) A(list[int]).check_value(arr) @requires_py39 def test_map_schema_accepts_ffi_dict(self) -> None: """Map[str, int] schema accepts tvm_ffi.Dict (C++ kOtherTypeIndex).""" d = tvm_ffi.Dict({"a": 1, "b": 2}) A(tvm_ffi.Map[str, int]).check_value(d) @requires_py39 def test_dict_schema_accepts_ffi_map(self) -> None: """Dict[str, int] schema accepts tvm_ffi.Map (C++ kOtherTypeIndex).""" m = tvm_ffi.Map({"a": 1, "b": 2}) A(dict[str, int]).check_value(m) @requires_py39 def test_array_schema_converts_list_elements(self) -> None: """Array[float] converts elements from tvm_ffi.List[int].""" lst = tvm_ffi.List([1, 2, True]) result = _to_py_class_value(A(tuple[float, ...]).convert(lst)) assert list(result) == [1.0, 2.0, 1.0] assert all(type(x) is float for x in result) @requires_py39 def test_list_schema_converts_array_elements(self) -> None: """List[float] converts elements from tvm_ffi.Array[int].""" arr = tvm_ffi.Array([1, 2, True]) result = _to_py_class_value(A(list[float]).convert(arr)) assert list(result) == [1.0, 2.0, 1.0] assert all(type(x) is float for x in result) @requires_py39 def test_map_schema_converts_dict_values(self) -> None: """Map[str, float] converts values from tvm_ffi.Dict.""" d = tvm_ffi.Dict({"a": 1, "b": True}) result = _to_py_class_value(A(tvm_ffi.Map[str, float]).convert(d)) assert result["a"] == 1.0 assert result["b"] == 1.0 @requires_py39 def test_dict_schema_converts_map_values(self) -> None: """Dict[str, float] converts values from tvm_ffi.Map.""" m = tvm_ffi.Map({"a": 1, "b": True}) result = _to_py_class_value(A(dict[str, float]).convert(m)) assert result["a"] == 1.0 assert result["b"] == 1.0 @requires_py39 def test_cross_type_still_rejects_wrong_container(self) -> None: """Array schema still rejects non-sequence CObjects (e.g. Map).""" m = tvm_ffi.Map({"a": 1}) with pytest.raises(TypeError, match="expected Array"): A(tuple[int, ...]).check_value(m) @requires_py39 def test_cross_type_map_rejects_array(self) -> None: """Map schema still rejects sequence CObjects (e.g. Array).""" arr = tvm_ffi.Array([1, 2]) with pytest.raises(TypeError, match="expected Map"): A(tvm_ffi.Map[str, int]).check_value(arr) # --------------------------------------------------------------------------- # Category 39: tuple accepts list and CObject Array # --------------------------------------------------------------------------- class TestTupleAcceptsListAndArray: @requires_py39 def test_tuple_accepts_python_list(self) -> None: """tuple[int, str] accepts Python list input.""" result = _to_py_class_value(A(tuple[int, str]).convert([42, "hello"])) assert list(result) == [42, "hello"] @requires_py39 def test_tuple_list_with_conversion(self) -> None: """tuple[float, int] converts list elements (bool->float, bool->int).""" result = _to_py_class_value(A(tuple[float, int]).convert([True, False])) assert list(result) == [1.0, 0] assert type(result[0]) is float assert type(result[1]) is int @requires_py39 def test_tuple_rejects_wrong_length_list(self) -> None: """tuple[int, str] rejects list of wrong length.""" with pytest.raises(TypeError, match="length"): A(tuple[int, str]).check_value([1, "a", "b"]) @requires_py39 def test_tuple_accepts_ffi_array(self) -> None: """tuple[int, int] accepts tvm_ffi.Array (C++ Tuple accepts kTVMFFIArray).""" arr = tvm_ffi.Array([1, 2]) A(tuple[int, int]).check_value(arr) @requires_py39 def test_tuple_ffi_array_with_conversion(self) -> None: """tuple[float, float] converts tvm_ffi.Array elements.""" arr = tvm_ffi.Array([1, True]) result = _to_py_class_value(A(tuple[float, float]).convert(arr)) assert list(result) == [1.0, 1.0] assert all(type(x) is float for x in result) @requires_py39 def test_tuple_ffi_array_wrong_length(self) -> None: """tuple[int, int] rejects tvm_ffi.Array of wrong length.""" arr = tvm_ffi.Array([1, 2, 3]) with pytest.raises(TypeError, match="length"): A(tuple[int, int]).check_value(arr) @requires_py39 def test_tuple_rejects_ffi_map(self) -> None: """Tuple schema rejects Map CObject.""" m = tvm_ffi.Map({"a": 1}) with pytest.raises(TypeError, match="expected tuple"): A(tuple[int]).check_value(m) def test_untyped_tuple_accepts_list(self) -> None: """Tuple (no args) accepts any list as-is.""" # Untyped tuple has tuple_len=0, so it just checks the container type # but doesn't validate elements A(tuple).check_value([1, "a", None]) def test_untyped_tuple_accepts_ffi_array(self) -> None: """Tuple (no args) accepts tvm_ffi.Array as-is.""" arr = tvm_ffi.Array([1, 2, 3]) A(tuple).check_value(arr) def test_typed_empty_tuple_rejects_non_empty_list(self) -> None: """Explicit empty tuple schema enforces length 0.""" schema = TypeSchema("tuple", ()) with pytest.raises(TypeError, match="length 0"): schema.check_value([1]) def test_untyped_tuple_converts_ffi_list_to_array(self) -> None: """Tuple (no args) normalizes tvm_ffi.List input to tvm_ffi.Array.""" lst = tvm_ffi.List([1, 2, 3]) result = _to_py_class_value(A(tuple).convert(lst)) assert isinstance(result, tvm_ffi.Array) assert list(result) == [1, 2, 3] assert not result.same_as(lst) # --------------------------------------------------------------------------- # Category 40: dtype string parse errors # --------------------------------------------------------------------------- class TestDtypeParseErrors: def test_check_value_bad_dtype_raises_error(self) -> None: """check_value should raise TypeError for invalid dtype.""" with pytest.raises(TypeError, match="dtype"): A(tvm_ffi.core.DataType).check_value("not_a_valid_dtype_xyz") def test_convert_bad_dtype_raises_type_error_2(self) -> None: """Convert should raise TypeError for invalid dtype string.""" with pytest.raises(TypeError, match="dtype"): A(tvm_ffi.core.DataType).convert("not_a_valid_dtype_xyz") def test_convert_bad_dtype_raises_type_error(self) -> None: """Convert should raise TypeError for invalid dtype string.""" with pytest.raises(TypeError, match="dtype"): A(tvm_ffi.core.DataType).convert("not_a_valid_dtype_xyz") def test_valid_dtype_string_still_works(self) -> None: """Valid dtype strings should still convert successfully.""" result = _to_py_class_value(A(tvm_ffi.core.DataType).convert("float32")) assert str(result) == "float32" def test_convert_valid_dtype(self) -> None: """Convert with valid dtype returns DataType.""" result = _to_py_class_value(A(tvm_ffi.core.DataType).convert("int8")) assert str(result) == "int8" # --------------------------------------------------------------------------- # Category 41: int64 boundary checking # --------------------------------------------------------------------------- class TestInt64Boundaries: """Verify int converter rejects values outside int64 range. The FFI marshals Python int to C++ int64_t. Values outside [-2^63, 2^63-1] would silently overflow at marshal time, so the converter rejects them early. """ def test_int64_max_accepted(self) -> None: """2^63-1 (INT64_MAX) is the largest valid int.""" A(int).check_value(2**63 - 1) def test_int64_min_accepted(self) -> None: """-2^63 (INT64_MIN) is the smallest valid int.""" A(int).check_value(-(2**63)) def test_int64_max_plus_one_rejected(self) -> None: """2^63 exceeds int64 range.""" with pytest.raises(TypeError, match="int64 range"): A(int).check_value(2**63) def test_int64_min_minus_one_rejected(self) -> None: """-2^63-1 exceeds int64 range.""" with pytest.raises(TypeError, match="int64 range"): A(int).check_value(-(2**63) - 1) def test_very_large_positive_rejected(self) -> None: """Very large positive integer rejected.""" with pytest.raises(TypeError, match="int64 range"): A(int).check_value(10**100) def test_very_large_negative_rejected(self) -> None: """Very large negative integer rejected.""" with pytest.raises(TypeError, match="int64 range"): A(int).check_value(-(10**100)) def test_convert_raises_type_error_for_overflow(self) -> None: """Convert raises TypeError for overflow.""" with pytest.raises(TypeError, match="int64 range"): A(int).convert(2**63) def test_bool_to_int_no_range_issue(self) -> None: """Bool -> int conversion (0 or 1) always fits.""" assert _to_py_class_value(A(int).convert(True)) == 1 assert _to_py_class_value(A(int).convert(False)) == 0 def test_int64_boundaries_in_float_conversion(self) -> None: """Float schema accepts large ints (float64 has wider range).""" # float64 can represent integers up to 2^53 exactly, # and larger values with precision loss (but no range error) A(float).check_value(2**63) A(float).check_value(-(2**63)) def test_int64_overflow_in_optional_int(self) -> None: """Optional[int] propagates int64 range check.""" with pytest.raises(TypeError, match="int64 range"): A(Optional[int]).check_value(2**63) @requires_py39 def test_int64_overflow_in_array_element(self) -> None: """Array[int] element overflow is caught with path.""" with pytest.raises(TypeError, match="int64 range"): A(tuple[int, ...]).check_value([1, 2**63, 3]) # --------------------------------------------------------------------------- # Category 42: Unknown origin errors (lazy converter construction) # --------------------------------------------------------------------------- class TestUnknownOriginErrors: """Converter is built lazily via cached_property. Unknown origins construct fine but raise TypeError on first convert/check_value. """ def test_unknown_origin_constructs_ok(self) -> None: """TypeSchema with unknown origin can be constructed.""" schema = S("not_a_real_type") # S: intentionally invalid origin assert schema.origin == "not_a_real_type" def test_unknown_origin_errors_on_check_value(self) -> None: """Unknown origin raises TypeError on check_value.""" schema = S("not_a_real_type") # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.check_value(42) def test_unknown_origin_errors_on_convert(self) -> None: """Unknown origin raises TypeError on convert.""" schema = S("not_a_real_type") # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.convert(42) def test_unknown_origin_errors_on_convert_2(self) -> None: """Unknown origin raises TypeError on convert (duplicate check).""" schema = S("not_a_real_type") # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.convert(42) def test_unknown_origin_errors_on_check_value_2(self) -> None: """Unknown origin raises TypeError on check_value (duplicate check).""" schema = S("not_a_real_type") # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.check_value(42) def test_typo_origin_errors(self) -> None: """Common typos are caught, not silently passed through.""" for typo in ("innt", "floot", "strr", "Int", "Float"): schema = S(typo) # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.check_value(42) def test_unknown_nested_in_optional_errors(self) -> None: """Unknown origin nested inside Optional errors on use.""" schema = S("Optional", S("not_a_real_type")) # S: intentionally invalid origin with pytest.raises(TypeError, match="unknown TypeSchema origin"): schema.check_value(42) # --------------------------------------------------------------------------- # Category 43: convert/check_value raise TypeError on errors # --------------------------------------------------------------------------- class TestConvertCheckValueErrors: """Verify convert and check_value raise TypeError on errors.""" def test_convert_catches_custom_integral_error(self) -> None: """Custom Integral whose __int__ raises is caught by convert.""" class BadInt: """Registered as Integral via ABC but __int__ raises.""" def __int__(self) -> int: raise RuntimeError("broken __int__") Integral.register(BadInt) with pytest.raises(TypeError, match="broken __int__"): A(int).convert(BadInt()) def test_check_value_catches_custom_integral_error(self) -> None: """Custom Integral whose __int__ raises is caught by check_value.""" class BadInt2: def __int__(self) -> int: raise ValueError("bad int conversion") Integral.register(BadInt2) with pytest.raises(TypeError, match="bad int conversion"): A(int).check_value(BadInt2()) def test_convert_unknown_origin_raises(self) -> None: """Convert with unknown origin raises TypeError.""" with pytest.raises(TypeError, match="unknown TypeSchema origin"): S("bogus_type").convert("anything") # S: intentionally invalid origin def test_check_value_unknown_origin_raises(self) -> None: """check_value with unknown origin raises TypeError.""" with pytest.raises(TypeError, match="unknown TypeSchema origin"): S("bogus_type").check_value("anything") # S: intentionally invalid origin # --------------------------------------------------------------------------- # Category 44: Schema arity validation (ValueError, not assert) # --------------------------------------------------------------------------- class TestSchemaArityValidation: """Verify arity checks use ValueError (not assert) so they work under -O.""" def test_union_too_few_args(self) -> None: """Union with < 2 args raises ValueError.""" with pytest.raises(ValueError, match="at least two"): S("Union", A(int)) def test_optional_wrong_arity(self) -> None: """Optional with != 1 arg raises ValueError.""" with pytest.raises(ValueError, match="exactly one"): S("Optional") with pytest.raises(ValueError, match="exactly one"): S("Optional", A(int), A(str)) def test_array_too_many_args(self) -> None: """Array with > 1 arg raises ValueError.""" with pytest.raises(ValueError, match="0 or 1"): S("Array", A(int), A(str)) def test_list_too_many_args(self) -> None: """List with > 1 arg raises ValueError.""" with pytest.raises(ValueError, match="0 or 1"): S("List", A(int), A(str)) def test_map_wrong_arity(self) -> None: """Map with 1 or 3 args raises ValueError.""" with pytest.raises(ValueError, match="0 or 2"): S("Map", A(str)) with pytest.raises(ValueError, match="0 or 2"): S("Map", A(str), A(int), A(float)) def test_dict_wrong_arity(self) -> None: """Dict with 1 or 3 args raises ValueError.""" with pytest.raises(ValueError, match="0 or 2"): S("Dict", A(str)) def test_lowercase_list_too_many_args(self) -> None: """Lowercase 'list' with > 1 arg raises ValueError.""" with pytest.raises(ValueError, match="0 or 1"): S("list", A(int), A(str)) # S: lowercase "list" is an internal origin def test_lowercase_dict_wrong_arity(self) -> None: """Lowercase 'dict' with 1 arg raises ValueError.""" with pytest.raises(ValueError, match="0 or 2"): S("dict", A(str)) # S: lowercase "dict" is an internal origin # --------------------------------------------------------------------------- # Category 45: from_type_index edge cases # --------------------------------------------------------------------------- class TestFromTypeIndexEdgeCases: """Verify from_type_index behavior for valid indices. Note: Unregistered type indices trigger a fatal C++ assertion (TVMFFIGetTypeInfo CHECK failure) that cannot be caught from Python. Only valid indices obtained from the type registry should be passed. """ def test_valid_pod_index_roundtrip(self) -> None: """POD type_index from TypeSchema.origin_type_index round-trips.""" int_schema = A(int) schema = TypeSchema.from_type_index(int_schema.origin_type_index) assert schema.origin == "int" schema.check_value(42) def test_valid_object_index_works(self) -> None: """Valid registered object type_index constructs fine.""" tindex = tvm_ffi.core._object_type_key_to_index("testing.TestIntPair") assert tindex is not None schema = TypeSchema.from_type_index(tindex) assert schema.origin == "testing.TestIntPair" def test_from_type_index_with_args(self) -> None: """from_type_index with type arguments creates parameterized schema.""" arr_schema = A(tvm_ffi.Array) schema = TypeSchema.from_type_index(arr_schema.origin_type_index, (A(int),)) assert schema.origin == "Array" schema.check_value([1, 2, 3]) # =========================================================================== # Protocol-based conversion tests (matching Python FFI marshal path) # =========================================================================== # --------------------------------------------------------------------------- # Category 46: __tvm_ffi_int__ protocol # --------------------------------------------------------------------------- class TestIntProtocol: """int schema accepts values with __tvm_ffi_int__ protocol.""" def test_int_protocol_accepted(self) -> None: """Object with __tvm_ffi_int__ passes int schema.""" class IntProto: def __tvm_ffi_int__(self) -> int: return 42 A(int).check_value(IntProto()) def test_int_protocol_check_value(self) -> None: """check_value succeeds for __tvm_ffi_int__ value.""" class IntProto: def __tvm_ffi_int__(self) -> int: return 10 A(int).check_value(IntProto()) def test_int_protocol_convert_returns_value(self) -> None: """Convert returns the protocol value as-is (marshal handles conversion).""" class IntProto: def __tvm_ffi_int__(self) -> int: return 99 obj = IntProto() result = _to_py_class_value(A(int).convert(obj)) assert result is not None def test_without_protocol_still_rejected(self) -> None: """Object without __tvm_ffi_int__ is still rejected by int schema.""" class NoProto: pass with pytest.raises(TypeError, match="expected int"): A(int).check_value(NoProto()) # --------------------------------------------------------------------------- # Category 47: __tvm_ffi_float__ protocol # --------------------------------------------------------------------------- class TestFloatProtocol: """float schema accepts values with __tvm_ffi_float__ protocol.""" def test_float_protocol_accepted(self) -> None: """Object with __tvm_ffi_float__ passes float schema.""" class FloatProto: def __tvm_ffi_float__(self) -> float: return 3.14 A(float).check_value(FloatProto()) def test_float_protocol_convert(self) -> None: """Convert returns protocol value as-is.""" class FloatProto: def __tvm_ffi_float__(self) -> float: return 2.0 obj = FloatProto() result = _to_py_class_value(A(float).convert(obj)) assert result is not None def test_without_protocol_still_rejected(self) -> None: """Object without __tvm_ffi_float__ is still rejected.""" class NoProto: pass with pytest.raises(TypeError, match="expected float"): A(float).check_value(NoProto()) # --------------------------------------------------------------------------- # Category 48: __tvm_ffi_opaque_ptr__ protocol # --------------------------------------------------------------------------- class TestOpaquePtrProtocol: """ctypes.c_void_p schema accepts __tvm_ffi_opaque_ptr__ protocol.""" def test_opaque_ptr_protocol_accepted(self) -> None: """Object with __tvm_ffi_opaque_ptr__ passes ctypes.c_void_p schema.""" class PtrProto: def __tvm_ffi_opaque_ptr__(self) -> int: return 0xDEAD A(ctypes.c_void_p).check_value(PtrProto()) def test_opaque_ptr_protocol_convert(self) -> None: """Convert returns protocol value as-is.""" class PtrProto: def __tvm_ffi_opaque_ptr__(self) -> int: return 0 obj = PtrProto() result = _to_py_class_value(A(ctypes.c_void_p).convert(obj)) assert result is not None # --------------------------------------------------------------------------- # Category 49: __dlpack_device__ protocol # --------------------------------------------------------------------------- class TestDeviceProtocol: """Device schema accepts __dlpack_device__ protocol.""" def test_dlpack_device_protocol_accepted(self) -> None: """Object with __dlpack_device__ passes Device schema.""" class DevProto: def __dlpack_device__(self) -> tuple[int, int]: return (1, 0) A(tvm_ffi.Device).check_value(DevProto()) def test_dlpack_device_protocol_convert(self) -> None: """Convert returns protocol value as-is.""" class DevProto: def __dlpack_device__(self) -> tuple[int, int]: return (2, 1) obj = DevProto() result = _to_py_class_value(A(tvm_ffi.Device).convert(obj)) assert result is not None def test_without_protocol_still_rejected(self) -> None: """Object without __dlpack_device__ is still rejected.""" class NoProto: pass with pytest.raises(TypeError, match="expected Device"): A(tvm_ffi.Device).check_value(NoProto()) # --------------------------------------------------------------------------- # Category 50: dtype protocols (torch.dtype, numpy.dtype, __dlpack_data_type__) # --------------------------------------------------------------------------- class TestDtypeProtocols: """dtype schema accepts torch.dtype, numpy.dtype, __dlpack_data_type__.""" def test_dlpack_data_type_protocol_accepted(self) -> None: """Object with __dlpack_data_type__ passes dtype schema.""" class DTypeProto: def __dlpack_data_type__(self) -> tuple[int, int, int]: return (2, 32, 1) # float32 A(tvm_ffi.core.DataType).check_value(DTypeProto()) def test_dlpack_data_type_protocol_convert(self) -> None: """Convert returns protocol value as-is.""" class DTypeProto: def __dlpack_data_type__(self) -> tuple[int, int, int]: return (0, 32, 1) obj = DTypeProto() result = _to_py_class_value(A(tvm_ffi.core.DataType).convert(obj)) assert result is not None def test_numpy_dtype_accepted(self) -> None: """numpy.dtype passes dtype schema (if numpy installed).""" numpy = pytest.importorskip("numpy") A(tvm_ffi.core.DataType).check_value(numpy.dtype("float32")) def test_numpy_dtype_convert(self) -> None: """Convert returns numpy.dtype as-is.""" numpy = pytest.importorskip("numpy") dt = numpy.dtype("int32") result = _to_py_class_value(A(tvm_ffi.core.DataType).convert(dt)) assert result is not None def test_torch_dtype_accepted(self) -> None: """torch.dtype passes dtype schema (if torch installed).""" torch = pytest.importorskip("torch") A(tvm_ffi.core.DataType).check_value(torch.float32) def test_torch_dtype_convert(self) -> None: """Convert returns torch.dtype as-is.""" torch = pytest.importorskip("torch") dt = torch.int64 result = _to_py_class_value(A(tvm_ffi.core.DataType).convert(dt)) assert result is not None # --------------------------------------------------------------------------- # Category 51: __dlpack_c_exchange_api__ protocol (Tensor) # --------------------------------------------------------------------------- class TestTensorProtocol: """Tensor schema accepts __dlpack_c_exchange_api__ protocol.""" def test_dlpack_c_exchange_api_accepted(self) -> None: """Object with a valid __dlpack_c_exchange_api__ passes Tensor schema.""" np = pytest.importorskip("numpy") tensor = tvm_ffi.from_dlpack(np.arange(4, dtype="int32")) wrapper = tvm_ffi.core.DLTensorTestWrapper(tensor) A(tvm_ffi.Tensor).check_value(wrapper) def test_dlpack_c_exchange_api_convert(self) -> None: """Valid exchange-api wrappers can be converted to Tensor.""" np = pytest.importorskip("numpy") tensor = tvm_ffi.from_dlpack(np.arange(4, dtype="int32")) wrapper = tvm_ffi.core.DLTensorTestWrapper(tensor) result = _to_py_class_value(A(tvm_ffi.Tensor).convert(wrapper)) assert isinstance(result, tvm_ffi.Tensor) def test_dlpack_still_accepted(self) -> None: """Object with __dlpack__ still accepted (existing behavior).""" np = pytest.importorskip("numpy") A(tvm_ffi.Tensor).check_value(np.arange(4, dtype="int32")) # --------------------------------------------------------------------------- # Category 52: __tvm_ffi_object__ protocol # --------------------------------------------------------------------------- class TestObjectProtocol: """Object schemas accept __tvm_ffi_object__ protocol.""" def test_object_protocol_generic_object(self) -> None: """__tvm_ffi_object__ returning a CObject passes generic Object schema.""" inner = TestIntPair(1, 2) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner A(tvm_ffi.core.Object).check_value(ObjProto()) def test_object_protocol_specific_type(self) -> None: """__tvm_ffi_object__ returning TestIntPair passes TestIntPair schema.""" inner = TestIntPair(3, 4) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner A(TestIntPair).check_value(ObjProto()) def test_object_protocol_convert_returns_cobject(self) -> None: """Convert returns the CObject from __tvm_ffi_object__().""" inner = TestIntPair(5, 6) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner result = _to_py_class_value(A(TestIntPair).convert(ObjProto())) assert result.same_as(inner) def test_object_protocol_wrong_type_rejected(self) -> None: """__tvm_ffi_object__ returning wrong type is rejected.""" inner = TestIntPair(1, 2) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner with pytest.raises( TypeError, match=r"expected testing\.TestCxxClassBase, got testing\.TestIntPair" ): A(_TestCxxClassBase).check_value(ObjProto()) def test_object_protocol_raises_caught(self) -> None: """__tvm_ffi_object__ that raises produces _ConvertError.""" class BadProto: def __tvm_ffi_object__(self) -> object: raise RuntimeError("broken") with pytest.raises(TypeError, match=r"__tvm_ffi_object__\(\) failed"): A(tvm_ffi.core.Object).check_value(BadProto()) def test_object_protocol_hierarchy(self) -> None: """__tvm_ffi_object__ returning derived passes base schema.""" derived = _TestCxxClassDerived(v_i64=1, v_i32=2, v_f64=3.0) class ObjProto: def __tvm_ffi_object__(self) -> object: return derived A(_TestCxxClassBase).check_value(ObjProto()) # --------------------------------------------------------------------------- # Category 53: ObjectConvertible protocol # --------------------------------------------------------------------------- class TestObjectConvertibleProtocol: """Object schemas accept ObjectConvertible subclass.""" def test_object_convertible_accepted(self) -> None: """ObjectConvertible with asobject() passes Object schema.""" inner = TestIntPair(10, 20) class MyConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner A(tvm_ffi.core.Object).check_value(MyConvertible()) def test_object_convertible_specific_type(self) -> None: """ObjectConvertible passes specific type schema.""" inner = TestIntPair(1, 2) class MyConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner A(TestIntPair).check_value(MyConvertible()) def test_object_convertible_convert_returns_cobject(self) -> None: """Convert returns the CObject from asobject().""" inner = TestIntPair(7, 8) class MyConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner result = _to_py_class_value(A(TestIntPair).convert(MyConvertible())) assert result.same_as(inner) def test_object_convertible_in_union(self) -> None: """Union dispatch unwraps ObjectConvertible before trying alternatives.""" inner = TestIntPair(9, 10) class MyConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner result = _to_py_class_value(A(Union[TestIntPair, int]).convert(MyConvertible())) assert result.same_as(inner) def test_object_convertible_wrong_type(self) -> None: """ObjectConvertible returning wrong type is rejected.""" inner = TestIntPair(1, 2) class MyConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner with pytest.raises( TypeError, match=r"type check failed for testing\.TestCxxClassBase: expected testing\.TestCxxClassBase, got testing\.TestIntPair", ): A(_TestCxxClassBase).check_value(MyConvertible()) def test_object_convertible_raises_caught(self) -> None: """asobject() that raises produces error, not exception.""" class BadConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: raise RuntimeError("broken asobject") with pytest.raises(TypeError, match=r"asobject\(\) failed"): A(tvm_ffi.core.Object).check_value(BadConvertible()) # --------------------------------------------------------------------------- # Category 54: __tvm_ffi_value__ protocol (recursive fallback) # --------------------------------------------------------------------------- class TestValueProtocol: """__tvm_ffi_value__ provides recursive conversion fallback.""" def test_value_protocol_int(self) -> None: """__tvm_ffi_value__ returning int passes int schema.""" class ValProto: def __tvm_ffi_value__(self) -> object: return 42 A(int).check_value(ValProto()) def test_value_protocol_float(self) -> None: """__tvm_ffi_value__ returning float passes float schema.""" class ValProto: def __tvm_ffi_value__(self) -> object: return 3.14 A(float).check_value(ValProto()) def test_value_protocol_convert(self) -> None: """Convert returns the unwrapped value from __tvm_ffi_value__.""" class ValProto: def __tvm_ffi_value__(self) -> object: return 42 result = _to_py_class_value(A(int).convert(ValProto())) assert result == 42 def test_value_protocol_nested(self) -> None: """Nested __tvm_ffi_value__ is recursively unwrapped.""" class ValProto: def __init__(self, v: object) -> None: self.v = v def __tvm_ffi_value__(self) -> object: return self.v # ValProto(ValProto(ValProto(10))) should unwrap to 10 wrapped = ValProto(ValProto(ValProto(10))) assert _to_py_class_value(A(int).convert(wrapped)) == 10 def test_value_protocol_object(self) -> None: """__tvm_ffi_value__ returning a CObject passes object schema.""" inner = TestIntPair(1, 2) class ValProto: def __tvm_ffi_value__(self) -> object: return inner A(TestIntPair).check_value(ValProto()) def test_value_protocol_still_fails_on_mismatch(self) -> None: """__tvm_ffi_value__ returning wrong type still fails.""" class ValProto: def __tvm_ffi_value__(self) -> object: return "not_an_int" with pytest.raises(TypeError, match="expected int"): A(int).check_value(ValProto()) def test_value_protocol_raises_uses_original_error(self) -> None: """If __tvm_ffi_value__ raises, the original error is returned.""" class BadValProto: def __tvm_ffi_value__(self) -> object: raise RuntimeError("broken") with pytest.raises(TypeError, match="expected int"): A(int).check_value(BadValProto()) def test_nested_optional_value_protocol_stall(self) -> None: """Optional[Optional[float]] reports the unwrapped target type.""" class SelfRef: def __tvm_ffi_value__(self) -> object: return self with pytest.raises(TypeError, match="expected float"): A(Optional[Optional[float]]).check_value(SelfRef()) def test_value_protocol_eventually_resolves(self) -> None: """Short __tvm_ffi_value__ chains still resolve successfully.""" class ChainStep: def __init__(self, remaining: int) -> None: self.remaining = remaining def __tvm_ffi_value__(self) -> object: if self.remaining > 0: return ChainStep(self.remaining - 1) return 42 assert _to_py_class_value(A(int).convert(ChainStep(5))) == 42 # --------------------------------------------------------------------------- # Category 55: Protocol values in containers # --------------------------------------------------------------------------- class TestProtocolsInContainers: """Protocol-accepting values work inside containers and composites.""" @requires_py39 def test_int_protocol_in_array(self) -> None: """Array[int] accepts elements with __tvm_ffi_int__.""" class IntProto: def __tvm_ffi_int__(self) -> int: return 1 A(tuple[int, ...]).check_value([1, IntProto(), 3]) def test_float_protocol_in_optional(self) -> None: """Optional[float] accepts __tvm_ffi_float__ value.""" class FloatProto: def __tvm_ffi_float__(self) -> float: return 1.0 A(Optional[float]).check_value(FloatProto()) A(Optional[float]).check_value(None) def test_object_protocol_in_union(self) -> None: """Union[testing.TestIntPair, int] accepts __tvm_ffi_object__ value.""" inner = TestIntPair(1, 2) class ObjProto: def __tvm_ffi_object__(self) -> object: return inner A(Union[TestIntPair, int]).check_value(ObjProto()) @requires_py39 def test_value_protocol_in_array(self) -> None: """Array[int] elements use __tvm_ffi_value__ fallback (recursive).""" class ValProto: def __tvm_ffi_value__(self) -> object: return 42 # __tvm_ffi_value__ fallback is applied recursively at every level, # matching the marshal path where TVMFFIPyArgSetterFactory_ is # called per-element. A(tuple[int, ...]).check_value([ValProto()]) @requires_py39 def test_object_convertible_in_array(self) -> None: """Array[Object] elements unwrap ObjectConvertible before element dispatch.""" inner = TestIntPair(3, 4) class Convertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return inner result = _to_py_class_value(A(tuple[tvm_ffi.core.Object, ...]).convert([Convertible()])) assert result[0].same_as(inner) @requires_py39 def test_device_protocol_in_map_value(self) -> None: """Map[str, Device] accepts __dlpack_device__ values.""" class DevProto: def __dlpack_device__(self) -> tuple[int, int]: return (1, 0) A(tvm_ffi.Map[str, tvm_ffi.Device]).check_value({"gpu": DevProto()}) # --------------------------------------------------------------------------- # Category 56: Nested __tvm_ffi_value__ in containers (recursive fallback) # --------------------------------------------------------------------------- class TestNestedValueProtocol: """__tvm_ffi_value__ fallback works recursively inside containers.""" @requires_py39 def test_value_in_array_elements(self) -> None: """Array[int] elements with __tvm_ffi_value__ are accepted.""" class VP: def __tvm_ffi_value__(self) -> object: return 42 A(tuple[int, ...]).check_value([1, VP(), 3]) @requires_py39 def test_value_in_map_values(self) -> None: """Map[str, int] values with __tvm_ffi_value__ are accepted.""" class VP: def __tvm_ffi_value__(self) -> object: return 99 A(tvm_ffi.Map[str, int]).check_value({"a": VP()}) @requires_py39 def test_value_in_map_keys(self) -> None: """Map[str, int] keys with __tvm_ffi_value__ are accepted.""" class VP: def __tvm_ffi_value__(self) -> object: return "key" A(tvm_ffi.Map[str, int]).check_value({VP(): 1}) @requires_py39 def test_value_in_tuple_positions(self) -> None: """tuple[int, str] positions with __tvm_ffi_value__ are accepted.""" class IntVP: def __tvm_ffi_value__(self) -> object: return 42 class StrVP: def __tvm_ffi_value__(self) -> object: return "hello" A(tuple[int, str]).check_value((IntVP(), StrVP())) def test_value_in_optional_inner(self) -> None: """Optional[int] inner with __tvm_ffi_value__ is accepted.""" class VP: def __tvm_ffi_value__(self) -> object: return 42 A(Optional[int]).check_value(VP()) def test_value_in_union_alternatives(self) -> None: """Union[int, str] with __tvm_ffi_value__ is accepted.""" class VP: def __tvm_ffi_value__(self) -> object: return "hello" A(Union[int, str]).check_value(VP()) @requires_py39 def test_multi_hop_value_in_container(self) -> None: """Nested __tvm_ffi_value__ unwrapping inside containers.""" class VP: def __init__(self, v: object) -> None: self.v = v def __tvm_ffi_value__(self) -> object: return self.v A(tuple[int, ...]).check_value([VP(VP(10))]) @requires_py39 def test_value_convert_in_array(self) -> None: """Convert returns unwrapped values in container.""" class VP: def __tvm_ffi_value__(self) -> object: return 42 result = _to_py_class_value(A(tuple[int, ...]).convert([VP()])) assert list(result) == [42] # --------------------------------------------------------------------------- # Category 57: __tvm_ffi_value__ cycle protection # --------------------------------------------------------------------------- class TestValueProtocolCycles: """Cycle protection in __tvm_ffi_value__ fallback.""" def test_self_cycle_returns_error(self) -> None: """__tvm_ffi_value__() returning self doesn't infinite-loop.""" class SelfCycle: def __tvm_ffi_value__(self) -> object: return self with pytest.raises(TypeError, match="expected int"): A(int).check_value(SelfCycle()) def test_any_self_cycle_returns_original_error(self) -> None: """Any also routes __tvm_ffi_value__ through the bounded fallback loop.""" call_count = 0 class SelfCycle: def __tvm_ffi_value__(self) -> object: nonlocal call_count call_count += 1 return self with pytest.raises(TypeError, match=r"failed to convert Any from .*SelfCycle"): A(typing.Any).convert(SelfCycle()) assert call_count == 1 def test_mutual_cycle_bounded(self) -> None: """Mutual cycle is bounded by explicit depth limit.""" class A: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other class B: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other a, b = A(), B() a.other = b b.other = a # Should not hang — bounded by depth limit in the fallback loop with pytest.raises(TypeError, match="cycle"): S("int").check_value(a) def test_any_mutual_cycle_bounded(self) -> None: """Any reports a bounded cycle instead of recursing in raw CAny packing.""" class Left: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other class Right: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other a, b = Left(), Right() a.other = b b.other = a with pytest.raises(TypeError, match="cycle"): A(typing.Any).convert(a) def test_value_protocol_deep_chain_hits_cycle_limit(self) -> None: """Long __tvm_ffi_value__ chains trip the explicit depth guard.""" class DeepChain: def __init__(self, depth: int) -> None: self.depth = depth def __tvm_ffi_value__(self) -> object: return DeepChain(self.depth + 1) if self.depth < 100 else 42 with pytest.raises(TypeError, match="cycle"): A(int).check_value(DeepChain(0)) # --------------------------------------------------------------------------- # Category 58: Object marshal fallback # --------------------------------------------------------------------------- class TestObjectConvertAttrRegistration: """Object targets register __ffi_convert__ consistently.""" def test_core_object_types_register_convert_attr(self) -> None: """Core object types register ``__ffi_convert__``.""" for type_key in ( "ffi.Object", "ffi.Function", "ffi.Error", "ffi.String", "ffi.Bytes", "ffi.Array", "ffi.List", "ffi.Map", "ffi.Dict", "ffi.Shape", "ffi.Tensor", ): type_index = _object_type_key_to_index(type_key) assert type_index is not None assert _lookup_type_attr(type_index, "__ffi_convert__") is not None def test_explicit_ref_registration_registers_convert_attr(self) -> None: """Explicit ``.ref()`` registration adds ``__ffi_convert__``.""" type_index = _object_type_key_to_index("testing.TestIntPair") assert type_index is not None assert _lookup_type_attr(type_index, "__ffi_convert__") is not None def test_reflected_object_without_ref_does_not_register_convert_attr(self) -> None: """Reflected object classes without ``.ref()`` do not auto-register.""" type_index = _object_type_key_to_index("testing.TestObjectBase") assert type_index is not None assert _lookup_type_attr(type_index, "__ffi_convert__") is None class TestObjectMarshalFallback: """Object schema accepts values that the marshal path converts to Objects.""" def test_exception_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts Exception (-> ffi.Error).""" A(tvm_ffi.core.Object).check_value(RuntimeError("test")) def test_exception_accepted_by_error_schema(self) -> None: """TypeSchema('ffi.Error') accepts Exception.""" A(tvm_ffi.core.Error).check_value(ValueError("oops")) def test_shape_accepted_from_python_list_via_convert(self) -> None: """TypeSchema('ffi.Shape') converts Python lists via __ffi_convert__.""" result = _to_py_class_value(S("ffi.Shape").convert([1, 2, 3])) assert isinstance(result, tvm_ffi.Shape) assert tuple(result) == (1, 2, 3) def test_shape_accepted_from_ffi_list_via_convert(self) -> None: """TypeSchema('ffi.Shape') converts ffi.List via __ffi_convert__.""" result = _to_py_class_value(S("ffi.Shape").convert(tvm_ffi.List([1, 2, 3]))) assert isinstance(result, tvm_ffi.Shape) assert tuple(result) == (1, 2, 3) def test_shape_accepted_from_ffi_array_via_convert(self) -> None: """TypeSchema('ffi.Shape') converts ffi.Array via __ffi_convert__.""" result = _to_py_class_value(S("ffi.Shape").convert(tvm_ffi.Array([1, 2, 3]))) assert isinstance(result, tvm_ffi.Shape) assert tuple(result) == (1, 2, 3) def test_shape_convert_repeated_conversions(self) -> None: """Repeated __ffi_convert__ conversions keep returning live owning objects.""" result1 = _to_py_class_value(S("ffi.Shape").convert([1, 2, 3])) result2 = _to_py_class_value(S("ffi.Shape").convert(tvm_ffi.List([1, 2, 3]))) assert isinstance(result1, tvm_ffi.Shape) assert isinstance(result2, tvm_ffi.Shape) assert tuple(result1) == (1, 2, 3) assert tuple(result2) == (1, 2, 3) def test_exception_rejected_by_array_schema(self) -> None: """Exception is NOT accepted by Array schema (Error !IS-A Array).""" with pytest.raises(TypeError, match="expected Array"): A(tvm_ffi.Array).check_value(RuntimeError("x")) def test_opaque_object_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts arbitrary Python objects (-> OpaquePyObject).""" class Custom: pass A(tvm_ffi.core.Object).check_value(Custom()) def test_plain_object_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts object().""" A(tvm_ffi.core.Object).check_value(object()) def test_opaque_rejected_by_specific_schema(self) -> None: """Specific schema rejects arbitrary Python object.""" class Custom: pass with pytest.raises(TypeError, match=r"got .*Custom"): A(TestIntPair).check_value(Custom()) @pytest.mark.xfail(reason="SmallStr -> ObjectRef conversion is not supported yet") def test_str_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts str (-> ffi.String IS-A Object).""" A(tvm_ffi.core.Object).check_value("hello") @pytest.mark.xfail(reason="SmallBytes -> ObjectRef conversion is not supported yet") def test_bytes_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts bytes (-> ffi.Bytes IS-A Object).""" A(tvm_ffi.core.Object).check_value(b"hello") def test_list_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts list (-> ffi.Array IS-A Object).""" A(tvm_ffi.core.Object).check_value([1, 2, 3]) def test_dict_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts dict (-> ffi.Map IS-A Object).""" A(tvm_ffi.core.Object).check_value({"a": 1}) def test_callable_accepted_by_object_schema(self) -> None: """TypeSchema('Object') accepts callable (-> ffi.Function IS-A Object).""" A(tvm_ffi.core.Object).check_value(lambda: None) def test_int_rejected_by_object_schema(self) -> None: """TypeSchema('Object') rejects int (int is a POD type, not Object).""" with pytest.raises(TypeError): A(tvm_ffi.core.Object).check_value(42) def test_float_rejected_by_object_schema(self) -> None: """TypeSchema('Object') rejects float (float is a POD, not Object).""" with pytest.raises(TypeError): A(tvm_ffi.core.Object).check_value(3.14) def test_none_rejected_by_object_schema(self) -> None: """TypeSchema('Object') rejects None (None is a POD, not Object).""" with pytest.raises(TypeError): A(tvm_ffi.core.Object).check_value(None) # --------------------------------------------------------------------------- # Category 59: __cuda_stream__ for ctypes.c_void_p # --------------------------------------------------------------------------- class TestCudaStreamProtocol: """ctypes.c_void_p schema accepts __cuda_stream__ protocol.""" def test_cuda_stream_accepted(self) -> None: """Object with __cuda_stream__ passes ctypes.c_void_p schema.""" class CUStream: def __cuda_stream__(self) -> tuple[int, int]: return (0, 0) A(ctypes.c_void_p).check_value(CUStream()) def test_cuda_stream_convert(self) -> None: """Convert returns __cuda_stream__ value as-is.""" class CUStream: def __cuda_stream__(self) -> tuple[int, int]: return (0, 123) obj = CUStream() result = _to_py_class_value(A(ctypes.c_void_p).convert(obj)) assert result is not None def test_cuda_stream_and_opaque_ptr(self) -> None: """Object with both __cuda_stream__ and __tvm_ffi_opaque_ptr__ accepted.""" class DualProto: def __cuda_stream__(self) -> tuple[int, int]: return (0, 0) def __tvm_ffi_opaque_ptr__(self) -> int: return 0 A(ctypes.c_void_p).check_value(DualProto()) # --------------------------------------------------------------------------- # Category 60: Device __dlpack__ guard # --------------------------------------------------------------------------- class TestDeviceDlpackGuard: """Device schema respects __dlpack__ precedence.""" def test_both_dlpack_and_device_rejected_by_device(self) -> None: """Object with both __dlpack__ and __dlpack_device__ rejected by Device.""" class TensorLike: def __dlpack__(self) -> object: return None def __dlpack_device__(self) -> tuple[int, int]: return (1, 0) with pytest.raises(TypeError): A(tvm_ffi.Device).check_value(TensorLike()) def test_both_dlpack_and_device_accepted_by_tensor(self) -> None: """Object with both __dlpack__ and __dlpack_device__ accepted by Tensor.""" np = pytest.importorskip("numpy") class TensorLike: def __init__(self) -> None: self.array = np.arange(4, dtype="int32") def __dlpack__(self) -> object: return self.array.__dlpack__() def __dlpack_device__(self) -> tuple[int, int]: return self.array.__dlpack_device__() A(tvm_ffi.Tensor).check_value(TensorLike()) def test_device_only_accepted_by_device(self) -> None: """Object with only __dlpack_device__ still accepted by Device.""" class DevOnly: def __dlpack_device__(self) -> tuple[int, int]: return (1, 0) A(tvm_ffi.Device).check_value(DevOnly()) def test_dlpack_only_rejected_by_device(self) -> None: """Object with only __dlpack__ rejected by Device schema.""" class DLPackOnly: def __dlpack__(self) -> object: return None with pytest.raises(TypeError): A(tvm_ffi.Device).check_value(DLPackOnly()) # --------------------------------------------------------------------------- # Category 61: SKIP_DLPACK_C_EXCHANGE_API env gate # --------------------------------------------------------------------------- class TestSkipDlpackEnvGate: """Tensor schema respects TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API.""" def test_exchange_api_accepted_by_default(self) -> None: """__dlpack_c_exchange_api__ accepted when env not set.""" os.environ.pop("TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API", None) np = pytest.importorskip("numpy") tensor = tvm_ffi.from_dlpack(np.arange(4, dtype="int32")) wrapper = tvm_ffi.core.DLTensorTestWrapper(tensor) A(tvm_ffi.Tensor).check_value(wrapper) def test_exchange_api_rejected_when_skipped(self) -> None: """__dlpack_c_exchange_api__ rejected when env=1.""" os.environ["TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API"] = "1" try: class ExchangeAPI: def __dlpack_c_exchange_api__(self) -> int: return 0 with pytest.raises(TypeError): A(tvm_ffi.Tensor).check_value(ExchangeAPI()) finally: del os.environ["TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API"] # --------------------------------------------------------------------------- # Category 62: from_type_index low-level indices # --------------------------------------------------------------------------- class TestFromTypeIndexLowLevel: """from_type_index handles all built-in type indices.""" def test_dl_tensor_ptr(self) -> None: """KTVMFFIDLTensorPtr maps to Tensor.""" s = TypeSchema.from_type_index(7) # kTVMFFIDLTensorPtr assert s.origin == "Tensor" def test_raw_str(self) -> None: """KTVMFFIRawStr maps to str.""" s = TypeSchema.from_type_index(8) # kTVMFFIRawStr assert s.origin == "str" def test_byte_array_ptr(self) -> None: """KTVMFFIByteArrayPtr maps to bytes.""" s = TypeSchema.from_type_index(9) # kTVMFFIByteArrayPtr assert s.origin == "bytes" def test_object_rvalue_ref(self) -> None: """KTVMFFIObjectRValueRef maps to Object.""" s = TypeSchema.from_type_index(10) # kTVMFFIObjectRValueRef assert s.origin == "Object" def test_small_str(self) -> None: """KTVMFFISmallStr maps to str.""" s = TypeSchema.from_type_index(11) # kTVMFFISmallStr assert s.origin == "str" def test_small_bytes(self) -> None: """KTVMFFISmallBytes maps to bytes.""" s = TypeSchema.from_type_index(12) # kTVMFFISmallBytes assert s.origin == "bytes" def test_all_low_level_schemas_usable(self) -> None: """Schemas from low-level indices can be used for conversion.""" for idx in (7, 8, 9, 11, 12): s = TypeSchema.from_type_index(idx) # Trigger converter build; some schemas raise TypeError for None try: s.convert(None) except TypeError: pass # --------------------------------------------------------------------------- # Category 63: STL origin parsing # --------------------------------------------------------------------------- class TestSTLOriginParsing: """C++ STL schema origins are correctly parsed.""" def test_std_vector(self) -> None: """std::vector maps to Array.""" s = TypeSchema.from_json_str('{"type":"std::vector","args":[{"type":"int"}]}') assert s.origin == "Array" def test_std_optional(self) -> None: """std::optional maps to Optional.""" s = TypeSchema.from_json_str('{"type":"std::optional","args":[{"type":"int"}]}') assert s.origin == "Optional" assert repr(s) == "int | None" def test_std_variant(self) -> None: """std::variant maps to Union.""" s = TypeSchema.from_json_str( '{"type":"std::variant","args":[{"type":"int"},{"type":"str"}]}' ) assert s.origin == "Union" assert repr(s) == "int | str" def test_std_tuple(self) -> None: """std::tuple maps to tuple.""" s = TypeSchema.from_json_str('{"type":"std::tuple","args":[{"type":"int"},{"type":"str"}]}') assert s.origin == "tuple" def test_std_map(self) -> None: """std::map maps to Map.""" s = TypeSchema.from_json_str('{"type":"std::map","args":[{"type":"str"},{"type":"int"}]}') assert s.origin == "Map" def test_std_unordered_map(self) -> None: """std::unordered_map maps to Map.""" s = TypeSchema.from_json_str( '{"type":"std::unordered_map","args":[{"type":"str"},{"type":"int"}]}' ) assert s.origin == "Map" def test_std_function(self) -> None: """std::function maps to Callable.""" s = TypeSchema.from_json_str( '{"type":"std::function","args":[{"type":"int"},{"type":"str"}]}' ) assert s.origin == "Callable" def test_object_rvalue_ref_origin(self) -> None: """ObjectRValueRef maps to Object.""" s = TypeSchema.from_json_str('{"type":"ObjectRValueRef","args":[]}') assert s.origin == "Object" # --------------------------------------------------------------------------- # Category 64: Zero-copy container conversion # --------------------------------------------------------------------------- class TestZeroCopyConversion: """Typed container conversion preserves identity when no elements change.""" @requires_py39 def test_array_int_exact_list(self) -> None: """Array[int] on exact Python list converts successfully.""" original = [1, 2, 3] result = _to_py_class_value(A(tuple[int, ...]).convert(original)) assert list(result) == original @requires_py39 def test_array_int_needs_conversion(self) -> None: """Array[int] on list needing bool->int returns converted list.""" original = [1, True, 3] result = _to_py_class_value(A(tuple[int, ...]).convert(original)) assert list(result) == [1, 1, 3] @requires_py39 def test_map_str_int_exact_dict(self) -> None: """Map[str, int] on exact dict converts successfully.""" original = {"a": 1, "b": 2} result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert(original)) assert dict(result) == original @requires_py39 def test_map_str_int_needs_conversion(self) -> None: """Map[str, int] on dict needing conversion returns converted dict.""" original = {"a": True, "b": 2} result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert(original)) assert result is not None @requires_py39 def test_tuple_exact_match(self) -> None: """tuple[int, str] on exact tuple converts successfully.""" original = (42, "hello") result = _to_py_class_value(A(tuple[int, str]).convert(original)) assert tuple(result) == original @requires_py39 def test_tuple_needs_conversion(self) -> None: """tuple[int, str] on tuple needing conversion returns converted tuple.""" original = (True, "hello") result = _to_py_class_value(A(tuple[int, str]).convert(original)) assert tuple(result) == (1, "hello") @requires_py39 def test_list_int_exact(self) -> None: """List[int] on exact list converts successfully.""" original = [10, 20] result = _to_py_class_value(A(list[int]).convert(original)) assert list(result) == original # --------------------------------------------------------------------------- # Category 65: Exception normalization in check_value/convert # --------------------------------------------------------------------------- class TestExceptionNormalization: """check_value/convert normalize custom __int__/__float__ failures.""" def test_broken_integral_convert(self) -> None: """Integral with broken __int__ caught by convert.""" class BadIntegral: def __int__(self) -> int: raise OverflowError("too big") Integral.register(BadIntegral) with pytest.raises(TypeError, match="too big"): A(int).convert(BadIntegral()) def test_broken_integral_check_value(self) -> None: """Integral with broken __int__ handled by check_value.""" class BrokenInt: def __int__(self) -> int: raise ValueError("broken") Integral.register(BrokenInt) # check_value should raise TypeError (wrapping the ValueError) with pytest.raises(TypeError, match="broken"): A(int).check_value(BrokenInt()) def test_broken_integral_bool_check_value(self) -> None: """Integral with broken __bool__ is normalized to TypeError.""" class BrokenBoolInt: def __int__(self) -> int: return 1 def __bool__(self) -> bool: raise RuntimeError("broken bool") Integral.register(BrokenBoolInt) with pytest.raises(TypeError, match="broken bool"): A(bool).check_value(BrokenBoolInt()) def test_union_falls_back_after_broken_bool(self) -> None: """Union keeps trying alternatives when bool conversion fails. A str subclass registered as Integral is a pathological case; the C packer may dispatch it as Integral rather than str, so we only guarantee that Union dispatch does not silently swallow the error — either a successful conversion or a raised error is acceptable. """ class BrokenBoolStr(str): def __bool__(self) -> bool: raise RuntimeError("broken bool") Integral.register(BrokenBoolStr) try: result = _to_py_class_value(A(Union[bool, str]).convert(BrokenBoolStr("hello"))) assert result == "hello" except (ValueError, RuntimeError, TypeError): pass # --------------------------------------------------------------------------- # Category 66: __tvm_ffi_value__ eager normalization # --------------------------------------------------------------------------- class TestValueProtocolPrecedence: """__tvm_ffi_value__ runs before schema-specific dispatch.""" def test_value_protocol_runs_before_int_protocol(self) -> None: """__tvm_ffi_value__ is applied before __tvm_ffi_int__.""" class Dual: def __tvm_ffi_int__(self) -> int: return 42 def __tvm_ffi_value__(self) -> object: return TestIntPair(1, 2) with pytest.raises(TypeError): A(int).check_value(Dual()) A(tvm_ffi.core.Object).check_value(Dual()) def test_value_protocol_runs_before_float_protocol(self) -> None: """__tvm_ffi_value__ is applied before __tvm_ffi_float__.""" class Dual: def __tvm_ffi_float__(self) -> float: return 1.0 def __tvm_ffi_value__(self) -> object: return TestIntPair(1, 2) with pytest.raises(TypeError): A(float).check_value(Dual()) A(tvm_ffi.core.Object).check_value(Dual()) def test_pure_value_protocol_still_works(self) -> None: """Class with ONLY __tvm_ffi_value__ still converts eagerly.""" class PureVP: def __tvm_ffi_value__(self) -> object: return 42 A(int).check_value(PureVP()) def test_value_protocol_runs_before_callable_dispatch(self) -> None: """Callable classes are normalized through __tvm_ffi_value__ first.""" class CallableVP: def __call__(self) -> None: pass def __tvm_ffi_value__(self) -> object: return 42 with pytest.raises(TypeError): A(Callable).check_value(CallableVP()) A(int).check_value(CallableVP()) def test_object_protocol_precedes_value_and_convertible(self) -> None: """__tvm_ffi_object__ wins over value and ObjectConvertible hooks.""" inner = TestIntPair(10, 20) class Both(ObjectConvertible): def __tvm_ffi_object__(self) -> object: return inner def __tvm_ffi_value__(self) -> object: return 999 def asobject(self) -> tvm_ffi.core.Object: return TestIntPair(99, 99) result = _to_py_class_value(A(tvm_ffi.core.Object).convert(Both())) assert result.same_as(inner) # --------------------------------------------------------------------------- # Category 67: Union single-call __tvm_ffi_value__ # --------------------------------------------------------------------------- class TestUnionValueProtocol: """Union dispatches __tvm_ffi_value__ once, not per-alternative.""" def test_union_value_protocol_once(self) -> None: """__tvm_ffi_value__ called once for Union.""" call_count = 0 class CountingVP: def __tvm_ffi_value__(self) -> object: nonlocal call_count call_count += 1 return 42 A(Union[str, int]).check_value(CountingVP()) assert call_count == 1 def test_union_value_protocol_mismatch(self) -> None: """__tvm_ffi_value__ returning wrong type fails Union.""" call_count = 0 class WrongVP: def __tvm_ffi_value__(self) -> object: nonlocal call_count call_count += 1 return object() with pytest.raises(TypeError): A(Union[int, str]).check_value(WrongVP()) assert call_count == 1 # --------------------------------------------------------------------------- # Category 68: from_json_obj robustness # --------------------------------------------------------------------------- class TestFromJsonObjRobustness: """from_json_obj handles non-dict args and malformed input.""" def test_non_dict_args_skipped(self) -> None: """Non-dict elements in args list are silently skipped.""" obj = {"type": "std::vector", "args": [{"type": "int"}, 42]} s = TypeSchema.from_json_obj(obj) assert s.origin == "Array" assert len(s.args) == 1 assert s.args[0].origin == "int" def test_malformed_input_raises_type_error(self) -> None: """Non-dict top-level raises TypeError, not AssertionError.""" with pytest.raises(TypeError, match="expected schema dict"): TypeSchema.from_json_obj("not_a_dict") # type: ignore[arg-type] def test_missing_type_key_raises_type_error(self) -> None: """Dict without 'type' key raises TypeError.""" with pytest.raises(TypeError, match="expected schema dict"): TypeSchema.from_json_obj({"args": []}) # --------------------------------------------------------------------------- # Category 69: Mutual-cycle RecursionError normalized # --------------------------------------------------------------------------- class TestMutualCycleNormalization: """Mutual __tvm_ffi_value__ cycles produce TypeError, not RecursionError.""" def test_mutual_cycle_check_value(self) -> None: """check_value normalizes mutual-cycle RecursionError to TypeError.""" class A: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other class B: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other a, b = A(), B() a.other = b b.other = a with pytest.raises(TypeError, match="cycle"): S("int").check_value(a) def test_mutual_cycle_convert(self) -> None: """Convert normalizes mutual-cycle RecursionError to TypeError.""" class A: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other class B: def __init__(self) -> None: self.other: object = None def __tvm_ffi_value__(self) -> object: return self.other a, b = A(), B() a.other = b b.other = a with pytest.raises(TypeError, match="cycle"): S("int").convert(a) # --------------------------------------------------------------------------- # Category 70: ObjectConvertible vs __tvm_ffi_value__ precedence # --------------------------------------------------------------------------- class TestObjectConvertiblePrecedence: """__tvm_ffi_value__ takes precedence over ObjectConvertible.""" def test_value_protocol_wins_over_convertible(self) -> None: """Class with both __tvm_ffi_value__ and ObjectConvertible uses fallback.""" pair = TestIntPair(10, 20) class DualProtocol(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return pair def __tvm_ffi_value__(self) -> object: return 42 # int schema: __tvm_ffi_value__ returns 42, accepted A(int).check_value(DualProtocol()) # Object schema: __tvm_ffi_value__ returns 42 (POD int, not Object), # should REJECT (not accept via ObjectConvertible) with pytest.raises(TypeError): A(tvm_ffi.core.Object).check_value(DualProtocol()) def test_pure_convertible_still_works(self) -> None: """ObjectConvertible without __tvm_ffi_value__ still accepted.""" pair = TestIntPair(1, 2) class PureConvertible(ObjectConvertible): def asobject(self) -> tvm_ffi.core.Object: return pair A(tvm_ffi.core.Object).check_value(PureConvertible()) A(TestIntPair).check_value(PureConvertible()) # --------------------------------------------------------------------------- # Category 71: from_json_obj non-iterable args # --------------------------------------------------------------------------- class TestFromJsonObjNonIterableArgs: """from_json_obj handles non-iterable args values gracefully.""" def test_non_iterable_args_treated_as_empty(self) -> None: """Non-list/tuple args value (e.g., int) treated as empty args.""" s = TypeSchema.from_json_obj({"type": "int", "args": 42}) assert s.origin == "int" assert s.args == () def test_string_args_treated_as_empty(self) -> None: """String args value treated as empty (not iterated char-by-char).""" s = TypeSchema.from_json_obj({"type": "int", "args": "bad"}) assert s.origin == "int" assert s.args == () # --------------------------------------------------------------------------- # CAny class tests # --------------------------------------------------------------------------- class TestCAny: """Tests for the CAny owned-value container.""" def test_cany_from_int(self) -> None: """convert(int) returns CAny with correct type_index.""" cany = A(int).convert(42) assert isinstance(cany, CAny) assert cany.type_index == 1 # kTVMFFIInt def test_cany_from_float(self) -> None: """convert(float) returns CAny with correct type_index.""" cany = A(float).convert(3.14) assert isinstance(cany, CAny) assert cany.type_index == 3 # kTVMFFIFloat def test_cany_from_bool(self) -> None: """convert(bool) returns CAny with correct type_index.""" cany = A(bool).convert(True) assert isinstance(cany, CAny) assert cany.type_index == 2 # kTVMFFIBool def test_cany_from_none(self) -> None: """convert(None) returns CAny with type_index 0.""" cany = A(None).convert(None) assert isinstance(cany, CAny) assert cany.type_index == 0 # kTVMFFINone def test_cany_from_str(self) -> None: """convert(str) returns CAny.""" cany = A(str).convert("hello") assert isinstance(cany, CAny) # Short strings have type_index=11 (SmallStr), longer ones have 65 (Str) assert cany.type_index in (11, 65) @requires_py39 def test_cany_from_array(self) -> None: """convert(Array) returns CAny with array type_index.""" cany = A(tuple[int, ...]).convert([1, 2, 3]) assert isinstance(cany, CAny) assert cany.type_index >= 64 # object type def test_to_py_int(self) -> None: """to_py() round-trips int correctly.""" result = _to_py_class_value(A(int).convert(42)) assert result == 42 assert type(result) is int def test_to_py_float(self) -> None: """to_py() round-trips float correctly.""" result = _to_py_class_value(A(float).convert(3.14)) assert result == 3.14 assert type(result) is float def test_to_py_bool(self) -> None: """to_py() round-trips bool correctly.""" assert _to_py_class_value(A(bool).convert(True)) is True assert _to_py_class_value(A(bool).convert(False)) is False def test_to_py_none(self) -> None: """to_py() round-trips None correctly.""" assert _to_py_class_value(A(None).convert(None)) is None def test_to_py_str(self) -> None: """to_py() round-trips str correctly.""" assert _to_py_class_value(A(str).convert("hello")) == "hello" @requires_py39 def test_to_py_array(self) -> None: """to_py() returns ffi.Array for Array convert.""" result = _to_py_class_value(A(tuple[int, ...]).convert([1, 2, 3])) assert isinstance(result, tvm_ffi.Array) assert list(result) == [1, 2, 3] @requires_py39 def test_to_py_list(self) -> None: """to_py() returns ffi.List for List convert.""" result = _to_py_class_value(A(list[int]).convert([1, 2, 3])) assert isinstance(result, tvm_ffi.List) assert list(result) == [1, 2, 3] @requires_py39 def test_to_py_map(self) -> None: """to_py() returns ffi.Map for Map convert.""" result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert({"a": 1})) assert isinstance(result, tvm_ffi.Map) @requires_py39 def test_to_py_dict(self) -> None: """to_py() returns ffi.Dict for Dict convert.""" result = _to_py_class_value(A(dict[str, int]).convert({"a": 1})) assert isinstance(result, tvm_ffi.Dict) def test_multiple_to_py_calls(self) -> None: """to_py() can be called multiple times safely.""" cany = A(int).convert(42) assert _to_py_class_value(cany) == 42 assert _to_py_class_value(cany) == 42 assert _to_py_class_value(cany) == 42 @requires_py39 def test_object_refcount_safety(self) -> None: """to_py() for objects properly IncRefs — no double-free.""" cany = A(tuple[int, ...]).convert([1, 2, 3]) py1 = _to_py_class_value(cany) py2 = _to_py_class_value(cany) del cany # CAny.__dealloc__ runs assert list(py1) == [1, 2, 3] assert list(py2) == [1, 2, 3] def test_repr_int(self) -> None: """Repr shows type and value for int.""" cany = A(int).convert(42) assert "int" in repr(cany) assert "42" in repr(cany) def test_repr_none(self) -> None: """Repr shows None.""" cany = A(None).convert(None) assert "None" in repr(cany) def test_repr_float(self) -> None: """Repr shows float value.""" cany = A(float).convert(3.14) assert "float" in repr(cany) @requires_py39 def test_repr_object(self) -> None: """Repr shows type_index for objects.""" cany = A(tuple[int, ...]).convert([1, 2, 3]) assert "type_index" in repr(cany) def test_convert_raises_type_error(self) -> None: """Convert still raises TypeError for incompatible values.""" with pytest.raises(TypeError): A(int).convert("hello") def test_check_value_does_not_return_cany(self) -> None: """check_value returns None (not CAny).""" result = A(int).check_value(42) assert result is None # --------------------------------------------------------------------------- # from_annotation structural equality tests # --------------------------------------------------------------------------- class TestFromAnnotationScalars: """Scalar types — from_annotation produces correct TypeSchema.""" def test_int(self) -> None: """Int annotation.""" assert A(int) == S("int") def test_float(self) -> None: """Float annotation.""" assert A(float) == S("float") def test_bool(self) -> None: """Bool annotation.""" assert A(bool) == S("bool") def test_str(self) -> None: """Str annotation.""" assert A(str) == S("str") def test_bytes(self) -> None: """Bytes annotation.""" assert A(bytes) == S("bytes") def test_none_type(self) -> None: """type(None) annotation.""" assert A(type(None)) == S("None") def test_none_literal(self) -> None: """None annotation.""" assert A(None) == S("None") def test_any(self) -> None: """typing.Any annotation.""" assert A(typing.Any) == S("Any") def test_tvm_ffi_string(self) -> None: """tvm_ffi.String maps to str schema.""" assert A(tvm_ffi.core.String) == S("str") def test_tvm_ffi_bytes(self) -> None: """tvm_ffi.Bytes maps to bytes schema.""" assert A(tvm_ffi.core.Bytes) == S("bytes") class TestFromAnnotationFFITypes: """FFI container and object types.""" def test_array(self) -> None: """tvm_ffi.Array → canonical origin 'Array'.""" assert A(tvm_ffi.Array) == S("Array") def test_list(self) -> None: """tvm_ffi.List → same as A(list).""" assert A(tvm_ffi.List) == A(list) def test_map(self) -> None: """tvm_ffi.Map → canonical origin 'Map'.""" assert A(tvm_ffi.Map) == S("Map") def test_dict(self) -> None: """tvm_ffi.Dict → same as A(dict).""" assert A(tvm_ffi.Dict) == A(dict) def test_function(self) -> None: """tvm_ffi.Function → same as A(Callable).""" assert A(tvm_ffi.core.Function) == A(Callable) def test_object(self) -> None: """tvm_ffi.Object → canonical origin 'Object'.""" assert A(tvm_ffi.core.Object) == S("Object") def test_tensor(self) -> None: """tvm_ffi.Tensor → canonical origin 'Tensor'.""" assert A(tvm_ffi.Tensor) == S("Tensor") def test_dtype(self) -> None: """tvm_ffi.core.DataType → canonical origin 'dtype'.""" assert A(tvm_ffi.core.DataType) == S("dtype") def test_device(self) -> None: """tvm_ffi.Device → canonical origin 'Device'.""" assert A(tvm_ffi.Device) == S("Device") def test_ctypes_c_void_p(self) -> None: """ctypes.c_void_p → canonical origin 'ctypes.c_void_p'.""" assert A(ctypes.c_void_p) == S("ctypes.c_void_p") @requires_py39 def test_array_parameterized(self) -> None: """tvm_ffi.Array[int] cross-equivalent to tuple[int, ...].""" assert A(tvm_ffi.Array[int]) == A(tuple[int, ...]) @requires_py39 def test_list_parameterized(self) -> None: """tvm_ffi.List[str] cross-equivalent to list[str].""" assert A(tvm_ffi.List[str]) == A(list[str]) @requires_py39 def test_map_parameterized(self) -> None: """tvm_ffi.Map[str, float].""" assert A(tvm_ffi.Map[str, float]) == S("Map", S("str"), S("float")) @requires_py39 def test_dict_parameterized(self) -> None: """tvm_ffi.Dict[str, int] cross-equivalent to dict[str, int].""" assert A(tvm_ffi.Dict[str, int]) == A(dict[str, int]) @requires_py39 def test_array_too_many_args(self) -> None: """tvm_ffi.Array[int, str] raises TypeError.""" with pytest.raises(TypeError, match="requires 1"): A(tvm_ffi.Array[int, str]) # type: ignore[type-arg] @requires_py39 def test_list_too_many_args(self) -> None: """tvm_ffi.List[int, str] raises TypeError.""" with pytest.raises(TypeError, match="requires 1"): A(tvm_ffi.List[int, str]) # type: ignore[type-arg] @requires_py39 def test_dict_one_arg(self) -> None: """tvm_ffi.Dict[str] raises TypeError.""" with pytest.raises(TypeError, match="requires 2"): A(tvm_ffi.Dict[str]) # type: ignore[type-arg] @requires_py39 def test_dict_three_args(self) -> None: """tvm_ffi.Dict[str, int, float] raises TypeError.""" with pytest.raises(TypeError, match="requires 2"): A(tvm_ffi.Dict[str, int, float]) # type: ignore[type-arg] @requires_py39 def test_map_one_arg(self) -> None: """tvm_ffi.Map[str] raises TypeError.""" with pytest.raises(TypeError, match="requires 2"): A(tvm_ffi.Map[str]) # type: ignore[type-arg] def test_unregistered_cobject_errors(self) -> None: """Unregistered CObject subclass raises TypeError.""" with pytest.raises(TypeError, match="not registered"): A(tvm_ffi.core.CObject) class TestFromAnnotationCallable: """Callable annotation tests.""" def test_bare(self) -> None: """Bare Callable.""" assert A(Callable) == S("Callable") def test_bare_collections_abc(self) -> None: """Bare collections.abc.Callable.""" assert A(collections.abc.Callable) == S("Callable") def test_params(self) -> None: """Callable[[int, str], bool].""" assert A(Callable[[int, str], bool]) == S("Callable", S("bool"), S("int"), S("str")) def test_ellipsis(self) -> None: """Callable[..., int].""" assert A(Callable[..., int]) == S("Callable", S("int")) def test_no_params(self) -> None: """Callable[[], int].""" assert A(Callable[[], int]) == S("Callable", S("int")) class TestFromAnnotationList: """list[T] → List tests.""" def test_bare(self) -> None: """Bare list.""" assert A(list).origin == "List" @requires_py39 def test_int(self) -> None: """list[int].""" assert A(list[int]) == S("List", S("int")) @requires_py39 def test_nested(self) -> None: """list[list[int]].""" assert A(list[list[int]]) == S("List", S("List", S("int"))) class TestFromAnnotationDict: """dict[K, V] → Dict tests.""" def test_bare(self) -> None: """Bare dict.""" assert A(dict).origin == "Dict" @requires_py39 def test_str_int(self) -> None: """dict[str, int].""" assert A(dict[str, int]) == S("Dict", S("str"), S("int")) class TestFromAnnotationArray: """tuple[T, ...] → Array tests.""" @requires_py39 def test_int(self) -> None: """tuple[int, ...].""" assert A(tuple[int, ...]) == S("Array", S("int")) @requires_py39 def test_float(self) -> None: """tuple[float, ...].""" assert A(tuple[float, ...]) == S("Array", S("float")) class TestFromAnnotationTuple: """tuple[T1, T2] (fixed) tests.""" def test_bare(self) -> None: """Bare tuple.""" assert A(tuple).origin == "tuple" @requires_py39 def test_int_str(self) -> None: """tuple[int, str].""" assert A(tuple[int, str]) == S("tuple", S("int"), S("str")) @requires_py39 def test_empty(self) -> None: """tuple[()] stays distinct from bare tuple.""" assert A(tuple[()]) == TypeSchema("tuple", ()) class TestFromAnnotationOptional: """Optional[T] tests.""" def test_int(self) -> None: """Optional[int].""" assert A(Optional[int]) == S("Optional", S("int")) def test_union_with_none_becomes_optional(self) -> None: """Union[int, None] normalizes to Optional[int].""" assert A(Union[int, None]) == S("Optional", S("int")) @requires_py310 def test_pipe_syntax(self) -> None: """Int | None.""" assert A(eval("int | None")) == S("Optional", S("int")) class TestFromAnnotationUnion: """Union[T1, T2] tests.""" def test_int_str(self) -> None: """Union[int, str].""" assert A(Union[int, str]) == S("Union", S("int"), S("str")) def test_nested_union_flattening(self) -> None: """Nested unions flatten to a single Union schema.""" assert A(Union[int, Union[str, float]]) == S("Union", S("int"), S("str"), S("float")) @requires_py310 def test_pipe_syntax(self) -> None: """Int | str.""" assert A(eval("int | str")) == S("Union", S("int"), S("str")) class TestFromAnnotationObject: """Registered CObject subclasses.""" def test_test_int_pair(self) -> None: """TestIntPair annotation.""" assert A(TestIntPair) == S("testing.TestIntPair") def test_cxx_class_base(self) -> None: """_TestCxxClassBase annotation.""" assert A(_TestCxxClassBase) == S("testing.TestCxxClassBase") class TestFromAnnotationErrors: """from_annotation raises TypeError for unsupported annotations.""" def test_unsupported_type(self) -> None: """Complex is not supported.""" with pytest.raises(TypeError, match="Cannot convert"): A(complex) @requires_py39 def test_list_too_many_args(self) -> None: """list[int, int, float] raises.""" with pytest.raises(TypeError, match="list takes at most 1"): A(list[int, int, float]) # type: ignore[type-arg] @requires_py39 def test_dict_one_arg(self) -> None: """dict[str] raises.""" with pytest.raises(TypeError, match="dict requires 0 or 2"): A(dict[str]) # type: ignore[type-arg] # --------------------------------------------------------------------------- # Convert returns FFI containers # --------------------------------------------------------------------------- import tvm_ffi as _tvm_ffi class TestConvertReturnFFIContainers: """convert() returns ffi.Array/List/Map/Dict.""" @requires_py39 def test_array_from_list(self) -> None: """Array convert from Python list.""" result = _to_py_class_value(A(tuple[float, ...]).convert([1, 2, 3])) assert isinstance(result, _tvm_ffi.Array) assert list(result) == [1.0, 2.0, 3.0] @requires_py39 def test_list_from_list(self) -> None: """List convert from Python list.""" result = _to_py_class_value(A(list[int]).convert([1, 2, 3])) assert isinstance(result, _tvm_ffi.List) assert list(result) == [1, 2, 3] @requires_py39 def test_dict_from_dict(self) -> None: """Dict convert from Python dict.""" result = _to_py_class_value(A(dict[str, int]).convert({"a": 1})) assert isinstance(result, _tvm_ffi.Dict) @requires_py39 def test_map_from_dict(self) -> None: """Map convert from Python dict.""" result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert({"a": 1})) assert isinstance(result, _tvm_ffi.Map) @requires_py39 def test_array_passthrough(self) -> None: """ffi.Array input passes through unchanged.""" arr = _tvm_ffi.Array([1, 2, 3]) result = _to_py_class_value(A(tuple[int, ...]).convert(arr)) assert result.same_as(arr) @requires_py39 def test_list_passthrough(self) -> None: """ffi.List input passes through unchanged.""" lst = _tvm_ffi.List([1, 2, 3]) result = _to_py_class_value(A(list[int]).convert(lst)) assert result.same_as(lst) @requires_py39 def test_array_subclass_passthrough(self) -> None: """ffi.Array subclasses pass through unchanged.""" class MyArray(_tvm_ffi.Array): pass arr = MyArray([1, 2, 3]) result = _to_py_class_value(A(tuple[int, ...]).convert(arr)) assert result.same_as(arr) @requires_py39 def test_list_subclass_passthrough(self) -> None: """ffi.List subclasses pass through unchanged.""" class MyList(_tvm_ffi.List): pass lst = MyList([1, 2, 3]) result = _to_py_class_value(A(list[int]).convert(lst)) assert result.same_as(lst) @requires_py39 def test_map_subclass_passthrough(self) -> None: """ffi.Map subclasses pass through unchanged for Map[Any, Any].""" class MyMap(_tvm_ffi.Map): pass m = MyMap({"a": 1}) result = _to_py_class_value(A(tvm_ffi.Map[typing.Any, typing.Any]).convert(m)) assert result.same_as(m) @requires_py39 def test_dict_subclass_passthrough(self) -> None: """ffi.Dict subclasses pass through unchanged for Dict[Any, Any].""" class MyDict(_tvm_ffi.Dict): pass d = MyDict({"a": 1}) result = _to_py_class_value(A(dict[typing.Any, typing.Any]).convert(d)) assert result.same_as(d) @requires_py39 def test_nested_array_convert(self) -> None: """Nested array conversion.""" result = _to_py_class_value(A(tuple[tuple[int, ...], ...]).convert([[1, 2], [3, 4]])) assert isinstance(result, _tvm_ffi.Array) assert isinstance(result[0], _tvm_ffi.Array) # --------------------------------------------------------------------------- # FFI type guarantees: convert() always returns tvm_ffi types # --------------------------------------------------------------------------- class TestConvertToFFITypes: """convert() returns canonical FFI types for all value kinds.""" def test_short_str_is_string(self) -> None: """Short str (SmallStr) promotes to tvm_ffi.String.""" result = _to_py_class_value(A(str).convert("hi")) assert isinstance(result, tvm_ffi.core.String) assert result == "hi" def test_long_str_is_string(self) -> None: """Long str (kTVMFFIStr object) is tvm_ffi.String.""" long_s = "x" * 200 result = _to_py_class_value(A(str).convert(long_s)) assert isinstance(result, tvm_ffi.core.String) assert result == long_s def test_empty_str_is_string(self) -> None: """Empty str is tvm_ffi.String.""" result = _to_py_class_value(A(str).convert("")) assert isinstance(result, tvm_ffi.core.String) assert result == "" def test_short_bytes_is_bytes(self) -> None: """Short bytes (SmallBytes) promotes to tvm_ffi.Bytes.""" result = _to_py_class_value(A(bytes).convert(b"hi")) assert isinstance(result, tvm_ffi.core.Bytes) assert result == b"hi" def test_long_bytes_is_bytes(self) -> None: """Long bytes (kTVMFFIBytes object) is tvm_ffi.Bytes.""" long_b = b"x" * 200 result = _to_py_class_value(A(bytes).convert(long_b)) assert isinstance(result, tvm_ffi.core.Bytes) assert result == long_b def test_empty_bytes_is_bytes(self) -> None: """Empty bytes is tvm_ffi.Bytes.""" result = _to_py_class_value(A(bytes).convert(b"")) assert isinstance(result, tvm_ffi.core.Bytes) assert result == b"" def test_bytearray_converts_to_ffi_bytes(self) -> None: """Bytearray converts to tvm_ffi.Bytes.""" result = _to_py_class_value(A(bytes).convert(bytearray(b"hello"))) assert isinstance(result, tvm_ffi.core.Bytes) assert result == b"hello" def test_callable_is_function(self) -> None: """Callable converts to tvm_ffi.Function.""" result = _to_py_class_value(A(Callable).convert(lambda x: x)) assert isinstance(result, tvm_ffi.core.Function) @requires_py39 def test_array_is_ffi_array(self) -> None: """Array[int] converts to tvm_ffi.Array.""" result = _to_py_class_value(A(tuple[int, ...]).convert([1, 2])) assert isinstance(result, _tvm_ffi.Array) @requires_py39 def test_list_is_ffi_list(self) -> None: """List[int] converts to tvm_ffi.List.""" result = _to_py_class_value(A(list[int]).convert([1, 2])) assert isinstance(result, _tvm_ffi.List) @requires_py39 def test_map_is_ffi_map(self) -> None: """Map[str, int] converts to tvm_ffi.Map.""" result = _to_py_class_value(A(tvm_ffi.Map[str, int]).convert({"a": 1})) assert isinstance(result, _tvm_ffi.Map) @requires_py39 def test_dict_is_ffi_dict(self) -> None: """Dict[str, int] converts to tvm_ffi.Dict.""" result = _to_py_class_value(A(dict[str, int]).convert({"a": 1})) assert isinstance(result, _tvm_ffi.Dict) def test_int_is_int(self) -> None: """Int stays as int.""" result = _to_py_class_value(A(int).convert(42)) assert type(result) is int assert result == 42 def test_float_is_float(self) -> None: """Float stays as float.""" result = _to_py_class_value(A(float).convert(3.14)) assert type(result) is float assert result == 3.14 def test_bool_is_bool(self) -> None: """Bool stays as bool.""" result = _to_py_class_value(A(bool).convert(True)) assert result is True def test_none_is_none(self) -> None: """None stays as None.""" result = _to_py_class_value(A(None).convert(None)) assert result is None def test_object_is_cobject(self) -> None: """Object converts to CObject subclass.""" obj = TestIntPair(1, 2) result = _to_py_class_value(A(TestIntPair).convert(obj)) assert isinstance(result, tvm_ffi.core.CObject) assert result.same_as(obj) # --------------------------------------------------------------------------- # Small string/bytes optimization: verify CAny uses inline storage # --------------------------------------------------------------------------- # SmallStr/SmallBytes threshold is sizeof(int64_t) - 1 = 7 bytes. _SMALL_STR_TYPE_INDEX = 11 # kTVMFFISmallStr _SMALL_BYTES_TYPE_INDEX = 12 # kTVMFFISmallBytes _STR_TYPE_INDEX = 65 # kTVMFFIStr (heap-allocated String object) _BYTES_TYPE_INDEX = 66 # kTVMFFIBytes (heap-allocated Bytes object) class TestSmallStringOptimization: """Verify that short str/bytes use SmallStr/SmallBytes inline storage.""" def test_short_str_is_small_str(self) -> None: """A 2-byte string should be stored inline as SmallStr.""" cany = A(str).convert("hi") assert cany.type_index == _SMALL_STR_TYPE_INDEX def test_empty_str_is_small_str(self) -> None: """An empty string should be stored inline as SmallStr.""" cany = A(str).convert("") assert cany.type_index == _SMALL_STR_TYPE_INDEX def test_7_byte_str_is_small_str(self) -> None: """A 7-byte ASCII string (max) should be stored inline as SmallStr.""" cany = A(str).convert("abcdefg") assert cany.type_index == _SMALL_STR_TYPE_INDEX def test_8_byte_str_is_heap_str(self) -> None: """An 8-byte string exceeds SmallStr capacity and uses heap String.""" cany = A(str).convert("abcdefgh") assert cany.type_index == _STR_TYPE_INDEX def test_long_str_is_heap_str(self) -> None: """A long string uses heap-allocated String object.""" cany = A(str).convert("x" * 200) assert cany.type_index == _STR_TYPE_INDEX def test_short_bytes_is_small_bytes(self) -> None: """A 2-byte value should be stored inline as SmallBytes.""" cany = A(bytes).convert(b"hi") assert cany.type_index == _SMALL_BYTES_TYPE_INDEX def test_empty_bytes_is_small_bytes(self) -> None: """Empty bytes should be stored inline as SmallBytes.""" cany = A(bytes).convert(b"") assert cany.type_index == _SMALL_BYTES_TYPE_INDEX def test_7_byte_bytes_is_small_bytes(self) -> None: """7-byte value (max) should be stored inline as SmallBytes.""" cany = A(bytes).convert(b"abcdefg") assert cany.type_index == _SMALL_BYTES_TYPE_INDEX def test_8_byte_bytes_is_heap_bytes(self) -> None: """An 8-byte value exceeds SmallBytes capacity and uses heap Bytes.""" cany = A(bytes).convert(b"abcdefgh") assert cany.type_index == _BYTES_TYPE_INDEX def test_long_bytes_is_heap_bytes(self) -> None: """A long bytes value uses heap-allocated Bytes object.""" cany = A(bytes).convert(b"x" * 200) assert cany.type_index == _BYTES_TYPE_INDEX def test_bytearray_short_is_small_bytes(self) -> None: """A short bytearray should produce SmallBytes.""" cany = A(bytes).convert(bytearray(b"hi")) assert cany.type_index == _SMALL_BYTES_TYPE_INDEX def test_bytearray_long_is_heap_bytes(self) -> None: """A long bytearray should produce heap Bytes.""" cany = A(bytes).convert(bytearray(b"x" * 200)) assert cany.type_index == _BYTES_TYPE_INDEX def test_small_str_roundtrips_as_string(self) -> None: """SmallStr round-trips to tvm_ffi.String via _to_py_class_value.""" cany = A(str).convert("hi") assert cany.type_index == _SMALL_STR_TYPE_INDEX result = _to_py_class_value(cany) assert isinstance(result, tvm_ffi.core.String) assert result == "hi" def test_small_bytes_roundtrips_as_bytes(self) -> None: """SmallBytes round-trips to tvm_ffi.Bytes via _to_py_class_value.""" cany = A(bytes).convert(b"hi") assert cany.type_index == _SMALL_BYTES_TYPE_INDEX result = _to_py_class_value(cany) assert isinstance(result, tvm_ffi.core.Bytes) assert result == b"hi" def test_ffi_string_short_repacks_as_small(self) -> None: """A short tvm_ffi.String is re-packed as SmallStr by CAny.""" s = tvm_ffi.core.String("hi") cany = A(str).convert(s) assert cany.type_index == _SMALL_STR_TYPE_INDEX def test_ffi_string_long_stays_heap(self) -> None: """A long tvm_ffi.String stays as heap kTVMFFIStr.""" s = tvm_ffi.core.String("x" * 200) cany = A(str).convert(s) assert cany.type_index == _STR_TYPE_INDEX def test_ffi_bytes_short_repacks_as_small(self) -> None: """A short tvm_ffi.Bytes is re-packed as SmallBytes by CAny.""" b = tvm_ffi.core.Bytes(b"hi") cany = A(bytes).convert(b) assert cany.type_index == _SMALL_BYTES_TYPE_INDEX def test_ffi_bytes_long_stays_heap(self) -> None: """A long tvm_ffi.Bytes stays as heap kTVMFFIBytes.""" b = tvm_ffi.core.Bytes(b"x" * 200) cany = A(bytes).convert(b) assert cany.type_index == _BYTES_TYPE_INDEX tvm-ffi-0.1.12/tests/python/test_typed_method.py000066400000000000000000000306101521067262500217320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for ``@tvm_ffi.method`` — opt-in TypeMethod registration on ``@py_class``-decorated classes. """ from __future__ import annotations import itertools from typing import Any import pytest from tvm_ffi import Object, method from tvm_ffi.core import TypeInfo from tvm_ffi.dataclasses import py_class _counter = itertools.count() def _unique_key(base: str) -> str: """Return a globally unique type key so tests can re-register freely.""" return f"testing.method_dec.{base}_{next(_counter)}" def _find_method(info: TypeInfo, name: str) -> Any: """Return the ``TypeMethod`` entry for ``name`` or :data:`None`.""" return next((m for m in info.methods if m.name == name), None) def _toy_method_resolve(obj: Any, ref: str, *args: Any, **kwargs: Any) -> Any: """Test helper: resolve a ``"$method:NAME"`` string ref and call it. Parses the ``$method:`` prefix, looks up ``NAME`` in the instance's ``TypeInfo.methods`` table, and invokes the resolved ``Function`` on ``obj`` (plus any extra args). A successful return value means the whole ``@method`` → ``TVMFFITypeRegisterMethod`` → FFI callable chain is intact. The ``$method:`` prefix is a test-local convention for demonstrating name-based method lookup; it has no production meaning on its own. """ prefix = "$method:" if not ref.startswith(prefix): raise ValueError(f"Not a $method: ref: {ref!r}") name = ref[len(prefix) :] info = type(obj).__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute] m = _find_method(info, name) if m is None: raise LookupError( f"{type(obj).__name__}.{name}: not in TypeInfo.methods — " "was the method decorated with ``@tvm_ffi.method``?", ) return m.func(obj, *args, **kwargs) # --------------------------------------------------------------------------- # Registration — ``@method``-marked methods land in TypeInfo.methods # --------------------------------------------------------------------------- class TestMethodRegistration: """``@method`` drops the function's signature into ``TVMFFITypeRegisterMethod``; the name is resolvable from any FFI consumer. """ def test_instance_method_registered_and_ffi_callable(self) -> None: """A plain instance-style method registers with ``is_static=False`` and the returned FFI Function accepts the instance as arg 0. """ @py_class(_unique_key("Node")) class Node(Object): x: int @method def label(self) -> str: return f"N({self.x})" m = _find_method(Node.__tvm_ffi_type_info__, "label") # ty: ignore[unresolved-attribute] assert m is not None assert m.is_static is False # FFI call routes through the C method table — proves the # registration landed on the C side, not just the Python attr. assert m.func(Node(x=7)) == "N(7)" def test_staticmethod_registered_with_is_static_true(self) -> None: """``@method`` on top of ``@staticmethod`` marks the underlying function; the unwrap happens inside ``_collect_py_methods``. """ @py_class(_unique_key("Nstat")) class Nstat(Object): x: int @method @staticmethod def constant() -> int: return 42 m = _find_method(Nstat.__tvm_ffi_type_info__, "constant") # ty: ignore[unresolved-attribute] assert m is not None assert m.is_static is True assert m.func() == 42 def test_multiple_methods_all_registered(self) -> None: """Every ``@method``-marked callable appears in ``info.methods``.""" @py_class(_unique_key("NodeMulti")) class NodeMulti(Object): x: int @method def kind(self) -> str: return "multi" @method def double(self) -> int: return self.x * 2 @method def prefixed(self, p: str) -> str: return f"{p}-{self.x}" names = {m.name for m in NodeMulti.__tvm_ffi_type_info__.methods} # ty: ignore[unresolved-attribute] assert {"kind", "double", "prefixed"}.issubset(names) def test_no_decorator_no_registration(self) -> None: """Without ``@method``, a class-body function is a plain Python attribute — nothing reaches ``TypeInfo.methods``. Protects the opt-in contract: users aren't surprised by accidental FFI registration of helper methods. """ @py_class(_unique_key("NodeBare")) class NodeBare(Object): x: int def helper(self) -> int: # no @method return self.x assert _find_method(NodeBare.__tvm_ffi_type_info__, "helper") is None # ty: ignore[unresolved-attribute] def test_python_attribute_still_callable(self) -> None: """Registration doesn't shadow the Python attribute — callers can still invoke the method normally as ``instance.name(...)``. """ @py_class(_unique_key("NodeKeep")) class NodeKeep(Object): x: int @method def doubled(self) -> int: return self.x * 2 assert NodeKeep(x=5).doubled() == 10 # --------------------------------------------------------------------------- # End-to-end: name-based method lookup via ``TypeInfo.methods`` # --------------------------------------------------------------------------- class TestNameBasedResolution: """Resolving a ``@method``-decorated method by its name through ``TypeInfo.methods`` and calling it via the FFI path. This is the same code path any cross-language consumer (C++, Rust, future reflection-driven tooling) takes to invoke a Python-defined method by name. """ def test_name_lookup_invokes_decorated_method(self) -> None: """Name lookup → FFI call round-trip: look up ``label`` in ``TypeInfo.methods`` and invoke it via the resolved Function. """ @py_class(_unique_key("Op")) class Op(Object): kind: str @method def label(self) -> str: return f"op:{self.kind}" result = _toy_method_resolve(Op(kind="add"), "$method:label") assert result == "op:add" def test_name_lookup_threads_extra_args(self) -> None: """Extra positional / keyword arguments thread through the FFI Function unchanged — covers multi-argument shapes any consumer would need. """ @py_class(_unique_key("PrologueOp")) class PrologueOp(Object): kind: str @method def print_prologue(self, printer: Any, frame: Any) -> str: # Use the extra args so a missing pass-through would show up. return f"{printer}-{self.kind}-{frame}" op = PrologueOp(kind="add") assert _toy_method_resolve(op, "$method:print_prologue", "PR", "FR") == "PR-add-FR" def test_name_lookup_missing_method_raises_clearly(self) -> None: """Looking up a method name that was NOT decorated with ``@method`` raises at resolution time — the failure mode a user hits when they forget the decorator. """ @py_class(_unique_key("OpMiss")) class OpMiss(Object): kind: str def unmarked(self) -> str: # no @method — not registered return self.kind with pytest.raises(LookupError, match=r"not in TypeInfo\.methods"): _toy_method_resolve(OpMiss(kind="x"), "$method:unmarked") # --------------------------------------------------------------------------- # Validation — reserved names / wrong wrappers rejected at decoration # --------------------------------------------------------------------------- class TestMethodValidation: """``@method`` + the registration path both raise with clear, class-scoped messages when a name or wrapper is reserved. """ def test_rejects_classmethod(self) -> None: """``@classmethod``'s first-arg is the class, not the instance — breaks the packed-call convention. Rejected at decoration time (before py_class even sees the method). """ with pytest.raises(TypeError, match=r"@classmethod is not supported"): class _Bad: @method @classmethod def maker(cls) -> int: return 0 def test_rejects_classmethod_method_decorator_order_swap(self) -> None: """The decorator catches ``@method @classmethod`` but a user can bypass it by writing ``@classmethod @method`` — @method runs first on the bare function, then classmethod wraps the marked function. The collector must surface this with a clear error; without the guard, the entry would silently fail to register (Python 3.11+) or register as a malformed instance method. """ with pytest.raises(TypeError, match=r"wrapped by @classmethod"): @py_class(_unique_key("CMOrderBad")) class _CMOrderBad(Object): x: int @classmethod @method def maker(cls) -> int: return 0 def test_rejects_manually_marked_classmethod(self) -> None: """The decorator can also be bypassed by marking a function directly (``fn.__ffi_method__ = True``) and then wrapping it in ``classmethod``. The collector's classmethod check fires on the descriptor regardless of how the marker got there. """ def _maker(cls: type) -> int: return 0 _maker.__ffi_method__ = True # ty: ignore[unresolved-attribute] cm = classmethod(_maker) with pytest.raises(TypeError, match=r"wrapped by @classmethod"): @py_class(_unique_key("CMManualBad")) class _CMManualBad(Object): x: int maker = cm def test_rejects_non_callable(self) -> None: """``@method`` applied to a bare value (not a callable) raises.""" with pytest.raises(TypeError, match=r"expected a callable"): method(42) def test_rejects_reserved_ffi_prefix(self) -> None: """``__ffi_*`` names are routed through TypeAttrColumn — using ``@method`` on them is surely a user error (would silently double-register), so ``_collect_py_methods`` raises. """ with pytest.raises(NameError, match=r"reserved ``__ffi_`` prefix"): @py_class(_unique_key("RFfiPfx")) class _RFfiPfx(Object): x: int @method def __ffi_custom__(self) -> int: return 0 def test_rejects_typeattrcolumn_name(self) -> None: """Decorating a TypeAttrColumn dunder with ``@method`` is rejected — those are routed to ``TVMFFITypeRegisterAttr`` already, never to TypeMethod. """ with pytest.raises(NameError, match=r"TypeAttrColumn"): @py_class(_unique_key("RAttr")) class _RAttr(Object): x: int @method def __ffi_repr__(self, fn_repr: Any) -> str: return "r" def test_rejects_python_protocol_dunder(self) -> None: """``__len__`` / ``__iter__`` / etc. are reserved for Python semantics — cannot be FFI TypeMethods. """ with pytest.raises(NameError, match=r"Python protocol dunder"): @py_class(_unique_key("RDun")) class _RDun(Object): x: int @method def __len__(self) -> int: return 0 tvm-ffi-0.1.12/tests/python/utils/000077500000000000000000000000001521067262500167745ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/python/utils/filelock_worker.py000066400000000000000000000033431521067262500225320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import random import sys import time from tvm_ffi.utils.lockfile import FileLock def worker(worker_id: int, lock_path: str, counter_file: str) -> int: """Worker function that tries to acquire lock and increment counter.""" try: with FileLock(lock_path): # Critical section - read, increment, write counter with open(counter_file) as f: # noqa: PTH123 current_value = int(f.read().strip()) time.sleep(random.uniform(0.01, 0.1)) # Simulate some work with open(counter_file, "w") as f: # noqa: PTH123 f.write(str(current_value + 1)) print(f"Worker {worker_id}: success") return 0 except Exception as e: print(f"Worker {worker_id}: error: {e}") return 1 if __name__ == "__main__": worker_id = int(sys.argv[1]) lock_path = sys.argv[2] counter_file = sys.argv[3] sys.exit(worker(worker_id, lock_path, counter_file)) tvm-ffi-0.1.12/tests/python/utils/test_embed_cubin.py000066400000000000000000000214411521067262500226430ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for the embed_cubin utility.""" from __future__ import annotations import subprocess import sys import tempfile from pathlib import Path import pytest from tvm_ffi.utils.embed_cubin import embed_cubin try: import torch except ImportError: torch = None # ty: ignore[invalid-assignment] def _is_cuda_available() -> bool: """Check if CUDA is available for testing.""" if torch is None: return False return torch.cuda.is_available() def _create_test_object_file(obj_path: Path) -> None: """Create a simple test object file with TVM_FFI_EMBED_CUBIN macro usage. We avoid including the full header to prevent CUDA dependency in tests. Instead, we manually declare the external symbols that TVM_FFI_EMBED_CUBIN would create. """ cpp_code = """ // Manually declare the external symbols that will be provided by embed_cubin // This is what TVM_FFI_EMBED_CUBIN(test_cubin) would do, but without CUDA headers extern "C" { extern const unsigned char __tvm_ffi__cubin_test_cubin[]; extern const unsigned char __tvm_ffi__cubin_test_cubin_end[]; } // Simple function that references the embedded CUBIN symbols const void* get_cubin_start() { return __tvm_ffi__cubin_test_cubin; } const void* get_cubin_end() { return __tvm_ffi__cubin_test_cubin_end; } """ # Write C++ code to a temporary file cpp_file = obj_path.with_suffix(".cc") cpp_file.write_text(cpp_code) # Compile to object file # Try to find a C++ compiler compilers = ["g++", "clang++", "c++"] compiler = None for c in compilers: try: subprocess.run( [c, "--version"], check=True, capture_output=True, ) compiler = c break except (subprocess.CalledProcessError, FileNotFoundError): continue if compiler is None: pytest.skip("No C++ compiler found (tried g++, clang++, c++)") # ty: ignore[invalid-argument-type, too-many-positional-arguments] assert isinstance(compiler, str), "Compiler is not a string" compile_cmd: list[str] = [compiler, "-c", "-fPIC", "-o", str(obj_path), str(cpp_file)] try: subprocess.run(compile_cmd, check=True, capture_output=True) except subprocess.CalledProcessError as e: pytest.skip(f"Failed to compile test object file: {e.stderr.decode('utf-8')}") # ty: ignore[invalid-argument-type, too-many-positional-arguments] finally: # Clean up temporary C++ file cpp_file.unlink(missing_ok=True) def _create_test_cubin(cubin_path: Path) -> None: """Create a dummy CUBIN file for testing.""" # Create a simple binary file with some recognizable content cubin_path.write_bytes(b"DUMMY_CUBIN_DATA_FOR_TESTING" * 10) def _check_symbols_in_object(obj_path: Path, expected_symbols: list[str]) -> bool: """Check if expected symbols exist in the object file using nm.""" try: result = subprocess.run( ["nm", str(obj_path)], check=True, capture_output=True, text=True, ) nm_output = result.stdout for symbol in expected_symbols: if symbol not in nm_output: return False return True except (subprocess.CalledProcessError, FileNotFoundError): pytest.skip("nm tool not available") # ty: ignore[invalid-argument-type, too-many-positional-arguments] return False @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") def test_embed_cubin_basic() -> None: """Test basic embed_cubin functionality.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) # Create test files input_obj = temp_path / "input.o" output_obj = temp_path / "output.o" cubin_file = temp_path / "test.cubin" _create_test_object_file(input_obj) _create_test_cubin(cubin_file) # Embed CUBIN into object file embed_cubin( cubin_path=cubin_file, input_obj_path=input_obj, output_obj_path=output_obj, name="test_cubin", verbose=False, ) # Check that output file exists assert output_obj.exists(), "Output object file was not created" assert output_obj.stat().st_size > 0, "Output object file is empty" # Check that the output is larger than the input (it should contain both) assert output_obj.stat().st_size > input_obj.stat().st_size, ( "Output object file should be larger than input" ) # Check that expected symbols are present expected_symbols = [ "__tvm_ffi__cubin_test_cubin", "__tvm_ffi__cubin_test_cubin_end", ] assert _check_symbols_in_object(output_obj, expected_symbols), ( f"Expected symbols {expected_symbols} not found in output object file" ) @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") def test_embed_cubin_different_names() -> None: """Test embedding CUBIN with different names.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) # Create test files input_obj = temp_path / "input.o" output_obj = temp_path / "output.o" cubin_file = temp_path / "test.cubin" _create_test_object_file(input_obj) _create_test_cubin(cubin_file) # Test with a different name custom_name = "my_custom_kernel" embed_cubin( cubin_path=cubin_file, input_obj_path=input_obj, output_obj_path=output_obj, name=custom_name, verbose=False, ) # Check that symbols with custom name are present expected_symbols = [ f"__tvm_ffi__cubin_{custom_name}", f"__tvm_ffi__cubin_{custom_name}_end", ] assert _check_symbols_in_object(output_obj, expected_symbols), ( f"Expected symbols {expected_symbols} not found in output object file" ) def test_embed_cubin_nonexistent_input() -> None: """Test that embed_cubin raises error for nonexistent input files.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) input_obj = temp_path / "nonexistent_input.o" output_obj = temp_path / "output.o" cubin_file = temp_path / "nonexistent.cubin" # Test with nonexistent CUBIN file with pytest.raises(FileNotFoundError): embed_cubin( cubin_path=cubin_file, input_obj_path=input_obj, output_obj_path=output_obj, name="test", verbose=False, ) @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux") @pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available") def test_embed_cubin_verbose_mode() -> None: """Test that verbose mode doesn't crash.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = Path(temp_dir) # Create test files input_obj = temp_path / "input.o" output_obj = temp_path / "output.o" cubin_file = temp_path / "test.cubin" _create_test_object_file(input_obj) _create_test_cubin(cubin_file) # Run with verbose mode (should not crash) embed_cubin( cubin_path=cubin_file, input_obj_path=input_obj, output_obj_path=output_obj, name="test_cubin", verbose=True, # Test verbose mode ) assert output_obj.exists(), "Output object file was not created in verbose mode" tvm-ffi-0.1.12/tests/python/utils/test_filelock.py000066400000000000000000000103601521067262500221750ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Tests for the FileLock utility.""" import subprocess import sys import tempfile from pathlib import Path import pytest from tvm_ffi.utils.lockfile import FileLock def test_basic_lock_acquire_and_release() -> None: """Test basic lock acquisition and release.""" with tempfile.TemporaryDirectory() as temp_dir: lock_path = Path(temp_dir) / "test.lock" lock = FileLock(str(lock_path)) # Test acquire assert lock.acquire() is True assert lock._file_descriptor is not None # Test release lock.release() assert lock._file_descriptor is None def test_context_manager() -> None: """Test FileLock as a context manager.""" with tempfile.TemporaryDirectory() as temp_dir: lock_path = Path(temp_dir) / "test.lock" with FileLock(str(lock_path)) as lock: assert lock._file_descriptor is not None # Lock should be acquired assert lock_path.exists() # Lock should be released after exiting context # Note: file may still exist but should be unlocked def test_multiple_acquire_attempts() -> None: """Test multiple acquire attempts on the same lock instance.""" with tempfile.TemporaryDirectory() as temp_dir: lock_path = Path(temp_dir) / "test.lock" lock = FileLock(str(lock_path)) # First acquire should succeed assert lock.acquire() is True # Second acquire on same instance should fail # (can't acquire same lock twice) assert lock.acquire() is False lock.release() def test_exception_in_context_manager() -> None: """Test that lock is properly released even when exception occurs.""" with tempfile.TemporaryDirectory() as temp_dir: lock_path = Path(temp_dir) / "test.lock" # Test that exception is propagated and lock is released with pytest.raises(ValueError, match="test exception"): with FileLock(str(lock_path)) as lock: assert lock._file_descriptor is not None raise ValueError("test exception") # Lock should be released, so we can acquire it again lock2 = FileLock(str(lock_path)) assert lock2.acquire() is True lock2.release() def test_concurrent_access() -> None: """Test concurrent access from multiple processes.""" with tempfile.TemporaryDirectory() as temp_dir: lock_path = Path(temp_dir) / "test.lock" counter_file = Path(temp_dir) / "counter.txt" # Initialize counter file counter_file.write_text("0") # Create worker script content # Write worker script to a temporary file worker_script_path = Path(__file__).parent / "filelock_worker.py" # Run multiple worker processes concurrently num_workers = 16 processes = [] for i in range(num_workers): p = subprocess.Popen( [sys.executable, str(worker_script_path), str(i), str(lock_path), str(counter_file)] ) processes.append(p) # Wait for all processes to complete for p in processes: p.wait(timeout=60.0) assert p.returncode == 0, f"Worker process failed with return code {p.returncode}" # Check final counter value final_count = int(counter_file.read_text().strip()) # Counter should equal number of workers (no race conditions) assert final_count == num_workers if __name__ == "__main__": pytest.main([__file__]) tvm-ffi-0.1.12/tests/python/utils/test_kwargs_wrapper.py000066400000000000000000000405211521067262500234450ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import dataclasses import inspect from typing import Any import pytest from tvm_ffi.utils.kwargs_wrapper import make_kwargs_wrapper, make_kwargs_wrapper_from_signature def test_basic_wrapper() -> None: """Test basic wrapper functionality with various argument combinations.""" def target(*args: Any) -> int: return sum(args) # No defaults - all required wrapper = make_kwargs_wrapper(target, ["a", "b", "c"]) assert wrapper(1, 2, 3) == 6 assert wrapper(a=1, b=2, c=3) == 6 assert wrapper(1, b=2, c=3) == 6 # Single default argument wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(10,)) assert wrapper(1, 2) == 13 # c=10 assert wrapper(1, 2, 3) == 6 # c=3 explicit assert wrapper(1, 2, c=5) == 8 # c=5 via keyword # Multiple defaults (right-aligned) wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(20, 30)) assert wrapper(1) == 51 # b=20, c=30 assert wrapper(1, 2) == 33 # b=2, c=30 assert wrapper(1, 2, 3) == 6 # all explicit # All defaults wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(1, 2, 3)) assert wrapper() == 6 assert wrapper(10) == 15 assert wrapper(10, 20, 30) == 60 # Bound methods class Calculator: def __init__(self, base: int) -> None: self.base = base def add(self, a: int, b: int) -> int: return self.base + a + b calc = Calculator(100) wrapper = make_kwargs_wrapper(calc.add, ["a", "b"], arg_defaults=(5,)) assert wrapper(1) == 106 def test_keyword_only_arguments() -> None: """Test wrapper with keyword-only arguments.""" def target(*args: Any) -> int: return sum(args) # Optional keyword-only arguments (with defaults) wrapper = make_kwargs_wrapper( target, ["a", "b"], arg_defaults=(), kwonly_names=["c", "d"], kwonly_defaults={"c": 100, "d": 200}, ) assert wrapper(1, 2) == 303 # c=100, d=200 assert wrapper(1, 2, c=10) == 213 # d=200 assert wrapper(1, 2, c=10, d=20) == 33 wrapper = make_kwargs_wrapper( target, ["a", "b"], arg_defaults=(), kwonly_names=["c", "d"], kwonly_defaults={} ) assert wrapper(1, 2, c=10, d=20) == 33 # c and d are required wrapper = make_kwargs_wrapper( target, ["a", "b"], arg_defaults=(), kwonly_names=["c", "d"], kwonly_defaults={"d": 100}, ) assert wrapper(1, 2, c=10) == 113 # c required, d=100 assert wrapper(1, 2, c=10, d=20) == 33 # both explicit wrapper = make_kwargs_wrapper( target, ["a", "b", "c"], arg_defaults=(10,), kwonly_names=["d", "e"], kwonly_defaults={"d": 20, "e": 30}, ) assert wrapper(1, 2) == 63 # c=10, d=20, e=30 assert wrapper(1, 2, 5, d=15) == 53 # c=5 explicit, e=30 def test_validation_errors() -> None: """Test input validation and error handling.""" target = lambda *args: sum(args) # Duplicate positional argument names with pytest.raises(ValueError, match="Duplicate argument names found"): make_kwargs_wrapper(target, ["a", "b", "a"]) # Duplicate keyword-only argument names with pytest.raises(ValueError, match="Duplicate keyword-only argument names found"): make_kwargs_wrapper(target, ["a"], kwonly_names=["b", "c", "b"]) # Invalid argument name types with pytest.raises(TypeError, match="Argument name must be a string"): make_kwargs_wrapper(target, ["a", 123]) # ty: ignore[invalid-argument-type] # Invalid Python identifiers with pytest.raises(ValueError, match="not a valid Python identifier"): make_kwargs_wrapper(target, ["a", "b-c"]) # Python keywords cannot be used as parameter names with pytest.raises( ValueError, match="is a Python keyword and cannot be used as a parameter name" ): make_kwargs_wrapper(target, ["a", "if"]) # arg_defaults not a tuple with pytest.raises(TypeError, match="arg_defaults must be a tuple"): make_kwargs_wrapper(target, ["a", "b"], arg_defaults=[10]) # ty: ignore[invalid-argument-type] # arg_defaults too long with pytest.raises(ValueError, match=r"arg_defaults has .* values but only"): make_kwargs_wrapper(target, ["a"], arg_defaults=(1, 2, 3)) # Overlap between positional and keyword-only with pytest.raises(ValueError, match="cannot be both positional and keyword-only"): make_kwargs_wrapper(target, ["a", "b"], kwonly_names=["b"]) # kwonly_defaults key not in kwonly_names with pytest.raises(ValueError, match="not in kwonly_names"): make_kwargs_wrapper(target, ["a", "b"], kwonly_names=["c"], kwonly_defaults={"d": 10}) # Internal name conflict with pytest.raises(ValueError, match="conflict with internal names"): make_kwargs_wrapper(target, ["__i_target_func", "b"]) def test_special_default_values() -> None: """Test wrapper with special default values like None and objects.""" def target(a: Any, b: Any, c: Any) -> tuple[Any, Any, Any]: return (a, b, c) # None as default wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(None, None)) assert wrapper(1) == (1, None, None) # Complex objects as defaults (verify object reference is preserved) default_list = [1, 2, 3] wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(default_list, None)) result = wrapper(1) assert result[1] is default_list def test_wrapper_with_signature() -> None: """Test make_kwargs_wrapper_from_signature.""" target = lambda *args: sum(args) def source_func(a: Any, b: Any, c: int = 10, d: int = 20) -> None: """Source function documentation.""" pass sig = inspect.signature(source_func) wrapper = make_kwargs_wrapper_from_signature(target, sig) assert wrapper(1, 2) == 33 # 1 + 2 + 10 + 20 assert wrapper(1, 2, 3) == 26 # 1 + 2 + 3 + 20 assert wrapper(1, 2, 3, 4) == 10 # 1 + 2 + 3 + 4 # Test metadata preservation when prototype is provided wrapper_with_metadata = make_kwargs_wrapper_from_signature(target, sig, source_func) assert wrapper_with_metadata.__name__ == "source_func" # ty: ignore[unresolved-attribute] assert wrapper_with_metadata.__doc__ == "Source function documentation." # With keyword-only arguments def source_kwonly(a: Any, b: Any, *, c: int = 10, d: int = 20) -> None: pass wrapper = make_kwargs_wrapper_from_signature(target, inspect.signature(source_kwonly)) assert wrapper(1, 2) == 33 assert wrapper(1, 2, c=5, d=6) == 14 # With required keyword-only arguments def source_required_kwonly(a: Any, b: Any, *, c: Any, d: int = 20) -> None: pass wrapper = make_kwargs_wrapper_from_signature(target, inspect.signature(source_required_kwonly)) assert wrapper(1, 2, c=10) == 33 # c required, d=20 assert wrapper(1, 2, c=10, d=5) == 18 # both explicit # Reject *args and **kwargs def with_varargs(a: Any, *args: Any) -> None: pass with pytest.raises(ValueError, match=r"\*args not supported"): make_kwargs_wrapper_from_signature(target, inspect.signature(with_varargs)) def with_kwargs(a: Any, **kwargs: Any) -> None: pass with pytest.raises(ValueError, match=r"\*\*kwargs not supported"): make_kwargs_wrapper_from_signature(target, inspect.signature(with_kwargs)) # Test exclude_arg_names - ignore certain arguments from the signature def source_with_skip(a: Any, b: Any, c: int = 10, d: int = 20) -> None: pass wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_with_skip), exclude_arg_names=["c"] ) # c is ignored, so wrapper should only have a, b, d assert wrapper(1, 2) == 23 # 1 + 2 + 20 (d=20) assert wrapper(1, 2, d=5) == 8 # 1 + 2 + 5 # Test ignoring multiple arguments wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_with_skip), exclude_arg_names=["b", "d"] ) # b and d are ignored, so wrapper should only have a, c assert wrapper(1) == 11 # 1 + 10 (c=10) assert wrapper(1, c=5) == 6 # 1 + 5 # Test ignoring keyword-only arguments def source_kwonly_skip(a: Any, b: Any, *, c: int = 10, d: int = 20) -> None: pass wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_kwonly_skip), exclude_arg_names=["c"] ) # c is skipped, so wrapper should only have a, b, d assert wrapper(1, 2) == 23 # 1 + 2 + 20 (d=20) assert wrapper(1, 2, d=5) == 8 # 1 + 2 + 5 # Test excluding a non-existent argument (should be silently ignored) wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_with_skip), exclude_arg_names=["non_existent"] ) # Should be the same as no exclusion assert wrapper(1, 2) == 33 # 1 + 2 + 10 + 20 assert wrapper(1, 2, 3, 4) == 10 # 1 + 2 + 3 + 4 # Test excluding both existing and non-existent arguments wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_with_skip), exclude_arg_names=["c", "non_existent", "also_missing"], ) # Only c should be excluded, non-existent names are ignored assert wrapper(1, 2) == 23 # 1 + 2 + 20 (d=20, c excluded) assert wrapper(1, 2, d=5) == 8 # 1 + 2 + 5 def test_exception_propagation() -> None: """Test that exceptions from the target function are properly propagated.""" def raising_func(a: int, b: int, c: str) -> int: if a == 0: raise ValueError("a cannot be zero") if b < 0: raise RuntimeError(f"b must be non-negative, got {b}") if c != "valid": raise TypeError(f"c must be 'valid', got {c!r}") return a + b # Test with positional defaults wrapper = make_kwargs_wrapper(raising_func, ["a", "b", "c"], arg_defaults=(10, "valid")) assert wrapper(5) == 15 with pytest.raises(ValueError, match="a cannot be zero"): wrapper(0) with pytest.raises(RuntimeError, match="b must be non-negative"): wrapper(1, -5) # Test with keyword-only arguments wrapper_kwonly = make_kwargs_wrapper( raising_func, ["a"], kwonly_names=["b", "c"], kwonly_defaults={"b": 10, "c": "valid"}, ) assert wrapper_kwonly(5) == 15 with pytest.raises(ValueError, match="a cannot be zero"): wrapper_kwonly(0) with pytest.raises(RuntimeError, match="b must be non-negative"): wrapper_kwonly(5, b=-5) with pytest.raises(TypeError, match="c must be 'valid'"): wrapper_kwonly(5, c="invalid") def test_metadata_preservation() -> None: """Test that function metadata is preserved when prototype is provided.""" def my_function(x: int, y: int = 10) -> int: """Document the function.""" return x + y target = lambda *args: sum(args) wrapper = make_kwargs_wrapper(target, ["x", "y"], arg_defaults=(10,), prototype=my_function) assert wrapper.__name__ == "my_function" # ty: ignore[unresolved-attribute] assert wrapper.__doc__ == "Document the function." assert wrapper.__annotations__ == my_function.__annotations__ assert wrapper(5) == 15 def test_optimized_default_types() -> None: """Test that None, bool, and str defaults work correctly. This test verifies the optimization where None and bool defaults are directly embedded in the generated signature, while str defaults use the MISSING sentinel for safety. """ def target(*args: Any) -> tuple[Any, ...]: return args # Test None default (should be optimized - directly embedded) wrapper = make_kwargs_wrapper(target, ["a", "b", "c"], arg_defaults=(None,)) assert wrapper(1, 2) == (1, 2, None) assert wrapper(1, 2, 3) == (1, 2, 3) assert wrapper(1, 2, c=None) == (1, 2, None) # Test bool defaults (should be optimized - directly embedded) wrapper = make_kwargs_wrapper(target, ["a", "flag", "debug"], arg_defaults=(True, False)) assert wrapper(1) == (1, True, False) assert wrapper(1, False) == (1, False, False) assert wrapper(1, flag=False, debug=True) == (1, False, True) # Test str default (should use MISSING sentinel for safety) wrapper = make_kwargs_wrapper(target, ["a", "b", "name"], arg_defaults=("default",)) assert wrapper(1, 2) == (1, 2, "default") assert wrapper(1, 2, "custom") == (1, 2, "custom") assert wrapper(1, 2, name="custom") == (1, 2, "custom") # Test keyword-only with None, bool, and str wrapper = make_kwargs_wrapper( target, ["a"], kwonly_names=["b", "flag", "name"], kwonly_defaults={"b": None, "flag": True, "name": "default"}, ) assert wrapper(1) == (1, None, True, "default") assert wrapper(1, b=2) == (1, 2, True, "default") assert wrapper(1, flag=False) == (1, None, False, "default") assert wrapper(1, name="custom") == (1, None, True, "custom") assert wrapper(1, b=2, flag=False, name="test") == (1, 2, False, "test") def test_map_dataclass_to_tuple() -> None: """Test map_dataclass_to_tuple in make_kwargs_wrapper.""" @dataclasses.dataclass class Config: x: int y: int @dataclasses.dataclass class Nested: value: int cfg: Config def target(*args: Any) -> tuple[Any, ...]: return args # Basic: one dataclass arg converted wrapper = make_kwargs_wrapper(target, ["a", "cfg"], map_dataclass_to_tuple=["cfg"]) result = wrapper(1, Config(x=10, y=20)) assert result == (1, (10, 20)) # Dataclass passed as keyword argument result = wrapper(a=1, cfg=Config(x=3, y=4)) assert result == (1, (3, 4)) # Multiple dataclass args wrapper = make_kwargs_wrapper(target, ["a", "b"], map_dataclass_to_tuple=["a", "b"]) result = wrapper(Config(x=1, y=2), Config(x=3, y=4)) assert result == ((1, 2), (3, 4)) # Nested dataclass (auto-recursion via type annotations) wrapper = make_kwargs_wrapper(target, ["a", "nested"], map_dataclass_to_tuple=["nested"]) result = wrapper(1, Nested(value=5, cfg=Config(x=10, y=20))) assert result == (1, (5, (10, 20))) # Mixed: some args converted, others not wrapper = make_kwargs_wrapper(target, ["a", "cfg", "b"], map_dataclass_to_tuple=["cfg"]) result = wrapper(1, Config(x=10, y=20), 3) assert result == (1, (10, 20), 3) # With defaults default_cfg = Config(x=0, y=0) wrapper = make_kwargs_wrapper( target, ["a", "cfg"], arg_defaults=(default_cfg,), map_dataclass_to_tuple=["cfg"] ) result = wrapper(1) assert result == (1, (0, 0)) result = wrapper(1, Config(x=5, y=6)) assert result == (1, (5, 6)) # With keyword-only dataclass arg wrapper = make_kwargs_wrapper( target, ["a"], kwonly_names=["cfg"], kwonly_defaults={"cfg": Config(x=0, y=0)}, map_dataclass_to_tuple=["cfg"], ) result = wrapper(1) assert result == (1, (0, 0)) result = wrapper(1, cfg=Config(x=7, y=8)) assert result == (1, (7, 8)) # Empty list: no conversion wrapper = make_kwargs_wrapper(target, ["a", "b"], map_dataclass_to_tuple=[]) cfg = Config(x=1, y=2) result = wrapper(1, cfg) assert result == (1, cfg) assert result[1] is cfg # not converted # Works with make_kwargs_wrapper_from_signature def source_func(a: int, cfg: Config) -> None: pass wrapper = make_kwargs_wrapper_from_signature( target, inspect.signature(source_func), map_dataclass_to_tuple=["cfg"] ) result = wrapper(1, Config(x=10, y=20)) assert result == (1, (10, 20)) tvm-ffi-0.1.12/tests/python/utils/test_unpack_dataclass.py000066400000000000000000000226531521067262500237150ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from __future__ import annotations import dataclasses import sys from typing import Any import pytest from tvm_ffi.utils.unpack_dataclass import ( _extract_dataclass_to_tuple_schema, _validate_dataclass_to_tuple_schema, unpack_dataclass_to_tuple, ) # Module-level dataclass definitions for schema extraction tests. # Must be at module level so typing.get_type_hints() can resolve # cross-references with PEP 563 (from __future__ import annotations). @dataclasses.dataclass class _ExtractConfig: x: int y: int @dataclasses.dataclass class _ExtractNested: value: int cfg: _ExtractConfig @dataclasses.dataclass class _ExtractWithAny: data: Any scale: int @dataclasses.dataclass class _ExtractWithList: items: list[_ExtractConfig] scale: int @dataclasses.dataclass class _ExtractWithLeafList: values: list[int] name: str @dataclasses.dataclass class _ExtractWithDict: mapping: dict[str, _ExtractConfig] count: int @dataclasses.dataclass class _ExtractWithLeafDict: mapping: dict[str, int] count: int @dataclasses.dataclass class _ExtractWithTuple: pair: tuple[_ExtractConfig, int] flag: bool @dataclasses.dataclass class _ExtractWithOptional: value: int | None name: str @dataclasses.dataclass class _ExtractWithLeafListInt: items: list[int] scale: int def test_unpack_dataclass_to_tuple() -> None: """Test unpack_dataclass_to_tuple JIT-compiled unpacking.""" @dataclasses.dataclass class Config: x: int y: int @dataclasses.dataclass class Nested: value: int cfg: Config @dataclasses.dataclass class Deep: nested: Nested flag: bool # Flat dataclass assert unpack_dataclass_to_tuple(Config(x=1, y=2)) == (1, 2) # Nested dataclass (auto-recurses based on type annotation) assert unpack_dataclass_to_tuple(Nested(value=5, cfg=Config(x=10, y=20))) == (5, (10, 20)) # Deep nesting assert unpack_dataclass_to_tuple( Deep(nested=Nested(value=5, cfg=Config(x=10, y=20)), flag=True) ) == ((5, (10, 20)), True) # Leaf passthrough assert unpack_dataclass_to_tuple(42) == 42 assert unpack_dataclass_to_tuple("hello") == "hello" assert unpack_dataclass_to_tuple(None) is None # List recursion: list of dataclasses -> tuple of tuples assert unpack_dataclass_to_tuple([Config(x=1, y=2), Config(x=3, y=4)]) == [(1, 2), (3, 4)] # Tuple recursion assert unpack_dataclass_to_tuple((Config(x=1, y=2), 5)) == ((1, 2), 5) # Dict recursion (recurses values) assert unpack_dataclass_to_tuple({"a": Config(x=1, y=2), "b": 3}) == {"a": (1, 2), "b": 3} # Leaf values are NOT copied (no deep copy) class Holder: pass @dataclasses.dataclass class WithObj: obj: Any val: int h = Holder() result = unpack_dataclass_to_tuple(WithObj(obj=h, val=1)) assert result == (h, 1) assert result[0] is h # same object reference, no copy # Dynamic dispatch: Any-typed field receives a dataclass at runtime @dataclasses.dataclass class WithAnyField: data: Any scale: int # data is Any -> schema marks it as "unpack" -> __dispatch called at runtime # When data is a dataclass, it should be recursively unpacked result = unpack_dataclass_to_tuple(WithAnyField(data=Config(x=1, y=2), scale=3)) assert result == ((1, 2), 3) # When data is a plain value, passthrough result = unpack_dataclass_to_tuple(WithAnyField(data=42, scale=3)) assert result == (42, 3) # When data is a list of dataclasses, recurse each element result = unpack_dataclass_to_tuple( WithAnyField(data=[Config(x=1, y=2), Config(x=3, y=4)], scale=5) ) assert result == ([(1, 2), (3, 4)], 5) # When data is a nested dataclass result = unpack_dataclass_to_tuple( WithAnyField(data=Nested(value=10, cfg=Config(x=1, y=2)), scale=5) ) assert result == ((10, (1, 2)), 5) # When data is a dict with dataclass values result = unpack_dataclass_to_tuple( WithAnyField(data={"a": Config(x=1, y=2), "b": Config(x=3, y=4)}, scale=5) ) assert result == ({"a": (1, 2), "b": (3, 4)}, 5) # Self-referential dataclass (linked list): should not infinite recurse @dataclasses.dataclass class Node: value: int next: Node | None # Build a short linked list node = Node(value=1, next=Node(value=2, next=None)) result = unpack_dataclass_to_tuple(node) # The 'next' field is typed as Node|None, which on 3.10+ resolves to # a UnionType. The self-reference is caught by memo -> UNPACK -> dynamic dispatch. # Dynamic dispatch recursively unpacks the nested Node. assert result == (1, (2, None)) def test_validate_dataclass_to_tuple_schema() -> None: """Test internal schema validation.""" # Valid schemas _validate_dataclass_to_tuple_schema({"x": None, "y": None}) _validate_dataclass_to_tuple_schema({"cfg": {"x": None, "y": None}, "scale": None}) # Invalid: not a dict with pytest.raises(TypeError, match="must be a dict"): _validate_dataclass_to_tuple_schema([1, 2]) # type: ignore[arg-type] # Invalid: non-string key with pytest.raises(TypeError, match="must be a string"): _validate_dataclass_to_tuple_schema({123: None}) # Invalid: not a valid identifier with pytest.raises(ValueError, match="not a valid Python identifier"): _validate_dataclass_to_tuple_schema({"not-valid": None}) # Invalid: Python keyword with pytest.raises(ValueError, match="is a Python keyword"): _validate_dataclass_to_tuple_schema({"class": None}) # Invalid: nested schema with bad key with pytest.raises(ValueError, match="not a valid Python identifier"): _validate_dataclass_to_tuple_schema({"cfg": {"x": None, "y!": None}}) @pytest.mark.xfail( sys.version_info < (3, 10), reason="list[X]/dict[X,Y]/int|None not evaluable by get_type_hints on Python < 3.10", ) def test_extract_dataclass_to_tuple_schema() -> None: """Test schema extraction from dataclass types.""" # Flat: all known leaf types -> None schema = _extract_dataclass_to_tuple_schema(_ExtractConfig) assert schema == {"x": None, "y": None} # Nested: known dataclass field -> nested schema schema = _extract_dataclass_to_tuple_schema(_ExtractNested) assert schema == {"value": None, "cfg": {"x": None, "y": None}} # Any field -> "unpack" (dynamic dispatch) schema = _extract_dataclass_to_tuple_schema(_ExtractWithAny) assert schema == {"data": "unpack", "scale": None} # list[Config] -> "unpack" (container with dataclass element) schema = _extract_dataclass_to_tuple_schema(_ExtractWithList) assert schema == {"items": "unpack", "scale": None} # list[int] -> "unpack" (list must be converted to tuple per contract) schema = _extract_dataclass_to_tuple_schema(_ExtractWithLeafList) assert schema == {"values": "unpack", "name": None} # list[int] standalone field also gets UNPACK schema = _extract_dataclass_to_tuple_schema(_ExtractWithLeafListInt) assert schema == {"items": "unpack", "scale": None} # dict[str, Config] -> "unpack" (container with dataclass value type) schema = _extract_dataclass_to_tuple_schema(_ExtractWithDict) assert schema == {"mapping": "unpack", "count": None} # dict[str, int] -> None (container with only known leaf types) schema = _extract_dataclass_to_tuple_schema(_ExtractWithLeafDict) assert schema == {"mapping": None, "count": None} # tuple[Config, int] -> "unpack" (tuple containing a dataclass) schema = _extract_dataclass_to_tuple_schema(_ExtractWithTuple) assert schema == {"pair": "unpack", "flag": None} # Optional[int] (int | None) -> None (Union of known leaves) schema = _extract_dataclass_to_tuple_schema(_ExtractWithOptional) assert schema == {"value": None, "name": None} # Non-dataclass raises with pytest.raises(TypeError, match="Expected a dataclass class"): _extract_dataclass_to_tuple_schema(int) # Locally-defined classes with built-in type annotations resolve fine @dataclasses.dataclass class LocalConfig: x: int y: int schema = _extract_dataclass_to_tuple_schema(LocalConfig) assert schema == {"x": None, "y": None} # Locally-defined classes referencing other local classes can't be resolved # by get_type_hints — all fields fall back to UNPACK @dataclasses.dataclass class LocalNested: val: int cfg: LocalConfig schema = _extract_dataclass_to_tuple_schema(LocalNested) # get_type_hints fails for the class -> all fields become "unpack" assert schema == {"val": "unpack", "cfg": "unpack"} tvm-ffi-0.1.12/tests/scripts/000077500000000000000000000000001521067262500160025ustar00rootroot00000000000000tvm-ffi-0.1.12/tests/scripts/benchmark_dlpack.py000066400000000000000000000413331521067262500216300ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Benchmark API overhead of different python FFI API calling overhead through DLPack API. Specifically, we would like to understand the overall overhead python/C++ API calls. The general goal is to understand the overall space and get a sense of what are the possible operations. We pick function f(x, y, z) where x, y, z are length 1 tensors. The benchmark is running in eager mode so we can see what is possible. It is orthogonal to other optimizations. For example cudagraph can eliminate these overheads completely. So the goal is to get a sense of what is possible under eager mode. Summary of some takeaways: - numpy.add roughly takes 0.36 us per call, which gives roughly what can be done in python env. - torch.add on gpu takes about 3.7us per call, giving us an idea of what roughly we need to get to in eager mode. """ from __future__ import annotations import time from typing import Any, Callable, NamedTuple import numpy as np import torch import tvm_ffi class TestFFITensor: """Test FFI Tensor that exposes __tvm_ffi_object__ protocol.""" def __init__(self, tensor: tvm_ffi.Tensor) -> None: """Initialize the TestFFITensor.""" self._tensor = tensor def __tvm_ffi_object__(self) -> tvm_ffi.Tensor: """Implement __tvm_ffi_object__ protocol.""" return self._tensor class TestNamedTuple(NamedTuple): """Test FFI NamedTuple.""" x: torch.Tensor y: torch.Tensor z: torch.Tensor def print_speed(name: str, speed: float) -> None: print(f"{name:<60} {speed} sec/call") def print_error(name: str, error: Any) -> None: print(f"{name:<60} {error}") def baseline_torch_add(repeat: int) -> None: """Run torch.add with one element.""" def run_bench(device: str) -> None: x = torch.arange(1, device=device) y = torch.arange(1, device=device) z = torch.arange(1, device=device) torch.add(x, y, out=z) if device == "cuda": torch.cuda.synchronize() start = time.time() for i in range(repeat): torch.add(x, y, out=z) # note we deliberately do not use torch.cuda.synchronize() # because we want to see the overhead of the FFI call. end = time.time() print_speed(f"torch.add[{device}]", (end - start) / repeat) # rough take away: add on cuda roughly takes 3e-6 sec/call run_bench("cpu") run_bench("cuda") def baseline_numpy_add(repeat: int) -> None: """Run numpy.add with one element.""" x = np.arange(1) y = np.arange(1) z = np.arange(1) np.add(x, y, out=z) start = time.time() for i in range(repeat): np.add(x, y, out=z) end = time.time() speed = (end - start) / repeat print_speed("numpy.add", speed) def baseline_cupy_add(repeat: int) -> None: """Run cupy.add with one element.""" try: import cupy # noqa: PLC0415 except ImportError: # skip if cupy is not installed return x = cupy.arange(1) y = cupy.arange(1) z = cupy.arange(1) cupy.add(x, y, out=z) start = time.time() for i in range(repeat): cupy.add(x, y, out=z) end = time.time() speed = (end - start) / repeat print_speed("cupy.add", speed) def tvm_ffi_nop(repeat: int) -> None: """Overhead of tvm FFI python call via calling a NOP. testing.nop is defined in c++ and do nothing. """ nop = tvm_ffi.get_global_func("testing.nop") x = tvm_ffi.from_dlpack(torch.arange(1)) y = tvm_ffi.from_dlpack(torch.arange(1)) z = tvm_ffi.from_dlpack(torch.arange(1)) nop(x, y, z) start = time.time() for i in range(repeat): nop(x, y, z) end = time.time() print_speed("tvm_ffi.nop", (end - start) / repeat) def bench_ffi_nop_from_dlpack(name: str, x: Any, y: Any, z: Any, repeat: int) -> None: """Run dlpack conversion + tvm_ffi.nop. Measures overhead of running dlpack for each args then invoke """ nop = tvm_ffi.get_global_func("testing.nop") tx = tvm_ffi.from_dlpack(x) ty = tvm_ffi.from_dlpack(y) tz = tvm_ffi.from_dlpack(z) nop(tx, ty, tz) start = time.time() for i in range(repeat): tx = tvm_ffi.from_dlpack(x) ty = tvm_ffi.from_dlpack(y) tz = tvm_ffi.from_dlpack(z) nop(tx, ty, tz) end = time.time() print_speed(name, (end - start) / repeat) def tvm_ffi_nop_from_torch_dlpack(repeat: int) -> None: """Run dlpack conversion + tvm_ffi.nop. Measures overhead of running dlpack for each args then invoke """ x = torch.arange(1) y = torch.arange(1) z = torch.arange(1) bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(torch)", x, y, z, repeat) def tvm_ffi_nop_from_numpy_dlpack(repeat: int) -> None: """Run dlpack conversion + tvm_ffi.nop. Measures overhead of running dlpack for each args then invoke """ x = np.arange(1) y = np.arange(1) z = np.arange(1) bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(numpy)", x, y, z, repeat) def tvm_ffi_self_dlpack_nop(repeat: int) -> None: """Run dlpack conversion + tvm_ffi.nop. Measures overhead of running dlpack for each args then invoke """ x = tvm_ffi.from_dlpack(torch.arange(1)) y = tvm_ffi.from_dlpack(torch.arange(1)) z = tvm_ffi.from_dlpack(torch.arange(1)) bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(tvm)", x, y, z, repeat) def tvm_ffi_nop_from_torch_utils_to_dlpack(repeat: int) -> None: """Measures overhead of running dlpack for each args then invoke but uses the legacy torch.utils.dlpack.to_dlpack API. This helps to measure possible implementation overhead of torch. """ nop = tvm_ffi.get_global_func("testing.nop") x = torch.arange(1) y = torch.arange(1) z = torch.arange(1) tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x)) ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y)) tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z)) nop(tx, ty, tz) start = time.time() for i in range(repeat): tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x)) ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y)) tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z)) nop(tx, ty, tz) end = time.time() speed = (end - start) / repeat print_speed("tvm_ffi.nop+from_dlpack(torch.utils)", speed) def bench_tvm_ffi_nop_autodlpack(name: str, x: Any, y: Any, z: Any, repeat: int) -> None: """Measures overhead of running dlpack via auto convert by directly take torch.Tensor as inputs. """ nop = tvm_ffi.get_global_func("testing.nop") nop(x, y, z) start = time.time() for i in range(repeat): nop(x, y, z) end = time.time() speed = (end - start) / repeat print_speed(name, speed) def bench_tvm_ffi_nop_autodlpack_tuple(name: str, args: TestNamedTuple, repeat: int) -> None: """Measures overhead of running dlpack via auto convert by directly take torch.Tensor as inputs. """ nop = tvm_ffi.get_global_func("testing.nop") nop(args) start = time.time() for i in range(repeat): nop(args) end = time.time() speed = (end - start) / repeat print_speed(name, speed) def tvm_ffi_nop_autodlpack_from_torch( repeat: int, device: str = "cpu", stream: bool = False ) -> None: """Measures overhead of running dlpack via auto convert by directly take torch.Tensor as inputs. """ # use larger to ensure alignment req is met x = torch.arange(1, device=device) y = torch.arange(1, device=device) z = torch.arange(1, device=device) if stream: with torch.cuda.stream(torch.cuda.Stream()): bench_tvm_ffi_nop_autodlpack( f"tvm_ffi.nop.autodlpack(torch[{device}][stream])", x, y, z, repeat ) else: bench_tvm_ffi_nop_autodlpack(f"tvm_ffi.nop.autodlpack(torch[{device}])", x, y, z, repeat) def tvm_ffi_nop_autodlpack_from_numpy(repeat: int) -> None: """Measures overhead of running dlpack via auto convert by directly take numpy.ndarray as inputs. """ # use larger to ensure alignment req is met x = np.arange(256) y = np.arange(256) z = np.arange(256) bench_tvm_ffi_nop_autodlpack("tvm_ffi.nop.autodlpack(numpy)", x, y, z, repeat) def tvm_ffi_nop_autodlpack_from_dltensor_test_wrapper(repeat: int, device: str) -> None: """Measures overhead of running dlpack via auto convert by directly take test wrapper as inputs. This effectively measure DLPack exchange in tvm ffi. """ x = tvm_ffi.from_dlpack(torch.arange(1, device=device)) y = tvm_ffi.from_dlpack(torch.arange(1, device=device)) z = tvm_ffi.from_dlpack(torch.arange(1, device=device)) x = tvm_ffi.core.DLTensorTestWrapper(x) y = tvm_ffi.core.DLTensorTestWrapper(y) z = tvm_ffi.core.DLTensorTestWrapper(z) bench_tvm_ffi_nop_autodlpack( f"tvm_ffi.nop.autodlpack(DLTensorTestWrapper[{device}])", x, y, z, repeat ) def tvm_ffi_nop_autodlpack_from_test_tensor_namedtuple(repeat: int, device: str) -> None: """Measures overhead of running dlpack via auto convert by directly take test wrapper as inputs. This effectively measure DLPack exchange in tvm ffi. """ x = torch.arange(1, device=device) y = torch.arange(1, device=device) z = torch.arange(1, device=device) args = TestNamedTuple(x=x, y=y, z=z) bench_tvm_ffi_nop_autodlpack_tuple( f"tvm_ffi.nop.autodlpack(NamedTuple[{device}])", args, repeat ) def bench_to_dlpack(x: Any, name: str, repeat: int) -> None: x.__dlpack__() start = time.time() for i in range(repeat): x.__dlpack__() end = time.time() speed = (end - start) / repeat print_speed(name, speed) def bench_to_dlpack_versioned( x: Any, name: str, repeat: int, max_version: tuple[int, int] = (1, 1) ) -> None: """Measures overhead of running dlpack with latest 1.1.""" try: x.__dlpack__(max_version=max_version) start = time.time() for i in range(repeat): x.__dlpack__(max_version=max_version) end = time.time() speed = (end - start) / repeat print_speed(name, speed) except Exception as e: print_error(name, e) def bench_torch_utils_to_dlpack(repeat: int) -> None: """Measures overhead of running torch.utils.dlpack.to_dlpack.""" x = torch.arange(1) torch.utils.dlpack.to_dlpack(x) start = time.time() for i in range(repeat): torch.utils.dlpack.to_dlpack(x) end = time.time() speed = (end - start) / repeat print_speed("torch.utils.dlpack.to_dlpack", speed) def torch_get_cuda_stream_native(device_id: int) -> int: return torch.cuda.current_stream(device_id).cuda_stream def load_torch_get_current_cuda_stream() -> Callable[[int], int]: """Create a faster get_current_cuda_stream for torch through cpp extension.""" from torch.utils import cpp_extension # noqa: PLC0415 if torch.version.cuda is not None: source = """ #include int64_t get_current_cuda_stream(int device_id) { at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(device_id); // fast invariant, default stream is always 0 if (stream.id() == 0) return 0; // convert to cudaStream_t return reinterpret_cast(static_cast(stream)); } """ elif torch.version.hip is not None: source = """ #include int64_t get_current_cuda_stream(int device_id) { at::hip::HIPStream stream = at::hip::getCurrentHIPStream(device_id); // fast invariant, default stream is always 0 if (stream.id() == 0) return 0; // convert to hipStream_t return reinterpret_cast(static_cast(stream)); } """ result = cpp_extension.load_inline( name="get_current_cuda_stream", cpp_sources=[source], cuda_sources=[], extra_cflags=["-O3"], extra_include_paths=cpp_extension.include_paths("cuda"), functions=["get_current_cuda_stream"], ) return result.get_current_cuda_stream def bench_torch_get_current_stream(repeat: int, name: str, func: Callable[[int], int]) -> None: """Measures overhead of running torch.cuda.current_stream.""" x = torch.arange(1, device="cuda") # noqa: F841 func(0) start = time.time() for i in range(repeat): func(0) end = time.time() speed = (end - start) / repeat print_speed(f"torch.cuda.current_stream[{name}]", speed) def populate_object_table(num_classes: int) -> None: nop = tvm_ffi.get_global_func("testing.nop") dummy_instances = [type(f"DummyClass{i}", (object,), {})() for i in range(num_classes)] for instance in dummy_instances: nop(instance) def main() -> None: # noqa: PLR0915 repeat = 10000 # measures impact of object dispatch table size # takeaway so far is that there is no impact on the performance num_classes = 0 populate_object_table(num_classes) print("-----------------------------") print("Benchmark f(x, y, z) overhead") print("-----------------------------") baseline_numpy_add(repeat) baseline_torch_add(repeat) baseline_cupy_add(repeat) tvm_ffi_nop_from_torch_dlpack(repeat) tvm_ffi_nop_from_numpy_dlpack(repeat) tvm_ffi_self_dlpack_nop(repeat) tvm_ffi_nop_from_torch_utils_to_dlpack(repeat) tvm_ffi_nop_autodlpack_from_torch(repeat, "cpu") tvm_ffi_nop_autodlpack_from_torch(repeat, "cuda") tvm_ffi_nop_autodlpack_from_torch(repeat, "cuda", stream=True) tvm_ffi_nop_autodlpack_from_numpy(repeat) tvm_ffi_nop_autodlpack_from_dltensor_test_wrapper(repeat, "cpu") tvm_ffi_nop_autodlpack_from_dltensor_test_wrapper(repeat, "cuda") tvm_ffi_nop_autodlpack_from_test_tensor_namedtuple(repeat, "cpu") tvm_ffi_nop_autodlpack_from_test_tensor_namedtuple(repeat, "cuda") tvm_ffi_nop(repeat) print("-------------------------------") print("Benchmark x.__dlpack__ overhead") print("-------------------------------") bench_torch_utils_to_dlpack(repeat) bench_to_dlpack(torch.arange(1), "torch.__dlpack__", repeat) bench_to_dlpack(np.arange(1), "numpy.__dlpack__", repeat) bench_to_dlpack(tvm_ffi.from_dlpack(torch.arange(1)), "tvm.__dlpack__", repeat) print("---------------------------------------------------") print("Benchmark x.__dlpack__(max_version=(1,1)) overhead") print("---------------------------------------------------") bench_to_dlpack_versioned(torch.arange(1), "torch.__dlpack__(max_version=(1,1))", repeat) bench_to_dlpack_versioned(np.arange(1), "numpy.__dlpack__(max_version=(1,1))", repeat) bench_to_dlpack_versioned( tvm_ffi.from_dlpack(torch.arange(1)), "tvm.__dlpack__(max_version=(1,1))", repeat, ) print("---------------------------------------------------") print("Benchmark torch.get_cuda_stream[default stream]") print("---------------------------------------------------") bench_torch_get_current_stream(repeat, "cpp-extension", load_torch_get_current_cuda_stream()) bench_torch_get_current_stream(repeat, "python", torch_get_cuda_stream_native) print("---------------------------------------------------") print("Benchmark torch.get_cuda_stream[non-default stream]") print("---------------------------------------------------") with torch.cuda.stream(torch.cuda.Stream()): bench_torch_get_current_stream( repeat, "cpp-extension", load_torch_get_current_cuda_stream() ) bench_torch_get_current_stream(repeat, "python", torch_get_cuda_stream_native) print("---------------------------------------------------") print("Debug information") print("---------------------------------------------------") tvm_ffi.core._print_debug_info() release_gil = tvm_ffi.get_global_func("testing.nop").release_gil print(f"TVM_FFI_RELEASE_GIL_BY_DEFAULT={int(release_gil)}") print("---------------------------------------------------") if __name__ == "__main__": main() tvm-ffi-0.1.12/tests/scripts/benchmark_kwargs_wrapper.py000066400000000000000000000057221521067262500234320ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Benchmark API overhead of kwargs wrapper.""" from __future__ import annotations import dataclasses import time from typing import Any from tvm_ffi.utils.kwargs_wrapper import make_kwargs_wrapper def print_speed(name: str, speed: float) -> None: print(f"{name:<60} {speed} sec/call") def target_func(*args: Any) -> None: pass @dataclasses.dataclass class Config: x: int y: int z: int def benchmark_kwargs_wrapper(repeat: int = 1000000) -> None: """Benchmark kwargs wrapper with integer arguments and dataclass unpacking.""" x = 1 y = 2 z = 3 cfg = Config(x=x, y=y, z=z) wrapper = make_kwargs_wrapper(target_func, ["x", "y", "z"], arg_defaults=(None, None)) wrapper_dc = make_kwargs_wrapper(target_func, ["cfg"], map_dataclass_to_tuple=["cfg"]) # Warm up JIT cache wrapper_dc(cfg) # Direct call (baseline) start = time.time() for _ in range(repeat): target_func(x, y, z) end = time.time() print_speed("target_func(x, y, z)", (end - start) / repeat) # Wrapper with positional args start = time.time() for _ in range(repeat): wrapper(x, y, z) end = time.time() print_speed("wrapper(x, y, z)", (end - start) / repeat) # Wrapper with kwargs start = time.time() for _ in range(repeat): wrapper(x=x, y=y, z=z) end = time.time() print_speed("wrapper(x=x, y=y, z=z)", (end - start) / repeat) # Wrapper with defaults start = time.time() for _ in range(repeat): wrapper(x) end = time.time() print_speed("wrapper(x) [with defaults]", (end - start) / repeat) # Wrapper with dataclass unpack start = time.time() for _ in range(repeat): wrapper_dc(cfg) end = time.time() print_speed("wrapper_dc(cfg) [map_dataclass_to_tuple]", (end - start) / repeat) # Manual unpack (best possible) start = time.time() for _ in range(repeat): target_func(cfg.x, cfg.y, cfg.z) end = time.time() print_speed("target_func(cfg.x, cfg.y, cfg.z) [manual]", (end - start) / repeat) if __name__ == "__main__": print("Benchmarking kwargs_wrapper overhead...") print("-" * 90) benchmark_kwargs_wrapper() print("-" * 90) tvm-ffi-0.1.12/tests/scripts/benchmark_pycallback.py000066400000000000000000000076111521067262500225000ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Benchmark C++ -> Python callback overhead with 3 torch.Tensor arguments. Both variants are invoked by the same C++ ``invoke_n`` loop so the per-call cost reflects only the callback-arg conversion path: 1. ``convert_func(cb, tensor_cls=torch.Tensor)`` — the DLPack exchange API is threaded into the closure, so each tensor arg is converted to a ``torch.Tensor`` by the C-level callback arg setter before the callback runs. 2. ``convert_func(cb)`` — callback receives an ``ffi.Tensor`` and calls ``torch.from_dlpack(x)`` explicitly inside the callback body for each arg. Arguments are 3 x ``torch.zeros(1, device="cuda:0")``. """ from __future__ import annotations import time import torch import tvm_ffi import tvm_ffi.cpp from benchmark_dlpack import print_speed _INVOKE_N_CPP_SOURCE = r""" #include void invoke_n(tvm::ffi::Function callback, int64_t n, tvm::ffi::AnyView a, tvm::ffi::AnyView b, tvm::ffi::AnyView c) { for (int64_t i = 0; i < n; ++i) { callback(a, b, c); } } """ def _load_invoke_n() -> object: mod = tvm_ffi.cpp.load_inline( name="benchmark_pycallback_invoke_n", cpp_sources=_INVOKE_N_CPP_SOURCE, functions=["invoke_n"], ) return mod.get_function("invoke_n") def bench_pycallback_tensor_cls_torch(invoke_n, a, b, c, repeat: int) -> None: # noqa: ANN001 """convert_func(cb, tensor_cls=torch.Tensor): callback sees torch.Tensor directly.""" def cb(_a, _b, _c) -> None: # noqa: ANN001 pass callback = tvm_ffi.convert_func(cb, tensor_cls=torch.Tensor) invoke_n(callback, 10, a, b, c) start = time.time() invoke_n(callback, repeat, a, b, c) end = time.time() print_speed("pycallback[tensor_cls=torch.Tensor]", (end - start) / repeat) def bench_pycallback_from_dlpack(invoke_n, a, b, c, repeat: int) -> None: # noqa: ANN001 """convert_func(cb): callback receives ffi.Tensor, does torch.from_dlpack(x) explicitly.""" def cb(_a, _b, _c) -> None: # noqa: ANN001 torch.from_dlpack(_a) torch.from_dlpack(_b) torch.from_dlpack(_c) callback = tvm_ffi.convert_func(cb) invoke_n(callback, 10, a, b, c) start = time.time() invoke_n(callback, repeat, a, b, c) end = time.time() print_speed("pycallback+from_dlpack", (end - start) / repeat) def main() -> None: if not hasattr(torch.Tensor, "__dlpack_c_exchange_api__"): raise SystemExit("torch.Tensor.__dlpack_c_exchange_api__ not available") repeat = 10000 invoke_n = _load_invoke_n() a = torch.zeros(1, device="cuda:0") b = torch.zeros(1, device="cuda:0") c = torch.zeros(1, device="cuda:0") print("---------------------------------------------------") print("Benchmark C++ -> Python callback with 3 torch.Tensor args") print('Arguments: 3 x torch.zeros(1, device="cuda:0")') print("---------------------------------------------------") bench_pycallback_tensor_cls_torch(invoke_n, a, b, c, repeat) bench_pycallback_from_dlpack(invoke_n, a, b, c, repeat) print("---------------------------------------------------") if __name__ == "__main__": main() tvm-ffi-0.1.12/tests/scripts/benchmark_unpack_dataclass.py000066400000000000000000000117631521067262500236760ustar00rootroot00000000000000# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Benchmark unpack_dataclass_to_tuple vs dataclasses.astuple.""" from __future__ import annotations import dataclasses import time from typing import Any from tvm_ffi.utils.unpack_dataclass import unpack_dataclass_to_tuple def print_speed(name: str, speed: float) -> None: print(f"{name:<60} {speed} sec/call") def benchmark_unpack(unpack_fn: Any, val: Any, name: str, repeat: int = 1000000) -> None: """Benchmark an unpack function on a single value.""" unpack_fn(val) start = time.time() for _ in range(repeat): unpack_fn(val) end = time.time() print_speed(name, (end - start) / repeat) @dataclasses.dataclass class Config: x: int y: int @dataclasses.dataclass class Nested: value: int cfg: Config @dataclasses.dataclass class Wide: a: int b: int c: int d: int e: int @dataclasses.dataclass class Deep: nested: Nested flag: bool @dataclasses.dataclass class WithList: items: list[Config] scale: int @dataclasses.dataclass class WithAny: data: Any scale: int @dataclasses.dataclass class ConfigAny: x: Any y: Any def noop(x: Any) -> Any: return x if __name__ == "__main__": cfg = Config(x=1, y=2) nested = Nested(value=5, cfg=Config(x=10, y=20)) wide = Wide(a=1, b=2, c=3, d=4, e=5) deep = Deep(nested=Nested(value=5, cfg=Config(x=10, y=20)), flag=True) with_list = WithList(items=[Config(x=1, y=2), Config(x=3, y=4), Config(x=5, y=6)], scale=7) with_any_dc = WithAny(data=Config(x=1, y=2), scale=3) with_any_int = WithAny(data=42, scale=3) astuple = dataclasses.astuple # Correctness validation assert unpack_dataclass_to_tuple(cfg) == astuple(cfg) == (1, 2) assert unpack_dataclass_to_tuple(nested) == astuple(nested) == (5, (10, 20)) assert unpack_dataclass_to_tuple(wide) == astuple(wide) == (1, 2, 3, 4, 5) assert unpack_dataclass_to_tuple(deep) == astuple(deep) == ((5, (10, 20)), True) assert unpack_dataclass_to_tuple(with_list) == ([(1, 2), (3, 4), (5, 6)], 7) assert unpack_dataclass_to_tuple(with_any_dc) == ((1, 2), 3) assert unpack_dataclass_to_tuple(with_any_int) == (42, 3) assert unpack_dataclass_to_tuple(42) == 42 print("Benchmarking unpack_dataclass_to_tuple vs dataclasses.astuple...") print("-" * 90) benchmark_unpack(noop, cfg, "noop(Config) [baseline]") benchmark_unpack(noop, 42, "noop(int) [baseline]") benchmark_unpack(unpack_dataclass_to_tuple, 42, "unpack_dataclass_to_tuple(int) [leaf]") benchmark_unpack(unpack_dataclass_to_tuple, cfg, "unpack_dataclass_to_tuple(Config)") benchmark_unpack(astuple, cfg, "dataclasses.astuple(Config)") benchmark_unpack(unpack_dataclass_to_tuple, nested, "unpack_dataclass_to_tuple(Nested)") benchmark_unpack(astuple, nested, "dataclasses.astuple(Nested)") benchmark_unpack(unpack_dataclass_to_tuple, wide, "unpack_dataclass_to_tuple(Wide)") benchmark_unpack(astuple, wide, "dataclasses.astuple(Wide)") benchmark_unpack(unpack_dataclass_to_tuple, deep, "unpack_dataclass_to_tuple(Deep)") benchmark_unpack(astuple, deep, "dataclasses.astuple(Deep)") benchmark_unpack(unpack_dataclass_to_tuple, with_list, "unpack_dataclass_to_tuple(WithList)") benchmark_unpack(astuple, with_list, "dataclasses.astuple(WithList)") benchmark_unpack(unpack_dataclass_to_tuple, with_any_dc, "unpack(WithAny(Config)) [dynamic]") benchmark_unpack(astuple, with_any_dc, "dataclasses.astuple(WithAny(Config))") benchmark_unpack(unpack_dataclass_to_tuple, with_any_int, "unpack(WithAny(int)) [dynamic]") benchmark_unpack(astuple, with_any_int, "dataclasses.astuple(WithAny(int))") cfg_any = ConfigAny(x=1, y=2) cfg_any_nested = ConfigAny(x=Config(x=1, y=2), y=3) assert unpack_dataclass_to_tuple(cfg_any) == (1, 2) assert unpack_dataclass_to_tuple(cfg_any_nested) == ((1, 2), 3) benchmark_unpack(unpack_dataclass_to_tuple, cfg_any, "unpack(ConfigAny(1,2)) [all Any, leaf]") benchmark_unpack(astuple, cfg_any, "dataclasses.astuple(ConfigAny(1,2))") benchmark_unpack( unpack_dataclass_to_tuple, cfg_any_nested, "unpack(ConfigAny(Config,3)) [Any, nested]" ) benchmark_unpack(astuple, cfg_any_nested, "dataclasses.astuple(ConfigAny(Config,3))") print("-" * 90) tvm-ffi-0.1.12/tests/scripts/task_cpp_tests.sh000077500000000000000000000023531521067262500213720ustar00rootroot00000000000000#!/bin/bash # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. set -euxo pipefail BUILD_TYPE=RelWithDebugInfo rm -rf build/cpp_tests/CMakeCache.txt mkdir -p build/cpp_tests cmake -G Ninja -S . -B build/cpp_tests -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \ -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_CXX_COMPILER_LAUNCHER=ccache cmake --build build/cpp_tests --clean-first --config ${BUILD_TYPE} --target tvm_ffi_tests GTEST_COLOR=1 ctest -V -C ${BUILD_TYPE} --test-dir build/cpp_tests --output-on-failure