Return std::nullopt from xla::ifrt::LoadedExecutable::devices() for portable executables

Portable executables are not bound to any specific device list, so it makes sense to be explicit about it rather than letting each implementation return an unspecified value.

PiperOrigin-RevId: 837188096

Author: junwhanahn

xla/backends/cpu/nanort/ifrt_client.cc

@@ -1009,7 +1009,9 @@ class NanoExecutable final
     return client_->addressable_devices();
   }
 
-  const ifrt::DeviceListRef& devices() const override { return devices_; }
+  std::optional<ifrt::DeviceListRef> devices() const override {
+    return devices_;
+  }
 
   static char ID;  // NOLINT
 

xla/backends/cpu/nanort/ifrt_client_test.cc

@@ -337,6 +337,8 @@ int main(int argc, char** argv) {
       "*LoadedExecutableImplTest.GetHloModules*:"
       "*LoadedExecutableImplTest.ProgramText*:"
       "*LoadedExecutableImplTest.Analysis*:"
+      // NanoRT does not support portable execution.
+      "*LoadedExecutableImplTest.CompileAndExecutePortable*:"
       // Serialization is not implemented.
       "*SerializeAndLoad*";
   xla::ifrt::test_util::SetTestFilterIfNotUserSpecified(kFilter);

xla/python/ifrt/executable.h

@@ -279,8 +279,9 @@ class LoadedExecutable
       std::optional<DeviceListRef> devices) = 0;
 
   // Returns the list of devices where the executable has been compiled and
-  // loaded onto.
-  virtual const DeviceListRef& devices() const = 0;
+  // loaded onto. Returns `std::nullopt` if the executable is not bound to a
+  // particular device list, e.g., portable executables.
+  virtual std::optional<DeviceListRef> devices() const = 0;
 
   // The following APIs are taken from xla::PjRtLoadedExecutable for fast
   // prototyping.

xla/python/ifrt/mock.h

@@ -317,13 +317,6 @@ class MockExecutable : public llvm::RTTIExtends<MockExecutable, Executable> {
 class MockLoadedExecutable
     : public llvm::RTTIExtends<MockLoadedExecutable, LoadedExecutable> {
  public:
-  MockLoadedExecutable() {
-    static absl::NoDestructor<DeviceListRef> kEmptyDeviceList(
-        BasicDeviceList::Create({}));
-    ON_CALL(*this, devices())
-        .WillByDefault(testing::ReturnRef(*kEmptyDeviceList));
-  }
-
   MOCK_METHOD(Client*, client, (), (const, final));
   MOCK_METHOD(absl::string_view, name, (), (const, final));
   MOCK_METHOD(absl::StatusOr<std::optional<std::string>>, Fingerprint, (),
@@ -363,7 +356,7 @@ class MockLoadedExecutable
               (final));
   MOCK_METHOD(absl::Span<Device* const>, addressable_devices, (),
               (const, final));
-  MOCK_METHOD(const DeviceListRef&, devices, (), (const, final));
+  MOCK_METHOD(std::optional<DeviceListRef>, devices, (), (const, final));
 
   static char ID;  // NOLINT
 };

xla/python/ifrt_proxy/client/compiler.cc

@@ -177,8 +177,10 @@ absl::StatusOr<xla::ifrt::LoadedExecutableRef> Compiler::CompileAndLoad(
       devices.push_back(device);
     }
   }
