def get_selector_from_legacy_operator_selection_list( selected_op_list_path: str, ): from tools.autograd.utils import load_op_list_and_strip_overload selected_op_list = load_op_list_and_strip_overload( None, selected_op_list_path, ) # Internal build doesn't use this flag any more. Only used by OSS # build now. Every operator should be considered a root operator # (hence generating unboxing code for it, which is consistent with # the current behaviour), and also be considered as used for # training, since OSS doesn't support training on mobile for now. # is_root_operator = True is_used_for_training = True from tools.codegen.selective_build.selector import SelectiveBuilder selector: SelectiveBuilder = SelectiveBuilder.get_nop_selector() if selected_op_list is not None: selector = SelectiveBuilder.from_legacy_op_registration_allow_list( selected_op_list, is_root_operator, is_used_for_training, ) return selector
def get_custom_build_selector( provided_op_registration_allowlist: Optional[List[str]], op_selection_yaml_path: Optional[str]) -> SelectiveBuilder: assert not (provided_op_registration_allowlist is not None and op_selection_yaml_path is not None), ( "Both provided_op_registration_allowlist and " + "op_selection_yaml_path can NOT be provided at the " + "same time.") op_registration_allowlist: Optional[Set[str]] = None if provided_op_registration_allowlist is not None: op_registration_allowlist = set(provided_op_registration_allowlist) if op_registration_allowlist is not None: selector = SelectiveBuilder.from_legacy_op_registration_allow_list( op_registration_allowlist, True, False, ) elif op_selection_yaml_path is not None: selector = SelectiveBuilder.from_yaml_path(op_selection_yaml_path) else: selector = SelectiveBuilder.get_nop_selector() return selector
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") parsed_yaml = parse_native_yaml(native_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 run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str] = None) -> None: # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join( pytorch_root, 'aten/src/ATen/native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace class_name = parsed_backend_yaml.class_name backend_indices = parsed_backend_yaml.backend_indices selector = SelectiveBuilder.get_nop_selector() if backend_key is None: # This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet. return if class_name is None: # class_name is an optional argument to backend yaml file. # if specified it allows an external backend to override # the name of the class that all generated kernel definitions live under. # if not specified, its value is given as native_function_class_name. class_name = backend_indices[backend_key].native_function_class_name() assert class_name is not None if impl_path is not None: error_on_missing_kernels(native_functions, backend_indices, backend_key, autograd_key, class_name, impl_path) gen_dispatchkey_nativefunc_headers(fm, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, autograd_key) for dispatch_key in [backend_key] if autograd_key is None else [ backend_key, autograd_key ]: gen_dispatcher_registrations(fm, output_dir, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, dispatch_key, selector)
def computeSchemaRegister() -> Dict[str, object]: schema_registrations = list(mapMaybe( compute_type_method(None, target=Target.REGISTRATION, selector=SelectiveBuilder.get_nop_selector(), def_only=True), native_functions)) return { 'schema_registrations': schema_registrations, }
def get_selector_from_legacy_operator_selection_list( selected_op_list_path: str, ) -> Any: with open(selected_op_list_path, 'r') as f: # strip out the overload part # It's only for legacy config - do NOT copy this code! selected_op_list = { opname.split('.', 1)[0] for opname in yaml.load(f, Loader=YamlLoader) } # Internal build doesn't use this flag any more. Only used by OSS # build now. Every operator should be considered a root operator # (hence generating unboxing code for it, which is consistent with # the current behaviour), and also be considered as used for # training, since OSS doesn't support training on mobile for now. # is_root_operator = True is_used_for_training = True from tools.codegen.selective_build.selector import SelectiveBuilder selector = SelectiveBuilder.from_legacy_op_registration_allow_list( selected_op_list, is_root_operator, is_used_for_training, ) return selector
def main() -> None: parser = argparse.ArgumentParser(description='Generate backend stub files') parser.add_argument( '-s', '--source_yaml', help='path to source yaml file containing operator external definitions') parser.add_argument( '-o', '--output_dir', help='output directory') parser.add_argument( '--dry_run', type=bool, default=False, help='output directory') options = parser.parse_args() # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.dry_run) fm = make_file_manager(options.output_dir) native_yaml_path = os.path.join(pytorch_root, 'aten/src/ATen/native/native_functions.yaml') grouped_native_functions = get_grouped_native_functions(native_yaml_path) cpp_namespace, external_backend_functions = parse_backend_yaml(options.source_yaml, grouped_native_functions) native_functions = parse_native_yaml(native_yaml_path) selector = SelectiveBuilder.get_nop_selector() generated_comment = 'Autogenerated file by gen_backend_stubs.py. Do not edit directly!' fm.write('aten_xla_type.h', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, 'dispatch_xla_declarations': list(concatMap(dest.compute_native_function_declaration, external_backend_functions)), }) fm.write('aten_xla_type_default.h', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, 'dispatch_aten_fallback_declarations': list(concatMap( dest.GenExternalAtenFallback(Target.NAMESPACED_DECLARATION), external_backend_functions )), }) fm.write('aten_xla_type_default.cpp', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, # TODO: after cpu fallbacks are moved to a boxed kernel, # merge registrations / definitions into RegisterDispatchKey 'dispatch_aten_fallback_definitions': list(concatMap( dest.GenExternalAtenFallback(Target.NAMESPACED_DEFINITION), external_backend_functions )), 'dispatch_registrations': list(concatMap( dest.GenExternalAtenFallback(Target.REGISTRATION), [e for e in external_backend_functions if not e.is_autograd_kernel] )), 'dispatch_autograd_registrations': list(concatMap( dest.GenExternalAtenFallback(Target.REGISTRATION), [e for e in external_backend_functions if e.is_autograd_kernel] )), })
def generate_code(ninja_global=None, declarations_path=None, nn_path=None, native_functions_path=None, install_dir=None, subset=None, disable_autograd=False, force_schema_registration=False, operator_selector=None): from tools.autograd.gen_autograd import gen_autograd, gen_autograd_python from tools.autograd.gen_annotated_fn_args import gen_annotated from tools.jit.gen_unboxing_wrappers import gen_unboxing_wrappers from tools.codegen.selective_build.selector import SelectiveBuilder # Build ATen based Variable classes if install_dir is None: install_dir = 'torch/csrc' python_install_dir = 'torch/testing/_internal/generated' else: python_install_dir = install_dir autograd_gen_dir = os.path.join(install_dir, 'autograd', 'generated') jit_gen_dir = os.path.join(install_dir, 'jit', 'generated') for d in (autograd_gen_dir, jit_gen_dir, python_install_dir): if not os.path.exists(d): os.makedirs(d) runfiles_dir = os.environ.get("RUNFILES_DIR", None) data_dir = os.path.join(runfiles_dir, 'pytorch') if runfiles_dir else '' autograd_dir = os.path.join(data_dir, 'tools', 'autograd') tools_jit_templates = os.path.join(data_dir, 'tools', 'jit', 'templates') if subset == "pybindings" or not subset: gen_autograd_python(declarations_path or DECLARATIONS_PATH, native_functions_path or NATIVE_FUNCTIONS_PATH, autograd_gen_dir, autograd_dir) if operator_selector is None: operator_selector = SelectiveBuilder.get_nop_selector() if subset == "libtorch" or not subset: gen_autograd( declarations_path or DECLARATIONS_PATH, native_functions_path or NATIVE_FUNCTIONS_PATH, autograd_gen_dir, autograd_dir, disable_autograd=disable_autograd, operator_selector=operator_selector, ) gen_unboxing_wrappers( declarations_path or DECLARATIONS_PATH, jit_gen_dir, tools_jit_templates, disable_autograd=disable_autograd, operator_selector=operator_selector, force_schema_registration=force_schema_registration) if subset == "python" or not subset: gen_annotated(declarations_path or DECLARATIONS_PATH, python_install_dir, autograd_dir)
def get_selector(selected_op_list_path, operators_yaml_path): # 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 tools.codegen.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(operators_yaml_path)
def main(): parser = argparse.ArgumentParser(description='Generate JIT op dispatch') parser.add_argument('declarations', metavar='DECL', help='path to Declarations.yaml') parser.add_argument('out', metavar='OUT', help='path to output directory') parser.add_argument('template_path', metavar='TEMPLATE_PATH', help='path to templates directory') args = parser.parse_args() gen_unboxing_wrappers(args.declarations, args.out, args.template_path, SelectiveBuilder.get_nop_selector())
def main() -> None: parser = argparse.ArgumentParser( description='Generate autograd C++ files script') parser.add_argument('native_functions', metavar='NATIVE', help='path to native_functions.yaml') parser.add_argument('out', metavar='OUT', help='path to output directory') parser.add_argument('autograd', metavar='AUTOGRAD', help='path to autograd directory') args = parser.parse_args() gen_autograd(args.native_functions, args.out, args.autograd, SelectiveBuilder.get_nop_selector())
def write_selected_mobile_ops_with_all_dtypes( output_file_path: str, root_ops: Set[str], ) -> None: with open(output_file_path, "wb") as out_file: body_parts = [selected_mobile_ops_preamble] body_parts.append("#define TORCH_OPERATOR_WHITELIST " + (";".join(sorted(root_ops))) + ";\n\n") selective_builder = SelectiveBuilder.get_nop_selector() body_parts.append(get_selected_kernel_dtypes_code(selective_builder)) header_contents = "".join(body_parts) out_file.write(header_contents.encode("utf-8"))
def filter_decls(jit_decls, disable_autograd, operator_selector: SelectiveBuilder, force_schema_registration): result = [] for decl in jit_decls: if disable_autograd and is_backward_op(decl): continue op_name = op_name_with_overload(decl) if operator_selector.is_root_operator(op_name): result.append(decl) else: if force_schema_registration: decl['emit_dummy_placeholder'] = True result.append(decl) return result
def run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str]) -> None: # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join( pytorch_root, 'aten/src/ATen/native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace backend_indices = parsed_backend_yaml.backend_indices selector = SelectiveBuilder.get_nop_selector() assert backend_key is not None class_name = backend_indices[backend_key].native_function_class_name() if impl_path is not None: error_on_missing_kernels(native_functions, backend_indices, backend_key, autograd_key, impl_path) gen_dispatchkey_nativefunc_headers(fm, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, autograd_key) for dispatch_key in [backend_key] if autograd_key is None else [ backend_key, autograd_key ]: gen_dispatcher_registrations(fm, output_dir, cpp_namespace, backend_indices, grouped_native_functions, backend_key, dispatch_key, selector)
def run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str], gen_ts_lowerings: bool, node_base: str, node_base_hdr: Optional[str], tensor_class: str, tensor_class_hdr: str) -> None: # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join( pytorch_root, 'aten/src/ATen/native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) def sort_native_function( f: Union[NativeFunctionsGroup, NativeFunction]) -> str: """ We sort the native function because of the note in concat_map_codegen. TODO(alanwaketan): Remove this sorting hack once all ops are grouped properly. """ func = f.functional.func if isinstance( f, NativeFunctionsGroup) else f.func return str(func.name.name) grouped_native_functions = sorted(grouped_native_functions, key=sort_native_function) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace backend_indices = parsed_backend_yaml.backend_indices full_codegen = parse_full_codegen_ops(source_yaml, grouped_native_functions) def concat_map_codegen( func: Callable[[NativeFunction], Sequence[str]], xs: Iterable[Union[NativeFunctionsGroup, NativeFunction]], *, codegenInplaceVariant: bool = False) -> Iterator[str]: """ We code-gen for the functional variant, which is all we need for IR classes/lowerings/shape inferences, but we only code-gen additional entries for the inplace variant for the native functions. Note: If xs is not sorted, there may be an edge case when generating IR classes. Considering relu and relu_, if we encounter relu_ before relu. we will then generate an IR class with op = at::aten::relu_ for both relu and relu_ which will cause problems for relu. TODO(alanwaketan): Once all ops are grouped properly, we should no longer need this hack. """ generated = set() def gen_key(func: FunctionSchema) -> Tuple[str, str]: # we want to generate unique entries for overloads of functional variants, # but not for inplace variants unless explicitly told `codegenInplaceVariant` return (func.name.name.base, func.name.overload_name) for x in xs: f = x.functional if isinstance(x, NativeFunctionsGroup) else x # For the 'or'd terms: # 1. codegenInplaceVariant means we can generate the in-place variant corresponding items. # 2. not f.func.name.name.inplace means the op is not a in-place variant, so we can generate the item. # 3. f.func.name.name.base not in generated means even for in-place ops we still need to generate the item # as if they were the functional variants for one time. if f.func.name in full_codegen and \ (codegenInplaceVariant or not f.func.name.name.inplace or gen_key(f.func) not in generated): generated.add(gen_key(f.func)) for r in func(f): yield r selector = SelectiveBuilder.get_nop_selector() assert backend_key is not None class_name = backend_indices[backend_key].native_function_class_name() if impl_path is not None: error_on_missing_kernels(native_functions, backend_indices, backend_key, autograd_key, impl_path, full_codegen) assert class_name is not None # Generate nativefunction declarations gen_dispatchkey_nativefunc_headers(fm, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, autograd_key) # Generate Dispatcher registrations which hook up the nativefunctions for dispatch_key in [backend_key] if autograd_key is None else [ backend_key, autograd_key ]: gen_dispatcher_registrations(fm, output_dir, cpp_namespace, backend_indices, grouped_native_functions, backend_key, dispatch_key, selector) # Generate native function impls that build IR nodes fm.write_with_template( f'{backend_key}NativeFunctions.cpp', 'DispatchKeyNativeFunctions.cpp', lambda: { 'includes': [ f'#include <{path}>' for path in [ tensor_class_hdr, "ATen/MetaFunctions.h", "torch/csrc/lazy/core/metrics.h", "torch/csrc/lazy/core/shape.h", "lazy_tensor_core/csrc/aten_ltc_bridge.h", "lazy_tensor_core/csrc/lazy_graph_executor.h", f"{output_dir}/{backend_key}NativeFunctions.h", f"{output_dir}/{backend_key}LazyIr.h", f"{output_dir}/{backend_key}ShapeInference.h", ] ], 'native_functions_include': '', 'backend_namespace': 'torch_lazy_tensors', # this is wrong 'native_function_definitions': list( concat_map_codegen(dest.GenLazyNativeFuncDefinition( f'{backend_key}NativeFunctions', backend_indices[ backend_key], tensor_class), grouped_native_functions, codegenInplaceVariant=True)), }) # Generate headers for shape/dtype funcs for non-meta kernels fm.write_with_template( f'{backend_key}ShapeInference.h', 'ShapeInference.h', lambda: { 'lazy_ir_sysinc': [ f'#include <{path}>' for path in [ "ATen/Tensor.h", "c10/core/ScalarType.h", "c10/util/Optional.h", "torch/csrc/lazy/core/ir.h", "torch/csrc/lazy/core/shape.h", "vector", ] ], 'lazy_ir_inc': [], 'DispatchKey': backend_key, 'dispatch_namespace': backend_key.lower(), 'func_declarations': list( concat_map_codegen( dest.GenLazyShapeInferenceDefinition( backend_indices[backend_key], tensor_class), grouped_native_functions)), }) # Generate IR node classes fm.write_with_template( f'{backend_key}LazyIr.h', 'LazyIr.h', lambda: { 'lazy_ir_sysinc': [ f'#include <{path}>' for path in [ "ATen/core/Formatting.h", "c10/core/ScalarType.h", "c10/util/Optional.h", "torch/csrc/lazy/core/hash.h", "torch/csrc/lazy/core/ir.h", "vector", ] ], 'lazy_ir_inc': [ f'#include "{path}"' for path in [node_base_hdr if node_base_hdr is not None else None] if path is not None ], 'external_backend_headers': f'#include "{output_dir}/{backend_key}NativeFunctions.h"', 'namespaced_headers': '', 'DispatchKey': backend_key, 'dispatch_namespace': backend_key.lower(), 'ir_declarations': list( concat_map_codegen( dest.LazyIR(backend_indices[backend_key], node_base), grouped_native_functions)), })
def main() -> None: parser = argparse.ArgumentParser(description='Generate ATen source files') parser.add_argument('-s', '--source-path', help='path to source directory for ATen', default='aten/src/ATen') parser.add_argument( '-o', '--output-dependencies', help='output a list of dependencies into the given file and exit') parser.add_argument('-d', '--install_dir', help='output directory', default='build/aten/src/ATen') parser.add_argument( '--rocm', action='store_true', help='reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly') # TODO: --op_registration_whitelist will be removed when all call-sites # for gen.py are moved over to using the operator YAML file for mobile # custom build. parser.add_argument( '--op_registration_whitelist', nargs='*', help='filter op registrations by the whitelist (if set); ' 'each item is `namespace`::`operator name` without overload name; ' 'e.g.: aten::empty aten::conv2d ...') 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::)') parser.add_argument( '--backend_whitelist', nargs='*', help='filter dispatch backend by the whitelist (if set), ' 'e.g.: CPU CUDA QuantizedCPU ...') parser.add_argument( '--static_dispatch_backend', help='generate static dispatch code for the specific backend (if set)') parser.add_argument( '--force_schema_registration', action='store_true', help= 'force it to generate schema-only registrations for all ops, including' 'those that are not listed on --op_registration_whitelist') options = parser.parse_args() selector = get_custom_build_selector( options.op_registration_whitelist, options.op_selection_yaml_path, ) native_functions = parse_native_yaml( os.path.join(options.source_path, 'native/native_functions.yaml')) pre_grouped_native_functions: Dict[FunctionSchema, Dict[SchemaKind, NativeFunction]] pre_grouped_native_functions = defaultdict(dict) for f in native_functions: d = pre_grouped_native_functions[f.func.signature()] assert f.func.kind() not in d d[f.func.kind()] = f def flatten_pre_group( d: Dict[SchemaKind, NativeFunction] ) -> Sequence[Union[NativeFunction, NativeFunctionsGroup]]: r = NativeFunctionsGroup.from_dict(d) if r is None: return list(d.values()) else: return [r] # TODO: how come ValuesView isn't a Sequence lol grouped_native_functions = list( concatMap(flatten_pre_group, list(pre_grouped_native_functions.values()))) structured_native_functions = [ g for g in grouped_native_functions if isinstance(g, NativeFunctionsGroup) ] template_dir = os.path.join(options.source_path, "templates") # NB: It is mandatory to NOT use os.path.join here, as the install directory # will eventually be ingested by cmake, which does not respect Windows style # path slashes. If you switch this to use os.path.join, you'll get an error # like: # # Syntax error in cmake code when parsing string # # C:/Jenkins/workspace/pytorch-builds/pytorch-win-ws2016-cuda9-cudnn7-py3-build/build/aten/src/ATen\core/TensorMethods.h # # Invalid character escape '\c'. core_install_dir = f'{options.install_dir}/core' pathlib.Path(core_install_dir).mkdir(parents=True, exist_ok=True) def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.output_dependencies) core_fm = make_file_manager(core_install_dir) cpu_fm = make_file_manager(options.install_dir) cuda_fm = make_file_manager(options.install_dir) extra_cuda_headers = '''\ #include <c10/cuda/CUDAGuard.h> #include <ATen/cuda/ATenCUDAGeneral.h> #include <ATen/cuda/CUDADevice.h> #include <ATen/cuda/CUDAContext.h>''' if options.rocm: extra_cuda_headers = '''\ #include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h> #include <ATen/hip/ATenHIPGeneral.h> #include <ATen/hip/HIPDevice.h> #include <ATen/hip/HIPContext.h>''' dispatch_keys = [ DispatchKey.CPU, DispatchKey.SparseCPU, DispatchKey.SparseCsrCPU, DispatchKey.MkldnnCPU, DispatchKey.CUDA, DispatchKey.SparseCUDA, DispatchKey.SparseCsrCUDA, DispatchKey.QuantizedCPU, DispatchKey.QuantizedCUDA, DispatchKey.CompositeImplicitAutograd, DispatchKey.CompositeExplicitAutograd, # Meta is a magic key: it is automatically generated for structured # kernels DispatchKey.Meta, ] # Only a limited set of dispatch keys get CPUFunctions.h headers generated # for them; this is the set functions_keys = { DispatchKey.CPU, DispatchKey.CUDA, DispatchKey.CompositeImplicitAutograd, DispatchKey.CompositeExplicitAutograd, } if options.backend_whitelist: dispatch_keys = [ k for k in dispatch_keys if is_generic_dispatch_key(k) or str(k) in options.backend_whitelist ] static_dispatch_backend: Optional[DispatchKey] = None if options.static_dispatch_backend: static_dispatch_backend = DispatchKey.parse( options.static_dispatch_backend) for dispatch_key in dispatch_keys: fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm fm.write_with_template( f'Register{dispatch_key}.cpp', 'RegisterDispatchKey.cpp', lambda: { 'extra_cuda_headers': extra_cuda_headers if is_cuda_dispatch_key(dispatch_key) else '', 'legacy_th_headers': '#include <ATen/LegacyTHFunctionsCPU.h>' if dispatch_key == DispatchKey.CPU else '#include <ATen/LegacyTHFunctionsCUDA.h>' if dispatch_key == DispatchKey.CUDA else '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_definitions': list( concatMap( dest.RegisterDispatchKey(dispatch_key, Target.NAMESPACED_DEFINITION, selector, rocm=options.rocm), grouped_native_functions)), 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey(dispatch_key, Target.ANONYMOUS_DEFINITION, selector, rocm=options.rocm), grouped_native_functions)), 'dispatch_registrations': list( concatMap( dest.RegisterDispatchKey(dispatch_key, Target.REGISTRATION, selector, rocm=options.rocm), grouped_native_functions)), }) if dispatch_key in functions_keys: fm.write_with_template( f'{dispatch_key}Functions.h', 'DispatchKeyFunctions.h', lambda: { 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_declarations': list( concatMap( dest.RegisterDispatchKey( dispatch_key, Target.NAMESPACED_DECLARATION, selector, rocm=options.rocm), grouped_native_functions)), }) del fm # BackendSelect is generated specially cpu_fm.write( 'RegisterBackendSelect.cpp', lambda: { 'backend_select_method_definitions': list( mapMaybe(ComputeBackendSelect(Target.DEFINITION), native_functions)), 'backend_select_function_registrations': list( mapMaybe(ComputeBackendSelect(Target.REGISTRATION), native_functions)), }) cpu_fm.write( 'MetaFunctions.h', lambda: { 'declarations': list( mapMaybe(compute_meta_function_declaration, structured_native_functions)), }) schema_selector = selector if options.force_schema_registration: schema_selector = SelectiveBuilder.get_nop_selector() cpu_fm.write( 'RegisterSchema.cpp', lambda: { 'schema_registrations': list(mapMaybe(RegisterSchema(schema_selector), native_functions)), }) cpu_fm.write( 'Functions.h', lambda: { 'function_declarations': list( mapMaybe( ComputeFunction( Target.DECLARATION, static_dispatch_backend=static_dispatch_backend, is_redispatching_fn=False), native_functions)), }) cpu_fm.write( 'Functions.cpp', lambda: { 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_backend), 'function_definitions': list( mapMaybe( ComputeFunction( Target.DEFINITION, static_dispatch_backend=static_dispatch_backend, is_redispatching_fn=False), native_functions)), }) cpu_fm.write( 'RedispatchFunctions.h', lambda: { 'function_redispatch_declarations': list( mapMaybe( ComputeFunction( Target.DECLARATION, static_dispatch_backend=static_dispatch_backend, is_redispatching_fn=True), native_functions)), }) cpu_fm.write( 'RedispatchFunctions.cpp', lambda: { 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_backend), 'function_redispatch_definitions': list( mapMaybe( ComputeFunction( Target.DEFINITION, static_dispatch_backend=static_dispatch_backend, is_redispatching_fn=True), native_functions)), }) core_fm.write( 'TensorBody.h', lambda: { 'tensor_method_declarations': list( mapMaybe( ComputeTensorMethod(Target.DECLARATION, static_dispatch_backend= static_dispatch_backend), native_functions)), }) core_fm.write( 'TensorMethods.cpp', lambda: { 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_backend), 'tensor_method_definitions': list( mapMaybe( ComputeTensorMethod(Target.DEFINITION, static_dispatch_backend= static_dispatch_backend), native_functions)), }) core_fm.write( 'ATenOpList.cpp', lambda: { 'aten_ops': list(mapMaybe(compute_aten_op, native_functions)), }) cpu_fm.write( 'NativeFunctions.h', lambda: { 'native_function_declarations': list( concatMap(dest.compute_native_function_declaration, grouped_native_functions)), }) cpu_fm.write( 'Declarations.yaml', lambda: format_yaml( [compute_declaration_yaml(f) for f in native_functions])) cpu_fm.write( 'RegistrationDeclarations.h', lambda: { 'registration_declarations': [compute_registration_declarations(f) for f in native_functions], }) if options.output_dependencies: cpu_fm.write_outputs(options.output_dependencies) core_fm.write_outputs(f"{options.output_dependencies}-core") cuda_fm.write_outputs(f"{options.output_dependencies}-cuda")
def run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str]) -> None: # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join( pytorch_root, 'aten/src/ATen/native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace backend_indices = parsed_backend_yaml.backend_indices selector = SelectiveBuilder.get_nop_selector() # TODO: handle cases when yaml contains zero ops properly in a later PR. if backend_key is not None and autograd_key is not None: backend_dispatch_key: DispatchKey = backend_key autograd_dispatch_key: DispatchKey = autograd_key class_name = backend_indices[ backend_dispatch_key].native_function_class_name() if impl_path is not None: error_on_missing_kernels(native_functions, backend_indices, backend_key, autograd_key, impl_path) assert class_name is not None generated_comment = 'Autogenerated file by gen_backend_stubs.py. Do not edit directly!' fm.write_with_template( f'{backend_dispatch_key}NativeFunctions.h', 'DispatchKeyNativeFunctions.h', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, 'class_name': class_name, # Convert to a set first to remove duplicate kernel names. # Backends are allowed to repeat kernel names; only generate the declaration once! 'dispatch_declarations': list( set( concatMap( lambda f: dest.compute_native_function_declaration( f, backend_indices[backend_dispatch_key]), grouped_native_functions))) + list( set( concatMap( lambda f: dest.compute_native_function_declaration( f, backend_indices[autograd_dispatch_key]), grouped_native_functions))), }) for dispatch_key in [backend_dispatch_key, autograd_dispatch_key]: fm.write_with_template( f'Register{dispatch_key}.cpp', 'RegisterDispatchKey.cpp', lambda: { 'extra_cuda_headers': '', 'external_backend_headers': f'#include "{output_dir}/{backend_key}NativeFunctions.h"', 'namespaced_headers': '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_helpers': dest.gen_registration_helpers(backend_indices[dispatch_key] ), 'dispatch_namespaced_definitions': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.NAMESPACED_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name= f'{backend_dispatch_key}NativeFunctions'), grouped_native_functions)), 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.ANONYMOUS_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name= f'{backend_dispatch_key}NativeFunctions'), grouped_native_functions)), 'dispatch_registrations': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.REGISTRATION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name= f'{backend_dispatch_key}NativeFunctions'), grouped_native_functions)), })
def run(source_yaml: str, output_dir: str, dry_run: bool) -> None: # Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join( pytorch_root, 'aten/src/ATen/native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace backend_indices = parsed_backend_yaml.backend_indices selector = SelectiveBuilder.get_nop_selector() # TODO: handle cases when yaml contains zero ops properly in a later PR. if backend_key is not None and autograd_key is not None: backend_dispatch_key: DispatchKey = backend_key autograd_dispatch_key: DispatchKey = autograd_key generated_comment = 'Autogenerated file by gen_backend_stubs.py. Do not edit directly!' fm.write( 'aten_xla_type.h', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, # Convert to a set first to remove duplicate kernel names. # Backends are allowed to repeat kernel names; only generate the declaration once! 'dispatch_xla_declarations': list( set( concatMap( lambda f: dest.compute_native_function_declaration( f, backend_indices[backend_dispatch_key]), grouped_native_functions))) + list( set( concatMap( lambda f: dest.compute_native_function_declaration( f, backend_indices[autograd_dispatch_key]), grouped_native_functions))), }) external_backend_headers = '''\ #include <tensorflow/compiler/xla/xla_client/debug_macros.h> #include <tensorflow/compiler/xla/xla_client/metrics.h> #include <tensorflow/compiler/xla/xla_client/tf_logging.h> #include <torch_xla/csrc/function_call_tracker.h> #include <torch_xla/csrc/aten_xla_type.h> #include <torch_xla/csrc/aten_xla_type_default.h>''' for dispatch_key in [backend_dispatch_key, autograd_dispatch_key]: fm.write_with_template( f'Register{dispatch_key}.cpp', 'RegisterDispatchKey.cpp', lambda: { 'extra_cuda_headers': '', 'legacy_th_headers': '', 'external_backend_headers': external_backend_headers, 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_definitions': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.NAMESPACED_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace), grouped_native_functions)), 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.ANONYMOUS_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace), grouped_native_functions)), 'dispatch_registrations': list( concatMap( dest.RegisterDispatchKey( backend_indices[dispatch_key], Target.REGISTRATION, selector, rocm=False, cpp_namespace=cpp_namespace), grouped_native_functions)), }) fm.write( 'aten_xla_type_default.h', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, 'dispatch_aten_fallback_declarations': list( concatMap( dest.GenExternalAtenFallback( Target.NAMESPACED_DECLARATION, backend_indices[ backend_dispatch_key]), grouped_native_functions)), }) fm.write( 'aten_xla_type_default.cpp', lambda: { 'generated_comment': generated_comment, 'cpp_namespace': cpp_namespace, # TODO: after cpu fallbacks are moved to a boxed kernel, # merge registrations / definitions into RegisterDispatchKey 'dispatch_aten_fallback_definitions': list( concatMap( dest.GenExternalAtenFallback( Target.NAMESPACED_DEFINITION, backend_indices[ backend_dispatch_key]), grouped_native_functions)), 'dispatch_registrations': list( concatMap( dest.GenExternalAtenFallback( Target.REGISTRATION, backend_indices[ backend_dispatch_key]), grouped_native_functions)), })
def run_gen_lazy_tensor( aten_path: str, source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str], node_base: str = default_args.node_base, node_base_hdr: Optional[str] = default_args.node_base_hdr, tensor_class: str = default_args.tensor_class, tensor_class_hdr: str = default_args.tensor_class_hdr, shape_inference_hdr: str = default_args.shape_inference_hdr, lazy_ir_cls: Type[LazyIR] = default_args.lazy_ir_cls, # build_in_tree is true for TS backend and affects include paths build_in_tree: bool = False, # per_operator_headers changes whether ATen/Functions.h or individual operator headers are used # it must match how ATen was built per_operator_headers: bool = False, backend_name: str = default_args.backend_name, gen_forced_fallback_code: bool = False) -> None: template_dir = os.path.join(aten_path, "templates") def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=dry_run) fm = make_file_manager(output_dir) native_yaml_path = os.path.join(aten_path, 'native/native_functions.yaml') parsed_yaml = parse_native_yaml(native_yaml_path) native_functions, backend_indices = parsed_yaml.native_functions, parsed_yaml.backend_indices grouped_native_functions = get_grouped_native_functions(native_functions) def sort_native_function( f: Union[NativeFunctionsGroup, NativeFunction]) -> str: """ We sort the native function because of the note in concat_map_codegen. TODO(alanwaketan): Remove this sorting hack once all ops are grouped properly. """ func = f.functional.func if isinstance( f, NativeFunctionsGroup) else f.func return str(func.name.name) grouped_native_functions = sorted(grouped_native_functions, key=sort_native_function) parsed_backend_yaml = parse_backend_yaml(source_yaml, grouped_native_functions, backend_indices) backend_key = parsed_backend_yaml.backend_key autograd_key = parsed_backend_yaml.autograd_key cpp_namespace = parsed_backend_yaml.cpp_namespace backend_indices = parsed_backend_yaml.backend_indices full_codegen = parse_full_codegen_ops(source_yaml, grouped_native_functions) def concat_map_codegen( func: Callable[[NativeFunction], Sequence[str]], xs: Iterable[Union[NativeFunctionsGroup, NativeFunction]], *, codegenInplaceVariant: bool = False) -> Iterator[str]: """ We code-gen for the functional variant, which is all we need for IR classes/lowerings/shape inferences, but we only code-gen additional entries for the inplace variant for the native functions. Note: If xs is not sorted, there may be an edge case when generating IR classes. Considering relu and relu_, if we encounter relu_ before relu. we will then generate an IR class with op = at::aten::relu_ for both relu and relu_ which will cause problems for relu. TODO(alanwaketan): Once all ops are grouped properly, we should no longer need this hack. """ generated = set() def gen_key(func: FunctionSchema) -> Tuple[str, str]: # we want to generate unique entries for overloads of functional variants, # but not for inplace variants unless explicitly told `codegenInplaceVariant` return (func.name.name.base, func.name.overload_name) for x in xs: f = x.functional if isinstance(x, NativeFunctionsGroup) else x # For the 'or'd terms: # 1. codegenInplaceVariant means we can generate the in-place variant corresponding items. # 2. not f.func.name.name.inplace means the op is not a in-place variant, so we can generate the item. # 3. f.func.name.name.base not in generated means even for in-place ops we still need to generate the item # as if they were the functional variants for one time. if f.func.name in full_codegen and \ (codegenInplaceVariant or not f.func.name.name.inplace or gen_key(f.func) not in generated): generated.add(gen_key(f.func)) for r in func(f): yield r selector = SelectiveBuilder.get_nop_selector() assert backend_key is not None class_name = backend_indices[backend_key].native_function_class_name() if impl_path is not None: error_on_missing_kernels(native_functions, backend_indices, backend_key, autograd_key, class_name, impl_path, full_codegen) """ Validate Shape Inference Definitions Generated lazy native functions all perform shape inference, by first using a meta:: kernel if available for that op, and otherwise using a 'compute_shape_{op}' function instead. The generator knows the call signature for compute_shape_{op} becuase it matches the nativefunction (and meta::) signature, so it just has to check whether the op is structured and generate a call for one or the other. It's up to the dev to supply the missing compute_shape_{op} function, but the codegen at least warns you about this and provides the expected signature which can be copy-pasted into shape_inference.h. compute_shape_{op} functions are handwritten and should be replaced over time as ops get ported to structured kernels. See torch/csrc/lazy/core/shape_inference.cpp #READ THIS! for more information. """ if shape_inference_hdr is not None: expected_shape_infr_decls = list( concat_map_codegen(dest.GenLazyShapeInferenceDefinition( backend_indices[backend_key], tensor_class), grouped_native_functions, codegenInplaceVariant=True)) validate_shape_inference_header(shape_inference_hdr, expected_shape_infr_decls) assert class_name is not None # Generate nativefunction declarations # Note, eager registrations is set to False for the lazy TS backend as another LTC backend # may want to register their own lazy kernels instead of registering the TS ones. # The registration will lazily happen when init_ts_backend is called. gen_dispatchkey_nativefunc_headers(fm, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, autograd_key, backend_name) # Generate Dispatcher registrations which hook up the nativefunctions for dispatch_key in [backend_key] if autograd_key is None else [ backend_key, autograd_key ]: gen_dispatcher_registrations(fm, output_dir, class_name, cpp_namespace, backend_indices, grouped_native_functions, backend_key, dispatch_key, selector, build_in_tree=build_in_tree, per_operator_headers=per_operator_headers, backend_name=backend_name, eager_registration=False) # Generate native function impls that build IR nodes ns_helper = NamespaceHelper(cpp_namespace) fm.write_with_template( f'{backend_key}NativeFunctions.cpp', 'DispatchKeyNativeFunctions.cpp', lambda: { 'includes': [ f'#include <{path}>' for path in [ tensor_class_hdr, shape_inference_hdr, "ATen/Functions.h", "ATen/MetaFunctions.h", "ATen/Operators.h", "ATen/native/CPUFallback.h", "torch/csrc/lazy/core/lazy_graph_executor.h", "torch/csrc/lazy/core/metrics.h", "torch/csrc/lazy/core/shape.h", f"{output_dir}/{backend_key}NativeFunctions.h", f"{output_dir}/LazyIr.h", ] + (["torch/csrc/lazy/ts_backend/ts_eager_fallback.h"] if gen_forced_fallback_code else []) ], 'native_functions_include': '', 'namespace_prologue': ns_helper.prologue, 'namespace_epilogue': ns_helper.epilogue, 'native_function_definitions': list( concat_map_codegen(dest.GenLazyNativeFuncDefinition( f'{backend_key}NativeFunctions', backend_indices[ backend_key], tensor_class, gen_forced_fallback_code), grouped_native_functions, codegenInplaceVariant=True)), }) # Generate IR node classes fm.write_with_template( 'LazyIr.h', 'LazyIr.h', lambda: { 'lazy_ir_sysinc': [ f'#include <{path}>' for path in [ "ATen/core/Formatting.h", "c10/core/ScalarType.h", "c10/util/Optional.h", "torch/csrc/lazy/core/hash.h", "torch/csrc/lazy/core/ir.h", "torch/csrc/lazy/core/shape.h", "vector", ] ], 'lazy_ir_inc': [ f'#include "{path}"' for path in [node_base_hdr if node_base_hdr is not None else None] if path is not None ], 'ir_declarations': list( concat_map_codegen( lazy_ir_cls(backend_indices[backend_key], node_base), grouped_native_functions)), 'namespace_prologue': ns_helper.prologue, 'namespace_epilogue': ns_helper.epilogue, })
def main() -> None: parser = argparse.ArgumentParser(description='Generate ATen source files') parser.add_argument('-s', '--source-path', help='path to source directory for ATen', default='aten/src/ATen') parser.add_argument( '-o', '--output-dependencies', help='output a list of dependencies into the given file and exit') parser.add_argument('-d', '--install_dir', help='output directory', default='build/aten/src/ATen') parser.add_argument( '--rocm', action='store_true', help='reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly') # TODO: --op_registration_whitelist will be removed when all call-sites # for gen.py are moved over to using the operator YAML file for mobile # custom build. parser.add_argument( '--op_registration_whitelist', nargs='*', help='filter op registrations by the whitelist (if set); ' 'each item is `namespace`::`operator name` without overload name; ' 'e.g.: aten::empty aten::conv2d ...') 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::)') parser.add_argument( '--backend_whitelist', nargs='*', help='filter dispatch backend by the whitelist (if set), ' 'e.g.: CPU CUDA QuantizedCPU ...') parser.add_argument( '--force_schema_registration', action='store_true', help= 'force it to generate schema-only registrations for all ops, including' 'those that are not listed on --op_registration_whitelist') options = parser.parse_args() selector = get_custom_build_selector( options.op_registration_whitelist, options.op_selection_yaml_path, ) native_functions = parse_native_yaml( os.path.join(options.source_path, 'native/native_functions.yaml')) pre_grouped_native_functions: Dict[FunctionSchema, Dict[SchemaKind, NativeFunction]] pre_grouped_native_functions = defaultdict(dict) for f in native_functions: d = pre_grouped_native_functions[f.func.signature()] assert f.func.kind() not in d d[f.func.kind()] = f grouped_native_functions = [ NativeFunctionGroup.from_dict(v) for v in pre_grouped_native_functions.values() ] # NB: At the moment, grouped_native_functions isn't used by anything, # this code lives here to help potential future consumers; for a live # example see https://github.com/pytorch/pytorch/pull/45277 template_dir = os.path.join(options.source_path, "templates") # NB: It is mandatory to NOT use os.path.join here, as the install directory # will eventually be ingested by cmake, which does not respect Windows style # path slashes. If you switch this to use os.path.join, you'll get an error # like: # # Syntax error in cmake code when parsing string # # C:/Jenkins/workspace/pytorch-builds/pytorch-win-ws2016-cuda9-cudnn7-py3-build/build/aten/src/ATen\core/TensorMethods.h # # Invalid character escape '\c'. core_install_dir = f'{options.install_dir}/core' pathlib.Path(core_install_dir).mkdir(parents=True, exist_ok=True) def make_file_manager(install_dir: str) -> FileManager: return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.output_dependencies) core_fm = make_file_manager(core_install_dir) cpu_fm = make_file_manager(options.install_dir) cuda_fm = make_file_manager(options.install_dir) extra_cuda_headers = '''\ #include <ATen/cuda/ATenCUDAGeneral.h> #include <ATen/cuda/CUDADevice.h> #include <ATen/cuda/CUDAContext.h>''' if options.rocm: extra_cuda_headers = '''\ #include <ATen/hip/ATenHIPGeneral.h> #include <ATen/hip/HIPDevice.h> #include <ATen/hip/HIPContext.h>''' # NB: substrings in these dispatch keys matter, we do tests to see if # a key contains, e.g., CUDA to classify it as a CUDA backend dispatch_keys = [ "CPU", "SparseCPU", "MkldnnCPU", "CUDA", "SparseCUDA", "QuantizedCPU", "QuantizedCUDA", "Math", "DefaultBackend", ] if options.backend_whitelist: dispatch_keys = [ k for k in dispatch_keys if is_generic_dispatch_key(k) or k in options.backend_whitelist ] for dispatch_key in dispatch_keys: cpp_template = 'RegisterDispatchKey.cpp' fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm fm.write_with_template( f'Register{dispatch_key}.cpp', cpp_template, lambda: { 'extra_cuda_headers': extra_cuda_headers if is_cuda_dispatch_key(dispatch_key) else '', 'legacy_th_headers': '#include <ATen/LegacyTHFunctionsCPU.h>' if dispatch_key == "CPU" else '#include <ATen/LegacyTHFunctionsCUDA.h>' if dispatch_key == "CUDA" else '', 'DispatchKey': dispatch_key, 'dispatch_definitions': list( mapMaybe( RegisterDispatchKey(dispatch_key, Target.DEFINITION, selector), native_functions)), 'dispatch_registrations': list( mapMaybe( RegisterDispatchKey(dispatch_key, Target.REGISTRATION, selector), native_functions)), }) del fm # BackendSelect is generated specially cpu_fm.write( 'RegisterBackendSelect.cpp', lambda: { 'backend_select_method_definitions': list( mapMaybe(ComputeBackendSelect(Target.DEFINITION), native_functions)), 'backend_select_function_registrations': list( mapMaybe(ComputeBackendSelect(Target.REGISTRATION), native_functions)), }) schema_selector = selector if options.force_schema_registration: schema_selector = SelectiveBuilder.get_nop_selector() cpu_fm.write( 'RegisterSchema.cpp', lambda: { 'schema_registrations': list(mapMaybe(RegisterSchema(schema_selector), native_functions)), }) cpu_fm.write( 'Functions.h', lambda: { 'function_declarations': list( mapMaybe(ComputeFunction(Target.DECLARATION), native_functions) ), }) cpu_fm.write( 'Functions.cpp', lambda: { 'function_definitions': list(mapMaybe(ComputeFunction(Target.DEFINITION), native_functions) ), }) core_fm.write( 'TensorBody.h', lambda: { 'tensor_method_declarations': list( mapMaybe(ComputeTensorMethod(Target.DECLARATION), native_functions)), }) core_fm.write( 'TensorMethods.cpp', lambda: { 'tensor_method_definitions': list( mapMaybe(ComputeTensorMethod(Target.DEFINITION), native_functions)), }) core_fm.write( 'ATenOpList.cpp', lambda: { 'aten_ops': list(mapMaybe(compute_aten_op, native_functions)), }) cpu_fm.write( 'NativeFunctions.h', lambda: { 'native_function_declarations': list( concatMap(compute_native_function_declaration, native_functions )), }) cpu_fm.write( 'Declarations.yaml', lambda: format_yaml( [compute_declaration_yaml(f) for f in native_functions])) cpu_fm.write( 'RegistrationDeclarations.h', lambda: { 'registration_declarations': [compute_registration_declarations(f) for f in native_functions], }) if options.output_dependencies: cpu_fm.write_outputs(options.output_dependencies) core_fm.write_outputs(f"{options.output_dependencies}-core") cuda_fm.write_outputs(f"{options.output_dependencies}-cuda")
def main(argv: List[Any]) -> None: """This binary generates 3 files: 1. selected_mobile_ops.h: Primary operators used by templated selective build and Kernel Function dtypes captured by tracing 2. selected_operators.yaml: Selected root and non-root operators (either via tracing or static analysis) """ parser = argparse.ArgumentParser(description="Generate operator lists") parser.add_argument( "--output_dir", help= ("The directory to store the output yaml files (selected_mobile_ops.h, " + "selected_kernel_dtypes.h, selected_operators.yaml)"), required=True, ) parser.add_argument( "--model_file_list_path", help=( "Path to a file that contains the locations of individual " + "model YAML files that contain the set of used operators. This " + "file path must have a leading @-symbol, which will be stripped " + "out before processing."), required=True, ) parser.add_argument( "--allow_include_all_overloads", help= ("Flag to allow operators that include all overloads. " + "If not set, operators registered without using the traced style will" + "break the build."), action="store_true", default=False, required=False, ) options = parser.parse_args() if (os.path.isfile(options.model_file_list_path)): print("Processing model file: ", options.model_file_list_path) model_dicts = [] model_dict = yaml.safe_load(open(options.model_file_list_path)) model_dicts.append(model_dict) else: print("Processing model directory: ", options.model_file_list_path) assert options.model_file_list_path[0] == "@" model_file_list_path = options.model_file_list_path[1:] model_dicts = [] with open(model_file_list_path) as model_list_file: model_file_names = model_list_file.read().split() for model_file_name in model_file_names: with open(model_file_name, "rb") as model_file: model_dict = yaml.safe_load(model_file) model_dicts.append(model_dict) selective_builders = list( map( lambda m: SelectiveBuilder.from_yaml_dict(m), model_dicts, )) # While we have the model_dicts generate the supported mobile models api gen_supported_mobile_models(model_dicts, options.output_dir) # We may have 0 selective builders since there may not be any viable # pt_operator_library rule marked as a dep for the pt_operator_registry rule. # This is potentially an error, and we should probably raise an assertion # failure here. However, this needs to be investigated further. selective_builder = SelectiveBuilder.from_yaml_dict({}) if len(selective_builders) > 0: selective_builder = reduce( combine_selective_builders, selective_builders, ) if not options.allow_include_all_overloads: throw_if_any_op_includes_overloads(selective_builder) with open(os.path.join(options.output_dir, "selected_operators.yaml"), "wb") as out_file: out_file.write( yaml.safe_dump(selective_builder.to_dict(), default_flow_style=False).encode("utf-8"), ) write_selected_mobile_ops( os.path.join(options.output_dir, "selected_mobile_ops.h"), selective_builder, )