Closed ptoulme-aws closed 3 months ago
Thanks ! Do you plan to provide a patch for this?
I can submit PR that moves this to behind an optional config? Default will still be IndexParallel to maintain existing flow
I think making these things configurable makes sense.
In the SPMD partitioner the gather/scatter handler hardcodes the cost to {0,0} for IndexParallel strategy. This is not ideal for all hardware, especially for hardware with 2D Torus topology. This should be refactored to allow the strategy to hardcode to be passed as config.
Location of hardcode cost - https://github.com/openxla/xla/blob/eaed933666ca4b44ea96b6bdae13631c1edfea00/xla/service/spmd/gather_scatter_handler.cc#L736