[XLA:GPU][Coalescing] Fix a typo: operands index was passed instead o…

…f root index in thread_id map computation for reduce.

PiperOrigin-RevId: 609017178

third_party/xla/xla/service/gpu/fusions/reduction.cc

@@ -1453,7 +1453,7 @@ std::optional<IndexingMap> ReductionFusion::ComputeThreadIdToInputIndexing(
     mlir::MLIRContext* ctx) const {
   const auto& groups = reduction_codegen_info_.GetIndexGroups();
 
-  auto* hero = analysis_.fusion_heroes()[hero_operand_index];
+  auto* hero = analysis_.fusion_heroes()[root_index];
   if (groups.is_reduction_root[root_index] &&
       hero_operand_index >= hero->operand_count() / 2) {
     // We don't have indexing for the init values.

third_party/xla/xla/service/gpu/model/coalescing_analysis_test.cc

@@ -278,7 +278,7 @@ TEST_F(CoalescingTest, ColumnReduction) {
   EXPECT_THAT(IsReadCoalescedPerOperand(ir), ElementsAre(true));
 }
 
-TEST_F(CoalescingTest, VariadicReduce) {
+TEST_F(CoalescingTest, VariadicReduceViaLoopEmitter) {
   absl::string_view ir = R"(
     HloModule module
     max {
@@ -310,6 +310,39 @@ TEST_F(CoalescingTest, VariadicReduce) {
               ElementsAre(true, true, true, true));
 }
 
+TEST_F(CoalescingTest, VariadicReduceViaReductionEmitter) {
+  absl::string_view ir = R"(
+    HloModule module
+    max {
+      p0 = s32[] parameter(0)
+      p1 = s32[] parameter(1)
+      p2 = s32[] parameter(2)
+      p3 = s32[] parameter(3)
+      max01 = s32[] maximum(p0, p1)
+      max23 = s32[] maximum(p2, p3)
+      ROOT max = (s32[], s32[]) tuple(max01, max23)
+    }
+    fusion {
+      p0 = s32[32,40] parameter(0)
+      p1 = s32[32,40] parameter(1)
+      p2 = s32[] parameter(2)
+      p3 = s32[] parameter(3)
+      ROOT reduce = (s32[32], s32[32])
+        reduce(s32[32,40] p0, s32[32,40] p1, s32[] p2, s32[] p3),
+        dimensions={1}, to_apply=max
+    }
+    ENTRY entry {
+      p0 = s32[32,40] parameter(0)
+      p1 = s32[32,40] parameter(1)
+      p2 = s32[] parameter(2)
+      p3 = s32[] parameter(3)
+      ROOT f = (s32[32], s32[32]) fusion(p0, p1, p2, p3),
+          kind=kInput, calls=fusion
+    })";
+  EXPECT_THAT(IsReadCoalescedPerOperand(ir),
+              ElementsAre(true, true, true, true));
+}
+
 TEST_F(CoalescingTest, UnusedParameter) {
   Shape shape = ShapeUtil::MakeShape(F32, {100000});