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
|
# 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
)
|
