Add a compiled_using_pjrt field to persistent cache key.

This is to differentiate the executables that would be built by XLA and PJRT clients when they are being serialized to disk for persistence.

There would be no changes to the filenames of XLA serialized executables. PJRT serialized executables will have '__pjrt' appended to their filenames.

PiperOrigin-RevId: 504338570

Author: swachhandl

tensorflow/compiler/jit/BUILD

@@ -1299,6 +1299,8 @@ cc_library(
         ":xla_device_compiler_client",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/xla:util",
+        "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/pjrt:pjrt_client",
         "//tensorflow/compiler/xla/service:hlo_proto_cc",
         "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
@@ -1393,6 +1395,7 @@ tf_cc_test(
     deps = [
         ":device_compiler_client",
         ":device_executable_persistor",
+        ":pjrt_device_compiler_client",
         ":xla_compilation_cache_proto_cc",
         ":xla_cpu_device",
         ":xla_cpu_jit",
@@ -1403,10 +1406,13 @@ tf_cc_test(
         "//tensorflow/compiler/xla/client:client_library",
         "//tensorflow/compiler/xla/client:executable_build_options",
         "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/pjrt:pjrt_client",
+        "//tensorflow/compiler/xla/pjrt:tfrt_cpu_pjrt_client",
         "//tensorflow/core:test",
         "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:status_matchers",
         "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/tfrt/common:pjrt_util",
         "@com_google_googletest//:gtest_main",
     ],
 )

tensorflow/compiler/jit/device_executable_persistor.h

@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
 #include "tensorflow/compiler/jit/xla_device_compiler_client.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
 #include "tensorflow/compiler/xla/service/hlo.pb.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/framework/device.h"
@@ -100,6 +101,10 @@ class DeviceExecutablePersistor {
   XlaSerializedCacheKey BuildSerializedCacheKey(
       uint64 signature_hash, const xla::HloModuleProto& hlo_module) const;
 
+  XlaSerializedCacheKey BuildSerializedCacheKey(
+      uint64 signature_hash, const xla::HloModuleProto& hlo_module,
+      bool compiled_using_pjrt) const;
+
   // Serializes the signature and its corresponding entry to a proto message.
   StatusOr<XlaSerializedCacheEntry> SerializeEntry(
       uint64 signature_hash, const XlaCompiler::Options& options,
@@ -154,7 +159,10 @@ std::string DeviceExecutablePersistor<ExecutableType, ClientType>::
       key.prefix(), key.prefix().empty() ? "" : kXlaSerializedCacheKeySeparator,
       key.signature_fingerprint(), kXlaSerializedCacheKeySeparator,
       key.cluster_fingerprint(), kXlaSerializedCacheKeySeparator,
-      key.device_type());
+      key.device_type(),
+      key.compiled_using_pjrt()
+          ? absl::StrCat(kXlaSerializedCacheKeySeparator, "pjrt")
+          : "");
 }
 
 template <typename ExecutableType, typename ClientType>
@@ -165,17 +173,35 @@ std::string DeviceExecutablePersistor<ExecutableType, ClientType>::GetFilePath(
   return io::JoinPath(persistent_cache_directory_, file_name);
 }
 
+template <typename ExecutableType, typename ClientType>
+XlaSerializedCacheKey
+DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
+    uint64 signature_hash, const xla::HloModuleProto& hlo_module,
+    bool compiled_using_pjrt) const {
+  XlaSerializedCacheKey key;
+  key.set_signature_fingerprint(signature_hash);
+  key.set_cluster_fingerprint(DeterministicProtoHash64(hlo_module));
+  key.set_device_type(device_type().type_string());
+  key.set_prefix(persistence_prefix());
+  key.set_compiled_using_pjrt(compiled_using_pjrt);
+  return key;
+}
+
 template <typename ExecutableType, typename ClientType>
 XlaSerializedCacheKey
 DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
     uint64 signature_hash, const xla::HloModuleProto& hlo_module) const {
-  XlaSerializedCacheKey serialized_cache_key;
-  serialized_cache_key.set_signature_fingerprint(signature_hash);
-  serialized_cache_key.set_cluster_fingerprint(
-      DeterministicProtoHash64(hlo_module));
-  serialized_cache_key.set_device_type(device_type().type_string());
-  serialized_cache_key.set_prefix(persistence_prefix());
-  return serialized_cache_key;
+  return BuildSerializedCacheKey(signature_hash, hlo_module, false);
+}
+
+// This template specialization sets compiled_using_prjt to true in the cache
+// key when the template arguments are PjRtLoadedExecutable and PjRtClient.
+template <>
+inline XlaSerializedCacheKey
+DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>::
+    BuildSerializedCacheKey(uint64 signature_hash,
+                            const xla::HloModuleProto& hlo_module) const {
+  return BuildSerializedCacheKey(signature_hash, hlo_module, true);
 }
 
 template <typename ExecutableType, typename ClientType>