From db4f3a0c68af20f6cf550ca3048a8a34558fb2fe Mon Sep 17 00:00:00 2001 From: Jinjie Liu Date: Wed, 28 Jan 2026 00:54:18 +0800 Subject: [PATCH] setup the skeleton Signed-off-by: Jinjie Liu --- .gitignore | 16 ++++++++++++++ CMakeLists.txt | 25 ++++++++++++++++++++++ README.md | 13 ++++++++++++ pyproject.toml | 19 +++++++++++++++++ python/triton_tvm_ffi/__init__.py | 3 +++ python/triton_tvm_ffi/compiler.py | 7 +++++++ python/triton_tvm_ffi/driver.py | 20 ++++++++++++++++++ python/triton_tvm_ffi/utils/__init__.py | 11 ++++++++++ python/triton_tvm_ffi/utils/_ffi_api.py | 26 +++++++++++++++++++++++ src/CMakeLists.txt | 28 +++++++++++++++++++++++++ src/utils.cc | 8 +++++++ 11 files changed, 176 insertions(+) create mode 100644 .gitignore create mode 100644 CMakeLists.txt create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 python/triton_tvm_ffi/__init__.py create mode 100644 python/triton_tvm_ffi/compiler.py create mode 100644 python/triton_tvm_ffi/driver.py create mode 100644 python/triton_tvm_ffi/utils/__init__.py create mode 100644 python/triton_tvm_ffi/utils/_ffi_api.py create mode 100644 src/CMakeLists.txt create mode 100644 src/utils.cc diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ae0bd4a --- /dev/null +++ b/.gitignore @@ -0,0 +1,16 @@ +# Python-generated files +__pycache__/ +*.py[oc] +build/ +dist/ +wheels/ +*.egg-info + +# Virtual environments +.venv + +.cache +.clangd +.ruff_cache +.python-version +uv.lock diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..508829b --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,25 @@ +cmake_minimum_required(VERSION 3.18) + +if(DEFINED SKBUILD_PROJECT_NAME) + project(${SKBUILD_PROJECT_NAME}) +else() + project(triton-tvm-ffi) +endif() + +if(CMAKE_BUILD_TYPE STREQUAL "Debug") + set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +else(CMAKE_BUILD_TYPE STREQUAL "Release") +endif() + +find_package(Python COMPONENTS Interpreter REQUIRED) + +execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_VARIABLE TVM_FFI_CMAKEDIR + OUTPUT_STRIP_TRAILING_WHITESPACE +) +list(APPEND CMAKE_PREFIX_PATH "${TVM_FFI_CMAKEDIR}") + +find_package(tvm_ffi CONFIG REQUIRED) + +add_subdirectory(${PROJECT_SOURCE_DIR}/src) diff --git a/README.md b/README.md new file mode 100644 index 0000000..251643c --- /dev/null +++ b/README.md @@ -0,0 +1,13 @@ +# Triton-TVM-FFI + +## Instructions + +### Debug Install +```bash +SKBUILD_BUILD_DIR="build" SKBUILD_CMAKE_BUILD_TYPE=Debug uv pip install --no-build-isolation -ve . +``` + +### Format +```bash +find python -name "*.py" | xargs ruff format +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9b51f9f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,19 @@ +[project] +name = "triton-tvm-ffi" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +dependencies = [ + "apache-tvm-ffi", +] + +[build-system] +requires = ["apache-tvm-ffi", "scikit-build-core"] +build-backend = "scikit_build_core.build" + +[project.entry-points."triton.backends"] +nvidia = "triton_tvm_ffi" + +[tool.scikit-build] +wheel.install-dir = "triton_tvm_ffi" +wheel.packages = ["python/triton_tvm_ffi"] diff --git a/python/triton_tvm_ffi/__init__.py b/python/triton_tvm_ffi/__init__.py new file mode 100644 index 0000000..03235a6 --- /dev/null +++ b/python/triton_tvm_ffi/__init__.py @@ -0,0 +1,3 @@ +from . import utils + +__all__ = ["utils"] diff --git a/python/triton_tvm_ffi/compiler.py b/python/triton_tvm_ffi/compiler.py new file mode 100644 index 0000000..68f171a --- /dev/null +++ b/python/triton_tvm_ffi/compiler.py @@ -0,0 +1,7 @@ +from triton.backends.nvidia.compiler import CUDABackend + + +class TVMFFIBackend(CUDABackend): ... + + +del CUDABackend diff --git a/python/triton_tvm_ffi/driver.py b/python/triton_tvm_ffi/driver.py new file mode 100644 index 0000000..26c011d --- /dev/null +++ b/python/triton_tvm_ffi/driver.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from typing import Type +from triton.backends.nvidia.driver import CudaDriver + + +class TVMFFIUtils(object): + def __new__(cls: Type[TVMFFIUtils]) -> TVMFFIUtils: + if not hasattr(cls, "instance"): + cls.instance = super().__new__(cls) + return cls.instance + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + +class TVMFFIDriver(CudaDriver): ... + + +del CudaDriver diff --git a/python/triton_tvm_ffi/utils/__init__.py b/python/triton_tvm_ffi/utils/__init__.py new file mode 100644 index 0000000..04e1469 --- /dev/null +++ b/python/triton_tvm_ffi/utils/__init__.py @@ -0,0 +1,11 @@ +# 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) diff --git a/python/triton_tvm_ffi/utils/_ffi_api.py b/python/triton_tvm_ffi/utils/_ffi_api.py new file mode 100644 index 0000000..d2f4994 --- /dev/null +++ b/python/triton_tvm_ffi/utils/_ffi_api.py @@ -0,0 +1,26 @@ +# 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 tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB +from typing import TYPE_CHECKING +# 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("triton_tvm_ffi", "utils") +# tvm-ffi-stubgen(begin): global/triton_tvm_ffi.utils +# fmt: off +_FFI_INIT_FUNC("triton_tvm_ffi.utils", __name__) +if TYPE_CHECKING: + def hello() -> None: ... +# fmt: on +# tvm-ffi-stubgen(end) + +__all__ = [ + # tvm-ffi-stubgen(begin): __all__ + "LIB", + "hello", + # tvm-ffi-stubgen(end) +] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt new file mode 100644 index 0000000..81ef820 --- /dev/null +++ b/src/CMakeLists.txt @@ -0,0 +1,28 @@ +add_library( + utils + SHARED + ${CMAKE_CURRENT_SOURCE_DIR}/utils.cc +) + +target_include_directories( + utils + PRIVATE ${PROJECT_SOURCE_DIR}/include +) +target_compile_options( + utils + PRIVATE + $<$:-O0 -g -DDEBUG> + $<$:-O3 -DNDEBUG> +) + +tvm_ffi_configure_target( + utils + STUB_DIR "${CMAKE_SOURCE_DIR}/python" + STUB_INIT ON +) + +install( + TARGETS utils + LIBRARY DESTINATION . +) +tvm_ffi_install(utils DESTINATION .) diff --git a/src/utils.cc b/src/utils.cc new file mode 100644 index 0000000..af5f6db --- /dev/null +++ b/src/utils.cc @@ -0,0 +1,8 @@ +#include + +void hello() { std::cout << "Hello, world!\n"; } + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("triton_tvm_ffi.utils.hello", hello); +}