From cb146ebdcf55d2249844084ee2fe6ca694f99c21 Mon Sep 17 00:00:00 2001 From: Zhang Shuo <52872288+fuyou4546@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:39:48 +0000 Subject: [PATCH 1/2] feat(triton): add JIT backend --- CMakeLists.txt | 6 + scripts/generate_wrappers.py | 116 ++++++++++++++-- src/CMakeLists.txt | 34 +++++ src/config.h | 8 ++ src/triton/jit/cache.h | 226 ++++++++++++++++++++++++++++++++ src/triton/jit/compile.py | 132 +++++++++++++++++++ src/triton/jit/compiler.cc | 151 +++++++++++++++++++++ src/triton/jit/jit.cc | 136 +++++++++++++++++++ src/triton/jit/jit.h | 247 +++++++++++++++++++++++++++++++++++ src/triton/ops/add/add.py | 52 ++++++++ src/triton/ops/add/jit.h | 101 ++++++++++++++ 11 files changed, 1197 insertions(+), 12 deletions(-) create mode 100644 src/triton/jit/cache.h create mode 100644 src/triton/jit/compile.py create mode 100644 src/triton/jit/compiler.cc create mode 100644 src/triton/jit/jit.cc create mode 100644 src/triton/jit/jit.h create mode 100644 src/triton/ops/add/add.py create mode 100644 src/triton/ops/add/jit.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 293b15dce..26ae3ba80 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index f69ba1396..624a1a6fc 100644 --- a/scripts/generate_wrappers.py +++ b/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` in @@ -580,6 +582,53 @@ def _is_data_type_spelling(spelling): return spelling.rsplit("::", maxsplit=1)[-1] == "DataType" +def _uses_config_extension(impl_paths): + pattern = re.compile(r"\bconfig_t\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 config_from_py_dict(const py::dict& d) { + namespace py = pybind11; + auto ext = std::make_shared(); + if (d.contains("autotune")) { + ext->autotune = true; + py::dict at = d["autotune"].cast(); + if (at.contains("warmup")) ext->warmup = at["warmup"].cast(); + if (at.contains("rep")) ext->rep = at["rep"].cast(); + if (at.contains("configs")) { + for (auto cand : at["configs"].cast()) { + config_t c; + py::dict cd = cand.cast(); + if (cd.contains("num_warps")) c.num_warps = cd["num_warps"].cast(); + if (cd.contains("num_stages")) c.num_stages = cd["num_stages"].cast(); + for (auto item : cd) { + std::string key = item.first.cast(); + if (key != "num_warps" && key != "num_stages") + c.constexprs.emplace_back(key, item.second.cast()); + } + ext->configs.push_back(std::move(c)); + } + } + } else { + if (d.contains("num_warps")) ext->num_warps = d["num_warps"].cast(); + if (d.contains("num_stages")) ext->num_stages = d["num_stages"].cast(); + for (auto item : d) { + std::string key = item.first.cast(); + if (key != "num_warps" && key != "num_stages") + ext->constexprs.emplace_back(key, item.second.cast()); + } + } + return ext; + }""") + def _generate_pybind11(operator): optional_tensor_params = _find_optional_tensor_params(operator.name) @@ -774,7 +823,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,12 +842,23 @@ 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 config_dict" + extra_config_init = ( + " if (config_dict.has_value()) {\n" + " config.set_extension(config_from_py_dict(*config_dict));\n" + " }\n" + ) + extra_pybind = ', py::arg("config") = py::none()' + params = ( f"{call_params}, std::uintptr_t stream, " - "std::optional implementation_index" + f"std::optional implementation_index{extra_params}" if call_params - else "std::uintptr_t stream, " - "std::optional implementation_index" + else f"std::uintptr_t stream, std::optional implementation_index{extra_params}" ) py_args = _generate_py_args(call) py_args_str = f"{py_args}, " if py_args else "" @@ -806,6 +866,14 @@ def _generate_call(op_name, call, method=True): 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 +890,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 +939,19 @@ 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' + "namespace infini::ops {\n" + _generate_triton_jit_config_parser() + "\n}\n" + ) + else: + jit_include = "" return f"""#ifndef INFINI_OPS_BINDINGS_{op_name.upper()}_H_ #define INFINI_OPS_BINDINGS_{op_name.upper()}_H_ @@ -886,7 +965,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 +1331,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 +1803,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 +1860,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 +2025,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 +2045,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( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1edeebebd..2d248eafc 100644 --- a/src/CMakeLists.txt +++ b/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,12 @@ if(GENERATE_PYTHON_BINDINGS) target_include_directories(ops PRIVATE ${INFINI_OPS_NINETOOTHED_INCLUDE_DIRS}) endif() + + if(WITH_TRITON) + target_include_directories(ops PRIVATE + ${INFINIOPS_TRITON_INCLUDE_DIRS}) + endif() + target_link_libraries(ops PRIVATE infiniops) # Cambricon generated dispatch is compiled into the Python extension and @@ -1244,6 +1265,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 diff --git a/src/config.h b/src/config.h index a8b59a4fd..15bb430ca 100644 --- a/src/config.h +++ b/src/config.h @@ -2,6 +2,7 @@ #define INFINI_OPS_CONFIG_H_ #include +#include namespace infini::ops { @@ -13,8 +14,15 @@ class Config { 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/triton/jit/cache.h b/src/triton/jit/cache.h new file mode 100644 index 000000000..e9b4480ed --- /dev/null +++ b/src/triton/jit/cache.h @@ -0,0 +1,226 @@ +#ifndef INFINI_OPS_TRITON_JIT_CACHE_H_ +#define INFINI_OPS_TRITON_JIT_CACHE_H_ + +#include +#include +#include +#include +#include +#include + +#include "jit.h" + +namespace infini::ops { + +// ---- file helpers ---- + +inline bool file_exists(const char* path) { + FILE* f = fopen(path, "rb"); + if (f != nullptr) { + fclose(f); + return true; + } + return false; +} + +inline std::string read_file(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; +} + +inline bool cache_complete(const std::string& cubin_path, + const std::string& meta_path) { + return file_exists(cubin_path.c_str()) && file_exists(meta_path.c_str()); +} + +// ---- json field extraction ---- + +inline int json_get_int(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 json_get_string(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); +} + +// ---- kernel cache ---- + +inline std::string generate_desc(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 cache_mem_key(const char* op_name, const char* signature_str, + unsigned num_warps, unsigned num_stages, + int arch, int dev_id) { + return generate_desc(op_name, signature_str, num_warps, num_stages, arch) + + "|dev" + std::to_string(dev_id); +} + +inline std::string cache_file_key(const char* op_name, + const char* signature_str, unsigned num_warps, + unsigned num_stages, int arch) { + return std::to_string(std::hash{}( + generate_desc(op_name, signature_str, num_warps, num_stages, arch))); +} + +struct kernel_cache_entry_t { + void* func; + unsigned shared; +}; + +struct kernel_cache_t { + std::mutex mutex; + std::unordered_map map; +}; + +inline kernel_cache_t& kernel_cache() { + static kernel_cache_t c; + return c; +} + +inline bool kernel_cache_lookup(const std::string& key, + kernel_cache_entry_t* out) { + auto& c = kernel_cache(); + std::lock_guard lk(c.mutex); + auto it = c.map.find(key); + if (it == c.map.end()) return false; + *out = it->second; + return true; +} + +inline void kernel_cache_insert(const std::string& key, + kernel_cache_entry_t entry) { + auto& c = kernel_cache(); + std::lock_guard lk(c.mutex); + c.map[key] = entry; +} + +struct cache_query_result_t { + bool mem_hit; + void* func; + unsigned shared; + std::string out_prefix; + std::string mem_key; +}; + +inline cache_query_result_t cache_query(const char* op, const char* sig, + unsigned num_warps, unsigned num_stages, + int arch, int dev_id) { + auto mem_key = cache_mem_key(op, sig, num_warps, num_stages, arch, dev_id); + kernel_cache_entry_t entry; + if (kernel_cache_lookup(mem_key, &entry)) + return {true, entry.func, entry.shared, "", mem_key}; + auto desc = generate_desc(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}; +} + +struct autotune_cache_t { + std::mutex mutex; + std::unordered_map map; +}; + +inline autotune_cache_t& autotune_cache() { + static autotune_cache_t c; + return c; +} + +inline std::string autotune_cache_file_path(const std::string& key) { + return std::string(TRITON_JIT_CACHE_DIR) + "/" + + std::to_string(std::hash{}(key)) + ".autotune"; +} + +inline std::string serialize_config(const config_t& 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; +} + +inline bool deserialize_config(const std::string& content, config_t* 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 autotune_cache_lookup(const std::string& key, config_t* out) { + auto& c = autotune_cache(); + 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 = autotune_cache_file_path(key); + if (file_exists(path.c_str())) { + config_t parsed; + if (deserialize_config(read_file(path.c_str()), &parsed)) { + c.map[key] = parsed; + *out = parsed; + return true; + } + } + return false; +} + +inline void autotune_cache_insert(const std::string& key, + const config_t& config) { + auto& c = autotune_cache(); + std::lock_guard lk(c.mutex); + c.map[key] = config; + std::string path = autotune_cache_file_path(key); + std::string content = serialize_config(config); + FILE* f = fopen(path.c_str(), "w"); + if (f) { + fwrite(content.data(), 1, content.size(), f); + fclose(f); + } +} + +} // namespace infini::ops + +#endif diff --git a/src/triton/jit/compile.py b/src/triton/jit/compile.py new file mode 100644 index 000000000..c420327ee --- /dev/null +++ b/src/triton/jit/compile.py @@ -0,0 +1,132 @@ +import importlib.util +import json +from pathlib import Path + +import torch +import triton + + +_JIT_DIR = Path(__file__).resolve().parent +_OPS_DIR = _JIT_DIR.parent / "ops" + + +def _do_compile( + op_name, + out_prefix, + num_warps, + num_stages, + device_id, + signature, +): + + 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)}" + ) + + 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]] + else: + sig_dict[name] = part + + src = triton.compiler.ASTSource( + fn=fn, signature=sig_dict, constexprs=const_dict, attrs=attr_dict + ) + + with torch.cuda.device(device_id): + target = triton.runtime.driver.active.get_current_target() + ccinfo = triton.compile( + src, + target=target, + options={"num_warps": num_warps, "num_stages": num_stages}, + ) + + Path(out_prefix).parent.mkdir(parents=True, exist_ok=True) + backend = triton.compiler.make_backend(target) + bin_ext = backend.binary_ext + cubin = ccinfo.asm[bin_ext] + with open(out_prefix + ".cubin", "wb") as f: + f.write(cubin) + + meta = { + "name": getattr(ccinfo.metadata, "name", fn.__name__), + "shared": getattr(ccinfo.metadata, "shared", 0), + "num_warps": getattr(ccinfo.metadata, "num_warps", num_warps), + "arch": target.arch if hasattr(target, "arch") else 80, + "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, + } + with open(out_prefix + ".json", "w") as f: + json.dump(meta, f) + + +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 _do_autotune(op_name, configs, args, grids, warmup, rep, device_id): + fn = _load_kernel_fn(op_name) + best_idx = 0 + best_time = float("inf") + with torch.cuda.device(device_id): + 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"] + ) + + 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) + + 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 diff --git a/src/triton/jit/compiler.cc b/src/triton/jit/compiler.cc new file mode 100644 index 000000000..2c19d0856 --- /dev/null +++ b/src/triton/jit/compiler.cc @@ -0,0 +1,151 @@ +#include + +#include +#include + +#include "cache.h" +#include "jit.h" + +namespace infini::ops { + +bool compiler_init() { + 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 compile_kernel(const char* op_name, const char* out_prefix, int num_warps, + int num_stages, int device_id, const char* signature) { + if (!compiler_init()) 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, + device_id, signature); + return 0; + } catch (const py::error_already_set& e) { + fprintf(stderr, "jit compile: %s\n", e.what()); + return -2; + } +} + +config_t autotune_bench(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, int device_id) { + config_t cached; + if (autotune_cache_lookup(key, &cached)) return cached; + + namespace py = pybind11; + if (!compiler_init()) return configs.empty() ? config_t{} : configs[0]; + py::gil_scoped_acquire gil; + try { + py::module_ mod = py::module_::import("infini.triton.jit.compile"); + + device_info_t dev = current_device(); + + 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) + "/" + + cache_file_key(op_name, full_sig.c_str(), c.num_warps, + c.num_stages, dev.arch); + + cands.append(cd); + } + + 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)); + } + } + } + + 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); + } + + int best_idx = mod.attr("_do_autotune")(op_name, cands, args, grids_list, + warmup, rep, device_id) + .cast(); + if (best_idx < 0 || best_idx >= static_cast(configs.size())) + best_idx = 0; + config_t winner = configs[best_idx]; + autotune_cache_insert(key, winner); + return winner; + } catch (const py::error_already_set& e) { + fprintf(stderr, "jit autotune: %s\n", e.what()); + return configs.empty() ? config_t{} : configs[0]; + } +} + +} // namespace infini::ops diff --git a/src/triton/jit/jit.cc b/src/triton/jit/jit.cc new file mode 100644 index 000000000..abd1adb0b --- /dev/null +++ b/src/triton/jit/jit.cc @@ -0,0 +1,136 @@ +#include "jit.h" + +#include + +#include +#include + +#include "cache.h" + +namespace infini::ops { + +device_info_t current_device() { + device_info_t info; + CUdevice dev; + if (cuCtxGetDevice(&dev) != CUDA_SUCCESS) return info; + info.id = static_cast(dev); + int major = 0, minor = 0; + cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, + dev); + cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, + dev); + info.arch = major * 10 + minor; + return info; +} + +static CUresult load_cubin(const char* cubin_path, const char* meta_path, + CUfunction* out_func, unsigned* out_shared, + CUmodule* out_mod) { + if (!file_exists(meta_path)) return CUDA_ERROR_FILE_NOT_FOUND; + std::string meta_json = read_file(meta_path); + int shared = json_get_int(meta_json, "shared", 0); + if (shared < 0) shared = 0; + std::string fn_name = json_get_string(meta_json, "name", "kernel"); + + int global_scratch = json_get_int(meta_json, "global_scratch_size", 0); + int profile_scratch = json_get_int(meta_json, "profile_scratch_size", 0); + if (global_scratch > 0 || profile_scratch > 0) { + fprintf(stderr, "triton jit: scratch not supported yet\n"); + return CUDA_ERROR_NOT_SUPPORTED; + } + + CUmodule mod; + CUresult err = cuModuleLoad(&mod, cubin_path); + if (err != CUDA_SUCCESS) return err; + CUfunction func; + err = cuModuleGetFunction(&func, mod, fn_name.c_str()); + if (err != CUDA_SUCCESS) { + cuModuleUnload(mod); + return err; + } + + if (shared > 49152) { + CUdevice dev; + err = cuCtxGetDevice(&dev); + if (err != CUDA_SUCCESS) { + cuModuleUnload(mod); + return err; + } + int optin = 0; + cuDeviceGetAttribute( + &optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, dev); + int st = 0; + cuFuncGetAttribute(&st, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, func); + if (shared > optin - st) { + cuModuleUnload(mod); + return CUDA_ERROR_INVALID_VALUE; + } + cuFuncSetCacheConfig(func, CU_FUNC_CACHE_PREFER_SHARED); + err = cuFuncSetAttribute( + func, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, optin - st); + if (err != CUDA_SUCCESS) { + cuModuleUnload(mod); + return err; + } + } + + *out_func = func; + *out_shared = static_cast(shared); + *out_mod = mod; + return CUDA_SUCCESS; +} + +void* get_kernel(const char* op_name, const char* signature_str, void* stream, + const config_t& opts, unsigned* out_shared) { + device_info_t dev = current_device(); + + auto r = cache_query(op_name, signature_str, opts.num_warps, opts.num_stages, + dev.arch, dev.id); + if (r.mem_hit) { + *out_shared = r.shared; + return r.func; + } + + std::string cubin_path = r.out_prefix + ".cubin"; + std::string meta_path = r.out_prefix + ".json"; + + if (!cache_complete(cubin_path, meta_path)) { + int ret = compile_kernel(op_name, r.out_prefix.c_str(), opts.num_warps, + opts.num_stages, dev.id, signature_str); + if (ret != 0) return nullptr; + } + + CUfunction func; + unsigned shared; + CUmodule mod; + CUresult err = + load_cubin(cubin_path.c_str(), meta_path.c_str(), &func, &shared, &mod); + if (err != CUDA_SUCCESS) return nullptr; + + kernel_cache_entry_t mine{static_cast(func), shared}; + kernel_cache_entry_t winner; + if (kernel_cache_lookup(r.mem_key, &winner)) { + cuModuleUnload(mod); + func = static_cast(winner.func); + shared = winner.shared; + } else { + kernel_cache_insert(r.mem_key, mine); + } + + *out_shared = shared; + return static_cast(func); +} + +int launch_kernel(const char* op_name, const char* signature_str, void* stream, + grid_t grid, config_t opts, void** args) { + CUstream cu_stream = static_cast(stream); + unsigned shared = 0; + void* func_ptr = get_kernel(op_name, signature_str, stream, opts, &shared); + if (func_ptr == nullptr) return static_cast(CUDA_ERROR_UNKNOWN); + + return static_cast(cuLaunchKernel( + static_cast(func_ptr), grid.x, grid.y, grid.z, + opts.num_warps * 32, 1, 1, shared, cu_stream, args, nullptr)); +} + +} // namespace infini::ops diff --git a/src/triton/jit/jit.h b/src/triton/jit/jit.h new file mode 100644 index 000000000..dd4445ac2 --- /dev/null +++ b/src/triton/jit/jit.h @@ -0,0 +1,247 @@ +#ifndef INFINI_OPS_TRITON_JIT_H_ +#define INFINI_OPS_TRITON_JIT_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "config.h" +#include "data_type.h" +#include "tensor.h" + +namespace infini::ops { + +struct config_t : Config { + config_t() = default; + config_t(unsigned num_warps, unsigned num_stages, + std::vector> constexprs) + : num_warps(num_warps), + num_stages(num_stages), + constexprs(std::move(constexprs)) {} + + unsigned num_warps = 4; + unsigned num_stages = 3; + std::vector> constexprs; + + bool autotune = false; + std::vector configs; + int warmup = 5; + int rep = 50; + + bool is_autotune() const { return autotune; } + + 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 apply_defaults(const config_t& 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 grid_t { + unsigned x = 1, y = 1, z = 1; +}; + +struct device_info_t { + int id = 0; + int arch = 0; +}; + +bool compiler_init(); + +int compile_kernel(const char* op_name, const char* out_prefix, int num_warps, + + int num_stages, int device_id, const char* signature); + +int launch_kernel(const char* op_name, const char* signature_str, void* stream, + grid_t grid, config_t config, void** args); + +void* get_kernel(const char* op_name, const char* signature_str, void* stream, + const config_t& config, unsigned* out_shared); + +device_info_t current_device(); + +config_t autotune_bench(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, int device_id); + +// ---- specialization ---- + +inline const char* spec_ptr(uintptr_t v) { return v % 16 == 0 ? ":16" : ""; } + +template +const char* spec_int(T v) { + if (v == 1) return ":1"; + if ((v & 15) == 0) return ":16"; + return ""; +} + +// ---- DataType → Triton string ---- + +inline const char* dtype_to_ttype(DataType dt) { + switch (dt) { + case DataType::kFloat16: + return "fp16"; + case DataType::kBFloat16: + return "bf16"; + case DataType::kFloat32: + return "fp32"; + case DataType::kFloat64: + return "fp64"; + case DataType::kInt8: + return "i8"; + case DataType::kInt16: + return "i16"; + case DataType::kInt32: + return "i32"; + case DataType::kInt64: + return "i64"; + case DataType::kUInt8: + return "u8"; + case DataType::kUInt16: + return "u16"; + case DataType::kUInt32: + return "u32"; + case DataType::kUInt64: + return "u64"; + } + return "fp32"; +} + +// ---- C++ scalar type → Triton string ---- + +template +const char* cstype_to_ttype() { + 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" : "i32"; + } + return "i32"; +} + +// ---- arguments parser ---- + +struct arg_pack_t { + std::vector ptrs; + std::deque storage; + std::string sig; + + template + void* store(T v) { + static_assert(sizeof(T) <= sizeof(uint64_t), + "scalar arg wider than 8 bytes"); + uint64_t slot = 0; + std::memcpy(&slot, &v, sizeof(T)); + storage.push_back(slot); + return &storage.back(); + } +}; + +inline void _push_arg(const Tensor& t, arg_pack_t& pack) { + auto ptr = reinterpret_cast(t.data()); + pack.sig += + std::string("*") + dtype_to_ttype(t.dtype()) + spec_ptr(ptr) + ","; + pack.ptrs.push_back(pack.store(ptr)); +} + +template , int> = 0> +void _push_arg(T v, arg_pack_t& pack) { + const char* s = spec_int(v); + pack.sig += std::string(cstype_to_ttype()) + s + ","; + if (std::strcmp(s, ":1") != 0) pack.ptrs.push_back(pack.store(v)); +} + +inline void _push_arg(float v, arg_pack_t& pack) { + pack.ptrs.push_back(pack.store(v)); + pack.sig += "fp32,"; +} + +inline void _push_arg(double v, arg_pack_t& pack) { + pack.ptrs.push_back(pack.store(v)); + pack.sig += "fp64,"; +} + +// ---- launch wrapper ---- + +template +int launch_jit(const char* op, void* stream, grid_t grid, config_t config, + Args&&... args) { + arg_pack_t pack; + pack.sig.reserve(256); + (_push_arg(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(); + + // triton need + void* scratch = pack.store(0); + pack.ptrs.push_back(scratch); + pack.ptrs.push_back(scratch); + + return launch_kernel(op, pack.sig.c_str(), stream, grid, config, + pack.ptrs.data()); +} + +template +int launch_jit_autotune(const char* op, void* stream, const config_t& config, + const std::vector& key_dims, + DataType dtype, GridFn grid_fn, Args&&... args) { + std::string cache_key = op; + for (auto d : key_dims) cache_key += "|" + std::to_string(d); + cache_key += "|dt=" + std::to_string(static_cast(dtype)); + + arg_pack_t pack; + pack.sig.reserve(256); + (_push_arg(std::forward(args), pack), ...); + + std::vector grids; + grids.reserve(config.configs.size()); + for (const auto& c : config.configs) grids.push_back(grid_fn(c)); + + config_t best = autotune_bench(op, config.configs, pack.sig, pack.ptrs, grids, + config.warmup, config.rep, cache_key.c_str(), + current_device().id); + + grid_t grid = grid_fn(best); + + for (const auto& [name, val] : best.constexprs) + pack.sig += name + "=" + std::to_string(val) + ","; + if (!pack.sig.empty()) pack.sig.pop_back(); + + void* scratch = pack.store(0); + pack.ptrs.push_back(scratch); + pack.ptrs.push_back(scratch); + + return launch_kernel(op, pack.sig.c_str(), stream, grid, best, + pack.ptrs.data()); +} + +} // namespace infini::ops + +#endif diff --git a/src/triton/ops/add/add.py b/src/triton/ops/add/add.py new file mode 100644 index 000000000..eb3e35724 --- /dev/null +++ b/src/triton/ops/add/add.py @@ -0,0 +1,52 @@ +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, + 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, 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, mask=mask) diff --git a/src/triton/ops/add/jit.h b/src/triton/ops/add/jit.h new file mode 100644 index 000000000..42bcf6502 --- /dev/null +++ b/src/triton/ops/add/jit.h @@ -0,0 +1,101 @@ +#ifndef INFINI_OPS_TRITON_JIT_ADD_H_ +#define INFINI_OPS_TRITON_JIT_ADD_H_ + +#include + +#include +#include +#include + +#include "base/add.h" +#include "data_type.h" +#include "triton/jit/jit.h" + +namespace infini::ops { + +template <> +class Operator : public Add { + public: + using Add::Add; + using Add::operator(); + + static config_t default_config() { return {4u, 3u, {{"BLOCK_SIZE", 1024}}}; } + + static std::vector autotune_configs() { + return { + {4u, 3u, {{"BLOCK_SIZE", 256}}}, + {4u, 3u, {{"BLOCK_SIZE", 512}}}, + {8u, 4u, {{"BLOCK_SIZE", 1024}}}, + {8u, 4u, {{"BLOCK_SIZE", 2048}}}, + }; + } + + void operator()(const Tensor input, const Tensor other, + Tensor out) const override { + const int ndim = static_cast(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(out_shape_[i]); + h_meta[1 * ndim + i] = static_cast(input_strides_[i]); + h_meta[2 * ndim + i] = static_cast(other_strides_[i]); + h_meta[3 * ndim + i] = static_cast(out_strides_[i]); + } + const size_t meta_bytes = h_meta.size() * sizeof(int64_t); + CUdeviceptr d_meta; + cuMemAlloc(&d_meta, meta_bytes); + cuMemcpyHtoD(d_meta, h_meta.data(), meta_bytes); + const size_t stride_bytes = ndim * sizeof(int64_t); + + auto meta_shape = + std::vector{static_cast(std::max(ndim, 1))}; + Tensor d_out_shape{reinterpret_cast(d_meta + 0 * stride_bytes), + meta_shape, DataType::kInt64, out.device()}; + Tensor d_input_strides{reinterpret_cast(d_meta + 1 * stride_bytes), + meta_shape, DataType::kInt64, out.device()}; + Tensor d_other_strides{reinterpret_cast(d_meta + 2 * stride_bytes), + meta_shape, DataType::kInt64, out.device()}; + Tensor d_out_strides{reinterpret_cast(d_meta + 3 * stride_bytes), + meta_shape, DataType::kInt64, out.device()}; + + const size_t n_elements = out.numel(); + + auto extension = config_.extension(); + static const config_t defaults = default_config(); + const auto* config_ptr = static_cast(extension.get()); + config_t config = config_ptr ? *config_ptr : defaults; + if (extension) config.apply_defaults(defaults); + + int result; + if (config.is_autotune()) { + if (config.configs.empty()) config.configs = autotune_configs(); + for (auto& c : config.configs) c.apply_defaults(defaults); + result = launch_jit_autotune( + "add", stream_, config, {n_elements}, out.dtype(), + [&](const config_t& c) { + int block_size = c.at("BLOCK_SIZE"); + return grid_t{static_cast((n_elements + block_size - 1) / + block_size)}; + }, + input, other, out, d_out_shape, d_input_strides, d_other_strides, + d_out_strides, is_input_contiguous_, is_other_contiguous_, + is_out_contiguous_, ndim, n_elements); + } else { + const int block_size = config.at("BLOCK_SIZE"); + grid_t grid{ + static_cast((n_elements + block_size - 1) / block_size)}; + result = launch_jit( + "add", stream_, grid, config, input, other, out, d_out_shape, + d_input_strides, d_other_strides, d_out_strides, is_input_contiguous_, + is_other_contiguous_, is_out_contiguous_, ndim, n_elements); + } + + cuMemFreeAsync(d_meta, static_cast(stream_)); + + assert(result == 0 && "Triton JIT `Add` launch failed"); + } +}; + +} // namespace infini::ops + +#endif From 47b565277412fbfa04e3a520b2082e04b539b96d Mon Sep 17 00:00:00 2001 From: Zhang Shuo <52872288+fuyou4546@users.noreply.github.com> Date: Wed, 5 Aug 2026 06:21:46 +0000 Subject: [PATCH 2/2] refactor(triton): modify naming and move JIT to slot 10 --- scripts/generate_wrappers.py | 53 ++++++------- src/CMakeLists.txt | 9 +-- src/triton/jit/cache.h | 121 +++++++++++++++-------------- src/triton/jit/compiler.cc | 37 ++++----- src/triton/jit/jit.cc | 54 ++++++------- src/triton/jit/jit.h | 145 +++++++++++++++++++---------------- src/triton/ops/add/add.py | 5 +- src/triton/ops/add/jit.h | 48 ++++++------ 8 files changed, 245 insertions(+), 227 deletions(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 624a1a6fc..33ac2e98c 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -583,7 +583,7 @@ def _is_data_type_spelling(spelling): return spelling.rsplit("::", maxsplit=1)[-1] == "DataType" def _uses_config_extension(impl_paths): - pattern = re.compile(r"\bconfig_t\b") + pattern = re.compile(r"\bTritonConfig\b") for path in impl_paths: try: if pattern.search(path.read_text()): @@ -595,38 +595,37 @@ def _uses_config_extension(impl_paths): def _generate_triton_jit_config_parser(): return textwrap.dedent("""\ - inline std::shared_ptr config_from_py_dict(const py::dict& d) { - namespace py = pybind11; - auto ext = std::make_shared(); - if (d.contains("autotune")) { - ext->autotune = true; - py::dict at = d["autotune"].cast(); - if (at.contains("warmup")) ext->warmup = at["warmup"].cast(); - if (at.contains("rep")) ext->rep = at["rep"].cast(); - if (at.contains("configs")) { - for (auto cand : at["configs"].cast()) { - config_t c; - py::dict cd = cand.cast(); - if (cd.contains("num_warps")) c.num_warps = cd["num_warps"].cast(); - if (cd.contains("num_stages")) c.num_stages = cd["num_stages"].cast(); - for (auto item : cd) { + inline std::shared_ptr ConfigFromPyDict(const py::dict& config_dict) { + auto config = std::make_shared(); + if (config_dict.contains("autotune")) { + config->autotune = true; + 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("configs")) { + for (auto candidate : autotune_dict["configs"].cast()) { + TritonConfig 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") - c.constexprs.emplace_back(key, item.second.cast()); + candidate_config.constexprs.emplace_back(key, item.second.cast()); } - ext->configs.push_back(std::move(c)); + config->configs.push_back(std::move(candidate_config)); } } } else { - if (d.contains("num_warps")) ext->num_warps = d["num_warps"].cast(); - if (d.contains("num_stages")) ext->num_stages = d["num_stages"].cast(); - for (auto item : d) { + 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") - ext->constexprs.emplace_back(key, item.second.cast()); + config->constexprs.emplace_back(key, item.second.cast()); } } - return ext; + return config; }""") @@ -849,7 +848,7 @@ def _generate_call(op_name, call, method=True, uses_config=False): extra_params = ", std::optional config_dict" extra_config_init = ( " if (config_dict.has_value()) {\n" - " config.set_extension(config_from_py_dict(*config_dict));\n" + " config.set_extension(ConfigFromPyDict(*config_dict));\n" " }\n" ) extra_pybind = ', py::arg("config") = py::none()' @@ -947,8 +946,10 @@ def _overload_order_key(node): ) if supports_triton: jit_include = ( - '\n#include "triton/jit/jit.h"\n' - "namespace infini::ops {\n" + _generate_triton_jit_config_parser() + "\n}\n" + '\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 = "" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 2d248eafc..6f5f8707c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1206,11 +1206,6 @@ if(GENERATE_PYTHON_BINDINGS) ${INFINI_OPS_NINETOOTHED_INCLUDE_DIRS}) endif() - if(WITH_TRITON) - target_include_directories(ops PRIVATE - ${INFINIOPS_TRITON_INCLUDE_DIRS}) - endif() - target_link_libraries(ops PRIVATE infiniops) # Cambricon generated dispatch is compiled into the Python extension and @@ -1268,8 +1263,8 @@ if(GENERATE_PYTHON_BINDINGS) 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/. + # 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/" diff --git a/src/triton/jit/cache.h b/src/triton/jit/cache.h index e9b4480ed..871684e53 100644 --- a/src/triton/jit/cache.h +++ b/src/triton/jit/cache.h @@ -14,7 +14,7 @@ namespace infini::ops { // ---- file helpers ---- -inline bool file_exists(const char* path) { +inline bool FileExists(const char* path) { FILE* f = fopen(path, "rb"); if (f != nullptr) { fclose(f); @@ -23,7 +23,7 @@ inline bool file_exists(const char* path) { return false; } -inline std::string read_file(const char* path) { +inline std::string ReadFile(const char* path) { FILE* f = fopen(path, "rb"); if (f == nullptr) return {}; fseek(f, 0, SEEK_END); @@ -40,15 +40,15 @@ inline std::string read_file(const char* path) { return buf; } -inline bool cache_complete(const std::string& cubin_path, - const std::string& meta_path) { - return file_exists(cubin_path.c_str()) && file_exists(meta_path.c_str()); +inline bool CacheComplete(const std::string& cubin_path, + const std::string& meta_path) { + return FileExists(cubin_path.c_str()) && FileExists(meta_path.c_str()); } -// ---- json field extraction ---- +// ---- JSON field extraction ---- -inline int json_get_int(const std::string& json, const char* key, - int fallback = 0) { +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; @@ -57,8 +57,8 @@ inline int json_get_int(const std::string& json, const char* key, return std::atoi(json.c_str() + pos); } -inline std::string json_get_string(const std::string& json, const char* key, - const char* fallback) { +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; @@ -73,45 +73,46 @@ inline std::string json_get_string(const std::string& json, const char* key, // ---- kernel cache ---- -inline std::string generate_desc(const char* op, const char* sig, - unsigned num_warps, unsigned num_stages, - int arch) { +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 cache_mem_key(const char* op_name, const char* signature_str, - unsigned num_warps, unsigned num_stages, - int arch, int dev_id) { - return generate_desc(op_name, signature_str, num_warps, num_stages, 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 cache_file_key(const char* op_name, - const char* signature_str, unsigned num_warps, - unsigned num_stages, int arch) { +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{}( - generate_desc(op_name, signature_str, num_warps, num_stages, arch))); + GenerateDesc(op_name, signature_str, num_warps, num_stages, arch))); } -struct kernel_cache_entry_t { +struct KernelCacheEntry { void* func; + unsigned shared; }; -struct kernel_cache_t { +struct KernelCache { std::mutex mutex; - std::unordered_map map; + + std::unordered_map map; }; -inline kernel_cache_t& kernel_cache() { - static kernel_cache_t c; +inline KernelCache& GetKernelCache() { + static KernelCache c; return c; } -inline bool kernel_cache_lookup(const std::string& key, - kernel_cache_entry_t* out) { - auto& c = kernel_cache(); +inline 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; @@ -119,51 +120,55 @@ inline bool kernel_cache_lookup(const std::string& key, return true; } -inline void kernel_cache_insert(const std::string& key, - kernel_cache_entry_t entry) { - auto& c = kernel_cache(); +inline void KernelCacheInsert(const std::string& key, KernelCacheEntry entry) { + auto& c = GetKernelCache(); std::lock_guard lk(c.mutex); c.map[key] = entry; } -struct cache_query_result_t { +struct CacheQueryResult { bool mem_hit; + void* func; + unsigned shared; + std::string out_prefix; + std::string mem_key; }; -inline cache_query_result_t cache_query(const char* op, const char* sig, - unsigned num_warps, unsigned num_stages, - int arch, int dev_id) { - auto mem_key = cache_mem_key(op, sig, num_warps, num_stages, arch, dev_id); - kernel_cache_entry_t entry; - if (kernel_cache_lookup(mem_key, &entry)) +inline 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 = generate_desc(op, sig, num_warps, num_stages, arch); + 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}; } -struct autotune_cache_t { +struct AutotuneCache { std::mutex mutex; - std::unordered_map map; + + std::unordered_map map; }; -inline autotune_cache_t& autotune_cache() { - static autotune_cache_t c; +inline AutotuneCache& GetAutotuneCache() { + static AutotuneCache c; return c; } -inline std::string autotune_cache_file_path(const std::string& key) { +inline std::string AutotuneCacheFilePath(const std::string& key) { return std::string(TRITON_JIT_CACHE_DIR) + "/" + std::to_string(std::hash{}(key)) + ".autotune"; } -inline std::string serialize_config(const config_t& config) { +inline std::string SerializeConfig(const TritonConfig& config) { std::string s = std::to_string(config.num_warps) + " " + std::to_string(config.num_stages); for (const auto& [name, val] : config.constexprs) @@ -171,7 +176,7 @@ inline std::string serialize_config(const config_t& config) { return s; } -inline bool deserialize_config(const std::string& content, config_t* out) { +inline bool DeserializeConfig(const std::string& content, TritonConfig* out) { std::istringstream iss(content); std::string line; if (!std::getline(iss, line)) return false; @@ -187,18 +192,18 @@ inline bool deserialize_config(const std::string& content, config_t* out) { return true; } -inline bool autotune_cache_lookup(const std::string& key, config_t* out) { - auto& c = autotune_cache(); +inline bool AutotuneCacheLookup(const std::string& key, TritonConfig* 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 = autotune_cache_file_path(key); - if (file_exists(path.c_str())) { - config_t parsed; - if (deserialize_config(read_file(path.c_str()), &parsed)) { + std::string path = AutotuneCacheFilePath(key); + if (FileExists(path.c_str())) { + TritonConfig parsed; + if (DeserializeConfig(ReadFile(path.c_str()), &parsed)) { c.map[key] = parsed; *out = parsed; return true; @@ -207,13 +212,13 @@ inline bool autotune_cache_lookup(const std::string& key, config_t* out) { return false; } -inline void autotune_cache_insert(const std::string& key, - const config_t& config) { - auto& c = autotune_cache(); +inline void AutotuneCacheInsert(const std::string& key, + const TritonConfig& config) { + auto& c = GetAutotuneCache(); std::lock_guard lk(c.mutex); c.map[key] = config; - std::string path = autotune_cache_file_path(key); - std::string content = serialize_config(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); diff --git a/src/triton/jit/compiler.cc b/src/triton/jit/compiler.cc index 2c19d0856..a58a3ffcf 100644 --- a/src/triton/jit/compiler.cc +++ b/src/triton/jit/compiler.cc @@ -8,7 +8,7 @@ namespace infini::ops { -bool compiler_init() { +bool CompilerInit() { static std::once_flag flag; static bool ready = false; @@ -40,9 +40,9 @@ bool compiler_init() { return ready; } -int compile_kernel(const char* op_name, const char* out_prefix, int num_warps, - int num_stages, int device_id, const char* signature) { - if (!compiler_init()) return -1; +int CompileKernel(const char* op_name, const char* out_prefix, int num_warps, + int num_stages, int device_id, const char* signature) { + if (!CompilerInit()) return -1; namespace py = pybind11; py::gil_scoped_acquire gil; @@ -57,21 +57,22 @@ int compile_kernel(const char* op_name, const char* out_prefix, int num_warps, } } -config_t autotune_bench(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, int device_id) { - config_t cached; - if (autotune_cache_lookup(key, &cached)) return cached; +TritonConfig 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, int device_id) { + TritonConfig cached; + if (AutotuneCacheLookup(key, &cached)) return cached; namespace py = pybind11; - if (!compiler_init()) return configs.empty() ? config_t{} : configs[0]; + if (!CompilerInit()) return configs.empty() ? TritonConfig{} : configs[0]; py::gil_scoped_acquire gil; try { py::module_ mod = py::module_::import("infini.triton.jit.compile"); - device_info_t dev = current_device(); + DeviceInfo dev = CurrentDevice(); py::list cands; for (const auto& c : configs) { @@ -93,8 +94,8 @@ config_t autotune_bench(const char* op_name, 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) + "/" + - cache_file_key(op_name, full_sig.c_str(), c.num_warps, - c.num_stages, dev.arch); + CacheFileKey(op_name, full_sig.c_str(), c.num_warps, + c.num_stages, dev.arch); cands.append(cd); } @@ -139,12 +140,12 @@ config_t autotune_bench(const char* op_name, .cast(); if (best_idx < 0 || best_idx >= static_cast(configs.size())) best_idx = 0; - config_t winner = configs[best_idx]; - autotune_cache_insert(key, winner); + TritonConfig 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() ? config_t{} : configs[0]; + return configs.empty() ? TritonConfig{} : configs[0]; } } diff --git a/src/triton/jit/jit.cc b/src/triton/jit/jit.cc index abd1adb0b..64d9d9d49 100644 --- a/src/triton/jit/jit.cc +++ b/src/triton/jit/jit.cc @@ -9,8 +9,8 @@ namespace infini::ops { -device_info_t current_device() { - device_info_t info; +DeviceInfo CurrentDevice() { + DeviceInfo info; CUdevice dev; if (cuCtxGetDevice(&dev) != CUDA_SUCCESS) return info; info.id = static_cast(dev); @@ -23,17 +23,17 @@ device_info_t current_device() { return info; } -static CUresult load_cubin(const char* cubin_path, const char* meta_path, - CUfunction* out_func, unsigned* out_shared, - CUmodule* out_mod) { - if (!file_exists(meta_path)) return CUDA_ERROR_FILE_NOT_FOUND; - std::string meta_json = read_file(meta_path); - int shared = json_get_int(meta_json, "shared", 0); +static CUresult LoadCubin(const char* cubin_path, const char* meta_path, + CUfunction* out_func, unsigned* out_shared, + CUmodule* out_mod) { + if (!FileExists(meta_path)) return CUDA_ERROR_FILE_NOT_FOUND; + std::string meta_json = ReadFile(meta_path); + int shared = JsonGetInt(meta_json, "shared", 0); if (shared < 0) shared = 0; - std::string fn_name = json_get_string(meta_json, "name", "kernel"); + std::string fn_name = JsonGetString(meta_json, "name", "kernel"); - int global_scratch = json_get_int(meta_json, "global_scratch_size", 0); - int profile_scratch = json_get_int(meta_json, "profile_scratch_size", 0); + int global_scratch = JsonGetInt(meta_json, "global_scratch_size", 0); + int profile_scratch = JsonGetInt(meta_json, "profile_scratch_size", 0); if (global_scratch > 0 || profile_scratch > 0) { fprintf(stderr, "triton jit: scratch not supported yet\n"); return CUDA_ERROR_NOT_SUPPORTED; @@ -80,12 +80,12 @@ static CUresult load_cubin(const char* cubin_path, const char* meta_path, return CUDA_SUCCESS; } -void* get_kernel(const char* op_name, const char* signature_str, void* stream, - const config_t& opts, unsigned* out_shared) { - device_info_t dev = current_device(); +void* GetKernel(const char* op_name, const char* signature_str, void* stream, + const TritonConfig& opts, unsigned* out_shared) { + DeviceInfo dev = CurrentDevice(); - auto r = cache_query(op_name, signature_str, opts.num_warps, opts.num_stages, - dev.arch, dev.id); + auto r = CacheQuery(op_name, signature_str, opts.num_warps, opts.num_stages, + dev.arch, dev.id); if (r.mem_hit) { *out_shared = r.shared; return r.func; @@ -94,9 +94,9 @@ void* get_kernel(const char* op_name, const char* signature_str, void* stream, std::string cubin_path = r.out_prefix + ".cubin"; std::string meta_path = r.out_prefix + ".json"; - if (!cache_complete(cubin_path, meta_path)) { - int ret = compile_kernel(op_name, r.out_prefix.c_str(), opts.num_warps, - opts.num_stages, dev.id, signature_str); + if (!CacheComplete(cubin_path, meta_path)) { + int ret = CompileKernel(op_name, r.out_prefix.c_str(), opts.num_warps, + opts.num_stages, dev.id, signature_str); if (ret != 0) return nullptr; } @@ -104,28 +104,28 @@ void* get_kernel(const char* op_name, const char* signature_str, void* stream, unsigned shared; CUmodule mod; CUresult err = - load_cubin(cubin_path.c_str(), meta_path.c_str(), &func, &shared, &mod); + LoadCubin(cubin_path.c_str(), meta_path.c_str(), &func, &shared, &mod); if (err != CUDA_SUCCESS) return nullptr; - kernel_cache_entry_t mine{static_cast(func), shared}; - kernel_cache_entry_t winner; - if (kernel_cache_lookup(r.mem_key, &winner)) { + KernelCacheEntry mine{static_cast(func), shared}; + KernelCacheEntry winner; + if (KernelCacheLookup(r.mem_key, &winner)) { cuModuleUnload(mod); func = static_cast(winner.func); shared = winner.shared; } else { - kernel_cache_insert(r.mem_key, mine); + KernelCacheInsert(r.mem_key, mine); } *out_shared = shared; return static_cast(func); } -int launch_kernel(const char* op_name, const char* signature_str, void* stream, - grid_t grid, config_t opts, void** args) { +int LaunchKernel(const char* op_name, const char* signature_str, void* stream, + Grid grid, TritonConfig opts, void** args) { CUstream cu_stream = static_cast(stream); unsigned shared = 0; - void* func_ptr = get_kernel(op_name, signature_str, stream, opts, &shared); + void* func_ptr = GetKernel(op_name, signature_str, stream, opts, &shared); if (func_ptr == nullptr) return static_cast(CUDA_ERROR_UNKNOWN); return static_cast(cuLaunchKernel( diff --git a/src/triton/jit/jit.h b/src/triton/jit/jit.h index dd4445ac2..596a0e364 100644 --- a/src/triton/jit/jit.h +++ b/src/triton/jit/jit.h @@ -15,33 +15,39 @@ namespace infini::ops { -struct config_t : Config { - config_t() = default; - config_t(unsigned num_warps, unsigned num_stages, - std::vector> constexprs) +struct TritonConfig : Config { + TritonConfig() = default; + + TritonConfig(unsigned num_warps, unsigned num_stages, + std::vector> constexprs) : num_warps(num_warps), num_stages(num_stages), constexprs(std::move(constexprs)) {} unsigned num_warps = 4; + unsigned num_stages = 3; + std::vector> constexprs; bool autotune = false; - std::vector configs; + + std::vector configs; + int warmup = 5; + int rep = 50; - bool is_autotune() const { return autotune; } + bool IsAutotune() const { return autotune; } - int at(const std::string& key) const { + int At(const std::string& key) const { for (const auto& [k, v] : constexprs) if (k == key) return v; - assert(false && "constexpr not found"); + assert(false && "`constexpr` not found"); return 0; } - void apply_defaults(const config_t& defaults) { + void ApplyDefaults(const TritonConfig& defaults) { for (const auto& [dk, dv] : defaults.constexprs) { bool found = false; for (const auto& [k, v] : constexprs) @@ -54,49 +60,54 @@ struct config_t : Config { } }; -struct grid_t { - unsigned x = 1, y = 1, z = 1; +struct Grid { + unsigned x = 1; + + unsigned y = 1; + + unsigned z = 1; }; -struct device_info_t { +struct DeviceInfo { int id = 0; + int arch = 0; }; -bool compiler_init(); +bool CompilerInit(); -int compile_kernel(const char* op_name, const char* out_prefix, int num_warps, +int CompileKernel(const char* op_name, const char* out_prefix, int num_warps, + int num_stages, int device_id, const char* signature); - int num_stages, int device_id, const char* signature); +int LaunchKernel(const char* op_name, const char* signature_str, void* stream, + Grid grid, TritonConfig config, void** args); -int launch_kernel(const char* op_name, const char* signature_str, void* stream, - grid_t grid, config_t config, void** args); +void* GetKernel(const char* op_name, const char* signature_str, void* stream, + const TritonConfig& config, unsigned* out_shared); -void* get_kernel(const char* op_name, const char* signature_str, void* stream, - const config_t& config, unsigned* out_shared); +DeviceInfo CurrentDevice(); -device_info_t current_device(); - -config_t autotune_bench(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, int device_id); +TritonConfig 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, int device_id); // ---- specialization ---- -inline const char* spec_ptr(uintptr_t v) { return v % 16 == 0 ? ":16" : ""; } +inline const char* SpecPtr(uintptr_t v) { return v % 16 == 0 ? ":16" : ""; } template -const char* spec_int(T v) { +const char* SpecInt(T v) { if (v == 1) return ":1"; if ((v & 15) == 0) return ":16"; return ""; } -// ---- DataType → Triton string ---- +// ---- `DataType` → Triton string ---- -inline const char* dtype_to_ttype(DataType dt) { +inline const char* DataTypeToTritonType(DataType dt) { switch (dt) { case DataType::kFloat16: return "fp16"; @@ -129,7 +140,7 @@ inline const char* dtype_to_ttype(DataType dt) { // ---- C++ scalar type → Triton string ---- template -const char* cstype_to_ttype() { +const char* ScalarTypeToTritonType() { if constexpr (std::is_same_v) return "fp64"; else if constexpr (std::is_same_v) @@ -147,15 +158,17 @@ const char* cstype_to_ttype() { // ---- arguments parser ---- -struct arg_pack_t { +struct ArgPack { std::vector ptrs; + std::deque storage; + std::string sig; template - void* store(T v) { + void* Store(T v) { static_assert(sizeof(T) <= sizeof(uint64_t), - "scalar arg wider than 8 bytes"); + "scalar arg wider than `uint64_t`"); uint64_t slot = 0; std::memcpy(&slot, &v, sizeof(T)); storage.push_back(slot); @@ -163,83 +176,83 @@ struct arg_pack_t { } }; -inline void _push_arg(const Tensor& t, arg_pack_t& pack) { +inline void PushArg(const Tensor& t, ArgPack& pack) { auto ptr = reinterpret_cast(t.data()); pack.sig += - std::string("*") + dtype_to_ttype(t.dtype()) + spec_ptr(ptr) + ","; - pack.ptrs.push_back(pack.store(ptr)); + std::string("*") + DataTypeToTritonType(t.dtype()) + SpecPtr(ptr) + ","; + pack.ptrs.push_back(pack.Store(ptr)); } template , int> = 0> -void _push_arg(T v, arg_pack_t& pack) { - const char* s = spec_int(v); - pack.sig += std::string(cstype_to_ttype()) + s + ","; - if (std::strcmp(s, ":1") != 0) pack.ptrs.push_back(pack.store(v)); +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)); } -inline void _push_arg(float v, arg_pack_t& pack) { - pack.ptrs.push_back(pack.store(v)); +inline void PushArg(float v, ArgPack& pack) { + pack.ptrs.push_back(pack.Store(v)); pack.sig += "fp32,"; } -inline void _push_arg(double v, arg_pack_t& pack) { - pack.ptrs.push_back(pack.store(v)); +inline void PushArg(double v, ArgPack& pack) { + pack.ptrs.push_back(pack.Store(v)); pack.sig += "fp64,"; } // ---- launch wrapper ---- template -int launch_jit(const char* op, void* stream, grid_t grid, config_t config, - Args&&... args) { - arg_pack_t pack; +int LaunchJit(const char* op, void* stream, Grid grid, TritonConfig config, + Args&&... args) { + ArgPack pack; pack.sig.reserve(256); - (_push_arg(std::forward(args), pack), ...); + (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(); - // triton need - void* scratch = pack.store(0); + // Triton needs two trailing scratch pointers. + void* scratch = pack.Store(0); pack.ptrs.push_back(scratch); pack.ptrs.push_back(scratch); - return launch_kernel(op, pack.sig.c_str(), stream, grid, config, - pack.ptrs.data()); + return LaunchKernel(op, pack.sig.c_str(), stream, grid, config, + pack.ptrs.data()); } template -int launch_jit_autotune(const char* op, void* stream, const config_t& config, - const std::vector& key_dims, - DataType dtype, GridFn grid_fn, Args&&... args) { +int LaunchJitAutotune(const char* op, void* stream, const TritonConfig& config, + const std::vector& key_dims, DataType dtype, + GridFn grid_fn, Args&&... args) { std::string cache_key = op; for (auto d : key_dims) cache_key += "|" + std::to_string(d); cache_key += "|dt=" + std::to_string(static_cast(dtype)); - arg_pack_t pack; + ArgPack pack; pack.sig.reserve(256); - (_push_arg(std::forward(args), pack), ...); + (PushArg(std::forward(args), pack), ...); - std::vector grids; + std::vector grids; grids.reserve(config.configs.size()); for (const auto& c : config.configs) grids.push_back(grid_fn(c)); - config_t best = autotune_bench(op, config.configs, pack.sig, pack.ptrs, grids, - config.warmup, config.rep, cache_key.c_str(), - current_device().id); + TritonConfig best = AutotuneBench(op, config.configs, pack.sig, pack.ptrs, + grids, config.warmup, config.rep, + cache_key.c_str(), CurrentDevice().id); - grid_t grid = grid_fn(best); + Grid grid = grid_fn(best); for (const auto& [name, val] : best.constexprs) pack.sig += name + "=" + std::to_string(val) + ","; if (!pack.sig.empty()) pack.sig.pop_back(); - void* scratch = pack.store(0); + void* scratch = pack.Store(0); pack.ptrs.push_back(scratch); pack.ptrs.push_back(scratch); - return launch_kernel(op, pack.sig.c_str(), stream, grid, best, - pack.ptrs.data()); + return LaunchKernel(op, pack.sig.c_str(), stream, grid, best, + pack.ptrs.data()); } } // namespace infini::ops diff --git a/src/triton/ops/add/add.py b/src/triton/ops/add/add.py index eb3e35724..f813cda47 100644 --- a/src/triton/ops/add/add.py +++ b/src/triton/ops/add/add.py @@ -16,6 +16,7 @@ def kernel( out_contig, ndim, n_elements, + alpha, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -25,7 +26,7 @@ def kernel( 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, 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) @@ -49,4 +50,4 @@ def kernel( 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, mask=mask) + tl.store(out_ptr + out_offs, x + y * alpha, mask=mask) diff --git a/src/triton/ops/add/jit.h b/src/triton/ops/add/jit.h index 42bcf6502..7d7861fcc 100644 --- a/src/triton/ops/add/jit.h +++ b/src/triton/ops/add/jit.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_TRITON_JIT_ADD_H_ -#define INFINI_OPS_TRITON_JIT_ADD_H_ +#ifndef INFINI_OPS_TRITON_OPS_ADD_JIT_H_ +#define INFINI_OPS_TRITON_OPS_ADD_JIT_H_ #include @@ -14,14 +14,16 @@ namespace infini::ops { template <> -class Operator : public Add { +class Operator : public Add { public: using Add::Add; using Add::operator(); - static config_t default_config() { return {4u, 3u, {{"BLOCK_SIZE", 1024}}}; } + static TritonConfig DefaultConfig() { + return {4u, 3u, {{"BLOCK_SIZE", 1024}}}; + } - static std::vector autotune_configs() { + static std::vector AutotuneConfigs() { return { {4u, 3u, {{"BLOCK_SIZE", 256}}}, {4u, 3u, {{"BLOCK_SIZE", 512}}}, @@ -30,7 +32,7 @@ class Operator : public Add { }; } - void operator()(const Tensor input, const Tensor other, + void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { const int ndim = static_cast(ndim_); @@ -61,33 +63,33 @@ class Operator : public Add { const size_t n_elements = out.numel(); auto extension = config_.extension(); - static const config_t defaults = default_config(); - const auto* config_ptr = static_cast(extension.get()); - config_t config = config_ptr ? *config_ptr : defaults; - if (extension) config.apply_defaults(defaults); + static const TritonConfig defaults = DefaultConfig(); + const auto* config_ptr = static_cast(extension.get()); + TritonConfig config = config_ptr ? *config_ptr : defaults; + if (extension) config.ApplyDefaults(defaults); int result; - if (config.is_autotune()) { - if (config.configs.empty()) config.configs = autotune_configs(); - for (auto& c : config.configs) c.apply_defaults(defaults); - result = launch_jit_autotune( + if (config.IsAutotune()) { + if (config.configs.empty()) config.configs = AutotuneConfigs(); + for (auto& c : config.configs) c.ApplyDefaults(defaults); + result = LaunchJitAutotune( "add", stream_, config, {n_elements}, out.dtype(), - [&](const config_t& c) { - int block_size = c.at("BLOCK_SIZE"); - return grid_t{static_cast((n_elements + block_size - 1) / - block_size)}; + [&](const TritonConfig& c) { + int block_size = c.At("BLOCK_SIZE"); + return Grid{static_cast((n_elements + block_size - 1) / + block_size)}; }, input, other, out, d_out_shape, d_input_strides, d_other_strides, d_out_strides, is_input_contiguous_, is_other_contiguous_, - is_out_contiguous_, ndim, n_elements); + is_out_contiguous_, ndim, n_elements, alpha); } else { - const int block_size = config.at("BLOCK_SIZE"); - grid_t grid{ + const int block_size = config.At("BLOCK_SIZE"); + Grid grid{ static_cast((n_elements + block_size - 1) / block_size)}; - result = launch_jit( + result = LaunchJit( "add", stream_, grid, config, input, other, out, d_out_shape, d_input_strides, d_other_strides, d_out_strides, is_input_contiguous_, - is_other_contiguous_, is_out_contiguous_, ndim, n_elements); + is_other_contiguous_, is_out_contiguous_, ndim, n_elements, alpha); } cuMemFreeAsync(d_meta, static_cast(stream_));