Preserve an empty selective primitive-operator set (#24815)

codegen/tools/combine_prim_ops_headers.py

@@ -7,6 +7,12 @@
 """
 Script to combine multiple selected_prim_ops.h header files into a single header.
 This is used by selected_prim_operators_genrule to merge prim ops headers from dependencies.
+
+The combiner only runs for a build that asked for a selected primitive-operator
+set, so the union it produces is that set, and an empty union means the build
+selected none of them. Two cases are errors instead, because neither is a
+selection this tool may silently narrow: a dependency that reports no selection
+at all, or one that does not enumerate required operator overloads.
 """
 
 import argparse
@@ -15,6 +21,12 @@
 from pathlib import Path
 from typing import List, Set
 
+# Written into a dependency's selected_prim_ops.h by
+# executorch/codegen/tools/gen_selected_prim_ops.py when a dependency does not
+# enumerate every required operator overload. Kept as a literal so this tool needs
+# none of that one's dependencies; the agreement is pinned by its test.
+UNBOUNDED_SELECTION_MACRO = "EXECUTORCH_PRIM_OPS_SELECTION_UNBOUNDED"
+
 
 def read_header_file(file_path: Path) -> Set[str]:
     """
@@ -25,23 +37,34 @@ def read_header_file(file_path: Path) -> Set[str]:
 
     Returns:
         macros_set where macros_set contains unique macro defines
+
+    Raises:
+        OSError: the header is missing or cannot be read. It carries this
+            dependency's selection, so skipping it would silently narrow the
+            combined set.
+        ValueError: the dependency does not enumerate all required overloads,
+            so no selected primitive-operator set can be derived from it.
     """
     macros = set()
 
-    try:
-        with open(file_path, "r") as f:
-            for line in f:
-                line = line.strip()
-
-                # Extract #define statements for prim ops
-                if line.startswith("#define INCLUDE_") and not line.startswith(
-                    "#define EXECUTORCH_ENABLE"
-                ):
-                    macros.add(line)
-    except FileNotFoundError:
-        print(f"Warning: Header file not found: {file_path}", file=sys.stderr)
-    except Exception as e:
-        print(f"Error reading {file_path}: {e}", file=sys.stderr)
+    with open(file_path, "r") as f:
+        for line in f:
+            line = line.strip()
+
+            if line == f"#define {UNBOUNDED_SELECTION_MACRO}":
+                raise ValueError(
+                    f"{file_path} selects all operators or unenumerated operator "
+                    "overloads, so a selected primitive-operator set cannot be "
+                    "derived from it. List exact overloads or build "
+                    "the full primitive-operator registry instead, by passing "
+                    "include_all_prim_ops = True."
+                )
+
+            # Extract #define statements for prim ops
+            if line.startswith("#define INCLUDE_") and not line.startswith(
+                "#define EXECUTORCH_ENABLE"
+            ):
+                macros.add(line)
 
     return macros
 
@@ -55,20 +78,13 @@ def combine_prim_ops_headers(header_file_paths: List[str], output_path: str) ->
         output_path: Path to output the combined header
     """
     all_macros = set()
-    has_selective_build = False
 
     # Read all header files and collect unique macros
     for header_file_path in header_file_paths:
         header_file = Path(header_file_path) / "selected_prim_ops.h"
-        if os.path.exists(header_file):
-            macros = read_header_file(header_file)
-            all_macros.update(macros)
-            if len(all_macros) > 0:
-                has_selective_build = True
-        else:
-            print(
-                f"Warning: Header file does not exist: {header_file}", file=sys.stderr
-            )
+        if not header_file.exists():
+            raise FileNotFoundError(header_file)
+        all_macros.update(read_header_file(header_file))
 
     # Generate combined header
     header_content = [
@@ -78,28 +94,14 @@ def combine_prim_ops_headers(header_file_paths: List[str], output_path: str) ->
         "",
         "#pragma once",
         "",
+        "// Enable selective build for prim ops",
+        "#define EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD",
+        "",
+        "// Combined prim ops macros from all dependencies",
     ]
 
-    if all_macros and has_selective_build:
-        header_content.extend(
-            [
-                "// Enable selective build for prim ops",
-                "#define EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD",
-                "",
-                "// Combined prim ops macros from all dependencies",
-            ]
-        )
-
-        # Sort macros for deterministic output
-        sorted_macros = sorted(all_macros)
-        header_content.extend(sorted_macros)
-    else:
-        header_content.extend(
-            [
-                "// No prim ops found in dependencies - all prim ops will be included",
-                "// Selective build is disabled",
-            ]
-        )
+    # Sort macros for deterministic output
+    header_content.extend(sorted(all_macros))
 
     header_content.append("")
 
@@ -146,7 +148,6 @@ def main():
     )
 
     args = parser.parse_args()
-    import os
 
     header_file_paths = _get_header_file_paths_from_query_output(args.header_files)
 

codegen/tools/gen_oplist.py

@@ -229,7 +229,7 @@ def gen_oplist(
         or include_all_operators
     ):
         # dump empty yaml file
-        _dump_yaml([], output_path)
+        _dump_yaml([], output_path, include_all_operators=include_all_operators)
         return
 
     assert output_path, "Need to provide output_path for dumped yaml file."

codegen/tools/gen_selected_prim_ops.py

@@ -5,8 +5,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-unsafe
-
 import argparse
 import os
 import re
@@ -15,6 +13,7 @@
 
 import yaml
 from torchgen.code_template import CodeTemplate  # type: ignore[import-not-found]
+from torchgen.selective_build.operator import SelectiveBuildOperator
 
 
 selected_prim_ops_h_template_str = """#pragma once
@@ -47,32 +46,85 @@ def normalize_op_name(op_name: str) -> str:
     return normalized
 
 
-def read_prim_op_names_from_yaml(op_selection_yaml_path: str) -> List[str]:
+UNBOUNDED_SELECTION_MACRO = "EXECUTORCH_PRIM_OPS_SELECTION_UNBOUNDED"
+
+UNBOUNDED_SELECTION_GUARD = f"""#define {UNBOUNDED_SELECTION_MACRO}
+#if defined(EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD)
+#error \\
+    "A dependency selects all operators or unenumerated operator overloads. List exact overloads or build the full primitive-operator registry instead."
+#endif"""
+
+
+def _load_selection(op_selection_yaml_path: str) -> dict:
+    """
+    Load a selected_operators.yaml.
+
+    An absent or malformed file is an error. Callers rely on the selection being
+    the operators this build chose, so "no file" and "chose nothing" must not
+    look alike.
+    """
+    if not os.path.isfile(op_selection_yaml_path):
+        raise FileNotFoundError(
+            f"selected operator yaml not found: {op_selection_yaml_path}"
+        )
+
     with open(op_selection_yaml_path, "r") as f:
-        selected_operators = yaml.safe_load(f) or {}
+        selected_operators = yaml.safe_load(f)
 
+    if selected_operators is None:
+        selected_operators = {}
+    if not isinstance(selected_operators, dict):
+        raise ValueError(
+            f"selected operator yaml is not a mapping: {op_selection_yaml_path}"
+        )
+    return selected_operators
+
+
+def read_prim_op_selection_from_yaml(
+    op_selection_yaml_path: str,
+) -> tuple[List[str], bool]:
+    """Read named operators and whether any required overloads are unenumerated."""
+    selected_operators = _load_selection(op_selection_yaml_path)
     prim_op_names = set()
-    operators = selected_operators.get("operators", {})
-    if operators:
-        prim_op_names.update(operators.keys())
+    unbounded = bool(selected_operators.get("include_all_operators"))
+    for op_name, metadata in (selected_operators.get("operators") or {}).items():
+        operator = SelectiveBuildOperator.from_yaml_dict(op_name, metadata)
+        if operator.include_all_overloads:
+            unbounded = True
+        else:
+            prim_op_names.add(op_name)
 
-    et_kernel_metadata = selected_operators.get("et_kernel_metadata", {})
-    if et_kernel_metadata:
-        prim_op_names.update(et_kernel_metadata.keys())
+    prim_op_names.update((selected_operators.get("et_kernel_metadata") or {}).keys())
+    return sorted(prim_op_names), unbounded
 
-    return sorted(prim_op_names)
 
+def selection_is_unbounded(op_selection_yaml_path: str) -> bool:
+    """Whether the selection omits required operator or overload names."""
+    return read_prim_op_selection_from_yaml(op_selection_yaml_path)[1]
 
-def write_selected_prim_ops(prim_op_names: List[str], output_dir: str) -> None:
+
+def read_prim_op_names_from_yaml(op_selection_yaml_path: str) -> List[str]:
+    """Read enumerated names; use the selection reader to also check completeness."""
+    return read_prim_op_selection_from_yaml(op_selection_yaml_path)[0]
+
+
+def write_selected_prim_ops(
+    prim_op_names: List[str], output_dir: str, *, unbounded_selection: bool = False
+) -> None:
     """
     Generate selected_prim_ops.h from a list of prim op names.
 
     Args:
         prim_op_names: List of prim op names like ["executorch_prim::et_view.default", "aten::sym_size.int"]
         output_dir: Directory where to write selected_prim_ops.h
+        unbounded_selection: Required operator overloads are not enumerated.
+            The header rejects its use as a selected set instead of claiming
+            that no primitive operators are needed.
     """
     # Generate #define statements for each op
     defines = []
+    if unbounded_selection:
+        defines.append(UNBOUNDED_SELECTION_GUARD)
     for op_name in prim_op_names:
         macro_name = normalize_op_name(op_name)
         defines.append(f"#define {macro_name}")
@@ -116,17 +168,23 @@ def main(argv: List[Any]) -> None:
         parser.error("one of --prim-op-names or --op-selection-yaml-path is required")
 
     prim_op_names: List[str] = []
+    unbounded = False
     if options.prim_op_names is not None:
         # Parse comma-separated prim op names
         prim_op_names.extend(
             name.strip() for name in options.prim_op_names.split(",") if name.strip()
         )
     if options.op_selection_yaml_path:
-        prim_op_names.extend(
-            read_prim_op_names_from_yaml(options.op_selection_yaml_path)
+        yaml_names, unbounded = read_prim_op_selection_from_yaml(
+            options.op_selection_yaml_path
         )
+        prim_op_names.extend(yaml_names)
 
-    write_selected_prim_ops(sorted(set(prim_op_names)), options.output_dir)
+    write_selected_prim_ops(
+        sorted(set(prim_op_names)),
+        options.output_dir,
+        unbounded_selection=unbounded,
+    )
 
 
 if __name__ == "__main__":

codegen/tools/targets.bzl

@@ -115,6 +115,22 @@ def define_common_targets(is_fbcode = False):
         _is_external_target = True,
     )
 
+    runtime.python_test(
+        name = "test_combine_prim_ops_headers",
+        srcs = [
+            "test/test_combine_prim_ops_headers.py",
+        ],
+        package_style = "inplace",
+        visibility = [
+            "PUBLIC",
+        ],
+        deps = [
+            ":combine_prim_ops_headers_lib",
+            ":gen_selected_prim_ops_lib",
+        ],
+        _is_external_target = True,
+    )
+
     runtime.python_test(
         name = "test_gen_all_oplist",
         srcs = [