This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new abb9dae08 [FEAT] Split LLVM code generation into configurable backend
libraries (#72)
abb9dae08 is described below
commit abb9dae083a332470966b9d12b48348b826f4a8c
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 25 14:48:00 2026 -0400
[FEAT] Split LLVM code generation into configurable backend libraries (#72)
Split LLVM code generation into a shared backend linked to the core
compiler. Source builds retain all targets by default; the main wheel
flow builds a native-only backend and offers a matching optional
all-target wheel. Select one backend at startup using package-scoped
discovery and local loading.
Keep LLVM source selection, target setup and installation in the LLVM
CMake module. Dispatch target canonicalization and internal accelerator
fallback factories through the FFI registry so the core compiler does
not depend on the LLVM backend.
Prepare matching package versions and exact optional dependency pins,
build both distributions with cibuildwheel, and upload their artifacts
separately with byte-size summaries. Keep package publishing behind
explicit selection and separate PyPI environments.
---
.github/actions/build-wheel-for-publish/action.yml | 48 ++++++++++++--
.github/workflows/publish_wheel.yml | 73 +++++++++++++++++++---
3rdparty/tvm-ffi | 2 +-
CMakeLists.txt | 9 +--
ci/wheel/check_llvm_targets_wheel.py | 64 +++++++++++++++++++
ci/wheel/prepare_release_metadata.py | 67 ++++++++++++++++++++
ci/wheel/tvm-tirx-llvm-targets/CMakeLists.txt | 32 ++++++++++
ci/wheel/tvm-tirx-llvm-targets/pyproject.toml | 44 +++++++++++++
.../python/tvm_tirx_llvm_targets/__init__.py | 32 ++++++++++
cmake/modules/LLVM.cmake | 60 +++++++++++++++---
pyproject.toml | 7 ++-
python/tvm_tirx/base.py | 49 +++++++++++++++
src/backend/cuda/codegen/cuda_fallback_module.cc | 9 ++-
src/backend/cuda/codegen/cuda_fallback_module.h | 12 ++--
.../hexagon/codegen/hexagon_fallback_module.cc | 8 ++-
.../hexagon/codegen/hexagon_fallback_module.h | 13 ++--
.../hexagon/codegen/llvm/codegen_hexagon.cc | 6 +-
.../hexagon/codegen/llvm/intrin_rule_hexagon.cc | 6 +-
src/backend/hexagon/codegen/target_kind.cc | 16 -----
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 6 +-
src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc | 2 +-
src/backend/rocm/codegen/rocm_fallback_module.cc | 9 ++-
src/backend/rocm/codegen/rocm_fallback_module.h | 12 ++--
src/backend/rocm/codegen/target_kind.cc | 12 ----
src/codegen/canonicalizer/arm_mprofile.cc | 10 +++
src/codegen/llvm/canonicalizer/arm_aprofile.cc | 28 +++++++++
src/codegen/llvm/canonicalizer/canonicalize.cc | 21 +++----
src/codegen/llvm/llvm_instance.cc | 14 +++++
28 files changed, 562 insertions(+), 109 deletions(-)
diff --git a/.github/actions/build-wheel-for-publish/action.yml
b/.github/actions/build-wheel-for-publish/action.yml
index c44969d10..f3a05ed7c 100644
--- a/.github/actions/build-wheel-for-publish/action.yml
+++ b/.github/actions/build-wheel-for-publish/action.yml
@@ -17,8 +17,8 @@
name: Build TIRx Wheel
description: >
- Build and test the standalone LLVM-enabled apache-tvm-tirx wheel for a given
OS/architecture
- combination using cibuildwheel.
+ Build and test matching native TIRx and optional LLVM targets wheels for one
+ OS/architecture using cibuildwheel.
inputs:
arch:
@@ -39,6 +39,18 @@ inputs:
runs:
using: "composite"
steps:
+ - name: Set up release metadata Python
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
+ with:
+ python-version: "3.12"
+ activate-environment: true
+
+ - name: Prepare matching release metadata
+ shell: bash
+ run: |
+ uv pip install tomlkit setuptools-scm
+ python ci/wheel/prepare_release_metadata.py
+
# The {project} placeholder is not expanded inside CIBW_ENVIRONMENT
values, so
# compute the forward-slash host path here (the Windows analog of the Linux
# /project hardcode).
@@ -62,11 +74,11 @@ runs:
# Keep per-platform feature settings and the CUDA runtime sidecar.
CIBW_ENVIRONMENT_MACOS: >-
MACOSX_DEPLOYMENT_TARGET="14.0"
- CMAKE_ARGS="${{ inputs.cmake_defines }}"
+ CMAKE_ARGS="-DTVM_TIRX_LLVM_NATIVE_ONLY=ON ${{ inputs.cmake_defines
}}"
CIBW_ENVIRONMENT_LINUX: >-
- CMAKE_ARGS="${{ inputs.cmake_defines }}
-DTVM_TIRX_PACKAGE_EXTRA_LIBS=/project/build-wheel-cuda/lib/libtvm_runtime_cuda.so"
+ CMAKE_ARGS="-DTVM_TIRX_LLVM_NATIVE_ONLY=ON ${{ inputs.cmake_defines
}}
-DTVM_TIRX_PACKAGE_EXTRA_LIBS=/project/build-wheel-cuda/lib/libtvm_runtime_cuda.so"
CIBW_ENVIRONMENT_WINDOWS: >-
- CMAKE_ARGS="${{ inputs.cmake_defines }}
-DTVM_TIRX_PACKAGE_EXTRA_LIBS=${{ env.TVM_CUDA_EXTRA_LIB }}"
+ CMAKE_ARGS="-DTVM_TIRX_LLVM_NATIVE_ONLY=ON ${{ inputs.cmake_defines
}} -DTVM_TIRX_PACKAGE_EXTRA_LIBS=${{ env.TVM_CUDA_EXTRA_LIB }}"
# Turns the wheel-specific assertions in
# tests/python/runtime/test_validate_wheel_library.py
# ON (they skip unless these are set). Every published wheel is
LLVM-enabled,
@@ -75,3 +87,29 @@ runs:
CIBW_TEST_ENVIRONMENT: >-
TVM_WHEEL_EXPECT_LLVM="1"
TVM_WHEEL_EXPECT_CUDA_RUNTIME="${{ inputs.include_cuda_runtime }}"
+
+ - name: Build and test LLVM targets wheel
+ uses: pypa/cibuildwheel@b3395e271a35f73dad77f21802abe9b986dfe299 #
v4.3.1
+ with:
+ package-dir: ci/wheel/tvm-tirx-llvm-targets
+ config-file: pyproject.toml
+ output-dir: wheelhouse-llvm-targets
+ env:
+ CIBW_BUILD: ${{ inputs.build }}
+ CIBW_ARCHS_LINUX: ${{ inputs.arch }}
+ CIBW_ARCHS_MACOS: ${{ inputs.arch }}
+ CIBW_ARCHS_WINDOWS: ${{ inputs.arch }}
+ CIBW_ENVIRONMENT_MACOS: >-
+ MACOSX_DEPLOYMENT_TARGET="14.0"
+ CMAKE_ARGS="${{ inputs.cmake_defines }}"
+ CIBW_ENVIRONMENT_LINUX: >-
+ CMAKE_ARGS="${{ inputs.cmake_defines }}"
+ CIBW_ENVIRONMENT_WINDOWS: >-
+ CMAKE_ARGS="${{ inputs.cmake_defines }}"
+ # before-test runs before cibuildwheel installs the optional wheel, so
+ # its exact main dependency resolves to the wheel from this same build.
+ CIBW_BEFORE_TEST: >-
+ python "{project}/ci/wheel/check_llvm_targets_wheel.py"
--install-core "{project}/wheelhouse"
+ CIBW_TEST_COMMAND: >-
+ python -I "{project}/ci/wheel/check_llvm_targets_wheel.py" --wheel
"{wheel}" &&
+ python -I -m pytest -vvs
"{project}/tests/python/codegen/llvm/test_basic.py::test_llvm_vadd_pipeline"
diff --git a/.github/workflows/publish_wheel.yml
b/.github/workflows/publish_wheel.yml
index 3f97192c2..85c504105 100644
--- a/.github/workflows/publish_wheel.yml
+++ b/.github/workflows/publish_wheel.yml
@@ -117,7 +117,7 @@ jobs:
build_wheels:
name: ${{ matrix.name }}
# All-or-nothing: a failed sidecar leg blocks the whole wheel matrix (a
publish
- # needs the complete 4-wheel set anyway).
+ # needs both complete wheel sets anyway).
needs: [build_cuda_runtime]
runs-on: ${{ matrix.os }}
strategy:
@@ -168,7 +168,7 @@ jobs:
name: tirx-cuda-runtime-${{ matrix.arch }}
path: build-wheel-cuda/lib
- - name: Build and test TIRx core wheel
+ - name: Build and test native and LLVM targets wheels
uses: ./.github/actions/build-wheel-for-publish
with:
arch: ${{ matrix.arch }}
@@ -176,6 +176,24 @@ jobs:
cmake_defines: ${{ matrix.cmake_defines }}
include_cuda_runtime: ${{ matrix.include_cuda_runtime }}
+ - name: Report wheel sizes
+ shell: bash
+ run: |
+ python - <<'PYTHON'
+ import os
+ from pathlib import Path
+
+ lines = ["| Package | Wheel | Bytes | MiB |", "|---|---|---:|---:|"]
+ for package, directory in [("Main", "wheelhouse"), ("LLVM targets",
"wheelhouse-llvm-targets")]:
+ for wheel in sorted(Path(directory).glob("*.whl")):
+ size = wheel.stat().st_size
+ lines.append(f"| {package} | {wheel.name} | {size:,} | {size
/ 2**20:.2f} |")
+ summary = "\n".join(lines) + "\n"
+ print(summary)
+ with open(os.environ["GITHUB_STEP_SUMMARY"], "a", encoding="utf-8")
as output:
+ output.write(summary)
+ PYTHON
+
- name: Upload wheel artifact
uses: actions/[email protected]
with:
@@ -183,12 +201,53 @@ jobs:
path: wheelhouse/*.whl
if-no-files-found: error
- upload_pypi:
- name: Upload package distributions
+ - name: Upload LLVM targets wheel artifact
+ uses: actions/[email protected]
+ with:
+ name: tirx-llvm-targets-wheel-${{ matrix.artifact_suffix }}
+ path: wheelhouse-llvm-targets/*.whl
+ if-no-files-found: error
+
+ upload_llvm_targets:
+ name: Upload LLVM targets distributions
needs: [build_wheels]
- if: ${{ inputs.publish_repository != 'none' }}
+ if: ${{ inputs.publish_repository == 'pypi' }}
+ runs-on: ubuntu-latest
+ environment: pypi-llvm-targets
+ permissions:
+ id-token: write
+ attestations: write
+ steps:
+ - uses: actions/[email protected]
+ with:
+ pattern: tirx-llvm-targets-wheel-*
+ path: dist
+ merge-multiple: true
+
+ # Print wheel sizes for visibility only. We do not fail on the PyPI 100
MB
+ # limit here -- the publish step's upload would surface that error
directly.
+ - name: Print wheel sizes
+ shell: bash
+ run: ls -alh dist/*.whl
+
+ - name: Generate artifact attestation for wheels
+ uses: actions/[email protected]
+ with:
+ subject-path: dist/*
+
+ - name: Publish package distributions to PyPI
+ if: ${{ inputs.publish_repository == 'pypi' }}
+ uses:
pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3
+ with:
+ attestations: true
+ verbose: true
+
+ upload_pypi:
+ name: Upload native TIRx distributions
+ needs: [build_wheels, upload_llvm_targets]
+ if: ${{ inputs.publish_repository == 'pypi' }}
runs-on: ubuntu-latest
- environment: ${{ inputs.publish_repository }}
+ environment: pypi
permissions:
id-token: write
attestations: write
@@ -212,7 +271,7 @@ jobs:
- name: Publish package distributions to PyPI
if: ${{ inputs.publish_repository == 'pypi' }}
- uses:
pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0
+ uses:
pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3
with:
attestations: true
verbose: true
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 830f2ebac..dbdcdacee 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 830f2ebac90928f1375c882601599aa4e4fd3116
+Subproject commit dbdcdaceeef5b0b03251e8dea87725a1c1e2cde6
diff --git a/CMakeLists.txt b/CMakeLists.txt
index d9f965638..22b608d0b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -43,6 +43,7 @@ tvm_tirx_option(TVM_TIRX_USE_CUSTOM_LOGGING BOOL OFF "Use
user-defined logging i
# ######### LLVM ##########
tvm_tirx_option(TVM_TIRX_USE_LLVM STRING ON "Build with LLVM (OFF, ON, or
llvm-config command)")
+tvm_tirx_option(TVM_TIRX_LLVM_NATIVE_ONLY BOOL OFF "Link only the native LLVM
target")
# ######### Z3 ##########
@@ -157,8 +158,6 @@ file(GLOB CODEGEN_SRCS CONFIGURE_DEPENDS
src/codegen/c/*.cc
src/codegen/*.cc
src/codegen/canonicalizer/*.cc
- src/codegen/llvm/target_kind.cc
- src/codegen/llvm/canonicalizer/*.cc
src/backend/cuda/codegen/*.cc
src/backend/cuda/op/*.cc
src/backend/cuda/transform/*.cc
@@ -241,10 +240,6 @@ add_library(tvm_tirx_compiler::shared ALIAS
tvm_tirx_compiler)
target_link_libraries(tvm_tirx_compiler PUBLIC tvm_runtime tvm_tirx_header)
target_link_libraries(tvm_tirx_compiler PRIVATE ${TVM_LINKER_LIBS})
-# Avoid GNU ld relaxation of LLVM's GOT calls into incorrect direct calls.
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT "${TVM_TIRX_USE_LLVM}" MATCHES
${IS_FALSE_PATTERN})
- target_link_options(tvm_tirx_compiler PRIVATE "-Wl,--no-relax"
"-Wl,--as-needed")
-endif()
include(cmake/modules/Logging.cmake)
include(cmake/modules/Rust.cmake)
@@ -258,7 +253,7 @@ foreach(_target IN LISTS TVM_RUNTIME_BACKEND_LIBS)
tvm_ffi_add_prefix_map(${_target}_objs "${CMAKE_CURRENT_SOURCE_DIR}")
tvm_ffi_add_prefix_map(${_target} "${CMAKE_CURRENT_SOURCE_DIR}")
endforeach()
-foreach(_target tvm_runtime tvm_tirx_compiler ${TVM_RUNTIME_BACKEND_LIBS})
+foreach(_target tvm_runtime tvm_tirx_compiler ${TVM_RUNTIME_BACKEND_LIBS}
${TVM_LLVM_BACKEND_LIB})
if(TVM_TIRX_BUILD_PYTHON_MODULE)
if(APPLE)
set_property(TARGET ${_target} APPEND PROPERTY INSTALL_RPATH
"@loader_path/../../tvm_ffi/lib")
diff --git a/ci/wheel/check_llvm_targets_wheel.py
b/ci/wheel/check_llvm_targets_wheel.py
new file mode 100644
index 000000000..ef5f12051
--- /dev/null
+++ b/ci/wheel/check_llvm_targets_wheel.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Install the local core wheel or check the repaired optional wheel in CI."""
+
+import argparse
+import subprocess
+import sys
+import zipfile
+from pathlib import Path
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--install-core", type=Path)
+ parser.add_argument("--wheel", type=Path)
+ args = parser.parse_args()
+ if args.install_core:
+ wheels = list(args.install_core.glob("apache_tvm_tirx-*.whl"))
+ if len(wheels) != 1:
+ raise RuntimeError(f"Expected one local core wheel, found
{wheels}")
+ subprocess.run([sys.executable, "-m", "pip", "install",
str(wheels[0])], check=True)
+ return
+ if args.wheel is None:
+ parser.error("Pass --install-core or --wheel")
+
+ # The optional wheel must not bundle FFI or a second native LLVM backend.
+ with zipfile.ZipFile(args.wheel) as wheel:
+ forbidden = ("tvm_ffi", "llvm_backend_native")
+ for name in wheel.namelist():
+ if any(part in Path(name).name for part in forbidden):
+ raise RuntimeError(f"Optional wheel bundles an
FFI/native-backend library: {name}")
+
+ from importlib.metadata import version
+
+ import tvm_ffi
+ from tvm_tirx_llvm_targets import get_lib_path
+
+ from tvm_tirx import base
+
+ assert version("apache-tvm-tirx") ==
version("apache-tvm-tirx-llvm-targets")
+ loaded = Path(base._LOADED_LIBS["tvm_tirx_llvm_backend"]._name).resolve()
+ assert loaded == get_lib_path().resolve(), (loaded, get_lib_path())
+ targets = tvm_ffi.get_global_func("codegen.llvm_get_targets")()
+ assert {"x86_64", "aarch64"}.issubset(targets), targets
+ print(f"Selected optional LLVM backend {loaded}; {len(targets)} targets
available")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/ci/wheel/prepare_release_metadata.py
b/ci/wheel/prepare_release_metadata.py
new file mode 100644
index 000000000..40348f78c
--- /dev/null
+++ b/ci/wheel/prepare_release_metadata.py
@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Pin the optional LLVM wheel to the main wheel's resolved release version."""
+
+import argparse
+from pathlib import Path
+
+import tomlkit
+
+
+def resolve_version(root, metadata):
+ """Resolve the fixed version or the configured setuptools-scm source
once."""
+ project = metadata["project"]
+ if "version" not in project.get("dynamic", []):
+ return str(project["version"])
+ from setuptools_scm import get_version
+
+ return get_version(
+ **{
+ "relative_to": root / "pyproject.toml",
+ "dist_name": project["name"],
+ **metadata["tool"]["setuptools_scm"],
+ }
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("root", nargs="?", type=Path,
default=Path(__file__).resolve().parents[2])
+ root = parser.parse_args().root.resolve()
+ main_path = root / "pyproject.toml"
+ extra_path = root / "ci/wheel/tvm-tirx-llvm-targets/pyproject.toml"
+ main = tomlkit.parse(main_path.read_text(encoding="utf-8"))
+ extra = tomlkit.parse(extra_path.read_text(encoding="utf-8"))
+ version = resolve_version(root, main)
+ # Freeze SCM metadata before generated files can make the checkout dirty.
+ if "version" in main["project"].get("dynamic", []):
+ main["project"]["dynamic"].remove("version")
+ main["tool"]["scikit-build"].get("metadata", {}).pop("version", None)
+ main["project"]["version"] = version
+ main["project"]["optional-dependencies"]["all-targets"] = [
+ f"apache-tvm-tirx-llvm-targets=={version}"
+ ]
+ extra["project"]["version"] = version
+ extra["project"]["dependencies"] = [f"apache-tvm-tirx=={version}"]
+ main_path.write_text(tomlkit.dumps(main), encoding="utf-8")
+ extra_path.write_text(tomlkit.dumps(extra), encoding="utf-8")
+ print(f"Prepared main and LLVM targets wheels at version {version}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/ci/wheel/tvm-tirx-llvm-targets/CMakeLists.txt
b/ci/wheel/tvm-tirx-llvm-targets/CMakeLists.txt
new file mode 100644
index 000000000..199eed3cc
--- /dev/null
+++ b/ci/wheel/tvm-tirx-llvm-targets/CMakeLists.txt
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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.26)
+project(tvm_tirx_llvm_targets LANGUAGES C CXX)
+
+# Build the all-target library and install only its component. The matching
+# runtime/compiler comes from the main distribution dependency.
+set(TVM_TIRX_BUILD_PYTHON_MODULE OFF CACHE BOOL "Build Python package
artifacts" FORCE)
+set(TVM_TIRX_LLVM_NATIVE_ONLY OFF CACHE BOOL "Build only the native LLVM
target" FORCE)
+set(TVM_TIRX_BUILD_TESTS OFF CACHE BOOL "Build C++ tests")
+set(CMAKE_INSTALL_LIBDIR lib CACHE PATH "Library installation directory" FORCE)
+add_subdirectory("${CMAKE_CURRENT_SOURCE_DIR}/../../.." tvm-tirx)
+
+install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/../../../LICENSE"
+ "${CMAKE_CURRENT_SOURCE_DIR}/../../../NOTICE"
+ DESTINATION . COMPONENT tvm_tirx_llvm_backend)
+install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../../../licenses/"
+ DESTINATION licenses COMPONENT tvm_tirx_llvm_backend)
diff --git a/ci/wheel/tvm-tirx-llvm-targets/pyproject.toml
b/ci/wheel/tvm-tirx-llvm-targets/pyproject.toml
new file mode 100644
index 000000000..c64d71fb7
--- /dev/null
+++ b/ci/wheel/tvm-tirx-llvm-targets/pyproject.toml
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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-tirx-llvm-targets"
+version = "0.1.0a2"
+description = "Additional LLVM code generation targets for TVM TIRx"
+license = "Apache-2.0"
+requires-python = ">=3.10"
+dependencies = ["apache-tvm-tirx==0.1.0a2"]
+
+[build-system]
+requires = [
+ "scikit-build-core>=0.11",
+ "mlc-llvm-static[all]==23.1.1",
+ "mlc-z3-static==5.3.1.post0",
+]
+build-backend = "scikit_build_core.build"
+
+[tool.scikit-build]
+cmake.build-type = "Release"
+wheel.py-api = "py3"
+build-dir = "build"
+build.targets = ["tvm_tirx_llvm_backend"]
+install.components = ["tvm_tirx_llvm_backend"]
+wheel.packages = ["python/tvm_tirx_llvm_targets"]
+wheel.install-dir = "tvm_tirx_llvm_targets"
+
+[tool.scikit-build.cmake.define]
+TVM_TIRX_BUILD_PYTHON_MODULE = "OFF"
+TVM_TIRX_LLVM_NATIVE_ONLY = "OFF"
diff --git
a/ci/wheel/tvm-tirx-llvm-targets/python/tvm_tirx_llvm_targets/__init__.py
b/ci/wheel/tvm-tirx-llvm-targets/python/tvm_tirx_llvm_targets/__init__.py
new file mode 100644
index 000000000..b3ec369ea
--- /dev/null
+++ b/ci/wheel/tvm-tirx-llvm-targets/python/tvm_tirx_llvm_targets/__init__.py
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""All-target LLVM library selected when TVM TIRx is first imported."""
+
+import sys
+from pathlib import Path
+
+
+def get_lib_path() -> Path:
+ """Return the all-target LLVM library path without loading the library."""
+ name = (
+ "tvm_tirx_llvm_backend.dll"
+ if sys.platform == "win32"
+ else "libtvm_tirx_llvm_backend.dylib"
+ if sys.platform == "darwin"
+ else "libtvm_tirx_llvm_backend.so"
+ )
+ return Path(__file__).resolve().parent / "lib" / name
diff --git a/cmake/modules/LLVM.cmake b/cmake/modules/LLVM.cmake
index 3d660f864..2475cff3f 100644
--- a/cmake/modules/LLVM.cmake
+++ b/cmake/modules/LLVM.cmake
@@ -25,9 +25,20 @@ add_definitions(-DNDEBUG=1)
# TODO(@jroesch, @tkonolige): if we actually use targets we can do this.
# target_compile_definitions(tvm_compiler PRIVATE NDEBUG=1)
+# Target schema and canonicalizer dispatch do not depend on LLVM libraries.
+list(APPEND COMPILER_SRCS
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/codegen/llvm/target_kind.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/codegen/llvm/canonicalizer/canonicalize.cc")
+
+set(TVM_LLVM_BACKEND_LIB "")
+
# Test if ${TVM_TIRX_USE_LLVM} is not an explicit boolean false
# It may be a boolean or a string
if(NOT ${TVM_TIRX_USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
+ set(TVM_LLVM_BACKEND_LIB tvm_tirx_llvm_backend)
+ if(TVM_TIRX_LLVM_NATIVE_ONLY)
+ set(TVM_LLVM_BACKEND_LIB tvm_tirx_llvm_backend_native)
+ endif()
find_llvm(${TVM_TIRX_USE_LLVM})
if (${TVM_LLVM_VERSION} LESS 150)
message(FATAL_ERROR "LLVM version 15.0 or greater is required.")
@@ -52,16 +63,51 @@ if(NOT ${TVM_TIRX_USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
# Set flags that are only needed for LLVM target
add_definitions(-DTVM_LLVM_VERSION=${TVM_LLVM_VERSION})
add_definitions(-DTVM_LLVM_HAS_AARCH64_TARGET=${TVM_LLVM_HAS_AARCH64_TARGET})
- file(GLOB COMPILER_LLVM_SRCS CONFIGURE_DEPENDS src/codegen/llvm/*.cc
- src/backend/cuda/codegen/llvm/*.cc
- src/backend/rocm/codegen/llvm/*.cc
- src/backend/hexagon/codegen/llvm/*.cc
- )
+ file(GLOB COMPILER_LLVM_SRCS CONFIGURE_DEPENDS src/codegen/llvm/*.cc)
+ if(NOT TVM_TIRX_LLVM_NATIVE_ONLY)
+ file(GLOB ACCELERATOR_LLVM_SRCS CONFIGURE_DEPENDS
+ src/backend/cuda/codegen/llvm/*.cc
+ src/backend/rocm/codegen/llvm/*.cc
+ src/backend/hexagon/codegen/llvm/*.cc
+ )
+ list(APPEND COMPILER_LLVM_SRCS ${ACCELERATOR_LLVM_SRCS})
+ endif()
list(REMOVE_ITEM COMPILER_LLVM_SRCS
"${CMAKE_CURRENT_SOURCE_DIR}/src/codegen/llvm/target_kind.cc")
- list(APPEND TVM_LINKER_LIBS ${LLVM_LIBS})
- list(APPEND COMPILER_SRCS ${COMPILER_LLVM_SRCS})
if(NOT MSVC)
set_source_files_properties(${COMPILER_LLVM_SRCS}
PROPERTIES COMPILE_FLAGS "-fno-rtti")
endif()
+ # Keep the Arm canonicalizer with its LLVM implementation. With LLVM OFF,
+ # it remains in the core compiler for target parsing without LLVM.
+ list(APPEND COMPILER_LLVM_SRCS
+
"${CMAKE_CURRENT_SOURCE_DIR}/src/codegen/llvm/canonicalizer/arm_aprofile.cc")
+
+ # Native initialization references only the host target, so static LLVM
+ # archives contribute native code; otherwise initialize all configured
targets.
+ add_library(${TVM_LLVM_BACKEND_LIB} SHARED ${COMPILER_LLVM_SRCS})
+ target_link_libraries(${TVM_LLVM_BACKEND_LIB} PRIVATE
+ tvm_tirx_compiler tvm_llvm_header ${LLVM_LIBS})
+ target_compile_definitions(${TVM_LLVM_BACKEND_LIB} PRIVATE
+ TVM_TIRX_LLVM_NATIVE_ONLY=$<BOOL:${TVM_TIRX_LLVM_NATIVE_ONLY}>)
+ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+ # Avoid GNU ld relaxation of LLVM's GOT calls into incorrect direct calls.
+ target_link_options(${TVM_LLVM_BACKEND_LIB} PRIVATE "-Wl,--no-relax"
"-Wl,--as-needed")
+ endif()
+
+ tvm_tirx_configure_target(${TVM_LLVM_BACKEND_LIB})
+ tvm_ffi_add_prefix_map(${TVM_LLVM_BACKEND_LIB} "${CMAKE_CURRENT_SOURCE_DIR}")
+
+ # The optional package puts its DLL beside, rather than inside, tvm_tirx.
+ if(APPLE)
+ set_property(TARGET ${TVM_LLVM_BACKEND_LIB} APPEND PROPERTY INSTALL_RPATH
"@loader_path/../../tvm_tirx/lib" "@loader_path/../../tvm_ffi/lib")
+ elseif(UNIX)
+ set_property(TARGET ${TVM_LLVM_BACKEND_LIB} APPEND PROPERTY INSTALL_RPATH
"$ORIGIN/../../tvm_tirx/lib" "$ORIGIN/../../tvm_ffi/lib")
+ endif()
+ install(TARGETS ${TVM_LLVM_BACKEND_LIB}
+ ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx_llvm_backend
+ LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx_llvm_backend
+ RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT
tvm_tirx_llvm_backend)
+else()
+ list(APPEND COMPILER_SRCS
+
"${CMAKE_CURRENT_SOURCE_DIR}/src/codegen/llvm/canonicalizer/arm_aprofile.cc")
endif()
diff --git a/pyproject.toml b/pyproject.toml
index 76bcf3c33..7c47448c5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
- "apache-tvm-ffi>=0.1.15.post0",
+ "apache-tvm-ffi>=0.1.15.post1",
"apache-tvm-ffi-orcjit>=0.1.2",
"ml_dtypes",
"numpy",
@@ -52,6 +52,8 @@ Repository = "https://github.com/apache/tvm";
"Bug Tracker" = "https://github.com/apache/tvm/issues";
[project.optional-dependencies]
+# The release preparation script pins the matching LLVM targets wheel.
+all-targets = []
torch = ["torch"]
cuda = ["cuda-bindings"]
popen-pool = ["psutil", "cloudpickle"]
@@ -113,7 +115,7 @@ wheel.py-api = "py3"
build-dir = "build"
# Wheel configuration
-install.components = ["tvm_tirx"]
+install.components = ["tvm_tirx", "tvm_tirx_llvm_backend"]
wheel.packages = ["python/tvm_tirx"]
wheel.install-dir = "tvm_tirx"
@@ -147,6 +149,7 @@ logging.level = "INFO"
[tool.scikit-build.cmake.define]
TVM_TIRX_BUILD_PYTHON_MODULE = "ON"
TVM_TIRX_BUILD_TESTS = "OFF"
+TVM_TIRX_LLVM_NATIVE_ONLY = "OFF"
[tool.pytest.ini_options]
testpaths = ["tests"]
diff --git a/python/tvm_tirx/base.py b/python/tvm_tirx/base.py
index af8665125..8d0370187 100644
--- a/python/tvm_tirx/base.py
+++ b/python/tvm_tirx/base.py
@@ -18,6 +18,7 @@
# pylint: disable=invalid-name, import-outside-toplevel
"""Base library for TVM."""
+import importlib.util
import os
from pathlib import Path
@@ -61,6 +62,54 @@ if not _RUNTIME_ONLY:
"tvm_tirx", "tvm_tirx_compiler", "RTLD_LOCAL",
extra_lib_paths=_lib_paths
)
+
+def _load_llvm_backend(lib_paths):
+ """Load one generic or native backend from the compiler installation."""
+ package = "tvm_tirx"
+ target_name = "tvm_tirx_llvm_backend"
+ try:
+ tvm_ffi_libinfo._find_library_by_basename(
+ package,
+ target_name,
+ lib_paths,
+ use_fallback_build_path=False,
+ use_env_path=False,
+ )
+ except RuntimeError:
+ # Only a native LLVM installation opts into the optional targets
package.
+ # LLVM-OFF installations must not load it even when it is installed.
+ target_name = "tvm_tirx_llvm_backend_native"
+ try:
+ tvm_ffi_libinfo._find_library_by_basename(
+ package,
+ target_name,
+ lib_paths,
+ use_fallback_build_path=False,
+ use_env_path=False,
+ )
+ except RuntimeError:
+ return None
+ if importlib.util.find_spec("tvm_tirx_llvm_targets") is not None:
+ package = "apache-tvm-tirx-llvm-targets"
+ target_name = "tvm_tirx_llvm_backend"
+ tvm_ffi_libinfo._find_library_by_basename(
+ package,
+ target_name,
+ use_fallback_build_path=False,
+ use_env_path=False,
+ )
+ return tvm_ffi_libinfo.load_lib_ctypes(
+ package, target_name, "RTLD_LOCAL", extra_lib_paths=lib_paths
+ )
+
+
+if not _RUNTIME_ONLY:
+ _llvm_backend = _load_llvm_backend(
+ [Path(_LOADED_LIBS["tvm_tirx_compiler"]._name).resolve().parent]
+ )
+ if _llvm_backend is not None:
+ _LOADED_LIBS["tvm_tirx_llvm_backend"] = _llvm_backend
+
if _RUNTIME_ONLY:
from tvm_ffi import registry as _tvm_ffi_registry
diff --git a/src/backend/cuda/codegen/cuda_fallback_module.cc
b/src/backend/cuda/codegen/cuda_fallback_module.cc
index 7e3f9ba97..6a6604eef 100644
--- a/src/backend/cuda/codegen/cuda_fallback_module.cc
+++ b/src/backend/cuda/codegen/cuda_fallback_module.cc
@@ -22,8 +22,8 @@
* \brief CUDAFallbackModuleObj — codegen-time placeholder used when the CUDA
* runtime is not linked. Mirrors `CUDAModuleObj`'s save/load format
* byte-for-byte, including the backend configuration string.
- * Always compiled (independent of USE_CUDA); never registered as an
- * FFI factory or loader.
+ * Always compiled (independent of USE_CUDA); registered only under the
+ * fallback factory name, with no runtime loader.
*/
#include "cuda_fallback_module.h"
@@ -31,6 +31,7 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/support/bytes_io.h>
#include <string>
@@ -123,5 +124,9 @@ ffi::Module CUDAFallbackModuleCreate(ffi::Bytes code,
ffi::String fmt,
return ffi::Module(n);
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("ffi.Module.create.cuda_fallback",
CUDAFallbackModuleCreate);
+}
+
} // namespace target
} // namespace tvm
diff --git a/src/backend/cuda/codegen/cuda_fallback_module.h
b/src/backend/cuda/codegen/cuda_fallback_module.h
index 89ea39964..a79396818 100644
--- a/src/backend/cuda/codegen/cuda_fallback_module.h
+++ b/src/backend/cuda/codegen/cuda_fallback_module.h
@@ -24,7 +24,7 @@
* `CUDAModuleCreateWithFallback` is the ONLY entry point codegen uses to
* construct a CUDA `ffi::Module`. It tries the runtime-registered factory
* "ffi.Module.create.cuda" via the FFI registry; on miss it constructs a
- * `CUDAFallbackModuleObj` directly. The fallback exists so that codegen
+ * `CUDAFallbackModuleObj` through its fallback factory. The fallback
exists so that codegen
* can succeed on a build where the CUDA runtime is not linked. This
* setup is helpful for cross compilation where we compile on one env
* and run on another.
@@ -46,7 +46,7 @@ namespace target {
/*!
* \brief Construct a `CUDAFallbackModuleObj` directly (no FFI registry
- * round-trip). Used by `CUDAModuleCreateWithFallback` and by tests
+ * round-trip). Registered as the fallback factory and used by tests
* that explicitly want a fallback instance.
*/
ffi::Module CUDAFallbackModuleCreate(ffi::Bytes code, ffi::String fmt,
@@ -68,8 +68,8 @@ inline ffi::Module CUDAModuleCreateWithFallback(ffi::Bytes
code, ffi::String fmt
ffi::Map<ffi::String,
runtime::FunctionInfo> fmap,
ffi::Map<ffi::String,
ffi::String> source) {
if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) {
- return CUDAFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return
(*ffi::Function::GetGlobal("ffi.Module.create.cuda_fallback"))(code, fmt, fmap,
source)
+ .cast<ffi::Module>();
}
// Registry: "ffi.Module.create.cuda" — real CUDA runtime factory.
// Grep hint: grep -rn 'ffi.Module.create.cuda' src/
@@ -77,8 +77,8 @@ inline ffi::Module CUDAModuleCreateWithFallback(ffi::Bytes
code, ffi::String fmt
if (fcreate.has_value()) {
return (*fcreate)(code, fmt, fmap, source).cast<ffi::Module>();
}
- return CUDAFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return (*ffi::Function::GetGlobal("ffi.Module.create.cuda_fallback"))(code,
fmt, fmap, source)
+ .cast<ffi::Module>();
}
} // namespace target
diff --git a/src/backend/hexagon/codegen/hexagon_fallback_module.cc
b/src/backend/hexagon/codegen/hexagon_fallback_module.cc
index a2fcdb316..44e6e2a04 100644
--- a/src/backend/hexagon/codegen/hexagon_fallback_module.cc
+++ b/src/backend/hexagon/codegen/hexagon_fallback_module.cc
@@ -21,12 +21,13 @@
* \file hexagon_fallback_module.cc
* \brief HexagonFallbackModuleObj — codegen-time placeholder used when the
* Hexagon runtime is not linked. Preserves the serialized module
- * format for external runtimes; never registers an FFI factory or
loader.
+ * format for external runtimes; registers its fallback factory, but no
runtime loader.
*/
#include "hexagon_fallback_module.h"
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/support/bytes_io.h>
#include <string>
@@ -118,5 +119,10 @@ ffi::Module HexagonFallbackModuleCreate(ffi::Bytes code,
ffi::String fmt,
return ffi::Module(n);
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("ffi.Module.create.hexagon_fallback",
+ HexagonFallbackModuleCreate);
+}
+
} // namespace target
} // namespace tvm
diff --git a/src/backend/hexagon/codegen/hexagon_fallback_module.h
b/src/backend/hexagon/codegen/hexagon_fallback_module.h
index 463a58a13..eb252cf52 100644
--- a/src/backend/hexagon/codegen/hexagon_fallback_module.h
+++ b/src/backend/hexagon/codegen/hexagon_fallback_module.h
@@ -24,7 +24,7 @@
* `HexagonModuleCreateWithFallback` is the ONLY entry point codegen uses
* to construct a Hexagon `ffi::Module`. It tries the runtime-registered
* factory "ffi.Module.create.hexagon" via the FFI registry; on miss it
- * constructs a `HexagonFallbackModuleObj` directly. The fallback exists
+ * constructs a `HexagonFallbackModuleObj` through its fallback factory.
The fallback exists
* so that codegen can succeed on a build where the Hexagon runtime is
* not linked. This setup is helpful for cross compilation where we
* compile on one env and run on another.
@@ -50,7 +50,7 @@ using HexagonSourceMap = ffi::Map<ffi::String,
ffi::Variant<ffi::String, ffi::By
/*!
* \brief Construct a `HexagonFallbackModuleObj` directly (no FFI registry
- * round-trip). Used by `HexagonModuleCreateWithFallback` and by
+ * round-trip). Registered as the fallback factory and used by
* tests that explicitly want a fallback instance.
*/
ffi::Module HexagonFallbackModuleCreate(ffi::Bytes code, ffi::String fmt,
@@ -74,8 +74,9 @@ inline ffi::Module HexagonModuleCreateWithFallback(
ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String,
runtime::FunctionInfo> fmap,
HexagonSourceMap source) {
if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) {
- return HexagonFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return
(*ffi::Function::GetGlobal("ffi.Module.create.hexagon_fallback"))(code, fmt,
fmap,
+
source)
+ .cast<ffi::Module>();
}
// Registry: "ffi.Module.create.hexagon" — real Hexagon runtime factory.
// Grep hint: grep -rn 'ffi.Module.create.hexagon' src/
@@ -83,8 +84,8 @@ inline ffi::Module HexagonModuleCreateWithFallback(
if (fcreate.has_value()) {
return (*fcreate)(code, fmt, fmap, source).cast<ffi::Module>();
}
- return HexagonFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return
(*ffi::Function::GetGlobal("ffi.Module.create.hexagon_fallback"))(code, fmt,
fmap, source)
+ .cast<ffi::Module>();
}
} // namespace target
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index 54539f7b1..26673b65a 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -588,11 +588,7 @@ ffi::Module BuildHexagon(IRModule mod, Target target) {
ExtractFuncInfo(mod),
std::move(source));
}
-void RegisterHexagonCodegen() {
- static bool registered = false;
- if (registered) return;
- registered = true;
-
+TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("target.build.hexagon", BuildHexagon)
diff --git a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
index 1b0469ff6..c168431aa 100644
--- a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
@@ -83,11 +83,7 @@ inline PrimExpr DispatchTVMQHLWrapperFp16(const PrimExpr& e)
{
.as_or_throw<PrimExpr>();
}
-void RegisterHexagonIntrinRules() {
- static bool registered = false;
- if (registered) return;
- registered = true;
-
+TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("prim.fma")
.set_attr<FLowerIntrinsic>("hexagon.FLowerIntrinsic",
diff --git a/src/backend/hexagon/codegen/target_kind.cc
b/src/backend/hexagon/codegen/target_kind.cc
index 1e711b145..8ced6ce52 100644
--- a/src/backend/hexagon/codegen/target_kind.cc
+++ b/src/backend/hexagon/codegen/target_kind.cc
@@ -29,18 +29,6 @@
#include "memory_scope.h"
-namespace tvm {
-
-#ifdef TVM_LLVM_VERSION
-namespace codegen {
-void RegisterHexagonCodegen();
-namespace llvm {
-void RegisterHexagonIntrinRules();
-} // namespace llvm
-} // namespace codegen
-#endif
-} // namespace tvm
-
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::backend::hexagon::RegisterMemoryScopes();
using namespace tvm;
@@ -53,8 +41,4 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_option<int64_t>("num-cores")
.def_option<int64_t>("vtcm-capacity")
.set_default_keys({"hexagon", "cpu"});
-#ifdef TVM_LLVM_VERSION
- tvm::codegen::llvm::RegisterHexagonIntrinRules();
- tvm::codegen::RegisterHexagonCodegen();
-#endif
}
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index c686ee77e..e8ff0c8ce 100644
--- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -341,11 +341,7 @@ ffi::Module BuildAMDGPU(IRModule mod, Target target) {
ExtractFuncInfo(mod),
std::move(source));
}
-void RegisterAMDGPUCodegen() {
- static bool registered = false;
- if (registered) return;
- registered = true;
-
+TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("target.build.rocm", BuildAMDGPU)
diff --git a/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
b/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
index 8dfb26e6e..38c1fdeab 100644
--- a/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
+++ b/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
@@ -117,7 +117,7 @@ inline PrimExpr DispatchShuffle(const PrimExpr& e) {
namespace llvm {
using tirx::FLowerIntrinsic;
-void RegisterROCMIntrinRules() {
+TVM_FFI_STATIC_INIT_BLOCK() {
// dummy because we don't have the activemask
// clang-format off
OpDef("tirx.gpu_warp_activemask")
diff --git a/src/backend/rocm/codegen/rocm_fallback_module.cc
b/src/backend/rocm/codegen/rocm_fallback_module.cc
index f40ba56f1..505bf4d0c 100644
--- a/src/backend/rocm/codegen/rocm_fallback_module.cc
+++ b/src/backend/rocm/codegen/rocm_fallback_module.cc
@@ -22,13 +22,14 @@
* \brief ROCmFallbackModuleObj — codegen-time placeholder used when the ROCm
* runtime is not linked. Mirrors `ROCMModuleObj`'s save/load format
* byte-for-byte; see one-way comment in `SaveToBytes` below.
- * Always compiled (independent of USE_ROCM); never registered as an
- * FFI factory or loader.
+ * Always compiled (independent of USE_ROCM); registered only under the
+ * fallback factory name, with no runtime loader.
*/
#include "rocm_fallback_module.h"
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/support/bytes_io.h>
#include <string>
@@ -121,5 +122,9 @@ ffi::Module ROCmFallbackModuleCreate(ffi::Bytes code,
ffi::String fmt,
return ffi::Module(n);
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("ffi.Module.create.rocm_fallback",
ROCmFallbackModuleCreate);
+}
+
} // namespace target
} // namespace tvm
diff --git a/src/backend/rocm/codegen/rocm_fallback_module.h
b/src/backend/rocm/codegen/rocm_fallback_module.h
index 763e49c44..90fab8549 100644
--- a/src/backend/rocm/codegen/rocm_fallback_module.h
+++ b/src/backend/rocm/codegen/rocm_fallback_module.h
@@ -24,7 +24,7 @@
* `ROCmModuleCreateWithFallback` is the ONLY entry point codegen uses to
* construct a ROCm `ffi::Module`. It tries the runtime-registered factory
* "ffi.Module.create.rocm" via the FFI registry; on miss it constructs a
- * `ROCmFallbackModuleObj` directly. The fallback exists so that codegen
+ * `ROCmFallbackModuleObj` through its fallback factory. The fallback
exists so that codegen
* can succeed on a build where the ROCm runtime is not linked. This
* setup is helpful for cross compilation where we compile on one env
* and run on another.
@@ -46,7 +46,7 @@ namespace target {
/*!
* \brief Construct a `ROCmFallbackModuleObj` directly (no FFI registry
- * round-trip). Used by `ROCmModuleCreateWithFallback` and by tests
+ * round-trip). Registered as the fallback factory and used by tests
* that explicitly want a fallback instance.
*/
ffi::Module ROCmFallbackModuleCreate(ffi::Bytes code, ffi::String fmt,
@@ -68,8 +68,8 @@ inline ffi::Module ROCmModuleCreateWithFallback(ffi::Bytes
code, ffi::String fmt
ffi::Map<ffi::String,
runtime::FunctionInfo> fmap,
ffi::Map<ffi::String,
ffi::String> source) {
if (tvm::support::GetEnv<bool>("TVM_TIRX_COMPILE_FORCE_FALLBACK", false)) {
- return ROCmFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return
(*ffi::Function::GetGlobal("ffi.Module.create.rocm_fallback"))(code, fmt, fmap,
source)
+ .cast<ffi::Module>();
}
// Registry: "ffi.Module.create.rocm" — real ROCm runtime factory.
// Grep hint: grep -rn 'ffi.Module.create.rocm' src/
@@ -77,8 +77,8 @@ inline ffi::Module ROCmModuleCreateWithFallback(ffi::Bytes
code, ffi::String fmt
if (fcreate.has_value()) {
return (*fcreate)(code, fmt, fmap, source).cast<ffi::Module>();
}
- return ROCmFallbackModuleCreate(std::move(code), std::move(fmt),
std::move(fmap),
- std::move(source));
+ return (*ffi::Function::GetGlobal("ffi.Module.create.rocm_fallback"))(code,
fmt, fmap, source)
+ .cast<ffi::Module>();
}
} // namespace target
diff --git a/src/backend/rocm/codegen/target_kind.cc
b/src/backend/rocm/codegen/target_kind.cc
index d8c90de2c..dae3d8e78 100644
--- a/src/backend/rocm/codegen/target_kind.cc
+++ b/src/backend/rocm/codegen/target_kind.cc
@@ -111,14 +111,6 @@ ffi::Map<ffi::String, ffi::Any>
UpdateROCmAttrs(ffi::Map<ffi::String, ffi::Any>
} // namespace rocm
} // namespace backend
-namespace codegen {
-#ifdef TVM_LLVM_VERSION
-void RegisterAMDGPUCodegen();
-namespace llvm {
-void RegisterROCMIntrinRules();
-} // namespace llvm
-#endif
-} // namespace codegen
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -139,8 +131,4 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_option<int64_t>("thread_warp_size", refl::DefaultValue(64))
.set_default_keys({"rocm", "gpu"})
.set_target_canonicalizer(UpdateROCmAttrs);
-#ifdef TVM_LLVM_VERSION
- tvm::codegen::llvm::RegisterROCMIntrinRules();
- tvm::codegen::RegisterAMDGPUCodegen();
-#endif
}
diff --git a/src/codegen/canonicalizer/arm_mprofile.cc
b/src/codegen/canonicalizer/arm_mprofile.cc
index 86012f7ea..51b047bb0 100644
--- a/src/codegen/canonicalizer/arm_mprofile.cc
+++ b/src/codegen/canonicalizer/arm_mprofile.cc
@@ -24,6 +24,8 @@
#include "arm_mprofile.h"
+#include <tvm/ffi/reflection/registry.h>
+
#include <string>
namespace tvm {
@@ -139,6 +141,14 @@ ffi::Map<ffi::String, ffi::Any>
Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
return target;
}
+// CPU-only targets may omit mtriple, including existing Cortex-M
configurations.
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("codegen.llvm.canonicalize_target.default",
+ [](ffi::Map<ffi::String, ffi::Any> attrs) {
+ return IsArch(attrs) ?
Canonicalize(attrs) : attrs;
+ });
+}
+
} // namespace mprofile
} // namespace canonicalizer
} // namespace target
diff --git a/src/codegen/llvm/canonicalizer/arm_aprofile.cc
b/src/codegen/llvm/canonicalizer/arm_aprofile.cc
index 3f8529b28..cf1b3d78b 100644
--- a/src/codegen/llvm/canonicalizer/arm_aprofile.cc
+++ b/src/codegen/llvm/canonicalizer/arm_aprofile.cc
@@ -24,6 +24,7 @@
#include "arm_aprofile.h"
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <memory>
@@ -154,6 +155,33 @@ ffi::Map<ffi::String, ffi::Any>
Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
return target;
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ auto canonicalize = [](ffi::Map<ffi::String, ffi::Any> attrs) {
+ // Cortex-M handling remains in the core and does not require LLVM.
+ auto default_canonicalize =
+ ffi::Function::GetGlobal("codegen.llvm.canonicalize_target.default");
+ attrs = (*default_canonicalize)(attrs).cast<ffi::Map<ffi::String,
ffi::Any>>();
+ return IsArch(attrs) ? Canonicalize(attrs) : attrs;
+ };
+ // Preserve the Arm architecture spellings accepted by the existing parser.
+ for (const char* arch :
+ {"aarch64", "aarch64_be", "aarch64_32", "arm", "arm64",
"armeb",
+ "armv4", "armv4t", "armv5", "armv5t", "armv5te",
"armv5tej",
+ "armv6", "armv6k", "armv6t2", "armv6kz", "armv6m",
"armv7",
+ "armv7a", "armv7l", "armv7ve", "armv7r", "armv7m",
"armv7em",
+ "armv7s", "armv7k", "armv8", "armv8a", "armv8l",
"armv8r",
+ "armv8m.base", "armv8m.main", "armv8.1", "armv8.1a",
"armv8.1m.main", "armv8.2",
+ "armv8.2a", "armv8.3", "armv8.3a", "armv8.4", "armv8.4a",
"armv8.5",
+ "armv8.5a", "armv8.6", "armv8.6a", "armv8.7", "armv8.7a",
"armv8.8",
+ "armv8.8a", "armv8.9", "armv8.9a", "armv9", "armv9a",
"armv9.1",
+ "armv9.1a", "armv9.2", "armv9.2a", "armv9.3", "armv9.3a",
"armv9.4",
+ "armv9.4a", "armv9.5", "armv9.5a", "armv9.6", "armv9.6a",
"armv9.7",
+ "armv9.7a"}) {
+ std::string name = std::string("codegen.llvm.canonicalize_target.") + arch;
+ ffi::reflection::GlobalDef().def(name.c_str(), canonicalize);
+ }
+}
+
} // namespace aprofile
} // namespace llvm
} // namespace canonicalizer
diff --git a/src/codegen/llvm/canonicalizer/canonicalize.cc
b/src/codegen/llvm/canonicalizer/canonicalize.cc
index 9b22c0e2a..cae191e5f 100644
--- a/src/codegen/llvm/canonicalizer/canonicalize.cc
+++ b/src/codegen/llvm/canonicalizer/canonicalize.cc
@@ -20,21 +20,14 @@
#include <string>
-#include "../../canonicalizer/arm_mprofile.h"
-#include "arm_aprofile.h"
-
namespace tvm {
namespace target {
namespace canonicalizer {
namespace llvm {
ffi::Optional<ffi::String> DetectSystemTriple() {
-#ifdef TVM_LLVM_VERSION
auto pf = tvm::ffi::Function::GetGlobal("codegen.llvm_get_system_triple");
- TVM_FFI_ICHECK(pf.has_value()) << "The target llvm_get_system_triple was not
found, "
- "please compile with USE_LLVM = ON";
- return (*pf)().cast<ffi::String>();
-#endif
+ if (pf) return (*pf)().cast<ffi::String>();
return {};
}
@@ -51,12 +44,14 @@ ffi::Map<ffi::String, ffi::Any>
Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
target.Set("mtriple", system_triple);
}
- if (mprofile::IsArch(target)) {
- return mprofile::Canonicalize(target);
+ std::string triple =
target.Get("mtriple").value_or(ffi::String("")).as_or_throw<ffi::String>();
+ std::string arch = triple.substr(0, triple.find('-'));
+ auto canonicalize =
ffi::Function::GetGlobal("codegen.llvm.canonicalize_target." + arch);
+ if (!canonicalize) {
+ canonicalize =
ffi::Function::GetGlobal("codegen.llvm.canonicalize_target.default");
}
-
- if (aprofile::IsArch(target)) {
- return aprofile::Canonicalize(target);
+ if (canonicalize) {
+ return (*canonicalize)(target).cast<ffi::Map<ffi::String, ffi::Any>>();
}
return target;
diff --git a/src/codegen/llvm/llvm_instance.cc
b/src/codegen/llvm/llvm_instance.cc
index 8c8a5e440..14935ef76 100644
--- a/src/codegen/llvm/llvm_instance.cc
+++ b/src/codegen/llvm/llvm_instance.cc
@@ -108,11 +108,17 @@ namespace {
bool InitializeLLVM() {
static std::atomic_flag initialized = ATOMIC_FLAG_INIT;
if (!initialized.test_and_set()) {
+#if TVM_TIRX_LLVM_NATIVE_ONLY
+ TVM_FFI_ICHECK(!llvm::InitializeNativeTarget());
+ TVM_FFI_ICHECK(!llvm::InitializeNativeTargetAsmParser());
+ TVM_FFI_ICHECK(!llvm::InitializeNativeTargetAsmPrinter());
+#else
llvm::InitializeAllTargetInfos();
llvm::InitializeAllTargets();
llvm::InitializeAllTargetMCs();
llvm::InitializeAllAsmParsers();
llvm::InitializeAllAsmPrinters();
+#endif
}
return true;
}
@@ -393,7 +399,15 @@ static const llvm::Target* CreateLLVMTargetInstance(const
std::string triple,
const llvm::Target* llvm_instance =
llvm::TargetRegistry::lookupTarget(triple, error);
#endif
if (!allow_missing && !llvm_instance) {
+#if TVM_TIRX_LLVM_NATIVE_ONLY
+ TVM_FFI_THROW(ValueError) << "LLVM target '" << triple
+ << "' is unavailable in this native-only build. "
+ << "You can try to install pip install "
+ "\"apache-tvm-tirx[all-targets]\" to gain
more targets. "
+ << "LLVM: " << error;
+#else
TVM_FFI_ICHECK(llvm_instance) << "LLVM instance error: `" << error << "`";
+#endif
}
return llvm_instance;
