基于pytorch1.9原生解释,重点介绍native_functions.yaml中关键参数

文档资料:

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=[]表示。

  • 0
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值