setup the skeleton

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

.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

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)

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
+```

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"]

python/triton_tvm_ffi/__init__.py

@@ -0,0 +1,3 @@
+from . import utils
+
+__all__ = ["utils"]

python/triton_tvm_ffi/compiler.py

@@ -0,0 +1,7 @@
+from triton.backends.nvidia.compiler import CUDABackend
+
+
+class TVMFFIBackend(CUDABackend): ...
+
+
+del CUDABackend

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

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)

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)
+]

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
+    $<$<CONFIG:Debug>:-O0 -g -DDEBUG>
+    $<$<CONFIG:Release>:-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 .)

src/utils.cc

@@ -0,0 +1,8 @@
+#include <tvm/ffi/tvm_ffi.h>
+
+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);
+}