[Dist dialect] Add pass for dist_reshape and nd_mesh_all2all in PIR

[pkuzyc]

PR Category

Auto Parallel

PR Types

New features

Description

Pcard-67164
Add pass for dist_reshape and nd_mesh_all2all in PIR.

Major changes:

1. dist_reshape pass: replace the dist op "dist_op.dist_reshape" with reshape op for further execution.
2. nd_mesh_all2all: use all-to-all communication when performing shard-to-shard reshard on multi-dimensional process mesh.
   2.1 This PR add a s_to_s reshard function in PIR.
   2.2 This PR modify the reshard pass, when resharding from [Replicate(), Shard(1)] to [Replicate(), Shard(0)], it will use all-to-all communication.

Minor:

1. Fix the bug in setting no_need_buffer values.
2. Fix the bug about getting wrong communication group in parsing communication context.
3. Fix the bug in PIR sub_to_global mesh reshard function.

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

[paddle-ci-bot]

Sorry to inform you that a2a7c36's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[paddle-ci-bot]

Sorry to inform you that 0953f87's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[paddle-ci-bot]

Sorry to inform you that 8c22dfd's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[pkuzyc]

all changes except the dist_reshape pass have been moved to #71264 and #71320