feat(triton): add JIT backend with add operator

CMakeLists.txt

@@ -25,6 +25,8 @@ option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
 
 option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF)
 
+option(WITH_TRITON "Enable Triton-generated kernels" OFF)
+
 # Custom `AscendC` kernels under `src/native/ascend/custom/`. `ON` by default
 # so CI and routine dev builds always exercise `implementation_index=1/2`
 # for `RmsNorm` / `AddRmsNorm`. Gated by `WITH_ASCEND` in
@@ -334,6 +336,10 @@ 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.")
+endif()
+
 if(WITH_NVIDIA)
     add_compile_definitions(WITH_NVIDIA=1)
     enable_language(CUDA)

scripts/generate_wrappers.py

@@ -419,6 +419,8 @@ def __init__(self, name, constructors, calls):
 
         self.calls = calls
 
+        self.impl_paths = []
+
 
 def _find_optional_tensor_params(op_name):
     """Return a set of parameter names declared as `std::optional<Tensor>` in
@@ -580,6 +582,52 @@ def _is_data_type_spelling(spelling):
 
     return spelling.rsplit("::", maxsplit=1)[-1] == "DataType"
 
+def _uses_config_extension(impl_paths):
+    pattern = re.compile(r"\bTritonConfig\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<TritonConfig> ConfigFromPyDict(const py::dict& config_dict) {
+      auto config = std::make_shared<TritonConfig>();
+      if (config_dict.contains("autotune")) {
+        config->autotune = true;
+        py::dict autotune_dict = config_dict["autotune"].cast<py::dict>();
+        if (autotune_dict.contains("warmup")) config->warmup = autotune_dict["warmup"].cast<unsigned>();
+        if (autotune_dict.contains("rep")) config->rep = autotune_dict["rep"].cast<unsigned>();
+        if (autotune_dict.contains("configs")) {
+          for (auto candidate : autotune_dict["configs"].cast<py::list>()) {
+            TritonConfig candidate_config;
+            py::dict candidate_dict = candidate.cast<py::dict>();
+            if (candidate_dict.contains("num_warps")) candidate_config.num_warps = candidate_dict["num_warps"].cast<unsigned>();
+            if (candidate_dict.contains("num_stages")) candidate_config.num_stages = candidate_dict["num_stages"].cast<unsigned>();
+            for (auto item : candidate_dict) {
+              std::string key = item.first.cast<std::string>();
+              if (key != "num_warps" && key != "num_stages")
+                candidate_config.constexprs.emplace_back(key, item.second.cast<int>());
+            }
+            config->configs.push_back(std::move(candidate_config));
+          }
+        }
+      } else {
+        if (config_dict.contains("num_warps")) config->num_warps = config_dict["num_warps"].cast<unsigned>();
+        if (config_dict.contains("num_stages")) config->num_stages = config_dict["num_stages"].cast<unsigned>();
+        for (auto item : config_dict) {
+          std::string key = item.first.cast<std::string>();
+          if (key != "num_warps" && key != "num_stages")
+            config->constexprs.emplace_back(key, item.second.cast<int>());
+        }
+      }
+      return config;
+    }""")
+
 
 def _generate_pybind11(operator):
     optional_tensor_params = _find_optional_tensor_params(operator.name)
@@ -774,7 +822,7 @@ def _generate_py_args(node):
 
         return ", ".join(parts)
 
-    def _generate_call(op_name, call, method=True):
+    def _generate_call(op_name, call, method=True, uses_config=False):
         call_params = _generate_params(call)
         call_args = _generate_arguments(call)
 
@@ -793,19 +841,38 @@ def _generate_call(op_name, call, method=True):
                 call_args = _generate_arguments(
                     call, first_tensor_arg, converted_first_tensor_name
                 )
