From 1d4cc1bac908b8884d3b353c0957a6a42a89c8a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 24 Sep 2026 07:38:15 -0700
Subject: [PATCH] Parse split-proto merge records from a riegeli::Chain without
 flattening.

`HandleProtoMergeRecord` read each record into a flat `absl::string_view` and
then called `MergeFromString`. For records larger than a riegeli buffer (e.g.
an HLO module carrying hundreds of MB of constants) the decoded record spans
several blocks, and producing a flat view forces riegeli to copy it into a
contiguous scratch buffer first.

Read the record as a `riegeli::Chain` instead, which shares the decoded
blocks, and parse it with `riegeli::ParseMessage(..., set_merge(true))`,
which consumes the chain directly.

PiperOrigin-RevId: 987496597
---
 third_party/xla/xla/util/split_proto/BUILD      |  2 ++
 .../xla/util/split_proto/split_proto_reader.cc  | 17 ++++++++++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/util/split_proto/BUILD b/third_party/xla/xla/util/split_proto/BUILD
index 3f8d0dc579c..8a512236254 100644
--- a/third_party/xla/xla/util/split_proto/BUILD
+++ b/third_party/xla/xla/util/split_proto/BUILD
@@ -74,7 +74,9 @@ cc_library(
         "@com_google_absl//absl/strings:string_view",
         "@com_google_protobuf//:protobuf",
         "@riegeli//riegeli/base:any",
+        "@riegeli//riegeli/base:chain",
         "@riegeli//riegeli/bytes:reader",
+        "@riegeli//riegeli/messages:parse_message",
         "@riegeli//riegeli/records:record_position",
         "@riegeli//riegeli/records:record_reader",
     ],
diff --git a/third_party/xla/xla/util/split_proto/split_proto_reader.cc b/third_party/xla/xla/util/split_proto/split_proto_reader.cc
index 47fc1419aa0..3e3a18284cc 100644
--- a/third_party/xla/xla/util/split_proto/split_proto_reader.cc
+++ b/third_party/xla/xla/util/split_proto/split_proto_reader.cc
@@ -27,7 +27,9 @@ limitations under the License.
 #include "google/protobuf/message.h"
 #include "google/protobuf/reflection.h"
 #include "riegeli/base/any.h"
+#include "riegeli/base/chain.h"
 #include "riegeli/bytes/reader.h"
+#include "riegeli/messages/parse_message.h"
 #include "riegeli/records/record_position.h"
 #include "riegeli/records/record_reader.h"
 #include "xla/tsl/platform/errors.h"
@@ -52,15 +54,20 @@ absl::Status ReadRecord(riegeli::RecordReader<Src>& record_reader, T& record) {
 template <typename Src>
 absl::Status HandleProtoMergeRecord(riegeli::RecordReader<Src>& record_reader,
                                     google::protobuf::Message& proto) {
-  absl::string_view record_data;
+  // Read the record as a `Chain` rather than a flat `string_view`: large
+  // records (e.g. HLO modules with big constants) span multiple blocks of the
+  // decoded chunk, and reading them as a flat view would force an extra copy
+  // to make them contiguous. The `Chain` shares the decoded blocks instead and
+  // the proto parser can consume it directly.
+  riegeli::Chain record_data;
   TF_RETURN_WITH_CONTEXT_IF_ERROR(ReadRecord(record_reader, record_data),
                                   "failed to read proto merge record data for ",
                                   proto.GetTypeName());
 
-  if (!proto.MergeFromString(record_data)) {
-    return absl::InternalError(absl::StrFormat(
-        "Failed to parse proto merge record for %s", proto.GetTypeName()));
-  }
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      riegeli::ParseMessage(record_data, proto,
+                            riegeli::ParseMessageOptions().set_merge(true)),
+      "Failed to parse proto merge record for ", proto.GetTypeName());
   return absl::OkStatus();
 }