Move host-compute-offloading wrapping after SPMD sharding to shard ca…

…lls executed on host PiperOrigin-RevId: 635998816
tensorflow · May 22, 2024 · 6dd6123 · 6dd6123
1 parent 4bc3f5a
commit 6dd6123
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -2283,9 +2283,11 @@ Status SpmdPartitioningVisitor::HandleCall(HloInstruction* hlo) {
                                                 call_graph_)
                          .status());
   SetPartitionedHlo(hlo, [&] {
-    return b_.AddInstruction(HloInstruction::CreateCall(
+    auto* call = b_.AddInstruction(HloInstruction::CreateCall(
         MakePartitionedShape(hlo->shape(), hlo->sharding()), call_args,
         hlo->called_computations()[0]));
+    call->set_raw_backend_config_string(hlo->raw_backend_config_string());
+    return call;
   });
   return OkStatus();
 }