Skip to content

Commit

Permalink
[xla:gpu] Remove GatherSimplifier and ScatterSimplifier from `run…
Browse files Browse the repository at this point in the history
…SPMDPasses`, to avoid breaking the parallel dimensions pattern for propagation and partitioning in GSPMD.

Those passes are applied in `RunOptimizationPasses` which is after partitioning.

PiperOrigin-RevId: 625246878
  • Loading branch information
tensorflower-gardener authored and copybara-github committed Apr 16, 2024
1 parent 6ffb49c commit c7dbf7e
Showing 1 changed file with 0 additions and 2 deletions.
2 changes: 0 additions & 2 deletions xla/service/gpu/gpu_compiler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -651,10 +651,8 @@ absl::Status RunSPMDPasses(

spmd_simplify.AddPass<SortSimplifier>();
spmd_simplify.AddPass<TupleSimplifier>();
spmd_simplify.AddPass<ScatterSimplifier>();
spmd_simplify.AddPass<ScatterExpander>(
ScatterExpander::kEliminateSimpleScatters);
spmd_simplify.AddPass<GatherSimplifier>();
spmd_simplify.AddPass<GatherExpander>(
GatherExpander::kEliminateSimpleGathers);
spmd_simplify.AddPass<WhileLoopConstantSinking>();
Expand Down

0 comments on commit c7dbf7e

Please sign in to comment.