[NVIDIA] Add fixes for supporting determinism expander for high-dimensional scatter operation and a flag to disable it

[serach24]

This PR is the 2nd step (out of 2) to improve the performance of deterministic scatter. Originally, the scatter op will be expanded to be deterministic in xla/service/ScatterExpander.cc. However, since it took a while-loop-based approach, the performance is extremely poor. We designed and implemented a prefix-scan-based approach to rewrite the scatter operation to be an efficient deterministic scatter. This PR completes the optimization of deterministic scatter operations with non-scalar indices and updates.

The change of this PR is on top of #17886, and has fixed issues reported in the reverted PR #18326. The issue was that the changes in #18326 were not able to handle different kinds of complicated but realistic scatter dimension numbers. Specifically, this PR unifies the implementation of 1D and multi-dimensional scatter operation to make the code easier to maintain, adds multiple tests for various scatter dimension numbers, and thoroughly handles all cases of different kinds of dimension numbers.

Moreover, this PR also adds an option xla_gpu_enable_scatter_determinism_expander, the default value of which is set to be true. This option could make sure that although unlikely, if anything happens with changes in this PR, the user can easily disable the scatter_determinism_expander pass without getting blocked.

Design doc: https://docs.google.com/document/d/1K204VZR3OP0SUDOPsGUYgIIDf2ucTKEC4yQj8XRG2SA/edit

Bugs resolved: jax-ml/jax#17844

[akuegel]

Thanks. With the help of the flag, it would also be possible to turn it off by default if there turns out to be another bug, so we don't have to revert the full change.

[akuegel]

Actually, it seems this is already used by xla_gpu_experimental_stream_annotation
Please use 343 and update Next id: 344

[serach24]

Changed

[AyanmoI]

@akuegel Does this need another approval?

[akuegel]

@akuegel Does this need another approval?

It would have needed a rebase. I did that internally, but then that needs to be approved by someone else. I asked our current Openxla PR oncall to approve, but he is in a different time zone than me

[akuegel]

It seems there is still a bug, now it is failing after the new pass with a HloVerifier error:

[hlo verifier]: Updates tensor must be of rank 2; got 1.
	, for instruction %scatter.33785 = f32[1,384,384]{2,1,0} scatter(f32[1,384,384]{2,1,0} %broadcast.33640, s32[384,3]{1,0} %select.8153, f32[384]{0} %select.8152), update_window_dims={}, inserted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=2, indices_are_sorted=true, unique_indices=true, to_apply=%region_70.14671

I will turn off the expander by default.

[akuegel]

Actually it was already disabled in badb11c

[akuegel]

Reproducer (as a unit test):

TEST_F(ScatterDeterminismExpanderTest, Reproducer) {
  const char* const kModuleStr = R"(
    HloModule scatter_determinism_expander

    scatter_computation {
      arg1.173 = f32[] parameter(1)
      arg0.172 = f32[] parameter(0)
      ROOT add.48 = f32[] add(arg0.172, arg1.173)
    }

    ENTRY scatter_add_computation {
      operand = f32[1,128,128]{2,1,0} parameter(0)
      indices = s32[1,128,3]{2,1,0} parameter(1)
      updates = f32[1,128]{1,0} parameter(2)
      ROOT %scatter.33751 = f32[1,128,128]{2,1,0} scatter(operand, indices, updates), update_window_dims={}, inserted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=2, to_apply=scatter_computation
    }
  )";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          ParseAndReturnVerifiedModule(kModuleStr));

  ScatterDeterminismExpander scatter_determinism_expander;
  TF_ASSERT_OK_AND_ASSIGN(
      bool result, RunHloPass(&scatter_determinism_expander, module.get()));

  EXPECT_TRUE(result);
}