Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
from functools import cache
|
||||
from typing import cast
|
||||
|
||||
# This handles collecting registration of all native ops
|
||||
# Also need to import DSL utils to make sure DSL registration is ok
|
||||
from . import cutedsl_utils, dsl_registry, ops, registry, triton_utils
|
||||
|
||||
|
||||
@cache
|
||||
def get_user_ordering_fn() -> registry.UserOrderingFn | None:
|
||||
"""
|
||||
Get a user-supplied graph-ordering function if specified.
|
||||
|
||||
Pass in a `package.submodule.fn` string to the env variable
|
||||
`TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN` that implements the
|
||||
calling API described in `torch/_native/README.md`. This function
|
||||
must be part of an importable path.
|
||||
|
||||
Return either the imported function or `None`
|
||||
"""
|
||||
env_var = os.getenv("TORCH_PYTHON_NATIVE_USER_GRAPH_ORDER_FN")
|
||||
|
||||
if not env_var:
|
||||
return None
|
||||
|
||||
try:
|
||||
import importlib
|
||||
|
||||
# Split into "package.submodule.fn_name
|
||||
module_name, fn_name = env_var.rsplit(".", 1)
|
||||
|
||||
module = importlib.import_module(module_name)
|
||||
fn = getattr(module, fn_name)
|
||||
|
||||
if not callable(fn):
|
||||
raise TypeError(f"{env_var} does not describe a callable function")
|
||||
|
||||
# Cast needed: getattr returns object, but we've verified fn is callable with correct signature
|
||||
return cast(registry.UserOrderingFn, fn)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Could not resolve {env_var} into an importable & callable function"
|
||||
) from e
|
||||
|
||||
|
||||
user_order_fn = get_user_ordering_fn()
|
||||
if user_order_fn:
|
||||
registry.reorder_graphs_from_user_function(user_order_fn)
|
||||
|
||||
|
||||
# Actually perform all registrations
|
||||
registry._register_all_overrides()
|
||||
@@ -0,0 +1,64 @@
|
||||
import importlib
|
||||
import importlib.metadata
|
||||
import os
|
||||
from functools import cache
|
||||
|
||||
from torch._vendor.packaging import version as _packaging_version
|
||||
|
||||
|
||||
@cache
|
||||
def check_native_jit_disabled() -> bool:
|
||||
"""
|
||||
Single point to check if native DSL ops are disabled globally,
|
||||
checked via:
|
||||
TORCH_DISABLE_NATIVE_JIT=1
|
||||
"""
|
||||
return int(os.getenv("TORCH_DISABLE_NATIVE_JIT", 0)) == 1
|
||||
|
||||
|
||||
def _unavailable_reason(deps: list[tuple[str, str]]) -> None | str:
|
||||
"""
|
||||
Check availability of required packages - cuteDSL & deps,
|
||||
informing user what (if anything) is missing
|
||||
|
||||
NOTE: Doesn't actually import anything.
|
||||
"""
|
||||
for package_name, module_name in deps:
|
||||
# Note this doesn't actually import the packages
|
||||
if importlib.util.find_spec(module_name) is None:
|
||||
return (
|
||||
f"missing optional dependency `{package_name}` "
|
||||
f"(importlib.util.find_spec({package_name}) failed)"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _available_version(package: str) -> _packaging_version.Version | None:
|
||||
"""
|
||||
Get the installed version of a package as (major, minor, patch).
|
||||
|
||||
Handles pre-release suffixes like "0.7.0rc1" or "3.1.0.post1" by
|
||||
stripping non-numeric tails from each component. Returns None on
|
||||
parse failure.
|
||||
"""
|
||||
try:
|
||||
version = importlib.metadata.version(package)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return None
|
||||
|
||||
try:
|
||||
v = _packaging_version.parse(version)
|
||||
except _packaging_version.InvalidVersion:
|
||||
return None
|
||||
|
||||
return v
|
||||
|
||||
|
||||
@cache
|
||||
def check_native_version_skip() -> bool:
|
||||
"""
|
||||
Single point to check if native DSL version gating should be skipped,
|
||||
checked via:
|
||||
TORCH_NATIVE_SKIP_VERSION_CHECK=1
|
||||
"""
|
||||
return int(os.getenv("TORCH_NATIVE_SKIP_VERSION_CHECK", 0)) == 1
|
||||
@@ -0,0 +1,132 @@
|
||||
import functools
|
||||
import logging
|
||||
import sys
|
||||
from typing import cast
|
||||
|
||||
from torch._vendor.packaging.version import Version
|
||||
|
||||
from ..backends import cuda as _cuda
|
||||
from .common_utils import (
|
||||
_available_version,
|
||||
_unavailable_reason,
|
||||
check_native_jit_disabled,
|
||||
check_native_version_skip,
|
||||
)
|
||||
from .dsl_registry import dsl_registry, DSLModuleProtocol
|
||||
from .registry import (
|
||||
_OpFn,
|
||||
deregister_op_overrides as _deregister_op_overrides_impl,
|
||||
register_op_override as _register_op_override_impl,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_CUTEDSL_DSL_NAME = "cutedsl"
|
||||
_CUTEDSL_REQUIRED_VERSIONS: set[Version] = {
|
||||
# Current version - Note Version.from_part(release=(4.4.1)) is better
|
||||
# but > v26 of packaging.
|
||||
Version(f"{4}.{4}.{1}"),
|
||||
Version(f"{4}.{4}.{2}"),
|
||||
}
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _check_runtime_available() -> tuple[bool, Version | None]:
|
||||
"""
|
||||
Check if cutedsl (and deps) are available.
|
||||
|
||||
NOTE: Doesn't import at this point
|
||||
"""
|
||||
# Skip all checks if running on CPU-only binary
|
||||
if not _cuda.is_built():
|
||||
return (False, None)
|
||||
|
||||
deps = [
|
||||
("nvidia_cutlass_dsl", "cutlass"),
|
||||
("apache_tvm_ffi", "tvm_ffi"),
|
||||
]
|
||||
reason = _unavailable_reason(deps)
|
||||
if reason is None:
|
||||
available = True
|
||||
version = _available_version("nvidia_cutlass_dsl")
|
||||
else:
|
||||
log.warning(
|
||||
"CuTeDSL operators require optional Python packages "
|
||||
"`nvidia-cutlass-dsl` and `apache-tvm-ffi`; "
|
||||
"%s",
|
||||
reason,
|
||||
)
|
||||
available = False
|
||||
version = None
|
||||
return available, version
|
||||
|
||||
|
||||
def runtime_available() -> bool:
|
||||
available, _ = _check_runtime_available()
|
||||
return available
|
||||
|
||||
|
||||
def runtime_version() -> None | Version:
|
||||
_, version = _check_runtime_available()
|
||||
return version
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _version_is_ok() -> bool:
|
||||
_, version = _check_runtime_available()
|
||||
if check_native_version_skip() or (version in _CUTEDSL_REQUIRED_VERSIONS):
|
||||
return True
|
||||
|
||||
log.warning(
|
||||
"cutedsl version %s is not known-good (ok: %s); "
|
||||
"set TORCH_NATIVE_SKIP_VERSION_CHECK=1 to override",
|
||||
version,
|
||||
_CUTEDSL_REQUIRED_VERSIONS,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def deregister_op_overrides() -> None:
|
||||
"""
|
||||
Deregister all ops through cuteDSL
|
||||
"""
|
||||
_deregister_op_overrides_impl(disable_dsl_names=_CUTEDSL_DSL_NAME)
|
||||
|
||||
|
||||
def register_op_override(
|
||||
lib_symbol: str,
|
||||
op_symbol: str,
|
||||
dispatch_key: str,
|
||||
impl: _OpFn,
|
||||
*,
|
||||
allow_multiple_override: bool = False,
|
||||
unconditional_override: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
See torch/_native/registry.py for the underlying implementation
|
||||
and arguments. This is a thin, DSL-checking wrapper over
|
||||
_register_op_override_impl
|
||||
"""
|
||||
available, version = _check_runtime_available()
|
||||
if (not available) or check_native_jit_disabled():
|
||||
return
|
||||
|
||||
if not _version_is_ok():
|
||||
return
|
||||
|
||||
_register_op_override_impl(
|
||||
_CUTEDSL_DSL_NAME,
|
||||
lib_symbol,
|
||||
op_symbol,
|
||||
dispatch_key,
|
||||
impl,
|
||||
allow_multiple_override=allow_multiple_override,
|
||||
unconditional_override=unconditional_override,
|
||||
)
|
||||
|
||||
|
||||
# Register this DSL module with the registry
|
||||
# Note: Import-time registration ensures DSL is available when module is loaded
|
||||
dsl_registry.register_dsl("cutedsl", cast(DSLModuleProtocol, sys.modules[__name__]))
|
||||
@@ -0,0 +1,153 @@
|
||||
# Owner(s): ["module: dsl-native-ops"]
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from typing import Protocol
|
||||
|
||||
from torch._vendor.packaging.version import Version
|
||||
|
||||
from .registry import _OpFn
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DSLModuleProtocol(Protocol):
|
||||
"""Complete interface for DSL utility modules"""
|
||||
|
||||
def runtime_available(self) -> bool: ...
|
||||
def runtime_version(self) -> Version | None: ...
|
||||
|
||||
def deregister_op_overrides(self) -> None: ...
|
||||
|
||||
def register_op_override(
|
||||
self,
|
||||
lib_symbol: str,
|
||||
op_symbol: str,
|
||||
dispatch_key: str,
|
||||
impl: _OpFn,
|
||||
*,
|
||||
allow_multiple_override: bool = False,
|
||||
unconditional_override: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class DSLRegistry:
|
||||
"""Registry for DSL modules - calls their existing API functions dynamically"""
|
||||
|
||||
def __init__(self):
|
||||
self._dsl_modules: dict[str, DSLModuleProtocol] = {}
|
||||
|
||||
def _validate_dsl_name(self, name: str) -> None:
|
||||
"""Validate DSL name at runtime"""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(f"DSL name must be string, got {type(name).__name__}")
|
||||
|
||||
if not name.strip():
|
||||
raise ValueError("DSL name cannot be empty or whitespace")
|
||||
|
||||
def register_dsl(self, name: str, dsl_module: DSLModuleProtocol) -> None:
|
||||
"""Register a DSL module with required interface"""
|
||||
# Runtime validation for name and module interface
|
||||
self._validate_dsl_name(name)
|
||||
|
||||
# Validate that module implements the protocol
|
||||
required_methods = [
|
||||
"runtime_available",
|
||||
"runtime_version",
|
||||
"register_op_override",
|
||||
"deregister_op_overrides",
|
||||
]
|
||||
missing_methods = [
|
||||
method for method in required_methods if not hasattr(dsl_module, method)
|
||||
]
|
||||
if missing_methods:
|
||||
raise TypeError(
|
||||
f"DSL module '{name}' missing required methods: {missing_methods}"
|
||||
)
|
||||
|
||||
# Handle duplicate registration case
|
||||
if name in self._dsl_modules:
|
||||
existing_module = self._dsl_modules[name]
|
||||
if existing_module is dsl_module:
|
||||
# Same module re-registering - this is OK (import-time registration)
|
||||
log.debug(
|
||||
"DSL '%s' re-registered with same module",
|
||||
name,
|
||||
)
|
||||
return
|
||||
else:
|
||||
# Different module object but same name - warn and allow (for testing)
|
||||
# This can happen when tests import modules directly
|
||||
log.warning(
|
||||
"DSL '%s' re-registered with different module object (possibly from test imports)",
|
||||
name,
|
||||
)
|
||||
# Continue to allow the registration
|
||||
|
||||
# No cast needed - already properly typed
|
||||
self._dsl_modules[name] = dsl_module
|
||||
|
||||
# Clear caches to prevent stale results after registration
|
||||
self.is_dsl_available.cache_clear()
|
||||
self.get_dsl_version.cache_clear()
|
||||
self.list_available_dsls.cache_clear()
|
||||
self.list_all_dsls.cache_clear()
|
||||
|
||||
log.info("Successfully registered DSL: %s", name)
|
||||
|
||||
@functools.cache # noqa: B019
|
||||
def is_dsl_available(self, dsl_name: str) -> bool:
|
||||
"""Check if DSL is available by calling its runtime_available()"""
|
||||
dsl_module = self._dsl_modules.get(dsl_name)
|
||||
if dsl_module is None:
|
||||
return False
|
||||
try:
|
||||
return dsl_module.runtime_available()
|
||||
except ImportError:
|
||||
log.debug("DSL %s import error", dsl_name, exc_info=True)
|
||||
return False
|
||||
except Exception:
|
||||
log.exception("Error checking availability for DSL %s", dsl_name)
|
||||
return False
|
||||
|
||||
@functools.cache # noqa: B019
|
||||
def get_dsl_version(self, dsl_name: str) -> Version | None:
|
||||
"""Get DSL version by calling its runtime_version()"""
|
||||
dsl_module = self._dsl_modules.get(dsl_name)
|
||||
if dsl_module is None:
|
||||
return None
|
||||
try:
|
||||
return dsl_module.runtime_version()
|
||||
except Exception:
|
||||
log.debug("Error getting version for DSL %s", dsl_name, exc_info=True)
|
||||
return None
|
||||
|
||||
@functools.cache # noqa: B019
|
||||
def list_available_dsls(self) -> tuple[str, ...]:
|
||||
"""Get names of currently available DSLs"""
|
||||
available = []
|
||||
for name in self._dsl_modules:
|
||||
if self.is_dsl_available(name): # Use cached method
|
||||
available.append(name)
|
||||
return tuple(available)
|
||||
|
||||
@functools.cache # noqa: B019
|
||||
def list_all_dsls(self) -> tuple[str, ...]:
|
||||
"""Get all registered DSL names (available or not)"""
|
||||
return tuple(self._dsl_modules.keys())
|
||||
|
||||
def get_dsl_module(self, name: str) -> DSLModuleProtocol | None:
|
||||
"""Get a registered DSL module by name.
|
||||
|
||||
Args:
|
||||
name: Name of the DSL to retrieve.
|
||||
|
||||
Returns:
|
||||
The DSL module if registered, None otherwise.
|
||||
"""
|
||||
return self._dsl_modules.get(name)
|
||||
|
||||
|
||||
# Global registry instance
|
||||
dsl_registry = DSLRegistry()
|
||||
@@ -0,0 +1,807 @@
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,131 @@
|
||||
import functools
|
||||
import logging
|
||||
import sys
|
||||
from typing import cast
|
||||
|
||||
from torch._vendor.packaging.version import Version
|
||||
|
||||
from ..backends import cuda as _cuda
|
||||
from .common_utils import (
|
||||
_available_version,
|
||||
_unavailable_reason,
|
||||
check_native_jit_disabled,
|
||||
check_native_version_skip,
|
||||
)
|
||||
from .dsl_registry import dsl_registry, DSLModuleProtocol
|
||||
from .registry import (
|
||||
_OpFn,
|
||||
deregister_op_overrides as _deregister_op_overrides_impl,
|
||||
register_op_override as _register_op_override_impl,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_TRITON_DSL_NAME = "triton"
|
||||
_TRITON_REQUIRED_VERSION_MAJOR = 3
|
||||
_TRITON_MINIMUM_VERSION_MINOR = 6
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _check_runtime_available() -> tuple[bool, Version | None]:
|
||||
"""
|
||||
Check if triton is available
|
||||
|
||||
NOTE: must not import at this point
|
||||
"""
|
||||
# Skip all checks if running on CPU-only binary
|
||||
if not _cuda.is_built():
|
||||
return (False, None)
|
||||
|
||||
deps = [
|
||||
("triton", "triton"),
|
||||
]
|
||||
reason = _unavailable_reason(deps)
|
||||
if reason is None:
|
||||
available = True
|
||||
version = _available_version("triton")
|
||||
else:
|
||||
log.warning("triton native DSL ops require: `triton` %s", reason)
|
||||
available = False
|
||||
version = None
|
||||
return available, version
|
||||
|
||||
|
||||
def runtime_available() -> bool:
|
||||
available, _ = _check_runtime_available()
|
||||
return available
|
||||
|
||||
|
||||
def runtime_version() -> None | Version:
|
||||
_, version = _check_runtime_available()
|
||||
return version
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _version_is_sufficient() -> bool:
|
||||
_, version = _check_runtime_available()
|
||||
|
||||
if version is None:
|
||||
return False
|
||||
|
||||
# Either exact version, or same major
|
||||
major_ok = version.major == _TRITON_REQUIRED_VERSION_MAJOR
|
||||
minor_ok = version.minor >= _TRITON_MINIMUM_VERSION_MINOR
|
||||
|
||||
if (major_ok and minor_ok) or check_native_version_skip():
|
||||
return True
|
||||
|
||||
log.warning(
|
||||
"triton version %s is not sufficient (>= (%s.%s.*)); "
|
||||
"set TORCH_NATIVE_SKIP_VERSION_CHECK=1 to override",
|
||||
version,
|
||||
_TRITON_REQUIRED_VERSION_MAJOR,
|
||||
_TRITON_MINIMUM_VERSION_MINOR,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def deregister_op_overrides() -> None:
|
||||
"""
|
||||
Deregister all ops through triton
|
||||
"""
|
||||
_deregister_op_overrides_impl(disable_dsl_names=_TRITON_DSL_NAME)
|
||||
|
||||
|
||||
def register_op_override(
|
||||
lib_symbol: str,
|
||||
op_symbol: str,
|
||||
dispatch_key: str,
|
||||
impl: _OpFn,
|
||||
*,
|
||||
allow_multiple_override: bool = False,
|
||||
unconditional_override: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
See torch/_native/registry.py for the underlying implementation
|
||||
and arguments. This is a thin, DSL-checking wrapper over
|
||||
_register_op_override_impl
|
||||
"""
|
||||
available, version = _check_runtime_available()
|
||||
if (not available) or check_native_jit_disabled():
|
||||
return
|
||||
|
||||
if not _version_is_sufficient():
|
||||
return
|
||||
|
||||
_register_op_override_impl(
|
||||
_TRITON_DSL_NAME,
|
||||
lib_symbol,
|
||||
op_symbol,
|
||||
dispatch_key,
|
||||
impl,
|
||||
allow_multiple_override=allow_multiple_override,
|
||||
unconditional_override=unconditional_override,
|
||||
)
|
||||
|
||||
|
||||
# Register this DSL module with the registry
|
||||
# Note: Import-time registration ensures DSL is available when module is loaded
|
||||
dsl_registry.register_dsl("triton", cast(DSLModuleProtocol, sys.modules[__name__]))
|
||||
Reference in New Issue
Block a user