[XLA] Remove limitation that probably was just copy-pasted in collect…

…ive-pipeliner. PiperOrigin-RevId: 609171910
tensorflow · Feb 22, 2024 · ad16e54 · ad16e54
1 parent 4b08a92
commit ad16e54
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/third_party/xla/xla/service/collective_pipeliner.cc b/third_party/xla/xla/service/collective_pipeliner.cc
@@ -517,7 +517,7 @@ std::optional<std::vector<HloInstruction*>> CollectChainsToPushBackwards(
     HloInstruction* instr, int64_t loop_iter, const HloComputation* while_body,
     int64_t level_to_operate_on,
     const absl::flat_hash_set<const HloInstruction*>& loop_invariant_params) {
-  if (instr->user_count() != 1 || instr->HasControlDependencies()) {
+  if (instr->HasControlDependencies()) {
     return std::nullopt;
   }
   return CollectIndependentOperandChain(instr, loop_iter,