Enabling stablehlo-complex-math-expander pass is platform dependent

[pearu]

Consider JAX expression

jax.numpy.log1p(jax.numpy.complex64(0))

and the corresponding compilation paths for CPU and CUDA platforms. Currently, the CUDA path never calls MlirToXlaComputation which is unexpected per #20853 (comment) .
In particular, the difference of CPU and CUDA compilation paths is as follows:

CPU:
xla::ifrt::PjRtLoadedExecutable::Create 
  -> xla::TfrtCpuClient::Compile 
  -> xla::MlirToXlaComputation

CUDA:
xla::ifrt::PjRtLoadedExecutable::Create 
  -> xla::PjRtCApiClient::Compile 
  -> xla::Serialize 
  -> xla::SerializeUsingVersionedStablehlo

In fact, no XlaComputation instance is created in the CUDA compilation path which indicates that calling MlirToXlaComputation is platform dependent.

The aim of this issue is to define a platform independent compilation path for the stablehlo-complex-math-expander pass that preferably uses MlirToXlaComputation.

Originally posted by @pearu in #20853 (comment)

[akuegel]

Was this fixed with #20853 ?

[pearu]

No, #20853 just revealed this issue.

[GleasonK]

There are a few StableHLO passes we want to run for all XLA PJRT plugins (complex math expanders, quantized math expanders, etc), and there isn't a good single place we can add these passes. Seems like plugins are using multiple methods to go from StableHLO->HLO which is surprising. This ticket is about consolidating to a single HLO conversion path for plugins.