文档资料:
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/README.md
1)func:
- func: func_name[.overload_name](ArgType arg0[=default], ArgType arg1[=default], ...) -> Return
Tensor -》 const at::tensor&
Tensor(a) -》 alias of this tensor
Tensor(a!) -》 inplace of this tensor
Tensor(a! -> a|b) -》 Tensor is in seta, written to, and after the write is in setaANDb
Tensor[] -》 TensorList
* -》 is a special sentinel argument, which doesn't translate into an actual argument, but indicates that in the Python bindings, any subsequent arguments must be specified as keyword arguments
? -》 is trailing question mark that annotates an argument to be an optional type
2)variants:
function, method
method的定义表示: self.for(), function的定义表示: at::for()
3)annotations:
abs(Tensor self) -> Tensor stays the same as it will always allocate new memory.
abs_(Tensor(a!) self) -> Tensor(a!) self may be written to and returned. Further, the annotation indicates that the return value may alias the input. This indicates an inplace function and by convention ends in a single '_'.
abs(Tensor self, *, Tensor(a!) out) -> Tensor(a!) In the Python frontend out can be passed as a keyword argument and may be written to. In this case it indicates the schema for a function that must accept out as this does not provide a default argument. The idea behind representing this as a optional argument is to document the intended usage. This maps to the legacy abs_out(Tensor out, Tensor self) -> Tensor. As with the legacy _out function you must call the argument Tensor out or Tensor out0, Tensor out1 in the context of multiple arguments.
4)dispatch
Available backend options can be found by searching dispatch_keys
in codegen.
dispatch: // 分发到对应cpu或者cuda的设备上
CPU: func_cpu
CUDA: func_cuda
三个通用的dispatch key:
CompositeExplicitAutograd: 表示这个kernel使用于所有的算子,但是需要在derivatives.yaml文件中显示地声明反向函数。这类算子典型的应用场景是通过stub将复杂的任务分发给单独的kernel,而本身只进行简单的工作。 优先级高于CompositeImplicitAutograd。
native_functions.yaml
- func: sgn(Tensor self) -> Tensor
variants: function, method
dispatch:
CompositeExplicitAutograd: sgn
derivatives.yaml
- name: sgn(Tensor self) -> Tensor
self: sgn_backward(result, grad, self)
c++
REGISTER_DISPATCH(sgn_stub, &sgn_kernel_cuda);
CompositeImplicitAutograd: 表示这个kernel使用于所有的算子,且天然支持反向处理。(默认行为)
- func: instance_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool use_input_stats, float momentum, float eps, bool cudnn_enabled) -> Tensor
variants: function
- func: addr(Tensor self, Tensor vec1, Tensor vec2, *, Scalar beta=1, Scalar alpha=1) -> Tensor
variants: function, method
dispatch:
CPU, CUDA: addr
CompositeImplicitAutograd: math_addr // you still want to provide a numerically stable gradient formula instead of using autograd, 单独添加反向函数
CompositeExplicitAutogradNonFunctional: 1.9分支未看到使用
4)manual_kernel_registration
如果设置为True,将不会自动生成c++端代码,而是通过TypeDefault进行分发。1.9分支已经没有使用
5)其他
5.1)通过python 方式注册dispatch。Validate the computed dispatch table matches what you want. You can use PythonDispatcher
provided in torch/_python_dispatcher.py.
dispatcher = PythonDispatcher()
dispatcher.register(["CPU", "XLA", "AutogradCPU", "CompositeImplicitAutograd"])
print(dispatcher.dispatchTable()) # Tells you exactly which kernel is used for certain backend.
5.2) 当注册提供了CompositeImplicitAutograd,并且不添加任何的dispatch方式,将默认自动支持autograd;
5.3)如下情况,需要单独给非out接口定义反向,其中sign_和sign.out不支持反向。
- func: sign(Tensor self) -> Tensor
device_check: NoCheck # TensorIterator
structured_delegate: sign.out
variants: function, method
dispatch:
CompositeExplicitAutograd: sign
- func: sign_(Tensor(a!) self) -> Tensor(a!)
device_check: NoCheck # TensorIterator
structured_delegate: sign.out
variants: method
dispatch:
CompositeExplicitAutograd: sign_
- func: sign.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # TensorIterator
structured: True
structured_inherits: TensorIteratorBase
dispatch:
CPU, CUDA: sign_out
5.4)undefined tensor可以通过Tensor?
or Tensor? x=[]表示。