[Dygraph API] Fix merged_momentum, provide actual inplace operations …

[RuohengMa]

PR types

Bug fixes

PR changes

APIs

Description

在维持#58204功能的前提下修复遗留的问题，使fallback到CPU后，merged_momentum kernel里输入输出算子在CPU上的地址是一样的

以merged_momentum为例：
output部分：
修复前：

  std::tuple<std::vector<Tensor>&, std::vector<Tensor>&, paddle::optional<std::vector<Tensor>>&> api_output{param, velocity, master_param};
  auto kernel_out_0 = SetInplaceVectorKernelOutput(param.size(), &std::get<0>(api_output));
  if (kernel_result.has_fallback_cpu) {
    TransDataBackend(kernel_out_0, actual_kernel_backend, kernel_out_0);
  }
  auto kernel_out_1 = SetInplaceVectorKernelOutput(param.size(), &std::get<1>(api_output));
  if (kernel_result.has_fallback_cpu) {
    TransDataBackend(kernel_out_1, actual_kernel_backend, kernel_out_1);
  }
  auto kernel_out_2 = SetInplaceOptionalVectorKernelOutput(param.size(), std::get<2>(api_output));
  if (kernel_result.has_fallback_cpu) {
    TransDataBackend(kernel_out_2, actual_kernel_backend, kernel_out_2);
  }

修复后：

  std::tuple<std::vector<Tensor>&, std::vector<Tensor>&, paddle::optional<std::vector<Tensor>>&> api_output{param, velocity, master_param};
  auto kernel_out_0 = SetInplaceVectorKernelOutput(param.size(), &std::get<0>(api_output));
  if (kernel_result.has_fallback_cpu) {
    for (size_t i = 0; i < kernel_out_0.size(); ++i) {
      kernel_out_0[i] = const_cast<phi::DenseTensor*>(input_param[i]);
    }
  }
  auto kernel_out_1 = SetInplaceVectorKernelOutput(param.size(), &std::get<1>(api_output));
  if (kernel_result.has_fallback_cpu) {
    for (size_t i = 0; i < kernel_out_1.size(); ++i) {
      kernel_out_1[i] = const_cast<phi::DenseTensor*>(input_velocity[i]);
    }
  }
  auto kernel_out_2 = SetInplaceOptionalVectorKernelOutput(param.size(), std::get<2>(api_output));
  if (kernel_result.has_fallback_cpu) {
    for (size_t i = 0; i < kernel_out_2.size(); ++i) {
      kernel_out_2[i] = const_cast<phi::DenseTensor*>(input_master_param->at(i));
    }
  }

return前的copy back部分：
修复前：

  if (kernel_result.has_fallback_cpu) {

    TransDataBackend(kernel_out_0, kernel_backend, kernel_out_0);
    TransDataBackend(kernel_out_1, kernel_backend, kernel_out_1);
    TransDataBackend(kernel_out_2, kernel_backend, kernel_out_2);

  }

修复后：

  if (kernel_result.has_fallback_cpu) {

    TransDataBackend(kernel_out_0, kernel_backend, kernel_out_0);
    for (size_t i = 0; i < param.size(); ++i) {
      auto target_ptr = static_cast<phi::DenseTensor*>(param.at(i).impl().get());
      *target_ptr = *kernel_out_0.at(i);
    }
    TransDataBackend(kernel_out_1, kernel_backend, kernel_out_1);
    for (size_t i = 0; i < velocity.size(); ++i) {
      auto target_ptr = static_cast<phi::DenseTensor*>(velocity.at(i).impl().get());
      *target_ptr = *kernel_out_1.at(i);
    }
    TransDataBackend(kernel_out_2, kernel_backend, kernel_out_2);
    if (master_param) {
      for (size_t i = 0; i < master_param->size(); ++i) {
        auto target_ptr = static_cast<phi::DenseTensor*>(master_param->at(i).impl().get());
        *target_ptr = *kernel_out_2.at(i);
      }
    }

  }

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[CLAassistant]

All committers have signed the CLA.

[paddle-bot]

❌ The PR is not created using PR's template. You can refer to this Demo.
Please use PR's template, it helps save our maintainers' time so that more developers get helped.

[qili93]

LGTM for const_cast

[YuanRisheng]

本身就是inplace，这里为何要再次把input的指针赋值给out

[RuohengMa]

因为如果发生了fallback到cpu的情况的话，输入的tensor是会被从xpu copy到cpu上；因为算子是inplace，输入输出的指针应该是同一个，所以需要把kernel_out指向copy到cpu上的数据而不是指向xpu上的数据~

[zyfncg]

这里等号左右的tensor不是同一个吗？

[YuanRisheng]

这里不是inplace吗，为何又要把输出给了输入

[RuohengMa]

因为fallback到cpu之后，out是在cpu上的，所以需要重新把out的值写回xpu上的input~

[cqulilujia]

LGTM