[Prim/Comp] Fix place setting to avoid redundant H2D/D2H copy for backward decomposition

[HydrogenSulfate]

PR Category

Performance Optimization

PR Types

Bug fixes

Description

Pcard-75624

问题描述

现有动态图/静态图组合算子大量使用了带有place信息的基础算子，如full、full_scalar、full_with_tensor，这些算子末尾的参数place决定了首先在哪个设备上创建张量。但目前所有的组合算子调用方式中均未手动指定这些张量的place参数，导致全部默认创建在 CPU 上，

Paddle/paddle/fluid/eager/backward.cc

Lines 379 to 382 in d774689

	node_input_buffers_dict[next_node]->add(edge_rank.first,
	edge_rank.second,
	grad_output_tensor,
	create_graph);

继而在backward.cc（以动态图为例）的梯度累加过程中，由于累加时会对 source 和 target Tensor 所在设备进行检测并将 target 拷贝到 source 设备上，因此会不断地触发 H2D 或者 D2H 拷贝，拖慢反向运行速度

Paddle/paddle/fluid/imperative/gradient_accumulator.cc

Lines 194 to 197 in 7aae194

	// if src and dst are in different place, copy dst to src's place
	if (dst_tensor->place() != place) {
	paddle::framework::TensorCopySync(*dst_tensor, place, dst_tensor);
	}

解决方案

本 PR 修复了所有动态图、静态图组合反向算子里的相关调用，指定了place与输入张量保持一致，避免了大量的H2D/D2H拷贝操作。

上方为修复后的timeline，下方为修复前的timeline，可以看到修复后基本看不到红色的拷贝操作

上方为修复前的H2D/D2H拷贝耗时情况，下方为修复后的拷贝耗时情况，可以看到修复后原本异常的拷贝耗时，恢复正常

案例测试结果：

1. Allen_Cahn（二阶微分方程+一阶反向，性能提升 100%，显存增加1.2%）

2. DeePMD-kit DPA2（21s -> 18s，提升 14%）

3. ldc_2d(二阶微分方程+一阶反向，性能提升63%）

TODO：修复静态图前向组合composite.h中的place问题，下一个PR修复

[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.