def main() -> None: parser = argparse.ArgumentParser(description="Generate unboxing source files") parser.add_argument( "-s", "--source-path", help="path to source directory for ATen", default="aten/src/ATen", ) parser.add_argument( "-d", "--install_dir", help="output directory", default="build/aten/src/ATen" ) parser.add_argument( "-o", "--output-dependencies", help="output a list of dependencies into the given file and exit", ) parser.add_argument( "--dry-run", action="store_true", help="run without writing any files (still updates outputs)", ) parser.add_argument( "--op_selection_yaml_path", help="Provide a path to the operator selection (for custom build) YAML " "that contains the information about the set of selected operators " "and their categories (training, ...). Each operator is either a " "full operator name with overload or just a bare operator name. " "The operator names also contain the namespace prefix (e.g. aten::)", ) options = parser.parse_args() if options.op_selection_yaml_path is not None: selector = SelectiveBuilder.from_yaml_path(options.op_selection_yaml_path) else: selector = SelectiveBuilder.get_nop_selector() native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml") tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml") parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) native_functions, backend_indices = ( parsed_yaml.native_functions, parsed_yaml.backend_indices, ) cpu_fm = make_file_manager(options=options) gen_unboxing(native_functions=native_functions, cpu_fm=cpu_fm, selector=selector) if options.output_dependencies: depfile_path = pathlib.Path(options.output_dependencies).resolve() depfile_name = depfile_path.name depfile_stem = depfile_path.stem path = depfile_path.parent / depfile_name cpu_fm.write_outputs(depfile_stem, str(path))
def get_selector( selected_op_list_path: Optional[str], operators_yaml_path: Optional[str], ) -> Any: # cwrap depends on pyyaml, so we can't import it earlier root = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, root) from torchgen.selective_build.selector import SelectiveBuilder assert not (selected_op_list_path is not None and operators_yaml_path is not None), ( "Expected at most one of selected_op_list_path and " + "operators_yaml_path to be set.") if selected_op_list_path is None and operators_yaml_path is None: return SelectiveBuilder.get_nop_selector() elif selected_op_list_path is not None: return get_selector_from_legacy_operator_selection_list( selected_op_list_path) else: return SelectiveBuilder.from_yaml_path(cast(str, operators_yaml_path))