Change unique_ptr to shared_ptr for latency_estimator and async_tracker in latency_hiding_scheduler, so that we can potentially reuse them across different users without needing to re-construct new ones.

FUTURE_COPYBARA_INTEGRATE_REVIEW=#25166 from openxla:add_slow_arg_init_alarm faa4fd1
PiperOrigin-RevId: 747709188

Author: Google-ML-Automation

xla/service/BUILD

@@ -1182,7 +1182,6 @@ cc_library(
         "//xla/hlo/analysis:hlo_alias_analysis",
         "//xla/hlo/analysis:hlo_reachability",
         "//xla/hlo/ir:hlo",
-        "//xla/hlo/ir:ptrvec",
         "//xla/hlo/pass:hlo_pass",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
@@ -5858,6 +5857,7 @@ cc_library(
         "//xla/hlo/ir:ptrvec",
         "//xla/hlo/pass:hlo_pass",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",

xla/service/latency_hiding_scheduler.cc

@@ -2926,15 +2926,16 @@ absl::StatusOr<bool> LatencyHidingScheduler::Run(
       saved_schedules[computation] = std::move(new_schedule);
     }
   }
-  LOG(INFO) << "LatencyHidingScheduler current memory usage: "
+  LOG(INFO) << "[" << name() << "]"
+            << " LatencyHidingScheduler current memory usage: "
             << scheduler_core_->GetMemoryPeak()
             << " bytes. Current limit: " << scheduler_core_->GetMemoryLimit();
   for (HloComputation* computation : computations_to_schedule) {
-    VLOG(1) << "Statistics before scheduling:";
+    VLOG(1) << "[" << name() << "] Statistics before scheduling:";
     LogScheduleStatistics(computation);
     module->schedule().set_sequence(
         computation, absl::MakeConstSpan(saved_schedules[computation]));
-    VLOG(1) << "Statistics after scheduling:";
+    VLOG(1) << "[" << name() << "] Statistics after scheduling:";
     LogScheduleStatistics(computation);
   }
   if (debug_options.xla_dump_latency_hiding_schedule()) {

xla/service/latency_hiding_scheduler.h

@@ -43,7 +43,6 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_schedule.h"
-#include "xla/hlo/ir/ptrvec.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/map_util.h"
 #include "xla/service/hlo_buffer.h"
@@ -1224,8 +1223,8 @@ class LatencyHidingScheduler : public HloModulePass {
   };
 
   LatencyHidingScheduler(
-      std::unique_ptr<LatencyEstimator> latency_estimator,
-      std::unique_ptr<AsyncTracker> async_tracker,
+      std::shared_ptr<LatencyEstimator> latency_estimator,
+      std::shared_ptr<AsyncTracker> async_tracker,
       std::unique_ptr<SchedulerCore> scheduler_core,
       const HloCostAnalysis::ShapeSizeFunction& shape_size_bytes)
       : latency_estimator_(std::move(latency_estimator)),
@@ -1252,8 +1251,8 @@ class LatencyHidingScheduler : public HloModulePass {
   virtual void LogScheduleStatistics(const HloComputation* computation);
 
  private:
-  std::unique_ptr<LatencyEstimator> latency_estimator_;
-  std::unique_ptr<AsyncTracker> async_tracker_;
+  std::shared_ptr<LatencyEstimator> latency_estimator_;
+  std::shared_ptr<AsyncTracker> async_tracker_;
   std::unique_ptr<SchedulerCore> scheduler_core_;
   const HloCostAnalysis::ShapeSizeFunction shape_size_bytes_;
   absl::flat_hash_set<HloComputation*> computations_to_schedule_;

xla/service/legalize_scheduling_annotations.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
@@ -48,7 +49,8 @@ namespace {
 // Given a group of annotated instructions (sources), find all reachable
 // instructions from them in the same computation.
 absl::flat_hash_set<HloInstruction*> PropagateAnnotationFromSources(
-    const std::vector<HloInstruction*>& sources, HloComputation* computation) {
+    const std::vector<HloInstruction*>& sources,
+    const HloComputation* computation) {
   absl::flat_hash_set<HloInstruction*> to_annotate;
   auto reachability = HloReachabilityMap::Build(computation);
   // worklist contains instructions that can reach any source instruction.
@@ -107,28 +109,13 @@ absl::Status AttachAnnotation(
           absl::StrCat(annotation_id) + " to " + std::string(instr->name()) +
           " but it has an existing annotation: " + absl::StrCat(*id));
     }
+    LOG(INFO) << "Propagating annotation " << annotation_id << " to "
+              << instr->name();
     SetSchedulingAnnotation(instr, annotation_id);
   }
   return absl::OkStatus();
 }
 
-absl::StatusOr<bool> PropagateAnnotations(
-    HloComputation* computation,
-    const absl::flat_hash_map<int64_t, std::vector<HloInstruction*>>&
-        annotation_id_to_instructions) {
-  bool changed = false;
-  for (auto& [annotation_id, sources] : annotation_id_to_instructions) {
-    absl::flat_hash_set<HloInstruction*> to_annotate =
-        PropagateAnnotationFromSources(sources, computation);
-    changed |= (!to_annotate.empty());
-    auto status = AttachAnnotation(annotation_id, to_annotate);
-    if (!status.ok()) {
-      return status;
-    }
-  }
-  return changed;
-}
-
 absl::StatusOr<int64_t> ExtractAnnotation(
     const ::google::protobuf::Map<std::string, std::string>& attrs,
     absl::string_view instr_name) {
@@ -261,6 +248,23 @@ absl::Status CheckGapBetweenAnnotatedInstructions(
 
 }  // namespace
 
+absl::StatusOr<bool> LegalizeSchedulingAnnotations::PropagateAnnotations(
+    const HloComputation* computation,
+    const absl::btree_map<int64_t, std::vector<HloInstruction*>>&
+        annotation_id_to_instructions) {
+  bool changed = false;
+  for (auto& [annotation_id, sources] : annotation_id_to_instructions) {
+    absl::flat_hash_set<HloInstruction*> to_annotate =
+        PropagateAnnotationFromSources(sources, computation);
+    changed |= (!to_annotate.empty());
+    auto status = AttachAnnotation(annotation_id, to_annotate);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  return changed;
+}
+
 bool LegalizeSchedulingAnnotations::KeepSchedulingAnnotation(
     HloInstruction* instr) {
   return IsSupportedAsyncOp(instr) || config_.keep_sync_annotation(instr);
@@ -348,10 +352,12 @@ absl::StatusOr<bool> LegalizeSchedulingAnnotations::Run(
     return false;
   }
 
-  auto status = CheckStartDoneAnnotationConsistency(annotation_to_instructions,
-                                                    annotation);
-  if (!status.ok()) {
-    return status;
+  if (config_.check_start_done_annotation_consistency) {
+    auto status = CheckStartDoneAnnotationConsistency(
+        annotation_to_instructions, annotation);
+    if (!status.ok()) {
+      return status;
+    }
   }
 
   bool changed = false;
@@ -362,7 +368,7 @@ absl::StatusOr<bool> LegalizeSchedulingAnnotations::Run(
     // same annotation ID.
     for (HloComputation* computation :
          module->MakeNonfusionComputations(execution_threads)) {
-      absl::flat_hash_map<int64_t, std::vector<HloInstruction*>>
+      absl::btree_map<int64_t, std::vector<HloInstruction*>>
           per_computation_annotation_to_instructions;
       for (const auto& [annotation_id, comp_inst_vector] :
            annotation_to_instructions) {

xla/service/legalize_scheduling_annotations.h

@@ -16,11 +16,16 @@ limitations under the License.
 #ifndef XLA_SERVICE_LEGALIZE_SCHEDULING_ANNOTATIONS_H_
 #define XLA_SERVICE_LEGALIZE_SCHEDULING_ANNOTATIONS_H_
 
+#include <cstdint>
 #include <utility>
+#include <vector>
 
+#include "absl/container/btree_map.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/util.h"
@@ -34,13 +39,20 @@ class LegalizeSchedulingAnnotations : public HloModulePass {
   struct Config {
     HloPredicate keep_sync_annotation = HloPredicateTrue;
     bool propagate_annotation = false;
+    bool check_start_done_annotation_consistency = true;
   };
 
   explicit LegalizeSchedulingAnnotations(Config config)
       : config_(std::move(config)) {}
   absl::string_view name() const override {
     return "legalize-scheduling-annotations";
   }
+
+  static absl::StatusOr<bool> PropagateAnnotations(
+      const HloComputation* computation,
+      const absl::btree_map<int64_t, std::vector<HloInstruction*>>&
+          annotation_id_to_instructions);
+
   using HloPassInterface::Run;
   absl::StatusOr<bool> Run(
       HloModule* module,

xla/tools/multihost_hlo_runner/BUILD

@@ -168,6 +168,7 @@ cc_library(
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_module_util",
         "//xla/service:hlo_proto_cc",
+        "//xla/service:slow_operation_alarm",
         "//xla/tests:test_utils",
         "//xla/tools:hlo_control_flow_flattening",
         "//xla/tsl/platform:env",

xla/tools/multihost_hlo_runner/functional_hlo_runner.cc

@@ -65,6 +65,7 @@ limitations under the License.
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_module_util.h"
+#include "xla/service/slow_operation_alarm.h"
 #include "xla/shape_layout.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
@@ -1161,6 +1162,11 @@ FunctionalHloRunner::CreateArgumentsOnDevice(
         client, executable, running_options, flatten_arguments);
   }
 
+  SlowOperationAlarm alarm(
+      absl::Seconds(5),
+      absl::StrFormat("Argument initialization is slow. Consider changing "
+                      "--hlo_argument_mode."));
+
   absl::Span<PjRtDevice* const> addressable_devices =
       executable->addressable_devices();
   size_t num_addressable_devices = addressable_devices.size();
@@ -1185,7 +1191,7 @@ FunctionalHloRunner::CreateArgumentsOnDevice(
           ModuleArgumentMode::kUseZerosAsInput;
 
   for (int i = 0; i < num_addressable_devices; ++i) {
-    VLOG(3) << "Creating fake argument for device " << i;
+    VLOG(3) << "Creating fake arguments for device " << i;
     LiteralVec& argument_literals =
         per_device_argument_literals[addressable_devices[i]->id()];
     int executable_idx = hlo_modules.size() == 1