PR #23681: Add an API to get a computation's caller(s).

Imported from GitHub PR #23681

This should replace CallGraph in most cases, and adds an alternative to the deprecated .*CallInstruction functions.
Copybara import of the project:

--
4168abb by Johannes Reifferscheid <jreiffers@nvidia.com>:

Add an API to get a computation's caller(s).

This should replace CallGraph in most cases, and adds an alternative to
the deprecated .*CallInstruction functions.

Merging this change closes #23681

FUTURE_COPYBARA_INTEGRATE_REVIEW=#23681 from jreiffers:caller-instructions 4168abb
PiperOrigin-RevId: 736448470

Author: jreiffers

xla/hlo/ir/hlo_computation.cc

@@ -204,6 +204,9 @@ HloComputation::~HloComputation() {
   }
   CHECK(caller_computations_.empty());
 
+  // Delete the map from caller instructions to count, if it exists.
+  delete GetCallersMap();
+
   for (const auto& i : instructions_) {
     delete i.inst();
   }
@@ -277,8 +280,6 @@ static void IncrementCount(
   ++map[key];
 }
 
-// Returns true if the callee was present and its count was decremented; returns
-// false if the callee was not present.
 static void DecrementCount(
     absl::btree_map<HloComputation*, int, HloComputation::UniqueIdComparator>&
         map,
@@ -292,17 +293,75 @@ static void DecrementCount(
   }
 }
 
