Remove more code paths that can create >1 module module groups.

hawkinsp · Google-ML-Automation · commit 0b573a004707 · 2025-10-04T06:05:49.000-07:00
Make HloModuleGroup::push_back() private.

PiperOrigin-RevId: 814460024
diff --git a/xla/hlo/ir/hlo_module_group.h b/xla/hlo/ir/hlo_module_group.h
@@ -34,18 +34,9 @@ namespace xla {
 // concurrently across different devices.
 class HloModuleGroup {
  public:
-  // Construct an empty module group.
-  explicit HloModuleGroup(absl::string_view name) : name_(name) {}
-
   // Construct a module group containing a single module.
   explicit HloModuleGroup(std::unique_ptr<HloModule> module);
 
-  // Construct a module group containing any number of modules.
-  HloModuleGroup(absl::string_view name,
-                 absl::Span<std::unique_ptr<HloModule>> modules);
-  HloModuleGroup(absl::string_view name,
-                 std::vector<std::unique_ptr<HloModule>>&& modules);
-
   HloModuleGroup(const HloModuleGroup& other) = delete;
   HloModuleGroup(HloModuleGroup&& other) = default;
   HloModuleGroup& operator=(const HloModuleGroup& other) = delete;
@@ -57,9 +48,6 @@ class HloModuleGroup {
   // Returns a module at a particular index.
   HloModule& module(int index) const { return *module_ptrs_.at(index); }
 
-  // Add a module to the back of vector of modules in the group.
-  void push_back(std::unique_ptr<HloModule> module);
-
   // Replaces the existing module at the given index with the given module. The
   // existing module is discarded.
   void ReplaceModule(int index, std::unique_ptr<HloModule> module);
@@ -105,6 +93,19 @@ class HloModuleGroup {
   }
 
  private:
+  // Construct an empty module group.
+  explicit HloModuleGroup(absl::string_view name) : name_(name) {}
+
+  // Construct a module group containing any number of modules.
+  HloModuleGroup(absl::string_view name,
+                 absl::Span<std::unique_ptr<HloModule>> modules);
+  HloModuleGroup(absl::string_view name,
+                 std::vector<std::unique_ptr<HloModule>>&& modules);
+
+  // Add a module to the back of vector of modules in the group. Private
+  // because we no longer want to support > 1 module per group.
+  void push_back(std::unique_ptr<HloModule> module);
+
   std::string name_;
 
   // Vector of modules as std::unique_ptrs.
diff --git a/xla/hlo/pass/hlo_pass_fix_test.cc b/xla/hlo/pass/hlo_pass_fix_test.cc
@@ -182,85 +182,5 @@ TEST_F(HloPassFixTest, RunModuleToNonDefaultEarlyExit) {
   EXPECT_EQ(root->literal().GetFirstElement<int32_t>(), 20);
 }
 
-TEST_F(HloPassFixTest, RunModuleGroupToFixedPoint) {
-  constexpr absl::string_view kModule0 = R"(
-    HloModule First
-
-    ENTRY main {
-      ROOT c = s32[] constant(5)
-    }
-  )";
-
-  constexpr absl::string_view kModule1 = R"(
-    HloModule Second
-
-    ENTRY main {
-      ROOT c = s32[] constant(3)
-    }
-  )";
-
-  constexpr absl::string_view kModule2 = R"(
-    HloModule Second
-
-    ENTRY main {
-      ROOT c = s32[] constant(0)
-    }
-  )";
-
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module0,
-                          ParseAndReturnVerifiedModule(kModule0));
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module1,
-                          ParseAndReturnVerifiedModule(kModule1));
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module2,
-                          ParseAndReturnVerifiedModule(kModule2));
-  HloModuleGroup module_group("group");
-  module_group.push_back(std::move(module0));
-  module_group.push_back(std::move(module1));
-  module_group.push_back(std::move(module2));
-
-  HloPassFix<DecrementPositiveConstants> pass;
-  TF_ASSERT_OK_AND_ASSIGN(bool changed, pass.RunOnModuleGroup(&module_group));
-  EXPECT_TRUE(changed);
-  HloInstruction* root0 =
-      module_group.module(0).entry_computation()->root_instruction();
-  ASSERT_EQ(root0->opcode(), HloOpcode::kConstant);
-  EXPECT_EQ(root0->literal().GetFirstElement<int32_t>(), 0);
-  HloInstruction* root1 =
-      module_group.module(1).entry_computation()->root_instruction();
-  ASSERT_EQ(root1->opcode(), HloOpcode::kConstant);
-  EXPECT_EQ(root1->literal().GetFirstElement<int32_t>(), 0);
-  HloInstruction* root2 =
-      module_group.module(2).entry_computation()->root_instruction();
-  ASSERT_EQ(root2->opcode(), HloOpcode::kConstant);
-  EXPECT_EQ(root2->literal().GetFirstElement<int32_t>(), 0);
-}
-
-TEST_F(HloPassFixTest, OscillationsStillTerminate) {
-  constexpr absl::string_view kModule = R"(
-    HloModule Oscillating
-
-    ENTRY main {
-      a = f32[4] parameter(0)
-      b = f32[4] parameter(1)
-      ROOT c = f32[4] add(a, b)
-    }
-  )";
-
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(kModule));
-  HloPassFix<FlipAddSubtract> pass;
-
-  // We expect this to terminate and report that the module did not change.
-  TF_ASSERT_OK_AND_ASSIGN(bool changed, pass.Run(module.get()));
-  EXPECT_FALSE(changed);
-
-  // But don't lie when crash_on_hlo_pass_silent_hlo_change is set.
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_unsupported_crash_on_hlo_pass_silent_hlo_change(true);
-  TF_ASSERT_OK_AND_ASSIGN(changed, pass.Run(module.get()));
-  EXPECT_TRUE(changed);
-}
-
 }  // namespace
 }  // namespace xla
