search

openxla / xla

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

Add SPMD config option to specify zero cost method for gather/scatter. #14862

Closed ptoulme-aws closed 1 month ago

ptoulme-aws commented 1 month ago

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.

ptoulme-aws commented 1 month ago

@frgossen please review. Thank you.

thomasjoerg commented 1 month ago

@frgossen friendly reminder!

ptoulme-aws commented 1 month ago

Clang-format also fixed some formatting.

These PRs will be auto-formatted on import so those changes might disappear again. Generally, it will be easier to review if you don't change formatting.

Thanks for making this configurable!

Okay, I will not format any future changes. Thank you for reviewing.

I see my PR is not merging due to ROCm and Windows Bazel build failures in copy-bara import to Tensorflow. I cannot view the ROCm failure. The Windows failure is on this target - tensorflow/tools/ci_build/builds/gen_win_out.exe. How do I debug this unless these builds are broken on TOT right now?