HloFusionAnalysis: Don't peek outside the fusion boundary.

Currently, we use the FindNonTrivialHero overload that doesn't take a boundary
function. This isn't always correct.

PiperOrigin-RevId: 572576780

Author: jreiffers

xla/service/gpu/BUILD

@@ -3300,6 +3300,20 @@ cc_library(
     ],
 )
 
+xla_cc_test(
+    name = "hlo_fusion_analysis_test",
+    srcs = ["hlo_fusion_analysis_test.cc"],
+    deps = [
+        ":backend_configs_cc",
+        ":gpu_device_info_for_tests",
+        ":hlo_fusion_analysis",
+        ":hlo_traversal",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:xla_internal_test_main",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
 cc_library(
     name = "gpu_performance_model",
     srcs = ["gpu_performance_model.cc"],

xla/service/gpu/hlo_fusion_analysis.cc

@@ -262,7 +262,7 @@ StatusOr<HloFusionAnalysis> HloFusionAnalysis::Create(
   std::vector<const HloInstruction*> heroes;
   heroes.reserve(hlo_roots.size());
   for (auto* root : hlo_roots) {
-    heroes.push_back(&FindNonTrivialHero(*root));
+    heroes.push_back(&FindNonTrivialHero(*root, boundary_fn));
   }
 
   std::vector<const HloInstruction*> fusion_arguments;
@@ -306,28 +306,26 @@ HloFusionAnalysis::EmitterFusionKind HloFusionAnalysis::GetEmitterFusionKind()
     return EmitterFusionKind::kTriton;
   }
 #endif
-  const auto& roots = fusion_roots();
-
-  if (absl::c_any_of(roots, [](const HloInstruction* root) {
-        return IsRealReductionHero(*root, FindNonTrivialHero(*root));
-      })) {
-    return EmitterFusionKind::kReduction;
+  for (auto [root, hero] : llvm::zip(fusion_roots_, fusion_heroes_)) {
+    if (IsRealReductionHero(*root, *hero)) {
+      return EmitterFusionKind::kReduction;
+    }
   }
 
   // We expect that the last dimension is swapped with a different dimension.
   if (HasConsistentTransposeHeros() && tiled_transpose_->permutation[2] != 2) {
     return EmitterFusionKind::kTranspose;
   }
 
-  if (roots.size() > 1) {
-    if (IsInputFusibleNonStridedSlices(roots) &&
-        AllSliceInputsAreCompatible(roots)) {
+  if (fusion_roots_.size() > 1) {
+    if (IsInputFusibleNonStridedSlices(fusion_roots_) &&
+        AllSliceInputsAreCompatible(fusion_roots_)) {
       return EmitterFusionKind::kInputSlices;
     }
     return EmitterFusionKind::kLoop;
   }
 
-  if (roots[0]->opcode() == HloOpcode::kScatter) {
+  if (fusion_roots_[0]->opcode() == HloOpcode::kScatter) {
     return EmitterFusionKind::kScatter;
   }
 

xla/service/gpu/hlo_fusion_analysis_test.cc

@@ -0,0 +1,65 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "xla/service/gpu/hlo_fusion_analysis.h"
+
+#include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/gpu_device_info_for_tests.h"
+#include "xla/service/gpu/hlo_traversal.h"
+#include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla::gpu {
+namespace {
+
+class HloFusionAnalysisTest : public HloTestBase {};
+
+TEST_F(HloFusionAnalysisTest, DoesNotPeekOutsideBoundary) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+    HloModule test_module
+
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+
+    ENTRY main {
+      %p0 = f32[1024] parameter(0)
+      %p1 = f32[] parameter(1)
+      %reduce = f32[] reduce(%p0, %p1), dimensions={0}, to_apply=add
+      ROOT %bitcast = s32[] bitcast(%reduce)
+    })")
+                    .value();
+
+  auto device_info = TestGpuDeviceInfo::RTXA6000DeviceInfo();
+
+  auto* root = module->entry_computation()->root_instruction();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto analysis, HloFusionAnalysis::Create(
+                         FusionBackendConfig::default_instance(), {root},
+                         MakeSingleInstructionFusion(*root), &device_info));
+  EXPECT_EQ(analysis.GetEmitterFusionKind(),
+            HloFusionAnalysis::EmitterFusionKind::kLoop);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto analysis_fused,
+      HloFusionAnalysis::Create(FusionBackendConfig::default_instance(), {root},
+                                DefaultFusionBoundaryFn, &device_info));
+  EXPECT_EQ(analysis_fused.GetEmitterFusionKind(),
+            HloFusionAnalysis::EmitterFusionKind::kReduction);
+}
+
+}  // namespace
+}  // namespace xla::gpu

xla/service/gpu/hlo_traversal.cc

@@ -66,6 +66,13 @@ FusionBoundaryFn MakeProducerConsumerFusion(
   };
 }
 
+FusionBoundaryFn MakeSingleInstructionFusion(const HloInstruction& root) {
+  if (root.opcode() == HloOpcode::kFusion) {
+    return DefaultFusionBoundaryFn;
+  }
+  return [](const HloInstruction&, const HloInstruction&) { return true; };
+}
+
 void HloBfsConsumersFirstTraversal(
     absl::Span<const HloInstruction* const> roots,
     const FusionBoundaryFn& boundary,

xla/service/gpu/hlo_traversal.h

@@ -46,6 +46,12 @@ bool DefaultFusionBoundaryFn(const HloInstruction& producer,
 FusionBoundaryFn MakeProducerConsumerFusion(
     const HloInstruction& fused_producer, const HloInstruction& fused_consumer);
 
+// Creates a fusion boundary function for a fusion consisting only of `root`. If
+// `root` is a fusion, the result is the same as `DefaultFusionBuondaryFn`. If
+// `root` is the root of a fusion, the result is just that root, not the entire
+// computation.
+FusionBoundaryFn MakeSingleInstructionFusion(const HloInstruction& root);
+
 // Visit the HLO nodes starting from `roots` in BFS order (consumers before
 // producers). Each node will be visited exactly once. The graph is not
 // traversed along edges for which `boundary` returns true.

xla/service/gpu/hlo_traversal_test.cc

@@ -311,6 +311,38 @@ TEST_F(HloTraversalTest, FuseFusionConsumerAndProducer) {
   EXPECT_THAT(params, ElementsAre("negate", "p0"));
 }
 
+TEST_F(HloTraversalTest, SingleInstructionFusionOfFusion) {
+  auto module = ParseAndReturnVerifiedModule(kTwoFusions).value();
+  auto* fusion =
+      module->entry_computation()->GetInstructionWithName("fusion.1");
+
+  auto boundary = MakeSingleInstructionFusion(*fusion);
+  std::vector<std::string> nodes;
+  HloBfsConsumersFirstTraversal({fusion}, boundary,
+                                [&](const HloInstruction& node) {
+                                  nodes.emplace_back(node.name());
+                                  return TraversalResult::kVisitOperands;
+                                });
+
+  EXPECT_THAT(nodes,
+              ElementsAre("fusion.1", "reduce.1", "mul", "p0.1", "p1.1"));
+}
+
+TEST_F(HloTraversalTest, SingleInstructionFusionOfInstruction) {
+  auto module = ParseAndReturnVerifiedModule(kTwoFusions).value();
+  auto* negate = module->entry_computation()->GetInstructionWithName("negate");
+
+  auto boundary = MakeSingleInstructionFusion(*negate);
+  std::vector<std::string> nodes;
+  HloBfsConsumersFirstTraversal({negate}, boundary,
+                                [&](const HloInstruction& node) {
+                                  nodes.emplace_back(node.name());
+                                  return TraversalResult::kVisitOperands;
+                                });
+
+  EXPECT_THAT(nodes, ElementsAre("negate"));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla