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)
),
})
