808 lines
24 KiB
Python
808 lines
24 KiB
Python
import logging
|
|
from collections.abc import Callable, Iterable
|
|
from dataclasses import dataclass, field
|
|
from typing import Concatenate, ParamSpec, TypeVar
|
|
|
|
import torch.library
|
|
|
|
|
|
__all__ = [
|
|
"UserOrderingFn",
|
|
"register_op_override",
|
|
"reorder_graphs_from_user_function",
|
|
"reenable_op_overrides",
|
|
"deregister_op_overrides",
|
|
"get_dsl_operations",
|
|
]
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
P = ParamSpec("P")
|
|
R = TypeVar("R")
|
|
|
|
_OpOverrideFn = Callable[Concatenate[torch.DispatchKeySet, P], R]
|
|
_OpReplaceFn = Callable[P, R]
|
|
|
|
_OpFn = _OpOverrideFn | _OpReplaceFn
|
|
|
|
|
|
@dataclass
|
|
class _OverrideNode:
|
|
"""Track function override data."""
|
|
|
|
dsl_name: str
|
|
op_symbol: str
|
|
dispatch_key: str
|
|
override_fn: _OpFn
|
|
unconditional_override: bool = False
|
|
active: bool = True
|
|
|
|
|
|
UserOrderingFn = Callable[[str, str, list[_OverrideNode]], list[_OverrideNode]]
|
|
|
|
|
|
@dataclass
|
|
class _FilterState:
|
|
"""Manages filtering state for override nodes."""
|
|
|
|
_dsl_names: set[str] = field(default_factory=set)
|
|
_op_symbols: set[str] = field(default_factory=set)
|
|
_dispatch_keys: set[str] = field(default_factory=set)
|
|
|
|
def check_enabled(self, node: _OverrideNode) -> bool:
|
|
"""
|
|
Check if a node is enabled based on current filter state.
|
|
|
|
Args:
|
|
node: The override node to check
|
|
|
|
Returns:
|
|
bool: True if the node should be enabled, False if filtered out
|
|
"""
|
|
if node.dsl_name in self._dsl_names:
|
|
return False
|
|
|
|
if node.op_symbol in self._op_symbols:
|
|
return False
|
|
|
|
if node.dispatch_key in self._dispatch_keys:
|
|
return False
|
|
|
|
return True
|
|
|
|
def update(
|
|
self,
|
|
dsl_names: str | Iterable[str] | None,
|
|
op_symbols: str | Iterable[str] | None,
|
|
dispatch_keys: str | Iterable[str] | None,
|
|
remove_keys: bool = False,
|
|
) -> None:
|
|
"""
|
|
Update filter sets as (current | new) or (current ~ new).
|
|
|
|
Args:
|
|
dsl_names: DSL names to add/remove from filter
|
|
op_symbols: Operation symbols to add/remove from filter
|
|
dispatch_keys: Dispatch keys to add/remove from filter
|
|
remove_keys: If True, remove keys from filter; if False, add them
|
|
|
|
Note:
|
|
Uses set.discard as it doesn't raise an exception if the element
|
|
wasn't in the set to begin with.
|
|
"""
|
|
if remove_keys:
|
|
self._dsl_names -= set(_resolve_iterable(dsl_names))
|
|
self._op_symbols -= set(_resolve_iterable(op_symbols))
|
|
self._dispatch_keys -= set(_resolve_iterable(dispatch_keys))
|
|
else:
|
|
self._dsl_names |= set(_resolve_iterable(dsl_names))
|
|
self._op_symbols |= set(_resolve_iterable(op_symbols))
|
|
self._dispatch_keys |= set(_resolve_iterable(dispatch_keys))
|
|
|
|
def build_disable_key_set(self) -> set[tuple[str, str]]:
|
|
"""
|
|
Build a set of dictionary keys based on the current filter state.
|
|
|
|
Returns:
|
|
set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples
|
|
"""
|
|
return _build_key_set(
|
|
self._dsl_names,
|
|
self._op_symbols,
|
|
self._dispatch_keys,
|
|
)
|
|
|
|
def __str__(self) -> str:
|
|
"""Return string representation of filter state."""
|
|
s = ""
|
|
s += "Filter State:\n"
|
|
s += " === DSL: ===\n"
|
|
for i, dsl in enumerate(self._dsl_names):
|
|
s += f" {i}: {dsl}\n"
|
|
s += " === OP SYMBOL: ===\n"
|
|
for i, op in enumerate(self._op_symbols):
|
|
s += f" {i}: {op}\n"
|
|
s += " === DISPATCH KEYS: ===\n"
|
|
for i, key in enumerate(self._dispatch_keys):
|
|
s += f" {i}: {key}\n"
|
|
|
|
return s
|
|
|
|
|
|
# Store the global override filtering state
|
|
_filter_state: _FilterState = _FilterState()
|
|
|
|
# Store torch.library.Library instances
|
|
_libs: dict[tuple[str, str], torch.library.Library] = {}
|
|
|
|
# store graph structures
|
|
_GraphsType = dict[tuple[str, str], list[_OverrideNode]]
|
|
_graphs: _GraphsType = {}
|
|
|
|
_MappingType = dict[str, list[tuple[str, str]]]
|
|
|
|
# map a {dsl, op, dispatch_key} to keys to all graphs that contain it
|
|
_dsl_name_to_lib_graph: _MappingType = {}
|
|
_dispatch_key_to_lib_graph: _MappingType = {}
|
|
_op_symbol_to_lib_graph: _MappingType = {}
|
|
|
|
|
|
def _build_key_set(
|
|
dsl_names: str | Iterable[str] | None,
|
|
op_symbols: str | Iterable[str] | None,
|
|
dispatch_keys: str | Iterable[str] | None,
|
|
) -> set[tuple[str, str]]:
|
|
"""
|
|
Build a set of dictionary keys based on filter criteria.
|
|
|
|
Args:
|
|
dsl_names: DSL names to include in key set
|
|
op_symbols: Operation symbols to include in key set
|
|
dispatch_keys: Dispatch keys to include in key set
|
|
|
|
Returns:
|
|
set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples
|
|
"""
|
|
key_set: set[tuple[str, str]] = set()
|
|
|
|
def _append_to_set(
|
|
entries: str | Iterable[str] | None, graph_lib_dict: _MappingType
|
|
) -> None:
|
|
"""Helper to add matching keys from graph_lib_dict to key_set."""
|
|
resolved_entries = _resolve_iterable(entries)
|
|
|
|
for entry in resolved_entries:
|
|
if entry in graph_lib_dict:
|
|
for key in graph_lib_dict[entry]:
|
|
key_set.add(key)
|
|
|
|
_append_to_set(dsl_names, _dsl_name_to_lib_graph)
|
|
_append_to_set(op_symbols, _op_symbol_to_lib_graph)
|
|
_append_to_set(dispatch_keys, _dispatch_key_to_lib_graph)
|
|
|
|
return key_set
|
|
|
|
|
|
def _print_override_graphs(*, print_inactive: bool = False) -> None:
|
|
"""
|
|
Print all override graphs for debugging purposes.
|
|
|
|
Args:
|
|
print_inactive: Whether to print inactive nodes
|
|
"""
|
|
for (op, key), node_list in _graphs.items():
|
|
print(f"{op=}, {key=}")
|
|
|
|
for i, node in enumerate(node_list):
|
|
if node.active or print_inactive:
|
|
s: str = f" {i}: {node.dsl_name=}, {node.unconditional_override=}"
|
|
if print_inactive:
|
|
s += f" {node.active=}"
|
|
|
|
print(s)
|
|
|
|
|
|
def _get_or_create_library(op_symbol: str, dispatch_key: str) -> torch.library.Library:
|
|
"""
|
|
Get or create a torch.library.Library instance for the given key.
|
|
|
|
Args:
|
|
op_symbol: The operation symbol
|
|
dispatch_key: The dispatch key
|
|
|
|
Returns:
|
|
torch.library.Library: The library instance
|
|
"""
|
|
global _libs
|
|
|
|
key = (op_symbol, dispatch_key)
|
|
if key not in _libs:
|
|
_libs[key] = torch.library.Library("aten", "IMPL", dispatch_key)
|
|
|
|
return _libs[key]
|
|
|
|
|
|
def _register_node_impl(
|
|
lib: torch.library.Library, node: _OverrideNode, dispatch_key: str
|
|
) -> None:
|
|
"""
|
|
Register a single node implementation with the library.
|
|
|
|
Args:
|
|
lib: The torch.library.Library instance
|
|
node: The override node to register
|
|
dispatch_key: The dispatch key for registration
|
|
"""
|
|
lib.impl(
|
|
node.op_symbol,
|
|
node.override_fn,
|
|
dispatch_key,
|
|
with_keyset=not node.unconditional_override,
|
|
allow_override=True,
|
|
)
|
|
|
|
|
|
def _resolve_iterable(iterable: str | Iterable[str] | None) -> Iterable[str]:
|
|
"""
|
|
Resolve various input types to a consistent iterable of strings.
|
|
|
|
Args:
|
|
iterable: String, iterable of strings, or None
|
|
|
|
Returns:
|
|
Iterable[str]: Consistent iterable output
|
|
"""
|
|
if iterable is None:
|
|
return []
|
|
|
|
if not isinstance(iterable, Iterable) or isinstance(iterable, str):
|
|
return (iterable,)
|
|
|
|
return iterable
|
|
|
|
|
|
def reenable_op_overrides(
|
|
*,
|
|
enable_dsl_names: str | list[str] | None = None,
|
|
enable_op_symbols: str | list[str] | None = None,
|
|
enable_dispatch_keys: str | list[str] | None = None,
|
|
) -> None:
|
|
"""
|
|
Re-enable overrides by removing them from filter state and reregistering.
|
|
|
|
Args:
|
|
enable_dsl_names: DSL names to re-enable
|
|
enable_op_symbols: Operation symbols to re-enable
|
|
enable_dispatch_keys: Dispatch keys to re-enable
|
|
|
|
Note:
|
|
This function uses reverse filter state management (removing from
|
|
filters to enable).
|
|
"""
|
|
log.info(
|
|
"Re-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %s",
|
|
enable_dsl_names,
|
|
enable_op_symbols,
|
|
enable_dispatch_keys,
|
|
)
|
|
|
|
# Update the filters - note `remove_keys=True` because
|
|
# we are removing keys from the filters (vs. adding them)
|
|
_filter_state.update(
|
|
enable_dsl_names,
|
|
enable_op_symbols,
|
|
enable_dispatch_keys,
|
|
remove_keys=True,
|
|
)
|
|
|
|
# Get the set of keys that need to be reprocessed
|
|
key_set: set[tuple[str, str]] = _build_key_set(
|
|
enable_dsl_names,
|
|
enable_op_symbols,
|
|
enable_dispatch_keys,
|
|
)
|
|
|
|
# Process each affected graph with updated filter state
|
|
for key in key_set:
|
|
op_symbol, dispatch_key = key
|
|
|
|
if key in _graphs:
|
|
# Note: We don't need to cleanup and recreate the library here
|
|
# since we're just updating the registration with new filter state
|
|
_register_overrides_from_graph(
|
|
op_symbol, dispatch_key, _graphs[key], filter_state=_filter_state
|
|
)
|
|
|
|
|
|
def deregister_op_overrides(
|
|
*,
|
|
disable_dsl_names: str | list[str] | None = None,
|
|
disable_op_symbols: str | list[str] | None = None,
|
|
disable_dispatch_keys: str | list[str] | None = None,
|
|
) -> None:
|
|
"""
|
|
De-register overrides by updating filter state and reregistering graphs.
|
|
|
|
Args:
|
|
disable_dsl_names: DSL names to disable
|
|
disable_op_symbols: Operation symbols to disable
|
|
disable_dispatch_keys: Dispatch keys to disable
|
|
|
|
Note:
|
|
This function uses filter state management to selectively disable
|
|
operations.
|
|
"""
|
|
log.info(
|
|
"De-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %s",
|
|
disable_dsl_names,
|
|
disable_op_symbols,
|
|
disable_dispatch_keys,
|
|
)
|
|
|
|
# Update filter state to disable specified entries
|
|
_filter_state.update(disable_dsl_names, disable_op_symbols, disable_dispatch_keys)
|
|
|
|
# Get the set of keys that need to be reprocessed
|
|
key_set: set[tuple[str, str]] = _filter_state.build_disable_key_set()
|
|
|
|
# Process each affected graph with filter state
|
|
for key in key_set:
|
|
op_symbol, dispatch_key = key
|
|
|
|
if key in _graphs:
|
|
_cleanup_and_reregister_graph(
|
|
op_symbol,
|
|
dispatch_key,
|
|
_graphs[key],
|
|
filter_state=_filter_state,
|
|
)
|
|
|
|
|
|
def get_dsl_operations(dsl_name: str) -> list[str]:
|
|
"""Get list of operations registered by a specific DSL.
|
|
|
|
Args:
|
|
dsl_name: Name of the DSL to query.
|
|
|
|
Returns:
|
|
Sorted list of operation names registered by the DSL.
|
|
"""
|
|
operations = set()
|
|
for (op_symbol, _), nodes in _graphs.items():
|
|
for node in nodes:
|
|
if node.dsl_name == dsl_name:
|
|
operations.add(op_symbol)
|
|
break
|
|
return sorted(operations)
|
|
|
|
|
|
def _update_registration_maps(
|
|
dsl_name: str,
|
|
op_symbol: str,
|
|
dispatch_key: str,
|
|
key: tuple[str, str],
|
|
) -> None:
|
|
"""
|
|
Update the registration mapping dictionaries.
|
|
|
|
Args:
|
|
dsl_name: The DSL name
|
|
op_symbol: The operation symbol
|
|
dispatch_key: The dispatch key
|
|
key: The dictionary key tuple
|
|
"""
|
|
global _dsl_name_to_lib_graph
|
|
global _op_symbol_to_lib_graph
|
|
global _dispatch_key_to_lib_graph
|
|
|
|
def _get_new_entry_or_append(
|
|
registration: dict[str, list[tuple[str, str]]],
|
|
symbol: str,
|
|
key: tuple[str, str],
|
|
) -> None:
|
|
"""Helper to add key to registration list or create new entry."""
|
|
entry_list = registration.get(symbol)
|
|
|
|
if entry_list is None:
|
|
entry_list = [key]
|
|
registration[symbol] = entry_list
|
|
else:
|
|
entry_list.append(key)
|
|
|
|
_get_new_entry_or_append(_dsl_name_to_lib_graph, dsl_name, key)
|
|
_get_new_entry_or_append(_op_symbol_to_lib_graph, op_symbol, key)
|
|
_get_new_entry_or_append(_dispatch_key_to_lib_graph, dispatch_key, key)
|
|
|
|
|
|
def register_op_override(
|
|
backend: str,
|
|
lib_symbol: str,
|
|
op_symbol: str,
|
|
dispatch_key: str,
|
|
impl: _OpOverrideFn | _OpReplaceFn,
|
|
*,
|
|
allow_multiple_override: bool = False,
|
|
unconditional_override: bool = False,
|
|
) -> None:
|
|
"""
|
|
Register a passed override function to the dispatcher.
|
|
|
|
Actually a graph-building operation; real registration happens later.
|
|
|
|
Args:
|
|
backend: The backend name (DSL name)
|
|
lib_symbol: Library you're overriding symbols in (must be "aten")
|
|
op_symbol: Name of the operation you're overriding
|
|
dispatch_key: Dispatch key to override
|
|
impl: Implementation function for the override
|
|
allow_multiple_override: Allow overriding an existing override
|
|
unconditional_override: Implementation doesn't have a fallback and
|
|
doesn't require torch.DispatchKeySet as the first argument
|
|
|
|
Raises:
|
|
ValueError: If lib_symbol is not "aten"
|
|
"""
|
|
if lib_symbol != "aten":
|
|
raise ValueError(f'Unsupported lib_symbol (must be "aten", got: "{lib_symbol}"')
|
|
|
|
key = (op_symbol, dispatch_key)
|
|
|
|
global _graphs
|
|
op_graph = _graphs.get(key, [])
|
|
|
|
op_graph.append(
|
|
_OverrideNode(
|
|
dsl_name=backend,
|
|
op_symbol=op_symbol,
|
|
dispatch_key=dispatch_key,
|
|
override_fn=impl,
|
|
unconditional_override=unconditional_override,
|
|
)
|
|
)
|
|
_graphs[key] = op_graph
|
|
# Build additional maps helpful for de-registration
|
|
_update_registration_maps(backend, op_symbol, dispatch_key, key=key)
|
|
|
|
|
|
def _should_reregister_graph(
|
|
original_graph: list[_OverrideNode],
|
|
new_graph: list[_OverrideNode],
|
|
*,
|
|
force_reregister: bool = False,
|
|
) -> bool:
|
|
"""
|
|
Determine if a graph needs reregistration based on changes.
|
|
|
|
Args:
|
|
original_graph: The original graph before modification
|
|
new_graph: The graph after modification
|
|
force_reregister: If True, always reregister regardless of changes
|
|
|
|
Returns:
|
|
bool: True if reregistration is needed
|
|
"""
|
|
if force_reregister:
|
|
return True
|
|
|
|
# Check if the graph structure has changed
|
|
return original_graph != new_graph
|
|
|
|
|
|
def _cleanup_and_reregister_graph(
|
|
op_symbol: str,
|
|
dispatch_key: str,
|
|
graph: list[_OverrideNode],
|
|
*,
|
|
filter_state: _FilterState | None = None,
|
|
) -> None:
|
|
"""
|
|
Clean up existing library and reregister a graph.
|
|
|
|
This is the common pattern used across reorder, deregister, and reenable operations.
|
|
|
|
Args:
|
|
op_symbol: The operation symbol
|
|
dispatch_key: The dispatch key
|
|
graph: The graph to register
|
|
filter_state: Optional filter state for conditional registration
|
|
"""
|
|
key = (op_symbol, dispatch_key)
|
|
|
|
# Remove existing library if it exists
|
|
if key in _libs:
|
|
del _libs[key]
|
|
|
|
# Only create a library if the graph has nodes
|
|
# Empty graphs (disabled operations) shouldn't get libraries
|
|
if graph:
|
|
_register_overrides_from_graph(
|
|
op_symbol,
|
|
dispatch_key,
|
|
graph,
|
|
filter_state=filter_state,
|
|
)
|
|
|
|
|
|
def _apply_graph_transformation(
|
|
transformation_fn: UserOrderingFn,
|
|
*,
|
|
keys_to_process: set[tuple[str, str]] | None = None,
|
|
reregister_overrides: bool = False,
|
|
filter_state: _FilterState | None = None,
|
|
) -> None:
|
|
"""
|
|
Apply a transformation function to graphs and optionally reregister.
|
|
|
|
This is the core pattern used by reorder_graphs_from_user_function and
|
|
can be reused for other graph transformation operations.
|
|
|
|
Args:
|
|
transformation_fn: Function to transform each graph
|
|
keys_to_process: Keys to process, or None for all graphs
|
|
reregister_overrides: Whether to reregister changed graphs
|
|
filter_state: Optional filter state for conditional registration
|
|
|
|
Note:
|
|
If transformation_fn raises an exception for a specific graph, that graph
|
|
will be skipped and processing will continue with remaining graphs.
|
|
"""
|
|
global _graphs
|
|
|
|
# Determine which graphs to process
|
|
target_keys = (
|
|
keys_to_process if keys_to_process is not None else set(_graphs.keys())
|
|
)
|
|
|
|
# Process each graph
|
|
for op_symbol, dispatch_key in list(target_keys):
|
|
if (op_symbol, dispatch_key) not in _graphs:
|
|
continue # Skip if graph doesn't exist
|
|
|
|
original_graph = list(_graphs[(op_symbol, dispatch_key)])
|
|
|
|
# Apply the transformation with error handling
|
|
try:
|
|
new_graph = transformation_fn(op_symbol, dispatch_key, original_graph)
|
|
except (TypeError, ValueError, AttributeError, RuntimeError):
|
|
log.warning(
|
|
"Graph transformation failed for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
exc_info=True,
|
|
)
|
|
continue
|
|
except Exception:
|
|
log.exception(
|
|
"Unexpected error in graph transformation for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
)
|
|
continue
|
|
|
|
# Validate that the transformation returned a valid result
|
|
if not isinstance(new_graph, list):
|
|
log.warning(
|
|
"Graph transformation returned invalid type %s for %s/%s. Expected list. Preserving original graph.",
|
|
type(new_graph).__name__,
|
|
op_symbol,
|
|
dispatch_key,
|
|
)
|
|
continue
|
|
|
|
# Update the graph
|
|
_graphs[(op_symbol, dispatch_key)] = new_graph
|
|
|
|
# Reregister if needed
|
|
if reregister_overrides and _should_reregister_graph(
|
|
original_graph, new_graph, force_reregister=False
|
|
):
|
|
_cleanup_and_reregister_graph(
|
|
op_symbol,
|
|
dispatch_key,
|
|
new_graph,
|
|
filter_state=filter_state,
|
|
)
|
|
|
|
|
|
def _register_overrides_from_graph(
|
|
op_symbol: str,
|
|
dispatch_key: str,
|
|
graph: list[_OverrideNode],
|
|
*,
|
|
filter_state: _FilterState | None = None,
|
|
) -> None:
|
|
"""
|
|
Register all overrides in a single graph.
|
|
|
|
Args:
|
|
op_symbol: The operation symbol
|
|
dispatch_key: The dispatch key
|
|
graph: List of override nodes to register
|
|
filter_state: Optional filter state for conditional registration
|
|
"""
|
|
key = (op_symbol, dispatch_key)
|
|
lib = _get_or_create_library(*key)
|
|
|
|
for node in graph:
|
|
enable = True
|
|
if filter_state:
|
|
enable = filter_state.check_enabled(node)
|
|
|
|
if enable:
|
|
_register_node_impl(lib, node, dispatch_key)
|
|
node.active = True
|
|
else:
|
|
node.active = False
|
|
|
|
|
|
def _register_all_overrides() -> None:
|
|
"""
|
|
Perform all registration calls from previously-built override graphs.
|
|
"""
|
|
for key, graph in _graphs.items():
|
|
op_symbol, dispatch_key = key
|
|
|
|
_register_overrides_from_graph(
|
|
op_symbol,
|
|
dispatch_key,
|
|
graph,
|
|
)
|
|
|
|
|
|
def reorder_graphs_from_user_function(
|
|
fn: UserOrderingFn,
|
|
*,
|
|
reregister_overrides: bool = False,
|
|
) -> None:
|
|
"""
|
|
Reorder override graphs using a user-provided ordering function.
|
|
|
|
Args:
|
|
fn: User-provided function that takes (op_symbol, dispatch_key, graph)
|
|
and returns a reordered graph
|
|
reregister_overrides: Whether to reregister graphs that have changed
|
|
|
|
Note:
|
|
This function uses the common graph transformation pattern and can serve
|
|
as an example for other graph manipulation operations.
|
|
"""
|
|
_apply_graph_transformation(
|
|
transformation_fn=fn,
|
|
reregister_overrides=reregister_overrides,
|
|
)
|
|
|
|
|
|
def _apply_graph_filter(
|
|
filter_fn: Callable[[str, str, _OverrideNode], bool],
|
|
*,
|
|
reregister_overrides: bool = False,
|
|
) -> None:
|
|
"""
|
|
Apply a filter function to remove nodes from graphs.
|
|
|
|
This is a convenience function that uses the graph transformation pattern
|
|
to filter out unwanted nodes.
|
|
|
|
Args:
|
|
filter_fn: Function that takes (op_symbol, dispatch_key, node) and
|
|
returns True to keep the node, False to remove it
|
|
reregister_overrides: Whether to reregister modified graphs
|
|
|
|
Example:
|
|
# Remove all nodes with "deprecated" in the DSL name
|
|
_apply_graph_filter(
|
|
lambda op, dk, node: "deprecated" not in node.dsl_name,
|
|
reregister_overrides=True
|
|
)
|
|
|
|
Note:
|
|
If filter_fn raises an exception for a specific graph, the original
|
|
graph will be preserved and processing will continue.
|
|
"""
|
|
|
|
def filtering_transformation(
|
|
op_symbol: str, dispatch_key: str, graph: list[_OverrideNode]
|
|
) -> list[_OverrideNode]:
|
|
"""Apply filter_fn to graph with error handling."""
|
|
try:
|
|
return [node for node in graph if filter_fn(op_symbol, dispatch_key, node)]
|
|
except (TypeError, ValueError, AttributeError, RuntimeError):
|
|
log.warning(
|
|
"Graph transformation failed for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
exc_info=True,
|
|
)
|
|
return graph
|
|
except Exception:
|
|
log.exception(
|
|
"Unexpected error in graph transformation for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
)
|
|
return graph
|
|
|
|
_apply_graph_transformation(
|
|
transformation_fn=filtering_transformation,
|
|
reregister_overrides=reregister_overrides,
|
|
)
|
|
|
|
|
|
def _apply_selective_reordering(
|
|
condition_fn: Callable[[str, str], bool],
|
|
ordering_fn: UserOrderingFn,
|
|
*,
|
|
reregister_overrides: bool = False,
|
|
) -> None:
|
|
"""
|
|
Apply reordering only to graphs that match a condition.
|
|
|
|
This allows for more targeted reordering operations.
|
|
|
|
Args:
|
|
condition_fn: Function that takes (op_symbol, dispatch_key) and
|
|
returns True if the graph should be reordered
|
|
ordering_fn: Ordering function to apply to matching graphs
|
|
reregister_overrides: Whether to reregister modified graphs
|
|
|
|
Example:
|
|
# Only reorder CUDA operations
|
|
_apply_selective_reordering(
|
|
condition_fn=lambda op, dk: dk == "CUDA",
|
|
ordering_fn=lambda op, dk, g: sorted(g, key=lambda n: n.dsl_name),
|
|
reregister_overrides=True
|
|
)
|
|
|
|
Note:
|
|
If condition_fn or ordering_fn raises an exception for a specific graph,
|
|
the original graph will be preserved and processing will continue.
|
|
"""
|
|
|
|
def conditional_transformation(
|
|
op_symbol: str, dispatch_key: str, graph: list[_OverrideNode]
|
|
) -> list[_OverrideNode]:
|
|
"""Apply ordering_fn conditionally based on condition_fn result."""
|
|
try:
|
|
should_reorder = condition_fn(op_symbol, dispatch_key)
|
|
except (TypeError, ValueError, AttributeError, RuntimeError):
|
|
log.warning(
|
|
"Graph transformation failed for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
exc_info=True,
|
|
)
|
|
return graph
|
|
except Exception:
|
|
log.exception(
|
|
"Unexpected error in graph transformation for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
)
|
|
return graph
|
|
|
|
if should_reorder:
|
|
try:
|
|
return ordering_fn(op_symbol, dispatch_key, graph)
|
|
except (TypeError, ValueError, AttributeError, RuntimeError):
|
|
log.warning(
|
|
"Graph transformation failed for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
exc_info=True,
|
|
)
|
|
return graph
|
|
except Exception:
|
|
log.exception(
|
|
"Unexpected error in graph transformation for %s/%s. Preserving original graph.",
|
|
op_symbol,
|
|
dispatch_key,
|
|
)
|
|
return graph
|
|
|
|
return graph # Return unchanged if condition doesn't match
|
|
|
|
_apply_graph_transformation(
|
|
transformation_fn=conditional_transformation,
|
|
reregister_overrides=reregister_overrides,
|
|
)
|