Add a utility for finding HLO use chains

This is needed for proper epilogue indexing computations, e.g.
in fusions like this:

```
          reduce
         /      \
   broadcast    log
        |        |
       neg    bitcast
         \      /
           ROOT
```

The current assumption in `ComputeEpilogueInputToOutputIndexing`
that we can just take the first user is incorrect here - the
reduce is both part of the side output's computation and the hero
of the fusion. Fusions like this make absolutely no sense, but they
exist.

PiperOrigin-RevId: 629994503

Author: jreiffers

third_party/xla/xla/service/gpu/BUILD

@@ -5732,6 +5732,7 @@ xla_cc_test(
         ":hlo_traversal",
         "//xla/hlo/ir:hlo",
         "//xla/tests:hlo_test_base",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
     ],
 )

third_party/xla/xla/service/gpu/hlo_traversal.cc

@@ -14,13 +14,15 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/hlo_traversal.h"
 
+#include <algorithm>
 #include <functional>
 #include <iterator>
 #include <memory>
 #include <optional>
 #include <queue>
 #include <sstream>
 #include <string>
+#include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
@@ -482,5 +484,29 @@ std::optional<const HloInstruction*> HloFindIf(
   return std::nullopt;
 }
 
+std::vector<HloInstructionAdaptor> HloFindUseChain(HloInstructionAdaptor parent,
+                                                   HloInstructionAdaptor root) {
+  absl::flat_hash_set<HloInstructionAdaptor> visited;
+  std::vector<HloInstructionAdaptor> result;
+  std::function<bool(HloInstructionAdaptor)> visit;
+  visit = [&](HloInstructionAdaptor node) {
+    if (node == root) return true;
+    for (const auto& user : node.GetUsers()) {
+      if (visited.insert(user).second && visit(user)) {
+        result.push_back(user);
+        return true;
+      }
+    }
+    return false;
+  };
+  if (visit(parent)) {
+    result.push_back(parent);
+    std::reverse(result.begin(), result.end());
+  } else {
+    result.clear();
+  }
+  return result;
+}
+
 }  // namespace gpu
 }  // namespace xla

third_party/xla/xla/service/gpu/hlo_traversal.h

@@ -186,6 +186,11 @@ void FindFusionArguments(
     const HloFusionAdaptor& fusion,
     const std::function<void(HloInstructionAdaptor producer)>& visit);
 
+// Find a use chain from `parent` to `root`. Empty if no chain exists.
+// `[parent]` if `parent` is `root`.
+std::vector<HloInstructionAdaptor> HloFindUseChain(HloInstructionAdaptor parent,
+                                                   HloInstructionAdaptor root);
+
 }  // namespace gpu
 }  // namespace xla
 

third_party/xla/xla/service/gpu/hlo_traversal_test.cc

@@ -21,6 +21,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/tests/hlo_test_base.h"
@@ -30,6 +31,7 @@ namespace gpu {
 namespace {
 
 using ::testing::ElementsAre;
+using ::testing::IsEmpty;
 
 MATCHER_P(InstructionAdaptorName, name, "") { return arg.name() == name; }
 
@@ -288,7 +290,7 @@ const char kTwoFusions[] = R"(
       sum = f32[128] add(p1, p1)
       negate = f32[128] negate(sum)
       fusion.1 = f32[] fusion(p0, negate), kind=kLoop, calls=fused_computation_1
-      fusion.2 = f32[] fusion(fusion.1, negate), kind=kLoop, calls=fused_computation_2 
+      fusion.2 = f32[] fusion(fusion.1, negate), kind=kLoop, calls=fused_computation_2
       ROOT difference = f32[] subtract(fusion.2, p0)
     })";
 
@@ -497,6 +499,49 @@ TEST_F(HloTraversalTest, MakeInstructionsPostOrder_TwoMultiOutputFusions) {
                                  InstructionAdaptorName("reduce.2")));
 }
 
+TEST_F(HloTraversalTest, HloFindUseChain) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+    fusion {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      negate = f32[] negate(p0)
+      log = f32[] log(p0)
+      sum = f32[] add(p0, log)
+      exp = f32[] exponential(p1)
+      ROOT call = f32[] custom-call(negate, exp, sum), custom_call_target="it"
+    }
+
+    ENTRY main {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT fusion = f32[] fusion(p0, p1), kind=kLoop, calls=fusion
+    }
+    )")
+                    .value();
+
+  auto* fusion_computation = module->GetComputationWithName("fusion");
+  auto fusion = HloFusionAdaptor::ForComputation(fusion_computation);
+  auto get = [&](absl::string_view name) {
+    return HloInstructionAdaptor{
+        *fusion_computation->GetInstructionWithName(name), fusion.get()};
+  };
+  auto p0 = get("p0");
+  auto p1 = get("p1");
+  auto log = get("log");
+  auto sum = get("sum");
+  auto negate = get("negate");
+  auto exp = get("exp");
+  auto call = get("call");
+
+  EXPECT_THAT(HloFindUseChain(p0, p0), ElementsAre(p0));
+  EXPECT_THAT(HloFindUseChain(p0, p1), IsEmpty());
+  EXPECT_THAT(HloFindUseChain(p0, call), ElementsAre(p0, negate, call));
+  EXPECT_THAT(HloFindUseChain(p0, sum), ElementsAre(p0, log, sum));
+  EXPECT_THAT(HloFindUseChain(p1, exp), ElementsAre(p1, exp));
+  EXPECT_THAT(HloFindUseChain(negate, exp), IsEmpty());
+  EXPECT_THAT(HloFindUseChain(call, p0), IsEmpty());
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla