what's the purpose of this pattern rewriter?

We will have

B = All-Gather(A)
C = Reshape(B)
D = OptimizationBarrier(C)
E = Slice(D)

graph pattern lowered from TorchAcc, to prevent standalone Reshape Op, which will cause extra global memory read & write, we add this rule to change this graph into

B = All-Gather(A)
C = OptimizationBarrier(B)
D = Reshape(C)
E = Slice(D)

@Yancey1989

Thank you for the clarification. After discussing with @eedalong in the office, it appears that the reshape->optimization_barrier->slice pattern is not a commonly simplified rule. In the XLA FSDP ( implementation, this pattern serves as a workaround for TPUs. However, it is not required for NVIDIA GPUs. Perhaps we could set _shard_size_multiple=1, which would eliminate the need for this pattern. Ref: :https://github.com/AlibabaPAI/xla/blob/67edb354372a8e5cbce41f20ecfca68b328635c6/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py#L1435 .
cc @anw90