More spmd fixes (#1843)

* More spmd fixes

* fix

Author: wsmoses

patches/xla_spmd2.patch

@@ -0,0 +1,22 @@
+diff --git a/xla/service/spmd/spmd_partitioner.cc b/xla/service/spmd/spmd_partitioner.cc
+index 54fdd9a0fe..08c5f900a5 100644
+--- a/xla/service/spmd/spmd_partitioner.cc
++++ b/xla/service/spmd/spmd_partitioner.cc
+@@ -570,6 +570,17 @@ PartitionedHlo PartitionedHlo::ReshardNoCache(
+     return PartitionedHlo(partitioned, base_shape_, state_);
+   }
+ 
++  if (state_.module->config().debug_options().xla_enable_enzyme_comms_opt()) {
++    if (hlo_->opcode() == HloOpcode::kBroadcast && hlo_->operand(0)->shape().dimensions().size() == 0 &&
++       hlo_->operand(0)->IsConstant()) {
++      HloInstruction* new_broadcast =
++          state_.b->AddInstruction(HloInstruction::CreateBroadcast(
++              hlo_->shape(), hlo_->mutable_operand(0), {}));
++      new_broadcast->set_sharding(target);
++      return PartitionedHlo(new_broadcast, base_shape_, state_);
++    }
++  }
++
+   if (CanReshardWithCollectivePermute(sharding(), target)) {
+     return ReshardWithCollectivePermute(target);
+   }

third_party/xla/workspace.bzl

@@ -17,6 +17,6 @@ def repo(extra_patches = [], override_commit = ""):
         strip_prefix = "openxla-xla-{commit}".format(commit = commit[:7]),
         urls = ["https://api.github.com/repos/openxla/xla/tarball/{commit}".format(commit = commit)],
         patch_cmds = XLA_PATCHES + extra_patches,
-        patches = ["//:patches/xla.patch", "//:patches/xla_spmd.patch"],
+        patches = ["//:patches/xla.patch", "//:patches/xla_spmd.patch", "//:patches/xla_spmd2.patch"],
         patch_args = ["-p1"],
     )