search

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.75k stars 440 forks source link

[XLA:Collective] Support normalizing all-reduce #19819

Open copybara-service[bot] opened 6 days ago

copybara-service[bot] commented 6 days ago

[XLA:Collective] Support normalizing all-reduce

  1. Add a normalizer to normalier unsupported All-reduce into All-To-All + Reduce + All-Gather
  2. Add common utility functions into (i) hlo_sharding_util (ii). collective_util (iii). shape