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)