-void HloComputation::AddCallee(HloComputation* callee) {
+void HloComputation::AddCallee(const HloInstruction* caller,
+                               HloComputation* callee) {
   IncrementCount(callee_computations_, callee);
   IncrementCount(callee->caller_computations_, this);
+
+  if (auto* map = callee->GetCallersMap()) {
+    ++(*map)[caller];
+  } else if (callee->callers_ == 0) {
+    callee->callers_ = reinterpret_cast<uintptr_t>(caller);
+  } else {
+    // Convert the single instruction to a map.
+    auto* current_caller = reinterpret_cast<const HloInstruction*>(
+        callee->callers_ & ~kCallerTypeMask);
+    auto* map = new absl::flat_hash_map<const HloInstruction*, int>();
+    (*map)[current_caller] = 1;
+    ++(*map)[caller];
+    callee->callers_ = reinterpret_cast<uintptr_t>(map) |
+                       static_cast<uintptr_t>(CallersType::kCallerCountHashMap);
+  }
+
   if (parent() != nullptr && callee->parent() == parent()) {
     parent()->topological_sort_.AddEdge(this, callee);
   }
 }
 
-void HloComputation::RemoveCallee(HloComputation* callee) {
+void HloComputation::RemoveCallee(const HloInstruction* caller,
+                                  HloComputation* callee) {
+  CHECK(caller);
+  CHECK(callee);
   DecrementCount(callee_computations_, callee);
   DecrementCount(callee->caller_computations_, this);
+
+  if (callee->callers_ == reinterpret_cast<uintptr_t>(caller)) {
+    // The callee had just this single caller, so we reset it to 0 (no caller).
+    callee->callers_ = 0;
+  } else {
+    auto* map = callee->GetCallersMap();
+    CHECK(map) << "Attempted to remove a caller " << caller->name()
+               << " that did not call the computation " << name() << "."
+               << callee->callers_;
+    auto it = map->find(caller);
+    CHECK(it != map->end())
+        << "Attempted to remove a caller " << caller->name()
+        << " that did not call the computation " << name() << ".";
+    --it->second;
+    // We don't convert back to the inline representation, since this case
+    // should be rare.
+  }
+}
+
+absl::flat_hash_map<const HloInstruction*, int>*
+HloComputation::GetCallersMap() {
+  if (static_cast<CallersType>(callers_ & kCallerTypeMask) ==
+      CallersType::kCallerCountHashMap) {
+    return reinterpret_cast<absl::flat_hash_map<const HloInstruction*, int>*>(
+        callers_ & ~kCallerTypeMask);
+  }
+  return nullptr;
+}
+
+absl::flat_hash_map<const HloInstruction*, int>* const
+HloComputation::GetCallersMap() const {
+  if (static_cast<CallersType>(callers_ & kCallerTypeMask) ==
+      CallersType::kCallerCountHashMap) {
+    return reinterpret_cast<
+        absl::flat_hash_map<const HloInstruction*, int>* const>(
+        callers_ & ~kCallerTypeMask);
+  }
+  return nullptr;
 }
 
 HloInstruction* HloComputation::AddInstructionInternal(
@@ -329,7 +388,7 @@ HloInstruction* HloComputation::AddInstructionInternal(
     CHECK(parent() == nullptr || called_computation->parent() == parent())
         << "Called computation " << called_computation->name()
         << " is not in the same module as " << name();
-    AddCallee(called_computation);
+    AddCallee(pinst, called_computation);
   }
   return pinst;
 }

xla/hlo/ir/hlo_computation.h

@@ -982,11 +982,34 @@ class HloComputation {
     return caller_computations_;
   }
 
+  // The returned callers are in no particular order.
+  absl::InlinedVector<const HloInstruction*, 1> caller_instructions() const {
+    if (const auto* map = GetCallersMap()) {
+      absl::InlinedVector<const HloInstruction*, 1> result;
+      for (const auto& [instr, _] : *map) {
+        result.push_back(instr);
+      }
+      return result;
+    }
+
+    if (callers_ == 0) {
+      return {};
+    }
+    return {
+        reinterpret_cast<const HloInstruction*>(callers_ & ~kCallerTypeMask)};
+  }
+
   void ClearCalledComputations();
 
  private:
   friend class HloModule;
 
+  enum class CallersType : uint8_t {
+    kHloInstruction = 0,
+    kCallerCountHashMap = 1,
+  };
+  static constexpr uintptr_t kCallerTypeMask = 0b1;
+
   explicit HloComputation(
       const std::string& name, int parameter_count,
       std::vector<std::unique_ptr<HloInstruction>>* instructions,
@@ -1051,8 +1074,14 @@ class HloComputation {
   void SetUniqueIdHelper(int64_t id);
 
   friend class HloInstruction;
-  void AddCallee(HloComputation* callee);
-  void RemoveCallee(HloComputation* callee);
+  // Add/remove call from `caller`, which must be in this computation, to
+  // `callee`.
+  void AddCallee(const HloInstruction* caller, HloComputation* callee);
+  void RemoveCallee(const HloInstruction* caller, HloComputation* callee);
+
+  // Returns nullptr if `callers_` is not a map.
+  absl::flat_hash_map<const HloInstruction*, int>* GetCallersMap();
+  absl::flat_hash_map<const HloInstruction*, int>* const GetCallersMap() const;
 
   // Unique ID of this computation.
   // This is set to -1 if the computation is not in a module. Should only be
@@ -1067,6 +1096,11 @@ class HloComputation {
   // The respective type in the least significant three bits.
   uintptr_t instruction_and_type_ = 0;
 
+  // Contains an HloInstruction* or an absl::flat_hash_map<HloInstruction*,
+  // /*count=*/int> in the high bits and a CallersType in the least significant
+  // bit.
+  uintptr_t callers_ = 0;
+
   // If this computation is an async computation, this field points to the
   // first async instruction (async-start) in the asynchronous op chain that
   // calls this computation.

xla/hlo/ir/hlo_instruction.cc

@@ -220,7 +220,7 @@ void HloInstruction::AppendComputation(HloComputation* computation) {
   // of T and hlo_instruction.h does not include hlo_computation.h.
   mutable_rare()->called_computations.push_back(computation);
   if (parent()) {
-    parent()->AddCallee(computation);
+    parent()->AddCallee(this, computation);
   }
 }
 
@@ -235,10 +235,10 @@ void HloInstruction::set_called_computation(int index,
   std::swap(old_computation, mutable_rare()->called_computations[index]);
   if (parent()) {
     if (old_computation) {
-      parent()->RemoveCallee(old_computation);
+      parent()->RemoveCallee(this, old_computation);
     }
     if (computation) {
-      parent()->AddCallee(computation);
+      parent()->AddCallee(this, computation);
     }
   }
 }
@@ -255,7 +255,7 @@ void HloInstruction::ClearCalledComputations() {
     if (parent()) {
       for (HloComputation* computation : called_computations()) {
         if (computation) {
-          parent()->RemoveCallee(computation);
+          parent()->RemoveCallee(this, computation);
         }
       }
     }

xla/hlo/ir/hlo_module_test.cc

@@ -43,6 +43,8 @@ namespace xla {
 namespace {
 
 using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+using ::testing::UnorderedElementsAre;
 
 TEST(HloModuleTest, AbslHashValue) {
   HloModule module1("temp_module", HloModuleConfig());
@@ -477,41 +479,95 @@ TEST(HloModuleTest, CheckToStringHonorsDebugOptions) {
 }
 
 TEST(HloModuleTest, TestCallersAndCallees) {
-  // Check that the debug options xla_dump_large_constants,
-  // xla_syntax_sugar_async_ops are honored.
   const char* hlo = R"(
     HloModule jit_h
 
     f {
-      Arg_0.3 = f32[] parameter(0)
-      ROOT sine.4 = f32[] sine(Arg_0.3)
+      p0 = f32[] parameter(0)
+      ROOT sine.4 = f32[] sine(p0)
     }
 
     g {
-      Arg_0.13 = f32[] parameter(0)
-      call.14 = f32[] call(Arg_0.13), to_apply=f
-      ROOT call.15 = f32[] call(call.14), to_apply=f
+      p0 = f32[] parameter(0)
+      call.f.0 = f32[] call(p0), to_apply=f
+      ROOT call.f.1 = f32[] call(call.f.0), to_apply=f
+    }
+
+    h {
+      ROOT p0 = f32[] parameter(0)
+    }
+
+    uncalled {
+      p0 = f32[] parameter(0)
+      ROOT call.h = f32[] call(p0), to_apply=h
     }
 
     ENTRY main {
       Arg_0.1 = f32[] parameter(0)
-      call.5 = f32[] call(Arg_0.1), to_apply=f
-      call.16 = f32[] call(call.5), to_apply=g
-      ROOT call.27 = f32[] call(call.16), to_apply=g
+      call.f.2 = f32[] call(Arg_0.1), to_apply=f
+      call.g.0 = f32[] call(call.f.2), to_apply=g
+      ROOT call.g.1 = f32[] call(call.g.0), to_apply=g
     })";
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(hlo));
-  EXPECT_EQ(module->computation_count(), 3);
+  EXPECT_EQ(module->computation_count(), 5);
   HloComputation* main = module->GetComputationWithName("main");
   HloComputation* f = module->GetComputationWithName("f");
   HloComputation* g = module->GetComputationWithName("g");
+  HloComputation* h = module->GetComputationWithName("h");
+  HloComputation* uncalled = module->GetComputationWithName("uncalled");
   EXPECT_THAT(main->callee_computations(),
               ElementsAre(std::make_pair(f, 1), std::make_pair(g, 2)));
   EXPECT_THAT(f->callee_computations(), ElementsAre());
   EXPECT_THAT(g->callee_computations(), ElementsAre(std::make_pair(f, 2)));
   EXPECT_THAT(f->caller_computations(),
               ElementsAre(std::make_pair(g, 2), std::make_pair(main, 1)));
   EXPECT_THAT(g->caller_computations(), ElementsAre(std::make_pair(main, 2)));
+
+  HloInstruction* call_f_0 = g->GetInstructionWithName("call.f.0");
+  HloInstruction* call_f_1 = g->GetInstructionWithName("call.f.1");
+  HloInstruction* call_f_2 = main->GetInstructionWithName("call.f.2");
+  HloInstruction* call_g_0 = main->GetInstructionWithName("call.g.0");
+  HloInstruction* call_g_1 = main->GetInstructionWithName("call.g.1");
+  HloInstruction* call_h = uncalled->GetInstructionWithName("call.h");
+
+  EXPECT_THAT(f->caller_instructions(),
+              UnorderedElementsAre(call_f_0, call_f_1, call_f_2));
+  EXPECT_THAT(g->caller_instructions(),
+              UnorderedElementsAre(call_g_0, call_g_1));
+  EXPECT_THAT(h->caller_instructions(), ElementsAre(call_h));
+  EXPECT_THAT(uncalled->caller_instructions(), IsEmpty());
+}
+
+TEST(HloModuleTest, MultipleCallsFromOneInstruction) {
+  const char* hlo = R"(
+    f {
+      tparam = f32[4] parameter(0)
+      ROOT tuple = (f32[4]) tuple(tparam)
+    }
+
+    g {
+      fparam = f32[4] parameter(0)
+      ROOT tuple = (f32[4]) tuple(fparam)
+    }
+
+    ENTRY main {
+      p0 = f32[4] parameter(0)
+      b0 = s32[] parameter(1)
+      ROOT conditional = (f32[4]) conditional(b0, p0, p0, p0),
+        branch_computations={f, f, g}
+    })";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnUnverifiedModule(hlo));
+  EXPECT_EQ(module->computation_count(), 3);
+  HloComputation* main = module->GetComputationWithName("main");
+  HloComputation* f = module->GetComputationWithName("f");
+  HloComputation* g = module->GetComputationWithName("g");
+
+  HloInstruction* conditional = main->GetInstructionWithName("conditional");
+
+  EXPECT_THAT(f->caller_instructions(), ElementsAre(conditional));
+  EXPECT_THAT(g->caller_instructions(), ElementsAre(conditional));
 }
 
 }  // namespace