+            extra_params = ""
+            extra_config_init = ""
+            extra_pybind = ""
+            if uses_config:
+                extra_params = ", std::optional<py::dict> config_dict"
+                extra_config_init = (
+                    "    if (config_dict.has_value()) {\n"
+                    "      config.set_extension(ConfigFromPyDict(*config_dict));\n"
+                    "    }\n"
+                )
+                extra_pybind = ', py::arg("config") = py::none()'
+
             params = (
                 f"{call_params}, std::uintptr_t stream, "
-                "std::optional<std::size_t> implementation_index"
+                f"std::optional<std::size_t> implementation_index{extra_params}"
                 if call_params
-                else "std::uintptr_t stream, "
-                "std::optional<std::size_t> implementation_index"
+                else f"std::uintptr_t stream, std::optional<std::size_t> implementation_index{extra_params}"
             )
             py_args = _generate_py_args(call)
             py_args_str = f"{py_args}, " if py_args else ""
             default_impl_index = _default_impl_index_expr(
                 call, converted_first_tensor_name
             )
 
+            if uses_config:
+                dispatch = (
+                    f"    auto op = generated_dispatch::Make{symbol_name}(config, {call_args});\n"
+                    f"    (*op)(handle, {call_args});"
+                )
+            else:
+                dispatch = f"    return generated_dispatch::Call{symbol_name}(handle, config, {call_args});"
+
             return (
                 f'  m.def("{op_name}", []({params}) {{\n'
                 f"    [[maybe_unused]] HostRangeScope host_range_binding_body{{\n"
@@ -822,8 +889,9 @@ def _generate_call(op_name, call, method=True):
                 f"      config.set_implementation_index(\n"
                 f"          {default_impl_index});\n"
                 f"    }}\n"
-                f"    return generated_dispatch::Call{symbol_name}(handle, config, {call_args});\n"
-                f'  }}, {py_args_str}py::kw_only(), py::arg("stream") = 0, py::arg("implementation_index") = py::none());'
+                f"{extra_config_init}"
+                f"{dispatch}\n"
+                f'  }}, {py_args_str}py::kw_only(), py::arg("stream") = 0, py::arg("implementation_index") = py::none(){extra_pybind});'
             )
 
         # The first lambda parameter is conventionally named `self`, but
