请教一个cast grad的相关问题

[houj04]

请提出你的问题 Please ask your question

背景：某个项目需求，我在梳理某硬件设备上能跑的飞桨算子，发现该硬件上面实现了cast算子，但是没有实现cast_grad算子。
可是这个硬件设备能正常跑训练，如果缺算子的话应该跑不起来啊。于是去翻代码，发现CPU和GPU都有这个反向算子的实现，分别位于：
paddle/phi/kernels/gpu/cast_grad_kernel.cu
paddle/phi/kernels/cpu/cast_grad_kernel.cc
然后我又找到有这个东西
paddle/phi/ops/yaml/backward.yaml
其中有一段是这样写的：

- backward_op : cast_grad
  forward : cast (Tensor x, DataType dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  invoke : cast (out_grad, x.dtype())
  composite: cast_grad(x, out_grad, x_grad)
  no_need_buffer : x

从字面意思上看，我直观理解，当遇到cast grad算子的时候，就跑去执行cast的前向算子。
那问题来了：
1、想请教下这个backward.yaml的作用，是不是上面我直观理解的那样。有没有更详细的解释呢（例如，它是如何发挥作用的，这个yaml文件是怎么影响算子构建和计算图执行的）。
2、都有这个backward.yaml了，为啥CPU和GPU上还单独有一个cast grad的kernel呢？

[ZhangX-21]

1. backward.yaml文件描述反向算子的注册信息。
   例如- backward_op : cast_grad
   forward : cast (Tensor x, DataType dtype) -> Tensor(out)对应的cast正向实现。
   args : (Tensor x, Tensor out_grad) 对应cast_grad输入参数。
   output : Tensor(x_grad) 对应cast_grad输出参数。
   等等信息。

2. cast_grad反向函数根据链式求导法则复用了cast正向kernel实现：
   paddle/phi/kernels/gpu/cast_grad_kernel.cu中

CastKernel<T, Context>(dev_ctx, out_grad, x.dtype(), x_grad);

paddle/phi/kernels/cpu/cast_grad_kernel.cc

CastKernelImpl<T, data_t>(dev_ctx, out_grad, x_grad->dtype(), x_grad);

输入参数变成都out_grad，x.dtype()，x_grad，其中x.dtype()和x_grad->dtype()等价。

[houj04]

我写代码试了一下，问题在于，CastGradKernel这个东西似乎并没有执行，而是像我猜的那样，直接走的CastKernel。
基于当前非常新的develop分支：

我是这样改的代码：

测试代码如下：

import paddle
x = paddle.to_tensor([2, 3, 4], 'float64', stop_gradient=False)
y = paddle.cast(x, 'uint8')
print(y)
y.backward()

跑出来的结果如下，注意它调用了两次CastKernel，而不是CastGradKernel。

[bqwwen]

该链接介绍了算子Yaml配置规则：
https://www.paddlepaddle.org.cn/documentation/docs/zh/dev_guides/api_contributing_guides/new_cpp_op_cn.html#yaml
当前paddle/phi/ops/yaml/backward.yaml

- backward_op : cast_grad
  forward : cast (Tensor x, DataType dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  invoke : cast (out_grad, x.dtype())
  composite: cast_grad(x, out_grad, x_grad)
  no_need_buffer : x

• backward_op：反向算子名称为cast_grad。
• forward：对应前向算子的名称cast。
• args：反向算子输入参数。
• output：反向算子输出，顺序需要与前向输入 Tensor 一致。
• invoke：复用已有的算子接口或实现自定义的 C++ API，配置时以函数调用的形式配置即可。即，描述了如何调用反向传播操作，具体为cast (out_grad, x.dtype())。
• composite：定义了组合操作cast_grad(x, out_grad, x_grad)。
• no_need_buffer：可选配置，标记的 Tensor 变量在前向运行完成后，持有的内存或显存会被释放，以减少训练过程中的内存使用。

此时，此时正向和反向都会调用CastKernel。
生成的paddle/fluid/eager/api/generated/eager_generated/backwards/nodes.cc中的CastGradNode和invoke描述的调用已有的cast函数实现反向函数一致：

  // Call grad_api function

  if (trace_backward) {
    auto api_output = cast_ad_func (out_grad, x.dtype());
    *api_output_0 = api_output;
  } else {
    auto api_output = paddle::experimental::cast (out_grad, x.dtype());
    *api_output_0 = api_output;
  }

如果希望正向调用CastKernel，反向调用CastGradKernel，可以做如下改动（改完之后需要重新cmake+ninja）：

- backward_op : cast_grad
  forward : cast (Tensor x, DataType dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  infer_meta :
    func : UnchangedInferMeta
    param: [x]
  kernel :
    func : cast_grad
    data_type : out_grad

• backward_op：反向算子名称为cast_grad。
• forward：对应前向算子的名称cast。
• args：反向算子输入参数。
• output：反向算子输出，顺序需要与前向输入 Tensor 一致。
• infer_meta：InferMeta 函数负责根据输入变量推断返回 Tensor 的维度与类型，这里是对算子使用的 InferMeta 函数进行配置。其中func指定了推断函数为UnchangedInferMeta，这意味着在推断元数据时，输入张量x的元数据（如形状、数据类型等）将保持不变，即输出张量out的元数据将与输入张量x相同（除了数据类型可能因cast操作而改变）。param指定了推断函数关注的参数为[x]。
• kernel：定义了实际执行计算的部分。func给出了算子对应kernel的函数为cast_grad，data_type表示根据指定参数out_grad推导调用 kernel 的 data_type。

此时，正向调用CastKernel，反向调用CastGradKernel。
生成的paddle/fluid/eager/api/generated/eager_generated/backwards/nodes.cc中的CastGradNode与kernel定义的用cast_grad函数实际执行计算部分一致：

  // Call grad_api function

  paddle::experimental::cast_grad(x, out_grad, api_output_0);

[houj04]

https://www.paddlepaddle.org.cn/documentation/docs/zh/dev_guides/api_contributing_guides/new_cpp_op_cn.html#yaml 中介绍的，当前backward.yaml 、、、

* backward_op : cast_grad
  forward : cast (Tensor x, DataType dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  invoke : cast (out_grad, x.dtype())
  composite: cast_grad(x, out_grad, x_grad)
  no_need_buffer : x
  、、、


* backward_op：反向算子名称为cast_grad

* forward：对应前向算子的名称cast

* args：反向算子输入参数

* output：反向算子输出，顺序需要与前向输入 Tensor 一致

* invoke：复用已有的算子接口或实现自定义的 C++ API，配置时以函数调用的形式配置即可。即，描述了如何调用反向传播操作，具体为cast (out_grad, x.dtype())。

- backward_op : cast_grad
  forward : cast (Tensor x, DataType dtype) -> Tensor(out)
  args : (Tensor x, Tensor out_grad)
  output : Tensor(x_grad)
  infer_meta :
    func : UnchangedInferMeta
    param: [x]
  kernel :
    func : cast_grad
    data_type : out_grad

按照这里描述的，在需要计算cast的反向的时候，就会“复用已有的算子接口”，调用前向的cast算子。那和我实验的结果是一致的。
那么下一个问题是：这是不是说明，paddle/phi/kernels/gpu/cast_grad_kernel.cu以及paddle/phi/kernels/cpu/cast_grad_kernel.cc这两个文件，不再起作用了？那是不是可以把它们移除掉。