Repository: NVIDIA/Fuser
Pull request #4325: Cleaning up scheduling of static repeat
Author: login: naoyam
Labels: name: rope
Created: 2025-04-26T19:19:11Z | Merged: 2025-04-30T20:06:43Z

The resize scheduler automatically detects a sequence of ops to repeat a tensor at a certain ID and slightly modifies the scheduling to reduce redundant computations. This PR makes the analysis a little more flexible so that it also works with a pattern appearing in a Llama forward RoPE module.

Specifically, previously, the specific scheduling is only applied when a sequence of `BroadcastOp, ExpandOp and ViewOp` are detected in this order, just because that's how a repetition of a tensor is commonly represented. However, the only op that is absolutely necessary is the final reshape. As long as it meets the patterns for repetition, it should be sufficient to apply the scheduling. In fact, in a segment of a Llama RoPE forward, there's a segment input that has a broadcast ID, which is then expanded inside the segment and merged to realize a repetition. This case is not detectable as the segment lacks a BroadcastOp in the current main but is detected with this PR.


Diff:
diff --git a/csrc/scheduler/resize.cpp b/csrc/scheduler/resize.cpp
index 92a02fb59ac..d449e3267ed 100644
--- a/csrc/scheduler/resize.cpp
+++ b/csrc/scheduler/resize.cpp
@@ -315,6 +315,47 @@ void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
       tvs_with_extra_transforms, ref_tv->getLogicalDomain());
 }
 
+// Partition a given set of tensors to two disjoint sets based on a
+// given iter domain and reachability from the iter domain. Returns two
+// vectors of tensors, first of which contains all tensors that has an
+// iter domain that is reachable from the given iter domain, whereas
+// the rest of tensors are all grouped into the second
+// list. Reachability is determined by using the permissive BFS
+// traversal on a given graph.
+std::pair<std::vector<TensorView*>, std::vector<TensorView*>> partitionTvsById(
+    const std::vector<TensorView*> tvs,
+    IterDomain* id,
+    const ValGraph& graph) {
+  ValGroups target_groups;
+  for (auto tv : tvs) {
+    target_groups.pushBack(graph.toGroups(tv->getLogicalDomain()));
+  }
+
+  const auto reachable_groups = getReachableValsFrom<ValGraphPermissiveBFS>(
+      {graph.toGroup(id)},
+      target_groups.vector(),
+      /*allowed_direction=*/Direction::Undefined,
+      graph);
+  const std::unordered_set<ValGroup> reachable_group_set{
+      reachable_groups.begin(), reachable_groups.end()};
+
+  std::vector<TensorView*> reachable_tvs;
+  std::vector<TensorView*> unreachable_tvs;
+
+  for (auto tv : tvs) {
+    if (std::ranges::any_of(
+            tv->getLogicalDomain(), [&](IterDomain* logical_id) {
+              return reachable_group_set.contains(graph.toGroup(logical_id));
+            })) {
+      reachable_tvs.push_back(tv);
+    } else {
+      unreachable_tvs.push_back(tv);
+    }
+  }
+
+  return std::make_pair(reachable_tvs, unreachable_tvs);
+}
+
 } // namespace
 
 void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
@@ -384,7 +425,7 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   id_model->buildExactGraph();
 
   // Detect an ending repeat
-  auto static_repeat_info = scheduler_tools::getMaybeStaticRepeatInfo(ref_tv);
+  auto repeat_info = scheduler_tools::getMaybeStaticRepeatInfo(ref_tv);
 
   // Just simple scheduling for now.
   // TODO: Do something smarter. Can just use the pointwise scheduler?
@@ -425,35 +466,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
   // scheduling.
   bool repeat_id_moved_to_outermost = false;
