Skip to content

Commit 8a3ef38

Browse files
ivayloenfacebook-github-bot
authored andcommitted
Preserve an empty selective primitive-operator set (#24815)
Summary: Fix `include_all_prim_ops=False` to preserve an empty selective primitive-operator set instead of registering every primitive. Derive selections from explicit names and YAML in one read, rejecting unenumerated overloads in selective builds. The default full-registry path is unchanged. AI-assisted. Reviewed By: rascani Differential Revision: D128664365
1 parent dd38e81 commit 8a3ef38

9 files changed

Lines changed: 630 additions & 72 deletions

File tree

‎codegen/tools/combine_prim_ops_headers.py‎

Lines changed: 46 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,12 @@
77
"""
88
Script to combine multiple selected_prim_ops.h header files into a single header.
99
This is used by selected_prim_operators_genrule to merge prim ops headers from dependencies.
10+
11+
The combiner only runs for a build that asked for a selected primitive-operator
12+
set, so the union it produces is that set, and an empty union means the build
13+
selected none of them. Two cases are errors instead, because neither is a
14+
selection this tool may silently narrow: a dependency that reports no selection
15+
at all, or one that does not enumerate required operator overloads.
1016
"""
1117

1218
import argparse
@@ -15,6 +21,12 @@
1521
from pathlib import Path
1622
from typing import List, Set
1723

24+
# Written into a dependency's selected_prim_ops.h by
25+
# executorch/codegen/tools/gen_selected_prim_ops.py when a dependency does not
26+
# enumerate every required operator overload. Kept as a literal so this tool needs
27+
# none of that one's dependencies; the agreement is pinned by its test.
28+
UNBOUNDED_SELECTION_MACRO = "EXECUTORCH_PRIM_OPS_SELECTION_UNBOUNDED"
29+
1830

1931
def read_header_file(file_path: Path) -> Set[str]:
2032
"""
@@ -25,23 +37,34 @@ def read_header_file(file_path: Path) -> Set[str]:
2537
2638
Returns:
2739
macros_set where macros_set contains unique macro defines
40+
41+
Raises:
42+
OSError: the header is missing or cannot be read. It carries this
43+
dependency's selection, so skipping it would silently narrow the
44+
combined set.
45+
ValueError: the dependency does not enumerate all required overloads,
46+
so no selected primitive-operator set can be derived from it.
2847
"""
2948
macros = set()
3049

31-
try:
32-
with open(file_path, "r") as f:
33-
for line in f:
34-
line = line.strip()
35-
36-
# Extract #define statements for prim ops
37-
if line.startswith("#define INCLUDE_") and not line.startswith(
38-
"#define EXECUTORCH_ENABLE"
39-
):
40-
macros.add(line)
41-
except FileNotFoundError:
42-
print(f"Warning: Header file not found: {file_path}", file=sys.stderr)
43-
except Exception as e:
44-
print(f"Error reading {file_path}: {e}", file=sys.stderr)
50+
with open(file_path, "r") as f:
51+
for line in f:
52+
line = line.strip()
53+
54+
if line == f"#define {UNBOUNDED_SELECTION_MACRO}":
55+
raise ValueError(
56+
f"{file_path} selects all operators or unenumerated operator "
57+
"overloads, so a selected primitive-operator set cannot be "
58+
"derived from it. List exact overloads or build "
59+
"the full primitive-operator registry instead, by passing "
60+
"include_all_prim_ops = True."
61+
)
62+
63+
# Extract #define statements for prim ops
64+
if line.startswith("#define INCLUDE_") and not line.startswith(
65+
"#define EXECUTORCH_ENABLE"
66+
):
67+
macros.add(line)
4568

4669
return macros
4770

@@ -55,20 +78,13 @@ def combine_prim_ops_headers(header_file_paths: List[str], output_path: str) ->
5578
output_path: Path to output the combined header
5679
"""
5780
all_macros = set()
58-
has_selective_build = False
5981

6082
# Read all header files and collect unique macros
6183
for header_file_path in header_file_paths:
6284
header_file = Path(header_file_path) / "selected_prim_ops.h"
63-
if os.path.exists(header_file):
64-
macros = read_header_file(header_file)
65-
all_macros.update(macros)
66-
if len(all_macros) > 0:
67-
has_selective_build = True
68-
else:
69-
print(
70-
f"Warning: Header file does not exist: {header_file}", file=sys.stderr
71-
)
85+
if not header_file.exists():
86+
raise FileNotFoundError(header_file)
87+
all_macros.update(read_header_file(header_file))
7288

7389
# Generate combined header
7490
header_content = [
@@ -78,28 +94,14 @@ def combine_prim_ops_headers(header_file_paths: List[str], output_path: str) ->
7894
"",
7995
"#pragma once",
8096
"",
97+
"// Enable selective build for prim ops",
98+
"#define EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD",
99+
"",
100+
"// Combined prim ops macros from all dependencies",
81101
]
82102

83-
if all_macros and has_selective_build:
84-
header_content.extend(
85-
[
86-
"// Enable selective build for prim ops",
87-
"#define EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD",
88-
"",
89-
"// Combined prim ops macros from all dependencies",
90-
]
91-
)
92-
93-
# Sort macros for deterministic output
94-
sorted_macros = sorted(all_macros)
95-
header_content.extend(sorted_macros)
96-
else:
97-
header_content.extend(
98-
[
99-
"// No prim ops found in dependencies - all prim ops will be included",
100-
"// Selective build is disabled",
101-
]
102-
)
103+
# Sort macros for deterministic output
104+
header_content.extend(sorted(all_macros))
103105

104106
header_content.append("")
105107

@@ -146,7 +148,6 @@ def main():
146148
)
147149

148150
args = parser.parse_args()
149-
import os
150151

151152
header_file_paths = _get_header_file_paths_from_query_output(args.header_files)
152153

‎codegen/tools/gen_oplist.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ def gen_oplist(
229229
or include_all_operators
230230
):
231231
# dump empty yaml file
232-
_dump_yaml([], output_path)
232+
_dump_yaml([], output_path, include_all_operators=include_all_operators)
233233
return
234234

235235
assert output_path, "Need to provide output_path for dumped yaml file."

‎codegen/tools/gen_selected_prim_ops.py‎

Lines changed: 73 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
# This source code is licensed under the BSD-style license found in the
66
# LICENSE file in the root directory of this source tree.
77

8-
# pyre-unsafe
9-
108
import argparse
119
import os
1210
import re
@@ -15,6 +13,7 @@
1513

1614
import yaml
1715
from torchgen.code_template import CodeTemplate # type: ignore[import-not-found]
16+
from torchgen.selective_build.operator import SelectiveBuildOperator
1817

1918

2019
selected_prim_ops_h_template_str = """#pragma once
@@ -47,32 +46,85 @@ def normalize_op_name(op_name: str) -> str:
4746
return normalized
4847

4948

50-
def read_prim_op_names_from_yaml(op_selection_yaml_path: str) -> List[str]:
49+
UNBOUNDED_SELECTION_MACRO = "EXECUTORCH_PRIM_OPS_SELECTION_UNBOUNDED"
50+
51+
UNBOUNDED_SELECTION_GUARD = f"""#define {UNBOUNDED_SELECTION_MACRO}
52+
#if defined(EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD)
53+
#error \\
54+
"A dependency selects all operators or unenumerated operator overloads. List exact overloads or build the full primitive-operator registry instead."
55+
#endif"""
56+
57+
58+
def _load_selection(op_selection_yaml_path: str) -> dict:
59+
"""
60+
Load a selected_operators.yaml.
61+
62+
An absent or malformed file is an error. Callers rely on the selection being
63+
the operators this build chose, so "no file" and "chose nothing" must not
64+
look alike.
65+
"""
66+
if not os.path.isfile(op_selection_yaml_path):
67+
raise FileNotFoundError(
68+
f"selected operator yaml not found: {op_selection_yaml_path}"
69+
)
70+
5171
with open(op_selection_yaml_path, "r") as f:
52-
selected_operators = yaml.safe_load(f) or {}
72+
selected_operators = yaml.safe_load(f)
5373

74+
if selected_operators is None:
75+
selected_operators = {}
76+
if not isinstance(selected_operators, dict):
77+
raise ValueError(
78+
f"selected operator yaml is not a mapping: {op_selection_yaml_path}"
79+
)
80+
return selected_operators
81+
82+
83+
def read_prim_op_selection_from_yaml(
84+
op_selection_yaml_path: str,
85+
) -> tuple[List[str], bool]:
86+
"""Read named operators and whether any required overloads are unenumerated."""
87+
selected_operators = _load_selection(op_selection_yaml_path)
5488
prim_op_names = set()
55-
operators = selected_operators.get("operators", {})
56-
if operators:
57-
prim_op_names.update(operators.keys())
89+
unbounded = bool(selected_operators.get("include_all_operators"))
90+
for op_name, metadata in (selected_operators.get("operators") or {}).items():
91+
operator = SelectiveBuildOperator.from_yaml_dict(op_name, metadata)
92+
if operator.include_all_overloads:
93+
unbounded = True
94+
else:
95+
prim_op_names.add(op_name)
5896

59-
et_kernel_metadata = selected_operators.get("et_kernel_metadata", {})
60-
if et_kernel_metadata:
61-
prim_op_names.update(et_kernel_metadata.keys())
97+
prim_op_names.update((selected_operators.get("et_kernel_metadata") or {}).keys())
98+
return sorted(prim_op_names), unbounded
6299

63-
return sorted(prim_op_names)
64100

101+
def selection_is_unbounded(op_selection_yaml_path: str) -> bool:
102+
"""Whether the selection omits required operator or overload names."""
103+
return read_prim_op_selection_from_yaml(op_selection_yaml_path)[1]
65104

66-
def write_selected_prim_ops(prim_op_names: List[str], output_dir: str) -> None:
105+
106+
def read_prim_op_names_from_yaml(op_selection_yaml_path: str) -> List[str]:
107+
"""Read enumerated names; use the selection reader to also check completeness."""
108+
return read_prim_op_selection_from_yaml(op_selection_yaml_path)[0]
109+
110+
111+
def write_selected_prim_ops(
112+
prim_op_names: List[str], output_dir: str, *, unbounded_selection: bool = False
113+
) -> None:
67114
"""
68115
Generate selected_prim_ops.h from a list of prim op names.
69116
70117
Args:
71118
prim_op_names: List of prim op names like ["executorch_prim::et_view.default", "aten::sym_size.int"]
72119
output_dir: Directory where to write selected_prim_ops.h
120+
unbounded_selection: Required operator overloads are not enumerated.
121+
The header rejects its use as a selected set instead of claiming
122+
that no primitive operators are needed.
73123
"""
74124
# Generate #define statements for each op
75125
defines = []
126+
if unbounded_selection:
127+
defines.append(UNBOUNDED_SELECTION_GUARD)
76128
for op_name in prim_op_names:
77129
macro_name = normalize_op_name(op_name)
78130
defines.append(f"#define {macro_name}")
@@ -116,17 +168,23 @@ def main(argv: List[Any]) -> None:
116168
parser.error("one of --prim-op-names or --op-selection-yaml-path is required")
117169

118170
prim_op_names: List[str] = []
171+
unbounded = False
119172
if options.prim_op_names is not None:
120173
# Parse comma-separated prim op names
121174
prim_op_names.extend(
122175
name.strip() for name in options.prim_op_names.split(",") if name.strip()
123176
)
124177
if options.op_selection_yaml_path:
125-
prim_op_names.extend(
126-
read_prim_op_names_from_yaml(options.op_selection_yaml_path)
178+
yaml_names, unbounded = read_prim_op_selection_from_yaml(
179+
options.op_selection_yaml_path
127180
)
181+
prim_op_names.extend(yaml_names)
128182

129-
write_selected_prim_ops(sorted(set(prim_op_names)), options.output_dir)
183+
write_selected_prim_ops(
184+
sorted(set(prim_op_names)),
185+
options.output_dir,
186+
unbounded_selection=unbounded,
187+
)
130188

131189

132190
if __name__ == "__main__":

‎codegen/tools/targets.bzl‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,22 @@ def define_common_targets(is_fbcode = False):
115115
_is_external_target = True,
116116
)
117117

118+
runtime.python_test(
119+
name = "test_combine_prim_ops_headers",
120+
srcs = [
121+
"test/test_combine_prim_ops_headers.py",
122+
],
123+
package_style = "inplace",
124+
visibility = [
125+
"PUBLIC",
126+
],
127+
deps = [
128+
":combine_prim_ops_headers_lib",
129+
":gen_selected_prim_ops_lib",
130+
],
131+
_is_external_target = True,
132+
)
133+
118134
runtime.python_test(
119135
name = "test_gen_all_oplist",
120136
srcs = [

0 commit comments

Comments
 (0)