-  TF_ASSIGN_OR_RETURN(DeviceListRef device_list,
-                      client_->MakeDeviceList(devices));
+  std::optional<DeviceListRef> device_list;
+  if (!devices.empty()) {
+    TF_ASSIGN_OR_RETURN(device_list, client_->MakeDeviceList(devices));
+  }
 
   return std::make_unique<LoadedExecutable>(
       client_, rpc_helper_, response->loaded_executable_handle(),

xla/python/ifrt_proxy/client/executable.cc

@@ -337,7 +337,8 @@ class LoadedExecutable::OutputSpecCache {
 
 LoadedExecutable::LoadedExecutable(
     xla::ifrt::Client* client, std::shared_ptr<RpcHelper> rpc_helper,
-    uint64_t handle, std::string name, int num_devices, DeviceListRef devices,
+    uint64_t handle, std::string name, int num_devices,
+    std::optional<DeviceListRef> devices,
     std::vector<xla::ifrt::Device*> addressable_devices,
     absl::StatusOr<std::optional<std::string>> fingerprint,
     tsl::Future<> ready_future,
@@ -843,7 +844,9 @@ LoadedExecutable::Execute(absl::Span<xla::ifrt::ArrayRef> args,
   return result;
 }
 
-const DeviceListRef& LoadedExecutable::devices() const { return devices_; }
+std::optional<DeviceListRef> LoadedExecutable::devices() const {
+  return devices_;
+}
 
 absl::Span<xla::ifrt::Device* const> LoadedExecutable::addressable_devices()
     const {

xla/python/ifrt_proxy/client/executable.h

@@ -59,7 +59,8 @@ class LoadedExecutable final
  public:
   LoadedExecutable(xla::ifrt::Client* client,
                    std::shared_ptr<RpcHelper> rpc_helper, uint64_t handle,
-                   std::string name, int num_devices, DeviceListRef devices,
+                   std::string name, int num_devices,
+                   std::optional<DeviceListRef> devices,
                    std::vector<xla::ifrt::Device*> addressable_devices,
                    absl::StatusOr<std::optional<std::string>> fingerprint,
                    tsl::Future<> ready_future,
@@ -111,7 +112,7 @@ class LoadedExecutable final
       absl::Span<xla::ifrt::ArrayRef> args, const ExecuteOptions& options,
       std::optional<xla::ifrt::DeviceListRef> devices) override;
 
-  const DeviceListRef& devices() const override;
+  std::optional<DeviceListRef> devices() const override;
   absl::Span<xla::ifrt::Device* const> addressable_devices() const override;
 
   static char ID;  // NOLINT
@@ -148,7 +149,7 @@ class LoadedExecutable final
   const uint64_t handle_;
   const std::string name_;
   const int num_devices_;
-  const DeviceListRef devices_;
+  const std::optional<DeviceListRef> devices_;
   const std::vector<xla::ifrt::Device*> addressable_devices_;
   const absl::StatusOr<std::optional<std::string>> fingerprint_;
   const tsl::Future<> ready_future_;

xla/python/ifrt_proxy/server/ifrt_backend.cc

@@ -1456,8 +1456,12 @@ tsl::Future<BackendInterface::Response> IfrtBackend::HandleCompileRequest(
     for (const auto* device : executable->addressable_devices()) {
       compile_resp->add_addressable_device_ids(device->Id().value());
     }
-    for (const auto* device : executable->devices()->devices()) {
-      compile_resp->add_device_ids(device->Id().value());
+    if (std::optional<xla::ifrt::DeviceListRef> device_list =
+            executable->devices();
+        device_list.has_value()) {
+      for (const auto* device : (*device_list)->devices()) {
+        compile_resp->add_device_ids(device->Id().value());
+      }
     }
     // TODO(b/282757875): Consider making fingerprint calculation asynchronous
     // if it is expected to take long.

xla/python/ifrt_proxy/server/ifrt_backend_test.cc

@@ -1157,7 +1157,8 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccess) {
   auto executable = std::make_unique<MockLoadedExecutable>();
   EXPECT_CALL(*executable, name()).WillOnce(Return("executable_name"));
   EXPECT_CALL(*executable, num_devices()).WillOnce(Return(4));
-  EXPECT_CALL(*executable, devices()).WillOnce(ReturnRef(device_list));
+  EXPECT_CALL(*executable, devices())
+      .WillOnce(Return(std::make_optional(device_list)));
   EXPECT_CALL(*executable, addressable_devices())
       .WillOnce(Return(absl::MakeSpan(addressable_devices)));
   EXPECT_CALL(*executable, Fingerprint()).WillOnce(Return("fingerprint"));

xla/python/pjrt_ifrt/pjrt_executable.h

@@ -316,7 +316,14 @@ class PjRtLoadedExecutable final
       absl::Span<ArrayRef> args, const ExecuteOptions& options,
       std::optional<DeviceListRef> devices) override;
 
-  const DeviceListRef& devices() const override { return devices_; }
+  std::optional<DeviceListRef> devices() const override {
+    if (pjrt_loaded_executable_->addressable_devices().empty()) {
+      // Portable executable.
+      return std::nullopt;
+    } else {
+      return devices_;
+    }
+  }
 
   absl::Span<Device* const> addressable_devices() const override {
     DCHECK(this);