File size: 8,058 Bytes
c61ccee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import contextlib
import functools
import warnings
from typing import Callable, Optional

import torch
from torch._library.utils import Kernel, RegistrationHandle


class AbstractImplHolder:
    """A holder where one can register an abstract 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 abstract impl.



        Returns a RegistrationHandle that one can use to de-register this

        abstract impl.

        """
        if self.kernel is not None:
            raise RuntimeError(
                f"impl_abstract(...): the operator {self.qualname} "
                f"already has an abstract impl registered at "
                f"{self.kernel.source}."
            )
        if torch._C._dispatch_has_kernel_for_dispatch_key(self.qualname, "Meta"):
            raise RuntimeError(
                f"impl_abstract(...): 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"impl_abstract."
            )

        if torch._C._dispatch_has_kernel_for_dispatch_key(
            self.qualname, "CompositeImplicitAutograd"
        ):
            raise RuntimeError(
                f"impl_abstract(...): 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 abstract "
                f"impl; "
                f"instead, the operator will decompose into its constituents "
                f"and those "
                f"can have abstract impls defined on them."
            )

        # Store the kernel in this holder
        self.kernel = Kernel(func, source)

        # Also register the abstract 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_abstract_impl():
            if self.lib:
                self.lib._destroy()
                self.lib = None
            self.kernel = None

        return RegistrationHandle(deregister_abstract_impl)


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 abstract implementations for custom operators.

    """

    def __init__(self, _shape_env, _op):
        self._shape_env = _shape_env
        self._op = _op

    def create_unbacked_symint(self, *, min=2, max=None) -> torch.SymInt:
        warnings.warn(
            "create_unbacked_symint is deprecated, please use new_dynamic_size instead"
        )
        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 abstract 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.impl_abstract("mymodule::custom_nonzero")

            >>> def custom_nonzero_abstract(x):

            >>>     # Number of nonzero-elements is data-dependent.

            >>>     # Since we cannot peek at the data in an abstract 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 custom_nonzero_cpu(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