[XLA:GPU] Support proto payload in split gpu executable

PiperOrigin-RevId: 901065155

Author: thcmbs

xla/util/split_proto/BUILD

@@ -110,6 +110,7 @@ cc_library(
         "//xla:sort_json",
         "//xla:xla_proto_cc",
         "//xla/service:hlo_proto_cc",
+        "//xla/service:hlo_proto_util",
         "//xla/service/gpu:gpu_executable_proto_cc",
         "//xla/tsl/platform:errors",
         "@com_google_absl//absl/status",

xla/util/split_proto/split_gpu_executable_writer.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include "riegeli/bytes/writer.h"
 #include "riegeli/records/record_writer.h"
 #include "xla/service/hlo.pb.h"
+#include "xla/service/hlo_proto_util.h"
 #include "xla/sort_json.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/util/split_proto/split_proto.pb.h"
@@ -69,21 +70,46 @@ SplitProtoManifest BuildManifest(int32_t num_of_contants) {
 // deterministic because the order of keys in json is not guaranteed.
 //
 // If the backend config is not a json string, it is not modified.
-void NormalizeBackendConfig(gpu::GpuExecutableProto& executable) {
+absl::Status NormalizeBackendConfig(gpu::GpuExecutableProto& executable) {
   for (HloComputationProto& computation :
        *executable.mutable_hlo_module_with_config()
             ->mutable_hlo_module()
             ->mutable_computations()) {
     for (HloInstructionProto& instruction :
          *computation.mutable_instructions()) {
-      absl::StatusOr<std::string> normalized_backend_config =
-          SortJson(instruction.backend_config());
-      if (normalized_backend_config.ok()) {
-        instruction.set_backend_config(*normalized_backend_config);
+      TF_ASSIGN_OR_RETURN(
+          std::string backend_config_str,
+          GetBackendConfigString(
+              instruction, &executable.hlo_module_with_config().hlo_module()));
+      auto normalized_or = SortJson(backend_config_str);
+      if (!normalized_or.ok()) {
+        continue;
+      }
+      std::string normalized = std::move(normalized_or).value();
+      if (normalized == backend_config_str) {
+        continue;
+      }
+      if (instruction.has_backend_config_payload()) {
+        Payload* payload = instruction.mutable_backend_config_payload();
+        switch (payload->payload_source_case()) {
+          case Payload::kId: {
+            int id = static_cast<int>(payload->id());
+            auto* module = executable.mutable_hlo_module_with_config()
+                               ->mutable_hlo_module();
+            *module->mutable_payloads(id) = normalized;
+            break;
+          }
+          case Payload::kValue:
+          case Payload::PAYLOAD_SOURCE_NOT_SET:
+            payload->set_value(normalized);
+            break;
+        }
+      } else {
+        instruction.set_backend_config(normalized);
       }
-      // If the backend config is not a json string, then do nothing.
     }
   }
+  return absl::OkStatus();
 }
 
 }  // namespace
@@ -115,7 +141,7 @@ absl::Status WriteSplitGpuExecutable(gpu::GpuExecutableProto executable,
   executable.clear_constants();
 
   // The rest of the fields (i.e. the non-offloaded fields)
-  NormalizeBackendConfig(executable);
+  TF_RETURN_IF_ERROR(NormalizeBackendConfig(executable));
   // Module IDs are created via a static counter when deserializing, and they
   // can cause non-determinism, so we don't preserve them.
   executable.mutable_hlo_module_with_config()->mutable_hlo_module()->clear_id();

xla/util/split_proto/split_gpu_executable_writer_test.cc

@@ -113,12 +113,69 @@ TEST(SplitGpuExecutableWriterTest, NonJsonBackendConfigIsAccepted) {
        ->mutable_hlo_module()
        ->add_computations()
        ->add_instructions()
-       ->mutable_backend_config() = "not-json";
+       ->mutable_backend_config() = "x-json";
 
   std::string serialized1;
   ASSERT_OK(WriteSplitGpuExecutable(
       proto1, std::make_unique<riegeli::StringWriter<>>(&serialized1)));
 }
 
+TEST(SplitGpuExecutableWriterTest, JsonBackendConfigPayloadIsNormalized) {
+  GpuExecutableProto proto1;
+  proto1.mutable_hlo_module_with_config()
+      ->mutable_hlo_module()
+      ->add_computations()
+      ->add_instructions()
+      ->mutable_backend_config_payload()
+      ->set_value(R"json({"a": 1, "b": 2, "c": 3})json");
+
+  GpuExecutableProto proto2;
+  proto2.mutable_hlo_module_with_config()
+      ->mutable_hlo_module()
+      ->add_computations()
+      ->add_instructions()
+      ->mutable_backend_config_payload()
+      ->set_value(R"json({"c": 3, "b": 2, "a": 1})json");
+
+  std::string serialized1;
+  ASSERT_OK(WriteSplitGpuExecutable(
+      proto1, std::make_unique<riegeli::StringWriter<>>(&serialized1)));
+
+  std::string serialized2;
+  ASSERT_OK(WriteSplitGpuExecutable(
+      proto2, std::make_unique<riegeli::StringWriter<>>(&serialized2)));
+
+  EXPECT_EQ(serialized1, serialized2);
+}
+
+TEST(SplitGpuExecutableWriterTest,
+     JsonBackendConfigExternalPayloadIsNormalized) {
+  GpuExecutableProto proto1;
+  auto* module1 = proto1.mutable_hlo_module_with_config()->mutable_hlo_module();
+  module1->add_payloads(R"json({"a": 1, "b": 2, "c": 3})json");
+  module1->add_computations()
+      ->add_instructions()
+      ->mutable_backend_config_payload()
+      ->set_id(0);
+
+  GpuExecutableProto proto2;
+  auto* module2 = proto2.mutable_hlo_module_with_config()->mutable_hlo_module();
+  module2->add_payloads(R"json({"c": 3, "b": 2, "a": 1})json");
+  module2->add_computations()
+      ->add_instructions()
+      ->mutable_backend_config_payload()
+      ->set_id(0);
+
+  std::string serialized1;
+  ASSERT_OK(WriteSplitGpuExecutable(
+      proto1, std::make_unique<riegeli::StringWriter<>>(&serialized1)));
+
+  std::string serialized2;
+  ASSERT_OK(WriteSplitGpuExecutable(
+      proto2, std::make_unique<riegeli::StringWriter<>>(&serialized2)));
+
+  EXPECT_EQ(serialized1, serialized2);
+}
+
 }  // namespace
 }  // namespace xla