def setUp(self) -> None:
self.native_functions: List[NativeFunction] = []
self.backend_indices: Dict[DispatchKey,
Dict[OperatorName,
BackendMetadata]] = defaultdict(dict)
yaml_entry = """
- func: op(Tensor self) -> Tensor
dispatch:
CompositeExplicitAutograd: op
autogen: op.out
"""
es = yaml.load(yaml_entry, Loader=LineLoader)
self.one_return_func, m = NativeFunction.from_yaml(es[0],
loc=Location(
__file__, 1),
valid_tags=set())
BackendIndex.grow_index(self.backend_indices, m)
self.two_returns_func, two_returns_backend_index = NativeFunction.from_yaml(
{
"func": "op_2() -> (Tensor, Tensor)",
"dispatch": {
"CPU": "kernel_1"
},
"autogen": "op_2.out",
},
loc=torchgen.model.Location(__file__, 1),
valid_tags=set(),
)
BackendIndex.grow_index(self.backend_indices,
two_returns_backend_index)
def setUp(self) -> None:
self.op_1_native_function, op_1_backend_index = NativeFunction.from_yaml(
{
"func": "op_1() -> bool",
"dispatch": {
"CPU": "kernel_1"
}
},
loc=torchgen.model.Location(__file__, 1),
valid_tags=set(),
)
self.op_2_native_function, op_2_backend_index = NativeFunction.from_yaml(
{
"func": "op_2() -> bool",
"dispatch": {
"CPU": "kernel_2",
"QuantizedCPU": "custom::kernel_3"
},
},
loc=torchgen.model.Location(__file__, 1),
valid_tags=set(),
)
backend_indices: Dict[DispatchKey,
Dict[OperatorName, BackendMetadata]] = {
DispatchKey.CPU: {},
DispatchKey.QuantizedCPU: {},
}
BackendIndex.grow_index(backend_indices, op_1_backend_index)
BackendIndex.grow_index(backend_indices, op_2_backend_index)
self.backend_indices = {
k: BackendIndex(
dispatch_key=k,
use_out_as_primary=True,
external=False,
symint=False,
device_guard=False,
index=backend_indices[k],
)
for k in backend_indices
}
def test_custom_namespace_selected_correctly(self):
yaml_config = """
operators:
aten::add.int:
is_used_for_training: No
is_root_operator: Yes
include_all_overloads: No
custom::add:
is_used_for_training: Yes
is_root_operator: No
include_all_overloads: Yes
"""
selector = SelectiveBuilder.from_yaml_str(yaml_config)
native_function, _ = NativeFunction.from_yaml(
{"func": "custom::add() -> Tensor"},
loc=Location(__file__, 1),
valid_tags=set(),
)
self.assertTrue(selector.is_native_function_selected(native_function))
def generate_function(
f: NativeFunction, k: SchemaKind
) -> Tuple[NativeFunction, Dict[DispatchKey, Dict["OperatorName", "BackendMetadata"]]]:
from torchgen.api import cpp
if k == SchemaKind.functional:
assert f.func.kind() != SchemaKind.functional
# The new "functional" NativeFunction has:
# - any mutable arguments have been converted into (immutable) returns.
# (if a mutable argument was not also a return, it gets converted to one)
# - "_functional" appended to the base name, ONLY IF this op has a mutable variant.
# See Note [Overload Ambiguity With Functional Variants]
# The default grouping logic in signature() actually already does this,
# so we can piggy-back off it (but we still want return names)
func = f.func.signature(keep_return_names=True).with_name(
OperatorName(
name=BaseOperatorName(
base=f.func.name.name.base,
inplace=False,
dunder_method=f.func.name.name.dunder_method,
# See Note [Overload Ambiguity With Functional Variants]
functional_overload=f.func.kind() == SchemaKind.mutable,
),
overload_name=f.func.name.overload_name,
)
)
elif k == SchemaKind.out:
# We generate out= ops mostly just so that we can pair up NativeFunctions into groups easily,
# but at least today, there is no good reason to actually use them.
# we'll generate a dispatcher entry for them, but won't actually register any kernels for them.
if f.func.kind() == SchemaKind.inplace:
func = self_to_out_signature(f.func)
elif f.func.kind() == SchemaKind.mutable:
func = mutable_to_out_signature(f.func)
elif f.func.kind() == SchemaKind.functional:
func = functional_to_out_signature(f.func)
else:
raise AssertionError(
"We only bother generating out= functions from either inplace or mutable or functional variants"
)
else:
raise AssertionError(
"We currently only generate either functional or out= NativeFunctions"
)
# Generated kernel naming convention for out: <op_name>_<overload_name>. The reason for this is to
# disambiguate operator with the same name but different overload name, e.g., `randn.names_out` and
# `randn.generator_with_names_out`.
kernel_name = (
func.name.unambiguous_name()
if func.kind() == SchemaKind.out
else cpp.name(func)
)
backend_metadata = {
DispatchKey.CompositeExplicitAutograd: {
func.name: BackendMetadata(
kernel=kernel_name,
structured=False,
cpp_namespace=DEFAULT_KERNEL_NAMESPACE,
)
}
}
return (
NativeFunction(
func=func,
use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors,
# These generated fn's aren't meant to be user friendly- don't generate methods.
variants=set([Variant.function]),
structured=False,
structured_delegate=None,
structured_inherits=None,
precomputed=None,
autogen=[],
ufunc_inner_loop={},
manual_kernel_registration=False,
manual_cpp_binding=False,
python_module=None,
category_override=None,
device_guard=False,
device_check=DeviceCheckType.NoCheck,
loc=f.loc,
cpp_no_default_args=set(),
is_abstract=f.is_abstract,
has_composite_implicit_autograd_kernel=False,
has_composite_explicit_autograd_kernel=True,
has_composite_explicit_autograd_non_functional_kernel=False,
# Every generated NativeFunction gets a "generated" tag, so it's easy to tell
# which NativeFunction objects did not come directly from native_functions.yaml.
tags=set(["generated"]) | (f.tags & {"nondeterministic_seeded"}),
namespace=f.namespace,
),
backend_metadata,
)
def generate_function(
f: NativeFunction, k: SchemaKind
) -> Tuple[NativeFunction, Dict[DispatchKey, Dict["OperatorName",
"BackendMetadata"]]]:
from torchgen.api import cpp
if k == SchemaKind.functional:
assert f.func.kind() != SchemaKind.functional
gets_composite_kernel = True
# The new "functional" NativeFunction has:
# - any mutable arguments have been converted into (immutable) returns.
# (if a mutable argument was not also a return, it gets converted to one)
# - a "functional" overload name.
# The default grouping logic in signature() actually already does this,
# so we can piggy-back off it (but we still want return names)
func = f.func.signature(keep_return_names=True).with_name(
f.func.name.remove_inplace().with_overload(
"functional" if not f.func.name.overload_name else
f"{f.func.name.overload_name}_functional"))
elif k == SchemaKind.out:
# We generate out= ops mostly just so that we can pair up NativeFunctions into groups easily,
# but at least today, there is no good reason to actually use them.
# we'll generate a dispatcher entry for them, but won't actually register any kernels for them.
gets_composite_kernel = False
if f.func.kind() == SchemaKind.inplace:
func = self_to_out_signature(f.func)
elif f.func.kind() == SchemaKind.mutable:
func = mutable_to_out_signature(f.func)
else:
raise AssertionError(
"We only bother generating out= functions from either inplace or mutable variants"
)
else:
raise AssertionError(
"We currently only generate either functional or out= NativeFunctions"
)
if gets_composite_kernel:
backend_metadata = {
DispatchKey.CompositeExplicitAutograd: {
func.name: BackendMetadata(cpp.name(func), structured=False)
}
}
else:
backend_metadata = {}
return (
NativeFunction(
func=func,
use_const_ref_for_mutable_tensors=f.
use_const_ref_for_mutable_tensors,
# These generated fn's aren't meant to be user friendly- don't generate methods.
variants=set([Variant.function]),
structured=False,
structured_delegate=None,
structured_inherits=None,
precomputed=None,
autogen=[],
ufunc_inner_loop={},
manual_kernel_registration=False,
manual_cpp_binding=False,
python_module=None,
category_override=None,
device_guard=False,
device_check=DeviceCheckType.NoCheck,
loc=f.loc,
cpp_no_default_args=set(),
is_abstract=f.is_abstract,
has_composite_implicit_autograd_kernel=False,
has_composite_explicit_autograd_kernel=gets_composite_kernel,
# Every generated NativeFunction gets a "generated" tag, so it's easy to tell
# which NativeFunction objects did not come directly from native_functions.yaml.
tags=set(["generated"]),
),
backend_metadata,
)