Skip to content

Replaces the use of RemapInstructionIds with HloModule::CreateFromProto without preserving unique id in xprof. #1856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 27, 2025
+128 −29
Merged

Replaces the use of RemapInstructionIds with HloModule::CreateFromProto without preserving unique id in xprof. #1856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 28 additions & 25 deletions WORKSPACE
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,39 @@ http_archive(
url = "https://github.com/bazelbuild/rules_java/releases/download/8.7.0/rules_java-8.7.0.tar.gz",
)

# Toolchains for ML projects
# Details: https://github.com/google-ml-infra/rules_ml_toolchain
http_archive(
name = "rules_ml_toolchain",
sha256 = "2a5591ec7543c8b37aead3cb681eb2b93c9616ce94abdf3aedcf391b372d4007",
strip_prefix = "rules_ml_toolchain-b2b08356ac30353c49587b0e8dfe65aabb35e78d",
urls = [
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/b2b08356ac30353c49587b0e8dfe65aabb35e78d.tar.gz",
],
)

load(
"@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
"cc_toolchain_deps",
)

cc_toolchain_deps()

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")

load("@rules_ml_toolchain//gpu/sycl:sycl_configure.bzl", "sycl_configure")
load("@rules_ml_toolchain//gpu/sycl:sycl_init_repository.bzl", "sycl_init_repository")

http_archive(
name = "xla",
patch_args = ["-p1"],
patches = ["//third_party:xla.patch"],
sha256 = "a106290c8a1f522d57feed0be31496c571c2a50545cc92a1cdb32aef2309270b",
strip_prefix = "xla-845061f0e1162559fabf5dc6555b85a31bd96cb9",
sha256 = "82160211319100b8c1d55e016c426b2999ccb9c3091f699ac55a2d536c784630",
strip_prefix = "xla-68b8314049f2a7256aea628a7e3377a00278345a",
urls = [
"https://github.com/openxla/xla/archive/845061f0e1162559fabf5dc6555b85a31bd96cb9.zip",
"https://github.com/openxla/xla/archive/68b8314049f2a7256aea628a7e3377a00278345a.zip",
],
)

Expand Down Expand Up @@ -105,28 +130,6 @@ load("@xla//:workspace0.bzl", "xla_workspace0")

xla_workspace0()

# Toolchains for ML projects
# Details: https://github.com/google-ml-infra/rules_ml_toolchain
http_archive(
name = "rules_ml_toolchain",
sha256 = "d1a64a54b1688446619364dac25ff5bcef65c6ffb6984f82128986f5f66129f6",
strip_prefix = "rules_ml_toolchain-b42dc53b80d7f4da1e12abca7503a264e96de98e",
urls = [
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/b42dc53b80d7f4da1e12abca7503a264e96de98e.tar.gz",
],
)

load(
"@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
"cc_toolchain_deps",
)

cc_toolchain_deps()

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")

load(
"@xla//third_party/py:python_wheel.bzl",
"python_wheel_version_suffix_repository",
Expand Down
6 changes: 3 additions & 3 deletions xprof/utils/hlo_proto_to_module.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,10 @@ absl::StatusOr<std::unique_ptr<xla::HloModule>> ConvertHloProtoToModule(
const xla::HloModuleProto& module_proto = hlo_proto.hlo_module();
TF_ASSIGN_OR_RETURN(auto config, xla::HloModule::CreateModuleConfigFromProto(
module_proto, xla::DebugOptions()));
TF_ASSIGN_OR_RETURN(xla::HloModuleProto remapped_module_proto,
xla::HloModule::RemapInstructionIds(module_proto));
TF_ASSIGN_OR_RETURN(auto module, xla::HloModule::CreateFromProto(
remapped_module_proto, config));
module_proto, config,
/*buffer_assignment_proto=*/nullptr,
/*preserve_instruction_ids=*/false));
return module;
}

Expand Down
98 changes: 97 additions & 1 deletion xprof/utils/hlo_proto_to_module_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ TEST(HloProtoToModuleTest, FixNonConsecutiveInstructionIds) {
}
}
id: 4294967303
operand_ids: 6
operand_ids: 1
}
id: 1
root_id: 4294967303
Expand Down Expand Up @@ -88,6 +88,102 @@ TEST(HloProtoToModuleTest, FixNonConsecutiveInstructionIds) {
ElementsAre(Property(&xla::HloInstruction::local_id, 0),
Property(&xla::HloInstruction::local_id, 1),
Property(&xla::HloInstruction::local_id, 2)));
// Check correct operand translation
EXPECT_EQ(module->entry_computation()->parameter_instruction(0)->name(),
"arg0.1");
EXPECT_EQ(module->entry_computation()->parameter_instruction(0)->local_id(),
0);
EXPECT_THAT(
module->entry_computation()->root_instruction()->operands(),
ElementsAre(module->entry_computation()->parameter_instruction(0)));
}

TEST(HloProtoToModuleTest, FixNonConsecutiveInstructionIdsForModule) {
xla::HloProto hlo_proto;
ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
R"pb(
hlo_module {
name: "some_module"
entry_computation_name: "some_module"
computations {
name: "some_module"
instructions {
name: "arg0.1"
opcode: "parameter"
shape {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
id: 4294967297
}
instructions {
name: "arg1.1"
opcode: "parameter"
shape {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
parameter_number: 1
id: 4294967298
}
instructions {
name: "XLA_Retvals.1"
opcode: "tuple"
shape {
element_type: TUPLE
tuple_shapes {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
}
id: 4294967303
operand_ids: 1
}
id: 1
root_id: 4294967303
}
host_program_shape {
parameters {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
parameters {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
result {
element_type: TUPLE
tuple_shapes {
element_type: S32
layout { tail_padding_alignment_in_elements: 1 }
}
}
parameter_names: "arg0"
parameter_names: "arg1"
}
id: 1
entry_computation_id: 1
}
)pb",
&hlo_proto));


ASSERT_OK_AND_ASSIGN(auto module,
ConvertHloProtoToModule(hlo_proto));
EXPECT_EQ(module->entry_computation()->instruction_count(), 3);
// Check that ids are consecutive
EXPECT_THAT(module->entry_computation()->instructions(),
ElementsAre(Property(&xla::HloInstruction::local_id, 0),
Property(&xla::HloInstruction::local_id, 1),
Property(&xla::HloInstruction::local_id, 2)));
// Check correct operand translation
EXPECT_EQ(module->entry_computation()->parameter_instruction(0)->name(),
"arg0.1");
EXPECT_EQ(module->entry_computation()->parameter_instruction(0)->local_id(),
0);
EXPECT_THAT(
module->entry_computation()->root_instruction()->operands(),
ElementsAre(module->entry_computation()->parameter_instruction(0)));
}

} // namespace
Expand Down