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)
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)
metadata[op_name] = m
return BackendIndex(dispatch_key=dispatch_key,
use_out_as_primary=use_out_as_primary,
external=True,
device_guard=use_device_guard,
index=metadata)