File size: 7,889 Bytes
d1ceb73 |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
# mypy: allow-untyped-defs
import contextlib
import functools
from typing import Callable, Optional
from typing_extensions import deprecated
import torch
from torch._library.utils import Kernel, RegistrationHandle
class AbstractImplHolder:
"""A holder where one can register an fake impl to."""
def __init__(self, qualname: str):
self.qualname: str = qualname
self.kernel: Optional[Kernel] = None
self.lib: Optional[torch.library.Library] = None
def register(self, func: Callable, source: str) -> RegistrationHandle:
"""Register an fake impl.
Returns a RegistrationHandle that one can use to de-register this
fake impl.
"""
if self.kernel is not None:
raise RuntimeError(
f"register_fake(...): the operator {self.qualname} "
f"already has an fake impl registered at "
f"{self.kernel.source}."
)
if torch._C._dispatch_has_kernel_for_dispatch_key(self.qualname, "Meta"):
raise RuntimeError(
f"register_fake(...): the operator {self.qualname} "
f"already has an DispatchKey::Meta implementation via a "
f"pre-existing torch.library or TORCH_LIBRARY registration. "
f"Please either remove that registration or don't call "
f"register_fake."
)
if torch._C._dispatch_has_kernel_for_dispatch_key(
self.qualname, "CompositeImplicitAutograd"
):
raise RuntimeError(
f"register_fake(...): the operator {self.qualname} "
f"already has an implementation for this device type via a "
f"pre-existing registration to "
f"DispatchKey::CompositeImplicitAutograd."
f"CompositeImplicitAutograd operators do not need an fake "
f"impl; "
f"instead, the operator will decompose into its constituents "
f"and those "
f"can have fake impls defined on them."
)
# Store the kernel in this holder
self.kernel = Kernel(func, source)
# Also register the fake impl to Meta key
if self.lib is None:
ns = self.qualname.split("::")[0]
self.lib = torch.library.Library(ns, "FRAGMENT")
meta_kernel = construct_meta_kernel(self.qualname, self)
self.lib.impl(self.qualname, meta_kernel, "Meta")
def deregister_fake_class():
if self.lib:
self.lib._destroy()
self.lib = None
self.kernel = None
return RegistrationHandle(deregister_fake_class)
def construct_meta_kernel(
qualname: str, abstract_impl_holder: AbstractImplHolder
) -> Callable:
assert abstract_impl_holder.kernel is not None
@functools.wraps(abstract_impl_holder.kernel.func)
def meta_kernel(*args, **kwargs):
assert abstract_impl_holder.kernel is not None
source = abstract_impl_holder.kernel.source
def error_on_ctx():
raise RuntimeError(
f"Attempted to call get_ctx() for the meta implementation "
f"for {qualname} (implemented at {source})"
f"You have presumably called get_ctx() because the operator "
f"has a data-dependent output shape; if so, there is no "
f"such meta implementation and this error is the correct "
f"behavior."
)
with set_ctx_getter(error_on_ctx):
return abstract_impl_holder.kernel(*args, **kwargs)
return meta_kernel
def get_none():
return None
global_ctx_getter: Callable = get_none
@contextlib.contextmanager
def set_ctx_getter(ctx_getter):
global global_ctx_getter
prev = global_ctx_getter
try:
global_ctx_getter = ctx_getter
yield
finally:
global_ctx_getter = prev
class AbstractImplCtx:
"""
Context object for writing fake implementations for custom operators.
"""
def __init__(self, _fake_mode, _op):
self._fake_mode = _fake_mode
self._shape_env = _fake_mode.shape_env
self._op = _op
@deprecated(
"`create_unbacked_symint` is deprecated, please use `new_dynamic_size` instead",
category=FutureWarning,
)
def create_unbacked_symint(self, *, min=2, max=None) -> torch.SymInt:
return self.new_dynamic_size(min=min, max=max)
def new_dynamic_size(self, *, min=0, max=None) -> torch.SymInt:
"""Constructs a new symint (symbolic int) representing a data-dependent value.
This is useful for writing the fake implementation (which is necessary
for torch.compile) for a CustomOp where an output Tensor has a size
that depends on the data of the input Tensors.
Args:
min (int): A statically known inclusive lower bound for this symint. Default: 0
max (Optional[int]): A statically known inclusive upper bound for this
symint. Default: None
.. warning:
It is important that the ``min`` and ``max`` (if not None) values are set
correctly, otherwise, there will be undefined behavior under
torch.compile. The default value of ``min`` is 2 due to torch.compile
specializing on 0/1 sizes.
You must also verify that your implementation on concrete Tensors
(e.g. CPU/CUDA) only returns Tensors where the size that corresponds
to the symint also has respects these constraint.
The easiest way to do this is to add an assertion in the CPU/CUDA/etc
implementation that the size follows these bounds.
Example::
>>> # An operator with data-dependent output shape
>>> lib = torch.library.Library("mymodule", "FRAGMENT")
>>> lib.define("mymodule::custom_nonzero(Tensor x) -> Tensor")
>>>
>>> @torch.library.register_fake("mymodule::custom_nonzero")
>>> def _(x):
>>> # Number of nonzero-elements is data-dependent.
>>> # Since we cannot peek at the data in an fake impl,
>>> # we use the ctx object to construct a new symint that
>>> # represents the data-dependent size.
>>> ctx = torch.library.get_ctx()
>>> nnz = ctx.new_dynamic_size()
>>> shape = [nnz, x.dim()]
>>> result = x.new_empty(shape, dtype=torch.int64)
>>> return result
>>>
>>> @torch.library.impl(lib, "custom_nonzero", "CPU")
>>> def _(x):
>>> x_np = x.numpy()
>>> res = np.stack(np.nonzero(x_np), axis=1)
>>> return torch.tensor(res, device=x.device)
"""
if (
self._shape_env is None
or not self._shape_env.allow_dynamic_output_shape_ops
):
raise torch._subclasses.fake_tensor.DynamicOutputShapeException(self._op)
if isinstance(min, torch.SymInt) or isinstance(max, torch.SymInt):
raise ValueError(
f"ctx.new_dynamic_size(min={min}, max={max}): expected "
f"min and max to be statically known ints but got SymInt. "
f"This is not supported."
)
if min < 0:
raise ValueError(
f"ctx.new_dynamic_size(min={min}, ...): expected min to be "
f"greater than or equal to 0: this API can only create "
f"non-negative sizes."
)
result = self._shape_env.create_unbacked_symint()
torch.fx.experimental.symbolic_shapes._constrain_range_for_size(
result, min=min, max=max
)
return result
|