|
|
|
|
|
import inspect |
|
|
|
|
|
from torch._custom_op.impl import ( |
|
|
_custom_op_with_schema, |
|
|
_find_custom_op, |
|
|
infer_schema, |
|
|
parse_qualname, |
|
|
validate_namespace, |
|
|
) |
|
|
from torch.library import get_ctx |
|
|
|
|
|
|
|
|
__all__ = [ |
|
|
"custom_op", |
|
|
"impl", |
|
|
"impl_abstract", |
|
|
"get_ctx", |
|
|
"impl_save_for_backward", |
|
|
"impl_backward", |
|
|
] |
|
|
|
|
|
|
|
|
def custom_op(qualname, func_or_schema=None): |
|
|
r"""Register a new custom operator |
|
|
|
|
|
In PyTorch, defining an op (short for "operator") is a two step-process: |
|
|
- we need to define the op (by providing an operator name and schema) |
|
|
- we need to implement behavior for how the operator interacts with |
|
|
various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc. |
|
|
|
|
|
This entrypoint defines the custom operator (the first step) |
|
|
you must then perform the second step by calling various |
|
|
``impl_*`` APIs. |
|
|
|
|
|
This API may be used as a decorator (see examples). |
|
|
|
|
|
For a detailed guide on custom ops, please see |
|
|
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk |
|
|
|
|
|
Arguments: |
|
|
qualname (str): Should be a string that looks like |
|
|
"namespace::operator_name". Operators in PyTorch need a namespace to |
|
|
avoid name collisions; a given operator may only be created once. |
|
|
If you are writing a Python library, we recommend the namespace to |
|
|
be the name of your top-level module. |
|
|
func_or_schema (Union[Callable, str]): Each PyTorch operator needs a |
|
|
schema that tells PyTorch the types of the inputs/outputs. |
|
|
If this is a Callable, we will automatically infer the schema from |
|
|
the type annotations on the function (see examples). Otherwise, |
|
|
if you don't want to use type annotations, you may provide us the |
|
|
schema string. |
|
|
|
|
|
Example:: |
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) |
|
|
>>> import torch |
|
|
>>> import numpy as np |
|
|
>>> from torch import Tensor |
|
|
>>> |
|
|
>>> # Step 1: define the custom op. |
|
|
>>> # We need to provide the API a "prototype function" |
|
|
>>> # (a function that returns NotImplementedError), from which |
|
|
>>> # we will infer the types of the inputs and outputs. |
|
|
>>> @torch._custom_ops.custom_op("mylibrary::numpy_sin") |
|
|
>>> def numpy_sin(x: Tensor) -> Tensor: |
|
|
>>> raise NotImplementedError |
|
|
>>> |
|
|
>>> # The custom op is now accessible via the torch.ops module: |
|
|
>>> torch.ops.mylibrary.numpy_sin |
|
|
>>> |
|
|
>>> # Step 2: Register an implementation for various PyTorch subsystems |
|
|
>>> |
|
|
>>> # Register an implementation for CPU tensors |
|
|
>>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cpu") |
|
|
>>> def numpy_sin_impl_cpu(x): |
|
|
>>> return torch.from_numpy(np.sin(x.numpy())) |
|
|
>>> |
|
|
>>> # Register an implementation for CUDA tensors |
|
|
>>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cuda") |
|
|
>>> def numpy_sin_impl_cuda(x): |
|
|
>>> return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device) |
|
|
>>> |
|
|
>>> x = torch.randn(3) |
|
|
>>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cpu |
|
|
>>> |
|
|
>>> x_cuda = x.cuda() |
|
|
>>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cuda |
|
|
|
|
|
""" |
|
|
ns, name = parse_qualname(qualname) |
|
|
validate_namespace(ns) |
|
|
|
|
|
def inner(func): |
|
|
if not inspect.isfunction(func): |
|
|
raise ValueError( |
|
|
f"custom_op(...)(func): Expected `func` to be a Python " |
|
|
f"function, got: {type(func)}" |
|
|
) |
|
|
|
|
|
if func.__name__ != name: |
|
|
raise ValueError( |
|
|
f"custom_op(qualname='{qualname}', ...)(func): expected `func` " |
|
|
f"to have name '{name}' but got '{func.__name__}'. " |
|
|
f"Please either change the name of `func` or the qualname that " |
|
|
f"is passed to `custom_op`" |
|
|
) |
|
|
|
|
|
schema = infer_schema(func, mutates_args=()) |
|
|
_custom_op_with_schema(qualname, schema) |
|
|
return func |
|
|
|
|
|
if func_or_schema is None: |
|
|
return inner |
|
|
if isinstance(func_or_schema, str): |
|
|
_custom_op_with_schema(qualname, func_or_schema) |
|
|
else: |
|
|
return inner(func_or_schema) |
|
|
|
|
|
|
|
|
def impl(qualname, *, device_types=("cpu", "cuda"), func=None): |
|
|
r"""Register an implementation for a device type for this custom op. |
|
|
|
|
|
If the op is passed multiple Tensor inputs with different device |
|
|
types, it will dispatch to the registered implementation for the highest |
|
|
priority device type among those present. |
|
|
The supported device types, in order of priority, are {'cuda', 'cpu'}. |
|
|
|
|
|
This API may be used as a decorator (see examples). |
|
|
|
|
|
For a detailed guide on custom ops, please see |
|
|
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk |
|
|
|
|
|
Arguments: |
|
|
device_types (str or Iterable[str]): the device type(s) to register the function for. |
|
|
|
|
|
Example:: |
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) |
|
|
>>> import torch |
|
|
>>> import numpy as np |
|
|
>>> from torch import Tensor |
|
|
>>> |
|
|
>>> # Step 1: define the custom op. |
|
|
>>> # We need to provide the API a "prototype function" |
|
|
>>> # (a function that returns NotImplementedError), from which |
|
|
>>> # we will infer the types of the inputs and outputs. |
|
|
>>> @torch._custom_ops.custom_op("mylibrary::numpy_cos") |
|
|
>>> def numpy_cos(x: Tensor) -> Tensor: |
|
|
>>> raise NotImplementedError |
|
|
>>> |
|
|
>>> # The custom op is now accessible via the torch.ops module: |
|
|
>>> torch.ops.mylibrary.numpy_cos |
|
|
>>> |
|
|
>>> # Step 2: Register an implementation for various PyTorch subsystems |
|
|
>>> |
|
|
>>> # Register an implementation for CPU tensors |
|
|
>>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cpu") |
|
|
>>> def numpy_cos_impl_cpu(x): |
|
|
>>> return torch.from_numpy(np.cos(x.numpy())) |
|
|
>>> |
|
|
>>> # Register an implementation for CUDA tensors |
|
|
>>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cuda") |
|
|
>>> def numpy_cos_impl_cuda(x): |
|
|
>>> return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device) |
|
|
>>> |
|
|
>>> x = torch.randn(3) |
|
|
>>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cpu |
|
|
>>> |
|
|
>>> x_cuda = x.cuda() |
|
|
>>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cuda |
|
|
|
|
|
""" |
|
|
|
|
|
def inner(func): |
|
|
custom_op = _find_custom_op(qualname, also_check_torch_library=True) |
|
|
custom_op.impl(device_types, _stacklevel=3)(func) |
|
|
return func |
|
|
|
|
|
if func is None: |
|
|
return inner |
|
|
return inner(func) |
|
|
|
|
|
|
|
|
def impl_abstract(qualname, *, func=None): |
|
|
r"""Register an abstract implementation for this operator. |
|
|
|
|
|
An "abstract implementation" specifies the behavior of this operator on |
|
|
Tensors that carry no data. Given some input Tensors with certain properties |
|
|
(sizes/strides/storage_offset/device), it specifies what the properties of |
|
|
the output Tensors are. |
|
|
|
|
|
The abstract implementation has the same signature as the operator. |
|
|
It is run for both FakeTensors and meta tensors. To write an abstract |
|
|
implementation, assume that all Tensor inputs to the operator are |
|
|
regular CPU/CUDA/Meta tensors, but they do not have storage, and |
|
|
you are trying to return regular CPU/CUDA/Meta tensor(s) as output. |
|
|
The abstract implementation must consist of only PyTorch operations |
|
|
(and may not directly access the storage or data of any input or |
|
|
intermediate Tensors). |
|
|
|
|
|
This API may be used as a decorator (see examples). |
|
|
|
|
|
For a detailed guide on custom ops, please see |
|
|
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk |
|
|
|
|
|
Examples:: |
|
|
>>> import numpy as np |
|
|
>>> from torch import Tensor |
|
|
>>> |
|
|
>>> # Example 1: an operator without data-dependent output shape |
|
|
>>> @torch._custom_ops.custom_op("mylibrary::custom_linear") |
|
|
>>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor: |
|
|
>>> raise NotImplementedError |
|
|
>>> |
|
|
>>> @torch._custom_ops.impl_abstract("mylibrary::custom_linear") |
|
|
>>> def custom_linear_abstract(x, weight): |
|
|
>>> assert x.dim() == 2 |
|
|
>>> assert weight.dim() == 2 |
|
|
>>> assert bias.dim() == 1 |
|
|
>>> assert x.shape[1] == weight.shape[1] |
|
|
>>> assert weight.shape[0] == bias.shape[0] |
|
|
>>> assert x.device == weight.device |
|
|
>>> |
|
|
>>> return (x @ weight.t()) + bias |
|
|
>>> |
|
|
>>> # Example 2: an operator with data-dependent output shape |
|
|
>>> @torch._custom_ops.custom_op('mylibrary::custom_nonzero') |
|
|
>>> def custom_nonzero(x: Tensor) -> Tensor: |
|
|
>>> ... |
|
|
>>> |
|
|
>>> @torch._custom_ops.impl_abstract("mylibrary::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._custom_ops.get_ctx() |
|
|
>>> nnz = ctx.create_unbacked_symint() |
|
|
>>> shape = [x.dim(), nnz] |
|
|
>>> result = x.new_empty(shape, dtype=torch.long) |
|
|
>>> return result |
|
|
>>> |
|
|
>>> @torch._custom_ops.impl("mylibrary::custom_nonzero") |
|
|
>>> def custom_nonzero_impl(x): |
|
|
>>> x_np = to_numpy(x) |
|
|
>>> res = np.stack(np.nonzero(x_np), axis=1) |
|
|
>>> # unbacked symbolic ints in PyTorch must be >= 2, so we |
|
|
>>> # constrain the range to at least 2 |
|
|
>>> if res.shape[0] <= 1: |
|
|
>>> raise RuntimeError("not supported") |
|
|
>>> return torch.tensor(res, device=x.device) |
|
|
|
|
|
""" |
|
|
import torch.library |
|
|
|
|
|
return torch.library.register_fake(qualname, func, _stacklevel=2) |
|
|
|
|
|
|
|
|
def impl_save_for_backward(qualname, *, func=None): |
|
|
r"""Register a function that tells us what to save for backward. |
|
|
|
|
|
Please see :func:`impl_backward` for more details. |
|
|
""" |
|
|
|
|
|
def inner(func): |
|
|
custom_op = _find_custom_op(qualname, also_check_torch_library=True) |
|
|
custom_op.impl_save_for_backward(_stacklevel=3)(func) |
|
|
return func |
|
|
|
|
|
if func is None: |
|
|
return inner |
|
|
return inner(func) |
|
|
|
|
|
|
|
|
def impl_backward(qualname, output_differentiability=None, *, func=None): |
|
|
r"""Registers a backward formula for an operator. |
|
|
|
|
|
In order for an operator to work with autograd, you need to register |
|
|
a backward formula. There are two pieces to this: |
|
|
1. You must give us a function to specify what to save for backward. |
|
|
Call this the "save for backward" function. |
|
|
2. You must give us a function that computes gradients. Call this the |
|
|
"backward" function. |
|
|
|
|
|
Use `impl_save_for_backward` to define a "save for backward" function |
|
|
that specifies what gets saved for backward. The function should accept |
|
|
two arguments ``(inputs, output)`` and return the quantities to be saved |
|
|
for backward. |
|
|
|
|
|
During runtime, when you call the operator in a forwards pass, PyTorch |
|
|
will invoke the "save for backward" function with the inputs and output |
|
|
of the operator. |
|
|
|
|
|
Use `impl_backward` to define the "backward" function. The backward |
|
|
function must accept ``(ctx, saved, *grads)``: |
|
|
- ``ctx`` is a context object where we may provide information |
|
|
- ``saved`` is exactly what gets returned from the "save for backward" |
|
|
function |
|
|
- ``grads`` is one or more gradients. The number of gradients matches |
|
|
the number of outputs of the operator. |
|
|
|
|
|
The backward function must return a dict that maps the name of |
|
|
an input to the operator to its corresponding gradient. All inputs that |
|
|
were declared to be Tensors in the operator definition must be accounted |
|
|
for in the dict. The gradient may be a Tensor or None. |
|
|
|
|
|
For a detailed guide on custom ops, please see |
|
|
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk |
|
|
|
|
|
""" |
|
|
|
|
|
def inner(func): |
|
|
custom_op = _find_custom_op(qualname, also_check_torch_library=True) |
|
|
custom_op.impl_backward(output_differentiability, _stacklevel=3)(func) |
|
|
return func |
|
|
|
|
|
if func is None: |
|
|
return inner |
|
|
return inner(func) |
|
|
|
|
|
|
|
|
def _destroy(qualname): |
|
|
"""De-registers a custom op. For testing purposes only""" |
|
|
custom_op = _find_custom_op(qualname) |
|
|
custom_op._destroy() |
|
|
|