-  if (static_repeat_info.has_value()) {
-    NVF_ERROR(ref_tv == static_repeat_info->repeat_output_tv);
-    auto ref_repeat_id_it = std::find_if(
+  if (repeat_info.has_value()) {
+    auto ref_factor_id_it = std::find_if(
         ref_tv->getLoopDomain().begin(),
         ref_tv->getLoopDomain().end(),
         [&](IterDomain* loop_id) {
           return id_model->idGraph(IdMappingMode::EXACT)
               .disjointValSets()
-              .strictAreMapped(loop_id, static_repeat_info->reshape_repeat_id);
+              .strictAreMapped(loop_id, repeat_info->factor_id);
         });
-    // Gives up if the repeat ID is not found. Unclear if this could
-    // actually happen, though.
-    if (ref_repeat_id_it != ref_tv->getLoopDomain().end()) {
-      auto repeat_id_pos =
-          std::distance(ref_tv->getLoopDomain().begin(), ref_repeat_id_it);
+    // The factor ID should be found in the loop domain as the
+    // reshape should be cancelled at this point. Gives up if not.
+    if (ref_factor_id_it != ref_tv->getLoopDomain().end()) {
+      auto factor_id_pos =
+          std::distance(ref_tv->getLoopDomain().begin(), ref_factor_id_it);
       NVF_ERROR(
-          repeat_id_pos >= outermost_pos,
-          "Unexpected to have DID-parallelized repeat axis: ",
-          static_repeat_info->reshape_repeat_id->toString());
+          factor_id_pos >= outermost_pos,
+          "Unexpected to have DID-parallelized repeat factor axis: ",
+          repeat_info->factor_id->toString());
 
-      // [DID, ..., repeat_id, ...]
+      // [DID, ..., repeat_factor_id, ...]
       //        ^
       //        +--- outermost_pos
-      ref_tv->reorder(std::unordered_map<int64_t, int64_t>{{repeat_id_pos, 0}});
+      ref_tv->reorder(std::unordered_map<int64_t, int64_t>{{factor_id_pos, 0}});
       ++outermost_pos;
-      // [repeat_id, DID, ...]
-      //                   ^
-      //                   +--- outermost_pos
-
+      // [repeat_factor_id, DID, ...]
+      //                         ^
+      //                         +--- outermost_pos
       repeat_id_moved_to_outermost = true;
     }
   }
@@ -514,34 +553,25 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // post-repeat group, where only the latter group has the repeat
   // IDs. When propagating the loop domain of the reference tensor,
   // which has the repeat ID, the full loop domain is propagated only
-  // to the post-repeat group. For the pre-repeat group, the repeat ID
-  // is dropped and only the remaining loop domain is propagated.
+  // to the tensors that have IDs that are mapped with the repeat
+  // ID. For the rest of the tensros, the repeat ID is dropped and
+  // only the remaining loop domain is propagated.
   if (repeat_id_moved_to_outermost) {
-    // Divide all tvs to the pre and posgt repeat groups
-    auto all_tvs = fusion->allTvs();
-    std::vector<TensorView*> post_repeat_tvs;
-    post_repeat_tvs.reserve(static_repeat_info->repeat_tvs.size());
-    std::vector<TensorView*> pre_repeat_tvs;
-    pre_repeat_tvs.reserve(
-        all_tvs.size() - static_repeat_info->repeat_tvs.size());
-    for (auto tv : all_tvs) {
-      if (static_repeat_info->repeat_tvs.count(tv)) {
-        post_repeat_tvs.push_back(tv);
-      } else {
-        pre_repeat_tvs.push_back(tv);
-      }
-    }
+    const auto& [tvs_with_repeat_id, tvs_without_repeat_id] = partitionTvsById(
+        fusion->allTvs(),
+        repeat_info->factor_id,
+        id_model->maybeBuildGraph(IdMappingMode::BROADCAST));
 
     // The repeat ID should be located at the outermost position
     std::vector<IterDomain*> non_repeated_loop{
         ref_tv->getLoopDomain().begin() + 1, ref_tv->getLoopDomain().end()};
 
     scheduler_tools::scheduleLoopDomainsLike(
-        pre_repeat_tvs,
+        tvs_without_repeat_id,
         non_repeated_loop,
         /*update_loop_domain_only=*/true);
     scheduler_tools::scheduleLoopDomainsLike(
-        post_repeat_tvs,
+        tvs_with_repeat_id,
         ref_tv->getLoopDomain(),
         /*update_loop_domain_only=*/true);
   } else {
diff --git a/csrc/scheduler/tools/static_repeat.cpp b/csrc/scheduler/tools/static_repeat.cpp
index 331481aac80..a0f53276956 100644
--- a/csrc/scheduler/tools/static_repeat.cpp
+++ b/csrc/scheduler/tools/static_repeat.cpp
@@ -8,126 +8,40 @@
 
 #include <ir/all_nodes.h>
 #include <ir/utils.h>
+#include <logical_domain_map.h>
 #include <scheduler/tools/static_repeat.h>
 
 namespace nvfuser {
 namespace scheduler_tools {
 
 std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
-    TensorView* maybe_repeat_out) {
-  // The pattern to detect:
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // repeat_out = reshape(expand_out)
-  //
-  // Additionally, since maybe_repeat_out is commonly a fusion
-  // output, it is likely there's a cache tv between expand_out and
-  // repeat_out, so the following pattern should also be detected.
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // cache_of_repeat_out = reshape(expand_out)
-  // repeat_out = set(cache_of_repeat_out)
-
-  std::unordered_set<TensorView*> repeat_tvs;
-  repeat_tvs.insert(maybe_repeat_out);
-
-  auto reshape_out = maybe_repeat_out;
-
-  // Check if there's a cache
-  if (auto ldst = dynamic_cast<LoadStoreOp*>(maybe_repeat_out->definition());
-      ldst != nullptr && ldst->opType() == LoadStoreOpType::Set) {
-    reshape_out = ldst->in()->as<TensorView>();
-    repeat_tvs.insert(reshape_out);
+    TensorView* maybe_repeat_out_tv) {
+  // Skip set ops if any (e.g., inserted by caching). Only Set
+  // or SegmenterSet are considered.
+  while (auto ldst =
+             dynamic_cast<LoadStoreOp*>(maybe_repeat_out_tv->definition())) {
+    if (ldst->opType() != LoadStoreOpType::Set &&
+        ldst->opType() != LoadStoreOpType::SegmenterSet) {
+      break;
+    }
+    maybe_repeat_out_tv = ldst->in()->as<TensorView>();
   }
 
   // Detect reshape
-  auto reshape = dynamic_cast<ViewOp*>(reshape_out->definition());
+  auto reshape = dynamic_cast<ViewOp*>(maybe_repeat_out_tv->definition());
   if (reshape == nullptr) {
     return std::nullopt;
   }
 
-  // Detect expand
-  auto expand_out = reshape->in();
-  repeat_tvs.insert(expand_out);
-  auto expand = dynamic_cast<ExpandOp*>(expand_out->definition());
-  if (expand == nullptr) {
-    return std::nullopt;
-  }
-
-  // Detect broadcast
-  auto broadcast_out = expand->in();
-  repeat_tvs.insert(broadcast_out);
-  auto broadcast = dynamic_cast<BroadcastOp*>(broadcast_out->definition());
-  if (broadcast == nullptr) {
-    return std::nullopt;
-  }
-
-  auto inp_tv = broadcast->in();
-
-  // Not sure if this is really necessary to check, but assume there's
-  // only single chain of the ops and tensors from inp_tv to
-  // maybe_reshape_out
-  if (inp_tv->uses().size() > 1 &&
-      std::any_of(repeat_tvs.begin(), repeat_tvs.end(), [](TensorView* tv) {
-        return tv->uses().size() > 1;
-      })) {
-    return std::nullopt;
-  }
-
-  // Check if the ops match with the repeat pattern. Currently only
-  // one iter domain can be repeated
-  IterDomain* broadcast_id = nullptr;
-  int64_t broadcast_pos = -1;
-  for (const auto i : arange(broadcast_out->getLogicalDomain().size())) {
-    if (broadcast->getBroadcastDimFlags().at(i)) {
-      if (broadcast_id != nullptr) {
-        // Multiple broadcast IDs not supported
-        return std::nullopt;
-      }
-      broadcast_id = broadcast_out->getLogicalDomain().at(i);
-      broadcast_pos = (int64_t)i;
-    }
-  }
+  auto reshape_in = reshape->input(0)->as<TensorView>();
+  auto reshape_out = reshape->output(0)->as<TensorView>();
 
-  if (broadcast_id == nullptr) {
-    return std::nullopt;
-  }
-
-  // Check if and only if the broadcast ID is expanded
-  IterDomain* expanded_id = nullptr;
-  for (const auto i : arange(broadcast_out->getLogicalDomain().size())) {
-    auto p_id = broadcast_out->getLogicalDomain().at(i);
-    auto c_id = expand_out->getLogicalDomain().at(i);
-    if (p_id == broadcast_id && c_id->isBroadcast() &&
-        c_id->hasExpandedExtent()) {
-      expanded_id = c_id;
-    } else if (
-        p_id->isBroadcast() && !p_id->hasExpandedExtent() &&
-        c_id->isBroadcast() && c_id->hasExpandedExtent()) {
-      // Expanded but this broadcast was not introduced by the
-      // preceding broadcast op
-      return std::nullopt;
-    }
-  }
-
-  if (expanded_id == nullptr) {
-    return std::nullopt;
-  }
-
-  // Only a static repeat factor is considered
-  if (!expanded_id->expandedExtent()->isConstInt()) {
-    return std::nullopt;
-  }
-
-  // The expanded ID should be merged with the iter domain next to it,
-  // and that should be the only reshape expr
   auto reshape_exprs = DependencyCheck::getAllExprsBetween(
       {reshape_out->getRootDomain().begin(),
        reshape_out->getRootDomain().end()},
       {reshape_out->getLogicalDomain().begin(),
        reshape_out->getLogicalDomain().end()});
+
   if (reshape_exprs.size() != 1) {
     return std::nullopt;
   }
@@ -137,14 +51,6 @@ std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
     return std::nullopt;
   }
 
-  // The corresponding root ID of the outout tv should be one of the
-  // inputs of the merge
-  auto reshape_root_broadcast = reshape_out->getRootDomain().at(broadcast_pos);
-  if (reshape_merge->outer() != reshape_root_broadcast &&
-      reshape_merge->inner() != reshape_root_broadcast) {
-    return std::nullopt;
-  }
-
   // Reshape of an expanded broadcast always generates a concrete
   // non-broadcast ID, so this check is not necessary, but just in
   // case in the future that may change.
@@ -154,10 +60,39 @@ std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
   }
 
   StaticRepeatInfo info;
-  info.repeat_output_tv = maybe_repeat_out;
-  info.reshape_output_tv = reshape_out;
-  info.reshape_repeat_id = reshape_out->getRootDomain().at(broadcast_pos);
-  info.repeat_tvs = repeat_tvs;
+
+  info.output_id = reshape_merge->out();
+
+  const auto c2p =
+      PairwiseLogicalDomainMap(reshape_in, reshape_out).mapConsumerToProducer();
+
+  auto producer_merge_outer = c2p.at(reshape_merge->outer());
+  auto producer_merge_inner = c2p.at(reshape_merge->inner());
+  IterDomain* producer_factor_id = nullptr;
+
+  if (producer_merge_outer->isBroadcast() &&
+      producer_merge_outer->hasExpandedExtent() &&
+      !producer_merge_inner->isBroadcast()) {
+    // Inner ID is repeated by the factor of the outer extent
+    info.input_id = reshape_merge->inner();
+    info.factor_id = reshape_merge->outer();
+    producer_factor_id = producer_merge_outer;
+  } else if (
+      producer_merge_inner->isBroadcast() &&
+      producer_merge_inner->hasExpandedExtent() &&
+      !producer_merge_outer->isBroadcast()) {
+    // Outer ID is repeated by the factor of the inner extent
+    info.input_id = reshape_merge->outer();
+    info.factor_id = reshape_merge->inner();
+    producer_factor_id = producer_merge_inner;
+  } else {
+    return std::nullopt;
+  }
+
+  // Check if the expanded ID has a static expanded extent
+  if (!producer_factor_id->expandedExtent()->isConstInt()) {
+    return std::nullopt;
+  }
 
   return info;
 }
diff --git a/csrc/scheduler/tools/static_repeat.h b/csrc/scheduler/tools/static_repeat.h
index bfdd8d4f346..f20c02c78fd 100644
--- a/csrc/scheduler/tools/static_repeat.h
+++ b/csrc/scheduler/tools/static_repeat.h
@@ -34,7 +34,7 @@ namespace scheduler_tools {
 // repeated at the end of its computation.
 //
 // This can be problematic since the whole segment is scheduled based
-// on the repeated tensor whose size is largere than the rest of the
+// on the repeated tensor whose size is larger than the rest of the
 // tensors by the repetition factor. For example, if the it is
 // repeated twice, we would launch threads and blocks that are
 // required for the twice-larger tensor but most of the actual
@@ -53,20 +53,18 @@ namespace scheduler_tools {
 // TODO: Consider generalizing this heuristics to the other
 // schedulers.
 
+// Some of the relevant iter domains of the output tensor of the
+// reshape that realizes a repetition.
 struct StaticRepeatInfo {
-  // The final output tensor of the detected repeat pattern, e.g.,
-  // t3 in the above example case.
-  TensorView* repeat_output_tv = nullptr;
-  // The reshape output tensor, e.g., t3 in the above example case. It
-  // is not the same as repeat_output_tv when there's a cache.
-  TensorView* reshape_output_tv = nullptr;
-  // The ID of reshape output TV that corresponds to the
-  // expanded broadcast ID. In the above example case, this
-  // would be the root ID of t3 that corresponds to b2
-  IterDomain* reshape_repeat_id = nullptr;
-  // Output tensors of the detected broadcast, expand and reshape
-  // ops. In the above example case, this would consist of t1, t2 and t3.
-  std::unordered_set<TensorView*> repeat_tvs;
+  // Root ID that is repeated. In the above example, this corresponds
+  // to i1.
+  IterDomain* input_id = nullptr;
+  // Root ID that is originally an expanded broadcast. In the above example,
+  // this corresponds to b(2).
+  IterDomain* factor_id = nullptr;
+  // Logical repeated ID. In the above example, this corresponds
+  // to 2*i1.
+  IterDomain* output_id = nullptr;
 };
 
 // Check if the given tensor matches with the final reshape output
diff --git a/tests/cpp/test_rope.cpp b/tests/cpp/test_rope.cpp
index d4449c38560..0c37a2213ab 100644
--- a/tests/cpp/test_rope.cpp
+++ b/tests/cpp/test_rope.cpp
@@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
+  fusion.addOutput(tv3);
 
   auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
   auto t0 = at::randn(shape1, options);
@@ -1696,4 +1697,71 @@ TEST_F(RopeTest, EndingRepeat) {
   }
 }
 
+// Similar to EndingRepeat but with a broadcast ID already found in an
+// input tensor. A similar Pattern appears in the LitGPT Llama RoPE
+// module.
+TEST_F(RopeTest, EndingRepeatWithNoBroadcastOp) {
+  auto fusion_ptr = std::make_unique<Fusion>();
+  FusionGuard fg(fusion_ptr.get());
+  Fusion& fusion = *fusion_ptr;
+
+  std::vector<int64_t> shape1{3, 1, 200};
+
+  auto tv0 = makeContigConcreteTensor(shape1);
+  fusion.addInput(tv0);
+
+  auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
+  auto tv2 = expand(
+      tv1,
+      {IrBuilder::create<Val>(-1),
+       IrBuilder::create<Val>(2),
+       IrBuilder::create<Val>(-1)});
+  auto tv3 =
+      reshape(tv2, {IrBuilder::create<Val>(6), IrBuilder::create<Val>(-1)});
+  fusion.addOutput(tv3);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn(shape1, options);
+
+  FusionExecutorCache executor_cache(std::move(fusion_ptr));
+  auto outputs = executor_cache.runFusionWithInputs({t0});
+  testValidate(&fusion, outputs, {t0}, __LINE__, __FILE__);
+
+  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
+  EXPECT_FALSE(runtime->isSegmented());
+  const auto& heuristic_param =
+      runtime->schedulerHeuristics()->heuristicsList().front();
+  EXPECT_EQ(heuristic_param->scheduler_type, SchedulerType::Resize);
+  Fusion* scheduled_fusion = runtime->executors()
+                                 .at(0)
+                                 ->as<KernelExecutor>()
+                                 ->compiledKernel()
+                                 ->kernel();
+
+  // Similar to the EndingRepeat tensor, the repeat factor ID should
+  // be placed at the outermost position.
+  auto ref_tv = scheduled_fusion->outputs().at(0)->as<TensorView>();
+  // The outermost loop ID should be a Serial ID with an extent of 2.
+  EXPECT_EQ(
+      ref_tv->getLoopDomain().at(0)->getParallelType(), ParallelType::Serial);
+  EXPECT_TRUE(ref_tv->getLoopDomain().at(0)->extent()->isConstInt());
+  EXPECT_EQ(
+      ref_tv->getLoopDomain().at(0)->extent()->evaluate().as<int64_t>(), 2L);
+
+  IdModel id_model(scheduled_fusion, /*build_graphs=*/false);
+  const auto& exact_graph = id_model.buildExactGraph();
+
+  const auto ref_loop = exact_graph.toGroups(ref_tv->getLoopDomain());
+
+  // All of the tensors have a mapped ID as the factor ID, so they
+  // should all have the same loop ID groups.
+  for (auto tv : scheduled_fusion->allTvs()) {
+    if (tv->isFusionInput()) {
+      continue;
+    }
+    EXPECT_EQ(exact_graph.toGroups(tv->getLoopDomain()), ref_loop);
+    EXPECT_EQ(tv->getLoopDomain().size(), tv->getComputeAtPosition());
+  }
+}
+
 } // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T00:10:20Z
user: login: jjsjann123
body: LGTM in general, I have quite some questions, but those are mostly just for my own curiosity.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-30T09:11:46Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T13:20:17Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T13:36:55Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T13:40:57Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T13:42:03Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T13:50:13Z
user: login: protonu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T14:21:19Z
user: login: protonu
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-30T14:35:10Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T16:51:25Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T18:49:15Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T18:52:18Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T18:54:33Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T18:57:44Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T19:01:21Z
user: login: jjsjann123
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T19:46:30Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T20:02:05Z

Review comments:
diff_hunk: @@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
path: tests/cpp/test_rope.cpp
user: login: naoyam
body: Just to make sure the transformation is applied by ignoring plain set ops.
created_at: 2025-04-30T00:10:19Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,51 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Partition a given set of tensors to two disjoint sets based on a
+// given iter domain and rechability from the iter domain. Returns two
path: csrc/scheduler/resize.cpp
user: login: protonu
body: nit:reachability?
created_at: 2025-04-30T03:00:21Z
author_association: COLLABORATOR
diff_hunk: @@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
path: tests/cpp/test_rope.cpp
user: login: jjsjann123
body: Out of curiosity, is there any specific reason we are using a segment_set, instead of a plain set?
created_at: 2025-04-30T07:24:42Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
   // scheduling.
   bool repeat_id_moved_to_outermost = false;
-  if (static_repeat_info.has_value()) {
-    NVF_ERROR(ref_tv == static_repeat_info->repeat_output_tv);
-    auto ref_repeat_id_it = std::find_if(
+  if (repeat_info.has_value()) {
+    auto ref_factor_id_it = std::find_if(
         ref_tv->getLoopDomain().begin(),
         ref_tv->getLoopDomain().end(),
         [&](IterDomain* loop_id) {
           return id_model->idGraph(IdMappingMode::EXACT)
               .disjointValSets()
-              .strictAreMapped(loop_id, static_repeat_info->reshape_repeat_id);
+              .strictAreMapped(loop_id, repeat_info->factor_id);
         });
-    // Gives up if the repeat ID is not found. Unclear if this could
-    // actually happen, though.
-    if (ref_repeat_id_it != ref_tv->getLoopDomain().end()) {
-      auto repeat_id_pos =
-          std::distance(ref_tv->getLoopDomain().begin(), ref_repeat_id_it);
+    // The factor ID should be found in the loop domain as the
+    // reshape should be cancelled at this point. Gives up if not.
+    if (ref_factor_id_it != ref_tv->getLoopDomain().end()) {
+      auto factor_id_pos =
+          std::distance(ref_tv->getLoopDomain().begin(), ref_factor_id_it);
       NVF_ERROR(
-          repeat_id_pos >= outermost_pos,
-          "Unexpected to have DID-parallelized repeat axis: ",
-          static_repeat_info->reshape_repeat_id->toString());
+          factor_id_pos >= outermost_pos,
+          "Unexpected to have DID-parallelized repeat factor axis: ",
+          repeat_info->factor_id->toString());
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: out of curiosity, do we have any check rejecting this case?
created_at: 2025-04-30T08:22:22Z
author_association: COLLABORATOR
diff_hunk: @@ -440,34 +483,26 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // post-repeat group, where only the latter group has the repeat
   // IDs. When propagating the loop domain of the reference tensor,
   // which has the repeat ID, the full loop domain is propagated only
-  // to the post-repeat group. For the pre-repeat group, the repeat ID
+  // to the tensors that have IDs that are mapped with the repeat
+  // ID. For the rest of the tensros, the repeat ID
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: super nitpick, format on the comment with the line below.
created_at: 2025-04-30T08:23:44Z
author_association: COLLABORATOR
diff_hunk: @@ -440,34 +483,26 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // post-repeat group, where only the latter group has the repeat
   // IDs. When propagating the loop domain of the reference tensor,
   // which has the repeat ID, the full loop domain is propagated only
-  // to the post-repeat group. For the pre-repeat group, the repeat ID
+  // to the tensors that have IDs that are mapped with the repeat
+  // ID. For the rest of the tensros, the repeat ID
   // is dropped and only the remaining loop domain is propagated.
   if (repeat_id_moved_to_outermost) {
-    // Divide all tvs to the pre and posgt repeat groups
     auto all_tvs = fusion->allTvs();
-    std::vector<TensorView*> post_repeat_tvs;
-    post_repeat_tvs.reserve(static_repeat_info->repeat_tvs.size());
-    std::vector<TensorView*> pre_repeat_tvs;
-    pre_repeat_tvs.reserve(
-        all_tvs.size() - static_repeat_info->repeat_tvs.size());
-    for (auto tv : all_tvs) {
-      if (static_repeat_info->repeat_tvs.count(tv)) {
-        post_repeat_tvs.push_back(tv);
-      } else {
-        pre_repeat_tvs.push_back(tv);
-      }
-    }
+    const auto& [tvs_with_repeat_id, tvs_without_repeat_id] = partitionTvsById(
+        all_tvs,
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: nitpick:
```suggestion
        fusion->allTvs(),
```
created_at: 2025-04-30T08:25:46Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,51 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Partition a given set of tensors to two disjoint sets based on a
+// given iter domain and rechability from the iter domain. Returns two
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: ```suggestion
// given iter domain and reachability from the iter domain. Returns two
```
created_at: 2025-04-30T08:27:31Z
author_association: COLLABORATOR
diff_hunk: @@ -440,34 +483,26 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // post-repeat group, where only the latter group has the repeat
   // IDs. When propagating the loop domain of the reference tensor,
   // which has the repeat ID, the full loop domain is propagated only
-  // to the post-repeat group. For the pre-repeat group, the repeat ID
+  // to the tensors that have IDs that are mapped with the repeat
+  // ID. For the rest of the tensros, the repeat ID
   // is dropped and only the remaining loop domain is propagated.
   if (repeat_id_moved_to_outermost) {
-    // Divide all tvs to the pre and posgt repeat groups
     auto all_tvs = fusion->allTvs();
-    std::vector<TensorView*> post_repeat_tvs;
-    post_repeat_tvs.reserve(static_repeat_info->repeat_tvs.size());
-    std::vector<TensorView*> pre_repeat_tvs;
-    pre_repeat_tvs.reserve(
-        all_tvs.size() - static_repeat_info->repeat_tvs.size());
-    for (auto tv : all_tvs) {
-      if (static_repeat_info->repeat_tvs.count(tv)) {
-        post_repeat_tvs.push_back(tv);
-      } else {
-        pre_repeat_tvs.push_back(tv);
-      }
-    }
+    const auto& [tvs_with_repeat_id, tvs_without_repeat_id] = partitionTvsById(
+        all_tvs,
+        repeat_info->factor_id,
+        id_model->maybeBuildGraph(IdMappingMode::BROADCAST));
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: naive question for my own curiosity, do we need to map `BROADCAST`?

In the added example
```
1708   std::vector<int64_t> shape1{3, 1, 200};
1709 
1710   auto tv0 = makeContigConcreteTensor(shape1);
1711   fusion.addInput(tv0);
1712 
1713   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
1714   auto tv2 = expand(
1715       tv1,
1716       {IrBuilder::create<Val>(-1),
1717        IrBuilder::create<Val>(2),
1718        IrBuilder::create<Val>(-1)});
1719   auto tv3 =
1720       reshape(tv2, {IrBuilder::create<Val>(6), IrBuilder::create<Val>(-1)});
1721   fusion.addOutput(tv3);
```

Say for `tv1 [i0, b(1), i2]`, after the expand, we would have `tv2 [i0, b(2), i2]`
The two broadcast ID in tv1 and tv2 would have different extent.

Q1. IIUC, mapping with broadcast would allow us map those two together?
Q2. Does it matter for us to group `tv1` with the `tvs_with_repeat_id`, even though the it only contains the non-expanded factor_id?
created_at: 2025-04-30T08:44:04Z
author_association: COLLABORATOR
diff_hunk: @@ -154,10 +60,39 @@ std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
   }
 
   StaticRepeatInfo info;
-  info.repeat_output_tv = maybe_repeat_out;
-  info.reshape_output_tv = reshape_out;
-  info.reshape_repeat_id = reshape_out->getRootDomain().at(broadcast_pos);
-  info.repeat_tvs = repeat_tvs;
+
+  info.output_id = reshape_merge->out();
+
+  const auto c2p =
+      PairwiseLogicalDomainMap(reshape_in, reshape_out).mapConsumerToProducer();
+
+  auto producer_merge_outer = c2p.at(reshape_merge->outer());
+  auto producer_merge_inner = c2p.at(reshape_merge->inner());
+  IterDomain* producer_factor_id = nullptr;
+
+  if (producer_merge_outer->isBroadcast() &&
+      producer_merge_outer->hasExpandedExtent() &&
+      !producer_merge_inner->isBroadcast()) {
path: csrc/scheduler/tools/static_repeat.cpp
user: login: jjsjann123
body: > !producer_merge_inner->isBroadcast()

Do we need to check for this? What happens when we merge two broadcast ID both with expanded extent?
An earlier comment says

>  54   // Reshape of an expanded broadcast always generates a concrete
 55   // non-broadcast ID, so this check is not necessary, but just in
 56   // case in the future that may change.

created_at: 2025-04-30T09:11:23Z
author_association: COLLABORATOR
diff_hunk: @@ -8,126 +8,40 @@
 
 #include <ir/all_nodes.h>
 #include <ir/utils.h>
+#include <logical_domain_map.h>
 #include <scheduler/tools/static_repeat.h>
 
 namespace nvfuser {
 namespace scheduler_tools {
 
 std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
-    TensorView* maybe_repeat_out) {
-  // The pattern to detect:
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // repeat_out = reshape(expand_out)
-  //
-  // Additionally, since maybe_repeat_out is commonly a fusion
-  // output, it is likely there's a cache tv between expand_out and
-  // repeat_out, so the following pattern should also be detected.
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // cache_of_repeat_out = reshape(expand_out)
-  // repeat_out = set(cache_of_repeat_out)
-
-  std::unordered_set<TensorView*> repeat_tvs;
-  repeat_tvs.insert(maybe_repeat_out);
-
-  auto reshape_out = maybe_repeat_out;
-
-  // Check if there's a cache
-  if (auto ldst = dynamic_cast<LoadStoreOp*>(maybe_repeat_out->definition());
-      ldst != nullptr && ldst->opType() == LoadStoreOpType::Set) {
-    reshape_out = ldst->in()->as<TensorView>();
-    repeat_tvs.insert(reshape_out);
+    TensorView* maybe_repeat_out_tv) {
+  // Skip set ops if any (e.g., inserted by caching). Only Set
path: csrc/scheduler/tools/static_repeat.cpp
user: login: protonu
body: I thought the deleted comment here was helpful - the bit about skipping caching ops.
created_at: 2025-04-30T13:36:55Z
author_association: COLLABORATOR
diff_hunk: @@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
path: tests/cpp/test_rope.cpp
user: login: protonu
body: Is it to ensure this:

// It is especially important to recognize this pattern when it
// appears at the end of a pointwise fusion segment, where an output
// is used as the reference tensor of scheduling the segment. 


If so, should the caching ops be mixed inside the `repeat`? 
created_at: 2025-04-30T13:40:56Z
author_association: COLLABORATOR
diff_hunk: @@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
path: tests/cpp/test_rope.cpp
user: login: protonu
body: Though I guess this was the end of the fusion anyway - so that may not make much sense.

created_at: 2025-04-30T13:42:03Z
author_association: COLLABORATOR
diff_hunk: @@ -8,126 +8,40 @@
 
 #include <ir/all_nodes.h>
 #include <ir/utils.h>
+#include <logical_domain_map.h>
 #include <scheduler/tools/static_repeat.h>
 
 namespace nvfuser {
 namespace scheduler_tools {
 
 std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
-    TensorView* maybe_repeat_out) {
-  // The pattern to detect:
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // repeat_out = reshape(expand_out)
-  //
-  // Additionally, since maybe_repeat_out is commonly a fusion
-  // output, it is likely there's a cache tv between expand_out and
-  // repeat_out, so the following pattern should also be detected.
-  //
-  // broadcast_out = broadcast(input)
path: csrc/scheduler/tools/static_repeat.cpp
user: login: protonu
body: These 4 lines were helpful.
created_at: 2025-04-30T13:50:12Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
path: csrc/scheduler/resize.cpp
user: login: protonu
body: nit: move the def of `repeat_info` here - near the use.
created_at: 2025-04-30T14:21:17Z
author_association: COLLABORATOR
diff_hunk: @@ -1631,7 +1631,8 @@ TEST_F(RopeTest, EndingRepeat) {
 
   auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
   auto tv2 = repeat(tv1, {2, 1});
-  fusion.addOutput(tv2);
+  auto tv3 = segment_set(tv2);
path: tests/cpp/test_rope.cpp
user: login: naoyam
body: Just because I saw a segment ending with `segment_set`, preceded by a repetition. I think that's one segment of a LitGPT Llama forward.
created_at: 2025-04-30T16:51:25Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: I'd keep it there as `ref_tv` is going to be transformed after that. It shouldn't affect the analysis, but there's no need to introduce an additional complexity.
created_at: 2025-04-30T18:49:15Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
   // scheduling.
   bool repeat_id_moved_to_outermost = false;
-  if (static_repeat_info.has_value()) {
-    NVF_ERROR(ref_tv == static_repeat_info->repeat_output_tv);
-    auto ref_repeat_id_it = std::find_if(
+  if (repeat_info.has_value()) {
+    auto ref_factor_id_it = std::find_if(
         ref_tv->getLoopDomain().begin(),
         ref_tv->getLoopDomain().end(),
         [&](IterDomain* loop_id) {
           return id_model->idGraph(IdMappingMode::EXACT)
               .disjointValSets()
-              .strictAreMapped(loop_id, static_repeat_info->reshape_repeat_id);
+              .strictAreMapped(loop_id, repeat_info->factor_id);
         });
-    // Gives up if the repeat ID is not found. Unclear if this could
-    // actually happen, though.
-    if (ref_repeat_id_it != ref_tv->getLoopDomain().end()) {
-      auto repeat_id_pos =
-          std::distance(ref_tv->getLoopDomain().begin(), ref_repeat_id_it);
+    // The factor ID should be found in the loop domain as the
+    // reshape should be cancelled at this point. Gives up if not.
+    if (ref_factor_id_it != ref_tv->getLoopDomain().end()) {
+      auto factor_id_pos =
+          std::distance(ref_tv->getLoopDomain().begin(), ref_factor_id_it);
       NVF_ERROR(
-          repeat_id_pos >= outermost_pos,
-          "Unexpected to have DID-parallelized repeat axis: ",
-          static_repeat_info->reshape_repeat_id->toString());
+          factor_id_pos >= outermost_pos,
+          "Unexpected to have DID-parallelized repeat factor axis: ",
+          repeat_info->factor_id->toString());
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: It's our scheduling, so I'm not sure how to reject it.
created_at: 2025-04-30T18:52:17Z
author_association: COLLABORATOR
diff_hunk: @@ -8,126 +8,40 @@
 
 #include <ir/all_nodes.h>
 #include <ir/utils.h>
+#include <logical_domain_map.h>
 #include <scheduler/tools/static_repeat.h>
 
 namespace nvfuser {
 namespace scheduler_tools {
 
 std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
-    TensorView* maybe_repeat_out) {
-  // The pattern to detect:
-  //
-  // broadcast_out = broadcast(input)
-  // expand_out = expand(broadcast_out)
-  // repeat_out = reshape(expand_out)
-  //
-  // Additionally, since maybe_repeat_out is commonly a fusion
-  // output, it is likely there's a cache tv between expand_out and
-  // repeat_out, so the following pattern should also be detected.
-  //
-  // broadcast_out = broadcast(input)
path: csrc/scheduler/tools/static_repeat.cpp
user: login: naoyam
body: Not sure which 4 lines, but broadcast is no longer required.
created_at: 2025-04-30T18:54:33Z
author_association: COLLABORATOR
diff_hunk: @@ -154,10 +60,39 @@ std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
   }
 
   StaticRepeatInfo info;
-  info.repeat_output_tv = maybe_repeat_out;
-  info.reshape_output_tv = reshape_out;
-  info.reshape_repeat_id = reshape_out->getRootDomain().at(broadcast_pos);
-  info.repeat_tvs = repeat_tvs;
+
+  info.output_id = reshape_merge->out();
+
+  const auto c2p =
+      PairwiseLogicalDomainMap(reshape_in, reshape_out).mapConsumerToProducer();
+
+  auto producer_merge_outer = c2p.at(reshape_merge->outer());
+  auto producer_merge_inner = c2p.at(reshape_merge->inner());
+  IterDomain* producer_factor_id = nullptr;
+
+  if (producer_merge_outer->isBroadcast() &&
+      producer_merge_outer->hasExpandedExtent() &&
+      !producer_merge_inner->isBroadcast()) {
path: csrc/scheduler/tools/static_repeat.cpp
user: login: naoyam
body: I think then we should factor out both IDs, so we would want to have a slightly different scheduling approach. I'm not sure if that's worth the additional work.
created_at: 2025-04-30T18:57:44Z
author_association: COLLABORATOR
diff_hunk: @@ -440,34 +483,26 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // post-repeat group, where only the latter group has the repeat
   // IDs. When propagating the loop domain of the reference tensor,
   // which has the repeat ID, the full loop domain is propagated only
-  // to the post-repeat group. For the pre-repeat group, the repeat ID
+  // to the tensors that have IDs that are mapped with the repeat
+  // ID. For the rest of the tensros, the repeat ID
   // is dropped and only the remaining loop domain is propagated.
   if (repeat_id_moved_to_outermost) {
-    // Divide all tvs to the pre and posgt repeat groups
     auto all_tvs = fusion->allTvs();
-    std::vector<TensorView*> post_repeat_tvs;
-    post_repeat_tvs.reserve(static_repeat_info->repeat_tvs.size());
-    std::vector<TensorView*> pre_repeat_tvs;
-    pre_repeat_tvs.reserve(
-        all_tvs.size() - static_repeat_info->repeat_tvs.size());
-    for (auto tv : all_tvs) {
-      if (static_repeat_info->repeat_tvs.count(tv)) {
-        post_repeat_tvs.push_back(tv);
-      } else {
-        pre_repeat_tvs.push_back(tv);
-      }
-    }
+    const auto& [tvs_with_repeat_id, tvs_without_repeat_id] = partitionTvsById(
+        all_tvs,
+        repeat_info->factor_id,
+        id_model->maybeBuildGraph(IdMappingMode::BROADCAST));
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: `EXACT` should work too. Previously, we schedule those tensors like `tv1` and `tv2` together with `tv3`, so using `BROADCAST` keeps the same behavior. I don't think there should be any actual difference in final performances.
created_at: 2025-04-30T19:01:21Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
   // scheduling.
   bool repeat_id_moved_to_outermost = false;
-  if (static_repeat_info.has_value()) {
-    NVF_ERROR(ref_tv == static_repeat_info->repeat_output_tv);
-    auto ref_repeat_id_it = std::find_if(
+  if (repeat_info.has_value()) {
+    auto ref_factor_id_it = std::find_if(
         ref_tv->getLoopDomain().begin(),
         ref_tv->getLoopDomain().end(),
         [&](IterDomain* loop_id) {
           return id_model->idGraph(IdMappingMode::EXACT)
               .disjointValSets()
-              .strictAreMapped(loop_id, static_repeat_info->reshape_repeat_id);
+              .strictAreMapped(loop_id, repeat_info->factor_id);
         });
-    // Gives up if the repeat ID is not found. Unclear if this could
-    // actually happen, though.
-    if (ref_repeat_id_it != ref_tv->getLoopDomain().end()) {
-      auto repeat_id_pos =
-          std::distance(ref_tv->getLoopDomain().begin(), ref_repeat_id_it);
+    // The factor ID should be found in the loop domain as the
+    // reshape should be cancelled at this point. Gives up if not.
+    if (ref_factor_id_it != ref_tv->getLoopDomain().end()) {
+      auto factor_id_pos =
+          std::distance(ref_tv->getLoopDomain().begin(), ref_factor_id_it);
       NVF_ERROR(
-          repeat_id_pos >= outermost_pos,
-          "Unexpected to have DID-parallelized repeat axis: ",
-          static_repeat_info->reshape_repeat_id->toString());
+          factor_id_pos >= outermost_pos,
+          "Unexpected to have DID-parallelized repeat factor axis: ",
+          repeat_info->factor_id->toString());
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: If we initially have `factor_id` parallelized as `DID`, resize scheduler shouldn't take the fusion in that case.
created_at: 2025-04-30T19:46:30Z
author_association: COLLABORATOR
diff_hunk: @@ -352,35 +397,33 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // detected. The repeat ID then just remains there with no
   // scheduling.
   bool repeat_id_moved_to_outermost = false;
-  if (static_repeat_info.has_value()) {
-    NVF_ERROR(ref_tv == static_repeat_info->repeat_output_tv);
-    auto ref_repeat_id_it = std::find_if(
+  if (repeat_info.has_value()) {
+    auto ref_factor_id_it = std::find_if(
         ref_tv->getLoopDomain().begin(),
         ref_tv->getLoopDomain().end(),
         [&](IterDomain* loop_id) {
           return id_model->idGraph(IdMappingMode::EXACT)
               .disjointValSets()
-              .strictAreMapped(loop_id, static_repeat_info->reshape_repeat_id);
+              .strictAreMapped(loop_id, repeat_info->factor_id);
         });
-    // Gives up if the repeat ID is not found. Unclear if this could
-    // actually happen, though.
-    if (ref_repeat_id_it != ref_tv->getLoopDomain().end()) {
-      auto repeat_id_pos =
-          std::distance(ref_tv->getLoopDomain().begin(), ref_repeat_id_it);
+    // The factor ID should be found in the loop domain as the
+    // reshape should be cancelled at this point. Gives up if not.
+    if (ref_factor_id_it != ref_tv->getLoopDomain().end()) {
+      auto factor_id_pos =
+          std::distance(ref_tv->getLoopDomain().begin(), ref_factor_id_it);
       NVF_ERROR(
-          repeat_id_pos >= outermost_pos,
-          "Unexpected to have DID-parallelized repeat axis: ",
-          static_repeat_info->reshape_repeat_id->toString());
+          factor_id_pos >= outermost_pos,
+          "Unexpected to have DID-parallelized repeat factor axis: ",
+          repeat_info->factor_id->toString());
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: Generally speaking, I'm sure there are many issues with DID types with the resize scheduler, in particular when logical domains are split, which I haven't considered while developing this scheduler.

That said, in this case, I think that it should not happen since `factor_id` is a root ID of this reference tensor or even it is non-existent before this resize scheduling. So there's no way to parallelize the ID before this scheduling.
created_at: 2025-04-30T20:02:05Z
author_association: COLLABORATOR

PR comments:
user: login: naoyam
created_at: 2025-04-26T19:19:36Z
body: !test --diff
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-26T19:19:57Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/e803fa07776afb76c1091108448e3a47bde048c5*


### **Description**
- Enhanced static repeat detection to handle cases without BroadcastOp

- Added `partitionTvsById` function for tensor partitioning

- Updated `getMaybeStaticRepeatInfo` to detect reshape patterns more flexibly

- Modified test cases to include scenarios with no BroadcastOp


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>resize.cpp</strong><dd><code>Add tensor partitioning and update repeat ID handling</code>&nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/resize.cpp

<li>Added <code>partitionTvsById</code> function for tensor partitioning based on iter <br>domain reachability<br> <li> Updated <code>ResizeScheduler::schedule</code> to use <code>getMaybeStaticRepeatInfo</code> and <br>handle repeat ID reordering<br> <li> Modified tensor partitioning logic to group tensors with and without <br>repeat IDs


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427a">+67/-37</a>&nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>static_repeat.cpp</strong><dd><code>Simplify static repeat detection logic</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/tools/static_repeat.cpp

<li>Simplified <code>getMaybeStaticRepeatInfo</code> to detect reshape patterns without <br>requiring BroadcastOp<br> <li> Removed unnecessary checks for BroadcastOp, ExpandOp, and specific op <br>sequences<br> <li> Updated logic to identify factor and input IDs using <br><code>PairwiseLogicalDomainMap</code>


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-3131e2576f5042e9a2e1d8742887aa61a398953d65795c989a18a17ddbda8b77">+48/-113</a></td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>static_repeat.h</strong><dd><code>Update StaticRepeatInfo structure</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/tools/static_repeat.h

<li>Updated <code>StaticRepeatInfo</code> structure to include <code>input_id</code>, <code>factor_id</code>, and <br><code>output_id</code><br> <li> Removed redundant fields <code>repeat_output_tv</code>, <code>reshape_output_tv</code>, and <br><code>repeat_tvs</code>


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-4dcf140ce7a2ea861a324337f20641884c763b49e078a829fb0e26381c8b5bee">+12/-14</a>&nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_rope.cpp</strong><dd><code>Add test cases for static repeat without BroadcastOp</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

tests/cpp/test_rope.cpp

<li>Added new test case <code>EndingRepeatWithNoBroadcastOp</code> to verify static <br>repeat detection without BroadcastOp<br> <li> Modified existing test case <code>EndingRepeat</code> to include <code>segment_set</code> <br>operation


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-b2a05bd8e964824d77e2957d8c9c2d89ccb4122dd94e7880233d0a03c5542393">+69/-1</a>&nbsp; &nbsp; </td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>PR contains tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427aR318-R364'><strong>Logic Change</strong></a>

The logic for detecting static repeats has been significantly changed. Ensure that the new logic correctly identifies all valid static repeat patterns and does not introduce false positives or negatives.</summary>

```c++
// Partition a given set of tensors to two disjoint sets based on a
// given iter domain and reachability from the iter domain. Returns two
// vectors of tensors, first of which contains all tensors that has an
// iter domain that is reachable from the given iter domain, whereas
// the rest of tensors are all grouped into the second
// list. Reachability is determined by using the permissive BFS
// traversal on a given graph.
std::pair<std::vector<TensorView*>, std::vector<TensorView*>> partitionTvsById(
    const std::vector<TensorView*> tvs,
    IterDomain* id,
    const ValGraph& graph) {
  ValGroups target_groups;
  for (auto tv : tvs) {
    target_groups.pushBack(graph.toGroups(tv->getLogicalDomain()));
  }

  const auto reachable_groups = getReachableValsFrom<ValGraphPermissiveBFS>(
      {graph.toGroup(id)},
      target_groups.vector(),
      /*allowed_direction=*/Direction::Undefined,
      graph);
  const std::unordered_set<ValGroup> reachable_group_set{
      reachable_groups.begin(), reachable_groups.end()};

  std::vector<TensorView*> reachable_tvs;
  std::vector<TensorView*> unreachable_tvs;

  for (auto tv : tvs) {
    if (std::ranges::any_of(
            tv->getLogicalDomain(), [&](IterDomain* logical_id) {
              return reachable_group_set.contains(graph.toGroup(logical_id));
            })) {
      reachable_tvs.push_back(tv);
    } else {
      unreachable_tvs.push_back(tv);
    }
  }

  return std::make_pair(reachable_tvs, unreachable_tvs);
}

} // namespace

void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
  FUSER_PERF_SCOPE("ResizeScheduler::schedule");

  FusionGuard fg(fusion);
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-3131e2576f5042e9a2e1d8742887aa61a398953d65795c989a18a17ddbda8b77R1-R96'><strong>Simplification</strong></a>

The detection of static repeats has been simplified. Verify that the new detection logic is still robust and covers all necessary cases, especially those involving complex patterns.</summary>

```c++
// clang-format off
/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 */
// clang-format on

#include <ir/all_nodes.h>
#include <ir/utils.h>
#include <logical_domain_map.h>
#include <scheduler/tools/static_repeat.h>

namespace nvfuser {
namespace scheduler_tools {

std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
    TensorView* maybe_repeat_out_tv) {
  // Skip set ops if any (e.g., inserted by caching). Only Set
  // or SegmenterSet are considered.
  while (auto ldst =
             dynamic_cast<LoadStoreOp*>(maybe_repeat_out_tv->definition())) {
    if (ldst->opType() != LoadStoreOpType::Set &&
        ldst->opType() != LoadStoreOpType::SegmenterSet) {
      break;
    }
    maybe_repeat_out_tv = ldst->in()->as<TensorView>();
  }

  // Detect reshape
  auto reshape = dynamic_cast<ViewOp*>(maybe_repeat_out_tv->definition());
  if (reshape == nullptr) {
    return std::nullopt;
  }

  auto reshape_in = reshape->input(0)->as<TensorView>();
  auto reshape_out = reshape->output(0)->as<TensorView>();

  auto reshape_exprs = DependencyCheck::getAllExprsBetween(
      {reshape_out->getRootDomain().begin(),
       reshape_out->getRootDomain().end()},
      {reshape_out->getLogicalDomain().begin(),
       reshape_out->getLogicalDomain().end()});

  if (reshape_exprs.size() != 1) {
    return std::nullopt;
  }

  auto reshape_merge = dynamic_cast<Merge*>(reshape_exprs.at(0));
  if (reshape_merge == nullptr) {
    return std::nullopt;
  }

  // Reshape of an expanded broadcast always generates a concrete
  // non-broadcast ID, so this check is not necessary, but just in
  // case in the future that may change.
  if (reshape_merge->out()->isBroadcast() ||
      reshape_merge->out()->hasExpandedExtent()) {
    return std::nullopt;
  }

  StaticRepeatInfo info;

  info.output_id = reshape_merge->out();

  const auto c2p =
      PairwiseLogicalDomainMap(reshape_in, reshape_out).mapConsumerToProducer();

  auto producer_merge_outer = c2p.at(reshape_merge->outer());
  auto producer_merge_inner = c2p.at(reshape_merge->inner());
  IterDomain* producer_factor_id = nullptr;

  if (producer_merge_outer->isBroadcast() &&
      producer_merge_outer->hasExpandedExtent() &&
      !producer_merge_inner->isBroadcast()) {
    // Inner ID is repeated by the factor of the outer extent
    info.input_id = reshape_merge->inner();
    info.factor_id = reshape_merge->outer();
    producer_factor_id = producer_merge_outer;
  } else if (
      producer_merge_inner->isBroadcast() &&
      producer_merge_inner->hasExpandedExtent() &&
      !producer_merge_outer->isBroadcast()) {
    // Outer ID is repeated by the factor of the inner extent
    info.input_id = reshape_merge->outer();
    info.factor_id = reshape_merge->inner();
    producer_factor_id = producer_merge_inner;
  } else {
    return std::nullopt;
  }

  // Check if the expanded ID has a static expanded extent
  if (!producer_factor_id->expandedExtent()->isConstInt()) {
    return std::nullopt;
  }

```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-b2a05bd8e964824d77e2957d8c9c2d89ccb4122dd94e7880233d0a03c5542393R1628-R1767'><strong>New Test</strong></a>

A new test case has been added for a scenario without a broadcast op. Ensure that this test case adequately covers the new logic and that the test is comprehensive.</summary>

```c++

  auto tv0 = makeContigConcreteTensor(shape1);
  fusion.addInput(tv0);

  auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
  auto tv2 = repeat(tv1, {2, 1});
  auto tv3 = segment_set(tv2);
  fusion.addOutput(tv3);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape1, options);

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs({t0});
  testValidate(&fusion, outputs, {t0}, __LINE__, __FILE__);

  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
  EXPECT_FALSE(runtime->isSegmented());
  const auto& heuristic_param =
      runtime->schedulerHeuristics()->heuristicsList().front();
  EXPECT_EQ(heuristic_param->scheduler_type, SchedulerType::Resize);
  Fusion* scheduled_fusion = runtime->executors()
                                 .at(0)
                                 ->as<KernelExecutor>()
                                 ->compiledKernel()
                                 ->kernel();

  // Check the loop domain of the reference. It should look like:
  //
  // T4_g_float[iS19{2 ex 2}, iblockIdx.x22{8}, ithreadIdx.x23{128}] ca_pos( 3 )
  // produce_pos( 3 )
  //  logical domain : (iS17{( 2 * 8 )}, iS18{128})
  //  contiguity: t t
  //   Merge: iS20{8} and iS18{128} -> iS21{1024}
  //   Split: iS21{1024} by factor 128 -> iblockIdx.x22{8}, ithreadIdx.x23{128}
  //  loop domain : (iS19{2 ex 2}, iblockIdx.x22{8}, ithreadIdx.x23{128})
  //
  // iS19 is the repeat ID, which should be just a Serial ID with an
  // extent of 2.
  auto ref_tv = scheduled_fusion->outputs().at(0)->as<TensorView>();
  // The outermost loop ID should be a Serial ID with an extent of 2.
  EXPECT_EQ(
      ref_tv->getLoopDomain().at(0)->getParallelType(), ParallelType::Serial);
  EXPECT_TRUE(ref_tv->getLoopDomain().at(0)->extent()->isConstInt());
  EXPECT_EQ(
      ref_tv->getLoopDomain().at(0)->extent()->evaluate().as<int64_t>(), 2L);

  IdModel id_model(scheduled_fusion, /*build_graphs=*/false);
  const auto& exact_graph = id_model.buildExactGraph();

  const auto ref_loop = exact_graph.toGroups(ref_tv->getLoopDomain());

  // The other tensors, except for the pad output, should be fully inlined into
  // the reference tensor.
  for (auto tv : scheduled_fusion->allTvs()) {
    if (tv->isFusionInput()) {
      continue;
    }
    auto tv_loop = exact_graph.toGroups(tv->getLoopDomain());
    if (tv->definition() != nullptr && tv->definition()->isA<PadOp>()) {
      ValGroups ref_groups{ref_loop.begin() + 1, ref_loop.end()};
      // In the case of pad, the loop domain of the output tensor
      // should be mapped with the loop domain of the reference
      // without the outermost ID.
      EXPECT_EQ(tv_loop, ref_groups);
    } else {
      EXPECT_EQ(tv_loop, ref_loop);
      EXPECT_EQ(tv->getLoopDomain().size(), tv->getComputeAtPosition());
    }
  }
}

// Similar to EndingRepeat but with a broadcast ID already found in an
// input tensor. A similar Pattern appears in the LitGPT Llama RoPE
// module.
TEST_F(RopeTest, EndingRepeatWithNoBroadcastOp) {
  auto fusion_ptr = std::make_unique<Fusion>();
  FusionGuard fg(fusion_ptr.get());
  Fusion& fusion = *fusion_ptr;

  std::vector<int64_t> shape1{3, 1, 200};

  auto tv0 = makeContigConcreteTensor(shape1);
  fusion.addInput(tv0);

  auto tv1 = pad(tv0, {fusion.oneVal(), fusion.oneVal()});
  auto tv2 = expand(
      tv1,
      {IrBuilder::create<Val>(-1),
       IrBuilder::create<Val>(2),
       IrBuilder::create<Val>(-1)});
  auto tv3 =
      reshape(tv2, {IrBuilder::create<Val>(6), IrBuilder::create<Val>(-1)});
  fusion.addOutput(tv3);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn(shape1, options);

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs({t0});
  testValidate(&fusion, outputs, {t0}, __LINE__, __FILE__);

  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
  EXPECT_FALSE(runtime->isSegmented());
  const auto& heuristic_param =
      runtime->schedulerHeuristics()->heuristicsList().front();
  EXPECT_EQ(heuristic_param->scheduler_type, SchedulerType::Resize);
  Fusion* scheduled_fusion = runtime->executors()
                                 .at(0)
                                 ->as<KernelExecutor>()
                                 ->compiledKernel()
                                 ->kernel();

  // Similar to the EndingRepeat tensor, the repeat factor ID should
  // be placed at the outermost position.
  auto ref_tv = scheduled_fusion->outputs().at(0)->as<TensorView>();
  // The outermost loop ID should be a Serial ID with an extent of 2.
  EXPECT_EQ(
      ref_tv->getLoopDomain().at(0)->getParallelType(), ParallelType::Serial);
  EXPECT_TRUE(ref_tv->getLoopDomain().at(0)->extent()->isConstInt());
  EXPECT_EQ(
      ref_tv->getLoopDomain().at(0)->extent()->evaluate().as<int64_t>(), 2L);

  IdModel id_model(scheduled_fusion, /*build_graphs=*/false);
  const auto& exact_graph = id_model.buildExactGraph();

  const auto ref_loop = exact_graph.toGroups(ref_tv->getLoopDomain());

  // All of the tensors have a mapped ID as the factor ID, so they
  // should all have the same loop ID groups.
  for (auto tv : scheduled_fusion->allTvs()) {
    if (tv->isFusionInput()) {
      continue;
    }
    EXPECT_EQ(exact_graph.toGroups(tv->getLoopDomain()), ref_loop);
    EXPECT_EQ(tv->getLoopDomain().size(), tv->getComputeAtPosition());
  }
}

} // namespace nvfuser
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: naoyam
created_at: 2025-04-26T22:35:49Z
body: !test --diff
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-26T22:36:27Z
body: <!-- llm action combined persistent comment, dont modify this line -->

### **Description**
- Refactored static repeat detection logic

- Introduced `partitionTvsById` function for better TV partitioning

- Simplified and optimized `getMaybeStaticRepeatInfo` function

- Updated variable names for clarity


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>resize.cpp</strong><dd><code>Enhance static repeat scheduling</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/resize.cpp

<li>Added <code>partitionTvsById</code> function for partitioning TensorViews based on <br>IterDomain<br> <li> Updated <code>schedule</code> method to use <code>partitionTvsById</code> for TV partitioning<br> <li> Renamed <code>static_repeat_info</code> to <code>repeat_info</code> for clarity<br> <li> Updated logic to handle repeat ID movement and partitioning


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427a">+59/-32</a>&nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>static_repeat.cpp</strong><dd><code>Simplify static repeat detection</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/tools/static_repeat.cpp

<li>Simplified <code>getMaybeStaticRepeatInfo</code> function<br> <li> Removed unnecessary checks and assumptions<br> <li> Updated logic to detect static repeat patterns more accurately<br> <li> Introduced <code>PairwiseLogicalDomainMap</code> for better domain mapping


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-3131e2576f5042e9a2e1d8742887aa61a398953d65795c989a18a17ddbda8b77">+42/-112</a></td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>static_repeat.h</strong><dd><code>Update StaticRepeatInfo structure</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/tools/static_repeat.h

<li>Updated <code>StaticRepeatInfo</code> structure with clearer variable names<br> <li> Removed redundant fields and added new fields for better <br>representation


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4325/files#diff-4dcf140ce7a2ea861a324337f20641884c763b49e078a829fb0e26381c8b5bee">+8/-13</a>&nbsp; &nbsp; </td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>No relevant tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427aR250-R281'><strong>Possible Issue</strong></a>

The `partitionTvsById` function does not handle cases where the `ValGraph` might not contain all the necessary mappings, which could lead to incorrect partitioning of `TensorView`s.</summary>

```c++
std::pair<std::vector<TensorView*>, std::vector<TensorView*>> partitionTvsById(
    const std::vector<TensorView*> tvs,
    IterDomain* id,
    const ValGraph& graph) {
  ValGroups target_groups;
  for (auto tv : tvs) {
    target_groups.pushBack(graph.toGroups(tv->getLogicalDomain()));
  }

  const auto reachable_groups = getReachableValsFrom<ValGraphPermissiveBFS>(
      {graph.toGroup(id)},
      target_groups.vector(),
      /*allowed_direction=*/Direction::Undefined,
      graph);
  const std::unordered_set<ValGroup> reachable_group_set{
      reachable_groups.begin(), reachable_groups.end()};

  std::vector<TensorView*> reachable_tvs;
  std::vector<TensorView*> unreachable_tvs;

  for (auto tv : tvs) {
    if (std::ranges::any_of(
            tv->getLogicalDomain(), [&](IterDomain* logical_id) {
              return reachable_group_set.contains(graph.toGroup(logical_id));
            })) {
      reachable_tvs.push_back(tv);
    } else {
      unreachable_tvs.push_back(tv);
    }
  }

  return std::make_pair(reachable_tvs, unreachable_tvs);
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-3131e2576f5042e9a2e1d8742887aa61a398953d65795c989a18a17ddbda8b77R17-R90'><strong>Simplification</strong></a>

The new implementation of `getMaybeStaticRepeatInfo` is more concise but could benefit from additional comments to explain the logic, especially around the use of `PairwiseLogicalDomainMap`.</summary>

```c++
std::optional<StaticRepeatInfo> getMaybeStaticRepeatInfo(
    TensorView* maybe_repeat_out_tv) {
  // Skip a set if any (e.g., inserted by caching)
  if (auto ldst = dynamic_cast<LoadStoreOp*>(maybe_repeat_out_tv->definition());
      ldst != nullptr && ldst->opType() == LoadStoreOpType::Set) {
    maybe_repeat_out_tv = ldst->in()->as<TensorView>();
  }

  // Detect reshape
  auto reshape = dynamic_cast<ViewOp*>(maybe_repeat_out_tv->definition());
  if (reshape == nullptr) {
    return std::nullopt;
  }

  auto reshape_in = reshape->input(0)->as<TensorView>();
  auto reshape_out = reshape->output(0)->as<TensorView>();

  auto reshape_exprs = DependencyCheck::getAllExprsBetween(
      {reshape_out->getRootDomain().begin(),
       reshape_out->getRootDomain().end()},
      {reshape_out->getLogicalDomain().begin(),
       reshape_out->getLogicalDomain().end()});

  if (reshape_exprs.size() != 1) {
    return std::nullopt;
  }

  auto reshape_merge = dynamic_cast<Merge*>(reshape_exprs.at(0));
  if (reshape_merge == nullptr) {
    return std::nullopt;
  }

  // Reshape of an expanded broadcast always generates a concrete
  // non-broadcast ID, so this check is not necessary, but just in
  // case in the future that may change.
  if (reshape_merge->out()->isBroadcast() ||
      reshape_merge->out()->hasExpandedExtent()) {
    return std::nullopt;
  }

  StaticRepeatInfo info;

  info.output_id = reshape_merge->out();

  const auto c2p =
      PairwiseLogicalDomainMap(reshape_in, reshape_out).mapConsumerToProducer();

  auto producer_merge_outer = c2p.at(reshape_merge->outer());
  auto producer_merge_inner = c2p.at(reshape_merge->inner());
  IterDomain* producer_factor_id = nullptr;

  if (producer_merge_outer->isBroadcast() &&
      producer_merge_outer->hasExpandedExtent() &&
      !producer_merge_inner->isBroadcast()) {
    // Inner ID is repeated by the factor of the outer extent
    info.input_id = reshape_merge->inner();
    info.factor_id = reshape_merge->outer();
    producer_factor_id = producer_merge_outer;
  } else if (
      producer_merge_inner->isBroadcast() &&
      producer_merge_inner->hasExpandedExtent() &&
      !producer_merge_outer->isBroadcast()) {
    // Outer ID is repeated by the factor of the inner extent
    info.input_id = reshape_merge->outer();
    info.factor_id = reshape_merge->inner();
    producer_factor_id = producer_merge_inner;
  } else {
    return std::nullopt;
  }

  // Check if the expanded ID has a static expanded extent
  if (!producer_factor_id->expandedExtent()->isConstInt()) {
    return std::nullopt;
  }
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4325/files#diff-4dcf140ce7a2ea861a324337f20641884c763b49e078a829fb0e26381c8b5beeR52-R63'><strong>Documentation</strong></a>

The updated `StaticRepeatInfo` struct fields need better documentation to clarify their roles and relationships within the scheduling process.</summary>

```c
// propagated to the rest of the tensors.
//
// TODO: Consider generalizing this heuristics to the other
// schedulers.

struct StaticRepeatInfo {
  // Root ID that is repeated
  IterDomain* input_id = nullptr;
  // Root ID that is originally an expanded broadcast
  IterDomain* factor_id = nullptr;
  // Logical repeated ID
  IterDomain* output_id = nullptr;
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: naoyam
created_at: 2025-04-29T03:08:20Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-29T03:16:03Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-29T03:16:51Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-30T00:10:57Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-30T19:02:05Z
body: !build
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-30T19:44:17Z
body: !build
author_association: COLLABORATOR