Do not promote broadcast only groups (#4154)

This PR attempts to not promote a loop group when it only consists of
broadcast IDs. For example, consider a fusion as shown below:

```
  auto tv0 = makeContigConcreteTensor({-1, 1});
  fusion.addInput(tv0);
  auto tv1 = makeContigTensor(2);
  fusion.addInput(tv1);

  auto tv2 = set(tv0);
  auto tv3 = add(tv2, tv1);
  fusion.addOutput(tv3);

  for (auto tv : fusion.allTvs()) {
    tv->split(1, 1, false);
    tv->reorder({{0, 1}, {1, 0}});
  }

  for (auto tv : fusion.allTvs()) {
    tv->inlineAt(2);
  }
```
```
Inputs:
  T0_g_float[bS8{1}, iS0{i0}, bS9{1}]
  T1_g_float[iS12{1}, iS2{i4}, iS13{i5}]
Outputs:
  T3_g_float[iS14{1}, iS6{i0}, iS15{i5}] ca_pos( 2 ) produce_pos( 2 )

%kernel {
T2_l_float[bS10{1}, iS4{i0}, bS11{1}] ca_pos( 2 )
   = Set( T0_g_float[bS8{1}, iS0{i0}, bS9{1}], cache_op=Streaming )
T3_g_float[iS14{1}, iS6{i0}, iS15{i5}] ca_pos( 2 ) produce_pos( 2 )
   = T2_l_float[bS10{1}, iS4{i0}, bS11{1}] ca_pos( 2 )
   + T1_g_float[iS12{1}, iS2{i4}, iS13{i5}];

TransformPrinter :
T0_g_float[bS8{1}, iS0{i0}, bS9{1}]
 logical domain : (iS0{i0}, bS1{1})
 contiguity: t n
  Outer split: bS1{1} by factor 1 -> bS8{1}, bS9{1}
 loop domain : (bS8{1}, iS0{i0}, bS9{1})
T2_l_float[bS10{1}, iS4{i0}, bS11{1}] ca_pos( 2 )
 logical domain : (iS4{i0}, bS5{1})
 contiguity: t n
  Outer split: bS5{1} by factor 1 -> bS10{1}, bS11{1}
 loop domain : (bS10{1}, iS4{i0}, bS11{1})
T1_g_float[iS12{1}, iS2{i4}, iS13{i5}]
 logical domain : (iS2{i4}, iS3{i5})
 contiguity: t t
  Outer split: iS3{i5} by factor 1 -> iS12{1}, iS13{i5}
 loop domain : (iS12{1}, iS2{i4}, iS13{i5})
T3_g_float[iS14{1}, iS6{i0}, iS15{i5}] ca_pos( 2 ) produce_pos( 2 )
 logical domain : (iS6{i0}, iS7{i5})
 contiguity: t t
  Outer split: iS7{i5} by factor 1 -> iS14{1}, iS15{i5}
 loop domain : (iS14{1}, iS6{i0}, iS15{i5})
} // %kernel
```

Here, the interesting part is the innermost loop ID of `T2`, `bS11{1}`.
Because `bS5` is promoted to `iS3` (or `iS7`), `bS11` is also promoted
to a non-broadcast ID that is exact mapped with `iS13` and `iS15`.
However, in this case, `bS11` doesn't really need to be promoted. More
specifically, as long as a loop group only consists of broadcast IDs,
the group should not need to be promoted.

Currently, the generated CUDA kernel with `NVFUSER_ENABLE=id_model(all)`
looks like below:

```
__global__ void CUDAGeneratedKernel(Tensor<float, 2, 2> T0, Tensor<float, 2, 2> T1, Tensor<float, 2, 2> T3) {
  #pragma unroll 1
  for(nvfuser_index_t i0 = 0LL; i0 < T0.logical_size[0LL]; ++i0) {
    nvfuser_index_t i1;
    i1 = T1.logical_size[1LL] * i0;
    Array<float, 1LL, 1> T2;
    #pragma unroll 1
    for(nvfuser_index_t i2 = 0LL; i2 < T1.logical_size[1LL]; ++i2) {
      T2[0LL]
         = T0[i0];
    }
    #pragma unroll 1
    for(nvfuser_index_t i3 = 0LL; i3 < T1.logical_size[1LL]; ++i3) {
      nvfuser_index_t i4;
      i4 = i1 + i3;
      T3[i4]
        = T2[0LL]
        + T1[i4];
    }
  }
}
```

The code is not incorrect, but `T2` is redundantly defined over the loop
of `T1.logical_size[1]` because of the promotion of `bS11`. Note that
the allocation of `T2` is not affected because [broadcast IDs are
excluded](https://github.com/NVIDIA/Fuser/blob/main/csrc/device_lower/pass/allocation.cpp#L299)
before promotion.

In this PR, for loop groups that only consist of broadcast IDs,
promotion to non-broadcast is reverted. With the change, the above
fusion results in the kernel below:

```
__global__ void CUDAGeneratedKernel(Tensor<float, 2, 2> T0, Tensor<float, 2, 2> T1, Tensor<float, 2, 2> T3) {
  #pragma unroll 1
  for(nvfuser_index_t i0 = 0LL; i0 < T0.logical_size[0LL]; ++i0) {
    nvfuser_index_t i1;
    i1 = T1.logical_size[1LL] * i0;
    Array<float, 1LL, 1> T2;
    T2[0LL]
       = T0[i0];
    #pragma unroll 1
    for(nvfuser_index_t i2 = 0LL; i2 < T1.logical_size[1LL]; ++i2) {
      nvfuser_index_t i3;
      i3 = i1 + i2;
      T3[i3]
        = T2[0LL]
        + T1[i3];
    }
  }
}
```

Author: naoyam

csrc/id_model/loop_promotion.cpp

@@ -14,6 +14,8 @@
 #include <options.h>
 #include <val_graph_visitor.h>
 
+#include <algorithm>
+
 namespace nvfuser {
 
 std::string toString(const CoveredGroups& covered_groups) {
@@ -106,10 +108,8 @@ bool isEqualToOrSuperSetOf(
       covered_groups_y.begin(),
       covered_groups_y.end(),
       [&](const CoveredGroup& covered_group_y) {
-        return std::any_of(
-            covered_groups_x.begin(),
-            covered_groups_x.end(),
-            [&](const CoveredGroup& covered_group_x) {
+        return std::ranges::any_of(
+            covered_groups_x, [&](const CoveredGroup& covered_group_x) {
               return covered_group_x.isEqualToOrSuperSetOf(covered_group_y);
             });
       });
@@ -167,10 +167,9 @@ bool isDependencyOf(
     return true;
   }
 
-  if (std::any_of(
-          of->begin(), of->end(), [&](const CoveredGroup& covered_group) {
-            return isDependencyOf(dependency, covered_group);
-          })) {
+  if (std::ranges::any_of(*of, [&](const CoveredGroup& covered_group) {
+        return isDependencyOf(dependency, covered_group);
+      })) {
     return true;
   }
 
@@ -380,6 +379,29 @@ std::unordered_map<ValGroup, IterDomain*> LoopPromotionMapBuilder::build() {
     return buildWithNoBroadcast();
   }
 
+  // Keep track of IDs whose loop groups only have broadcast
+  // IDs. These IDs should not need to be promoted to non-broadcastg
+  // IDs. Note that we can't just remember these loop ValGroups as
+  // they might be replaced during the following analysis.
+  for (const auto& loop_group :
+       idGraph(IdMappingMode::LOOP).disjointValSets().disjointSets()) {
+    if (std::ranges::any_of(*loop_group, [](Val* val) {
+          return !val->as<IterDomain>()->isBroadcast();
+        })) {
+      continue;
+    }
+
+    // Currently, only exact-mapped loop groups are considered. This
+    // condition is required as we are going to replace promotion IDs
+    // with an arbitrary member ID.
+    if (idGraph(IdMappingMode::EXACT).toGroups(*loop_group).size() != 1) {
+      continue;
+    }
+
+    broadcast_only_loop_group_ids_.insert(
+        loop_group->begin(), loop_group->end());
+  }
+
   // Make an intersection of the exact and loop map. This will group together
   // entries in each loop group that are exact with each other. This provides a
   // better graph to do promotion and replays.
@@ -470,6 +492,7 @@ std::unordered_map<ValGroup, IterDomain*> LoopPromotionMapBuilder::build() {
   if (loop_promotion_map_to_propagate.empty()) {
     auto final_loop_promotion_map = updateValGroupIdMap(
         initial_loop_promotion_map, idGraph(IdMappingMode::LOOP));
+    revertBroadcastOnlyLoopGroups(final_loop_promotion_map);
     sanityCheckLoopPromotionMap(final_loop_promotion_map);
     return final_loop_promotion_map;
   }
@@ -537,6 +560,7 @@ std::unordered_map<ValGroup, IterDomain*> LoopPromotionMapBuilder::build() {
   final_loop_promotion_map = updateValGroupIdMap(
       final_loop_promotion_map, idGraph(IdMappingMode::LOOP));
 
+  revertBroadcastOnlyLoopGroups(final_loop_promotion_map);
   sanityCheckLoopPromotionMap(final_loop_promotion_map);
 
   if (callback_) {
@@ -692,10 +716,8 @@ Expr* findMatchingExpr(
   // iel_graph, it means the domain is just replayed and by definition
   // has no mapping with any existing domain, which means there's no
   // matching expr.
-  if (std::any_of(
-          maybe_promoted_inputs.begin(),
-          maybe_promoted_inputs.end(),
-          [&](IterDomain* maybe_promoted_input) -> bool {
+  if (std::ranges::any_of(
+          maybe_promoted_inputs, [&](IterDomain* maybe_promoted_input) -> bool {
             return !iel_graph.hasGroup(maybe_promoted_input);
           })) {
     return nullptr;
@@ -972,20 +994,14 @@ LoopPromotionMapBuilder::computeCoveredGroups(
 
     // Initialize broadcast groups to empty since broadcast domains
     // don't matter for indexing
-    if (std::any_of(id_group->begin(), id_group->end(), [&](Val* id) {
+    if (std::ranges::any_of(*id_group, [&](Val* id) {
           return id->as<IterDomain>()->isBroadcast();
         })) {
       covered_group_map[id_group] = std::make_shared<CoveredGroups>();
     }
   }
 
   ValGraphStmtSort exact_stmt_sort(exact_graph, input_groups_of_graph);
-#if 0
-  std::cerr << "Sorted exprs:\n";
-  for (const ExprGroup& exact_expr : exact_stmt_sort.exprs()) {
-    std::cerr << exact_expr->front()->toString();
-  }
-#endif
   for (const ExprGroup& exact_expr : exact_stmt_sort.exprs()) {
     const std::vector<ValGroup> input_groups =
         exact_graph.inputGroups(exact_expr);
@@ -1180,12 +1196,9 @@ VectorOfUniqueEntries<IterDomain*> LoopPromotionMapBuilder::
       // then it's a terminal ID
       bool all_outs_in_loop_group = true;
       for (auto use : uses_it->second) {
-        if (std::any_of(
-                use->outputs().begin(),
-                use->outputs().end(),
-                [&](Val* out) -> bool {
-                  return group != idGraph(IdMappingMode::LOOP).toGroup(out);
-                })) {
+        if (std::ranges::any_of(use->outputs(), [&](Val* out) -> bool {
+              return group != idGraph(IdMappingMode::LOOP).toGroup(out);
+            })) {
           all_outs_in_loop_group = false;
           break;
         }
@@ -1293,4 +1306,45 @@ std::unordered_map<ValGroup, IterDomain*> LoopPromotionMapBuilder::
   return map;
 }
 
+void LoopPromotionMapBuilder::revertBroadcastOnlyLoopGroups(
+    std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map) const {
+  // If a loop group originally only consisted of broadcast IDs
+  // and now is promoted to a non-broadcast ID, it should not need to
+  // be promoted.
+  for (auto& [loop_group, current_promotion_id] : loop_promotion_map) {
+    if (current_promotion_id->isBroadcast()) {
+      continue;
+    }
+
+    // As long as there's a single ID marked as broadcast only, this
+    // group originally consisted of broadcast IDs only. Note that,
+    // since new IDs were added as part of the promotion analysis, not
+    // all of the IDs are included in the broadcast only set.
+    IterDomain* original_broadcast_id = nullptr;
+    for (auto val : *loop_group) {
+      if (broadcast_only_loop_group_ids_.contains(val)) {
+        original_broadcast_id = val->as<IterDomain>();
+        break;
+      }
+    }
+    if (original_broadcast_id == nullptr) {
+      continue;
+    }
+
+    // Note that this promotion should be valid for the existing
+    // IDs that originate from the fusion, but the loop group also
+    // contains other non-broadcast IDs for loop promotion, e.g.,
+    // current_promotion_id. This replacement means those
+    // non-broadcast IDs are also promoted to the broadcast ID, which
+    // does not make sense. For example, in the case of
+    // IdModelTest.BroadcastOnlyNoLoopPromotion, the innermost loop ID
+    // of tv2 has no mapping in the original fusion, but its loop
+    // group gets additional IDs, iS17 and iS19, both of which are
+    // exact mapped with the innermost loop IDs of tv1 and tv3.
+    //
+    // TODO: Consider cleaning up the unused non-broadcast IDs.
+    current_promotion_id = original_broadcast_id;
+  }
+}
+
 } // namespace nvfuser

csrc/id_model/loop_promotion.h

@@ -293,10 +293,16 @@ class LoopPromotionMapBuilder {
       const std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map)
       const;
 
+  // Revert unnecessary promotions to non-broadcast IDs
+  void revertBroadcastOnlyLoopGroups(
+      std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map) const;
+
  private:
   IdModel& id_model_;
   const StatefulInliningInfo& inlining_info_;
   LoopPromotionMapBuilderCallback* callback_ = nullptr;
+  // Keep track of IDs of broadcast only loop groups
+  std::unordered_set<Val*> broadcast_only_loop_group_ids_;
 
   // (For debugging only) When force_full_loop_promotion_analysis_ is
   // true, it always performs the full loop promotion analysis even

tests/cpp/test_id_model.cpp

@@ -3084,4 +3084,49 @@ TEST_F(IdModelTest, InvalidLoopPromotion) {
   testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
 }
 
+// When a loop group only includes broadcast IDs, the group should not
+// need to be promoted
+TEST_F(IdModelTest, BroadcastOnlyNoLoopPromotion) {
+  auto fusion_ptr = std::make_unique<Fusion>();
+  auto& fusion = *fusion_ptr;
+  FusionGuard fg(fusion_ptr.get());
+
+  auto tv0 = makeContigConcreteTensor({-1, 1});
+  fusion.addInput(tv0);
+  auto tv1 = makeContigTensor(2);
+  fusion.addInput(tv1);
+
+  auto tv2 = set(tv0);
+  auto tv3 = add(tv2, tv1);
+  fusion.addOutput(tv3);
+
+  for (auto tv : fusion.allTvs()) {
+    tv->split(1, 1, false);
+    tv->reorder({{0, 1}, {1, 0}});
+  }
+
+  for (auto tv : fusion.allTvs()) {
+    tv->inlineAt(2);
+  }
+
+  // T2_l_float[bS10{1}, iS4{i0}, bS11{1}] ca_pos( 2 )
+  // = Set( T0_g_float[bS8{1}, iS0{i0}, bS9{1}], cache_op=Streaming )
+  // T3_g_float[iS14{1}, iS6{i0}, iS15{i5}] ca_pos( 2 ) produce_pos( 2 )
+  // = T2_l_float[bS10{1}, iS4{i0}, bS11{1}] ca_pos( 2 )
+  // + T1_g_float[iS12{1}, iS2{i4}, iS13{i5}];
+
+  // In this fusion, the innermost loop ID of tv2 is broadcast and is
+  // not inlined. While its producer ID is promoted to the concrete
+  // logical ID of tv3, it should not need to promote the loop ID as
+  // it's just a broadcast.
+
+  IdModel id_model(&fusion, /*build_graphs=*/true);
+
+  auto promotion_id = id_model.loopPromotionMap().at(
+      id_model.idGraph(IdMappingMode::LOOP).toGroup(tv2->axis(-1)));
+  EXPECT_TRUE(promotion_id->isBroadcast())
+      << "Should not be promoted a non-broadcast ID: "
+      << promotion_id->toString();
+}
+
 } // namespace nvfuser