diff --git a/CMakeLists.txt b/CMakeLists.txt index bc6eda3ea..290befede 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 @@ -323,6 +325,14 @@ 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 Triton AOT temporarily targets CUDA.") +endif() + +if(WITH_TRITON) + set(TRITON_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run Triton AOT code generation") +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 f4d0ea1cd..4b7a4f752 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -415,6 +415,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 @@ -512,6 +514,54 @@ def _find_params_with_defaults(op_name): return mapping +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) optional_non_tensor_params = _find_optional_non_tensor_params(operator.name) @@ -642,22 +692,41 @@ 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) if not method: + 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 "" default_impl_index = _default_impl_index_expr(call) + 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" Handle handle;\n" @@ -667,8 +736,10 @@ def _generate_call(op_name, call, method=True): f" Config config;\n" f" config.set_implementation_index(\n" f" implementation_index.value_or({default_impl_index}));\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, ' + f'py::arg("implementation_index") = py::none(){extra_pybind});' ) # The first lambda parameter is conventionally named `self`, but @@ -715,9 +786,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_ @@ -729,7 +810,7 @@ def _overload_order_key(node): #include "config.h" #include "generated/bindings/generated_dispatch.h" #include "handle.h" -#include "pybind11_utils.h" +#include "pybind11_utils.h"{jit_include} namespace py = pybind11; @@ -1056,9 +1137,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: @@ -1477,13 +1561,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 = {} @@ -1532,6 +1618,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( @@ -1622,6 +1709,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): @@ -1636,6 +1729,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/scripts/triton/aot.py b/scripts/triton/aot.py new file mode 100644 index 000000000..90729a9d5 --- /dev/null +++ b/scripts/triton/aot.py @@ -0,0 +1,253 @@ +import ast +from dataclasses import dataclass +import pathlib +from typing import Any, Sequence + +from triton.tools import link +from triton.tools.compile import CompileArgs, compile_kernel + + +@dataclass(frozen=True) +class Signature: + pointer_dtypes: dict[str, str] + pointer_alignments: dict[str, int | None] | None = None + scalar_dtypes: dict[str, str] | None = None + constexprs: dict[str, Any] | None = None + + +@dataclass(frozen=True) +class CompileConfig: + signature: Signature + grid: str + out_name: str + num_warps: int = 4 + num_stages: int = 3 + target: Any = None + + +def compile( + config: CompileConfig, + *, + path: pathlib.Path, + kernel_name: str, + out_dir: pathlib.Path, + kernel_args: Sequence[str], +) -> list[pathlib.Path]: + _, files = compile_kernel( + CompileArgs( + path=str(path), + kernel_name=kernel_name, + signature=_render_signature(config.signature, kernel_args), + grid=config.grid, + num_warps=config.num_warps, + num_stages=config.num_stages, + out_name=config.out_name, + out_path=out_dir / config.out_name, + target=config.target, + ) + ) + + return [path for path in files if path.suffix == ".h"] + + +def link_headers(headers: Sequence[pathlib.Path], out_base: pathlib.Path): + parser = link.HeaderParser() + for header in headers: + parser.extract_linker_meta(header.read_text()) + + first_meta = next(iter(parser.kernels.values()))[0] + backend_prelude = ( + pathlib.Path(link.__file__).parent / "extra" / parser.backend_name / "link.h" + ).read_text() + + out_base.with_suffix(".h").write_text( + backend_prelude + + "\n".join( + link.make_algo_decls(name, meta) for name, meta in parser.kernels.items() + ) + + "\n" + + link.make_get_num_algos_decl(first_meta) + + "\n" + + link.make_global_decl(first_meta) + ) + + names = list(parser.kernels) + defs = [ + link.make_kernel_hints_dispatcher(name, meta) + for name, meta in parser.kernels.items() + ] + + out_base.with_suffix(".c").write_text( + backend_prelude + + "#include \n#include \n\n" + + "\n".join(defs) + + "\n" + + link.make_func_pointers(names, first_meta) + + "\n" + + link.make_get_num_algos_def(first_meta) + + "\n" + + link.make_kernel_meta_const_dispatcher(first_meta) + + "\n" + + link.make_kernel_load_def(names, first_meta) + + "\n" + + link.make_default_algo_kernel(first_meta) + ) + + +def build( + configs: Sequence[CompileConfig], + *, + path: pathlib.Path, + kernel_name: str, + out_dir: pathlib.Path, + kernel_args: Sequence[str], +) -> pathlib.Path: + if not configs: + raise ValueError("empty compile configs") + + out_name = configs[0].out_name + out_dir.mkdir(parents=True, exist_ok=True) + + headers = [] + for config in configs: + headers.extend( + compile( + config, + path=path, + kernel_name=kernel_name, + out_dir=out_dir, + kernel_args=kernel_args, + ) + ) + + if not headers: + raise ValueError(f"no headers generated for {out_name}") + + out_base = out_dir / out_name + link_headers(headers, out_base) + + return out_base + + +def write_header( + headers: Sequence[pathlib.Path], + out_path: pathlib.Path, + *, + op_name: str, + configs: Sequence[CompileConfig], + kernel_args: Sequence[str], +): + guard = f"INFINI_OPS_GENERATED_{out_path.stem.upper()}_H_" + includes = "\n".join(f'#include "{header.name}"' for header in headers) + params = _dispatch_params(configs[0].signature, kernel_args) + param_decls = ", ".join(f"{ty} {name}" for ty, name in params) + param_names = ", ".join(name for _, name in params) + + body = f"#ifndef {guard}\n#define {guard}\n\n" + body += f'extern "C" {{\n{includes}\n}}\n\n' + body += '#include \n\n#include "data_type.h"\n\n' + body += "namespace infini::ops {\n\n" + + body += f"inline TT_ResultTy launch_infini_ops_triton_{op_name}(\n" + body += f" DataType dtype, TT_StreamTy stream, {param_decls}) {{\n" + body += " switch (dtype) {\n" + for config in configs: + dtype = _out_dtype(config.out_name) + body += f" case DataType::{_data_type(dtype)}:\n" + body += f" return {config.out_name}_default(stream, {param_names});\n" + body += " default:\n return TT_ERROR_INVALID_VALUE;\n }\n}\n\n" + + body += f"inline void load_infini_ops_triton_{op_name}(DataType dtype) {{\n" + body += " switch (dtype) {\n" + for config in configs: + dtype = _out_dtype(config.out_name) + body += f" case DataType::{_data_type(dtype)}: {{\n" + body += " static std::once_flag once;\n" + body += f" std::call_once(once, &load_{config.out_name});\n" + body += " return;\n }\n" + body += " default:\n return;\n }\n}\n\n" + + body += "} // namespace infini::ops\n\n#endif\n" + out_path.write_text(body) + + +def kernel_args(path, kernel_name): + tree = ast.parse(pathlib.Path(path).read_text()) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == kernel_name: + return tuple(arg.arg for arg in node.args.args) + raise ValueError(f"kernel {kernel_name} not found in {path}") + + +def _render_signature(signature: Signature, args: Sequence[str]) -> str: + pointer_alignments = signature.pointer_alignments or {} + scalar_dtypes = signature.scalar_dtypes or {} + constexprs = signature.constexprs or {} + + parts = [] + for arg in args: + if arg in constexprs: + parts.append(str(constexprs[arg])) + elif arg in signature.pointer_dtypes: + parts.append( + _ptr(signature.pointer_dtypes[arg], pointer_alignments.get(arg)) + ) + elif arg in scalar_dtypes: + parts.append(str(scalar_dtypes[arg])) + else: + raise ValueError(f"missing signature rule for {arg}") + + return ", ".join(parts) + + +def _dispatch_params(signature: Signature, args: Sequence[str]): + scalar_dtypes = signature.scalar_dtypes or {} + constexprs = signature.constexprs or {} + + params = [] + for arg in args: + if arg in constexprs: + continue + if arg in signature.pointer_dtypes: + params.append(("CUdeviceptr", arg)) + elif arg in scalar_dtypes: + params.append((_scalar_ctype(scalar_dtypes[arg]), arg)) + else: + raise ValueError(f"missing dispatch rule for {arg}") + return params + + +def _scalar_ctype(dtype): + return { + "i32": "int32_t", + "i64": "int64_t", + "u32": "uint32_t", + "u64": "uint64_t", + "fp32": "float", + "fp64": "double", + }[dtype] + + +def _out_dtype(out_name): + return out_name.rsplit("_", 1)[1] + + +def _data_type(dtype): + return { + "fp16": "kFloat16", + "bf16": "kBFloat16", + "fp32": "kFloat32", + "fp64": "kFloat64", + "i8": "kInt8", + "i16": "kInt16", + "i32": "kInt32", + "i64": "kInt64", + "u8": "kUInt8", + "u16": "kUInt16", + "u32": "kUInt32", + "u64": "kUInt64", + }[dtype] + + +def _ptr(dtype, alignment=None): + return f"*{dtype}" if alignment is None else f"*{dtype}:{alignment}" diff --git a/scripts/triton/generate_ops.py b/scripts/triton/generate_ops.py new file mode 100644 index 000000000..e4ab763c1 --- /dev/null +++ b/scripts/triton/generate_ops.py @@ -0,0 +1,118 @@ +import argparse +import importlib.util +import pathlib +import shutil +import sys + +import aot + +_PROJECT_DIR = pathlib.Path(__file__).resolve().parents[2] +if str(_PROJECT_DIR) not in sys.path: + sys.path.insert(0, str(_PROJECT_DIR)) + +_KERNEL_NAME = "kernel" +_OPS_DIR = _PROJECT_DIR / "src" / "triton" / "ops" + + +def _prepend_sys_path(path): + path = str(path) + if path not in sys.path: + sys.path.insert(0, path) + + +def _find_op_modules(): + return { + path.parent.name: path + for path in sorted(_OPS_DIR.glob("*/build.py")) + if path.is_file() + } + + +def _build_manifest(output_dir): + return sorted(str(path) for path in pathlib.Path(output_dir).rglob("*.c")) + + +def _write_cmake_manifest(output_dir, sources): + manifest_path = pathlib.Path(output_dir) / "manifest.cmake" + lines = ["set(INFINIOPS_TRITON_SOURCES"] + lines.extend(f' "{source}"' for source in sources) + lines.append(")") + lines.append("") + lines.append(f'set(INFINIOPS_TRITON_INCLUDE_DIRS "{output_dir}")') + lines.append("") + manifest_path.write_text("\n".join(lines) + "\n") + + +def _load_op_module(path): + _prepend_sys_path(path.parent) + spec = importlib.util.spec_from_file_location( + f"infiniops_triton_{path.parent.name}_build", + path, + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + + return module + + +def generate(ops, *, output_dir): + op_modules = _find_op_modules() + unknown_ops = tuple(op for op in ops if op not in op_modules) + + if unknown_ops: + raise ValueError(f"unsupported Triton ops: {', '.join(unknown_ops)}") + + output_dir = pathlib.Path(output_dir) + shutil.rmtree(output_dir, ignore_errors=True) + output_dir.mkdir(parents=True, exist_ok=True) + + for op in ops: + path = op_modules[op] + kernel_path = path.parent / f"{op}.py" + module = _load_op_module(path) + kernel_args = aot.kernel_args(kernel_path, _KERNEL_NAME) + headers = [] + dispatch_configs = [] + for configs in module.configs(): + out_base = aot.build( + configs, + path=kernel_path, + kernel_name=_KERNEL_NAME, + out_dir=output_dir / op, + kernel_args=kernel_args, + ) + headers.append(out_base.with_suffix(".h")) + dispatch_configs.append(configs[0]) + aot.write_header( + headers, + output_dir / op / f"infini_ops_triton_{op}.h", + op_name=op, + configs=dispatch_configs, + kernel_args=kernel_args, + ) + + sources = _build_manifest(output_dir) + _write_cmake_manifest(output_dir, sources) + + return sources + + +def _parse_args(): + parser = argparse.ArgumentParser( + description="Generate Triton operator sources for InfiniOps." + ) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--ops", nargs="+", default=tuple(_find_op_modules())) + + return parser.parse_args() + + +def main(): + args = _parse_args() + generate(args.ops, output_dir=args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..d32d59489 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -126,6 +126,44 @@ 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) + + if(TRITON_PYTHON_EXECUTABLE) + set(_triton_python "${TRITON_PYTHON_EXECUTABLE}") + elseif(_TORCH_PYTHON) + set(_triton_python "${_TORCH_PYTHON}") + else() + set(_triton_python "${Python_EXECUTABLE}") + endif() + message(STATUS "Triton codegen Python: ${_triton_python}") + + set(_triton_output_dir "${CMAKE_CURRENT_BINARY_DIR}/triton") + set(_triton_generator_args + "${PROJECT_SOURCE_DIR}/scripts/triton/generate_ops.py" + --output-dir "${_triton_output_dir}") + + execute_process( + COMMAND "${_triton_python}" ${_triton_generator_args} + WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}" + RESULT_VARIABLE _triton_generation_result + ) + + if(NOT _triton_generation_result EQUAL 0) + message(FATAL_ERROR "Generating Triton AOT operator sources failed with `${_triton_python}`. Set `TRITON_PYTHON_EXECUTABLE` to a Python with `triton` and CUDA dependencies installed.") + endif() + + enable_language(C) + + include("${_triton_output_dir}/manifest.cmake") + target_compile_definitions(infiniops PUBLIC WITH_TRITON=1 + TRITON_JIT_CACHE_DIR="/tmp/triton_jit_cache") + target_include_directories(infiniops PRIVATE ${INFINIOPS_TRITON_INCLUDE_DIRS} ${pybind11_INCLUDE_DIRS}) + target_link_libraries(infiniops PRIVATE pybind11::embed Python::Python) + target_sources(infiniops PRIVATE ${INFINIOPS_TRITON_SOURCES} triton/jit/jit.cc triton/jit/compiler.cc) +endif() + if(WITH_ILUVATAR) set(ILUVATAR_PATTERNS "native/cuda/*.cc" @@ -682,6 +720,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} @@ -1006,6 +1048,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 @@ -1060,6 +1108,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/aot.h b/src/triton/ops/add/aot.h new file mode 100644 index 000000000..7999e8867 --- /dev/null +++ b/src/triton/ops/add/aot.h @@ -0,0 +1,66 @@ +#ifndef INFINI_OPS_TRITON_AOT_ADD_H_ +#define INFINI_OPS_TRITON_AOT_ADD_H_ + +#include + +#include +#include +#include + +#include "add/infini_ops_triton_add.h" +#include "base/add.h" +#include "data_type.h" + +namespace infini::ops { + +template <> +class Operator : public Add { + public: + using Add::Add; + using Add::operator(); + + void operator()(const Tensor input, const Tensor other, + Tensor out) const override { + load_infini_ops_triton_add(out.dtype()); + + 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); + const size_t stride_bytes = ndim * sizeof(int64_t); + + CUdeviceptr d_meta; + cuMemAlloc(&d_meta, meta_bytes); + cuMemcpyHtoD(d_meta, h_meta.data(), meta_bytes); + + CUstream stream = static_cast(stream_); + auto x = reinterpret_cast(const_cast(input.data())); + auto y = reinterpret_cast(const_cast(other.data())); + auto o = reinterpret_cast(out.data()); + + int32_t n = static_cast(out.numel()); + int32_t ndim_val = static_cast(ndim); + int32_t x_contig = static_cast(is_input_contiguous_); + int32_t y_contig = static_cast(is_other_contiguous_); + int32_t out_contig = static_cast(is_out_contiguous_); + + auto result = launch_infini_ops_triton_add( + out.dtype(), stream, x, y, o, d_meta + 0 * stride_bytes, + d_meta + 1 * stride_bytes, d_meta + 2 * stride_bytes, + d_meta + 3 * stride_bytes, x_contig, y_contig, out_contig, ndim_val, n); + + cuMemFree(d_meta); + + assert(result == CUDA_SUCCESS && "Triton `Add` launch failed"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/triton/ops/add/build.py b/src/triton/ops/add/build.py new file mode 100644 index 000000000..70fcf18be --- /dev/null +++ b/src/triton/ops/add/build.py @@ -0,0 +1,52 @@ +from scripts.triton import aot + +_DTYPES = ( + "fp16", + "bf16", + "fp32", + "fp64", + "i8", + "i16", + "i32", + "i64", + "u8", + "u16", + "u32", + "u64", +) +_BLOCK_SIZES = (512, 1024) +_ALIGNMENTS = (16, None) +_NUM_WARPS = 4 +_NUM_STAGES = 3 +_DATA_PTRS = ("x_ptr", "y_ptr", "out_ptr") +_META_PTRS = ("out_shape_ptr", "x_stride_ptr", "y_stride_ptr", "out_stride_ptr") +_SCALARS = ("x_contig", "y_contig", "out_contig", "ndim", "n_elements") + + +def _signature(dtype, block_size, alignment): + return aot.Signature( + pointer_dtypes={ + **{name: dtype for name in _DATA_PTRS}, + **{name: "i64" for name in _META_PTRS}, + }, + pointer_alignments={name: alignment for name in _DATA_PTRS}, + scalar_dtypes={ + **{name: "i32" for name in _SCALARS}, + }, + constexprs={"BLOCK_SIZE": block_size}, + ) + + +def configs(): + for dtype in _DTYPES: + yield tuple( + aot.CompileConfig( + signature=_signature(dtype, block_size, alignment), + grid=f"(n_elements + {block_size} - 1) / {block_size}, 1, 1", + out_name=f"infini_ops_triton_add_{dtype}", + num_warps=_NUM_WARPS, + num_stages=_NUM_STAGES, + ) + for block_size in _BLOCK_SIZES + for alignment in _ALIGNMENTS + ) 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