From be63341e9847678433a37e00eeee384642e7a7a7 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 19:26:39 +0800 Subject: [PATCH 1/4] feat!: support polymorphic handles and configs --- src/cloneable.h | 21 ++++ src/config.h | 15 +-- src/handle.h | 9 ++ src/operator.h | 8 +- tests/test_cpp_api.py | 251 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 293 insertions(+), 11 deletions(-) create mode 100644 src/cloneable.h diff --git a/src/cloneable.h b/src/cloneable.h new file mode 100644 index 000000000..71e75ee36 --- /dev/null +++ b/src/cloneable.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_CLONEABLE_H_ +#define INFINI_OPS_CLONEABLE_H_ + +#include +#include + +namespace infini::ops { + +template +class Cloneable : public Base { + public: + std::unique_ptr Clone() const override { + static_assert(std::is_final_v, + "Cloneable requires a final derived class."); + return std::make_unique(static_cast(*this)); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/config.h b/src/config.h index 15bb430ca..e156497bd 100644 --- a/src/config.h +++ b/src/config.h @@ -4,25 +4,26 @@ #include #include +#include "cloneable.h" + namespace infini::ops { class Config { public: + virtual ~Config() = default; + + virtual std::unique_ptr Clone() const { + return std::make_unique(*this); + } + std::size_t implementation_index() const { return implementation_index_; } void set_implementation_index(std::size_t implementation_index) { implementation_index_ = implementation_index; } - void set_extension(std::shared_ptr extension) { - extension_ = std::move(extension); - } - - std::shared_ptr extension() const { return extension_; } - private: std::size_t implementation_index_{0}; - std::shared_ptr extension_{}; }; } // namespace infini::ops diff --git a/src/handle.h b/src/handle.h index 4deeb83c9..27211f0ec 100644 --- a/src/handle.h +++ b/src/handle.h @@ -2,11 +2,20 @@ #define INFINI_OPS_HANDLE_H_ #include +#include + +#include "cloneable.h" namespace infini::ops { class Handle { public: + virtual ~Handle() = default; + + virtual std::unique_ptr Clone() const { + return std::make_unique(*this); + } + void* stream() const { return stream_; } void* workspace() const { return workspace_; } diff --git a/src/operator.h b/src/operator.h index 9d868bb95..1ec5ed5d9 100644 --- a/src/operator.h +++ b/src/operator.h @@ -162,9 +162,9 @@ class OperatorBase { virtual std::size_t workspace_size_in_bytes() const { return 0; } - void set_handle(const Handle& handle) { handle_ = handle; } + void set_handle(const Handle& handle) { handle_ = handle.Clone(); } - void set_config(const Config& config) { config_ = config; } + void set_config(const Config& config) { config_ = config.Clone(); } void set_stream(void* stream) { stream_ = stream; } @@ -175,9 +175,9 @@ class OperatorBase { } protected: - Handle handle_; + std::unique_ptr handle_; - Config config_; + std::unique_ptr config_; void* stream_{nullptr}; diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 1cc60fe34..221f8a851 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -58,6 +58,58 @@ def test_cpp_returning_call_smoke(tmp_path): _run([str(binary)]) +def test_cpp_configless_calls_use_first_active_implementation(tmp_path): + install_prefix = _install_prefix() + include_dir = install_prefix / "include" + library_dir = _library_dir(install_prefix) + source = tmp_path / "configless_active_implementation.cc" + binary = tmp_path / "configless_active_implementation" + source.write_text(_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE) + + _run( + [ + _compiler("CXX", "c++"), + "-std=c++17", + "-Werror", + f"-I{include_dir}", + str(source), + f"-L{library_dir}", + "-linfiniops", + "-linfinirt", + f"-Wl,-rpath,{library_dir}", + "-o", + str(binary), + ] + ) + _run([str(binary)]) + + +def test_cpp_polymorphic_context_smoke(tmp_path): + install_prefix = _install_prefix() + include_dir = install_prefix / "include" + library_dir = _library_dir(install_prefix) + source = tmp_path / "polymorphic_context.cc" + binary = tmp_path / "polymorphic_context" + source.write_text(_POLYMORPHIC_CONTEXT_SOURCE) + + _run( + [ + _compiler("CXX", "c++"), + "-std=c++17", + "-Werror", + f"-I{include_dir}", + str(source), + f"-L{library_dir}", + "-linfiniops", + "-linfinirt", + f"-Wl,-rpath,{library_dir}", + "-o", + str(binary), + ] + ) + _run([str(binary)]) + + @pytest.mark.parametrize( "header", ( @@ -332,3 +384,202 @@ class OwningTensor { } """ ).lstrip() + + +_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + #include + + namespace infini::ops { + + class ConfiglessSelection : public Operator { + public: + ConfiglessSelection(const Tensor input, Tensor out) {} + + ConfiglessSelection(const std::vector inputs, Tensor out) {} + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + virtual void operator()(const std::vector inputs, + Tensor out) const = 0; + + template + static auto MakeReturnValue(const TensorLike& input) { + return TensorLike::Empty(input.shape(), input.dtype(), input.device()); + } + }; + + template <> + class Operator + : public ConfiglessSelection { + public: + using ConfiglessSelection::ConfiglessSelection; + + void operator()(const Tensor input, Tensor out) const override { + const auto* input_data = static_cast(input.data()); + auto* out_data = static_cast(out.data()); + out_data[0] = input_data[0] + 16.0f; + } + + void operator()(const std::vector inputs, + Tensor out) const override { + const auto* first_data = static_cast(inputs[0].data()); + const auto* second_data = static_cast(inputs[1].data()); + auto* out_data = static_cast(out.data()); + out_data[0] = first_data[0] + second_data[0] + 16.0f; + } + }; + + } // namespace infini::ops + + int main() { + float input_data = 1.0f; + float second_data = 2.0f; + float out_data = 0.0f; + const infini::ops::Tensor::Shape shape{1}; + const infini::ops::Device device{infini::ops::Device::Type::kCpu}; + const infini::ops::DataType dtype{infini::ops::DataType::kFloat32}; + infini::ops::Tensor input(&input_data, shape, dtype, device); + infini::ops::Tensor second(&second_data, shape, dtype, device); + infini::ops::Tensor out(&out_data, shape, dtype, device); + + auto tensor_op = infini::ops::ConfiglessSelection::Make(input, out); + (*tensor_op)(input, out); + if (std::fabs(out_data - 17.0f) > 1e-6f) { + return 1; + } + + out_data = 0.0f; + std::vector inputs{input, second}; + auto vector_op = infini::ops::ConfiglessSelection::Make(inputs, out); + (*vector_op)(inputs, out); + if (std::fabs(out_data - 19.0f) > 1e-6f) { + return 2; + } + + float cat_out_data[2] = {}; + const infini::ops::Tensor::Shape cat_shape{2}; + infini::ops::Tensor cat_out(cat_out_data, cat_shape, dtype, device); + auto cat_op = + infini::ops::Cat::Make(inputs, std::int64_t{0}, cat_out); + (*cat_op)(inputs, std::int64_t{0}, cat_out); + if (std::fabs(cat_out_data[0] - 1.0f) > 1e-6f || + std::fabs(cat_out_data[1] - 2.0f) > 1e-6f) { + return 3; + } + + float abs_input_data = -4.0f; + float abs_out_data = 0.0f; + infini::ops::Tensor abs_input(&abs_input_data, shape, dtype, device); + infini::ops::Tensor abs_out(&abs_out_data, shape, dtype, device); + auto abs_op = infini::ops::Abs::Make(abs_input, abs_out); + (*abs_op)(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 4; + } + + auto abs_op_from_rvalue = infini::ops::Abs::Make( + abs_input, + infini::ops::Tensor(&abs_out_data, shape, dtype, device)); + (*abs_op_from_rvalue)(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 5; + } + + abs_out_data = 0.0f; + infini::ops::Abs::Call(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 6; + } + + return 0; + } + """ +).lstrip() + + +_POLYMORPHIC_CONTEXT_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + + namespace infini::ops { + + class DerivedConfig final : public Cloneable { + public: + explicit DerivedConfig(int value) : value_{value} {} + + int value() const { return value_; } + + private: + int value_; + }; + + class DerivedHandle final : public Cloneable { + public: + explicit DerivedHandle(int value) : value_{value} {} + + int value() const { return value_; } + + private: + int value_; + }; + + class PolymorphicOwner final : public OperatorBase { + public: + int config_value() const { + return static_cast(*config_).value(); + } + + std::size_t implementation_index() const { + return config_->implementation_index(); + } + + int handle_value() const { + return static_cast(*handle_).value(); + } + + void* handle_stream() const { return handle_->stream(); } + }; + + } // namespace infini::ops + + int main() { + using namespace infini::ops; + + static_assert(std::has_virtual_destructor_v); + static_assert(std::has_virtual_destructor_v); + + PolymorphicOwner owner; + + { + DerivedConfig config{17}; + config.set_implementation_index(3); + owner.set_config(config); + config.set_implementation_index(9); + } + + if (owner.config_value() != 17) return 1; + if (owner.implementation_index() != 3) return 2; + + int stream; + { + DerivedHandle handle{23}; + handle.set_stream(&stream); + owner.set_handle(handle); + handle.set_stream(nullptr); + } + + if (owner.handle_value() != 23) return 3; + if (owner.handle_stream() != &stream) return 4; + + return 0; + } + """ +).lstrip() From 95bd47e35ae8564757e7650448faf7216616eba2 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 15:16:30 +0800 Subject: [PATCH 2/4] fix: support multilevel cloneable hierarchies --- src/cloneable.h | 8 ++++---- src/operator.h | 8 ++++---- tests/test_cpp_api.py | 34 ++++++++++++++++++++++++++-------- 3 files changed, 34 insertions(+), 16 deletions(-) diff --git a/src/cloneable.h b/src/cloneable.h index 71e75ee36..0517d90c9 100644 --- a/src/cloneable.h +++ b/src/cloneable.h @@ -2,16 +2,16 @@ #define INFINI_OPS_CLONEABLE_H_ #include -#include +#include namespace infini::ops { template class Cloneable : public Base { public: - std::unique_ptr Clone() const override { - static_assert(std::is_final_v, - "Cloneable requires a final derived class."); + using Pointer = decltype(std::declval().Clone()); + + Pointer Clone() const override { return std::make_unique(static_cast(*this)); } }; diff --git a/src/operator.h b/src/operator.h index 1ec5ed5d9..d77fe65a8 100644 --- a/src/operator.h +++ b/src/operator.h @@ -162,9 +162,9 @@ class OperatorBase { virtual std::size_t workspace_size_in_bytes() const { return 0; } - void set_handle(const Handle& handle) { handle_ = handle.Clone(); } + void set_handle(const Handle& handle) { handle_ptr_ = handle.Clone(); } - void set_config(const Config& config) { config_ = config.Clone(); } + void set_config(const Config& config) { config_ptr_ = config.Clone(); } void set_stream(void* stream) { stream_ = stream; } @@ -175,9 +175,9 @@ class OperatorBase { } protected: - std::unique_ptr handle_; + std::unique_ptr handle_ptr_; - std::unique_ptr config_; + std::unique_ptr config_ptr_; void* stream_{nullptr}; diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 221f8a851..28cce589a 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -511,21 +511,35 @@ class Operator namespace infini::ops { - class DerivedConfig final : public Cloneable { + class IntermediateConfig + : public Cloneable { + public: + virtual int value() const { return -1; } + }; + + class DerivedConfig + : public Cloneable { public: explicit DerivedConfig(int value) : value_{value} {} - int value() const { return value_; } + int value() const override { return value_; } private: int value_; }; - class DerivedHandle final : public Cloneable { + class IntermediateHandle + : public Cloneable { + public: + virtual int value() const { return -1; } + }; + + class DerivedHandle + : public Cloneable { public: explicit DerivedHandle(int value) : value_{value} {} - int value() const { return value_; } + int value() const override { return value_; } private: int value_; @@ -534,18 +548,18 @@ class DerivedHandle final : public Cloneable { class PolymorphicOwner final : public OperatorBase { public: int config_value() const { - return static_cast(*config_).value(); + return static_cast(*config_ptr_).value(); } std::size_t implementation_index() const { - return config_->implementation_index(); + return config_ptr_->implementation_index(); } int handle_value() const { - return static_cast(*handle_).value(); + return static_cast(*handle_ptr_).value(); } - void* handle_stream() const { return handle_->stream(); } + void* handle_stream() const { return handle_ptr_->stream(); } }; } // namespace infini::ops @@ -555,6 +569,10 @@ class PolymorphicOwner final : public OperatorBase { static_assert(std::has_virtual_destructor_v); static_assert(std::has_virtual_destructor_v); + static_assert( + std::is_same_v>); + static_assert( + std::is_same_v>); PolymorphicOwner owner; From 94bd8fd1f80c87885737db10f0404fc7a55fd030 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 12 Aug 2026 03:16:15 +0800 Subject: [PATCH 3/4] test: fix configless C++ API smoke --- tests/test_cpp_api.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 28cce589a..5f77f2afb 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -58,13 +58,13 @@ def test_cpp_returning_call_smoke(tmp_path): _run([str(binary)]) -def test_cpp_configless_calls_use_first_active_implementation(tmp_path): +def test_cpp_configless_calls_use_default_implementation(tmp_path): install_prefix = _install_prefix() include_dir = install_prefix / "include" library_dir = _library_dir(install_prefix) - source = tmp_path / "configless_active_implementation.cc" - binary = tmp_path / "configless_active_implementation" - source.write_text(_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE) + source = tmp_path / "configless_default_implementation.cc" + binary = tmp_path / "configless_default_implementation" + source.write_text(_CONFIGLESS_DEFAULT_IMPLEMENTATION_SOURCE) _run( [ @@ -386,7 +386,7 @@ class OwningTensor { ).lstrip() -_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE = textwrap.dedent( +_CONFIGLESS_DEFAULT_IMPLEMENTATION_SOURCE = textwrap.dedent( r""" #include @@ -414,7 +414,7 @@ class ConfiglessSelection : public Operator { }; template <> - class Operator + class Operator : public ConfiglessSelection { public: using ConfiglessSelection::ConfiglessSelection; @@ -422,7 +422,7 @@ class Operator void operator()(const Tensor input, Tensor out) const override { const auto* input_data = static_cast(input.data()); auto* out_data = static_cast(out.data()); - out_data[0] = input_data[0] + 16.0f; + out_data[0] = input_data[0]; } void operator()(const std::vector inputs, @@ -430,7 +430,7 @@ class Operator const auto* first_data = static_cast(inputs[0].data()); const auto* second_data = static_cast(inputs[1].data()); auto* out_data = static_cast(out.data()); - out_data[0] = first_data[0] + second_data[0] + 16.0f; + out_data[0] = first_data[0] + second_data[0]; } }; @@ -449,7 +449,7 @@ class Operator auto tensor_op = infini::ops::ConfiglessSelection::Make(input, out); (*tensor_op)(input, out); - if (std::fabs(out_data - 17.0f) > 1e-6f) { + if (std::fabs(out_data - 1.0f) > 1e-6f) { return 1; } @@ -457,7 +457,7 @@ class Operator std::vector inputs{input, second}; auto vector_op = infini::ops::ConfiglessSelection::Make(inputs, out); (*vector_op)(inputs, out); - if (std::fabs(out_data - 19.0f) > 1e-6f) { + if (std::fabs(out_data - 3.0f) > 1e-6f) { return 2; } From de9694aef9aef5bfb1c2705536b3d00592d6e318 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 12 Aug 2026 03:18:47 +0800 Subject: [PATCH 4/4] refactor(triton): simplify JIT backend architecture --- CMakeLists.txt | 23 +- docs/build.md | 22 ++ pyproject.toml | 1 + scripts/generate_wrappers.py | 185 ++++++----- src/CMakeLists.txt | 19 +- src/driver.h | 16 - src/operator.h | 6 +- src/triton/jit/backend.h | 153 +++++++++ src/triton/jit/backend_nvidia.h | 168 ++++++++++ src/triton/jit/base.h | 126 -------- src/triton/jit/cache.cc | 391 ++++++++++++++++++++++ src/triton/jit/cache.h | 297 +++++------------ src/triton/jit/compile.py | 414 ++++++++++++++++-------- src/triton/jit/compiler.cc | 487 ++++++++++++++++++++-------- src/triton/jit/compiler.h | 40 +++ src/triton/jit/jit.h | 537 +++++++++++++++++-------------- src/triton/jit/jit_config.cc | 13 + src/triton/jit/jit_config.h | 81 +++++ src/triton/jit/pybind11_config.h | 123 +++++++ src/triton/ops/add/add.py | 53 --- src/triton/ops/add/jit.cc | 341 ++++++++++++++++---- src/triton/ops/add/jit.h | 16 - src/triton/ops/add/kernel.py | 66 ++++ tests/test_add.py | 108 ++++++- tests/test_generate_wrappers.py | 181 ++++++++++- tests/test_triton_jit_compile.py | 198 ++++++++++++ 26 files changed, 2955 insertions(+), 1110 deletions(-) delete mode 100644 src/driver.h create mode 100644 src/triton/jit/backend.h create mode 100644 src/triton/jit/backend_nvidia.h delete mode 100644 src/triton/jit/base.h create mode 100644 src/triton/jit/cache.cc create mode 100644 src/triton/jit/compiler.h create mode 100644 src/triton/jit/jit_config.cc create mode 100644 src/triton/jit/jit_config.h create mode 100644 src/triton/jit/pybind11_config.h delete mode 100644 src/triton/ops/add/add.py create mode 100644 src/triton/ops/add/kernel.py create mode 100644 tests/test_triton_jit_compile.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 26ae3ba80..74ca9f424 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -25,7 +25,7 @@ option(WITH_TORCH "Enable PyTorch C++ backend" OFF) option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF) -option(WITH_TRITON "Enable Triton-generated kernels" OFF) +option(WITH_TRITON "Enable the Triton JIT backend" OFF) # Custom `AscendC` kernels under `src/native/ascend/custom/`. `ON` by default # so CI and routine dev builds always exercise `implementation_index=1/2` @@ -336,14 +336,29 @@ if(WITH_NINETOOTHED) set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation") endif() -if(WITH_TRITON AND NOT WITH_NVIDIA) - message(FATAL_ERROR "`WITH_TRITON` temporarily requires `WITH_NVIDIA=ON` because the Triton backend temporarily targets CUDA.") +if(WITH_TRITON) + if(NOT WITH_NVIDIA) + message( + FATAL_ERROR + "`WITH_TRITON` requires `WITH_NVIDIA=ON` because NVIDIA is the only implemented Triton JIT backend." + ) + endif() + if(NOT GENERATE_PYTHON_BINDINGS) + message( + FATAL_ERROR + "`WITH_TRITON` requires `GENERATE_PYTHON_BINDINGS=ON` because the runtime compiler is shipped in the Python package." + ) + endif() endif() if(WITH_NVIDIA) add_compile_definitions(WITH_NVIDIA=1) enable_language(CUDA) - find_package(CUDAToolkit REQUIRED) + if(WITH_TRITON) + find_package(CUDAToolkit 12.0 REQUIRED) + else() + find_package(CUDAToolkit REQUIRED) + endif() endif() # Iluvatar: CUDA-compatible device. CoreX clang++ works for `-x ivcore` diff --git a/docs/build.md b/docs/build.md index ff235b2c2..9a07d9256 100644 --- a/docs/build.md +++ b/docs/build.md @@ -18,6 +18,7 @@ entry is `python -m pip install` with CMake options passed through | `WITH_ASCEND` | Enable the Ascend backend. | `OFF` | | `WITH_TORCH` | Enable PyTorch C++ ATen-backed operators. | `OFF` | | `WITH_NINETOOTHED` | Enable NineToothed-generated kernels. | `OFF` | +| `WITH_TRITON` | Enable the NVIDIA Triton JIT backend. Requires Python bindings and CUDA Toolkit 12.0 or newer. | `OFF` | | `AUTO_DETECT_DEVICES` | Auto-detect available device files. | `OFF` | | `AUTO_DETECT_BACKENDS` | Auto-detect available backend packages. | `OFF` | | `GENERATE_OPERATOR_CALL_INSTANTIATIONS` | Generate explicit C++ operator call instantiations. | `ON` | @@ -50,6 +51,27 @@ python -m pip install .[dev] \ --config-settings=cmake.define.WITH_NVIDIA=ON ``` +Enable the Triton JIT implementation with the matching optional dependency: + +```bash +python -m pip install .[triton] \ + --config-settings=cmake.define.INFINI_RT_ROOT=/path/to/infini-rt-prefix \ + --config-settings=cmake.define.WITH_NVIDIA=ON \ + --config-settings=cmake.define.WITH_TRITON=ON +``` + +The InfiniOps JIT bridge and kernel sources are packaged only with the Python +wheel. It requires CUDA Toolkit 12.0 or newer. Standalone C++ installations do +not provide this runtime. + +Compiled kernels are cached in the platform cache directory. Set +`INFINI_OPS_TRITON_CACHE_DIR` to override that location. + +Python calls with an explicit Triton config construct an operator for that +call instead of entering the generic operator cache. The compiled kernel and +auto-tuning result are still cached using the complete Triton config +identity. + Full builds with both `WITH_NVIDIA=ON` and `WITH_TORCH=ON` include `flash_attn_with_kvcache` and require a compatible FlashAttention 2.7 Python distribution to be installed in the build environment. The distribution must diff --git a/pyproject.toml b/pyproject.toml index 288f9d5be..0a230f46c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ version = "0.1.0" [project.optional-dependencies] dev = ["pytest", "pytest-cov", "pytest-xdist", "ruff==0.15.22", "torch", "pyyaml"] +triton = ["torch", "triton>=3.5,<3.6"] [tool.scikit-build.wheel] install-dir = "infini" diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 8dc140e12..9d5614cae 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -1,5 +1,6 @@ import argparse import concurrent.futures +import dataclasses import functools import json import os @@ -411,15 +412,34 @@ def _strip_default_argument(param): return parts[0].strip() +@dataclasses.dataclass(frozen=True) +class _Implementation: + path: pathlib.Path + backend: str + + class _Operator: - def __init__(self, name, constructors, calls): + def __init__(self, name, constructors, calls, implementations=()): self.name = name self.constructors = constructors self.calls = calls - self.impl_paths = [] + self.implementations = tuple(implementations) + + def has_backend(self, backend): + return any( + implementation.backend == backend for implementation in self.implementations + ) + + +def _parameter_type_signature(node): + return tuple( + " ".join(arg.type.spelling.split()) + for arg in node.get_arguments() + if arg.spelling != "stream" + ) def _find_optional_tensor_params(op_name): @@ -582,56 +602,6 @@ def _is_data_type_spelling(spelling): return spelling.rsplit("::", maxsplit=1)[-1] == "DataType" -def _uses_config_extension(impl_paths): - pattern = re.compile(r"\bJitConfig\b") - for path in impl_paths: - try: - if pattern.search(path.read_text()): - return True - except (OSError, UnicodeDecodeError): - pass - return False - - -def _generate_triton_jit_config_parser(): - return textwrap.dedent("""\ - inline std::shared_ptr ConfigFromPyDict(const py::dict& config_dict) { - if (config_dict.contains("autotune")) { - auto config = std::make_shared(); - py::dict autotune_dict = config_dict["autotune"].cast(); - if (autotune_dict.contains("warmup")) config->warmup = autotune_dict["warmup"].cast(); - if (autotune_dict.contains("rep")) config->rep = autotune_dict["rep"].cast(); - if (autotune_dict.contains("key")) { - for (auto k : autotune_dict["key"].cast()) - config->key.push_back(k.cast()); - } - if (autotune_dict.contains("configs")) { - for (auto candidate : autotune_dict["configs"].cast()) { - JitConfig candidate_config; - py::dict candidate_dict = candidate.cast(); - if (candidate_dict.contains("num_warps")) candidate_config.num_warps = candidate_dict["num_warps"].cast(); - if (candidate_dict.contains("num_stages")) candidate_config.num_stages = candidate_dict["num_stages"].cast(); - for (auto item : candidate_dict) { - std::string key = item.first.cast(); - if (key != "num_warps" && key != "num_stages") - candidate_config.constexprs.emplace_back(key, item.second.cast()); - } - config->candidates.push_back(std::move(candidate_config)); - } - } - return config; - } - auto config = std::make_shared(); - if (config_dict.contains("num_warps")) config->num_warps = config_dict["num_warps"].cast(); - if (config_dict.contains("num_stages")) config->num_stages = config_dict["num_stages"].cast(); - for (auto item : config_dict) { - std::string key = item.first.cast(); - if (key != "num_warps" && key != "num_stages") - config->constexprs.emplace_back(key, item.second.cast()); - } - return config; - }""") - def _generate_pybind11(operator): optional_tensor_params = _find_optional_tensor_params(operator.name) @@ -736,7 +706,7 @@ def _generate_arguments( preconverted_tensor_arg is not None and arg.spelling == preconverted_tensor_arg.spelling ): - args.append(f"std::move({preconverted_tensor_name})") + args.append(preconverted_tensor_name) elif _is_optional_vector_tensor(arg): args.append(f"VectorOptionalTensorFromPybind11Handle({arg.spelling})") elif _is_optional_tensor(arg): @@ -826,7 +796,7 @@ def _generate_py_args(node): return ", ".join(parts) - def _generate_call(op_name, call, method=True, uses_config=False): + def _generate_call(op_name, call, method=True, supports_triton_config=False): call_params = _generate_params(call) call_args = _generate_arguments(call) @@ -848,11 +818,15 @@ def _generate_call(op_name, call, method=True, uses_config=False): extra_params = "" extra_config_init = "" extra_pybind = "" - if uses_config: + if supports_triton_config: extra_params = ", std::optional config_dict" extra_config_init = ( + " std::unique_ptr triton_config_ptr;\n" " if (config_dict.has_value()) {\n" - " config.set_extension(ConfigFromPyDict(*config_dict));\n" + " triton_config_ptr = " + "TritonJitConfigFromPyDict(*config_dict);\n" + " triton_config_ptr->set_implementation_index(\n" + " config.implementation_index());\n" " }\n" ) extra_pybind = ', py::arg("config") = py::none()' @@ -869,10 +843,13 @@ def _generate_call(op_name, call, method=True, uses_config=False): call, converted_first_tensor_name ) - if uses_config: + if supports_triton_config: dispatch = ( - f" auto op = generated_dispatch::Make{symbol_name}(config, {call_args});\n" - f" (*op)(handle, {call_args});" + " if (triton_config_ptr) {\n" + f" auto op = generated_dispatch::Make{symbol_name}(*triton_config_ptr, {call_args});\n" + f" return (*op)(handle, {call_args});\n" + " }\n" + f" return generated_dispatch::Call{symbol_name}(handle, config, {call_args});" ) else: dispatch = f" return generated_dispatch::Call{symbol_name}(handle, config, {call_args});" @@ -943,34 +920,40 @@ def _overload_order_key(node): inits = "\n".join(_generate_init(constructor) for constructor in constructors) calls = "\n".join(_generate_call(operator.name, call) for call in operator_calls) - supports_triton = _uses_config_extension(operator.impl_paths) - callers = "\n".join( - _generate_call(operator.name, call, method=False, uses_config=supports_triton) + supports_triton = operator.has_backend("triton") + constructor_signatures = { + _parameter_type_signature(constructor) for constructor in constructors + } + configurable_calls = [ + supports_triton and _parameter_type_signature(call) in constructor_signatures for call in operator_calls - ) - if supports_triton: - jit_include = ( - '\n#include "triton/jit/jit.h"\n\n' - "namespace infini::ops {\n\n" - + _generate_triton_jit_config_parser() - + "\n\n} // namespace infini::ops\n" + ] + callers = "\n".join( + _generate_call( + operator.name, + call, + method=False, + supports_triton_config=supports_config, ) - else: - jit_include = "" + for call, supports_config in zip(operator_calls, configurable_calls) + ) + triton_config_include = ( + '\n#include "triton/jit/pybind11_config.h"' if any(configurable_calls) else "" + ) return f"""#ifndef INFINI_OPS_BINDINGS_{op_name.upper()}_H_ #define INFINI_OPS_BINDINGS_{op_name.upper()}_H_ #include #include -#include +#include #include "base/{op_name}.h" #include "config.h" #include "generated/bindings/generated_dispatch.h" #include "handle.h" #include "host_range_profiler.h" -#include "pybind11_utils.h"{jit_include} +#include "pybind11_utils.h"{triton_config_include} namespace py = pybind11; @@ -1336,10 +1319,7 @@ def _append_optional_params(prefix, params): emitted_make_params = set() - make_nodes = list(operator.constructors) - if _uses_config_extension(operator.impl_paths): - make_nodes.extend(operator.calls) - for node in make_nodes: + for node in operator.constructors: params = _generate_params(node) args = _generate_arguments(node) make_params = _append_optional_params("const Config& config", params) @@ -1751,6 +1731,24 @@ def _matches_scan_dir(impl_path, scan_dirs): return any(part in scan_dirs for part in impl_path.parts) +def _implementation_backend(path): + for backend in ("triton", "torch", "ninetoothed"): + if backend in path.parts: + return backend + + return "native" + + +def _implementation_from_json(implementation): + if isinstance(implementation, str): + path = pathlib.Path(implementation) + return _Implementation(path, _implementation_backend(path)) + + return _Implementation( + pathlib.Path(implementation["path"]), implementation["backend"] + ) + + _OPERATOR_DECL_RE = re.compile( r"\bclass\s+Operator<\s*((?:[A-Za-z_][A-Za-z0-9_]*::)*[A-Za-z_][A-Za-z0-9_]*)\b" ) @@ -1774,7 +1772,10 @@ def _index_impl_headers(impl_roots, scan_dirs): text = impl_path.read_text() for match in _OPERATOR_DECL_RE.finditer(text): - by_operator.setdefault(match.group(1), []).append(impl_path) + implementation = _Implementation( + impl_path, _implementation_backend(impl_path) + ) + by_operator.setdefault(match.group(1), []).append(implementation) return by_operator @@ -1849,23 +1850,30 @@ def _get_all_ops(devices, with_torch=False, with_ninetoothed=False, with_triton= if op_name in ops: continue - impl_paths = list( + implementations = list( impl_headers_by_operator.get(_op_relative_type(op_name), ()) ) - if not impl_paths: + if not implementations: continue - ops[op_name] = impl_paths + ops[op_name] = implementations return ops def _generate_op_artifacts(item): - op_name, impl_paths = item + op_name, implementations = item + implementations = tuple(implementations) + impl_paths = [implementation.path for implementation in implementations] extractor = _OperatorExtractor() - operator = extractor(op_name) - operator.impl_paths = impl_paths + parsed_operator = extractor(op_name) + operator = _Operator( + parsed_operator.name, + parsed_operator.constructors, + parsed_operator.calls, + implementations, + ) header_name = f"{op_name}.h" legacy_c_source, legacy_c_header = _generate_legacy_c(operator, impl_paths) dispatch_declarations, dispatch_definitions = _generate_generated_dispatch_entries( @@ -2044,7 +2052,14 @@ def _dispatch_gen_batch_size(): ops_json = pathlib.Path("ops.json") if ops_json.exists(): - ops = json.loads(ops_json.read_text()) + raw_ops = json.loads(ops_json.read_text()) + ops = { + op_name: [ + _implementation_from_json(implementation) + for implementation in implementations + ] + for op_name, implementations in raw_ops.items() + } else: ops = _get_all_ops( args.devices, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index cf2884cfa..1f2a8f0dd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -146,7 +146,6 @@ if(WITH_NVIDIA) target_sources(infiniops PRIVATE ${NVIDIA_SOURCES}) target_link_libraries(infiniops PRIVATE nvidia::cutlass::cutlass) - find_package(CUDAToolkit REQUIRED) target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cuda_driver) list(APPEND DEVICE_LIST "nvidia") @@ -190,14 +189,13 @@ if(WITH_NINETOOTHED) endif() if(WITH_TRITON) - find_package(Python COMPONENTS Interpreter Development REQUIRED) - find_package(pybind11 CONFIG REQUIRED) - - target_compile_definitions(infiniops PUBLIC WITH_TRITON=1 - TRITON_JIT_CACHE_DIR="/tmp/triton_jit_cache") - target_include_directories(infiniops PRIVATE ${pybind11_INCLUDE_DIRS}) - target_link_libraries(infiniops PRIVATE pybind11::embed Python::Python) - target_sources(infiniops PRIVATE triton/jit/compiler.cc) + find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) + + target_link_libraries(infiniops PRIVATE Python::Python) + target_sources(infiniops PRIVATE + triton/jit/cache.cc + triton/jit/compiler.cc + triton/jit/jit_config.cc) file(GLOB_RECURSE TRITON_SOURCES CONFIGURE_DEPENDS "triton/ops/*/*.cc") target_sources(infiniops PRIVATE ${TRITON_SOURCES}) endif() @@ -1272,8 +1270,7 @@ if(GENERATE_PYTHON_BINDINGS) install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/triton/ops/" DESTINATION triton/ops FILES_MATCHING - PATTERN "*.py" - PATTERN "build.py" EXCLUDE) + PATTERN "*.py") endif() endif() diff --git a/src/driver.h b/src/driver.h deleted file mode 100644 index 9ea6a79f5..000000000 --- a/src/driver.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef INFINI_OPS_DRIVER_H_ -#define INFINI_OPS_DRIVER_H_ - -#include - -namespace infini::ops { - -template -using Driver = infini::rt::driver::Driver; - -template -using DeviceDriver = infini::rt::driver::DeviceDriver; - -} // namespace infini::ops - -#endif diff --git a/src/operator.h b/src/operator.h index d77fe65a8..d40f65d7a 100644 --- a/src/operator.h +++ b/src/operator.h @@ -206,7 +206,7 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Tensor tensor, Args&&... args) { - return Make({}, tensor, std::forward(args)...); + return Make(Config{}, tensor, std::forward(args)...); } template @@ -222,7 +222,7 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const std::vector tensors, Args&&... args) { - return Make({}, tensors, std::forward(args)...); + return Make(Config{}, tensors, std::forward(args)...); } template @@ -279,7 +279,7 @@ class Operator : public OperatorBase { template static void Call(const Tensor tensor, const Args&... args) { - return Call({}, {}, tensor, args...); + return Call(Handle{}, Config{}, tensor, args...); } template < diff --git a/src/triton/jit/backend.h b/src/triton/jit/backend.h new file mode 100644 index 000000000..806e5321c --- /dev/null +++ b/src/triton/jit/backend.h @@ -0,0 +1,153 @@ +#ifndef INFINI_OPS_TRITON_JIT_BACKEND_H_ +#define INFINI_OPS_TRITON_JIT_BACKEND_H_ + +#include +#include +#include + +#include "device.h" + +namespace infini::ops { + +struct Grid { + unsigned x{1}; + + unsigned y{1}; + + unsigned z{1}; +}; + +struct JitTarget { + std::string backend; + + std::string architecture; + + std::string performance_identity; + + int device_id{0}; + + int warp_size{0}; + + std::uint64_t context_id{0}; +}; + +enum class JitArgumentType { + kPointer, + kInt8, + kUInt8, + kInt16, + kUInt16, + kInt32, + kUInt32, + kInt64, + kUInt64, + kFloat32, + kFloat64, +}; + +struct JitArgument { + JitArgumentType type; + + std::uint64_t bits; +}; + +struct KernelMetadata { + std::string name; + + std::string binary_extension; + + unsigned shared_memory_size{0}; + + int global_scratch_size{0}; + + int profile_scratch_size{0}; +}; + +template +class JitBackend { + public: + static constexpr bool kSupported = false; +}; + +template +class ScopedJitDevice { + public: + using Backend = JitBackend; + + explicit ScopedJitDevice(int device_id) { + if (!Backend::GetCurrentDevice(&previous_device_id_)) { + return; + } + + if (previous_device_id_ == device_id) { + valid_ = true; + return; + } + + if (!Backend::SetCurrentDevice(device_id)) { + return; + } + + restore_ = true; + valid_ = true; + } + + ScopedJitDevice(const ScopedJitDevice&) = delete; + + ScopedJitDevice& operator=(const ScopedJitDevice&) = delete; + + ~ScopedJitDevice() { + if (!restore_) { + return; + } + + const bool restored = Backend::SetCurrentDevice(previous_device_id_); + assert(restored && "`ScopedJitDevice` failed to restore the device"); + (void)restored; + } + + bool valid() const { return valid_; } + + private: + int previous_device_id_ = 0; + + bool restore_ = false; + + bool valid_ = false; +}; + +template +class LoadedKernel { + public: + using Backend = JitBackend; + + using Function = typename Backend::Function; + + using Module = typename Backend::Module; + + LoadedKernel(Module module, Function function, unsigned shared_memory_size) + : module_(module), + function_(function), + shared_memory_size_(shared_memory_size) {} + + LoadedKernel(const LoadedKernel&) = delete; + + LoadedKernel& operator=(const LoadedKernel&) = delete; + + ~LoadedKernel() { Backend::UnloadModule(module_); } + + Function function() const { return function_; } + + unsigned shared_memory_size() const { return shared_memory_size_; } + + private: + Module module_; + + Function function_; + + unsigned shared_memory_size_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/triton/jit/backend_nvidia.h b/src/triton/jit/backend_nvidia.h new file mode 100644 index 000000000..1996b712e --- /dev/null +++ b/src/triton/jit/backend_nvidia.h @@ -0,0 +1,168 @@ +#ifndef INFINI_OPS_TRITON_JIT_BACKEND_NVIDIA_H_ +#define INFINI_OPS_TRITON_JIT_BACKEND_NVIDIA_H_ + +#include +#include + +#include +#include +#include +#include + +#include "runtime.h" +#include "triton/jit/backend.h" + +namespace infini::ops { + +template <> +class JitBackend { + private: + using Driver = infini::rt::driver::Driver; + + using NvidiaRuntime = Runtime; + + static bool GetDeviceIdentity(CUdevice device, std::string* identity) { + CUuuid uuid{}; + if (cuDeviceGetUuid(&uuid, device) != CUDA_SUCCESS) return false; + + constexpr char kHexDigits[] = "0123456789abcdef"; + identity->resize(sizeof(uuid.bytes) * 2); + for (std::size_t index = 0; index < sizeof(uuid.bytes); ++index) { + const auto byte = static_cast(uuid.bytes[index]); + (*identity)[index * 2] = kHexDigits[byte >> 4]; + (*identity)[index * 2 + 1] = kHexDigits[byte & 0xf]; + } + return true; + } + + public: + static constexpr bool kSupported = true; + + using Function = typename Driver::Function; + + using Module = typename Driver::Module; + + using Stream = typename Driver::Stream; + + static bool GetCurrentDevice(int* device_id) { + return NvidiaRuntime::GetDevice(device_id) == NvidiaRuntime::kSuccess; + } + + static bool SetCurrentDevice(int device_id) { + return NvidiaRuntime::SetDevice(device_id) == NvidiaRuntime::kSuccess; + } + + static bool GetCurrentTarget(JitTarget* target) { + int device_id = 0; + if (!GetCurrentDevice(&device_id)) return false; + + CUcontext context = nullptr; + if (cuCtxGetCurrent(&context) != CUDA_SUCCESS || context == nullptr) { + return false; + } + + CUdevice context_device; + if (cuCtxGetDevice(&context_device) != CUDA_SUCCESS || + static_cast(context_device) != device_id) { + return false; + } + + unsigned long long context_id = 0; + if (cuCtxGetId(context, &context_id) != CUDA_SUCCESS) return false; + + int major = 0; + if (NvidiaRuntime::DeviceGetAttribute( + &major, NvidiaRuntime::kDevAttrComputeCapabilityMajor, device_id) != + NvidiaRuntime::kSuccess) { + return false; + } + + int minor = 0; + if (NvidiaRuntime::DeviceGetAttribute( + &minor, NvidiaRuntime::kDevAttrComputeCapabilityMinor, device_id) != + NvidiaRuntime::kSuccess) { + return false; + } + + int warp_size = 0; + if (NvidiaRuntime::DeviceGetAttribute( + &warp_size, NvidiaRuntime::kDevAttrWarpSize, device_id) != + NvidiaRuntime::kSuccess) { + return false; + } + + std::string performance_identity; + if (!GetDeviceIdentity(context_device, &performance_identity)) return false; + + *target = {"nvidia", + std::to_string(major * 10 + minor), + std::move(performance_identity), + device_id, + warp_size, + static_cast(context_id)}; + return true; + } + + static bool LoadKernel(const JitTarget& target, const std::string& binary, + const KernelMetadata& metadata, Module* module, + Function* function) { + if (Driver::ModuleLoadData(module, binary.data()) != Driver::kSuccess) { + return false; + } + + if (Driver::ModuleGetFunction(function, *module, metadata.name.c_str()) != + Driver::kSuccess) { + Driver::ModuleUnload(*module); + return false; + } + + constexpr unsigned kDefaultSharedMemoryLimit = 48 * 1024; + if (metadata.shared_memory_size <= kDefaultSharedMemoryLimit) return true; + + int maximum_shared_memory = 0; + if (NvidiaRuntime::DeviceGetAttribute( + &maximum_shared_memory, + NvidiaRuntime::kDevAttrMaxSharedMemoryPerBlockOptin, + target.device_id) != NvidiaRuntime::kSuccess) { + Driver::ModuleUnload(*module); + return false; + } + + int static_shared_memory = 0; + if (Driver::FuncGetAttribute(&static_shared_memory, + Driver::kFuncAttributeSharedSizeBytes, + *function) != Driver::kSuccess || + maximum_shared_memory < static_shared_memory || + metadata.shared_memory_size > + static_cast(maximum_shared_memory - + static_shared_memory)) { + Driver::ModuleUnload(*module); + return false; + } + + if (Driver::FuncSetCacheConfig(*function, Driver::kFuncCachePreferShared) != + Driver::kSuccess || + Driver::FuncSetAttribute( + *function, Driver::kFuncAttributeMaxDynamicSharedSizeBytes, + maximum_shared_memory - static_shared_memory) != Driver::kSuccess) { + Driver::ModuleUnload(*module); + return false; + } + + return true; + } + + static void UnloadModule(Module module) { Driver::ModuleUnload(module); } + + static int Launch(Function function, Grid grid, unsigned num_warps, + int warp_size, unsigned shared_memory_size, void* stream, + void** arguments) { + return Driver::LaunchKernel( + function, grid.x, grid.y, grid.z, num_warps * warp_size, 1, 1, + shared_memory_size, static_cast(stream), arguments, nullptr); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/triton/jit/base.h b/src/triton/jit/base.h deleted file mode 100644 index 02bdd785d..000000000 --- a/src/triton/jit/base.h +++ /dev/null @@ -1,126 +0,0 @@ -#ifndef INFINI_OPS_TRITON_JIT_BASE_H_ -#define INFINI_OPS_TRITON_JIT_BASE_H_ - -#include -#include -#include - -#include "config.h" -#include "device.h" -#include "driver.h" - -namespace infini::ops { - -// ---- types ---- - -struct JitConfig : Config { - JitConfig() = default; - - JitConfig(unsigned num_warps, unsigned num_stages, - std::vector> constexprs) - : num_warps(num_warps), - num_stages(num_stages), - constexprs(std::move(constexprs)) {} - - bool autotune = false; - - unsigned num_warps = 4; - - unsigned num_stages = 3; - - std::vector> constexprs; - - int At(const std::string& key) const { - for (const auto& [k, v] : constexprs) - if (k == key) return v; - assert(false && "`constexpr` not found"); - return 0; - } - - void ApplyDefaults(const JitConfig& defaults) { - for (const auto& [dk, dv] : defaults.constexprs) { - bool found = false; - for (const auto& [k, v] : constexprs) - if (k == dk) { - found = true; - break; - } - if (!found) constexprs.push_back({dk, dv}); - } - } -}; - -struct AutotuneConfig : public JitConfig { - AutotuneConfig() { autotune = true; } - - std::vector key; - - std::vector candidates; - - int warmup = 25; - - int rep = 100; -}; - -struct Grid { - unsigned x = 1; - - unsigned y = 1; - - unsigned z = 1; -}; - -struct TargetInfo { - std::string type; - - int id = 0; - - int arch = 0; - - int warp_size = 0; -}; - -struct KernelMeta { - std::string name; - - std::string binary_ext; - - unsigned shared = 0; - - unsigned num_warps = 0; - - int global_scratch_size = 0; - - int profile_scratch_size = 0; -}; - -// ---- declarations ---- - -bool CompilerInit(); - -int CompileKernel(const TargetInfo& target, const char* op_name, - const char* out_prefix, int num_warps, int num_stages, - const char* signature); - -template -int LaunchKernel(const char* op_name, const char* signature_str, void* stream, - Grid grid, const JitConfig& config, void** args); - -template -typename Driver::Function GetKernel(const char* op_name, - const char* signature_str, - void* stream, const JitConfig& config, - unsigned* out_shared); - -template -TargetInfo CurrentTarget(); - -JitConfig AutotuneBench(const char* op_name, - const std::vector& configs, - const std::string& sig, const std::vector& ptrs, - const std::vector& grids, int warmup, int rep, - const char* key, const TargetInfo& target); - -} // namespace infini::ops - -#endif diff --git a/src/triton/jit/cache.cc b/src/triton/jit/cache.cc new file mode 100644 index 000000000..455599540 --- /dev/null +++ b/src/triton/jit/cache.cc @@ -0,0 +1,391 @@ +#include "triton/jit/cache.h" + +#define PY_SSIZE_T_CLEAN +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace infini::ops { +namespace { + +constexpr std::uint64_t kFnvOffsetBasis = 14695981039346656037ull; +constexpr std::uint64_t kFnvPrime = 1099511628211ull; + +void AppendField(std::string* identity, std::string_view value) { + *identity += std::to_string(value.size()); + identity->push_back(':'); + identity->append(value.data(), value.size()); + identity->push_back('|'); +} + +template +void AppendNumber(std::string* identity, T value) { + AppendField(identity, std::to_string(value)); +} + +void AppendConfig(std::string* identity, const TritonJitConfig& config) { + AppendNumber(identity, config.num_warps()); + AppendNumber(identity, config.num_stages()); + AppendNumber(identity, config.constexprs().size()); + for (const auto& [name, value] : config.constexprs()) { + AppendField(identity, name); + AppendNumber(identity, value); + } +} + +std::string StableHash(const std::string& value) { + std::uint64_t hash = kFnvOffsetBasis; + for (const unsigned char byte : value) { + hash ^= byte; + hash *= kFnvPrime; + } + + std::ostringstream output; + output << std::hex << std::setfill('0') << std::setw(16) << hash; + return output.str(); +} + +std::filesystem::path CacheDirectory() { + if (const char* path = std::getenv("INFINI_OPS_TRITON_CACHE_DIR"); + path != nullptr && path[0] != '\0') { + return std::filesystem::path{path}; + } + +#if defined(_WIN32) + if (const char* path = std::getenv("LOCALAPPDATA"); + path != nullptr && path[0] != '\0') { + return std::filesystem::path{path} / "infiniops" / "triton"; + } +#else + if (const char* path = std::getenv("XDG_CACHE_HOME"); + path != nullptr && path[0] != '\0') { + return std::filesystem::path{path} / "infiniops" / "triton"; + } + if (const char* path = std::getenv("HOME"); + path != nullptr && path[0] != '\0') { + return std::filesystem::path{path} / ".cache" / "infiniops" / "triton"; + } +#endif + + std::error_code error; + const auto temporary_directory = std::filesystem::temp_directory_path(error); + return error ? std::filesystem::path{} + : temporary_directory / "infiniops" / "triton"; +} + +class GilState { + public: + GilState() : state_(PyGILState_Ensure()) {} + + GilState(const GilState&) = delete; + + GilState& operator=(const GilState&) = delete; + + ~GilState() { PyGILState_Release(state_); } + + private: + PyGILState_STATE state_; +}; + +struct PythonObjectDeleter { + void operator()(PyObject* object) const { Py_DECREF(object); } +}; + +using PythonObject = std::unique_ptr; + +bool ReadPythonString(PyObject* dictionary, const char* key, + std::string* value) { + PyObject* object = PyDict_GetItemString(dictionary, key); + if (object == nullptr || !PyUnicode_Check(object)) return false; + + const char* text = PyUnicode_AsUTF8(object); + if (text == nullptr) return false; + + *value = text; + return true; +} + +bool ReadPythonInteger(PyObject* dictionary, const char* key, int* value) { + PyObject* object = PyDict_GetItemString(dictionary, key); + if (object == nullptr || !PyLong_Check(object)) return false; + + const long parsed = PyLong_AsLong(object); + if ((parsed == -1 && PyErr_Occurred()) || + parsed < std::numeric_limits::min() || + parsed > std::numeric_limits::max()) { + return false; + } + + *value = static_cast(parsed); + return true; +} + +bool ParseKernelMetadata(const std::string& json, + const std::string& expected_identity, + KernelMetadata* metadata) { + if (!Py_IsInitialized()) return false; + + const GilState gil; + PythonObject json_module{PyImport_ImportModule("json")}; + PythonObject document{ + json_module + ? PyObject_CallMethod(json_module.get(), "loads", "s#", json.data(), + static_cast(json.size())) + : nullptr}; + if (!document || !PyDict_Check(document.get())) { + PyErr_Clear(); + return false; + } + + std::string identity; + int shared_memory_size = 0; + const bool valid = + ReadPythonString(document.get(), "cache_identity", &identity) && + identity == expected_identity && + ReadPythonString(document.get(), "name", &metadata->name) && + ReadPythonString(document.get(), "binary_ext", + &metadata->binary_extension) && + ReadPythonInteger(document.get(), "shared", &shared_memory_size) && + ReadPythonInteger(document.get(), "global_scratch_size", + &metadata->global_scratch_size) && + ReadPythonInteger(document.get(), "profile_scratch_size", + &metadata->profile_scratch_size) && + shared_memory_size >= 0; + if (!valid) { + PyErr_Clear(); + return false; + } + + metadata->shared_memory_size = static_cast(shared_memory_size); + return true; +} + +bool ReadFile(const std::filesystem::path& path, std::string* content) { + if (path.empty()) return false; + + std::ifstream input{path, std::ios::binary}; + if (!input) return false; + + std::ostringstream output; + output << input.rdbuf(); + if (input.bad()) return false; + + *content = output.str(); + return true; +} + +std::string SerializeAutoTuningConfig(const AutoTuningCacheKey& key, + const TritonJitConfig& config) { + std::ostringstream output; + output << key.identity() << '\n'; + output << config.num_warps() << ' ' << config.num_stages() << '\n'; + for (const auto& [name, value] : config.constexprs()) { + output << name << ' ' << value << '\n'; + } + return output.str(); +} + +bool DeserializeAutoTuningConfig(const AutoTuningCacheKey& key, + const std::string& content, + TritonJitConfig* config) { + std::istringstream input{content}; + std::string identity; + if (!std::getline(input, identity) || identity != key.identity()) { + return false; + } + + unsigned num_warps = 0; + unsigned num_stages = 0; + if (!(input >> num_warps >> num_stages)) return false; + + TritonJitConfig::Constexprs constexprs; + std::string name; + int value = 0; + while (input >> name >> value) { + if (!constexprs.emplace(name, value).second) return false; + } + if (!input.eof()) return false; + + *config = TritonJitConfig{num_warps, num_stages, std::move(constexprs)}; + return true; +} + +bool WriteFileAtomically(const std::filesystem::path& path, + const std::string& content) { + if (path.empty()) return false; + + std::error_code error; + std::filesystem::create_directories(path.parent_path(), error); + if (error) return false; + + const auto nonce = + std::chrono::steady_clock::now().time_since_epoch().count(); + const auto thread_hash = + std::hash{}(std::this_thread::get_id()); + const std::filesystem::path temporary_path{path.string() + ".tmp." + + std::to_string(nonce) + "." + + std::to_string(thread_hash)}; + + bool write_succeeded = false; + { + std::ofstream output{temporary_path, std::ios::binary | std::ios::trunc}; + if (output) { + output.write(content.data(), + static_cast(content.size())); + write_succeeded = static_cast(output); + } + } + if (!write_succeeded) { + std::filesystem::remove(temporary_path, error); + return false; + } + + std::filesystem::rename(temporary_path, path, error); + if (error) { + std::filesystem::remove(temporary_path, error); + return false; + } + return true; +} + +} // namespace + +KernelCacheKey KernelCacheKey::Build(const JitTarget& target, + const std::string& compilation_fingerprint, + const std::string& operator_name, + const std::string& signature, + const TritonJitConfig& config) { + std::string identity; + AppendField(&identity, "kernel-v1"); + AppendField(&identity, target.backend); + AppendField(&identity, target.architecture); + AppendNumber(&identity, target.warp_size); + AppendField(&identity, compilation_fingerprint); + AppendField(&identity, operator_name); + AppendField(&identity, signature); + AppendConfig(&identity, config); + + std::string memory_identity = identity; + AppendNumber(&memory_identity, target.context_id); + return {std::move(identity), std::move(memory_identity)}; +} + +std::string KernelCacheKey::ArtifactPrefix() const { + const auto directory = CacheDirectory(); + return directory.empty() ? std::string{} + : (directory / StableHash(identity_)).string(); +} + +AutoTuningCacheKey AutoTuningCacheKey::Build( + const JitTarget& target, const std::string& compilation_fingerprint, + const std::string& operator_name, const std::string& signature, + const std::vector& key_names, + const std::vector& key_values, + const std::vector& candidates, + const std::vector& grids, int warmup_milliseconds, + int repetition_milliseconds) { + std::string identity; + AppendField(&identity, "auto-tuning-v3"); + AppendField(&identity, target.backend); + AppendField(&identity, target.architecture); + AppendNumber(&identity, target.warp_size); + AppendField(&identity, target.performance_identity); + AppendField(&identity, compilation_fingerprint); + AppendField(&identity, operator_name); + AppendField(&identity, signature); + AppendNumber(&identity, warmup_milliseconds); + AppendNumber(&identity, repetition_milliseconds); + + AppendNumber(&identity, key_names.size()); + for (const auto& name : key_names) AppendField(&identity, name); + + AppendNumber(&identity, key_values.size()); + for (const auto value : key_values) AppendNumber(&identity, value); + + AppendNumber(&identity, candidates.size()); + for (const auto& candidate : candidates) AppendConfig(&identity, candidate); + + AppendNumber(&identity, grids.size()); + for (const auto grid : grids) { + AppendNumber(&identity, grid.x); + AppendNumber(&identity, grid.y); + AppendNumber(&identity, grid.z); + } + return AutoTuningCacheKey{std::move(identity)}; +} + +std::string AutoTuningCacheKey::FilePath() const { + const auto directory = CacheDirectory(); + return directory.empty() + ? std::string{} + : (directory / (StableHash(identity_) + ".autotune")).string(); +} + +bool ReadKernelArtifact(const std::string& output_prefix, + const std::string& expected_identity, + KernelArtifact* artifact) { + if (output_prefix.empty()) return false; + + std::string metadata_json; + if (!ReadFile(output_prefix + ".json", &metadata_json) || + !ParseKernelMetadata(metadata_json, expected_identity, + &artifact->metadata)) { + return false; + } + + return ReadFile(output_prefix + "." + artifact->metadata.binary_extension, + &artifact->binary) && + !artifact->binary.empty(); +} + +AutoTuningCache& AutoTuningCache::Instance() { + static AutoTuningCache cache; + return cache; +} + +bool AutoTuningCache::Find(const AutoTuningCacheKey& key, + TritonJitConfig* config) { + { + const std::lock_guard lock(mutex_); + const auto it = entries_.find(key.identity()); + if (it != entries_.end()) { + *config = it->second; + return true; + } + } + + std::string content; + TritonJitConfig parsed; + if (!ReadFile(key.FilePath(), &content) || + !DeserializeAutoTuningConfig(key, content, &parsed)) { + return false; + } + + const std::lock_guard lock(mutex_); + const auto result = entries_.try_emplace(key.identity(), parsed); + *config = result.first->second; + return true; +} + +void AutoTuningCache::Insert(const AutoTuningCacheKey& key, + const TritonJitConfig& config) { + { + const std::lock_guard lock(mutex_); + entries_.insert_or_assign(key.identity(), config); + } + + WriteFileAtomically(key.FilePath(), SerializeAutoTuningConfig(key, config)); +} + +} // namespace infini::ops diff --git a/src/triton/jit/cache.h b/src/triton/jit/cache.h index 3a6a68c2c..ec4be05a5 100644 --- a/src/triton/jit/cache.h +++ b/src/triton/jit/cache.h @@ -1,257 +1,120 @@ #ifndef INFINI_OPS_TRITON_JIT_CACHE_H_ #define INFINI_OPS_TRITON_JIT_CACHE_H_ -#include -#include +#include +#include #include -#include #include #include +#include +#include -#include "base.h" +#include "triton/jit/backend.h" +#include "triton/jit/jit_config.h" namespace infini::ops { -// ---- file helpers ---- +struct KernelArtifact { + KernelMetadata metadata; -inline bool FileExists(const char* path) { - FILE* f = fopen(path, "rb"); - if (f != nullptr) { - fclose(f); - return true; - } - return false; -} - -inline std::string ReadFile(const char* path) { - FILE* f = fopen(path, "rb"); - if (f == nullptr) return {}; - fseek(f, 0, SEEK_END); - long sz = ftell(f); - if (sz < 0) { - fclose(f); - return {}; - } - fseek(f, 0, SEEK_SET); - std::string buf(static_cast(sz), '\0'); - size_t nread = fread(buf.data(), 1, static_cast(sz), f); - fclose(f); - buf.resize(nread); - return buf; -} - -// ---- JSON field extraction ---- - -inline int JsonGetInt(const std::string& json, const char* key, - int fallback = 0) { - std::string pat = std::string("\"") + key + "\":"; - auto pos = json.find(pat); - if (pos == std::string::npos) return fallback; - pos += pat.size(); - while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++; - return std::atoi(json.c_str() + pos); -} - -inline std::string JsonGetString(const std::string& json, const char* key, - const char* fallback) { - std::string pat = std::string("\"") + key + "\":"; - auto pos = json.find(pat); - if (pos == std::string::npos) return fallback; - pos += pat.size(); - while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++; - if (pos >= json.size() || json[pos] != '"') return fallback; - pos++; - auto end = json.find('"', pos); - if (end == std::string::npos) return fallback; - return json.substr(pos, end - pos); -} - -// ---- key generation ---- - -inline std::string GenerateDesc(const char* op, const char* sig, - unsigned num_warps, unsigned num_stages, - int arch) { - return std::string(op) + "|" + sig + "|" + std::to_string(num_warps) + "|" + - std::to_string(num_stages) + "|sm" + std::to_string(arch); -} - -inline std::string CacheMemKey(const char* op_name, const char* signature_str, - unsigned num_warps, unsigned num_stages, - int arch, int dev_id) { - return GenerateDesc(op_name, signature_str, num_warps, num_stages, arch) + - "|dev" + std::to_string(dev_id); -} - -inline std::string CacheFileKey(const char* op_name, const char* signature_str, - unsigned num_warps, unsigned num_stages, - int arch) { - return std::to_string(std::hash{}( - GenerateDesc(op_name, signature_str, num_warps, num_stages, arch))); -} + std::string binary; +}; -// ---- artifact reader ---- +class KernelCacheKey { + public: + static KernelCacheKey Build(const JitTarget& target, + const std::string& compilation_fingerprint, + const std::string& operator_name, + const std::string& signature, + const TritonJitConfig& config); -inline bool ReadArtifacts(const std::string& out_prefix, KernelMeta* meta, - std::string* binary_data) { - std::string meta_path = out_prefix + ".json"; - std::string meta_json = ReadFile(meta_path.c_str()); - if (meta_json.empty()) return false; + const std::string& identity() const { return identity_; } - meta->name = JsonGetString(meta_json, "name", ""); - meta->binary_ext = JsonGetString(meta_json, "binary_ext", ""); - if (meta->binary_ext.empty()) meta->binary_ext = "cubin"; - meta->shared = JsonGetInt(meta_json, "shared"); - meta->num_warps = JsonGetInt(meta_json, "num_warps"); - meta->global_scratch_size = JsonGetInt(meta_json, "global_scratch_size"); - meta->profile_scratch_size = JsonGetInt(meta_json, "profile_scratch_size"); + const std::string& memory_identity() const { return memory_identity_; } - std::string binary_path = out_prefix + "." + meta->binary_ext; - *binary_data = ReadFile(binary_path.c_str()); - return !binary_data->empty(); -} + std::string ArtifactPrefix() const; -// ---- kernel cache ---- + private: + KernelCacheKey(std::string identity, std::string memory_identity) + : identity_(std::move(identity)), + memory_identity_(std::move(memory_identity)) {} -template -struct KernelCacheEntry { - typename Driver::Function func; + std::string identity_; - unsigned shared; + std::string memory_identity_; }; -template -struct KernelCache { - std::mutex mutex; - - std::unordered_map> map; -}; +class AutoTuningCacheKey { + public: + static AutoTuningCacheKey Build( + const JitTarget& target, const std::string& compilation_fingerprint, + const std::string& operator_name, const std::string& signature, + const std::vector& key_names, + const std::vector& key_values, + const std::vector& candidates, + const std::vector& grids, int warmup_milliseconds, + int repetition_milliseconds); -template -KernelCache& GetKernelCache() { - static KernelCache c; - return c; -} + const std::string& identity() const { return identity_; } -template -bool KernelCacheLookup(const std::string& key, KernelCacheEntry* out) { - auto& c = GetKernelCache(); - std::lock_guard lk(c.mutex); - auto it = c.map.find(key); - if (it == c.map.end()) return false; - *out = it->second; - return true; -} + std::string FilePath() const; -template -void KernelCacheInsert(const std::string& key, KernelCacheEntry entry) { - auto& c = GetKernelCache(); - std::lock_guard lk(c.mutex); - c.map[key] = entry; -} + private: + explicit AutoTuningCacheKey(std::string identity) + : identity_(std::move(identity)) {} -template -struct CacheQueryResult { - bool mem_hit; - - typename Driver::Function func; + std::string identity_; +}; - unsigned shared; +bool ReadKernelArtifact(const std::string& output_prefix, + const std::string& expected_identity, + KernelArtifact* artifact); - std::string out_prefix; +template +class KernelCache { + public: + using Entry = std::unique_ptr>; - std::string mem_key; -}; + static KernelCache& Instance() { + // Backend contexts may already be gone during static destruction. + static auto* cache = new KernelCache; + return *cache; + } -template -CacheQueryResult CacheQuery(const char* op, const char* sig, - unsigned num_warps, unsigned num_stages, - int arch, int dev_id) { - auto mem_key = CacheMemKey(op, sig, num_warps, num_stages, arch, dev_id); - KernelCacheEntry entry; - if (KernelCacheLookup(mem_key, &entry)) - return {true, entry.func, entry.shared, "", mem_key}; - auto desc = GenerateDesc(op, sig, num_warps, num_stages, arch); - return {false, nullptr, 0, - std::string(TRITON_JIT_CACHE_DIR) + "/" + - std::to_string(std::hash{}(desc)), - mem_key}; -} + const LoadedKernel* Find(const KernelCacheKey& key) const { + const std::lock_guard lock(mutex_); + const auto it = entries_.find(key.memory_identity()); + return it == entries_.end() ? nullptr : it->second.get(); + } -// ---- autotune cache ---- + const LoadedKernel* InsertOrGet(const KernelCacheKey& key, + Entry candidate) { + const std::lock_guard lock(mutex_); + const auto result = + entries_.try_emplace(key.memory_identity(), std::move(candidate)); + return result.first->second.get(); + } -struct AutotuneCache { - std::mutex mutex; + private: + mutable std::mutex mutex_; - std::unordered_map map; + // There is no erase path, so returned pointee addresses remain stable. + std::unordered_map entries_; }; -inline AutotuneCache& GetAutotuneCache() { - static AutotuneCache c; - return c; -} +class AutoTuningCache { + public: + static AutoTuningCache& Instance(); -inline std::string AutotuneCacheFilePath(const std::string& key) { - return std::string{TRITON_JIT_CACHE_DIR} + "/" + - std::to_string(std::hash{}(key)) + ".autotune"; -} + bool Find(const AutoTuningCacheKey& key, TritonJitConfig* config); -inline std::string SerializeConfig(const JitConfig& config) { - std::string s = std::to_string(config.num_warps) + " " + - std::to_string(config.num_stages); - for (const auto& [name, val] : config.constexprs) - s += "\n" + name + " " + std::to_string(val); - return s; -} + void Insert(const AutoTuningCacheKey& key, const TritonJitConfig& config); -inline bool DeserializeConfig(const std::string& content, JitConfig* out) { - std::istringstream iss(content); - std::string line; - if (!std::getline(iss, line)) return false; - std::istringstream head(line); - if (!(head >> out->num_warps >> out->num_stages)) return false; - out->constexprs.clear(); - while (std::getline(iss, line)) { - std::istringstream ls(line); - std::string name; - int val; - if (ls >> name >> val) out->constexprs.push_back({name, val}); - } - return true; -} - -inline bool AutotuneCacheLookup(const std::string& key, JitConfig* out) { - auto& c = GetAutotuneCache(); - std::lock_guard lk(c.mutex); - auto it = c.map.find(key); - if (it != c.map.end()) { - *out = it->second; - return true; - } - std::string path = AutotuneCacheFilePath(key); - if (FileExists(path.c_str())) { - JitConfig parsed; - if (DeserializeConfig(ReadFile(path.c_str()), &parsed)) { - c.map[key] = parsed; - *out = parsed; - return true; - } - } - return false; -} + private: + std::mutex mutex_; -inline void AutotuneCacheInsert(const std::string& key, - const JitConfig& config) { - auto& c = GetAutotuneCache(); - std::lock_guard lk(c.mutex); - c.map[key] = config; - std::string path = AutotuneCacheFilePath(key); - std::string content = SerializeConfig(config); - FILE* f = fopen(path.c_str(), "w"); - if (f) { - fwrite(content.data(), 1, content.size(), f); - fclose(f); - } -} + std::unordered_map entries_; +}; } // namespace infini::ops diff --git a/src/triton/jit/compile.py b/src/triton/jit/compile.py index c1a6d8f77..b86b847b4 100644 --- a/src/triton/jit/compile.py +++ b/src/triton/jit/compile.py @@ -1,163 +1,319 @@ +import functools +import hashlib import importlib.util import json +import os +import tempfile +from contextlib import contextmanager from pathlib import Path import triton import triton.backends -_JIT_DIR = Path(__file__).resolve().parent -_OPS_DIR = _JIT_DIR.parent / "ops" +_OPS_DIR = Path(__file__).resolve().parent.parent / "ops" -_TRITON_BACKEND = { +_TRITON_BACKENDS = { "nvidia": "cuda", } -def _make_target(device, device_id, arch, warp_size): - triton.runtime.driver.set_active(triton.backends.backends[device].driver()) - triton.runtime.driver.active.set_current_device(device_id) - return triton.backends.compiler.GPUTarget(_TRITON_BACKEND[device], arch, warp_size) +def _get_kernel_source_path(op_name): + source_path = _OPS_DIR / op_name / "kernel.py" + if not source_path.is_file(): + raise FileNotFoundError(f"The Triton kernel `{source_path}` does not exist.") -def _do_compile( - op_name, - out_prefix, - num_warps, - num_stages, - device_id, - signature, - device, - arch, - warp_size, -): + return source_path + + +def get_compilation_fingerprint(op_name): + kernel = _load_kernel(op_name) + digest = hashlib.sha256() + digest.update(f"triton={triton.__version__}\n".encode()) + digest.update(b"compiler\0") + digest.update(Path(__file__).read_bytes()) + digest.update(b"triton-jit\0") + digest.update(kernel.cache_key.encode()) + + return digest.hexdigest() + + +@functools.cache +def _load_kernel(op_name): + source_path = _get_kernel_source_path(op_name) + module_spec = importlib.util.spec_from_file_location(source_path.stem, source_path) + + if module_spec is None or module_spec.loader is None: + raise ImportError(f"The Triton kernel `{source_path}` could not be loaded.") - source_path = _OPS_DIR / f"{op_name}/{op_name}.py" - spec = importlib.util.spec_from_file_location(source_path.stem, source_path) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - fn = getattr(mod, "kernel") - while not isinstance(fn, triton.runtime.JITFunction): - fn = fn.fn - - sig_parts = [p.strip() for p in signature.split(",")] if signature else [] - assert len(sig_parts) == len(fn.arg_names), ( - f"signature length {len(sig_parts)} != kernel param count {len(fn.arg_names)}" + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + kernel = getattr(module, "kernel") + + while not isinstance(kernel, triton.runtime.JITFunction): + kernel = kernel.fn + + return kernel + + +def _create_target(target): + backend = target["backend"] + compiler_backend = _TRITON_BACKENDS.get(backend) + if compiler_backend is None: + raise ValueError(f"Unsupported Triton backend: `{backend}`.") + + architecture = target["architecture"] + if backend == "nvidia": + architecture = int(architecture) + + return triton.backends.compiler.GPUTarget( + compiler_backend, architecture, target["warp_size"] ) - sig_dict = {} - const_dict = {} - attr_dict = {} - - constexprs = {} - for part in sig_parts: - if "=" in part: - name, val = part.split("=", 1) - constexprs[name.strip()] = int(val) - - for i, (name, param, part) in enumerate(zip(fn.arg_names, fn.params, sig_parts)): - if param.is_constexpr: - const_dict[(i,)] = constexprs[name] - sig_dict[name] = "constexpr" - elif part.endswith(":1"): - const_dict[(i,)] = 1 - sig_dict[name] = "constexpr" - elif part.endswith(":16"): - sig_dict[name] = part[:-3] - attr_dict[(i,)] = [["tt.divisibility", 16]] + +@contextmanager +def _use_device(target, stream): + previous_driver = triton.runtime.driver.active + previous_device = previous_driver.get_current_device() + active_driver = triton.backends.backends[target["backend"]].driver() + + try: + triton.runtime.driver.set_active(active_driver) + active_driver.set_current_device(target["device_id"]) + device_interface = active_driver.get_device_interface() + + if stream: + benchmark_stream = device_interface.ExternalStream( + stream, device=target["device_id"] + ) + else: + benchmark_stream = device_interface.default_stream(target["device_id"]) + + with device_interface.stream(benchmark_stream): + yield + finally: + triton.runtime.driver.set_active(previous_driver) + previous_driver.set_current_device(previous_device) + + +def _build_ast_source(kernel, runtime_signature, constexprs): + runtime_types = ( + [part.strip() for part in runtime_signature.split(",")] + if runtime_signature + else [] + ) + constexpr_names = { + name + for name, parameter in zip(kernel.arg_names, kernel.params) + if parameter.is_constexpr + } + missing_constexprs = constexpr_names - constexprs.keys() + if missing_constexprs: + names = ", ".join(sorted(missing_constexprs)) + raise ValueError(f"The config does not define constexprs: {names}.") + + unknown_constexprs = constexprs.keys() - constexpr_names + if unknown_constexprs: + names = ", ".join(sorted(unknown_constexprs)) + raise ValueError(f"The config defines unknown constexprs: {names}.") + + expected_runtime_count = len(kernel.arg_names) - len(constexpr_names) + if len(runtime_types) != expected_runtime_count: + raise ValueError( + f"The runtime signature has {len(runtime_types)} entries, but `kernel` " + f"has {expected_runtime_count} runtime parameters." + ) + + signature_types = {} + constants = {} + attributes = {} + + runtime_type_iterator = iter(runtime_types) + for index, (name, parameter) in enumerate(zip(kernel.arg_names, kernel.params)): + if parameter.is_constexpr: + constants[(index,)] = constexprs[name] + signature_types[name] = "constexpr" + continue + + runtime_type = next(runtime_type_iterator) + if runtime_type.endswith(":1"): + constants[(index,)] = 1 + signature_types[name] = "constexpr" + elif runtime_type.endswith(":16"): + signature_types[name] = runtime_type[:-3] + attributes[(index,)] = [["tt.divisibility", 16]] else: - sig_dict[name] = part + signature_types[name] = runtime_type - src = triton.compiler.ASTSource( - fn=fn, signature=sig_dict, constexprs=const_dict, attrs=attr_dict + return triton.compiler.ASTSource( + fn=kernel, + signature=signature_types, + constexprs=constants, + attrs=attributes, ) - target = _make_target(device, device_id, arch, warp_size) - ccinfo = triton.compile( - src, + +def _compile_loaded_kernel(kernel, signature, target, config): + source = _build_ast_source(kernel, signature, config["constexprs"]) + + return triton.compile( + source, target=target, - options={"num_warps": num_warps, "num_stages": num_stages}, + options={ + "num_warps": config["num_warps"], + "num_stages": config["num_stages"], + }, ) - Path(out_prefix).parent.mkdir(parents=True, exist_ok=True) - backend = triton.compiler.make_backend(target) - bin_ext = backend.binary_ext - binary = ccinfo.asm[bin_ext] - with open(out_prefix + "." + bin_ext, "wb") as f: - f.write(binary) - - meta = { - "name": getattr(ccinfo.metadata, "name", fn.__name__), - "binary_ext": bin_ext, - "shared": getattr(ccinfo.metadata, "shared", 0), - "num_warps": getattr(ccinfo.metadata, "num_warps", num_warps), - "device": device, - "target_backend": target.backend, - "arch": str(target.arch), - "global_scratch_size": getattr(ccinfo.metadata, "global_scratch_size", 0), - "profile_scratch_size": getattr(ccinfo.metadata, "profile_scratch_size", 0), - "op_name": op_name, - "signature": signature, + +def _write_artifacts(compiled_kernel, target, output_prefix, cache_identity): + output_prefix = Path(output_prefix) + output_prefix.parent.mkdir(parents=True, exist_ok=True) + + binary_extension = triton.compiler.make_backend(target).binary_ext + binary_path = Path(f"{output_prefix}.{binary_extension}") + metadata_path = Path(f"{output_prefix}.json") + metadata = { + "name": compiled_kernel.metadata.name, + "binary_ext": binary_extension, + "shared": compiled_kernel.metadata.shared, + "global_scratch_size": compiled_kernel.metadata.global_scratch_size, + "profile_scratch_size": compiled_kernel.metadata.profile_scratch_size, + "cache_identity": cache_identity, } - with open(out_prefix + ".json", "w") as f: - json.dump(meta, f) + temporary_paths = [] + + try: + with tempfile.NamedTemporaryFile( + mode="wb", dir=output_prefix.parent, delete=False + ) as binary_file: + binary_file.write(compiled_kernel.asm[binary_extension]) + binary_temporary_path = Path(binary_file.name) + temporary_paths.append(binary_temporary_path) + + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=output_prefix.parent, delete=False + ) as metadata_file: + json.dump(metadata, metadata_file, sort_keys=True) + metadata_temporary_path = Path(metadata_file.name) + temporary_paths.append(metadata_temporary_path) + + os.replace(binary_temporary_path, binary_path) + os.replace(metadata_temporary_path, metadata_path) + finally: + for temporary_path in temporary_paths: + temporary_path.unlink(missing_ok=True) + + +def compile_kernel(op_name, output_prefix, signature, config, target, cache_identity): + kernel = _load_kernel(op_name) + compilation_target = _create_target(target) + compiled_kernel = _compile_loaded_kernel( + kernel, signature, compilation_target, config + ) + _write_artifacts(compiled_kernel, compilation_target, output_prefix, cache_identity) + + +def _build_launch_arguments(kernel, runtime_arguments, constexprs): + runtime_argument_iterator = iter(runtime_arguments) + launch_arguments = [] + sentinel = object() + + for name, parameter in zip(kernel.arg_names, kernel.params): + if parameter.is_constexpr: + if name not in constexprs: + raise ValueError(f"The config does not define constexpr `{name}`.") + launch_arguments.append(constexprs[name]) + continue + + try: + launch_arguments.append(next(runtime_argument_iterator)) + except StopIteration: + raise ValueError( + f"No runtime argument was provided for `{name}`." + ) from None + + if next(runtime_argument_iterator, sentinel) is not sentinel: + raise ValueError("More runtime arguments were provided than expected.") + return launch_arguments -def _load_kernel_fn(op_name): - source_path = _OPS_DIR / f"{op_name}/{op_name}.py" - spec = importlib.util.spec_from_file_location(source_path.stem, source_path) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - fn = getattr(mod, "kernel") - while not isinstance(fn, triton.runtime.JITFunction): - fn = fn.fn - return fn + +def _benchmark_candidate( + kernel, + compiled_kernel, + arguments, + candidate, + warmup_milliseconds, + repetition_milliseconds, +): + kernel_call = functools.partial( + compiled_kernel[tuple(candidate["grid"])], + *_build_launch_arguments(kernel, arguments, candidate["constexprs"]), + ) + + return triton.testing.do_bench( + kernel_call, + warmup=warmup_milliseconds, + rep=repetition_milliseconds, + return_mode="median", + ) -def _do_autotune( +def auto_tune( op_name, - configs, - args, - grids, - warmup, - rep, - device_id, - device, - arch, - warp_size, + candidates, + arguments, + stream, + warmup_milliseconds, + repetition_milliseconds, + target, ): - fn = _load_kernel_fn(op_name) - best_idx = 0 + if not candidates: + raise ValueError("At least one auto-tuning candidate is required.") + + if warmup_milliseconds < 0: + raise ValueError("The auto-tuning warmup duration must not be negative.") + + if repetition_milliseconds <= 0: + raise ValueError("The auto-tuning repetition duration must be positive.") + + kernel = _load_kernel(op_name) + compilation_target = _create_target(target) + best_index = None best_time = float("inf") - _make_target(device, device_id, arch, warp_size) - for i, cand in enumerate(configs): - constexprs = {kv[0]: kv[1] for kv in cand["constexprs"]} - num_warps = cand["num_warps"] - num_stages = cand["num_stages"] - grid = tuple(grids[i]) - - out_prefix = cand["out_prefix"] - _do_compile( - op_name, - out_prefix, - num_warps, - num_stages, - device_id, - cand["full_sig"], - device, - arch, - warp_size, - ) - def _kernel_call(g=grid, a=args, ce=constexprs, nw=num_warps, ns=num_stages): - fn[g](*a, **ce, num_warps=nw, num_stages=ns) + with _use_device(target, stream): + for index, candidate in enumerate(candidates): + try: + compiled_kernel = _compile_loaded_kernel( + kernel, + candidate["signature"], + compilation_target, + candidate, + ) + elapsed_time = _benchmark_candidate( + kernel, + compiled_kernel, + arguments, + candidate, + warmup_milliseconds, + repetition_milliseconds, + ) + _write_artifacts( + compiled_kernel, + compilation_target, + candidate["output_prefix"], + candidate["cache_identity"], + ) + except (triton.CompilationError, triton.OutOfResources): + continue - try: - t = triton.testing.do_bench( - _kernel_call, warmup=warmup, rep=rep, quantiles=(0.5, 0.2, 0.8) - )[0] - if t < best_time: - best_time = t - best_idx = i - except Exception: - pass - return best_idx + if elapsed_time < best_time: + best_time = elapsed_time + best_index = index + + if best_index is None: + raise RuntimeError("No auto-tuning candidate completed successfully.") + + return best_index diff --git a/src/triton/jit/compiler.cc b/src/triton/jit/compiler.cc index f53d99015..b48cea4f3 100644 --- a/src/triton/jit/compiler.cc +++ b/src/triton/jit/compiler.cc @@ -1,151 +1,366 @@ -#include +#include "triton/jit/compiler.h" +#include + +#include #include +#include +#include #include - -#include "jit.h" +#include +#include +#include namespace infini::ops { +namespace { -bool CompilerInit() { - static std::once_flag flag; - static bool ready = false; - - std::call_once(flag, [] { - namespace py = pybind11; - - auto setup = [] { py::module_::import("infini.triton.jit.compile"); }; - - if (Py_IsInitialized()) { - py::gil_scoped_acquire gil; - try { - setup(); - ready = true; - } catch (const py::error_already_set& e) { - fprintf(stderr, "jit init: %s\n", e.what()); - } - } else { - py::initialize_interpreter(false); - try { - setup(); - ready = true; - } catch (const py::error_already_set& e) { - fprintf(stderr, "jit init: %s\n", e.what()); - } - (void)PyEval_SaveThread(); - } - }); - - return ready; -} - -int CompileKernel(const TargetInfo& target, const char* op_name, - const char* out_prefix, int num_warps, int num_stages, - const char* signature) { - if (!CompilerInit()) return -1; - - namespace py = pybind11; - py::gil_scoped_acquire gil; - try { - py::module_ mod = py::module_::import("infini.triton.jit.compile"); - mod.attr("_do_compile")(op_name, out_prefix, num_warps, num_stages, - target.id, signature, target.type, target.arch, - target.warp_size); - return 0; - } catch (const py::error_already_set& e) { - fprintf(stderr, "jit compile: %s\n", e.what()); - return -2; - } -} - -JitConfig AutotuneBench(const char* op_name, - const std::vector& configs, - const std::string& sig, const std::vector& ptrs, - const std::vector& grids, int warmup, int rep, - const char* key, const TargetInfo& target) { - JitConfig cached; - if (AutotuneCacheLookup(key, &cached)) return cached; - - namespace py = pybind11; - if (!CompilerInit()) return configs.empty() ? JitConfig{} : configs[0]; - py::gil_scoped_acquire gil; - try { - py::module_ mod = py::module_::import("infini.triton.jit.compile"); - - py::list cands; - for (const auto& c : configs) { - py::dict cd; - cd["num_warps"] = c.num_warps; - cd["num_stages"] = c.num_stages; - py::list ce; - for (const auto& [k, v] : c.constexprs) { - py::tuple kv(2); - kv[0] = k; - kv[1] = v; - ce.append(kv); - } - cd["constexprs"] = ce; - - std::string full_sig = sig; - for (const auto& [k, v] : c.constexprs) - full_sig += k + "=" + std::to_string(v) + ","; - if (!full_sig.empty() && full_sig.back() == ',') full_sig.pop_back(); - cd["full_sig"] = full_sig; - cd["out_prefix"] = std::string(TRITON_JIT_CACHE_DIR) + "/" + - CacheFileKey(op_name, full_sig.c_str(), c.num_warps, - c.num_stages, target.arch); - - cands.append(cd); - } +class GilGuard { + public: + GilGuard() : state_(PyGILState_Ensure()) {} + + GilGuard(const GilGuard&) = delete; + + GilGuard& operator=(const GilGuard&) = delete; + + ~GilGuard() { PyGILState_Release(state_); } + + private: + PyGILState_STATE state_; +}; + +struct PyObjectDeleter { + void operator()(PyObject* object) const { Py_DECREF(object); } +}; + +using PyObjectPtr = std::unique_ptr; + +bool ReportPythonFailure(const char* operation) { + std::fprintf(stderr, "Triton JIT %s failed", operation); + if (PyErr_Occurred() != nullptr) { + std::fputs(":\n", stderr); + PyErr_Print(); + } else { + std::fputs("\n", stderr); + } + return false; +} + +bool IsPythonReady() { + if (Py_IsInitialized() != 0) return true; + + std::fputs( + "Triton JIT requires an initialized Python interpreter; use it through " + "the InfiniOps Python package\n", + stderr); + return false; +} + +PyObjectPtr GetCompilerFunction(const char* name) { + PyObjectPtr module(PyImport_ImportModule("infini.triton.jit.compile")); + if (!module) return nullptr; - py::list args; - size_t ptr_idx = 0; - size_t pos = 0; - while (pos < sig.size()) { - size_t comma = sig.find(',', pos); - std::string part = sig.substr(pos, comma - pos); - pos = (comma == std::string::npos) ? sig.size() : comma + 1; - if (part.empty()) continue; - - if (part[0] == '*') { - uint64_t val = *static_cast(ptrs[ptr_idx++]); - args.append(static_cast(val)); - } else if (part.find(":1") != std::string::npos) { - args.append(1); - } else { - uint64_t val = *static_cast(ptrs[ptr_idx++]); - if (part.compare(0, 4, "fp32") == 0 || part.compare(0, 3, "f32") == 0) { - args.append(*reinterpret_cast(&val)); - } else if (part.compare(0, 4, "fp64") == 0) { - args.append(*reinterpret_cast(&val)); - } else { - args.append(static_cast(val)); - } - } + PyObjectPtr function(PyObject_GetAttrString(module.get(), name)); + if (!function) return nullptr; + if (PyCallable_Check(function.get()) == 0) { + PyErr_Format(PyExc_TypeError, + "The function `infini.triton.jit.compile.%s` is not callable.", + name); + return nullptr; + } + + return function; +} + +bool SetDictItem(PyObject* dictionary, const char* key, PyObjectPtr value) { + return value != nullptr && + PyDict_SetItemString(dictionary, key, value.get()) == 0; +} + +bool AppendListItem(PyObject* list, PyObjectPtr value) { + return value != nullptr && PyList_Append(list, value.get()) == 0; +} + +bool SetTupleItem(PyObject* tuple, Py_ssize_t index, PyObjectPtr value) { + if (!value) return false; + return PyTuple_SetItem(tuple, index, value.release()) == 0; +} + +PyObjectPtr BuildTarget(const JitTarget& target) { + PyObjectPtr result(PyDict_New()); + if (!result || + !SetDictItem(result.get(), "backend", + PyObjectPtr(PyUnicode_FromStringAndSize( + target.backend.data(), target.backend.size()))) || + !SetDictItem( + result.get(), "architecture", + PyObjectPtr(PyUnicode_FromStringAndSize( + target.architecture.data(), target.architecture.size()))) || + !SetDictItem(result.get(), "device_id", + PyObjectPtr(PyLong_FromLong(target.device_id))) || + !SetDictItem(result.get(), "warp_size", + PyObjectPtr(PyLong_FromLong(target.warp_size)))) { + return nullptr; + } + + return result; +} + +PyObjectPtr BuildConfig(const TritonJitConfig& config) { + PyObjectPtr constexprs(PyDict_New()); + if (!constexprs) return nullptr; + + for (const auto& [name, value] : config.constexprs()) { + if (!SetDictItem( + constexprs.get(), name.c_str(), + PyObjectPtr(PyLong_FromLongLong(static_cast(value))))) { + return nullptr; } + } + + PyObjectPtr result(PyDict_New()); + if (!result || + !SetDictItem(result.get(), "num_warps", + PyObjectPtr(PyLong_FromUnsignedLong(config.num_warps()))) || + !SetDictItem(result.get(), "num_stages", + PyObjectPtr(PyLong_FromUnsignedLong(config.num_stages()))) || + !SetDictItem(result.get(), "constexprs", std::move(constexprs))) { + return nullptr; + } + + return result; +} + +PyObjectPtr BuildGrid(const Grid& grid) { + PyObjectPtr result(PyTuple_New(3)); + if (!result || + !SetTupleItem(result.get(), 0, + PyObjectPtr(PyLong_FromUnsignedLong(grid.x))) || + !SetTupleItem(result.get(), 1, + PyObjectPtr(PyLong_FromUnsignedLong(grid.y))) || + !SetTupleItem(result.get(), 2, + PyObjectPtr(PyLong_FromUnsignedLong(grid.z)))) { + return nullptr; + } + + return result; +} + +PyObjectPtr BuildCandidate(const AutoTuningCandidate& candidate) { + PyObjectPtr result = BuildConfig(candidate.config); + if (!result || + !SetDictItem(result.get(), "grid", BuildGrid(candidate.grid)) || + !SetDictItem( + result.get(), "signature", + PyObjectPtr(PyUnicode_FromStringAndSize( + candidate.signature.data(), candidate.signature.size()))) || + !SetDictItem(result.get(), "output_prefix", + PyObjectPtr(PyUnicode_FromStringAndSize( + candidate.output_prefix.data(), + candidate.output_prefix.size()))) || + !SetDictItem(result.get(), "cache_identity", + PyObjectPtr(PyUnicode_FromStringAndSize( + candidate.cache_identity.data(), + candidate.cache_identity.size())))) { + return nullptr; + } - py::list grids_list; - for (const auto& g : grids) { - py::tuple t(3); - t[0] = g.x; - t[1] = g.y; - t[2] = g.z; - grids_list.append(t); + return result; +} + +template +T DecodeBits(std::uint64_t bits) { + T value; + static_assert(sizeof(value) <= sizeof(bits)); + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +PyObjectPtr BuildArgument(const JitArgument& argument) { + switch (argument.type) { + case JitArgumentType::kPointer: + case JitArgumentType::kUInt64: + return PyObjectPtr(PyLong_FromUnsignedLongLong(argument.bits)); + case JitArgumentType::kInt8: + return PyObjectPtr( + PyLong_FromLongLong(DecodeBits(argument.bits))); + case JitArgumentType::kUInt8: + return PyObjectPtr( + PyLong_FromUnsignedLongLong(DecodeBits(argument.bits))); + case JitArgumentType::kInt16: + return PyObjectPtr( + PyLong_FromLongLong(DecodeBits(argument.bits))); + case JitArgumentType::kUInt16: + return PyObjectPtr(PyLong_FromUnsignedLongLong( + DecodeBits(argument.bits))); + case JitArgumentType::kInt32: + return PyObjectPtr( + PyLong_FromLongLong(DecodeBits(argument.bits))); + case JitArgumentType::kUInt32: + return PyObjectPtr(PyLong_FromUnsignedLongLong( + DecodeBits(argument.bits))); + case JitArgumentType::kInt64: + return PyObjectPtr( + PyLong_FromLongLong(DecodeBits(argument.bits))); + case JitArgumentType::kFloat32: + return PyObjectPtr(PyFloat_FromDouble(DecodeBits(argument.bits))); + case JitArgumentType::kFloat64: + return PyObjectPtr(PyFloat_FromDouble(DecodeBits(argument.bits))); + } + + PyErr_SetString(PyExc_ValueError, "Unsupported Triton JIT argument type."); + return nullptr; +} + +PyObjectPtr BuildCandidates( + const std::vector& candidates) { + PyObjectPtr result(PyList_New(0)); + if (!result) return nullptr; + + for (const auto& candidate : candidates) { + if (!AppendListItem(result.get(), BuildCandidate(candidate))) + return nullptr; + } + + return result; +} + +PyObjectPtr BuildArguments(const std::vector& arguments) { + PyObjectPtr result(PyList_New(0)); + if (!result) return nullptr; + + for (const auto& argument : arguments) { + if (!AppendListItem(result.get(), BuildArgument(argument))) return nullptr; + } + + return result; +} + +} // namespace + +bool GetCompilationFingerprint(const std::string& op_name, + std::string* fingerprint) { + if (fingerprint == nullptr) return false; + + static std::mutex cache_mutex; + static std::unordered_map cache; + { + const std::lock_guard lock(cache_mutex); + const auto cached = cache.find(op_name); + if (cached != cache.end()) { + *fingerprint = cached->second; + return true; } + } - int best_idx = mod.attr("_do_autotune")(op_name, cands, args, grids_list, - warmup, rep, target.id, target.type, - target.arch, target.warp_size) - .cast(); - if (best_idx < 0 || best_idx >= static_cast(configs.size())) - best_idx = 0; - JitConfig winner = configs[best_idx]; - AutotuneCacheInsert(key, winner); - return winner; - } catch (const py::error_already_set& e) { - fprintf(stderr, "jit autotune: %s\n", e.what()); - return configs.empty() ? JitConfig{} : configs[0]; + if (!IsPythonReady()) return false; + + GilGuard gil; + PyObjectPtr function = GetCompilerFunction("get_compilation_fingerprint"); + PyObjectPtr python_op_name( + PyUnicode_FromStringAndSize(op_name.data(), op_name.size())); + if (!function || !python_op_name) { + return ReportPythonFailure("fingerprint lookup"); + } + + PyObjectPtr python_result(PyObject_CallFunctionObjArgs( + function.get(), python_op_name.get(), static_cast(nullptr))); + if (!python_result) return ReportPythonFailure("fingerprint lookup"); + + Py_ssize_t size = 0; + const char* value = PyUnicode_AsUTF8AndSize(python_result.get(), &size); + if (value == nullptr) return ReportPythonFailure("fingerprint lookup"); + + { + const std::lock_guard lock(cache_mutex); + const auto cached = + cache.try_emplace(op_name, value, static_cast(size)).first; + *fingerprint = cached->second; } + return true; +} + +bool CompileKernel(const JitTarget& target, const std::string& op_name, + const std::string& output_prefix, + const TritonJitConfig& config, const std::string& signature, + const std::string& cache_identity) { + if (!IsPythonReady()) return false; + + GilGuard gil; + PyObjectPtr function = GetCompilerFunction("compile_kernel"); + PyObjectPtr python_op_name( + PyUnicode_FromStringAndSize(op_name.data(), op_name.size())); + PyObjectPtr python_output_prefix( + PyUnicode_FromStringAndSize(output_prefix.data(), output_prefix.size())); + PyObjectPtr python_signature( + PyUnicode_FromStringAndSize(signature.data(), signature.size())); + PyObjectPtr python_config = BuildConfig(config); + PyObjectPtr python_target = BuildTarget(target); + PyObjectPtr python_cache_identity(PyUnicode_FromStringAndSize( + cache_identity.data(), cache_identity.size())); + if (!function || !python_op_name || !python_output_prefix || + !python_signature || !python_config || !python_target || + !python_cache_identity) { + return ReportPythonFailure("compilation"); + } + + PyObjectPtr python_result(PyObject_CallFunctionObjArgs( + function.get(), python_op_name.get(), python_output_prefix.get(), + python_signature.get(), python_config.get(), python_target.get(), + python_cache_identity.get(), static_cast(nullptr))); + if (!python_result) return ReportPythonFailure("compilation"); + + return true; +} + +bool RunAutoTuning(const JitTarget& target, const std::string& op_name, + const std::vector& candidates, + const std::vector& arguments, void* stream, + int warmup_milliseconds, int repetition_milliseconds, + TritonJitConfig* result) { + if (result == nullptr || candidates.empty() || warmup_milliseconds < 0 || + repetition_milliseconds <= 0 || !IsPythonReady()) { + return false; + } + + GilGuard gil; + PyObjectPtr function = GetCompilerFunction("auto_tune"); + PyObjectPtr python_op_name( + PyUnicode_FromStringAndSize(op_name.data(), op_name.size())); + PyObjectPtr python_candidates = BuildCandidates(candidates); + PyObjectPtr python_arguments = BuildArguments(arguments); + PyObjectPtr python_stream(PyLong_FromVoidPtr(stream)); + PyObjectPtr python_warmup_milliseconds(PyLong_FromLong(warmup_milliseconds)); + PyObjectPtr python_repetition_milliseconds( + PyLong_FromLong(repetition_milliseconds)); + PyObjectPtr python_target = BuildTarget(target); + if (!function || !python_op_name || !python_candidates || !python_arguments || + !python_stream || !python_warmup_milliseconds || + !python_repetition_milliseconds || !python_target) { + return ReportPythonFailure("auto-tuning"); + } + + PyObjectPtr python_result(PyObject_CallFunctionObjArgs( + function.get(), python_op_name.get(), python_candidates.get(), + python_arguments.get(), python_stream.get(), + python_warmup_milliseconds.get(), python_repetition_milliseconds.get(), + python_target.get(), static_cast(nullptr))); + if (!python_result) return ReportPythonFailure("auto-tuning"); + + const long best_index = PyLong_AsLong(python_result.get()); + if (best_index == -1 && PyErr_Occurred() != nullptr) { + return ReportPythonFailure("auto-tuning result decoding"); + } + if (best_index < 0 || + static_cast(best_index) >= candidates.size()) { + std::fprintf(stderr, + "Triton JIT auto-tuning returned invalid candidate index " + "%ld for %zu candidates\n", + best_index, candidates.size()); + return false; + } + + *result = candidates[static_cast(best_index)].config; + return true; } } // namespace infini::ops diff --git a/src/triton/jit/compiler.h b/src/triton/jit/compiler.h new file mode 100644 index 000000000..2128c30e6 --- /dev/null +++ b/src/triton/jit/compiler.h @@ -0,0 +1,40 @@ +#ifndef INFINI_OPS_TRITON_JIT_COMPILER_H_ +#define INFINI_OPS_TRITON_JIT_COMPILER_H_ + +#include +#include + +#include "triton/jit/backend.h" +#include "triton/jit/jit_config.h" + +namespace infini::ops { + +struct AutoTuningCandidate { + TritonJitConfig config; + + Grid grid; + + std::string signature; + + std::string output_prefix; + + std::string cache_identity; +}; + +bool GetCompilationFingerprint(const std::string& op_name, + std::string* fingerprint); + +bool CompileKernel(const JitTarget& target, const std::string& op_name, + const std::string& output_prefix, + const TritonJitConfig& config, const std::string& signature, + const std::string& cache_identity); + +bool RunAutoTuning(const JitTarget& target, const std::string& op_name, + const std::vector& candidates, + const std::vector& arguments, void* stream, + int warmup_milliseconds, int repetition_milliseconds, + TritonJitConfig* result); + +} // namespace infini::ops + +#endif diff --git a/src/triton/jit/jit.h b/src/triton/jit/jit.h index ab955e881..b5568f7b9 100644 --- a/src/triton/jit/jit.h +++ b/src/triton/jit/jit.h @@ -2,174 +2,64 @@ #define INFINI_OPS_TRITON_JIT_H_ #include -#include #include #include +#include +#include #include #include +#include #include -#include "cache.h" #include "data_type.h" -#include "runtime.h" #include "tensor.h" +#include "triton/jit/backend.h" +#include "triton/jit/cache.h" +#include "triton/jit/compiler.h" +#include "triton/jit/jit_config.h" -namespace infini::ops { - -// ---- device support ---- +#ifdef WITH_NVIDIA +#include "triton/jit/backend_nvidia.h" +#endif -template -inline constexpr bool kJitSupported = false; +namespace infini::ops { -template <> -inline constexpr bool kJitSupported = true; +template ::kSupported> +class JitOperatorBase; -template > -struct JitOperatorBase : Op { +template +class JitOperatorBase : public Op { + public: using Op::Op; -}; -template -struct JitOperatorBase {}; - -// ---- compilation & launch ---- - -template -TargetInfo CurrentTarget() { - TargetInfo target; - target.type = Device::StringFromType(kDev); - int dev_id = 0; - if (Runtime::GetDevice(&dev_id) != Runtime::kSuccess) - return target; - target.id = dev_id; - int major = 0, minor = 0; - Runtime::DeviceGetAttribute( - &major, Runtime::kDevAttrComputeCapabilityMajor, dev_id); - Runtime::DeviceGetAttribute( - &minor, Runtime::kDevAttrComputeCapabilityMinor, dev_id); - target.arch = major * 10 + minor; - Runtime::DeviceGetAttribute(&target.warp_size, - Runtime::kDevAttrWarpSize, dev_id); - return target; -} - -template -bool Load(const TargetInfo& target, const char* binary_data, size_t binary_size, - const KernelMeta& meta, typename Driver::Function& func, - typename Driver::Module& mod) { - (void)binary_size; - - if (Driver::ModuleLoadData(&mod, binary_data) != Driver::kSuccess) - return false; + protected: + const TritonJitConfig& jit_config( + const TritonJitConfig& default_config) const { + if (!this->config_ptr_) return default_config; - if (Driver::ModuleGetFunction(&func, mod, meta.name.c_str()) != - Driver::kSuccess) { - Driver::ModuleUnload(mod); - return false; + const auto* config = + dynamic_cast(this->config_ptr_.get()); + return config == nullptr ? default_config : *config; } +}; - if (meta.shared > 49152) { - int optin = 0; - Runtime::DeviceGetAttribute( - &optin, Runtime::kDevAttrMaxSharedMemoryPerBlockOptin, target.id); - int st = 0; - Driver::FuncGetAttribute( - &st, Driver::kFuncAttributeSharedSizeBytes, func); - if (optin < st || meta.shared > static_cast(optin - st)) { - Driver::ModuleUnload(mod); - return false; - } - Driver::FuncSetCacheConfig(func, - Driver::kFuncCachePreferShared); - if (Driver::FuncSetAttribute( - func, Driver::kFuncAttributeMaxDynamicSharedSizeBytes, - optin - st) != Driver::kSuccess) { - Driver::ModuleUnload(mod); - return false; - } - } - - return true; -} - -template -typename Driver::Function GetKernel(const char* op_name, - const char* signature_str, - void* stream, const JitConfig& opts, - unsigned* out_shared) { - TargetInfo target = CurrentTarget(); - - auto r = CacheQuery(op_name, signature_str, opts.num_warps, - opts.num_stages, target.arch, target.id); - if (r.mem_hit) { - *out_shared = r.shared; - return r.func; - } - - KernelMeta meta; - std::string binary_data; - if (!ReadArtifacts(r.out_prefix, &meta, &binary_data)) { - int ret = CompileKernel(target, op_name, r.out_prefix.c_str(), - opts.num_warps, opts.num_stages, signature_str); - if (ret != 0) return nullptr; - if (!ReadArtifacts(r.out_prefix, &meta, &binary_data)) return nullptr; - } - - if (meta.global_scratch_size > 0 || meta.profile_scratch_size > 0) { - fprintf(stderr, "triton jit: scratch not supported yet\n"); - return nullptr; - } - - typename Driver::Function func; - typename Driver::Module mod; - if (!Load(target, binary_data.data(), binary_data.size(), meta, func, - mod)) - return nullptr; - - unsigned shared = meta.shared; - KernelCacheEntry mine{func, shared}; - KernelCacheEntry winner; - if (KernelCacheLookup(r.mem_key, &winner)) { - Driver::ModuleUnload(mod); - func = winner.func; - shared = winner.shared; - } else { - KernelCacheInsert(r.mem_key, mine); - } - - *out_shared = shared; - return func; -} - -template -int LaunchKernel(const char* op_name, const char* signature_str, void* stream, - Grid grid, const JitConfig& config, void** args) { - unsigned shared = 0; - auto func = GetKernel(op_name, signature_str, stream, config, &shared); - if (!func) return -1; - - TargetInfo target = CurrentTarget(); - return Driver::LaunchKernel( - func, grid.x, grid.y, grid.z, config.num_warps * target.warp_size, 1, 1, - shared, static_cast::Stream>(stream), args, - nullptr); -} +template +class JitOperatorBase {}; -// ---- specialization helpers ---- +namespace triton_jit::detail { -inline const char* SpecPtr(uintptr_t v) { return v % 16 == 0 ? ":16" : ""; } +template +inline constexpr bool kAlwaysFalse = false; -template -const char* SpecInt(T v) { - if (v == 1) return ":1"; - if ((v & 15) == 0) return ":16"; - return ""; -} +struct ScalarTypeDescriptor { + const char* name; -// ---- `DataType` → Triton string ---- + JitArgumentType argument_type; +}; -inline const char* DataTypeToTritonType(DataType dt) { - switch (dt) { +inline const char* TritonTypeName(DataType data_type) { + switch (data_type) { case DataType::kFloat16: return "fp16"; case DataType::kBFloat16: @@ -195,130 +85,305 @@ inline const char* DataTypeToTritonType(DataType dt) { case DataType::kUInt64: return "u64"; } - return "fp32"; + return nullptr; } -// ---- C++ scalar type → Triton string ---- - template -const char* ScalarTypeToTritonType() { - if constexpr (std::is_same_v) - return "fp64"; - else if constexpr (std::is_same_v) - return "fp64"; - else if constexpr (std::is_same_v) - return "i32"; - else if constexpr (std::is_integral_v) { - if constexpr (sizeof(T) == 1) return std::is_signed_v ? "i8" : "u8"; - if constexpr (sizeof(T) == 2) return std::is_signed_v ? "i16" : "u16"; - if constexpr (sizeof(T) == 4) return std::is_signed_v ? "i32" : "u32"; - if constexpr (sizeof(T) == 8) return std::is_signed_v ? "i64" : "u64"; +constexpr ScalarTypeDescriptor ScalarDescriptor() { + using Value = std::remove_cv_t; + if constexpr (std::is_same_v) { + return {"i32", JitArgumentType::kInt32}; + } else if constexpr (std::is_integral_v && sizeof(Value) == 1) { + return std::is_signed_v + ? ScalarTypeDescriptor{"i8", JitArgumentType::kInt8} + : ScalarTypeDescriptor{"u8", JitArgumentType::kUInt8}; + } else if constexpr (std::is_integral_v && sizeof(Value) == 2) { + return std::is_signed_v + ? ScalarTypeDescriptor{"i16", JitArgumentType::kInt16} + : ScalarTypeDescriptor{"u16", JitArgumentType::kUInt16}; + } else if constexpr (std::is_integral_v && sizeof(Value) == 4) { + return std::is_signed_v + ? ScalarTypeDescriptor{"i32", JitArgumentType::kInt32} + : ScalarTypeDescriptor{"u32", JitArgumentType::kUInt32}; + } else if constexpr (std::is_integral_v && sizeof(Value) == 8) { + return std::is_signed_v + ? ScalarTypeDescriptor{"i64", JitArgumentType::kInt64} + : ScalarTypeDescriptor{"u64", JitArgumentType::kUInt64}; + } else { + static_assert(kAlwaysFalse, + "unsupported `Triton` scalar argument type"); } - return "i32"; } -// ---- arguments parser ---- +template +std::uint64_t ArgumentBits(const T& value) { + static_assert(sizeof(T) <= sizeof(std::uint64_t), + "`Triton` scalar argument is wider than `uint64_t`"); + std::uint64_t bits = 0; + std::memcpy(&bits, &value, sizeof(T)); + return bits; +} -struct ArgPack { - std::vector ptrs; +class ArgumentPack { + public: + bool Push(const Tensor& tensor) { + const char* type_name = TritonTypeName(tensor.dtype()); + if (type_name == nullptr) return false; + + const auto pointer = reinterpret_cast(tensor.data()); + signature_ += "*" + std::string{type_name}; + if (pointer % 16 == 0) signature_ += ":16"; + signature_ += ","; + + const auto bits = static_cast(pointer); + arguments_.push_back({JitArgumentType::kPointer, bits}); + launch_arguments_.push_back(Store(bits)); + return true; + } - std::deque storage; + template >, int> = 0> + bool Push(T value) { + constexpr auto descriptor = ScalarDescriptor(); + signature_ += descriptor.name; + + if constexpr (std::is_same_v, bool>) { + signature_ += ","; + const std::int32_t normalized = value; + const auto bits = ArgumentBits(normalized); + arguments_.push_back({descriptor.argument_type, bits}); + launch_arguments_.push_back(Store(normalized)); + } else { + bool compile_time_one = false; + if (value == 1) { + signature_ += ":1"; + compile_time_one = true; + } else if ((value & 15) == 0) { + signature_ += ":16"; + } + signature_ += ","; + + const auto bits = ArgumentBits(value); + arguments_.push_back({descriptor.argument_type, bits}); + if (!compile_time_one) launch_arguments_.push_back(Store(value)); + } + return true; + } - std::string sig; + bool Push(float value) { + signature_ += "fp32,"; + const auto bits = ArgumentBits(value); + arguments_.push_back({JitArgumentType::kFloat32, bits}); + launch_arguments_.push_back(Store(value)); + return true; + } + + bool Push(double value) { + signature_ += "fp64,"; + const auto bits = ArgumentBits(value); + arguments_.push_back({JitArgumentType::kFloat64, bits}); + launch_arguments_.push_back(Store(value)); + return true; + } + + std::string RuntimeSignature() const { + std::string signature = signature_; + if (!signature.empty()) signature.pop_back(); + return signature; + } + + const std::vector& arguments() const { return arguments_; } + + void AddScratchArguments() { + void* scratch = Store(0); + launch_arguments_.push_back(scratch); + launch_arguments_.push_back(scratch); + } + + void** launch_arguments() { return launch_arguments_.data(); } + private: template - void* Store(T v) { - static_assert(sizeof(T) <= sizeof(uint64_t), - "scalar arg wider than `uint64_t`"); - uint64_t slot = 0; - std::memcpy(&slot, &v, sizeof(T)); - storage.push_back(slot); - return &storage.back(); + void* Store(const T& value) { + storage_.push_back(ArgumentBits(value)); + return &storage_.back(); } + + std::deque storage_; + + std::vector launch_arguments_; + + std::vector arguments_; + + std::string signature_; }; -inline void PushArg(const Tensor& t, ArgPack& pack) { - auto ptr = reinterpret_cast(t.data()); - pack.sig += - std::string("*") + DataTypeToTritonType(t.dtype()) + SpecPtr(ptr) + ","; - pack.ptrs.push_back(pack.Store(ptr)); +template +bool PushArguments(ArgumentPack* pack, Args&&... args) { + return (pack->Push(std::forward(args)) && ...); } -template , int> = 0> -void PushArg(T v, ArgPack& pack) { - const char* s = SpecInt(v); - pack.sig += std::string(ScalarTypeToTritonType()) + s + ","; - if (std::strcmp(s, ":1") != 0) pack.ptrs.push_back(pack.Store(v)); -} +template +const LoadedKernel* GetOrLoadKernel( + const JitTarget& target, const std::string& compilation_fingerprint, + const std::string& operator_name, const std::string& signature, + const TritonJitConfig& config) { + using Backend = JitBackend; + auto key = KernelCacheKey::Build(target, compilation_fingerprint, + operator_name, signature, config); + auto& cache = KernelCache::Instance(); + if (const auto* loaded = cache.Find(key)) return loaded; + + const std::string output_prefix = key.ArtifactPrefix(); + if (output_prefix.empty()) return nullptr; + + KernelArtifact artifact; + bool artifact_from_cache = + ReadKernelArtifact(output_prefix, key.identity(), &artifact); + for (;;) { + if (!artifact_from_cache && + (!CompileKernel(target, operator_name, output_prefix, config, signature, + key.identity()) || + !ReadKernelArtifact(output_prefix, key.identity(), &artifact))) { + return nullptr; + } -inline void PushArg(float v, ArgPack& pack) { - pack.ptrs.push_back(pack.Store(v)); - pack.sig += "fp32,"; -} + if (artifact.metadata.global_scratch_size > 0 || + artifact.metadata.profile_scratch_size > 0) { + return nullptr; + } + + typename Backend::Module module{}; + typename Backend::Function function{}; + if (Backend::LoadKernel(target, artifact.binary, artifact.metadata, &module, + &function)) { + auto candidate = std::make_unique>( + module, function, artifact.metadata.shared_memory_size); + return cache.InsertOrGet(key, std::move(candidate)); + } -inline void PushArg(double v, ArgPack& pack) { - pack.ptrs.push_back(pack.Store(v)); - pack.sig += "fp64,"; + if (!artifact_from_cache) return nullptr; + artifact_from_cache = false; + } } -// ---- launch wrapper ---- +} // namespace triton_jit::detail -template -int LaunchJit(const char* op, void* stream, Grid grid, const JitConfig& config, - Args&&... args) { - ArgPack pack; - pack.sig.reserve(256); - (PushArg(std::forward(args), pack), ...); - for (const auto& [name, val] : config.constexprs) - pack.sig += name + "=" + std::to_string(val) + ","; - if (!pack.sig.empty()) pack.sig.pop_back(); +template +int LaunchJit(const std::string& operator_name, void* stream, Grid grid, + const TritonJitConfig& config, Args&&... args) { + using Backend = JitBackend; + static_assert(Backend::kSupported, + "the requested `Triton` JIT backend is not supported"); - void* scratch = pack.Store(0); - pack.ptrs.push_back(scratch); - pack.ptrs.push_back(scratch); + JitTarget target; + if (!Backend::GetCurrentTarget(&target)) return -1; - return LaunchKernel(op, pack.sig.c_str(), stream, grid, config, - pack.ptrs.data()); + std::string compilation_fingerprint; + if (!GetCompilationFingerprint(operator_name, &compilation_fingerprint)) { + return -1; + } + + triton_jit::detail::ArgumentPack arguments; + if (!triton_jit::detail::PushArguments(&arguments, + std::forward(args)...)) { + return -1; + } + + const std::string signature = arguments.RuntimeSignature(); + const auto* kernel = triton_jit::detail::GetOrLoadKernel( + target, compilation_fingerprint, operator_name, signature, config); + if (kernel == nullptr) return -1; + + arguments.AddScratchArguments(); + return Backend::Launch(kernel->function(), grid, config.num_warps(), + target.warp_size, kernel->shared_memory_size(), stream, + arguments.launch_arguments()); } -template -int LaunchJitAutotune(const char* op, void* stream, - const AutotuneConfig& config, - const std::vector& key, - const std::vector& dtype, GridFn grid_fn, - Args&&... args) { - TargetInfo target = CurrentTarget(); +template +int LaunchJitWithAutoTuning(const std::string& operator_name, void* stream, + const AutoTuningOptions& options, + const std::vector& key_values, + GridFunction grid_function, Args&&... args) { + using Backend = JitBackend; + static_assert(Backend::kSupported, + "the requested `Triton` JIT backend is not supported"); + if (options.candidates.empty() || options.keys.size() != key_values.size()) { + return -1; + } + + JitTarget target; + if (!Backend::GetCurrentTarget(&target)) return -1; - std::string cache_key = op; - for (auto d : key) cache_key += "|" + std::to_string(d); - for (auto dt : dtype) - cache_key += "|" + std::string(DataTypeToTritonType(dt)); - cache_key += "|sm" + std::to_string(target.arch); + std::string compilation_fingerprint; + if (!GetCompilationFingerprint(operator_name, &compilation_fingerprint)) { + return -1; + } - ArgPack pack; - pack.sig.reserve(256); - (PushArg(std::forward(args), pack), ...); + triton_jit::detail::ArgumentPack arguments; + if (!triton_jit::detail::PushArguments(&arguments, + std::forward(args)...)) { + return -1; + } std::vector grids; - grids.reserve(config.candidates.size()); - for (const auto& c : config.candidates) grids.push_back(grid_fn(c)); + grids.reserve(options.candidates.size()); + for (const auto& candidate : options.candidates) { + const auto grid = grid_function(candidate); + if (!grid.has_value()) return -1; - JitConfig best = - AutotuneBench(op, config.candidates, pack.sig, pack.ptrs, grids, - config.warmup, config.rep, cache_key.c_str(), target); + grids.push_back(*grid); + } + + std::vector unsigned_key_values; + unsigned_key_values.reserve(key_values.size()); + for (const auto value : key_values) { + unsigned_key_values.push_back(static_cast(value)); + } - Grid grid = grid_fn(best); + const auto auto_tuning_key = AutoTuningCacheKey::Build( + target, compilation_fingerprint, operator_name, + arguments.RuntimeSignature(), options.keys, unsigned_key_values, + options.candidates, grids, options.warmup_milliseconds, + options.repetition_milliseconds); + + TritonJitConfig best_config; + auto& auto_tuning_cache = AutoTuningCache::Instance(); + if (!auto_tuning_cache.Find(auto_tuning_key, &best_config)) { + std::vector candidates; + candidates.reserve(options.candidates.size()); + for (std::size_t index = 0; index < options.candidates.size(); ++index) { + const auto& config = options.candidates[index]; + const std::string signature = arguments.RuntimeSignature(); + const auto kernel_key = KernelCacheKey::Build( + target, compilation_fingerprint, operator_name, signature, config); + const std::string output_prefix = kernel_key.ArtifactPrefix(); + if (output_prefix.empty()) return -1; + candidates.push_back({config, grids[index], signature, output_prefix, + kernel_key.identity()}); + } + + if (!RunAutoTuning(target, operator_name, candidates, arguments.arguments(), + stream, options.warmup_milliseconds, + options.repetition_milliseconds, &best_config)) { + return -1; + } + auto_tuning_cache.Insert(auto_tuning_key, best_config); + } - for (const auto& [name, val] : best.constexprs) - pack.sig += name + "=" + std::to_string(val) + ","; - if (!pack.sig.empty()) pack.sig.pop_back(); + const auto grid = grid_function(best_config); + if (!grid.has_value()) return -1; - void* scratch = pack.Store(0); - pack.ptrs.push_back(scratch); - pack.ptrs.push_back(scratch); + const std::string signature = arguments.RuntimeSignature(); + const auto* kernel = triton_jit::detail::GetOrLoadKernel( + target, compilation_fingerprint, operator_name, signature, best_config); + if (kernel == nullptr) return -1; - return LaunchKernel(op, pack.sig.c_str(), stream, grid, best, - pack.ptrs.data()); + arguments.AddScratchArguments(); + return Backend::Launch(kernel->function(), *grid, best_config.num_warps(), + target.warp_size, kernel->shared_memory_size(), stream, + arguments.launch_arguments()); } } // namespace infini::ops diff --git a/src/triton/jit/jit_config.cc b/src/triton/jit/jit_config.cc new file mode 100644 index 000000000..ce1f8d3f3 --- /dev/null +++ b/src/triton/jit/jit_config.cc @@ -0,0 +1,13 @@ +#include "triton/jit/jit_config.h" + +namespace infini::ops { + +const AutoTuningOptions* TritonJitConfig::auto_tuning_options() const { + return nullptr; +} + +const AutoTuningOptions* AutoTuningConfig::auto_tuning_options() const { + return &options_; +} + +} // namespace infini::ops diff --git a/src/triton/jit/jit_config.h b/src/triton/jit/jit_config.h new file mode 100644 index 000000000..baf04182f --- /dev/null +++ b/src/triton/jit/jit_config.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_TRITON_JIT_JIT_CONFIG_H_ +#define INFINI_OPS_TRITON_JIT_JIT_CONFIG_H_ + +#include +#include +#include +#include +#include +#include + +#include "config.h" + +namespace infini::ops { + +struct AutoTuningOptions; + +class TritonJitConfig : public Cloneable { + public: + using Constexprs = std::map>; + + TritonJitConfig() = default; + + TritonJitConfig(unsigned num_warps, unsigned num_stages, + Constexprs constexprs) + : num_warps_(num_warps), + num_stages_(num_stages), + constexprs_(std::move(constexprs)) {} + + virtual const AutoTuningOptions* auto_tuning_options() const; + + unsigned num_warps() const { return num_warps_; } + + unsigned num_stages() const { return num_stages_; } + + const Constexprs& constexprs() const { return constexprs_; } + + const int* FindConstexpr(std::string_view name) const { + const auto it = constexprs_.find(name); + return it == constexprs_.end() ? nullptr : &it->second; + } + + TritonJitConfig WithDefaultConstexprs(const TritonJitConfig& defaults) const { + TritonJitConfig result = *this; + for (const auto& [name, value] : defaults.constexprs_) { + result.constexprs_.try_emplace(name, value); + } + return result; + } + + private: + unsigned num_warps_{4}; + + unsigned num_stages_{3}; + + Constexprs constexprs_; +}; + +struct AutoTuningOptions { + std::vector keys; + + std::vector candidates; + + int warmup_milliseconds{25}; + + int repetition_milliseconds{100}; +}; + +class AutoTuningConfig : public Cloneable { + public: + explicit AutoTuningConfig(AutoTuningOptions options) + : options_(std::move(options)) {} + + const AutoTuningOptions* auto_tuning_options() const override; + + private: + AutoTuningOptions options_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/triton/jit/pybind11_config.h b/src/triton/jit/pybind11_config.h new file mode 100644 index 000000000..e4dfb0131 --- /dev/null +++ b/src/triton/jit/pybind11_config.h @@ -0,0 +1,123 @@ +#ifndef INFINI_OPS_TRITON_JIT_PYBIND11_CONFIG_H_ +#define INFINI_OPS_TRITON_JIT_PYBIND11_CONFIG_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "triton/jit/jit_config.h" + +namespace infini::ops { + +namespace triton_jit::detail { + +inline bool IsKnownConfigField( + const std::string& field, + std::initializer_list known_fields) { + for (const auto known_field : known_fields) { + if (field == known_field) return true; + } + + return false; +} + +inline void ValidateConfigFields( + const pybind11::dict& config_dict, + std::initializer_list known_fields) { + for (const auto& item : config_dict) { + const auto field = item.first.cast(); + if (!IsKnownConfigField(field, known_fields)) { + throw pybind11::value_error("Unknown Triton JIT config field `" + field + + "`."); + } + } +} + +inline TritonJitConfig::Constexprs ParseConstexprs( + const pybind11::dict& constexprs_dict) { + TritonJitConfig::Constexprs constexprs; + + for (const auto& item : constexprs_dict) { + constexprs.emplace(item.first.cast(), item.second.cast()); + } + + return constexprs; +} + +inline TritonJitConfig ParseTritonJitCompileConfig( + const pybind11::dict& config_dict) { + ValidateConfigFields(config_dict, {"num_warps", "num_stages", "constexprs"}); + + const TritonJitConfig defaults; + const auto num_warps = config_dict.contains("num_warps") + ? config_dict["num_warps"].cast() + : defaults.num_warps(); + const auto num_stages = config_dict.contains("num_stages") + ? config_dict["num_stages"].cast() + : defaults.num_stages(); + auto constexprs = + config_dict.contains("constexprs") + ? ParseConstexprs(config_dict["constexprs"].cast()) + : TritonJitConfig::Constexprs{}; + + return TritonJitConfig{num_warps, num_stages, std::move(constexprs)}; +} + +inline AutoTuningOptions ParseAutoTuningOptions( + const pybind11::dict& options_dict) { + ValidateConfigFields( + options_dict, + {"warmup_milliseconds", "repetition_milliseconds", "keys", "candidates"}); + + AutoTuningOptions options; + + if (options_dict.contains("warmup_milliseconds")) { + options.warmup_milliseconds = + options_dict["warmup_milliseconds"].cast(); + } + + if (options_dict.contains("repetition_milliseconds")) { + options.repetition_milliseconds = + options_dict["repetition_milliseconds"].cast(); + } + + if (options_dict.contains("keys")) { + options.keys = options_dict["keys"].cast>(); + } + + if (options_dict.contains("candidates")) { + for (const auto& candidate : + options_dict["candidates"].cast()) { + options.candidates.push_back( + ParseTritonJitCompileConfig(candidate.cast())); + } + } + + return options; +} + +} // namespace triton_jit::detail + +inline std::unique_ptr TritonJitConfigFromPyDict( + const pybind11::dict& config_dict) { + if (!config_dict.contains("auto_tuning")) { + return std::make_unique( + triton_jit::detail::ParseTritonJitCompileConfig(config_dict)); + } + + triton_jit::detail::ValidateConfigFields(config_dict, {"auto_tuning"}); + auto options = triton_jit::detail::ParseAutoTuningOptions( + config_dict["auto_tuning"].cast()); + + return std::make_unique(std::move(options)); +} + +} // namespace infini::ops + +#endif diff --git a/src/triton/ops/add/add.py b/src/triton/ops/add/add.py deleted file mode 100644 index f813cda47..000000000 --- a/src/triton/ops/add/add.py +++ /dev/null @@ -1,53 +0,0 @@ -import triton -import triton.language as tl - - -@triton.jit -def kernel( - x_ptr, - y_ptr, - out_ptr, - out_shape_ptr, - x_stride_ptr, - y_stride_ptr, - out_stride_ptr, - x_contig, - y_contig, - out_contig, - ndim, - n_elements, - alpha, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = (pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)).to(tl.int64) - mask = offsets < n_elements - - if (x_contig != 0) and (y_contig != 0) and (out_contig != 0): - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - tl.store(out_ptr + offsets, x + y * alpha, mask=mask) - else: - x_offs = tl.zeros([BLOCK_SIZE], dtype=tl.int64) - y_offs = tl.zeros([BLOCK_SIZE], dtype=tl.int64) - out_offs = tl.zeros([BLOCK_SIZE], dtype=tl.int64) - tmp = offsets - - for i in range(ndim): - s = tl.load(out_shape_ptr + (ndim - 1 - i)) - d = tmp % s - tmp = tmp // s - x_offs += d * tl.load(x_stride_ptr + (ndim - 1 - i)) - y_offs += d * tl.load(y_stride_ptr + (ndim - 1 - i)) - out_offs += d * tl.load(out_stride_ptr + (ndim - 1 - i)) - - if x_contig != 0: - x_offs = offsets - if y_contig != 0: - y_offs = offsets - if out_contig != 0: - out_offs = offsets - - x = tl.load(x_ptr + x_offs, mask=mask) - y = tl.load(y_ptr + y_offs, mask=mask) - tl.store(out_ptr + out_offs, x + y * alpha, mask=mask) diff --git a/src/triton/ops/add/jit.cc b/src/triton/ops/add/jit.cc index 00c6b0842..0e6cb4b84 100644 --- a/src/triton/ops/add/jit.cc +++ b/src/triton/ops/add/jit.cc @@ -1,92 +1,293 @@ #include "triton/ops/add/jit.h" +#include #include -#include -#include +#include +#include +#include +#include +#include +#include +#include #include "runtime.h" #include "triton/jit/jit.h" namespace infini::ops { +namespace { + +enum class MetadataSection : std::size_t { + kOutputShape, + kInputStrides, + kOtherStrides, + kOutputStrides, + kCount, +}; + +std::size_t MetadataOffset(MetadataSection section, + std::size_t metadata_extent) { + return static_cast(section) * metadata_extent; +} + +const TritonJitConfig& DefaultTritonJitConfig() { + static const TritonJitConfig config{4u, 3u, {{"BLOCK_SIZE", 1024}}}; + return config; +} + +std::vector DefaultAutoTuningCandidates() { + return { + {4u, 3u, {{"BLOCK_SIZE", 256}}}, + {4u, 3u, {{"BLOCK_SIZE", 512}}}, + {8u, 4u, {{"BLOCK_SIZE", 1024}}}, + {8u, 4u, {{"BLOCK_SIZE", 2048}}}, + }; +} + +std::vector BuildMetadata(int ndim, const Tensor::Shape& output_shape, + const Tensor::Strides& input_strides, + const Tensor::Strides& other_strides, + const Tensor::Strides& output_strides) { + const auto metadata_extent = static_cast(std::max(ndim, 1)); + std::vector metadata( + static_cast(MetadataSection::kCount) * metadata_extent, 0); + + for (int dimension = 0; dimension < ndim; ++dimension) { + const auto index = static_cast(dimension); + metadata[MetadataOffset(MetadataSection::kOutputShape, metadata_extent) + + index] = static_cast(output_shape[index]); + metadata[MetadataOffset(MetadataSection::kInputStrides, metadata_extent) + + index] = static_cast(input_strides[index]); + metadata[MetadataOffset(MetadataSection::kOtherStrides, metadata_extent) + + index] = static_cast(other_strides[index]); + metadata[MetadataOffset(MetadataSection::kOutputStrides, metadata_extent) + + index] = static_cast(output_strides[index]); + } + + return metadata; +} + +template +class DeviceBuffer { + public: + using RuntimeType = Runtime; + + explicit DeviceBuffer(void* stream) + : stream_(static_cast(stream)) {} + + ~DeviceBuffer() { + if (data_ == nullptr) return; + + const auto status = RuntimeType::FreeAsync(data_, stream_); + assert(status == RuntimeType::kSuccess && + "Triton JIT `Add` failed to release its metadata buffer"); + (void)status; + } + + DeviceBuffer(const DeviceBuffer&) = delete; + + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + bool Upload(const void* source, std::size_t size) { + if (RuntimeType::Malloc(&data_, size) != RuntimeType::kSuccess) { + data_ = nullptr; + return false; + } + + return RuntimeType::Memcpy(data_, source, size, + RuntimeType::kMemcpyHostToDevice) == + RuntimeType::kSuccess; + } + + void* data() const { return data_; } + + private: + void* data_{nullptr}; + typename RuntimeType::Stream stream_; +}; + +struct MetadataTensorViews { + Tensor output_shape; + Tensor input_strides; + Tensor other_strides; + Tensor output_strides; +}; + +MetadataTensorViews MakeMetadataTensorViews(void* metadata, int ndim, + const Tensor& output) { + const auto metadata_extent = static_cast(std::max(ndim, 1)); + const Tensor::Shape metadata_shape{ + static_cast(metadata_extent)}; + auto* base = static_cast(metadata); + const auto make_tensor = [&](MetadataSection section) { + const auto byte_offset = + MetadataOffset(section, metadata_extent) * sizeof(int64_t); + char* tensor_data = nullptr; + if (base != nullptr) { + tensor_data = base + byte_offset; + } + return Tensor{tensor_data, metadata_shape, DataType::kInt64, + output.device()}; + }; + + return { + make_tensor(MetadataSection::kOutputShape), + make_tensor(MetadataSection::kInputStrides), + make_tensor(MetadataSection::kOtherStrides), + make_tensor(MetadataSection::kOutputStrides), + }; +} + +std::optional> ResolveAutoTuningKeys( + const std::vector& requested_keys, Tensor::Size n_elements, + Tensor::Size ndim) { + std::vector values; + values.reserve(requested_keys.size()); + for (const auto& key : requested_keys) { + if (key == "n_elements") { + values.push_back(n_elements); + } else if (key == "ndim") { + values.push_back(ndim); + } else { + return std::nullopt; + } + } + + return values; +} + +AutoTuningOptions NormalizeAutoTuningOptions( + const AutoTuningOptions& options, const TritonJitConfig& default_config) { + AutoTuningOptions normalized = options; + if (normalized.keys.empty()) { + normalized.keys = {"n_elements"}; + } + + if (normalized.candidates.empty()) { + normalized.candidates = DefaultAutoTuningCandidates(); + } + + for (auto& candidate : normalized.candidates) { + candidate = candidate.WithDefaultConstexprs(default_config); + } + + return normalized; +} + +std::optional MakeGrid(const TritonJitConfig& config, + Tensor::Size n_elements) { + const int* block_size = config.FindConstexpr("BLOCK_SIZE"); + if (block_size == nullptr || *block_size <= 0) return std::nullopt; + + const auto block_size_value = static_cast(*block_size); + const auto block_count = + n_elements / block_size_value + + static_cast(n_elements % block_size_value != 0); + if (block_count > std::numeric_limits::max()) { + return std::nullopt; + } + + return Grid{static_cast(block_count)}; +} + +} // namespace + template void Operator::operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const { + const Device device = out.device(); + if (input.device().type() != device.type() || + input.device().index() != device.index() || + other.device().type() != device.type() || + other.device().index() != device.index()) { + assert(false && + "Triton JIT `Add` inputs and output must be on the same device"); + return; + } + + const Tensor::Size n_elements = out.numel(); + if (n_elements == 0) return; + + ScopedJitDevice device_guard(device.index()); + if (!device_guard.valid()) { + assert(false && "Triton JIT `Add` failed to select the output device"); + return; + } + + if (this->ndim_ > + static_cast(std::numeric_limits::max())) { + assert(false && "Triton JIT `Add` does not support this tensor rank"); + return; + } const int ndim = static_cast(this->ndim_); - std::vector h_meta(4 * std::max(ndim, 1), 0); - for (int i = 0; i < ndim; ++i) { - h_meta[0 * ndim + i] = static_cast(this->out_shape_[i]); - h_meta[1 * ndim + i] = static_cast(this->input_strides_[i]); - h_meta[2 * ndim + i] = static_cast(this->other_strides_[i]); - h_meta[3 * ndim + i] = static_cast(this->out_strides_[i]); + DeviceBuffer metadata_buffer(this->stream_); + const bool needs_metadata = !this->is_input_contiguous_ || + !this->is_other_contiguous_ || + !this->is_out_contiguous_; + if (needs_metadata) { + const auto metadata = + BuildMetadata(ndim, this->out_shape_, this->input_strides_, + this->other_strides_, this->out_strides_); + if (!metadata_buffer.Upload(metadata.data(), + metadata.size() * sizeof(int64_t))) { + assert(false && "Triton JIT `Add` failed to upload tensor metadata"); + return; + } } - const size_t meta_bytes = h_meta.size() * sizeof(int64_t); - void* d_meta = nullptr; - Runtime::Malloc(&d_meta, meta_bytes); - Runtime::Memcpy(d_meta, h_meta.data(), meta_bytes, - Runtime::kMemcpyHostToDevice); - const size_t stride_bytes = ndim * sizeof(int64_t); - - std::vector meta_shape{ - static_cast(std::max(ndim, 1))}; - char* base = static_cast(d_meta); - Tensor d_out_shape{base + stride_bytes * 0, meta_shape, DataType::kInt64, - out.device()}; - Tensor d_input_strides{base + stride_bytes * 1, meta_shape, DataType::kInt64, - out.device()}; - Tensor d_other_strides{base + stride_bytes * 2, meta_shape, DataType::kInt64, - out.device()}; - Tensor d_out_strides{base + stride_bytes * 3, meta_shape, DataType::kInt64, - out.device()}; - - const size_t n_elements = out.numel(); - - static const JitConfig defaults = DefaultConfig(); - std::shared_ptr extension = this->config_.extension(); - auto cfg = std::static_pointer_cast(extension); - - const std::unordered_map args{ - {"n_elements", n_elements}, - {"ndim", ndim}, - }; + const auto metadata_views = + MakeMetadataTensorViews(metadata_buffer.data(), ndim, out); + + const auto& default_config = DefaultTritonJitConfig(); + const auto& config = this->jit_config(default_config); + const auto* auto_tuning_options = config.auto_tuning_options(); + + int result = -1; + if (auto_tuning_options != nullptr) { + const auto options = + NormalizeAutoTuningOptions(*auto_tuning_options, default_config); + if (options.warmup_milliseconds < 0 || + options.repetition_milliseconds <= 0) { + assert(false && "Triton JIT `Add` auto-tuning durations are invalid"); + return; + } + + const auto key_values = + ResolveAutoTuningKeys(options.keys, n_elements, this->ndim_); + if (!key_values.has_value()) { + assert(false && + "Triton JIT `Add` auto-tuning keys must be `n_elements` or " + "`ndim`"); + return; + } - int result; - if (cfg && cfg->autotune) { - auto tune = std::static_pointer_cast(extension); - if (tune->candidates.empty()) tune->candidates = AutotuneConfigs(); - for (auto& c : tune->candidates) c.ApplyDefaults(defaults); - - auto key_names = tune->key.empty() ? DefaultKey() : tune->key; - std::vector key_vals; - for (const auto& name : key_names) key_vals.push_back(args.at(name)); - - result = LaunchJitAutotune( - "add", this->stream_, *tune, key_vals, - {input.dtype(), other.dtype(), out.dtype()}, - [&](const JitConfig& c) { - int block_size = c.At("BLOCK_SIZE"); - return Grid{static_cast((n_elements + block_size - 1) / - block_size)}; + result = LaunchJitWithAutoTuning( + "add", this->stream_, options, *key_values, + [n_elements](const TritonJitConfig& candidate) { + return MakeGrid(candidate, n_elements); }, - input, other, out, d_out_shape, d_input_strides, d_other_strides, - d_out_strides, this->is_input_contiguous_, this->is_other_contiguous_, - this->is_out_contiguous_, ndim, n_elements, alpha); + input, other, out, metadata_views.output_shape, + metadata_views.input_strides, metadata_views.other_strides, + metadata_views.output_strides, this->is_input_contiguous_, + this->is_other_contiguous_, this->is_out_contiguous_, ndim, n_elements, + alpha); } else { - JitConfig config = cfg ? *cfg : defaults; - if (cfg) config.ApplyDefaults(defaults); - const int block_size = config.At("BLOCK_SIZE"); - Grid grid{ - static_cast((n_elements + block_size - 1) / block_size)}; - result = LaunchJit("add", this->stream_, grid, config, input, other, - out, d_out_shape, d_input_strides, d_other_strides, - d_out_strides, this->is_input_contiguous_, - this->is_other_contiguous_, - this->is_out_contiguous_, ndim, n_elements, alpha); - } + const auto effective_config = config.WithDefaultConstexprs(default_config); + const auto grid = MakeGrid(effective_config, n_elements); + if (!grid.has_value()) { + assert(false && + "Triton JIT `Add` requires a positive `BLOCK_SIZE` whose grid " + "fits in an unsigned integer"); + return; + } - Runtime::FreeAsync( - d_meta, static_cast::Stream>(this->stream_)); + result = LaunchJit( + "add", this->stream_, *grid, effective_config, input, other, out, + metadata_views.output_shape, metadata_views.input_strides, + metadata_views.other_strides, metadata_views.output_strides, + this->is_input_contiguous_, this->is_other_contiguous_, + this->is_out_contiguous_, ndim, n_elements, alpha); + } assert(result == 0 && "Triton JIT `Add` launch failed"); } diff --git a/src/triton/ops/add/jit.h b/src/triton/ops/add/jit.h index be5499072..79d83a71a 100644 --- a/src/triton/ops/add/jit.h +++ b/src/triton/ops/add/jit.h @@ -1,9 +1,6 @@ #ifndef INFINI_OPS_TRITON_OPS_ADD_JIT_H_ #define INFINI_OPS_TRITON_OPS_ADD_JIT_H_ -#include -#include - #include "base/add.h" #include "triton/jit/jit.h" @@ -16,19 +13,6 @@ class Operator : public JitOperatorBase { void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const; - - static JitConfig DefaultConfig() { return {4u, 3u, {{"BLOCK_SIZE", 1024}}}; } - - static std::vector DefaultKey() { return {"n_elements"}; } - - static std::vector AutotuneConfigs() { - return { - {4u, 3u, {{"BLOCK_SIZE", 256}}}, - {4u, 3u, {{"BLOCK_SIZE", 512}}}, - {8u, 4u, {{"BLOCK_SIZE", 1024}}}, - {8u, 4u, {{"BLOCK_SIZE", 2048}}}, - }; - } }; } // namespace infini::ops diff --git a/src/triton/ops/add/kernel.py b/src/triton/ops/add/kernel.py new file mode 100644 index 000000000..6cb90bd6f --- /dev/null +++ b/src/triton/ops/add/kernel.py @@ -0,0 +1,66 @@ +import triton +import triton.language as tl + + +@triton.jit +def kernel( + input_ptr, + other_ptr, + output_ptr, + output_shape_ptr, + input_strides_ptr, + other_strides_ptr, + output_strides_ptr, + is_input_contiguous, + is_other_contiguous, + is_output_contiguous, + ndim, + n_elements, + alpha, + BLOCK_SIZE: tl.constexpr, +): + program_id = tl.program_id(0) + linear_offsets = (program_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)).to(tl.int64) + mask = linear_offsets < n_elements + + if ( + (is_input_contiguous != 0) + and (is_other_contiguous != 0) + and (is_output_contiguous != 0) + ): + input_value = tl.load(input_ptr + linear_offsets, mask=mask) + other_value = tl.load(other_ptr + linear_offsets, mask=mask) + tl.store( + output_ptr + linear_offsets, + input_value + other_value * alpha, + mask=mask, + ) + else: + input_offsets = tl.zeros([BLOCK_SIZE], dtype=tl.int64) + other_offsets = tl.zeros([BLOCK_SIZE], dtype=tl.int64) + output_offsets = tl.zeros([BLOCK_SIZE], dtype=tl.int64) + remaining_index = linear_offsets + + for dimension_offset in range(ndim): + dimension = ndim - 1 - dimension_offset + dimension_size = tl.load(output_shape_ptr + dimension) + dimension_index = remaining_index % dimension_size + remaining_index = remaining_index // dimension_size + input_offsets += dimension_index * tl.load(input_strides_ptr + dimension) + other_offsets += dimension_index * tl.load(other_strides_ptr + dimension) + output_offsets += dimension_index * tl.load(output_strides_ptr + dimension) + + if is_input_contiguous != 0: + input_offsets = linear_offsets + if is_other_contiguous != 0: + other_offsets = linear_offsets + if is_output_contiguous != 0: + output_offsets = linear_offsets + + input_value = tl.load(input_ptr + input_offsets, mask=mask) + other_value = tl.load(other_ptr + other_offsets, mask=mask) + tl.store( + output_ptr + output_offsets, + input_value + other_value * alpha, + mask=mask, + ) diff --git a/tests/test_add.py b/tests/test_add.py index 01dd310ff..6f12eb976 100644 --- a/tests/test_add.py +++ b/tests/test_add.py @@ -64,12 +64,31 @@ for alpha in (0.0, 0.5, 1.0, 2.0) ) -_TEST_CASES = tuple( - (*case, *dtype_case) for case in _DEFAULT_ADD_CASES for dtype_case in _DTYPE_CASES -) + tuple( - (*case, *dtype_case) - for case in _ALPHA_ADD_CASES - for dtype_case in _ALPHA_DTYPE_CASES +_EMPTY_ADD_CASE = ( + (0,), + (0,), + (0,), + None, + None, + None, + None, + torch.float32, + 1e-7, + 1e-7, +) + +_TEST_CASES = ( + tuple( + (*case, *dtype_case) + for case in _DEFAULT_ADD_CASES + for dtype_case in _DTYPE_CASES + ) + + tuple( + (*case, *dtype_case) + for case in _ALPHA_ADD_CASES + for dtype_case in _ALPHA_DTYPE_CASES + ) + + (_EMPTY_ADD_CASE,) ) @@ -136,6 +155,83 @@ def test_add( ) +@pytest.mark.smoke +@pytest.mark.parametrize( + "config", + ( + { + "num_warps": 4, + "num_stages": 3, + "constexprs": {"BLOCK_SIZE": 128}, + }, + { + "auto_tuning": { + "warmup_milliseconds": 0, + "repetition_milliseconds": 1, + "keys": ["n_elements"], + "candidates": [ + { + "num_warps": 4, + "num_stages": 3, + "constexprs": {"BLOCK_SIZE": 128}, + }, + { + "num_warps": 4, + "num_stages": 3, + "constexprs": {"BLOCK_SIZE": 64}, + }, + ], + } + }, + ), + ids=("compile-config", "auto-tuning"), +) +def test_add_triton_config(config, device, implementation_index): + if implementation_index != 10: + pytest.skip("requires Triton implementation 10") + + input = torch.randn(257, device=device) + other = torch.randn_like(input) + out = torch.empty_like(input) + stream = torch.cuda.Stream(device=input.device) + + with torch.cuda.stream(stream): + infini.ops.add( + input, + other, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + config=config, + ) + + stream.synchronize() + + torch.testing.assert_close(out, input + other) + + +@pytest.mark.smoke +def test_add_triton_config_rejects_unknown_fields(device, implementation_index): + if implementation_index != 10: + pytest.skip("requires Triton implementation 10") + + input = torch.randn(1, device=device) + other = torch.randn_like(input) + out = torch.empty_like(input) + + with pytest.raises( + ValueError, match="Unknown Triton JIT config field `unexpected`" + ): + infini.ops.add( + input, + other, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + config={"unexpected": 1}, + ) + + def _add(input, other, out, *, alpha, implementation_index=0): kwargs = { "stream": get_stream(input.device), diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 10842c5c4..cd7c91c0e 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -3,6 +3,8 @@ import re import sys +import pytest + def _load_generator_module(): path = ( @@ -318,7 +320,7 @@ class Mul { assert ( "DefaultImplementationIndexForMul(converted_first_tensor.device().type()))" ) in text - assert "std::move(converted_first_tensor)" in text + assert "std::move(converted_first_tensor)" not in text assert text.count("DeviceFromPybind11Handle(input)") == 1 assert ( "config.set_implementation_index(" @@ -362,7 +364,7 @@ class Cat { "auto converted_first_tensor{VectorTensorFromPybind11Handle(inputs)};" in text ) assert "converted_first_tensor.at(0).device().type()" in text - assert "std::move(converted_first_tensor)" in text + assert "std::move(converted_first_tensor)" not in text assert "DeviceFromPybind11Handle(inputs.at(0))" not in text @@ -769,3 +771,178 @@ class MhaFwdKvcache : public Operator { assert 'py::arg("out") = py::none()' not in text assert "std::optional out" not in text assert "OptionalTensorFromPybind11Handle(out)" not in text + + +_TRITON_ADD_SOURCE = """ + + +namespace infini::ops { + +struct Tensor {}; + +class Add { + public: + Add(const Tensor input, const Tensor other, const double alpha, Tensor out) {} + Add(const Tensor input, const Tensor other, Tensor out) {} + + virtual void operator()(const Tensor input, const Tensor other, + const double alpha, Tensor out) const = 0; + virtual void operator()(const Tensor input, const Tensor other, + Tensor out) const = 0; +}; + +} // namespace infini::ops +""" + + +def _make_backend_add(module, tmp_path, monkeypatch, backend): + base_header = tmp_path / "base" / "add.h" + base_header.parent.mkdir() + base_header.write_text(_TRITON_ADD_SOURCE) + monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header) + + parsed = module._OperatorExtractor()("add") + implementation_path = tmp_path / backend / "ops" / "add" / "jit.h" + implementation_path.parent.mkdir(parents=True) + implementation_path.write_text("// JitConfig is intentionally irrelevant.\n") + + return module._Operator( + parsed.name, + parsed.constructors, + parsed.calls, + implementations=[ + module._Implementation(implementation_path, backend), + ], + ) + + +def test_triton_binding_uses_backend_metadata_and_shared_config_parser( + tmp_path, monkeypatch +): + module = _load_generator_module() + operator = _make_backend_add(module, tmp_path, monkeypatch, "triton") + + binding = module._generate_pybind11(operator) + + assert '#include "triton/jit/pybind11_config.h"' in binding + assert "std::unique_ptr triton_config_ptr;" in binding + assert "TritonJitConfigFromPyDict(*config_dict)" in binding + assert "triton_config_ptr->set_implementation_index(" in binding + assert "config.implementation_index()" in binding + assert "if (triton_config_ptr)" in binding + assert "generated_dispatch::MakeAdd(*triton_config_ptr," in binding + assert "return (*op)(handle," in binding + assert "generated_dispatch::CallAdd(handle, config," in binding + assert "std::move(converted_first_tensor)" not in binding + assert "ValidateConfigFields" not in binding + + declarations, _ = module._generate_generated_dispatch_entries(operator) + make_declarations = [ + declaration for declaration in declarations if " MakeAdd(" in declaration + ] + assert len(make_declarations) == len(operator.constructors) + + +def test_native_binding_does_not_scan_implementation_source_for_config_support( + tmp_path, monkeypatch +): + module = _load_generator_module() + operator = _make_backend_add(module, tmp_path, monkeypatch, "native") + + binding = module._generate_pybind11(operator) + + assert "triton/jit/pybind11_config.h" not in binding + assert "config_dict" not in binding + assert "TritonJitConfigFromPyDict" not in binding + + +def test_implementation_index_records_structured_backend_metadata(tmp_path): + module = _load_generator_module() + implementation_path = tmp_path / "triton" / "ops" / "add" / "jit.h" + implementation_path.parent.mkdir(parents=True) + implementation_path.write_text("class Operator;\n") + + index = module._index_impl_headers([tmp_path], {"triton"}) + + assert index["Add"] == [ + module._Implementation(implementation_path, "triton"), + ] + + +def test_triton_config_is_only_exposed_for_constructor_shaped_calls(): + module = _load_generator_module() + parsed = module._OperatorExtractor()("add_rms_norm") + operator = module._Operator( + parsed.name, + parsed.constructors, + parsed.calls, + implementations=[ + module._Implementation( + pathlib.Path("triton/ops/add_rms_norm/jit.h"), + "triton", + ) + ], + ) + assert len(operator.constructors) == 1 + assert len(operator.calls) == 2 + + binding = module._generate_pybind11(operator) + + declarations, _ = module._generate_generated_dispatch_entries(operator) + make_declarations = [ + declaration for declaration in declarations if " MakeAddRmsNorm(" in declaration + ] + + assert binding.count('py::arg("config") = py::none()') == 1 + assert binding.count("TritonJitConfigFromPyDict(*config_dict)") == 1 + assert len(make_declarations) == len(operator.constructors) + + +@pytest.mark.parametrize( + "serialized, expected_path, expected_backend", + ( + ( + "src/triton/ops/add/jit.h", + pathlib.Path("src/triton/ops/add/jit.h"), + "triton", + ), + ( + {"path": "custom/add.h", "backend": "custom"}, + pathlib.Path("custom/add.h"), + "custom", + ), + ), +) +def test_implementation_json_accepts_legacy_and_structured_entries( + serialized, expected_path, expected_backend +): + module = _load_generator_module() + + implementation = module._implementation_from_json(serialized) + + assert implementation.path == expected_path + assert implementation.backend == expected_backend + + +def test_shared_triton_config_parser_has_one_explicit_schema(): + parser_header = ( + pathlib.Path(__file__).resolve().parents[1] + / "src" + / "triton" + / "jit" + / "pybind11_config.h" + ).read_text() + + for field in ( + "auto_tuning", + "num_warps", + "num_stages", + "constexprs", + "warmup_milliseconds", + "repetition_milliseconds", + "keys", + "candidates", + ): + assert f'"{field}"' in parser_header + + assert "namespace triton_jit::detail" in parser_header diff --git a/tests/test_triton_jit_compile.py b/tests/test_triton_jit_compile.py new file mode 100644 index 000000000..1b54b16d3 --- /dev/null +++ b/tests/test_triton_jit_compile.py @@ -0,0 +1,198 @@ +import contextlib +import importlib.util +import pathlib +import sys +import types + +import pytest + + +def _load_compile_module(monkeypatch): + class CompilationError(Exception): + pass + + class OutOfResources(Exception): + pass + + fake_triton = types.ModuleType("triton") + fake_triton.__path__ = [] + fake_triton.__version__ = "3.5.1" + fake_triton.CompilationError = CompilationError + fake_triton.OutOfResources = OutOfResources + fake_triton.compiler = types.SimpleNamespace( + ASTSource=lambda **arguments: arguments + ) + fake_triton.runtime = types.SimpleNamespace(JITFunction=type("JITFunction", (), {})) + fake_triton.testing = types.SimpleNamespace() + + fake_backends = types.ModuleType("triton.backends") + fake_triton.backends = fake_backends + monkeypatch.setitem(sys.modules, "triton", fake_triton) + monkeypatch.setitem(sys.modules, "triton.backends", fake_backends) + + path = ( + pathlib.Path(__file__).resolve().parents[1] + / "src" + / "triton" + / "jit" + / "compile.py" + ) + module_spec = importlib.util.spec_from_file_location( + "triton_jit_compile_under_test", path + ) + assert module_spec is not None + assert module_spec.loader is not None + module = importlib.util.module_from_spec(module_spec) + monkeypatch.setitem(sys.modules, module_spec.name, module) + module_spec.loader.exec_module(module) + + return module, fake_triton + + +def test_constexprs_are_interleaved_in_ast_and_launch(monkeypatch): + module, fake_triton = _load_compile_module(monkeypatch) + kernel = types.SimpleNamespace( + arg_names=["output", "BLOCK_SIZE", "value", "AXIS"], + params=[ + types.SimpleNamespace(is_constexpr=False), + types.SimpleNamespace(is_constexpr=True), + types.SimpleNamespace(is_constexpr=False), + types.SimpleNamespace(is_constexpr=True), + ], + ) + constexprs = {"AXIS": 1, "BLOCK_SIZE": 128} + compilation_target = object() + compile_calls = [] + + def compile_kernel(source, *, target, options): + compile_calls.append((source, target, options)) + + fake_triton.compile = compile_kernel + module._compile_loaded_kernel( + kernel, + "*fp32,i32:16", + compilation_target, + { + "constexprs": constexprs, + "num_warps": 4, + "num_stages": 3, + }, + ) + + assert len(compile_calls) == 1 + source, target, options = compile_calls[0] + assert target is compilation_target + assert options == {"num_warps": 4, "num_stages": 3} + + assert source["signature"] == { + "output": "*fp32", + "BLOCK_SIZE": "constexpr", + "value": "i32", + "AXIS": "constexpr", + } + assert source["constexprs"] == {(1,): 128, (3,): 1} + assert source["attrs"] == {(2,): [["tt.divisibility", 16]]} + + launch_calls = [] + + class CompiledKernel: + def __getitem__(self, grid): + assert grid == (2, 3, 4) + return lambda *arguments: launch_calls.append(arguments) + + def do_bench(function, *, warmup, rep, return_mode): + assert (warmup, rep, return_mode) == (25, 100, "median") + function() + return 3.25 + + fake_triton.testing.do_bench = do_bench + elapsed_time = module._benchmark_candidate( + kernel, + CompiledKernel(), + [0x1000, -7], + {"grid": [2, 3, 4], "constexprs": constexprs}, + 25, + 100, + ) + + assert elapsed_time == 3.25 + assert launch_calls == [(0x1000, 128, -7, 1)] + + +def test_ast_source_rejects_invalid_runtime_and_constexpr_inputs(monkeypatch): + module, _ = _load_compile_module(monkeypatch) + kernel = types.SimpleNamespace( + arg_names=["output", "BLOCK_SIZE", "value"], + params=[ + types.SimpleNamespace(is_constexpr=False), + types.SimpleNamespace(is_constexpr=True), + types.SimpleNamespace(is_constexpr=False), + ], + ) + + with pytest.raises(ValueError, match="has 2 runtime parameters"): + module._build_ast_source(kernel, "*fp32", {"BLOCK_SIZE": 128}) + + with pytest.raises(ValueError, match="does not define constexprs: BLOCK_SIZE"): + module._build_ast_source(kernel, "*fp32,i32", {}) + + with pytest.raises(ValueError, match="unknown constexprs: UNKNOWN"): + module._build_ast_source( + kernel, + "*fp32,i32", + {"BLOCK_SIZE": 128, "UNKNOWN": 1}, + ) + + +def test_out_of_resources_does_not_publish_artifact(monkeypatch): + module, fake_triton = _load_compile_module(monkeypatch) + kernel = object() + compiled_kernel = object() + artifact_calls = [] + + monkeypatch.setattr(module, "_load_kernel", lambda op_name: kernel) + monkeypatch.setattr(module, "_create_target", lambda target: object()) + monkeypatch.setattr( + module, + "_compile_loaded_kernel", + lambda kernel, signature, target, config: compiled_kernel, + ) + monkeypatch.setattr( + module, + "_use_device", + lambda target, stream: contextlib.nullcontext(), + ) + + def raise_out_of_resources(*arguments): + raise fake_triton.OutOfResources + + monkeypatch.setattr(module, "_benchmark_candidate", raise_out_of_resources) + monkeypatch.setattr( + module, + "_write_artifacts", + lambda *arguments: artifact_calls.append(arguments), + ) + candidate = { + "signature": "*fp32,i32", + "constexprs": {"BLOCK_SIZE": 128}, + "grid": [1, 1, 1], + "num_warps": 4, + "num_stages": 3, + "output_prefix": "unused", + "cache_identity": "unused", + } + + with pytest.raises( + RuntimeError, match="No auto-tuning candidate completed successfully" + ): + module.auto_tune( + "add", + [candidate], + [0x1000, 7], + 0, + 25, + 100, + {"backend": "nvidia"}, + ) + + assert artifact_calls == []