Skip to content

Commit 81faaf8

Browse files
rpsilva-awsqihqi
rpsilva-aws
authored and
qihqi
committed
[Computation Hash] Introduce deterministic hash for user computations (#8539)
1 parent 10cee98 commit 81faaf8

File tree

8 files changed

+211
-17
lines changed

8 files changed

+211
-17
lines changed

‎torch_xla/csrc/init_python_bindings.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1169,9 +1169,8 @@ class PyLoweringContext {
11691169
// Create a serialized HloModule protobuf from a lowered graph
11701170
py::bytes GetHlo() {
11711171
const xla::HloModuleProto& proto = computation.proto();
1172-
std::string result;
1173-
proto.SerializeToString(&result);
1174-
return result;
1172+
return ConsumeValue(
1173+
runtime::util::GetDeterministicSerializedModuleProto(proto));
11751174
}
11761175

11771176
// Create human-readable HloModule protobuf text from a lowered graph

‎torch_xla/csrc/runtime/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ cc_library(
5050
":types",
5151
":util",
5252
":xla_coordinator",
53+
":xla_util",
5354
"//torch_xla/csrc:device",
5455
"//torch_xla/csrc:dtype",
5556
"@com_google_absl//absl/memory",
@@ -460,6 +461,7 @@ ptxla_cc_test(
460461
size = "small",
461462
srcs = ["xla_util_test.cc"],
462463
deps = [
464+
":debug_macros",
463465
":xla_util",
464466
"@com_google_absl//absl/types:span",
465467
"@com_google_googletest//:gtest_main",

‎torch_xla/csrc/runtime/computation_client.cc‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#include "torch_xla/csrc/runtime/debug_macros.h"
1414
#include "torch_xla/csrc/runtime/env_vars.h"
1515
#include "torch_xla/csrc/runtime/sys_util.h"
16+
#include "torch_xla/csrc/runtime/xla_util.h"
1617
#include "tsl/platform/stacktrace_handler.h"
1718
#include "xla/status_macros.h"
1819

@@ -194,5 +195,13 @@ metrics::Metric* ComputationClient::OutboundDataMetric() {
194195
return metric;
195196
}
196197

198+
::absl::StatusOr<torch::lazy::hash_t>
199+
ComputationClient::Computation::ComputeHash(const xla::HloModuleProto& proto,
200+
const std::string& name) {
201+
TF_ASSIGN_OR_RETURN(auto serialized_status,
202+
util::GetDeterministicSerializedModuleProto(proto));
203+
return torch::lazy::MHash(name, serialized_status);
204+
}
205+
197206
} // namespace runtime
198207
} // namespace torch_xla

‎torch_xla/csrc/runtime/computation_client.h‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,8 @@ class ComputationClient {
115115
computation_(std::move(computation)),
116116
devices_(std::move(devices)) {
117117
program_shape_ = ConsumeValue(computation_.GetProgramShape());
118-
hash_ =
119-
torch::lazy::MHash(name, computation_.proto().SerializeAsString());
118+
const xla::HloModuleProto& proto = computation_.proto();
119+
hash_ = ConsumeValue(ComputeHash(proto, name));
120120
}
121121

122122
Computation(std::string name, xla::XlaComputation computation,
@@ -159,7 +159,7 @@ class ComputationClient {
159159
// here.
160160
xla::XlaComputation move_computation() {
161161
if (computation_moved_) {
162-
XLA_ERROR() << "Compuation has been moved\n";
162+
XLA_ERROR() << "Computation has been moved\n";
163163
}
164164
computation_moved_ = true;
165165
return std::move(const_cast<Computation*>(this)->computation_);
@@ -206,6 +206,13 @@ class ComputationClient {
206206

207207
torch::lazy::hash_t hash_;
208208
std::string name_;
209+
210+
// Computes a hash for an HLO module using deterministic proto
211+
// serialization. It ensures consistent ordering of Map fields and repeated
212+
// elements during during serialization. The resulting hash combines the
213+
// serialized module with its computation name.
214+
static ::absl::StatusOr<torch::lazy::hash_t> ComputeHash(
215+
const xla::HloModuleProto& proto, const std::string& name);
209216
};
210217

211218
using ComputationPtr = std::shared_ptr<Computation>;

‎torch_xla/csrc/runtime/xla_util.cc‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include "tsl/platform/errors.h"
1515
#include "tsl/platform/stacktrace.h"
1616
#include "xla/shape_util.h"
17+
#include "xla/tsl/lib/strings/proto_serialization.h"
1718
#include "xla/util.h"
1819

1920
namespace torch_xla {
@@ -115,6 +116,26 @@ torch::lazy::hash_t ShapeHash(const xla::Shape& shape) {
115116
return hash;
116117
}
117118

119+
absl::StatusOr<std::string> GetDeterministicSerializedModuleProto(
120+
const xla::HloModuleProto& hlo_proto) {
121+
const size_t size = hlo_proto.ByteSizeLong();
122+
if (size == 0) {
123+
return std::string();
124+
}
125+
std::string serialized;
126+
// Pre-allocate the string buffer for the serialized result.
127+
serialized.resize(size);
128+
129+
// Perform deterministic serialization ensuring consistent ordering
130+
// of map fields and repeated elements
131+
if (!tsl::SerializeToBufferDeterministic(hlo_proto, serialized.data(),
132+
size)) {
133+
return absl::InvalidArgumentError("Could not serialize module proto");
134+
}
135+
136+
return serialized;
137+
}
138+
118139
} // namespace util
119140
} // namespace runtime
120141
} // namespace torch_xla

‎torch_xla/csrc/runtime/xla_util.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,12 @@ void CheckComputationStatus(
4040

4141
torch::lazy::hash_t ShapeHash(const xla::Shape& shape);
4242

43+
// Return the serialized module proto, using deterministic proto serialization.
44+
// It ensures consistent ordering of Map fields and repeated elements during
45+
// serialization.
46+
absl::StatusOr<std::string> GetDeterministicSerializedModuleProto(
47+
const xla::HloModuleProto& hlo_proto);
48+
4349
} // namespace util
4450
} // namespace runtime
4551
} // namespace torch_xla

‎torch_xla/csrc/runtime/xla_util_test.cc‎

Lines changed: 149 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@
33
#include <gmock/gmock.h>
44
#include <gtest/gtest.h>
55

6+
#include <random>
67
#include <set>
8+
#include <string>
79
#include <unordered_map>
810
#include <vector>
911

1012
#include "absl/status/status.h"
1113
#include "absl/types/span.h"
14+
#include "torch_xla/csrc/runtime/debug_macros.h"
1215
#include "tsl/platform/errors.h"
1316
#include "tsl/platform/protobuf.h"
1417
#include "tsl/platform/status_matchers.h"
@@ -46,7 +49,7 @@ absl::StatusOr<MessageType> ParseTextProto(const std::string& text_proto) {
4649
return parsed_proto;
4750
}
4851

49-
TEST(XlaUtilrest, CreateModule) {
52+
TEST(XlaUtilTest, CreateModule) {
5053
TF_ASSERT_OK_AND_ASSIGN(
5154
xla::HloModuleProto hlo_module_proto,
5255
ParseTextProto<xla::HloModuleProto>(
@@ -102,7 +105,7 @@ TEST(XlaUtilrest, CreateModule) {
102105
EXPECT_EQ((*got)->computation_count(), 1);
103106
}
104107

105-
TEST(XlaUtilrest, XlaToHlo) {
108+
TEST(XlaUtilTest, XlaToHlo) {
106109
xla::Shape input_shape =
107110
xla::ShapeUtil::MakeShape(xla::PrimitiveType::F32, {2, 2});
108111
xla::XlaBuilder builder("AddComputation");
@@ -116,6 +119,150 @@ TEST(XlaUtilrest, XlaToHlo) {
116119
HasSubstr("ROOT %add.3"))));
117120
}
118121

122+
TEST(XlaUtilTest, TestDeterministicModuleProtoSerializationEmptyProto) {
123+
xla::HloModuleProto empty_proto;
124+
auto result =
125+
::ConsumeValue(GetDeterministicSerializedModuleProto(empty_proto));
126+
// Verify that the result is an empty string
127+
EXPECT_TRUE(result.empty());
128+
}
129+
130+
TEST(XlaUtilTest, TestDeterministicModuleProtoSerialization) {
131+
// Create a test HLO module with a known structure
132+
TF_ASSERT_OK_AND_ASSIGN(
133+
xla::HloModuleProto hlo_module_proto,
134+
ParseTextProto<xla::HloModuleProto>(
135+
R"pb(
136+
name: "myname"
137+
id: 9
138+
entry_computation_name: "MyCustomName.9"
139+
entry_computation_id: 9
140+
computations {
141+
id: 9
142+
name: "MyCustomName.9"
143+
instructions: {
144+
name: "p0.1"
145+
id: 1
146+
opcode: "parameter"
147+
shape: {
148+
element_type: S64
149+
layout { tail_padding_alignment_in_elements: 1 }
150+
}
151+
metadata {
152+
op_type: "xla__device_data"
153+
op_name: "xla__device_data"
154+
source_file: "/ansible/pytorch/xla/small_test.py"
155+
source_line: 14
156+
stack_frame_id: 1
157+
}
158+
}
159+
instructions: {
160+
name: "p1.2"
161+
id: 2
162+
opcode: "parameter"
163+
parameter_number: 1
164+
shape: {
165+
element_type: S64
166+
layout { tail_padding_alignment_in_elements: 1 }
167+
}
168+
metadata {
169+
op_type: "xla__device_data"
170+
op_name: "xla__device_data"
171+
source_file: "/ansible/pytorch/xla/small_test.py"
172+
source_line: 13
173+
stack_frame_id: 2
174+
}
175+
}
176+
instructions: {
177+
name: "call.7"
178+
id: 7
179+
opcode: "call"
180+
shape: {
181+
element_type: S64
182+
layout { tail_padding_alignment_in_elements: 1 }
183+
}
184+
metadata {
185+
op_type: "xla___op_some_op"
186+
op_name: "xla___op_some_op"
187+
source_file: "/ansible/pytorch/xla/torch_xla/core/xla_op_registry.py"
188+
source_line: 44
189+
stack_frame_id: 4
190+
}
191+
called_computation_ids: 3
192+
operand_ids: 2
193+
operand_ids: 1
194+
}
195+
instructions: {
196+
name: "tuple.8"
197+
id: 8
198+
opcode: "tuple"
199+
shape: {
200+
element_type: TUPLE
201+
tuple_shapes {
202+
element_type: S64
203+
layout { tail_padding_alignment_in_elements: 1 }
204+
}
205+
}
206+
operand_ids: 7
207+
}
208+
root_id: 8
209+
}
210+
host_program_shape: {
211+
parameters {
212+
element_type: S64
213+
layout { tail_padding_alignment_in_elements: 1 }
214+
}
215+
parameters {
216+
element_type: S64
217+
layout { tail_padding_alignment_in_elements: 1 }
218+
}
219+
result {
220+
element_type: TUPLE
221+
tuple_shapes {
222+
element_type: S64
223+
layout { tail_padding_alignment_in_elements: 1 }
224+
}
225+
}
226+
parameter_names: "p0"
227+
parameter_names: "p1"
228+
}
229+
)pb"));
230+
231+
// Define a set of dummy fixed key-value pairs for frontend attributes.
232+
std::vector<std::pair<std::string, std::string>> attr_pairs = {
233+
{"key1", "value1"},
234+
{"key2", "value2"},
235+
{"key3", "value3"},
236+
{"key4", "value4"}};
237+
238+
auto shuffle_and_hash = [&attr_pairs](xla::HloModuleProto hlo_module_proto) {
239+
// Create a random number generator for shuffling.
240+
std::random_device random_device;
241+
std::mt19937 random_generator(random_device());
242+
243+
for (auto& computation : *hlo_module_proto.mutable_computations()) {
244+
for (auto& instruction : *computation.mutable_instructions()) {
245+
std::shuffle(attr_pairs.begin(), attr_pairs.end(), random_generator);
246+
auto* frontend_attrs = instruction.mutable_frontend_attributes();
247+
// Add the dummy shuffled pairs to the frontend attributes.
248+
for (const auto& pair : attr_pairs) {
249+
(*frontend_attrs->mutable_map())[pair.first] = pair.second;
250+
}
251+
}
252+
}
253+
std::string serialized_proto =
254+
::ConsumeValue(GetDeterministicSerializedModuleProto(hlo_module_proto));
255+
return torch::lazy::Hash(serialized_proto);
256+
};
257+
258+
// Compute hashes with different random orderings of attributes
259+
torch::lazy::hash_t hash1 = shuffle_and_hash(hlo_module_proto);
260+
torch::lazy::hash_t hash2 = shuffle_and_hash(hlo_module_proto);
261+
// Verify that different orderings produce the same hash
262+
ASSERT_EQ(hash1, hash2)
263+
<< "Hashes should match regardless of the frontend attribute ordering";
264+
}
265+
119266
} // namespace util
120267
} // namespace runtime
121268
} // namespace torch_xla

‎torch_xla/csrc/xla_graph_executor.cpp‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1200,12 +1200,13 @@ XLAGraphExecutor::LookupCachedCompile(const torch::lazy::hash_t& hash) {
12001200
TORCH_LAZY_COUNTER("UncachedCompile", 1);
12011201
return nullptr;
12021202
}
1203+
std::string serialized_computation =
1204+
ConsumeValue(runtime::util::GetDeterministicSerializedModuleProto(
1205+
cached_computation->computation->computation().proto()));
12031206
TF_VLOG(5) << "Graph hash " << torch::lazy::HashToString(hash)
12041207
<< " is computation hash "
1205-
<< torch::lazy::HashToString(torch::lazy::Hash(
1206-
cached_computation->computation->computation()
1207-
.proto()
1208-
.SerializeAsString()));
1208+
<< torch::lazy::HashToString(
1209+
torch::lazy::Hash(serialized_computation));
12091210
TORCH_LAZY_COUNTER("CachedCompile", 1);
12101211
return cached_computation;
12111212
}
@@ -1443,11 +1444,13 @@ XLAGraphExecutor::CompilationResult XLAGraphExecutor::Compile(
14431444
<< coll.device << " done!";
14441445
TF_VLOG(5) << "Compiled program shape "
14451446
<< computations.front()->program_shape().ToString() << std::endl;
1446-
TF_VLOG(5)
1447-
<< "Graph hash " << torch::lazy::HashToString(coll.hash)
1448-
<< " is computation hash "
1449-
<< torch::lazy::HashToString(torch::lazy::Hash(
1450-
computations.front()->computation().proto().SerializeAsString()));
1447+
std::string serialized_computation =
1448+
ConsumeValue(runtime::util::GetDeterministicSerializedModuleProto(
1449+
computations.front()->computation().proto()));
1450+
TF_VLOG(5) << "Graph hash " << torch::lazy::HashToString(coll.hash)
1451+
<< " is computation hash "
1452+
<< torch::lazy::HashToString(
1453+
torch::lazy::Hash(serialized_computation));
14511454

14521455
if (use_autosharding) {
14531456
const xla::HloModuleProto& computation_proto =

0 commit comments

Comments
 (0)