def parse_native_functions_keys(
backend_yaml_path: str,
grouped_native_functions: Sequence[Union[NativeFunction,
NativeFunctionsGroup]],
) -> Tuple[List[OperatorName], List[Any], 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", [])
non_native = yaml_values.pop("non_native", [])
ir_gen = yaml_values.pop("ir_gen", [])
assert isinstance(full_codegen, list)
assert isinstance(non_native, list)
assert isinstance(ir_gen, list)
full_codegen_opnames = [OperatorName.parse(name) for name in full_codegen]
ir_gen_opnames = [OperatorName.parse(name) for name in ir_gen]
return full_codegen_opnames, non_native, ir_gen_opnames
def create_backend_index(
backend_ops: List[str],
dispatch_key: DispatchKey,
*,
use_out_as_primary: bool,
use_device_guard: bool,
) -> 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,
cpp_namespace=cpp_namespace)
metadata[op_name] = m
return BackendIndex(
dispatch_key=dispatch_key,
use_out_as_primary=use_out_as_primary,
external=True,
symint=True, # TODO: make this configurable
device_guard=use_device_guard,
index=metadata,
)
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 print_op_str_if_not_supported(op_str):
op = OperatorName.parse(op_str)
packet = getattr(torch.ops.aten, str(op.name))
overload = getattr(packet, op.overload_name if op.overload_name else "default")
if any(overload in d for d in [meta_dispatch_skips, meta_dispatch_device_skips['cuda']]):
print(f"{overload} # SKIP")
if any(overload in d for d in [meta_dispatch_expected_failures, meta_dispatch_device_expected_failures['cuda']]):
print(overload)