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 gen_dispatchkey_nativefunc_headers( fm: FileManager, class_name: str, cpp_namespace: str, backend_indices: Dict[DispatchKey, BackendIndex], grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], backend_dispatch_key: DispatchKey, autograd_dispatch_key: Optional[DispatchKey]) -> None: assert class_name is not None generated_comment = 'Autogenerated file by gen_backend_stubs.py. Do not edit directly!' # Convert to a set first to remove duplicate kernel names. # Backends are allowed to repeat kernel names; only generate the declaration once! # Sort for deterministic output. backend_declarations = list(sorted(set(concatMap( lambda f: dest.compute_native_function_declaration(f, backend_indices[backend_dispatch_key]), grouped_native_functions)))) autograd_declarations = list(sorted(set(concatMap( lambda f: [] if autograd_dispatch_key is None else dest.compute_native_function_declaration(f, backend_indices[autograd_dispatch_key]), grouped_native_functions)))) ns_helper = NamespaceHelper(cpp_namespace) fm.write_with_template(f'{backend_dispatch_key}NativeFunctions.h', 'DispatchKeyNativeFunctions.h', lambda: { 'generated_comment': generated_comment, 'namespace_prologue': ns_helper.prologue, 'class_name': class_name, 'namespace_epilogue': ns_helper.epilogue, 'dispatch_declarations': backend_declarations + autograd_declarations, })
def gen_dispatcher_registrations(fm: FileManager, output_dir: str, cpp_namespace: str, backend_indices: Dict[DispatchKey, BackendIndex], grouped_native_functions: Sequence[Union[ NativeFunction, NativeFunctionsGroup]], backend_dispatch_key: DispatchKey, dispatch_key: DispatchKey, selector: 'SelectiveBuilder') -> None: backend_index = backend_indices[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_dispatch_key}NativeFunctions.h"', 'ops_headers': '#include <ATen/Functions.h>', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_headers': dest.gen_registration_headers( backend_index, per_operator_headers=False, rocm=False), 'dispatch_helpers': dest.gen_registration_helpers(backend_index), 'dispatch_namespaced_definitions': '', 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey( backend_index, 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_index, Target.REGISTRATION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name= f'{dispatch_key}NativeFunctions'), grouped_native_functions)), })
def get_grouped_native_functions( native_yaml_path: str ) -> Sequence[Union[NativeFunction, NativeFunctionsGroup]]: native_functions = parse_native_yaml(native_yaml_path) 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 return list( concatMap(flatten_pre_group, list(pre_grouped_native_functions.values())))
def compute_native_function_declaration(g: Union[NativeFunctionsGroup, NativeFunction]) -> List[str]: if isinstance(g, NativeFunctionsGroup): if g.structured: return gen_structured(g) else: return list(concatMap(gen_unstructured, g.functions())) else: return gen_unstructured(g)
def compute_native_function_declaration( g: Union[NativeFunctionsGroup, NativeFunction, ExternalBackendFunctionsGroup, ExternalBackendFunction] ) -> List[str]: if isinstance(g, ExternalBackendFunctionsGroup): if g.structured: raise AssertionError( "Structured external backend functions are not implemented yet." ) else: return list(concatMap(gen_unstructured_external, g.functions())) elif isinstance(g, ExternalBackendFunction): return gen_unstructured_external(g) elif isinstance(g, NativeFunctionsGroup): if g.structured: return gen_structured(g) else: return list(concatMap(gen_unstructured, g.functions())) else: return gen_unstructured(g)
def load_derivatives(derivatives_yaml_path: str, native_yaml_path: str) -> Sequence[DifferentiabilityInfo]: # Do some caching as this is a deterministic function global _GLOBAL_LOAD_DERIVATIVE_CACHE key = (derivatives_yaml_path, native_yaml_path) if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE: with open(derivatives_yaml_path, 'r') as f: definitions = yaml.load(f, Loader=YamlLoader) funcs = parse_native_yaml(native_yaml_path).native_functions # From the parsed native functions, separate out the (generated) view_copy functions, # so we can generate derivatives for them separately. native_functions_with_view_groups = get_grouped_by_view_native_functions( funcs) native_functions_without_view_copies = concatMap( # We need to pull out the view_inplace ops too, since they might have their own derivative entries. lambda g: [g] if isinstance(g, NativeFunction) else list( g.functions(include_copy=False)), native_functions_with_view_groups) view_groups = [ g for g in native_functions_with_view_groups if isinstance(g, NativeFunctionsViewGroup) ] # What's the difference between function schema v.s. signature? # function schema is the complete declaration including mutability annotation / default value and etc. # signature is the canonical schema for a group of functions (in-place/out/functional variants) # that are semantically related. functions_by_signature: Dict[FunctionSchema, List[NativeFunction]] = defaultdict(list) functions_by_schema: Dict[str, NativeFunction] = dict() for function in native_functions_without_view_copies: functions_by_signature[function.func.signature()].append(function) assert str(function.func) not in functions_by_schema functions_by_schema[str(function.func)] = function # Keep track of how many of which ops we've seen so we can # disambiguate them with a numeric suffix. op_counter = Counter[str]() infos = [ create_differentiability_info(defn, functions_by_signature, functions_by_schema, op_counter) for defn in definitions ] infos += add_view_copy_derivatives(infos, view_groups) _GLOBAL_LOAD_DERIVATIVE_CACHE[key] = infos return _GLOBAL_LOAD_DERIVATIVE_CACHE[key]
def jit_arguments(func: FunctionSchema) -> List[Argument]: def to_argument(a: Union[Argument, TensorOptionsArguments, SelfArgument]) -> List[Argument]: if isinstance(a, Argument): return [a] elif isinstance(a, SelfArgument): return [a.argument] elif isinstance(a, TensorOptionsArguments): return [a.dtype, a.layout, a.device, a.pin_memory] else: assert_never(a) return list(concatMap(to_argument, itertools.chain( func.arguments.positional, func.arguments.kwarg_only, func.arguments.out)))
def parse_backend_yaml( backend_yaml_path: str, grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]] ) -> Tuple[str, List[Union[ExternalBackendFunction, ExternalBackendFunctionsGroup]]]: with open(backend_yaml_path, 'r') as f: yaml_values = yaml.load(f, Loader=LineLoader) assert isinstance(yaml_values, dict) cpp_namespace = yaml_values.pop('cpp_namespace') backend = yaml_values.pop('backend') supported = yaml_values.pop('supported', []) assert isinstance(supported, list), f'expected "supported" to be a list, but got: {supported}' supported_autograd = yaml_values.pop('autograd', []) assert isinstance(supported, list), f'expected "autograd" to be a list, but got: {supported_autograd}' assert len(yaml_values.keys()) > 0, \ f'{backend_yaml_path} contains unexpected keys: {", ".join(yaml_values.keys())}' metadata: Dict[OperatorName, ExternalBackendMetadata] = {} for op in supported: op_name = OperatorName.parse(op) m = ExternalBackendMetadata(op_name, backend, is_autograd=False) metadata[m.operator] = m for op in supported_autograd: op_name = OperatorName.parse(op) m = ExternalBackendMetadata(op_name, backend, is_autograd=True) metadata[m.operator] = m native_functions_map: Dict[OperatorName, NativeFunction] = { f.func.name: f for f in concatMap(lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()), grouped_native_functions) } def native_to_external( g: Union[NativeFunction, NativeFunctionsGroup] ) -> Union[ExternalBackendFunction, ExternalBackendFunctionsGroup]: if isinstance(g, NativeFunction): f = g m = metadata.get(f.func.name, None) return ExternalBackendFunction(f, m) elif isinstance(g, NativeFunctionsGroup): return ExternalBackendFunctionsGroup.from_function_group(g, metadata) else: assert_never(g) for op_name in metadata.keys(): if op_name not in native_functions_map: raise AssertionError(f"Found an invalid operator name: {op_name}") return cpp_namespace, [native_to_external(g) for g in grouped_native_functions]
def parse_full_codegen_ops( backend_yaml_path: str, grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], ) -> List[OperatorName]: native_functions_map: Dict[OperatorName, NativeFunction] = { f.func.name: f for f in concatMap(lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()), grouped_native_functions) } with open(backend_yaml_path, 'r') as f: yaml_values = yaml.load(f, Loader=YamlLoader) assert isinstance(yaml_values, dict) full_codegen = yaml_values.pop('full_codegen', []) assert isinstance(full_codegen, list), f'expected "full_codegen" to be a list, but got: {full_codegen}' full_codegen = [OperatorName.parse(name) for name in full_codegen] return full_codegen
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 parse_backend_yaml( backend_yaml_path: str, grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], backend_indices: Dict[DispatchKey, BackendIndex]) -> ParsedExternalYaml: native_functions_map: Dict[OperatorName, NativeFunction] = { f.func.name: f for f in concatMap( lambda f: [f] if isinstance(f, NativeFunction) else list( f.functions()), grouped_native_functions) } with open(backend_yaml_path, 'r') as f: yaml_values = yaml.load(f, Loader=YamlLoader) assert isinstance(yaml_values, dict) valid_keys = [ 'backend', 'cpp_namespace', 'extra_headers', 'supported', 'autograd' ] backend = yaml_values.pop('backend', None) assert backend is not None, 'You must provide a value for "backend"' cpp_namespace = yaml_values.pop('cpp_namespace', None) assert cpp_namespace is not None, 'You must provide a value for "cpp_namespace"' supported = yaml_values.pop('supported', []) if supported is None: supported = [] # Allow an empty list of supported ops assert isinstance( supported, list ), f'expected "supported" to be a list, but got: {supported} (of type {type(supported)})' supported_autograd = yaml_values.pop('autograd', []) assert isinstance( supported, list ), f'expected "autograd" to be a list, but got: {supported_autograd}' assert len(yaml_values.keys()) == 0, \ f'{backend_yaml_path} contains unexpected keys: {", ".join(yaml_values.keys())}. \ Only the following keys are supported: {", ".join(valid_keys)}' def create_backend_index(backend_ops: List[str], dispatch_key: DispatchKey) -> BackendIndex: metadata: Dict[OperatorName, BackendMetadata] = {} for op in backend_ops: op_name = OperatorName.parse(op) assert op_name in native_functions_map, f"Found an invalid operator name: {op_name}" # See Note [External Backends Follow Dispatcher API] kernel_name = dispatcher.name(native_functions_map[op_name].func) # TODO: allow structured external backends later. m = BackendMetadata(kernel=kernel_name, structured=False) metadata[op_name] = m # TODO: currently hardcoding the fact that XLA implements out/inplace in terms of functional ops, # this should eventually be toggleable per-backend. return BackendIndex(dispatch_key=dispatch_key, use_out_as_primary=False, external=True, index=metadata) backend_key: Optional[DispatchKey] = None if len(supported) > 0: with context( lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.' ): backend_key = DispatchKey.parse(backend) backend_idx = create_backend_index(supported, backend_key) assert backend_key not in backend_indices backend_indices[backend_key] = backend_idx autograd_key: Optional[DispatchKey] = None if len(supported_autograd) > 0: with context( lambda: f'The "autograd" key was specified, which indicates that you would like to override \ the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.' ): autograd_key = DispatchKey.parse(f'Autograd{backend}') autograd_idx = create_backend_index(supported_autograd, autograd_key) assert autograd_key not in backend_indices backend_indices[autograd_key] = autograd_idx for g in grouped_native_functions: if isinstance(g, NativeFunction): forward_kernels = [] if backend_key is None else \ [m for m in [backend_indices[backend_key].get_kernel(g)] if m is not None] backward_kernels = [] if autograd_key is None else \ [m for m in [backend_indices[autograd_key].get_kernel(g)] if m is not None] else: forward_kernels = [] if backend_key is None else [ m for m in [ backend_indices[backend_key].get_kernel(f) for f in g.functions() ] if m is not None ] backward_kernels = [] if autograd_key is None else [ m for m in [ backend_indices[autograd_key].get_kernel(f) for f in g.functions() ] if m is not None ] forward_kernels = [f for f in forward_kernels if f is not None] backward_kernels = [f for f in backward_kernels if f is not None] assert len(forward_kernels) == 0 or len(backward_kernels) == 0, \ f'Currently, all variants of an op must either be registered to a backend key, or to a backend\'s \ autograd key. They cannot be mix and matched. If this is something you need, feel free to create an issue! \ {forward_kernels[0].kernel} is listed under "supported", but {backward_kernels[0].kernel} is listed under "autograd".' return ParsedExternalYaml(backend_key, autograd_key, cpp_namespace, backend_indices)
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 gen_dispatcher_registrations( fm: FileManager, output_dir: str, class_name: str, cpp_namespace: str, backend_indices: Dict[DispatchKey, BackendIndex], grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], backend_dispatch_key: DispatchKey, dispatch_key: DispatchKey, selector: 'SelectiveBuilder', # build_in_tree is true for lazy TS backend and affects include paths, not used for external backends build_in_tree: bool = False, per_operator_headers: bool = False) -> None: headers = [ f"{output_dir}/{backend_dispatch_key}NativeFunctions.h", ] if build_in_tree: external_backend_headers_str = "\n".join(f'#include <{h}>' for h in headers) else: external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers) assert class_name is not None backend_index = backend_indices[dispatch_key] fm.write_with_template( f'Register{dispatch_key}.cpp', 'RegisterDispatchKey.cpp', lambda: { 'extra_cuda_headers': '', 'external_backend_headers': external_backend_headers_str, 'ops_headers': '#include <ATen/Functions.h>' if not per_operator_headers else '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_headers': dest.gen_registration_headers(backend_index, per_operator_headers= per_operator_headers, rocm=False), 'dispatch_helpers': dest.gen_registration_helpers(backend_index), 'dispatch_namespaced_definitions': '', 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey(backend_index, Target.ANONYMOUS_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name=f'{class_name}', skip_dispatcher_op_registration= False), grouped_native_functions) ), 'dispatch_registrations': list( concatMap( dest.RegisterDispatchKey(backend_index, Target.REGISTRATION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name=f'{class_name}', skip_dispatcher_op_registration= False), 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 gen_dispatcher_registrations( fm: FileManager, output_dir: str, class_name: str, cpp_namespace: str, backend_indices: Dict[DispatchKey, BackendIndex], grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]], backend_dispatch_key: DispatchKey, dispatch_key: DispatchKey, selector: 'SelectiveBuilder', # build_in_tree is true for lazy TS backend and affects include paths, not used for external backends build_in_tree: bool = False, per_operator_headers: bool = False, backend_name: str = "", eager_registration: bool = True) -> None: headers = [ f"{output_dir}/{backend_dispatch_key}NativeFunctions.h", ] if build_in_tree: external_backend_headers_str = "\n".join(f'#include <{h}>' for h in headers) else: external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers) assert class_name is not None backend_index = backend_indices[dispatch_key] dispatch_registrations_body = list( concatMap( dest.RegisterDispatchKey(backend_index, Target.REGISTRATION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name=f'{class_name}', skip_dispatcher_op_registration=False), grouped_native_functions)) deferred_dispatch_registrations = "" static_init_dispatch_registrations = "" if eager_registration: static_template = CodeTemplate("""\ TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) { $dispatch_registrations_body };""") static_init_dispatch_registrations = static_template.substitute( dispatch_key=dispatch_key, dispatch_registrations_body=dispatch_registrations_body) else: deferred_template = CodeTemplate("""\ TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() { static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key); $dispatch_registrations_body }""") deferred_dispatch_registrations = deferred_template.substitute( backend_name=backend_name, dispatch_key=dispatch_key, dispatch_registrations_body=dispatch_registrations_body) fm.write_with_template( f'Register{dispatch_key}.cpp', 'RegisterDispatchKey.cpp', lambda: { 'static_init_dispatch_registrations': static_init_dispatch_registrations, 'deferred_dispatch_registrations': deferred_dispatch_registrations, 'extra_cuda_headers': '', 'external_backend_headers': external_backend_headers_str, 'ops_headers': '#include <ATen/Functions.h>' if not per_operator_headers else '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_headers': dest.gen_registration_headers(backend_index, per_operator_headers= per_operator_headers, rocm=False), 'dispatch_helpers': dest.gen_registration_helpers(backend_index), 'dispatch_namespaced_definitions': '', 'dispatch_anonymous_definitions': list( concatMap( dest.RegisterDispatchKey(backend_index, Target.ANONYMOUS_DEFINITION, selector, rocm=False, cpp_namespace=cpp_namespace, class_method_name=f'{class_name}', skip_dispatcher_op_registration= False), grouped_native_functions) ), })