@@ -870,9 +938,21 @@ 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) for call in operator_calls
+        _generate_call(operator.name, call, method=False, uses_config=supports_triton)
+        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"
+        )
+    else:
+        jit_include = ""
 
     return f"""#ifndef INFINI_OPS_BINDINGS_{op_name.upper()}_H_
 #define INFINI_OPS_BINDINGS_{op_name.upper()}_H_
@@ -886,7 +966,7 @@ def _overload_order_key(node):
 #include "generated/bindings/generated_dispatch.h"
 #include "handle.h"
 #include "host_range_profiler.h"
-#include "pybind11_utils.h"
+#include "pybind11_utils.h"{jit_include}
 
 namespace py = pybind11;
 
@@ -1252,9 +1332,12 @@ def _append_optional_params(prefix, params):
 
     emitted_make_params = set()
 
-    for constructor in operator.constructors:
-        params = _generate_params(constructor)
-        args = _generate_arguments(constructor)
+    make_nodes = list(operator.constructors)
+    if _uses_config_extension(operator.impl_paths):
+        make_nodes.extend(operator.calls)
+    for node in make_nodes:
+        params = _generate_params(node)
+        args = _generate_arguments(node)
         make_params = _append_optional_params("const Config& config", params)
 
         if make_params in emitted_make_params:
@@ -1721,13 +1804,15 @@ def _filter_ops(ops, op_allowlist, *, strict=False):
     return {op_name: ops[op_name] for op_name in op_allowlist if op_name in ops}
 
 
-def _get_all_ops(devices, with_torch=False, with_ninetoothed=False):
+def _get_all_ops(devices, with_torch=False, with_ninetoothed=False, with_triton=False):
     scan_dirs = set(devices)
 
     if with_torch:
         scan_dirs.add("torch")
     if with_ninetoothed:
         scan_dirs.add("ninetoothed")
+    if with_triton:
+        scan_dirs.add("triton")
 
     ops = {}
 
@@ -1776,6 +1861,7 @@ def _generate_op_artifacts(item):
     op_name, impl_paths = item
     extractor = _OperatorExtractor()
     operator = extractor(op_name)
+    operator.impl_paths = impl_paths
     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(
@@ -1940,6 +2026,12 @@ def _dispatch_gen_batch_size():
         help="Fail if `--ops` contains operators unavailable for the active devices.",
     )
 
+    parser.add_argument(
+        "--with-triton",
+        action="store_true",
+        help="Include Triton backend implementations.",
+    )
+
     args = parser.parse_args()
 
     for directory in (_BINDINGS_DIR, _GENERATED_SRC_DIR, _INCLUDE_DIR):
@@ -1954,6 +2046,7 @@ def _dispatch_gen_batch_size():
             args.devices,
             with_torch=args.with_torch,
             with_ninetoothed=args.with_ninetoothed,
+            with_triton=args.with_triton,
         )
 
     ops = _filter_ops(

src/CMakeLists.txt

@@ -189,6 +189,17 @@ if(WITH_NINETOOTHED)
     target_sources(infiniops PRIVATE ${INFINI_OPS_NINETOOTHED_SOURCES})
 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/jit.cc triton/jit/compiler.cc)
+endif()
+
 if(WITH_ILUVATAR)
     set(ILUVATAR_PATTERNS
         "native/cuda/*.cc"
@@ -840,6 +851,10 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
         list(APPEND GENERATOR_ARGS --with-ninetoothed)
     endif()
 
+    if(WITH_TRITON)
+        list(APPEND GENERATOR_ARGS --with-triton)
+    endif()
+
     execute_process(
         COMMAND ${CMAKE_COMMAND} -E env
                 INFINI_RT_INCLUDE_DIRS=${INFINI_RT_INCLUDE_DIRS_ENV}
@@ -1190,6 +1205,7 @@ if(GENERATE_PYTHON_BINDINGS)
         target_include_directories(ops PRIVATE
             ${INFINI_OPS_NINETOOTHED_INCLUDE_DIRS})
     endif()
+
     target_link_libraries(ops PRIVATE infiniops)
 
     # Cambricon generated dispatch is compiled into the Python extension and
@@ -1244,6 +1260,19 @@ if(GENERATE_PYTHON_BINDINGS)
         install(FILES "${PROJECT_SOURCE_DIR}/generated/torch_ops_metadata.json"
             DESTINATION .)
     endif()
+
+    if(WITH_TRITON)
+        # Ship the JIT compiler and kernel sources so Triton JIT operators
+        # can compile kernels at runtime. `compile.py` uses `__file__` to
+        # locate `ops/` relative to itself; both must live under `triton/`.
+        install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/triton/jit/compile.py"
+                DESTINATION triton/jit)
+        install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/triton/ops/"
+                DESTINATION triton/ops
+                FILES_MATCHING
+                PATTERN "*.py"
+                PATTERN "build.py" EXCLUDE)
+    endif()
 endif()
 
 install(TARGETS infiniops

src/config.h

@@ -2,6 +2,7 @@
 #define INFINI_OPS_CONFIG_H_
 
 #include <cstddef>
+#include <memory>
 
 namespace infini::ops {
 
@@ -13,8 +14,15 @@ class Config {
     implementation_index_ = implementation_index;
   }
 
+  void set_extension(std::shared_ptr<Config> extension) {
+    extension_ = std::move(extension);
+  }
+
+  std::shared_ptr<Config> extension() const { return extension_; }
+
  private:
   std::size_t implementation_index_{0};
+  std::shared_ptr<Config> extension_{};
 };
 
 }  // namespace infini::ops