|
5 | 5 | # This source code is licensed under the BSD-style license found in the |
6 | 6 | # LICENSE file in the root directory of this source tree. |
7 | 7 |
|
8 | | -# pyre-unsafe |
9 | | - |
10 | 8 | import argparse |
11 | 9 | import os |
12 | 10 | import re |
|
15 | 13 |
|
16 | 14 | import yaml |
17 | 15 | from torchgen.code_template import CodeTemplate # type: ignore[import-not-found] |
| 16 | +from torchgen.selective_build.operator import SelectiveBuildOperator |
18 | 17 |
|
19 | 18 |
|
20 | 19 | selected_prim_ops_h_template_str = """#pragma once |
@@ -47,32 +46,85 @@ def normalize_op_name(op_name: str) -> str: |
47 | 46 | return normalized |
48 | 47 |
|
49 | 48 |
|
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 | + |
51 | 71 | with open(op_selection_yaml_path, "r") as f: |
52 | | - selected_operators = yaml.safe_load(f) or {} |
| 72 | + selected_operators = yaml.safe_load(f) |
53 | 73 |
|
| 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) |
54 | 88 | 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) |
58 | 96 |
|
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 |
62 | 99 |
|
63 | | - return sorted(prim_op_names) |
64 | 100 |
|
| 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] |
65 | 104 |
|
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: |
67 | 114 | """ |
68 | 115 | Generate selected_prim_ops.h from a list of prim op names. |
69 | 116 |
|
70 | 117 | Args: |
71 | 118 | prim_op_names: List of prim op names like ["executorch_prim::et_view.default", "aten::sym_size.int"] |
72 | 119 | 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. |
73 | 123 | """ |
74 | 124 | # Generate #define statements for each op |
75 | 125 | defines = [] |
| 126 | + if unbounded_selection: |
| 127 | + defines.append(UNBOUNDED_SELECTION_GUARD) |
76 | 128 | for op_name in prim_op_names: |
77 | 129 | macro_name = normalize_op_name(op_name) |
78 | 130 | defines.append(f"#define {macro_name}") |
@@ -116,17 +168,23 @@ def main(argv: List[Any]) -> None: |
116 | 168 | parser.error("one of --prim-op-names or --op-selection-yaml-path is required") |
117 | 169 |
|
118 | 170 | prim_op_names: List[str] = [] |
| 171 | + unbounded = False |
119 | 172 | if options.prim_op_names is not None: |
120 | 173 | # Parse comma-separated prim op names |
121 | 174 | prim_op_names.extend( |
122 | 175 | name.strip() for name in options.prim_op_names.split(",") if name.strip() |
123 | 176 | ) |
124 | 177 | 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 |
127 | 180 | ) |
| 181 | + prim_op_names.extend(yaml_names) |
128 | 182 |
|
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 | + ) |
130 | 188 |
|
131 | 189 |
|
132 | 190 | if __name__ == "__main__": |
|
0 commit comments