From 50e18012ca01aca5c3f8167489bcdae8a29008ba Mon Sep 17 00:00:00 2001
From: Shriram <shriram00@hotmail.co.uk>
Date: Mon, 23 Mar 2026 03:12:37 +0000
Subject: [PATCH] Fix persistent compilation cache not working on TPU

The TPU backend (via PjRtCApiClient) implements LoadSerializedExecutable
but does not override DeserializeExecutable, which returns UNIMPLEMENTED
from the base PjRtClient class. This caused every cache load attempt to
fail, forcing recompilation on every run even when valid cached
executables existed on disk.

Fix by trying LoadSerializedExecutable first (the path implemented by
PJRT C API plugins like TPU), then falling back to the two-step
DeserializeExecutable + Load path for backends that implement that
instead.

Fixes #9094

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../csrc/runtime/pjrt_computation_client.cpp  | 38 ++++++++++++++-----
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cpp b/torch_xla/csrc/runtime/pjrt_computation_client.cpp
index aa0145481e3..66d59c877ba 100644
--- a/torch_xla/csrc/runtime/pjrt_computation_client.cpp
+++ b/torch_xla/csrc/runtime/pjrt_computation_client.cpp
@@ -697,17 +697,35 @@ std::string PjRtComputationClient::SerializeComputation(
 
 ComputationClient::ComputationPtr PjRtComputationClient::DeserializeComputation(
     const std::string& serialized) {
-  absl::StatusOr<std::unique_ptr<xla::PjRtExecutable>> executable_or =
-      client_->DeserializeExecutable(serialized, std::nullopt);
-  if (!executable_or.ok()) {
-    TF_LOG(WARNING) << "Failed to deserialize executable: "
-                    << executable_or.status();
-    return nullptr;
+  std::unique_ptr<xla::PjRtLoadedExecutable> loaded_executable;
+
+  // First, try LoadSerializedExecutable which directly produces a loaded
+  // executable. This is the path implemented by PJRT C API plugins (e.g. TPU).
+  absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
+      loaded_executable_or = client_->LoadSerializedExecutable(
+          serialized, std::nullopt, xla::LoadOptions());
+  if (loaded_executable_or.ok()) {
+    loaded_executable = std::move(loaded_executable_or.value());
+  } else {
+    // Fall back to the two-step DeserializeExecutable + Load path for backends
+    // that implement DeserializeExecutable instead of LoadSerializedExecutable.
+    absl::StatusOr<std::unique_ptr<xla::PjRtExecutable>> executable_or =
+        client_->DeserializeExecutable(serialized, std::nullopt);
+    if (!executable_or.ok()) {
+      TF_LOG(WARNING) << "Failed to deserialize executable: "
+                      << loaded_executable_or.status() << " ; "
+                      << executable_or.status();
+      return nullptr;
+    }
+    absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> load_or =
+        client_->Load(std::move(executable_or.value()), xla::LoadOptions());
+    if (!load_or.ok()) {
+      TF_LOG(WARNING) << "Failed to load deserialized executable: "
+                      << load_or.status();
+      return nullptr;
+    }
+    loaded_executable = std::move(load_or.value());
   }
-  std::unique_ptr<xla::PjRtExecutable> executable =
-      std::move(executable_or.value());
-  std::unique_ptr<xla::PjRtLoadedExecutable> loaded_executable =
-      client_->Load(std::move(executable), xla::LoadOptions()).value();
 
   auto hlo_modules = loaded_executable->GetHloModules();
   if (!hlo_modules.ok()) {