def gen_variable_type_func(
fn: NativeFunctionWithDifferentiabilityInfo) -> Dict[str, List[str]]:
f = fn.func
with native_function_manager(f):
name = cpp.name(f.func)
formals = gen_formals(f)
type_definition = METHOD_DEFINITION.substitute(
return_type=cpp.returns_type(f.func.returns).cpp_type(),
type_wrapper_name=type_wrapper_name(f),
type_definition_body=emit_body(fn),
formals=formals,
)
wrapper_registration = gen_wrapper_registration(f)
# See Note [Manual Backend kernels]
assert (name in MANUAL_BACKEND) == f.manual_kernel_registration
# If you want to register a kernel to Autograd, you must make the op abstract.
# In other words, this op must have dispatch section in native_functions.yaml.
if name in MANUAL_AUTOGRAD_AND_TRACER or (fn.info
and fn.info.has_derivatives):
msg = (
f'There\'s a formula for {name}(or its functional variant) in derivatives.yaml. '
f'It\'s required to add a dispatch section for it with explicit supported backends e.g CPU/CUDA '
f'or CompositeExplicitAutograd in native_functions.yaml. Please see '
f'https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native#choosing-the-right-dispatch-keyword '
f'for instructions to choose the right dispatch keyword.')
assert f.is_abstract, msg
return {
'type_derived_method_definitions': [type_definition],
'wrapper_registrations': [wrapper_registration],
}
def gen_variable_type_func(
fn: NativeFunctionWithDifferentiabilityInfo) -> Dict[str, List[str]]:
f = fn.func
with native_function_manager(f):
name = cpp.name(f.func)
formals = gen_formals(f)
if fn.info is None and not get_base_name(f) in RESET_GRAD_ACCUMULATOR \
and not get_base_name(f) in DONT_REQUIRE_DERIVATIVE \
and len(gen_differentiable_outputs(fn)) > 0 \
and not get_base_name(f) in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE \
and not get_base_name(f) in DONT_ENFORCE_STORAGE_IMPL_USE_COUNT \
and not get_base_name(f) in DONT_ENFORCE_TENSOR_IMPL_USE_COUNT:
# NOTE: [ Registering AutogradNotImplemented boxed kernel ]
#
# When there is no derivatives.yaml entry, we register a generic boxed
# NotImplemented kernel to set grad_fn to be NotImplemented, so that forward
# proceeds as usual but an error is properly produced on backward.
# TODO: it would be nice to not have these special cases
#
# There are several cases where still let codegen handle it:
# 1) ops that need to reset grad accumulator (we let codegen handle this case
# because) the list is (currently) only accessible in Python.
# 2) User explicitly specifies DONT_REQUIRE_DERIVATIVE. This basically makes
# autograd a fallthrough with NDEBUG checks. This can be useful for when all
# outputs are integral.
# 3) When there are no differentiable outputs. This is similar to (2).
# 4) There are certain ops where we skip certain NDEBUG checks. this is similar
# to (1).
type_definition = ""
wrapper_registration = AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION.substitute(
unqual_operator_name_with_overload=f.func.name)
else:
type_definition = METHOD_DEFINITION.substitute(
return_type=cpp.returns_type(f.func.returns).cpp_type(),
type_wrapper_name=type_wrapper_name(f),
type_definition_body=emit_body(fn),
formals=formals,
)
wrapper_registration = gen_wrapper_registration(f)
# See Note [Manual Backend kernels]
assert (name in MANUAL_BACKEND) == f.manual_kernel_registration
# If you want to register a kernel to Autograd, you must make the op abstract.
# In other words, this op must have dispatch section in native_functions.yaml.
if name in MANUAL_AUTOGRAD_AND_TRACER or (fn.info
and fn.info.has_derivatives):
msg = (
f'There\'s a formula for {name}(or its functional variant) in derivatives.yaml. '
f'It\'s required to add a dispatch section for it with explicit supported backends e.g CPU/CUDA '
f'or CompositeExplicitAutograd in native_functions.yaml. Please see '
f'https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native#choosing-the-right-dispatch-keyword '
f'for instructions to choose the right dispatch keyword.')
assert f.is_abstract, msg
return {
'type_derived_method_definitions': [type_definition],
'wrapper_registrations': [wrapper_registration],
}
def __call__(self, f: NativeFunction) -> Optional[str]:
# We unconditionally generate function variants of the redispatch API.
# This is mainly because we can namespace functions separately, but not methods,
if Variant.function not in f.variants and not self.is_redispatching_fn:
return None
with native_function_manager(f):
return self.callImpl(f)
def compute_meta_function_declaration(g: NativeFunctionsGroup) -> Optional[str]:
if not g.structured:
return None
with native_function_manager(g.out):
name = meta.name(g)
args = structured.meta_arguments(g)
args_str = ', '.join(a.decl() for a in args)
parent_class = g.out.structured_inherits
if parent_class is None:
parent_class = "at::impl::MetaBase"
return f"""\
def __call__(self, groups: Sequence[NativeFunctionsGroup]) -> str:
if not groups:
return ""
generated_type_variants = []
for g in groups:
with native_function_manager(g):
assert is_supported(g)
assert isinstance(g, NativeFunctionsGroup)
generated_type_variant = self.gen_structured_test_case(g)
generated_type_variants.append(generated_type_variant)
return "\n".join(generated_type_variants)
def gen_variable_type_shard(
fm: FileManager,
fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo],
template_name: str,
output_name: str,
) -> None:
type_definitions: List[str] = []
wrapper_registrations: List[str] = []
filtered_fns_with_diff_infos = list(
filter(use_derived, fns_with_diff_infos))
for fn in filtered_fns_with_diff_infos:
f = fn.func
with native_function_manager(f):
name = cpp.name(f.func)
formals = gen_formals(f)
type_definitions.append(
METHOD_DEFINITION.substitute(
return_type=cpp.returns_type(f.func.returns).cpp_type(),
type_wrapper_name=type_wrapper_name(f),
type_definition_body=emit_body(fn),
formals=formals,
))
wrapper_registrations.append(gen_wrapper_registration(f))
# See Note [Manual Backend kernels]
assert (name in MANUAL_BACKEND) == f.manual_kernel_registration
# If you want to register a kernel to Autograd, you must make the op abstract.
# In other words, this op must have dispatch section in native_functions.yaml.
if name in MANUAL_AUTOGRAD_AND_TRACER or (fn.info
and fn.info.has_derivatives):
msg = (
f'There\'s a formula for {name}(or its functional variant) in derivatives.yaml. '
f'It\'s required to add a dispatch section for it with explicit supported backends e.g CPU/CUDA '
f'or CompositeExplicitAutograd in native_functions.yaml. Please see '
f'https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native#choosing-the-right-dispatch-keyword '
f'for instructions to choose the right dispatch keyword.')
assert f.is_abstract, msg
fm.write_with_template(
output_name, template_name, lambda: {
'generated_comment': '@'
f'generated from {fm.template_dir}/{template_name}',
'type_derived_method_definitions': type_definitions,
'wrapper_registrations': wrapper_registrations,
})
def __call__(self, groups: Sequence[NativeFunctionsGroup]) -> str:
if not groups:
return ""
generated_type_variants = []
for g in groups:
with native_function_manager(g):
assert is_supported(g)
assert isinstance(g, NativeFunctionsGroup)
generated_type_variant = self.gen_structured(g)
generated_type_variants.append(generated_type_variant)
op_name = op_name_from_group(groups[0])
body = "\n".join(generated_type_variants)
generated = f"""
REGISTER_OPERATOR_FUNCTOR(
aten::{op_name},
aten_{op_name},
[](Node* n) -> SROperator {{
{body}
LogAndDumpSchema(n);
return nullptr;
}});
"""
return generated
def create_decl(f: NativeFunction) -> str:
with native_function_manager(f):
return DispatcherSignature.from_schema(f.func).decl()
def gen_unstructured(
self,
f: NativeFunction,
g: Optional[NativeFunctionsGroup] = None) -> Optional[str]:
with native_function_manager(f):
inplace_meta = False
gets_out_inplace_wrapper = False
if not self.backend_index.has_kernel(f):
if (self.backend_index.dispatch_key == DispatchKey.Meta
and f.func.kind() is SchemaKind.inplace and
# Defer to composites for meta implementation
not f.has_composite_kernel and
# Inplace list operations are not supported
len(f.func.returns) == 1):
inplace_meta = True
elif (not self.backend_index.use_out_as_primary
and g is not None and gets_generated_out_inplace_wrapper(
f, g, self.backend_index)):
# We want to generate inplace/out wrappers, that don't have a kernel for the backend.
gets_out_inplace_wrapper = True
else:
return None
if f.manual_kernel_registration:
return None
if self.target is Target.REGISTRATION and not self.selector.is_native_function_selected(
f):
return None
sig = self.wrapper_kernel_sig(f)
name = sig.name()
returns_type = sig.returns_type().cpp_type()
args = sig.arguments()
args_str = ', '.join(a.defn() for a in args)
# See Note [Direct dispatch bindings]
cpp_sig_group = CppSignatureGroup.from_native_function(
f, method=False, fallback_binding=False)
if self.target is Target.NAMESPACED_DECLARATION:
result = f"TORCH_API {cpp_sig_group.signature.decl()};\n"
if cpp_sig_group.faithful_signature is not None:
result += f"TORCH_API {cpp_sig_group.faithful_signature.decl()};\n"
return result
elif self.target is Target.NAMESPACED_DEFINITION:
def generate_defn(cpp_sig: CppSignature) -> str:
return f"""
{cpp_sig.defn()} {{
return {sig.name()}({', '.join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))});
}}
"""
result = generate_defn(cpp_sig_group.signature)
if cpp_sig_group.faithful_signature is not None:
result += generate_defn(cpp_sig_group.faithful_signature)
return result
elif self.target is Target.ANONYMOUS_DEFINITION:
# short circuit for inplace_meta
if inplace_meta:
assert f.func.arguments.self_arg is not None
self_arg_name = f.func.arguments.self_arg.argument.name
# TODO: handle in place on tensor list
return f"""
{returns_type} {name}({args_str}) {{
TORCH_CHECK_NOT_IMPLEMENTED({self_arg_name}.is_meta(),
"Cannot inplace into non-meta tensor with meta tensor argument");
return {self_arg_name};
}}
"""
# short circuit for generated inplace/out wrappers
if gets_out_inplace_wrapper:
return self.gen_out_inplace_wrapper(f, g)
metadata = self.backend_index.get_kernel(f)
if metadata is None:
return None
if self.class_method_name is None:
impl_name = f"{self.cpp_namespace}::{metadata.kernel}"
else:
impl_name = f"{self.cpp_namespace}::{self.class_method_name}::{metadata.kernel}"
args_exprs_str = ', '.join(a.name for a in args)
device_check = ' // No device check\n'
if is_cuda_dispatch_key(self.backend_index.dispatch_key):
device_check_args = itertools.chain(
f.func.arguments.out, f.func.arguments.flat_positional)
device_check = RegisterDispatchKey.gen_device_check(
f.device_check, list(device_check_args), name)
device_guard = "// DeviceGuard omitted" # default
if f.device_guard and is_cuda_dispatch_key(
self.backend_index.dispatch_key):
has_tensor_options = any(
isinstance(a.argument, TensorOptionsArguments)
for a in args)
if has_tensor_options:
# kernel is creating a tensor
device_guard = """globalContext().lazyInitCUDA();
const DeviceGuard device_guard(device_or_default(device));"""
else:
# kernel is operating on existing tensors
# There is precedence for which argument we use to do
# device guard. This describes the precedence order.
self_arg = [
f.func.arguments.self_arg.argument
] if f.func.arguments.self_arg is not None else []
candidate_args = itertools.chain(
self_arg, f.func.arguments.out,
f.func.arguments.flat_positional)
# Only tensor like arguments are eligible
device_of = next((f'{a.name}' for a in candidate_args
if a.type.is_tensor_like()), None)
if device_of is not None:
device_guard = f"const OptionalDeviceGuard device_guard(device_of({device_of}));"
return f"""\
namespace {{
{returns_type} {name}({args_str}) {{
{device_check}
{device_guard}
return {impl_name}({args_exprs_str});
}}
}} // anonymous namespace
"""
elif self.target is Target.REGISTRATION:
if f.manual_kernel_registration:
return None
else:
payload = f"TORCH_FN({name})"
return f'm.impl("{f.func.name}",\n{payload});\n'
else:
assert_never(self.target)
def wrapper(f: NFWDI) -> T:
with native_function_manager(f.func):
return func(f)
def is_supported(g: NativeFunctionsGroup) -> bool:
with native_function_manager(g):
return gen_structured.is_supported(g)