diff --git a/xla/hlo/pass/hlo_pass_pipeline_test.cc b/xla/hlo/pass/hlo_pass_pipeline_test.cc
@@ -27,7 +27,6 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
@@ -50,15 +49,10 @@ using ::testing::StrEq;
 
 class HloPassPipelineTest : public HloHardwareIndependentTestBase {
  protected:
-  absl::StatusOr<HloModuleGroup> ParseModuleGroup(
-      absl::Span<const std::string> hlo_strings) {
-    HloModuleGroup group(TestName());
-    for (const std::string& hlo_string : hlo_strings) {
-      TF_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_string));
-      group.push_back(std::move(module));
-    }
-    return group;
+  absl::StatusOr<HloModuleGroup> ParseModuleGroup(std::string hlo_string) {
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> module,
+                        ParseAndReturnVerifiedModule(hlo_string));
+    return HloModuleGroup(std::move(module));
   }
 };
 
@@ -276,36 +270,22 @@ ENTRY main {
   ROOT baz = f32[] multiply(a, b)
 }
 )";
-  const std::string module_1_str = R"(
-HloModule MixedPipeline.0
-
-ENTRY main {
-  a = f32[] parameter(0)
-  b = f32[] parameter(1)
-  ROOT foo = f32[] multiply(a, b)
-}
-)";
-
   TF_ASSERT_OK_AND_ASSIGN(HloModuleGroup module_group,
-                          ParseModuleGroup({module_0_str, module_1_str}));
+                          ParseModuleGroup(module_0_str));
 
   HloPassPipeline pipeline(TestName());
   pipeline.AddPass<BazToQuxModuleGroupPass>();
   pipeline.AddPass<FooToBarModulePass>();
 
   HloInstruction* root0 =
       module_group.module(0).entry_computation()->root_instruction();
-  HloInstruction* root1 =
-      module_group.module(1).entry_computation()->root_instruction();
   EXPECT_EQ(root0->name(), "baz");
-  EXPECT_EQ(root1->name(), "foo");
 
   TF_ASSERT_OK_AND_ASSIGN(bool changed,
                           pipeline.RunOnModuleGroup(&module_group));
   EXPECT_TRUE(changed);
 
   EXPECT_EQ(root0->name(), "qux");
-  EXPECT_EQ(root1->name(), "bar");
 }
 
 TEST_F(HloPassPipelineTest, InvariantChecker) {
@@ -383,15 +363,13 @@ ENTRY main {
 
 // Test that metadata is set when a module group goes through a pass pipeline.
 TEST_F(HloPassPipelineTest, SetHloModuleMetadata) {
-  HloModuleGroup module_group(TestName());
-  module_group.push_back(CreateNewVerifiedModule());
-  module_group.push_back(CreateNewVerifiedModule());
+  HloModuleGroup module_group(CreateNewVerifiedModule());
 
   HloPassPipeline pipeline(TestName());
   pipeline.AddPass<BazToQuxModuleGroupPass>();
   pipeline.AddPass<FooToBarModulePass>();
   TF_ASSERT_OK(pipeline.RunOnModuleGroup(&module_group).status());
-  ASSERT_THAT(module_group.modules(), SizeIs(2));
+  ASSERT_THAT(module_group.modules(), SizeIs(1));
 
   std::vector<std::string> pass_names = {"pipeline-start", "baz2qux",
                                          "foo2bar"};
@@ -410,8 +388,7 @@ TEST_F(HloPassPipelineTest, SetHloModuleMetadata) {
       EXPECT_FALSE(pass_metadata.module_changed());
       EXPECT_EQ(pass_metadata.module_id(), module->unique_id());
       EXPECT_THAT(pass_metadata.module_group_module_ids(),
-                  ElementsAre(module_group.module(0).unique_id(),
-                              module_group.module(1).unique_id()));
+                  ElementsAre(module_group.module(0).unique_id()));
       EXPECT_GT(pass_metadata.start_timestamp_usec(), 0);
       EXPECT_LE(pass_metadata.start_timestamp_usec(),
                 pass_metadata.end_timestamp_usec());
diff --git a/xla/service/BUILD b/xla/service/BUILD
@@ -1838,10 +1838,7 @@ xla_cc_test(
 xla_cc_test(
     name = "hlo_module_group_test",
     srcs = ["hlo_module_group_test.cc"],
-    # TODO(b/148211710) Test fails in OSS.
-    tags = ["no_oss"],
     deps = [
-        ":hlo_module_group_metadata",
         ":hlo_proto_cc",
         "//xla/hlo/ir:hlo_module_group",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
diff --git a/xla/service/hlo_module_group_test.cc b/xla/service/hlo_module_group_test.cc
diff --git a/xla/service/service.cc b/xla/service/service.cc
