[PjRt-IFRT] Create ifrt::PjRtExecutable only from ifrt::PjRtCompiler and CompileOnlyIfrtCompiler

This change migrates direct calls to `ifrt::PjRtExecutable::Create()` outside to use a public IFRT API `ifrt::PjRtCompiler::Compile()` instead.

This change should be no-op in practice. For PjRt-IFRT, it now performs IFRT device ID -> PjRt device ID conversion in `xla::CompileOptions::executable_build_options` (which was missing before) and thus can handle a client using a different device ID mapping.

PiperOrigin-RevId: 844980890

Author: hyeontaek

jaxlib/py_client.cc

@@ -499,27 +499,33 @@ PyClient::CompileAndLoadIfrtProgram(
     ifrt::DeviceListRef executable_devices, xla::CompileOptions options) {
   mlir::OwningOpRef<mlir::ModuleOp> clone(module.clone());
   module = *clone;
-  ifrt::ExecutableRef executable_ref;
+  ifrt::ExecutableRef ifrt_executable;
   {
     TF_ASSIGN_OR_RETURN(
         auto topology,
         client->ifrt_client()->GetTopologyForDevices(executable_devices));
     auto xla_options = std::make_unique<ifrt::XlaCompileOptions>(
         options, std::move(executable_devices));
-#if JAX_IFRT_VERSION_NUMBER >= 38
+#if JAX_IFRT_VERSION_NUMBER >= 42
     TF_ASSIGN_OR_RETURN(
-        executable_ref,
+        ifrt_executable,
+        client->ifrt_client()->GetDefaultCompiler()->Compile(
+            std::make_unique<xla::ifrt::HloProgram>(std::move(module)),
+            *topology, std::move(xla_options)));
+#elif JAX_IFRT_VERSION_NUMBER >= 38
+    TF_ASSIGN_OR_RETURN(
+        ifrt_executable,
         ifrt::PjRtExecutable::Create(std::move(module), std::move(options),
                                      *topology->description()));
 #else
     TF_ASSIGN_OR_RETURN(
         auto pjrt_executable,
         PjRtCompile(std::move(options), module, *topology->description()));
-    TF_ASSIGN_OR_RETURN(executable_ref, ifrt::PjRtExecutable::Create(
-                                            std::move(pjrt_executable)));
+    TF_ASSIGN_OR_RETURN(ifrt_executable, ifrt::PjRtExecutable::Create(
+                                             std::move(pjrt_executable)));
 #endif
   }
-  return make_nb_class<PyExecutable>(executable_ref);
+  return make_nb_class<PyExecutable>(ifrt_executable);
 }
 
 /* static */ absl::StatusOr<nb_class_ptr<PyLoadedExecutable>>

jaxlib/py_compile_only_client.cc

@@ -40,6 +40,7 @@ limitations under the License.
 #include "xla/python/compile_only_ifrt/client.h"
 #include "xla/python/ifrt/device_list.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/pjrt_ifrt/pjrt_executable.h"
 #include "xla/python/pjrt_ifrt/pjrt_topology.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
@@ -79,7 +80,13 @@ absl::StatusOr<nb_class_ptr<PyExecutable>> CompileOnlyPyClient::CompileUnloaded(
 
     auto xla_options = std::make_unique<ifrt::XlaCompileOptions>(
         options, std::move(executable_devices));
-#if JAX_IFRT_VERSION_NUMBER >= 38
+#if JAX_IFRT_VERSION_NUMBER >= 42
+    TF_ASSIGN_OR_RETURN(
+        ifrt_executable,
+        ifrt_client->GetDefaultCompiler()->Compile(
+            std::make_unique<xla::ifrt::HloProgram>(std::move(module)),
+            ifrt_client->topology(), std::move(xla_options)));
+#elif JAX_IFRT_VERSION_NUMBER >= 38
     TF_ASSIGN_OR_RETURN(
         ifrt_executable,
         ifrt::PjRtExecutable::Create(std::move(module), std::move(options),