Replace the use of xla::OkStatus with absl::OkStatus now that they're…

… the same. PiperOrigin-RevId: 635927531
tensorflow · May 21, 2024 · b46b3fd · b46b3fd
1 parent 714fa2c
commit b46b3fd
Show file tree

Hide file tree

Showing 19 changed files with 85 additions and 85 deletions.
diff --git a/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc b/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc
@@ -136,7 +136,7 @@ absl::Status GetFrontendAttributes(HloCollectivePermuteInstruction* cp,
       cp->frontend_attributes().map().find(kSendRecvValidationAttr);
   if (validation_it == cp->frontend_attributes().map().end() ||
       validation_it->second == "invalid") {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   auto statusor_bounds = ParseReplicaGroupsOnly(validation_it->second);
@@ -172,7 +172,7 @@ absl::Status GetFrontendAttributes(HloCollectivePermuteInstruction* cp,
   std::string cp2_validation_str = bounds_to_string(cp2_bounds);
   (*cp1_attr.mutable_map())[kSendRecvValidationAttr] = cp1_validation_str;
   (*cp2_attr.mutable_map())[kSendRecvValidationAttr] = cp2_validation_str;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Decomposes a CollectivePermute instruction with a cycle in its source-target
@@ -238,7 +238,7 @@ absl::Status DecomposeCollectivePermuteCycle(
   TF_RETURN_IF_ERROR(cp->ReplaceAllUsesWith(recv_data));
   TF_RETURN_IF_ERROR(computation->RemoveInstructionAndUnusedOperands(cp));
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 }  // namespace
 

diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc b/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
@@ -1610,7 +1610,7 @@ absl::Status RestoreFwdGraph(
   }
   TF_RETURN_IF_ERROR(
       comp->ReplaceInstruction(activation_gte, cloned_activation));
-  return OkStatus();
+  return absl::OkStatus();
 }
 }  // namespace
 

diff --git a/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc b/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc
@@ -41,7 +41,7 @@ class SparseDotRewriterImpl : public DfsHloRewriteVisitor {
     // Only handle sparse dots with a single RHS sparse descriptor.
     HloDotInstruction* dot = Cast<HloDotInstruction>(instr);
     if (dot->sparse_operands() != 1 || dot->sparsity().front().index() != 1) {
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     HloInstruction* lhs = dot->mutable_operand(0);

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion.cc b/third_party/xla/xla/service/gpu/gemm_fusion.cc
@@ -739,7 +739,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
               *Cast<HloDotInstruction>(dot),
               std::get<se::CudaComputeCapability>(gpu_version_)) &&
           !should_fuse) {
-        return OkStatus();
+        return absl::OkStatus();
       }
     }
 

diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
@@ -722,7 +722,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
                                m::Broadcast(m::Op(&d_scale)).WithOneUser()))) {
       return F8ScaleD(instr, existing_gemm, d_scale);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleAdd(HloInstruction *instr) override {

diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h
@@ -64,7 +64,7 @@ inline BufferAssigner::Colorer CollectiveColorer() {
         value->set_color(0);
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   };
 }
 

diff --git a/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc b/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc
@@ -91,7 +91,7 @@ absl::Status PostprocessP2PImpl(
       instr->frontend_attributes().map().find(kSendRecvValidationAttr);
   if (validation_it == instr->frontend_attributes().map().end() ||
       validation_it->second == "invalid") {
-    return OkStatus();
+    return absl::OkStatus();
   }
   auto statusor_bounds = ParseReplicaGroupsOnly(validation_it->second);
   if (!statusor_bounds.ok()) {
@@ -101,7 +101,7 @@ absl::Status PostprocessP2PImpl(
   xla::FrontendAttributes attributes = instr->frontend_attributes();
   (*attributes.mutable_map())[kSendRecvValidationAttr] = validation_attr;
   instr->set_frontend_attributes(attributes);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Modifies the loop iteration frontend attribute for the peeled off Send and

diff --git a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc
@@ -217,7 +217,7 @@ absl::Status ProcessWindowedEinsumLoopForActivationCaching(
   }
   TF_RETURN_IF_ERROR(root->ReplaceOperandWith(full_cache_buffer_index,
                                               new_full_buffer_output));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 class WindowedEinsumVisitor : public DfsHloRewriteVisitor {

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -256,7 +256,7 @@ absl::Status IrEmitterUnnested::EmitConditional(const HloInstruction* instr) {
   AddThunkToThunkSequence(std::unique_ptr<Thunk>(
       new ConditionalThunk(Thunk::ThunkInfo::WithProfileAnnotation(instr),
                            std::move(config), slice)));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 llvm::Value* IrEmitterUnnested::CreateLoad(llvm::Value* address,
@@ -632,7 +632,7 @@ absl::Status IrEmitterUnnested::EmitConvolutionThunk(
   AddThunkToThunkSequence(std::make_unique<ConvolutionThunk>(
       Thunk::ThunkInfo::WithProfileAnnotation(instr), std::move(config),
       std::move(operand_slices), std::move(result_slices), scratch_slice));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitterUnnested::EmitGemmThunk(

diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis.h b/third_party/xla/xla/service/gpu/triton_fusion_analysis.h
@@ -42,8 +42,8 @@ class TritonFusionAnalysis {
   static absl::StatusOr<TritonFusionAnalysis> Execute(
       const HloComputation& computation, int split_k = 1);
 
-  // Execute the analysis of a produce-consumer fusion. Returns OkStatus, if the
-  // analysis can find a valid tiling for the producer-consumer fusion.
+  // Execute the analysis of a produce-consumer fusion. Returns absl::OkStatus,
+  // if the analysis can find a valid tiling for the producer-consumer fusion.
   // `split_k` indicates whether this operation was converted to the split-K
   // form and tells the analysis how to interpret the batch dimensions.
   static absl::Status ExecuteForProducerConsumer(const HloInstruction& producer,

diff --git a/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc b/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc
@@ -48,7 +48,7 @@ class AllGatherCanonicalizeTest : public HloTestBase {
     HloPassPipeline pipeline("all-gather-cse");
     pipeline.AddPass<CanonicalizeAllGatherForCSE>();
     TF_RETURN_IF_ERROR(pipeline.Run(module).status());
-    return OkStatus();
+    return absl::OkStatus();
   }
 };
 

diff --git a/third_party/xla/xla/service/spmd/custom_call_handler.cc b/third_party/xla/xla/service/spmd/custom_call_handler.cc
@@ -252,7 +252,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK(
       hlo, PartitionedHlo(create_tuple, hlo->shape(), MakePartitioningState())
                .Reshard(hlo->sharding()));
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
@@ -282,7 +282,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
   amount %= full_size;
   if (amount == 0) {
     SetPartitionedHlo(hlo, input);
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // First step: rotate `amount` on padded data. E.g., before
@@ -345,7 +345,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
   HloInstruction* rotated0 = rotate_with_padding(amount);
   if (right_padding == 0) {
     SetPartitionedHlo(hlo, [&] { return rotated0; });
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Second step: perform another rotate from input, with `right_padding` added
@@ -381,7 +381,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
     return b_.AddInstruction(HloInstruction::CreateTernary(
         rotated0->shape(), HloOpcode::kSelect, pred, rotated1, rotated0));
   });
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 std::unique_ptr<HloInstruction> CreateCustomCallSPMDInternal_RotateRight(
@@ -409,7 +409,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
     auto copy = b_.AddInstruction(
         HloInstruction::CreateUnary(input->shape(), HloOpcode::kCopy, input));
     SetPartitionedHlo(hlo, [&] { return copy; });
-    return OkStatus();
+    return absl::OkStatus();
   }
   if (hlo->custom_call_target() == "SPMDShardToFullShape") {
     // This op switches from manual partitioning to auto partitioning.
@@ -420,7 +420,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
     CHECK(ShapeUtil::Compatible(
         copy->shape(), MakePartitionedShape(hlo->shape(), hlo->sharding())));
     SetPartitionedHlo(hlo, [&] { return copy; });
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (hlo->custom_call_target() == kSPMDOpRotateRight) {
@@ -452,7 +452,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
       }
       return instr;
     });
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (hlo->custom_call_target() == "TopK") {
@@ -477,7 +477,7 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
     SetPartitionedHlo(hlo, PartitionedHlo(move_to_device, hlo->shape(),
                                           MakePartitioningState())
                                .Reshard(hlo->sharding()));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   return DefaultAction(hlo);

diff --git a/third_party/xla/xla/service/spmd/dot_handler.cc b/third_party/xla/xla/service/spmd/dot_handler.cc
@@ -4286,7 +4286,7 @@ absl::Status SpmdPartitioningVisitor::HandleDotHelper(
                    num_partitions_, create_sharded_dot, conv_window, module_,
                    hlo, options_, &b_, &windowed_dot_general_loops_, this));
   SetPartitionedHlo(hlo, [&] { return partitioned_dot; });
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -4370,7 +4370,7 @@ absl::Status SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
   auto to_sink = std::move(input_nodes.first);
   auto new_operands = std::move(input_nodes.second);
   if (to_sink.empty()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   auto computation = loop->parent();
   // Replace the old operand with a tuple of the found small operands.
@@ -4452,7 +4452,7 @@ absl::Status SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
       TF_RETURN_IF_ERROR(body->RemoveInstruction(ou));
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Checks a condition holds true for all recursive operands of an hlo.
@@ -4591,7 +4591,7 @@ absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
   // If nothing is found, to_move could contain only original_output, or
   // cleared by the above code.
   if (to_move.size() <= 1) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // If there is a reduce that's dependent of another reduce, then we can't do
@@ -4604,7 +4604,7 @@ absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
       for (const HloInstruction* other_reduce : reduce_outputs) {
         if (reduce != other_reduce &&
             reachability->IsReachable(reduce, other_reduce)) {
-          return OkStatus();
+          return absl::OkStatus();
         }
       }
     }
@@ -4615,7 +4615,7 @@ absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
         return !absl::c_linear_search(reduce_outputs, inst);
       };
       if (!CheckOperandsRecursive(reduce, reduce_outputs_do_not_contain)) {
-        return OkStatus();
+        return absl::OkStatus();
       }
     }
   }
@@ -4949,7 +4949,7 @@ absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
     TF_RETURN_IF_ERROR(
         computation->RemoveInstructionAndUnusedOperands(reduce_outputs[i]));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
@@ -4978,7 +4978,7 @@ absl::Status SpmdPartitioningVisitor::DoCodeMotionForWindowedDotGeneralLoops(
               loop.while_loop, options));
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace spmd

diff --git a/third_party/xla/xla/service/spmd/fft_handler.cc b/third_party/xla/xla/service/spmd/fft_handler.cc
@@ -427,7 +427,7 @@ absl::Status SpmdPartitioningVisitor::HandleFft(HloInstruction* hlo) {
   auto partitioned_fft =
       PartitionedHlo(result, hlo->shape(), partitioned_input.state());
   SetPartitionedHlo(hlo, partitioned_fft);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace spmd

diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
@@ -846,7 +846,7 @@ absl::Status SpmdPartitioningVisitor::HandleGather(HloInstruction* hlo) {
                       gather->gather_slice_sizes(), this));
   SetPartitionedHlo(gather, PartitionedHlo(pgather, gather->shape(),
                                            MakePartitioningState()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1694,7 +1694,7 @@ absl::Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) {
   }
   SetPartitionedHlo(scatter, PartitionedHlo(pscatter, scatter->shape(),
                                             MakePartitioningState()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace spmd