diff --git a/docs/build.md b/docs/build.md index b4ecb73c7..4c6edbfc0 100644 --- a/docs/build.md +++ b/docs/build.md @@ -27,9 +27,60 @@ entry is `python -m pip install` with CMake options passed through | `INFINI_OPS_BUILD_DOCS` | Enable the Doxygen documentation target. | `OFF` | | `INFINI_RT_ROOT` | InfiniRT install prefix containing `include/` and `lib/`. | `$INFINI_RT_ROOT` | | `INFINI_OPS_SMOKE_BUILD` | Build only the smoke-test operator subset. | `OFF` | -| `INFINI_OPS_OPS` | Comma- or semicolon-separated operator allowlist. | empty | +| `INFINI_OPS_OPS` | Comma- or semicolon-separated operator allowlist, or a path to an `ops.json` implementation selection. | empty | | `INFINI_OPS_TORCH_OPS` | Comma- or semicolon-separated ATen operator allowlist. | empty | +An `ops.json` file selects operators and implementation slots with a top-level +operator mapping: + +```json +{ + "add": { + "implementations": "all" + }, + "argmax": { + "implementations": [8] + }, + "top_k_top_p_sampling_from_logits": { + "implementations": [16] + } +} +``` + +`"all"` keeps every available implementation for the operator. An integer +array keeps exactly those slots. Slots range from 0 through 31. The selection +is a set, not a priority order; the default dispatch selects the smallest +active slot. The selection controls generated wrappers, generated slot-8 ATen +implementations, and linked +provider resolution. Unselected linked providers do not require their external +libraries to be installed. + +Pass the file explicitly with +`-DINFINI_OPS_OPS=/path/to/ops.json`. For compatibility, +`${PROJECT_SOURCE_DIR}/ops.json` is read automatically when present. Relative +implementation header paths in legacy configurations are resolved from +`${PROJECT_SOURCE_DIR}`. An explicit inline `INFINI_OPS_OPS` allowlist takes +precedence over an implicit `${PROJECT_SOURCE_DIR}/ops.json`. When +`INFINI_OPS_TORCH_OPS` and an explicit JSON selection are both set, generated +ATen ops use their intersection. The string and string-array values supported +by the current generator remain available for checked-in implementation +headers. Structured descriptors preserve an explicit backend name, including +for implementations outside the standard backend directory layout. Generated +implementation header paths are not supported: + +```json +{ + "add": "src/native/cpu/ops/add/add.h", + "gemm": ["src/native/cpu/ops/gemm/gemm.h"], + "custom_add": [ + { + "path": "custom/add.h", + "backend": "custom" + } + ] +} +``` + Only one GPU backend should be enabled in a build. CPU may be enabled with the selected accelerator backend. diff --git a/docs/linked-operators.md b/docs/linked-operators.md index c06d71bc8..7e5afee9f 100644 --- a/docs/linked-operators.md +++ b/docs/linked-operators.md @@ -86,6 +86,12 @@ cmake -S . -B build \ -DINFINI_OPS_OPS=silu_and_mul ``` +To resolve only selected linked implementation slots, pass an `ops.json` file +through `INFINI_OPS_OPS`. The resolver reads each linked provider's slot from +its sibling C++ header before locating external libraries, so an unselected +provider does not add a package or shared-library dependency. See +[Build configuration](build.md) for the file format. + The `torch` transport uses the installed PyTorch C++ headers and libraries for `at::Tensor`, but it does not enable the standard `src/torch` operator backend. Provider and PyTorch C++ ABIs must match. Configuration fails before compilation diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 0c90cedad..6a93050c6 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -36,6 +36,9 @@ import yaml _SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent +sys.path.insert(0, str(_SCRIPTS_DIR)) +import ops_config # noqa: E402 + _REPO_ROOT = _SCRIPTS_DIR.parent _OPS_YAML_PATH = _SCRIPTS_DIR / "torch_ops.yaml" _BASE_DIR = _REPO_ROOT / "src" / "base" @@ -1644,6 +1647,30 @@ def _emit(name: str, ops: list[Op], *, emit_base: bool) -> set[pathlib.Path]: return emitted_paths +def _select_op_names(cli_ops, default_ops, config): + if config is None: + return cli_ops or default_ops + + aten_names_by_public_name = collections.defaultdict(list) + for op_name in default_ops: + aten_names_by_public_name[_public_op_name(op_name)].append(op_name) + + selected_public_names = ops_config.torch_op_names( + config, aten_names_by_public_name, _PYTORCH_SLOT + ) + selected = [ + aten_name + for public_name in selected_public_names + for aten_name in aten_names_by_public_name.get(public_name, (public_name,)) + ] + + if cli_ops: + allowed = set(cli_ops) + selected = [op_name for op_name in selected if op_name in allowed] + + return selected + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument( @@ -1651,6 +1678,11 @@ def main() -> int: nargs="*", help="Override the op allowlist. If omitted, reads `scripts/torch_ops.yaml`.", ) + parser.add_argument( + "--ops-config", + type=pathlib.Path, + help="Path to an `ops.json` operator and implementation selection.", + ) parser.add_argument( "--pytorch-version", default=os.environ.get("INFINI_OPS_PYTORCH_VERSION", _DEFAULT_PYTORCH_VERSION), @@ -1666,7 +1698,9 @@ def main() -> int: global _CLANG_FORMAT _CLANG_FORMAT = _find_clang_format() - op_names = args.ops or yaml.safe_load(_OPS_YAML_PATH.read_text()) + default_ops = yaml.safe_load(_OPS_YAML_PATH.read_text()) + config = ops_config.load_ops_config(args.ops_config) if args.ops_config else None + op_names = _select_op_names(args.ops, default_ops, config) aten_entries = _load_aten_entries(args.pytorch_version) skipped: list[tuple[str, str]] = [] diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 74818e9f8..0a188b2a0 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2,14 +2,17 @@ import concurrent.futures import dataclasses import functools -import json import os import pathlib import re import shutil import subprocess +import sys import textwrap +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +import ops_config # noqa: E402 + try: import clang.cindex from clang.cindex import CursorKind @@ -1848,6 +1851,52 @@ 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 _select_ops_from_config(ops, config, config_path): + selected = {} + + for op_name, selection in config.items(): + headers = selection["headers"] + + if headers is not None: + selected[op_name] = [ + _implementation_from_json(header) for header in headers + ] + continue + + if op_name not in ops: + raise ValueError( + f"{config_path}: operator {op_name!r} is not available for " + "the active devices" + ) + + slots = selection["implementations"] + + if slots is None: + selected[op_name] = ops[op_name] + continue + + headers_by_slot = {} + + for implementation in ops[op_name]: + slot = ops_config.implementation_slot(implementation.path) + headers_by_slot.setdefault(slot, []).append(implementation) + + missing = [slot for slot in slots if slot not in headers_by_slot] + + if missing: + formatted = ", ".join(str(slot) for slot in missing) + raise ValueError( + f"{config_path}: operator {op_name!r} has no active " + f"implementation at slot(s) {formatted}" + ) + + selected[op_name] = [ + implementation for slot in slots for implementation in headers_by_slot[slot] + ] + + return selected + + def _get_all_ops( devices, with_torch=False, @@ -2084,6 +2133,11 @@ def _dispatch_gen_batch_size(): type=str, help="Operator allowlist to generate. Accepts names separated by spaces or commas.", ) + parser.add_argument( + "--ops-config", + type=pathlib.Path, + help="Path to an `ops.json` operator and implementation selection.", + ) parser.add_argument( "--strict-ops", action="store_true", @@ -2101,25 +2155,18 @@ def _dispatch_gen_batch_size(): for directory in (_BINDINGS_DIR, _GENERATED_SRC_DIR, _INCLUDE_DIR): directory.mkdir(parents=True, exist_ok=True) - ops_json = pathlib.Path("ops.json") + config_path = args.ops_config + ops = _get_all_ops( + args.devices, + with_torch=args.with_torch, + with_ninetoothed=args.with_ninetoothed, + with_linked=args.with_linked, + with_triton=args.with_triton, + ) - if ops_json.exists(): - raw_ops = json.loads(ops_json.read_text()) - ops = { - op_name: [ - _implementation_from_json(implementation) - for implementation in implementations - ] - for op_name, implementations in raw_ops.items() - } - else: - ops = _get_all_ops( - args.devices, - with_torch=args.with_torch, - with_ninetoothed=args.with_ninetoothed, - with_linked=args.with_linked, - with_triton=args.with_triton, - ) + if config_path is not None: + config = ops_config.load_ops_config(config_path) + ops = _select_ops_from_config(ops, config, config_path) ops = _filter_ops( ops, diff --git a/scripts/ops_config.py b/scripts/ops_config.py new file mode 100644 index 000000000..b1e764fb8 --- /dev/null +++ b/scripts/ops_config.py @@ -0,0 +1,214 @@ +import json +import pathlib +import re + + +_OPERATOR_SPECIALIZATION_RE = re.compile( + r"\bclass\s+Operator<\s*[^,>]+\s*,\s*[^,>]+\s*" + r"(?:,\s*(\d+)\s*)?>" +) + + +class OpsConfigError(ValueError): + pass + + +class _StrictJsonObject(dict): + pass + + +def _strict_object(pairs): + value = _StrictJsonObject() + + for key, item in pairs: + if key in value: + raise OpsConfigError(f"duplicate key {key!r}") + value[key] = item + + return value + + +def load_ops_config(path): + path = pathlib.Path(path) + + try: + config = json.loads( + path.read_text(encoding="utf-8"), object_pairs_hook=_strict_object + ) + except (OSError, json.JSONDecodeError, OpsConfigError) as error: + raise OpsConfigError(f"failed to read {path}: {error}") from error + + if not isinstance(config, dict): + raise OpsConfigError(f"{path} must contain a JSON object") + + normalized = {} + + for op_name, value in config.items(): + if not isinstance(op_name, str) or not op_name.strip(): + raise OpsConfigError(f"{path}: operator names must be non-empty strings") + if op_name != op_name.strip(): + raise OpsConfigError( + f"{path}: operator name {op_name!r} contains surrounding whitespace" + ) + + normalized[op_name] = _normalize_selection(path, op_name, value) + + return normalized + + +def _normalize_selection(path, op_name, value): + if isinstance(value, str): + if not value.strip(): + raise OpsConfigError( + f"{path}: {op_name!r} implementation path must not be empty" + ) + + return {"headers": [value], "implementations": None} + + if isinstance(value, list): + if not value: + raise OpsConfigError( + f"{path}: {op_name!r} implementation paths must be a non-empty array" + ) + headers = [_normalize_header(path, op_name, item) for item in value] + identities = [ + header if isinstance(header, str) else (header["path"], header["backend"]) + for header in headers + ] + if len(identities) != len(set(identities)): + raise OpsConfigError( + f"{path}: {op_name!r} implementation paths contain duplicates" + ) + return {"headers": headers, "implementations": None} + + if not isinstance(value, dict): + raise OpsConfigError( + f"{path}: {op_name!r} must map to implementation path(s) or an object" + ) + + unknown_keys = sorted(set(value) - {"implementations"}) + if unknown_keys: + raise OpsConfigError( + f"{path}: {op_name!r} contains unknown keys: {', '.join(unknown_keys)}" + ) + if "implementations" not in value: + raise OpsConfigError( + f"{path}: {op_name!r} is missing required key 'implementations'" + ) + + implementations = value["implementations"] + + if implementations == "all": + implementations = None + elif isinstance(implementations, list): + if not implementations: + raise OpsConfigError( + f"{path}: {op_name!r} implementations must not be empty" + ) + if any(type(slot) is not int or not 0 <= slot < 32 for slot in implementations): + raise OpsConfigError( + f"{path}: {op_name!r} implementations must contain integers " + "between 0 and 31" + ) + if len(implementations) != len(set(implementations)): + raise OpsConfigError( + f"{path}: {op_name!r} implementations contain duplicates" + ) + implementations = tuple(implementations) + else: + raise OpsConfigError( + f"{path}: {op_name!r} implementations must be 'all' or an array" + ) + + return {"headers": None, "implementations": implementations} + + +def _normalize_header(path, op_name, value): + if isinstance(value, str): + if value.strip(): + return value + raise OpsConfigError( + f"{path}: {op_name!r} implementation path must not be empty" + ) + + if not isinstance(value, dict): + raise OpsConfigError( + f"{path}: {op_name!r} implementation entries must be paths or " + "structured descriptors" + ) + + unknown_keys = sorted(set(value) - {"path", "backend"}) + if unknown_keys: + raise OpsConfigError( + f"{path}: {op_name!r} implementation descriptor contains unknown " + f"keys: {', '.join(unknown_keys)}" + ) + + for key in ("path", "backend"): + if ( + key not in value + or not isinstance(value[key], str) + or not value[key].strip() + ): + raise OpsConfigError( + f"{path}: {op_name!r} implementation descriptor requires a " + f"non-empty string {key!r}" + ) + + return {"path": value["path"], "backend": value["backend"]} + + +def implementation_path(header): + return header if isinstance(header, str) else header["path"] + + +def selected_op_names(config): + return list(config) + + +def selected_slots(config, op_name): + selection = config.get(op_name) + + if selection is None or selection["headers"] is not None: + return None + + return selection["implementations"] + + +def implementation_slot(path): + path = pathlib.Path(path) + + try: + text = path.read_text(encoding="utf-8") + except OSError as error: + raise OpsConfigError(f"failed to read {path}: {error}") from error + + slots = { + int(match.group(1)) if match.group(1) is not None else 0 + for match in _OPERATOR_SPECIALIZATION_RE.finditer(text) + } + + if len(slots) != 1: + formatted = ", ".join(str(slot) for slot in sorted(slots)) or "none" + raise OpsConfigError( + f"{path} must declare exactly one implementation slot; found {formatted}" + ) + + return slots.pop() + + +def torch_op_names(config, default_ops=(), slot=8): + selected = [] + + for op_name, selection in config.items(): + if selection["headers"] is not None: + continue + + implementations = selection["implementations"] + + if (implementations is None and op_name in default_ops) or ( + implementations is not None and slot in implementations + ): + selected.append(op_name) + + return selected diff --git a/scripts/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index 03f63be9e..66d63d9a8 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -13,6 +13,9 @@ import yaml +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +import ops_config # noqa: E402 + _PROJECT_DIR = pathlib.Path(__file__).resolve().parents[1] _DEFAULT_SOURCE_ROOT = _PROJECT_DIR / "src" / "linked" _DEFAULT_OUTPUT_DIR = _PROJECT_DIR / "generated" / "linked" @@ -139,9 +142,11 @@ def _find_platform_dirs(source_root, device): return platform_dirs -def _load_libraries(platform_dir, device, transport): +def _load_libraries(platform_dir, device, transport, selected_libraries=None): libraries = {} for path in sorted(platform_dir.glob("*.yaml")): + if selected_libraries is not None and path.stem not in selected_libraries: + continue data = _load_yaml_mapping(path, _LIBRARY_KEYS) libraries[path.stem] = LibraryConfig( device=device, @@ -157,13 +162,60 @@ def _load_libraries(platform_dir, device, transport): return libraries -def _load_bindings(platform_dir, device, transport, selected_ops): +def _binding_is_selected(name, header, selected_ops, config): + if selected_ops is not None and name not in selected_ops: + return False + + if config is not None: + selection = config.get(name) + + if selection is None: + return False + + headers = selection["headers"] + + if headers is not None: + selected_headers = { + ( + _PROJECT_DIR / ops_config.implementation_path(selected_header) + ).resolve() + for selected_header in headers + } + + return header.resolve() in selected_headers + + slots = selection["implementations"] + + return slots is None or ops_config.implementation_slot(header) in slots + + return selected_ops is None or name in selected_ops + + +def _load_bindings(platform_dir, device, transport, selected_ops, config): bindings = [] for path in sorted((platform_dir / "ops").glob("*/*.yaml")): name = path.parent.name - if selected_ops is not None and name not in selected_ops: + header = path.with_suffix(".h") + + if config is not None and name not in config: + continue + if config is None and selected_ops is not None and name not in selected_ops: + continue + if not header.is_file(): + raise ResolutionError(f"{path}: missing sibling {header.name}") + + try: + selected = _binding_is_selected(name, header, selected_ops, config) + except ops_config.OpsConfigError as error: + raise ResolutionError(str(error)) from error + + if not selected: continue + source = path.with_suffix(".cc") + if not source.is_file(): + raise ResolutionError(f"{path}: missing sibling {source.name}") + data = _load_yaml_mapping(path, _BINDING_KEYS, {"library"}) symbols = data.get("required_symbols") operator_schema = data.get("operator_schema") @@ -199,13 +251,6 @@ def _load_bindings(platform_dir, device, transport, selected_ops): raise ResolutionError(f"{path}: operator_schema requires dispatch_key") dispatch_key = _require_string(data, "dispatch_key", path) - header = path.with_suffix(".h") - source = path.with_suffix(".cc") - if not header.is_file(): - raise ResolutionError(f"{path}: missing sibling {header.name}") - if not source.is_file(): - raise ResolutionError(f"{path}: missing sibling {source.name}") - bindings.append( BindingConfig( device=device, @@ -533,6 +578,7 @@ def _normalize_values(values): def resolve_linked_ops( devices, ops=None, + config_path=None, *, source_root=_DEFAULT_SOURCE_ROOT, output_dir=_DEFAULT_OUTPUT_DIR, @@ -546,16 +592,33 @@ def resolve_linked_ops( selected_ops = _normalize_values(ops) selected_op_set = set(selected_ops) if selected_ops is not None else None + try: + selection_config = ( + ops_config.load_ops_config(config_path) if config_path is not None else None + ) + except ops_config.OpsConfigError as error: + raise ResolutionError(str(error)) from error + bindings = [] library_configs = {} for device in devices: for transport, platform_dir in _find_platform_dirs(source_root, device): - libraries = _load_libraries(platform_dir, device, transport) - for name, config in libraries.items(): - library_configs[(transport, device, name)] = config - bindings.extend( - _load_bindings(platform_dir, device, transport, selected_op_set) + platform_bindings = _load_bindings( + platform_dir, + device, + transport, + selected_op_set, + selection_config, + ) + bindings.extend(platform_bindings) + libraries = _load_libraries( + platform_dir, + device, + transport, + {binding.library for binding in platform_bindings}, ) + for name, library_config in libraries.items(): + library_configs[(transport, device, name)] = library_config bindings.sort( key=lambda binding: ( @@ -677,6 +740,7 @@ def _parse_args(): ) parser.add_argument("--devices", nargs="+", required=True) parser.add_argument("--ops", nargs="*") + parser.add_argument("--ops-config", type=pathlib.Path) parser.add_argument("--source-root", default=_DEFAULT_SOURCE_ROOT) parser.add_argument("--output-dir", default=_DEFAULT_OUTPUT_DIR) parser.add_argument("--nm", default=os.environ.get("CMAKE_NM", "nm")) @@ -685,15 +749,28 @@ def _parse_args(): return parser.parse_args() +def _selection_from_environment(ops, config_path): + if ops is not None or config_path is not None: + return ops, config_path + + value = os.environ.get("INFINI_OPS_OPS") + + if value is None: + return ops, config_path + if pathlib.Path(value).suffix.lower() == ".json": + return None, pathlib.Path(value) + + return [value], None + + def main(): args = _parse_args() - ops = args.ops - if ops is None and "INFINI_OPS_OPS" in os.environ: - ops = [os.environ["INFINI_OPS_OPS"]] + ops, config_path = _selection_from_environment(args.ops, args.ops_config) try: resolve_linked_ops( args.devices, ops, + config_path, source_root=args.source_root, output_dir=args.output_dir, nm=args.nm, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d88058efb..36f946d8d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -505,7 +505,7 @@ if(WITH_ASCEND) endif() set(INFINI_OPS_OPS "" CACHE STRING - "Semicolon- or comma-separated operator allowlist for generated wrappers and bindings") + "Operator allowlist or path to an ops.json implementation selection") set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL "Build only the smoke-test operator subset") set(_infini_ops_smoke_ops @@ -532,7 +532,33 @@ if(INFINI_OPS_SMOKE_BUILD) endif() endif() -if(INFINI_OPS_OPS) +set(_infini_ops_ops_config "") +if(INFINI_OPS_OPS MATCHES "\\.json$") + get_filename_component(_infini_ops_ops_config "${INFINI_OPS_OPS}" + ABSOLUTE BASE_DIR "${PROJECT_SOURCE_DIR}") + if(NOT EXISTS "${_infini_ops_ops_config}") + message(FATAL_ERROR + "Operator selection `${_infini_ops_ops_config}` does not exist.") + endif() +elseif(NOT INFINI_OPS_OPS AND EXISTS "${PROJECT_SOURCE_DIR}/ops.json") + set(_infini_ops_ops_config "${PROJECT_SOURCE_DIR}/ops.json") +endif() + +if(_infini_ops_ops_config) + file(GLOB_RECURSE _infini_ops_implementation_headers CONFIGURE_DEPENDS + "${PROJECT_SOURCE_DIR}/src/native/*.h" + "${PROJECT_SOURCE_DIR}/src/torch/*.h" + "${PROJECT_SOURCE_DIR}/src/ninetoothed/*.h" + "${PROJECT_SOURCE_DIR}/src/linked/*.h" + "${PROJECT_SOURCE_DIR}/src/triton/*.h") + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${_infini_ops_ops_config}" + "${PROJECT_SOURCE_DIR}/scripts/ops_config.py" + ${_infini_ops_implementation_headers}) + message(STATUS "Operator selection: ${_infini_ops_ops_config}") +endif() + +if(INFINI_OPS_OPS AND NOT INFINI_OPS_OPS MATCHES "\\.json$") string(REPLACE "," ";" _infini_ops_op_allowlist "${INFINI_OPS_OPS}") message(STATUS "Wrapper op allowlist: ${_infini_ops_op_allowlist}") endif() @@ -581,7 +607,11 @@ if(WITH_LINKED) --nm "${_linked_nm}" --readelf "${_linked_readelf}" --cxxfilt "${_linked_cxxfilt}") - if(INFINI_OPS_OPS) + if(_infini_ops_ops_config) + list(APPEND _linked_resolver_args + --ops-config "${_infini_ops_ops_config}") + endif() + if(_infini_ops_op_allowlist) list(APPEND _linked_resolver_args --ops ${_infini_ops_op_allowlist}) endif() @@ -665,6 +695,10 @@ if(WITH_TORCH) # which we then glob below alongside any hand-written torch sources. find_package(Python COMPONENTS Interpreter REQUIRED) + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py" + "${PROJECT_SOURCE_DIR}/scripts/torch_ops.yaml") + # Pin codegen to the locally installed torch version so vendor # forks (Cambricon's `torch_mlu` 2.1.0, etc.) get a schema whose # `at::_out` overloads match the headers they ship. Without @@ -688,6 +722,10 @@ if(WITH_TORCH) set(_torch_codegen_args ${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py --pytorch-version ${_torch_version_tag}) + if(_infini_ops_ops_config) + list(APPEND _torch_codegen_args + --ops-config "${_infini_ops_ops_config}") + endif() if(INFINI_OPS_TORCH_OPS) string(REPLACE "," ";" _torch_op_allowlist "${INFINI_OPS_TORCH_OPS}") list(APPEND _torch_codegen_args --ops ${_torch_op_allowlist}) @@ -988,7 +1026,11 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS) # failures. set(GENERATOR_ARGS --devices ${DEVICE_LIST}) - if(INFINI_OPS_OPS) + if(_infini_ops_ops_config) + list(APPEND GENERATOR_ARGS + --ops-config "${_infini_ops_ops_config}") + endif() + if(_infini_ops_op_allowlist) list(APPEND GENERATOR_ARGS --ops ${_infini_ops_op_allowlist}) endif() if(WITH_TORCH) diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index 43dffc4bc..6dbed1f71 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -21,6 +21,41 @@ def _load_generator_module(): return module +def test_select_op_names_honors_slot_8_and_legacy_headers(): + module = _load_generator_module() + config = { + "default_all": {"headers": None, "implementations": None}, + "default_native": {"headers": None, "implementations": (0,)}, + "explicit_torch": {"headers": None, "implementations": (8,)}, + "legacy": {"headers": ["legacy.h"], "implementations": None}, + } + + assert module._select_op_names( + None, + ["default_all", "default_native"], + config, + ) == ["default_all", "explicit_torch"] + assert module._select_op_names( + ["explicit_torch"], + ["default_all", "default_native"], + config, + ) == ["explicit_torch"] + + +def test_select_op_names_maps_public_names_to_aten_names(): + module = _load_generator_module() + config = { + "div": {"headers": None, "implementations": (8,)}, + "internal_log_softmax": {"headers": None, "implementations": (8,)}, + } + + assert module._select_op_names( + None, + ["div", "div_", "_log_softmax"], + config, + ) == ["div", "div_", "_log_softmax"] + + def test_load_aten_entries_uses_packaged_torchgen(monkeypatch): module = _load_generator_module() entries = [{"func": "relu.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)"}] diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 6025f44c6..0f34483a0 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -600,6 +600,34 @@ def test_torch_system_compiler_receives_host_range_profile_definition(): ) in cmake +def test_explicit_op_allowlist_precedes_implicit_ops_json(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + wrapper = ( + pathlib.Path(__file__).parents[1] / "scripts" / "generate_wrappers.py" + ).read_text(encoding="utf-8") + + assert ( + 'elseif(NOT INFINI_OPS_OPS AND EXISTS "${PROJECT_SOURCE_DIR}/ops.json")' + in cmake + ) + assert 'if(INFINI_OPS_OPS AND NOT INFINI_OPS_OPS MATCHES "\\\\.json$")' in cmake + assert "implicit_config_path" not in wrapper + + +def test_operator_selection_tracks_implementation_header_changes(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + + for root in ("native", "torch", "ninetoothed", "linked", "triton"): + assert f'"${{PROJECT_SOURCE_DIR}}/src/{root}/*.h"' in cmake + assert "${_infini_ops_implementation_headers}" in cmake + assert '"${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py"' in cmake + assert '"${PROJECT_SOURCE_DIR}/scripts/torch_ops.yaml"' in cmake + + def test_generated_dispatch_calls_start_with_dispatch_profile_scope( tmp_path, monkeypatch ): @@ -704,6 +732,62 @@ def test_filter_ops_strict_rejects_unavailable_ops(): raise AssertionError("strict unknown ops should fail") +def test_select_ops_from_config_filters_implementation_slots(tmp_path): + module = _load_generator_module() + slot_0 = tmp_path / "slot_0.h" + slot_16 = tmp_path / "slot_16.h" + slot_0.write_text("class Operator : public Add {};") + slot_16.write_text("class Operator : public Add {};") + implementation_0 = module._Implementation(slot_0, "native") + implementation_16 = module._Implementation(slot_16, "native") + config = { + "add": {"headers": None, "implementations": (16,)}, + } + + assert module._select_ops_from_config( + {"add": [implementation_0, implementation_16]}, config, "ops.json" + ) == {"add": [implementation_16]} + + +def test_select_ops_from_config_rejects_missing_slot(tmp_path): + module = _load_generator_module() + slot_0 = tmp_path / "slot_0.h" + slot_0.write_text("class Operator : public Add {};") + implementation_0 = module._Implementation(slot_0, "native") + config = { + "add": {"headers": None, "implementations": (16,)}, + } + + try: + module._select_ops_from_config({"add": [implementation_0]}, config, "ops.json") + except ValueError as exc: + assert "slot(s) 16" in str(exc) + else: + raise AssertionError("unavailable implementation slot should fail") + + +def test_select_ops_from_config_preserves_path_and_backend_descriptors(): + module = _load_generator_module() + config = { + "add": { + "headers": [ + "src/native/cpu/ops/add/add.h", + {"path": "custom/add.h", "backend": "custom"}, + ], + "implementations": None, + }, + } + + assert module._select_ops_from_config({}, config, "ops.json") == { + "add": [ + module._Implementation( + pathlib.Path("src/native/cpu/ops/add/add.h"), "native" + ), + module._Implementation(pathlib.Path("custom/add.h"), "custom"), + ] + } + + def test_linked_implementations_require_explicit_scan_flag(monkeypatch, tmp_path): module = _load_generator_module() src_dir = tmp_path / "moore" / "src" diff --git a/tests/test_ops_config.py b/tests/test_ops_config.py new file mode 100644 index 000000000..550cf8f67 --- /dev/null +++ b/tests/test_ops_config.py @@ -0,0 +1,114 @@ +import pathlib +import sys + +import pytest + + +_SCRIPTS_DIR = pathlib.Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(_SCRIPTS_DIR)) + +import ops_config # noqa: E402 + + +def _load(tmp_path, text): + path = tmp_path / "ops.json" + path.write_text(text) + + return ops_config.load_ops_config(path) + + +def test_loads_legacy_paths_and_slot_selections(tmp_path): + config = _load( + tmp_path, + """{ + "add": "src/native/cpu/ops/add/add.h", + "gemm": [ + "first.h", + {"path": "custom/second.h", "backend": "custom"} + ], + "relu": {"implementations": "all"}, + "sampling": {"implementations": [0, 16]} +} +""", + ) + + assert config == { + "add": {"headers": ["src/native/cpu/ops/add/add.h"], "implementations": None}, + "gemm": { + "headers": ["first.h", {"path": "custom/second.h", "backend": "custom"}], + "implementations": None, + }, + "relu": {"headers": None, "implementations": None}, + "sampling": {"headers": None, "implementations": (0, 16)}, + } + assert ops_config.selected_op_names(config) == ["add", "gemm", "relu", "sampling"] + assert ops_config.selected_slots(config, "sampling") == (0, 16) + assert ops_config.torch_op_names(config, ["relu", "sampling"]) == ["relu"] + config["sampling"]["implementations"] = (8,) + assert ops_config.torch_op_names(config, ["relu"]) == ["relu", "sampling"] + + +@pytest.mark.parametrize( + ("declaration", "slot"), + ( + ("class Operator : public Add {};", 0), + ( + "class Operator : public Add {};", + 16, + ), + ), +) +def test_reads_implementation_slot(tmp_path, declaration, slot): + header = tmp_path / "implementation.h" + header.write_text(declaration) + + assert ops_config.implementation_slot(header) == slot + + +def test_reads_one_slot_declared_for_multiple_devices(tmp_path): + header = tmp_path / "implementation.h" + header.write_text( + "class Operator : public Add {};\n" + "class Operator : public Add {};\n" + ) + + assert ops_config.implementation_slot(header) == 16 + + +def test_rejects_multiple_slots_in_one_header(tmp_path): + header = tmp_path / "implementation.h" + header.write_text( + "class Operator : public Add {};\n" + "class Operator : public Add {};\n" + ) + + with pytest.raises( + ops_config.OpsConfigError, + match="exactly one implementation slot; found 16, 17", + ): + ops_config.implementation_slot(header) + + +@pytest.mark.parametrize( + ("text", "message"), + ( + ('{"add": {}, "add": {"implementations": "all"}}', "duplicate key"), + ('{"add": {}}', "missing required key"), + ('{"add": {"slots": [0]}}', "unknown keys"), + ('{"add": {"implementations": []}}', "must not be empty"), + ('{"add": {"implementations": [0, 0]}}', "contain duplicates"), + ('{"add": {"implementations": [32]}}', "between 0 and 31"), + ('{"add": {"implementations": [true]}}', "between 0 and 31"), + ('{"add": ["first.h", 1]}', "paths or structured descriptors"), + ('{"add": [{"path": "x.h"}]}', "non-empty string 'backend'"), + ('{"add": [{"backend": "custom"}]}', "non-empty string 'path'"), + ( + '{"add": [{"path": "x.h", "backend": "custom", "slot": 1}]}', + "unknown keys", + ), + ), +) +def test_rejects_invalid_config(tmp_path, text, message): + with pytest.raises(ops_config.OpsConfigError, match=message): + _load(tmp_path, text) diff --git a/tests/test_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 72a7b45d0..3444b008b 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -226,6 +226,90 @@ def test_resolve_supports_multiple_implementations_for_one_operator( ] +def test_resolve_filters_linked_dependencies_by_slot(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + platform, op_dir = _write_linked_config(source_root) + (op_dir / "vllm.h").write_text( + "class Operator {};\n" + ) + (platform / "unused.yaml").write_text( + "python_distribution_package: unused\nlibrary_glob: unused/_C*.so\n" + ) + (op_dir / "unused.yaml").write_text( + "library: unused\nrequired_symbols:\n - unused()\n" + ) + (op_dir / "unused.h").write_text( + "class Operator {};\n" + ) + (op_dir / "unused.cc").write_text("// definition\n") + config_path = tmp_path / "ops.json" + config_path.write_text('{"silu_and_mul": {"implementations": [16]}}\n') + library_path = tmp_path / "vllm" / "_C.so" + library_path.parent.mkdir() + library_path.touch() + located = [] + + def locate(config): + located.append(config.name) + return library_path + + monkeypatch.setattr(module, "_locate_distribution_library", locate) + exported = {"silu_and_mul(at::Tensor&, at::Tensor&)"} + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: (exported, exported), + ) + + payload = module.resolve_linked_ops( + ["metax"], + config_path=config_path, + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + assert located == ["vllm"] + assert [op["implementation"] for op in payload["operators"]] == ["vllm"] + + +def test_resolve_matches_structured_implementation_descriptor(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + _, op_dir = _write_linked_config(source_root) + config_path = tmp_path / "ops.json" + config_path.write_text( + '{"silu_and_mul": [{"path": "' + + (op_dir / "vllm.h").relative_to(tmp_path).as_posix() + + '", "backend": "linked"}]}\n' + ) + library_path = tmp_path / "vllm" / "_C.so" + library_path.parent.mkdir() + library_path.touch() + + monkeypatch.setattr(module, "_PROJECT_DIR", tmp_path) + monkeypatch.setattr( + module, + "_locate_distribution_library", + lambda _config: library_path, + ) + exported = {"silu_and_mul(at::Tensor&, at::Tensor&)"} + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: (exported, exported), + ) + + payload = module.resolve_linked_ops( + ["metax"], + config_path=config_path, + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + assert [op["implementation"] for op in payload["operators"]] == ["vllm"] + + @pytest.mark.parametrize( ("binding", "message"), ( @@ -699,3 +783,31 @@ def read_text(self, filename): return json.dumps({"url": url, "dir_info": {"editable": editable}}) assert module._locate_editable_distribution_root(FakeDistribution()) is None + + +def test_environment_selection_recognizes_ops_json(monkeypatch, tmp_path): + module = _load_resolver_module() + config_path = tmp_path / "ops.JSON" + monkeypatch.setenv("INFINI_OPS_OPS", str(config_path)) + + assert module._selection_from_environment(None, None) == (None, config_path) + + +def test_environment_selection_keeps_inline_allowlist(monkeypatch): + module = _load_resolver_module() + monkeypatch.setenv("INFINI_OPS_OPS", "add,gemm") + + assert module._selection_from_environment(None, None) == (["add,gemm"], None) + + +def test_explicit_selection_precedes_environment(monkeypatch, tmp_path): + module = _load_resolver_module() + env_config = tmp_path / "environment.json" + explicit_config = tmp_path / "explicit.json" + monkeypatch.setenv("INFINI_OPS_OPS", str(env_config)) + + assert module._selection_from_environment(["add"], None) == (["add"], None) + assert module._selection_from_environment(None, explicit_config) == ( + None, + explicit_config, + )