[GPU] Fix passing of key-value store handle from client to compiler.

[sergachev]

No description provided.

[PatriosTheGreat]

Thanks for this fix!

Just our of curiosity, which PJRT client you are using that goes directly to PjRtStreamExecutorClient::Compile?

I recently added the same fix for JAX to PyClient::CompileIfrtProgram: f548642

I'm curious if we need to keep only one of these fixes.

[sergachev]

Just our of curiosity, which PJRT client you are using that goes directly to PjRtStreamExecutorClient::Compile?

JAX uses this in the reproducer you provided for example. Not directly, but eventually it goes also through this method. Without this fix any attempt to use sharded autotuning fails at head.

[sergachev]

While this isn't merged, I realized, there is a better way to fix this: setting the key-value store in CompileInternal() is better because all paths merge at it. Added a 2nd test case for that.

[sergachev]

Rebased and resolved a conflict.

[reedwm]

I'm seeing internal test failures for tests //xla/pjrt/gpu:se_gpu_pjrt_client_test and //xla/pjrt:pjrt_stream_executor_client_test, but they don't fail externally. I'm also seeing a third internal-only test fail.

The issue with pjrt_stream_executor_client_test and the internal test seems to be calling key_value_store() in pjrt_stream_executor_client.cc without checking there is a value in the std::optional. I see the internal test using the InterpreterDevice which does not override key_value_store. Not sure about pjrt_stream_executor_client_test though, and why it passes externally. In any case, can you fix the case where key_value_store() is nullopt?

se_gpu_pjrt_client_test is failing internally with the error

INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc:1886) absl::StrContains( executable->GetHloModules()->front()->ToString(), "triton_gemm")

Any ideas what the issue is? If not, I can investigate internally.