Add option to remove singleton opt-barriers only.

xla/hlo/transforms/expanders/BUILD

@@ -40,7 +40,8 @@ cc_library(
     srcs = ["optimization_barrier_expander.cc"],
     hdrs = ["optimization_barrier_expander.h"],
     deps = [
-        ":op_expander_pass",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
         "//xla/tsl/platform:status_macros",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -49,6 +50,23 @@ cc_library(
     ],
 )
 
+xla_cc_test(
+    name = "optimization_barrier_expander_test",
+    srcs = ["optimization_barrier_expander_test.cc"],
+    deps = [
+        ":optimization_barrier_expander",
+        "//xla:shape_util",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/hlo/testlib:test",
+        "//xla/hlo/testlib:verified_hlo_module",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "comparison_expander",
     srcs = ["comparison_expander.cc"],

xla/hlo/transforms/expanders/optimization_barrier_expander.cc

@@ -23,6 +23,10 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/tsl/platform/status_macros.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 
 namespace xla {
 
@@ -34,10 +38,16 @@ absl::StatusOr<bool> OptimizationBarrierExpander::RunImpl(
        module->MakeNonfusionComputations(execution_threads)) {
     bool modified = false;
     for (HloInstruction* inst : computation->instructions()) {
-      if (inst->opcode() == HloOpcode::kOptimizationBarrier) {
-        barriers.push_back(inst);
-        modified = true;
+      if (inst->opcode() != HloOpcode::kOptimizationBarrier) {
+        continue;
       }
+      if (only_remove_singleton_opt_barriers_ && inst->operand_count() == 1 &&
+          inst->operand(0)->opcode() == HloOpcode::kTuple &&
+          inst->operand(0)->operand_count() > 1) {
+        continue;
+      }
+      barriers.push_back(inst);
+      modified = true;
     }
 
     if (modified && module->has_schedule()) {

xla/hlo/transforms/expanders/optimization_barrier_expander.h

@@ -19,21 +19,28 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "xla/hlo/transforms/expanders/op_expander_pass.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 
 namespace xla {
 
 // This pass removes the opt-barrier operation which is functionally a no-op.
 class OptimizationBarrierExpander : public HloModulePass {
  public:
   OptimizationBarrierExpander() = default;
+  explicit OptimizationBarrierExpander(bool only_remove_singleton_opt_barriers)
+      : only_remove_singleton_opt_barriers_(
+            only_remove_singleton_opt_barriers) {}
 
-  absl::string_view name() const override { return "cse_barrier_expander"; }
+  absl::string_view name() const override { return "opt-barrier-expander"; }
 
  protected:
   absl::StatusOr<bool> RunImpl(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
+
+ private:
+  bool only_remove_singleton_opt_barriers_ = false;
 };
 
 }  // namespace xla

xla/hlo/transforms/expanders/optimization_barrier_expander_test.cc

@@ -0,0 +1,100 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/hlo/transforms/expanders/optimization_barrier_expander.h"
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/log/log.h"
+#include "absl/types/span.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/hlo/testlib/test.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/shape.h"
+
+namespace xla {
+namespace {
+
+class HloBarrierInstruction : public HloInstruction {
+ public:
+  HloBarrierInstruction(const Shape& shape,
+                        absl::Span<HloInstruction* const> operands)
+      : HloInstruction(HloOpcode::kOptimizationBarrier, shape) {
+    for (HloInstruction* operand : operands) {
+      AppendOperand(operand);
+    }
+  }
+};
+
+class OptimizationBarrierExpanderTest : public HloHardwareIndependentTestBase {
+};
+
+TEST_F(OptimizationBarrierExpanderTest, RemovesOptimizationBarrier) {
+  const char* hlo = R"(
+HloModule module
+
+ENTRY main {
+  param0 = f32[10] parameter(0)
+  add0 = f32[10] add(param0, param0)
+  add1 = f32[10] add(param0, add0)
+  tuple  = (f32[10], f32[10]) tuple(add0, add1)
+  barrier = (f32[10], f32[10]) opt-barrier(tuple)
+  gte = f32[10] get-tuple-element(barrier), index=0
+  ROOT root = f32[10] add(gte, param0)
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo));
+
+  OptimizationBarrierExpander expander;
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_TRUE(
+      FindInstructions(module.get(), HloOpcode::kOptimizationBarrier).empty());
+  VLOG(1) << module->ToString();
+}
+
+TEST_F(OptimizationBarrierExpanderTest, RemovesOnlySingularOptBarrier) {
+  const char* hlo = R"(
+HloModule module
+
+ENTRY main {
+  param0 = f32[10] parameter(0)
+  param1 = f32[10] parameter(1)
+  add0 = f32[10] add(param0, param1)
+  barrier = f32[10] opt-barrier(add0)
+  ROOT add1 = f32[10] add(barrier, param0)
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo));
+
+  OptimizationBarrierExpander expander(
+      /*only_remove_singleton_opt_barriers=*/true);
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_TRUE(
+      FindInstructions(module.get(), HloOpcode::kOptimizationBarrier).empty());
+  VLOG(1) << module->ToString();
+}
+
+}  // namespace
+}  // namespace xla