File size: 2,723 Bytes
67e9774 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
# mypy: allow-untyped-defs
from typing import Callable, Optional
from .fake_impl import FakeImplHolder
from .utils import RegistrationHandle
__all__ = ["SimpleLibraryRegistry", "SimpleOperatorEntry", "singleton"]
class SimpleLibraryRegistry:
"""Registry for the "simple" torch.library APIs
The "simple" torch.library APIs are a higher-level API on top of the
raw PyTorch DispatchKey registration APIs that includes:
- fake impl
Registrations for these APIs do not go into the PyTorch dispatcher's
table because they may not directly involve a DispatchKey. For example,
the fake impl is a Python function that gets invoked by FakeTensor.
Instead, we manage them here.
SimpleLibraryRegistry is a mapping from a fully qualified operator name
(including the overload) to SimpleOperatorEntry.
"""
def __init__(self):
self._data = {}
def find(self, qualname: str) -> "SimpleOperatorEntry":
if qualname not in self._data:
self._data[qualname] = SimpleOperatorEntry(qualname)
return self._data[qualname]
singleton: SimpleLibraryRegistry = SimpleLibraryRegistry()
class SimpleOperatorEntry:
"""This is 1:1 to an operator overload.
The fields of SimpleOperatorEntry are Holders where kernels can be
registered to.
"""
def __init__(self, qualname: str):
self.qualname: str = qualname
self.fake_impl: FakeImplHolder = FakeImplHolder(qualname)
self.torch_dispatch_rules: GenericTorchDispatchRuleHolder = (
GenericTorchDispatchRuleHolder(qualname)
)
# For compatibility reasons. We can delete this soon.
@property
def abstract_impl(self):
return self.fake_impl
class GenericTorchDispatchRuleHolder:
def __init__(self, qualname):
self._data = {}
self.qualname = qualname
def register(
self, torch_dispatch_class: type, func: Callable
) -> RegistrationHandle:
if self.find(torch_dispatch_class):
raise RuntimeError(
f"{torch_dispatch_class} already has a `__torch_dispatch__` rule registered for {self.qualname}"
)
self._data[torch_dispatch_class] = func
def deregister():
del self._data[torch_dispatch_class]
return RegistrationHandle(deregister)
def find(self, torch_dispatch_class):
return self._data.get(torch_dispatch_class, None)
def find_torch_dispatch_rule(op, torch_dispatch_class: type) -> Optional[Callable]:
return singleton.find(op.__qualname__).torch_dispatch_rules.find(
torch_dispatch_class
)
|