feat(build): select operator implementations from ops.json

Author: voltjia

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.
 

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

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,13 +1647,42 @@ 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(
         "--ops",
         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]] = []

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,