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..33ac2e98c 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,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 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") + candidate_config.constexprs.emplace_back(key, item.second.cast()); + } + config->configs.push_back(std::move(candidate_config)); + } + } + } else { + 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) @@ -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,12 +841,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(ConfigFromPyDict(*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 +865,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 +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( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1edeebebd..6f5f8707c 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,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 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..871684e53 --- /dev/null +++ b/src/triton/jit/cache.h @@ -0,0 +1,231 @@ +#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 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; +} + +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 ---- + +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); +} + +// ---- kernel cache ---- + +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))); +} + +struct KernelCacheEntry { + void* func; + + unsigned shared; +}; + +struct KernelCache { + std::mutex mutex; + + std::unordered_map map; +}; + +inline KernelCache& GetKernelCache() { + static KernelCache c; + return c; +} + +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; + *out = it->second; + return true; +} + +inline void KernelCacheInsert(const std::string& key, KernelCacheEntry entry) { + auto& c = GetKernelCache(); + std::lock_guard lk(c.mutex); + c.map[key] = entry; +} + +struct CacheQueryResult { + bool mem_hit; + + void* func; + + unsigned shared; + + std::string out_prefix; + + std::string mem_key; +}; + +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 = 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 AutotuneCache { + std::mutex mutex; + + std::unordered_map map; +}; + +inline AutotuneCache& GetAutotuneCache() { + static AutotuneCache c; + return c; +} + +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 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) + s += "\n" + name + " " + std::to_string(val); + return s; +} + +inline bool DeserializeConfig(const std::string& content, TritonConfig* 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, 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 = AutotuneCacheFilePath(key); + if (FileExists(path.c_str())) { + TritonConfig parsed; + if (DeserializeConfig(ReadFile(path.c_str()), &parsed)) { + c.map[key] = parsed; + *out = parsed; + return true; + } + } + return false; +} + +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 = 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); + } +} + +} // 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..a58a3ffcf --- /dev/null +++ b/src/triton/jit/compiler.cc @@ -0,0 +1,152 @@ +#include + +#include +#include + +#include "cache.h" +#include "jit.h" + +namespace infini::ops { + +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 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; + 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; + } +} + +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 (!CompilerInit()) return configs.empty() ? TritonConfig{} : configs[0]; + py::gil_scoped_acquire gil; + try { + py::module_ mod = py::module_::import("infini.triton.jit.compile"); + + DeviceInfo dev = CurrentDevice(); + + 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, 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; + 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() ? TritonConfig{} : 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..64d9d9d49 --- /dev/null +++ b/src/triton/jit/jit.cc @@ -0,0 +1,136 @@ +#include "jit.h" + +#include + +#include +#include + +#include "cache.h" + +namespace infini::ops { + +DeviceInfo CurrentDevice() { + DeviceInfo 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 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 = JsonGetString(meta_json, "name", "kernel"); + + 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; + } + + 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* GetKernel(const char* op_name, const char* signature_str, void* stream, + const TritonConfig& opts, unsigned* out_shared) { + DeviceInfo dev = CurrentDevice(); + + 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; + } + + std::string cubin_path = r.out_prefix + ".cubin"; + std::string meta_path = r.out_prefix + ".json"; + + 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; + } + + CUfunction func; + unsigned shared; + CUmodule mod; + CUresult err = + LoadCubin(cubin_path.c_str(), meta_path.c_str(), &func, &shared, &mod); + if (err != CUDA_SUCCESS) return nullptr; + + 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 { + KernelCacheInsert(r.mem_key, mine); + } + + *out_shared = shared; + return static_cast(func); +} + +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 = GetKernel(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..596a0e364 --- /dev/null +++ b/src/triton/jit/jit.h @@ -0,0 +1,260 @@ +#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 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; + + int warmup = 5; + + int rep = 50; + + bool IsAutotune() 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 ApplyDefaults(const TritonConfig& 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 { + unsigned x = 1; + + unsigned y = 1; + + unsigned z = 1; +}; + +struct DeviceInfo { + int id = 0; + + int arch = 0; +}; + +bool CompilerInit(); + +int CompileKernel(const char* op_name, const char* out_prefix, int num_warps, + 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); + +void* GetKernel(const char* op_name, const char* signature_str, void* stream, + const TritonConfig& config, unsigned* out_shared); + +DeviceInfo CurrentDevice(); + +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* SpecPtr(uintptr_t v) { return v % 16 == 0 ? ":16" : ""; } + +template +const char* SpecInt(T v) { + if (v == 1) return ":1"; + if ((v & 15) == 0) return ":16"; + return ""; +} + +// ---- `DataType` → Triton string ---- + +inline const char* DataTypeToTritonType(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* 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" : "i32"; + } + return "i32"; +} + +// ---- arguments parser ---- + +struct ArgPack { + 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 `uint64_t`"); + uint64_t slot = 0; + std::memcpy(&slot, &v, sizeof(T)); + storage.push_back(slot); + return &storage.back(); + } +}; + +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 , 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)); +} + +inline void PushArg(float v, ArgPack& pack) { + pack.ptrs.push_back(pack.Store(v)); + pack.sig += "fp32,"; +} + +inline void PushArg(double v, ArgPack& pack) { + pack.ptrs.push_back(pack.Store(v)); + pack.sig += "fp64,"; +} + +// ---- launch wrapper ---- + +template +int LaunchJit(const char* op, void* stream, Grid grid, TritonConfig 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(); + + // Triton needs two trailing scratch pointers. + void* scratch = pack.Store(0); + pack.ptrs.push_back(scratch); + pack.ptrs.push_back(scratch); + + return LaunchKernel(op, pack.sig.c_str(), stream, grid, config, + pack.ptrs.data()); +} + +template +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)); + + ArgPack pack; + pack.sig.reserve(256); + (PushArg(std::forward(args), pack), ...); + + std::vector grids; + grids.reserve(config.configs.size()); + for (const auto& c : config.configs) grids.push_back(grid_fn(c)); + + TritonConfig best = AutotuneBench(op, config.configs, pack.sig, pack.ptrs, + grids, config.warmup, config.rep, + cache_key.c_str(), CurrentDevice().id); + + 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); + pack.ptrs.push_back(scratch); + pack.ptrs.push_back(scratch); + + return LaunchKernel(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..f813cda47 --- /dev/null +++ b/src/triton/ops/add/add.py @@ -0,0 +1,53 @@ +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.h b/src/triton/ops/add/jit.h new file mode 100644 index 000000000..7d7861fcc --- /dev/null +++ b/src/triton/ops/add/jit.h @@ -0,0 +1,103 @@ +#ifndef INFINI_OPS_TRITON_OPS_ADD_JIT_H_ +#define INFINI_OPS_TRITON_OPS_ADD_JIT_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 TritonConfig DefaultConfig() { + return {4u, 3u, {{"BLOCK_SIZE", 1024}}}; + } + + 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}}}, + }; + } + + void operator()(const Tensor input, const Tensor other, const double alpha, + 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 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.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 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, alpha); + } else { + const int block_size = config.At("BLOCK_SIZE"); + Grid grid{ + static_cast((n_elements + block_size - 1) / block_size)}; + 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, alpha); + } + + cuMemFreeAsync(d_meta, static_cast(stream_)); + + assert(result == 0 && "Triton JIT `Add` launch failed"); + } +}; + +} // namespace infini::ops + +#endif