Repository: NVIDIA/Fuser
Pull request #4259: Fix resize scheduler with reshapes appearing after ref tensor
Author: login: naoyam
Labels: name: rope
Created: 2025-04-16T19:05:15Z | Merged: 2025-04-30T17:05:51Z

Fix #4250

In the resize scheduler, the scheduling applied to the reference tensor is propagated to all the other tensors by a backward traversal from the reference loop domain over the exact graph. This means that the loop domains of the other tensors need to be reachable from the reference. 

The requirement does not hold when, for example, there's a reshape after the reference. See the added test for a concrete example. More generally, if the logical domain of a tensor is not reachable by just traversing backward from the reference logical domain, there must be something that should prevent the transform propagation. Note that since the resize scheduler does not allow no additional `Resize` ops between the ref and outputs, these additional ops between the ref and other tensors must come from reshapes and should be "cancelable", i.e., setting the loop domain as the producer of the reshaped logical domain. `prepareForBackwardTransformPropagation`, which is called just before propagating ref transformations, searches for those tensors that have additional transformations and make their loop domains look like the reference logical domain.


Linked issue #4250: LitGPT based Llama backward error around RoPE
I am not sure what extra logic is here but there is something in the backward pass in the full model that is triggering an error.

Error:
```
Error from segmentation group 12:  INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/scheduler/tools/loop_domain_scheduler.cpp":436, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Trying to update the current loop domain but could not find a valid path from the referenc
e: T167_l___bfloat[bS490{1}, iS491{2048}, iS496{1024}rf]. Missing target ID groups: idg{496 660}
Exception raised from getReplayPath at /opt/pytorch/nvfuser/csrc/scheduler/tools/loop_domain_scheduler.cpp:436 (most recent call first):
```

Repro:
```
# torch version: 2.7.0a0+79aa17489c.nvInternal
# cuda version: 12.9
# nvfuser version: 0.2.27+git802f042
import torch
from nvfuser import FusionDefinition, DataType

def nvfuser_fusion_id10(fd : FusionDefinition) -> None :
    T0 = fd.define_tensor(shape=[1, 2048, 128], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T1 = fd.define_tensor(shape=[1, 8, 2048, 128], contiguity=[None, True, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[3, 2, 1, 0])
    T2 = fd.define_tensor(shape=[1, 32, 2048, 128], contiguity=[None, True, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[3, 2, 1, 0])
    T3 = fd.define_tensor(shape=[1, 2048, 6144], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T4 = fd.define_tensor(shape=[1, 2048, 128], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T5 = fd.define_tensor(shape=[1, 8, 2048, 128], contiguity=[None, True, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[3, 2, 1, 0])
    T6 = fd.define_tensor(shape=[1, 2048, 4096], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T7 = fd.define_tensor(shape=[1, 2048, 4096], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T8 = fd.define_tensor(shape=[1, 2048, 14336], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T9 = fd.define_tensor(shape=[1, 2048, 4096], contiguity=[None, True, True], dtype=DataType.BFloat16, is_cpu=False, stride_order=[2, 1, 0])
    T15 = fd.ops.reshape(T0, new_shape=[1, 1, 2048, 128])
    T31 = fd.ops.slice(T1, start_indices=[0, 0, 0, 0], end_indices=[1, 8, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    T37 = fd.ops.broadcast_in_dim(T15, shape=[1, 8, 2048, 128], broadcast_dims=[0, 1, 2, 3])
    T53 = fd.ops.slice(T2, start_indices=[0, 0, 0, 0], end_indices=[1, 32, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    T59 = fd.ops.broadcast_in_dim(T15, shape=[1, 32, 2048, 128], broadcast_dims=[0, 1, 2, 3])
    T60 = fd.ops.cast(T31, dtype=DataType.Float)
    T61 = fd.ops.cast(T37, dtype=DataType.Float)
    T62 = fd.ops.cast(T53, dtype=DataType.Float)
    T63 = fd.ops.cast(T59, dtype=DataType.Float)
    T76 = fd.ops.slice(T3, start_indices=[0, 0, 0], end_indices=[1, 2048, 4096], strides=[1, 1, 1], manual_normalization=0)
    T89 = fd.ops.slice(T3, start_indices=[0, 0, 4096], end_indices=[1, 2048, 5120], strides=[1, 1, 1], manual_normalization=0)
    T90 = fd.ops.mul(T61, T60)
    T91 = fd.ops.mul(T63, T62)
    T97 = fd.ops.reshape(T76, new_shape=[1, 2048, 32, 128])
    T103 = fd.ops.reshape(T89, new_shape=[1, 2048, 8, 128])
    T104 = fd.ops.cast(T90, dtype=DataType.BFloat16)
    T105 = fd.ops.cast(T91, dtype=DataType.BFloat16)
    T106 = fd.ops.permute(T97, dims=[0, 2, 1, 3])
    T107 = fd.ops.permute(T103, dims=[0, 2, 1, 3])
    T123 = fd.ops.slice(T104, start_indices=[0, 0, 0, 0], end_indices=[1, 8, 2048, 64], strides=[1, 1, 1, 1], manual_normalization=0)
    T139 = fd.ops.slice(T105, start_indices=[0, 0, 0, 0], end_indices=[1, 32, 2048, 64], strides=[1, 1, 1, 1], manual_normalization=0)
    T155 = fd.ops.slice(T106, start_indices=[0, 0, 0, 64], end_indices=[1, 32, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    T171 = fd.ops.slice(T107, start_indices=[0, 0, 0, 64], end_indices=[1, 8, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    T172 = fd.ops.cast(T123, dtype=DataType.Float)
    T178 = fd.ops.reshape(T4, new_shape=[1, 1, 2048, 128])
    T179 = fd.ops.cast(T139, dtype=DataType.Float)
    T180 = fd.ops.cast(T155, dtype=DataType.Float)
    T181 = fd.ops.cast(T171, dtype=DataType.Float)
    T182 = fd.ops.neg(T172)
    T188 = fd.ops.broadcast_in_dim(T178, shape=[1, 8, 2048, 128], broadcast_dims=[0, 1, 2, 3])
    S189 = fd.define_scalar(0, dtype=DataType.Int)
    T195 = fd.ops.full(shape=[1, 8, 2048, 0], fill_value=S189, dtype=DataType.BFloat16)
    T196 = fd.ops.neg(T179)
    T202 = fd.ops.broadcast_in_dim(T178, shape=[1, 32, 2048, 128], broadcast_dims=[0, 1, 2, 3])
    S203 = fd.define_scalar(0, dtype=DataType.Int)
    T209 = fd.ops.full(shape=[1, 32, 2048, 0], fill_value=S203, dtype=DataType.BFloat16)
    T210 = fd.ops.neg(T180)
    T211 = fd.ops.neg(T181)
    T212 = fd.ops.cast(T182, dtype=DataType.BFloat16)
    T213 = fd.ops.cast(T188, dtype=DataType.Float)
    S214 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T224 = fd.ops.pad(T195, [0, 128, 0, 0, 0, 0, 0, 0], S214)
    T225 = fd.ops.cast(T196, dtype=DataType.BFloat16)
    T226 = fd.ops.cast(T202, dtype=DataType.Float)
    S227 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T237 = fd.ops.pad(T209, [0, 128, 0, 0, 0, 0, 0, 0], S227)
    T253 = fd.ops.slice(T106, start_indices=[0, 0, 0, 0], end_indices=[1, 32, 2048, 64], strides=[1, 1, 1, 1], manual_normalization=0)
    T254 = fd.ops.cast(T210, dtype=DataType.BFloat16)
    T270 = fd.ops.slice(T107, start_indices=[0, 0, 0, 0], end_indices=[1, 8, 2048, 64], strides=[1, 1, 1, 1], manual_normalization=0)
    T271 = fd.ops.cast(T211, dtype=DataType.BFloat16)
    T287 = fd.ops.slice(T104, start_indices=[0, 0, 0, 64], end_indices=[1, 8, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    S288 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T298 = fd.ops.pad(T212, [64, 0, 0, 0, 0, 0, 0, 0], S288)
    T299 = fd.ops.mul(T213, T60)
    T300 = fd.ops.cast(T224, dtype=DataType.Float)
    T316 = fd.ops.slice(T105, start_indices=[0, 0, 0, 64], end_indices=[1, 32, 2048, 128], strides=[1, 1, 1, 1], manual_normalization=0)
    S317 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T327 = fd.ops.pad(T225, [64, 0, 0, 0, 0, 0, 0, 0], S317)
    T328 = fd.ops.mul(T226, T62)
    T329 = fd.ops.cast(T237, dtype=DataType.Float)
    T330 = fd.ops.cat([T254, T253], dim=-1, manual_padding=0)
    T331 = fd.ops.cat([T271, T270], dim=-1, manual_padding=0)
    S332 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T342 = fd.ops.pad(T287, [0, 64, 0, 0, 0, 0, 0, 0], S332)
    T343 = fd.ops.cast(T298, dtype=DataType.Float)
    T344 = fd.ops.add(T300, T299)
    S345 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T355 = fd.ops.pad(T316, [0, 64, 0, 0, 0, 0, 0, 0], S345)
    T356 = fd.ops.cast(T327, dtype=DataType.Float)
    T357 = fd.ops.add(T329, T328)
    T358 = fd.ops.cast(T106, dtype=DataType.Float)
    T359 = fd.ops.cast(T107, dtype=DataType.Float)
    T360 = fd.ops.cast(T330, dtype=DataType.Float)
    T361 = fd.ops.cast(T331, dtype=DataType.Float)
    T362 = fd.ops.cast(T342, dtype=DataType.Float)
    T363 = fd.ops.add(T344, T343)
    T364 = fd.ops.cast(T355, dtype=DataType.Float)
    T365 = fd.ops.add(T357, T356)
    T366 = fd.ops.mul(T358, T62)
    T367 = fd.ops.mul(T359, T60)
    T368 = fd.ops.mul(T360, T62)
    T369 = fd.ops.mul(T361, T60)
    T370 = fd.ops.add(T363, T362)
    T371 = fd.ops.add(T365, T364)
    T372 = fd.ops.sum(T366, dims=[0, 1], keepdim=False, dtype=DataType.Null)
    T373 = fd.ops.sum(T367, dims=[0, 1], keepdim=False, dtype=DataType.Null)
    T374 = fd.ops.sum(T368, dims=[0, 1], keepdim=False, dtype=DataType.Null)
    T375 = fd.ops.sum(T369, dims=[0, 1], keepdim=False, dtype=DataType.Null)
    T376 = fd.ops.cast(T370, dtype=DataType.BFloat16)
    T377 = fd.ops.cast(T371, dtype=DataType.BFloat16)
    T378 = fd.ops.cast(T372, dtype=DataType.BFloat16)
    T379 = fd.ops.cast(T373, dtype=DataType.BFloat16)
    T380 = fd.ops.cast(T374, dtype=DataType.BFloat16)
    T381 = fd.ops.cast(T375, dtype=DataType.BFloat16)
    T382 = fd.ops.permute(T5, dims=[0, 2, 1, 3])
    T383 = fd.ops.permute(T376, dims=[0, 2, 1, 3])
    T384 = fd.ops.permute(T377, dims=[0, 2, 1, 3])
    T390 = fd.ops.broadcast_in_dim(T378, shape=[1, 1, 2048, 128], broadcast_dims=[2, 3])
    T396 = fd.ops.broadcast_in_dim(T379, shape=[1, 1, 2048, 128], broadcast_dims=[2, 3])
    T402 = fd.ops.broadcast_in_dim(T380, shape=[1, 1, 2048, 128], broadcast_dims=[2, 3])
    T408 = fd.ops.broadcast_in_dim(T381, shape=[1, 1, 2048, 128], broadcast_dims=[2, 3])
    T413 = fd.ops.reshape(T382, new_shape=[1, 2048, 1024])
    T418 = fd.ops.reshape(T383, new_shape=[1, 2048, 1024])
    T423 = fd.ops.reshape(T384, new_shape=[1, 2048, 4096])
    T428 = fd.ops.reshape(T390, new_shape=[1, 2048, 128])
    T433 = fd.ops.reshape(T396, new_shape=[1, 2048, 128])
    T438 = fd.ops.reshape(T402, new_shape=[1, 2048, 128])
    T443 = fd.ops.reshape(T408, new_shape=[1, 2048, 128])
    T444 = fd.ops.cat([T423, T418, T413], dim=-1, manual_padding=0)
    T445 = fd.ops.cast(T428, dtype=DataType.Float)
    T446 = fd.ops.cast(T433, dtype=DataType.Float)
    T447 = fd.ops.cast(T438, dtype=DataType.Float)
    T448 = fd.ops.cast(T443, dtype=DataType.Float)
    T452 = fd.ops.reshape(T444, new_shape=[2048, 6144])
    T453 = fd.ops.add(T446, T445)
    T454 = fd.ops.add(T448, T447)
    T458 = fd.ops.reshape(T6, new_shape=[2048, 4096])
    T462 = fd.ops.reshape(T7, new_shape=[2048, 4096])
    T466 = fd.ops.reshape(T8, new_shape=[2048, 14336])
    T470 = fd.ops.reshape(T9, new_shape=[2048, 4096])
    T471 = fd.ops.permute(T452, dims=[1, 0])
    T472 = fd.ops.cast(T453, dtype=DataType.BFloat16)
    T473 = fd.ops.cast(T454, dtype=DataType.BFloat16)
    fd.add_output(T473)
    fd.add_output(T472)
    fd.add_output(T452)
    fd.add_output(T471)
    fd.add_output(T470)
    fd.add_output(T466)
    fd.add_output(T462)
    fd.add_output(T458)

with FusionDefinition() as fd:
    nvfuser_fusion_id10(fd)

inputs = [
    torch.testing.make_tensor((1, 2048, 128), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 8, 2048, 128), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 32, 2048, 128), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 6144), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 128), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 8, 2048, 128), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 4096), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 4096), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 14336), dtype=torch.bfloat16, device='cuda:0'),
    torch.testing.make_tensor((1, 2048, 4096), dtype=torch.bfloat16, device='cuda:0'),
]
fd.execute(inputs)
```

Diff:
diff --git a/csrc/scheduler/resize.cpp b/csrc/scheduler/resize.cpp
index 59bf84319f7..92a02fb59ac 100644
--- a/csrc/scheduler/resize.cpp
+++ b/csrc/scheduler/resize.cpp
@@ -245,6 +245,78 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Before propagating loop transformations of the ref tensor, all
+// tensors need their loop domains to be reachable from the ref by
+// just a backward traversal in the exact graph. This is to make
+// tranformations to be consistently propagated through resize-induced
+// cycles in the exact graph.
+//
+// There may be tensors that have further transformations that do not
+// exist in the reference tensor, e.g., a reshape of the
+// reference. Here, we transform those tensors so that they become
+// ready to get the reference transformations propagated to.
+void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
+  Fusion* fusion = ref_tv->fusion();
+
+  // Tensors that appear before the ref should not need anything here
+  // as they should already be ready to get propagated from the ref
+  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
+  const std::unordered_set<Statement*> all_dep_stmt_set{
+      all_dep_stmts.begin(), all_dep_stmts.end()};
+
+  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
+        return all_dep_stmt_set.contains(tv);
+      })) {
+    return;
+  }
+
+  IdModel id_model(fusion);
+  const auto& graph = id_model.buildBroadcastGraph();
+
+  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());
+
+  std::vector<TensorView*> tvs_with_extra_transforms;
+
+  // Look for tensor whose logical domain is not rechable from ref by
+  // just a backward traversal
+  for (auto tv : fusion->allTvs()) {
+    if (all_dep_stmt_set.contains(tv)) {
+      continue;
+    }
+
+    const auto tv_logical =
+        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));
+
+    auto all_visited = ValGraphBFS::getExprGroupsBetween(
+                           graph,
+                           ref_logical,
+                           tv_logical,
+                           /*require_all_to_visited=*/false,
+                           Direction::Backward)
+                           .second;
+
+    if (all_visited) {
+      continue;
+    }
+
+    // This tensor is unreachable by just traversing backward
+    tvs_with_extra_transforms.push_back(tv);
+  }
+
+  if (tvs_with_extra_transforms.empty()) {
+    return;
+  }
+
+  // Make their loop domains look like the ref so that loop
+  // transformations of the ref can be propagated
+  scheduler_tools::scheduleLoopDomainsLike(
+      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
+}
+
+} // namespace
+
 void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   FUSER_PERF_SCOPE("ResizeScheduler::schedule");
 
@@ -254,12 +326,12 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
 
   scheduler_utils::clearMemorySpace(fusion);
 
-  auto ref_tv = getReferenceTensor(fusion);
-  NVF_ERROR(ref_tv != nullptr);
-
   scheduler_utils::cacheInputs(fusion, true);
   scheduler_utils::cacheAndForkOutputs(fusion, true);
 
+  auto ref_tv = getReferenceTensor(fusion);
+  NVF_ERROR(ref_tv != nullptr);
+
   auto resize_tensor_ops = ir_utils::getOpsOfType<SliceOp, PadOp>(fusion);
 
   std::unique_ptr<IdModel> id_model =
@@ -298,11 +370,12 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
         largest_input, /*skip_innermost_id=*/true);
   }
 
+  // Propagate Resize ops to producer tensors. This is safe as this
+  // scheduler only accepts pointwise ops with no reduction.
   for (auto expr : fusion->exprs()) {
-    if (!expr->isOneOf<SliceOp, PadOp>()) {
+    if (!scheduler_tools::isResizeBasedOp(expr)) {
       continue;
     }
-
     scheduler_tools::propagateResizeToInputs(expr);
   }
 
@@ -426,12 +499,13 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // [..., I0/bdimx/max_gdimx, max_gdimx(BIDx), bdimx(TIDx), vec_factor] or
   // [..., I0/bdimx(BIDx), bdimx(TIDx), vec_factor]
 
+  // Before propagating the reference transformations, make sure all
+  // tensors are ready to get transformations from the ref
+  prepareForBackwardTransformPropagation(ref_tv);
+
   // Propagate the reference to the other tensors. Note that the
   // update flag is enabled to workaround the resize propagation
-  // issue. This may not work if there's a tensor that is reshaped
-  // from the reference tensor, but that should not be the case as the
-  // reference is picked by the same routine used for the pointwise
-  // scheduler.
+  // issue.
   //
   // When an ending static repeat is detected and the repeat ID is
   // moved to the outermost position, propagation is done separately
diff --git a/tests/cpp/test_resize.cpp b/tests/cpp/test_resize.cpp
index 4aaac477a52..df505290230 100644
--- a/tests/cpp/test_resize.cpp
+++ b/tests/cpp/test_resize.cpp
@@ -6129,4 +6129,50 @@ TEST_F(ResizeTest, VectorizeOuterPad) {
   EXPECT_EQ((*vec_id_it)->extent()->evaluate(), 8);
 }
 
+// Repro of issue #4250
+TEST_F(ResizeTest, ReshapeAfterRef) {
+  auto fusion_ptr = std::make_unique<Fusion>();
+  auto& fusion = *fusion_ptr;
+  FusionGuard fg(fusion_ptr.get());
+
+  std::vector<int64_t> shape({2, 16, 100});
+
+  EnableOptionsGuard enable_options_guard;
+  EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});
+
+  auto tv0 = makeConcreteTensor(shape);
+  fusion.addInput(tv0);
+  auto tv1 = sin(tv0);
+  auto tv2 = slice(
+      tv1,
+      {{fusion.zeroVal(), tv1->getLogicalDomain().at(0)->extent()},
+       {fusion.zeroVal(), tv1->getLogicalDomain().at(1)->extent()},
+       {fusion.zeroVal(), IrBuilder::create<Val>(shape[2] / 2)}});
+  auto tv3 = sin(tv0);
+  auto tv4 = slice(
+      tv3,
+      {{fusion.zeroVal(), tv3->getLogicalDomain().at(0)->extent()},
+       {fusion.zeroVal(), tv3->getLogicalDomain().at(1)->extent()},
+       {IrBuilder::create<Val>(shape[2] / 2),
+        IrBuilder::create<Val>(shape[2])}});
+  auto tv5 = cat({tv4, tv2}, 2);
+  auto tv6 = add(tv0, tv5);
+  fusion.addOutput(tv6);
+
+  auto tv7 = transpose(tv6, 0, 1);
+  auto tv8 = reshape(tv7, {IrBuilder::create<Val>(-1)});
+  fusion.addOutput(tv8);
+
+  // tv6 will be picked as the reference. The resize scheduler will
+  // try to update tv7 and t8 by propagating transformations from
+  // tv6. Before doing so, the loop domain of tv8 needs to be updated
+  // to match with tv6.
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn(shape, options);
+
+  auto outputs = scheduleAndRun(&fusion, SchedulerType::Resize, {t0});
+  testValidate(&fusion, outputs.outputs, {t0}, __LINE__, __FILE__);
+}
+
 } // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T20:13:56Z
user: login: jjsjann123
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T17:26:11Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T01:02:02Z
user: login: jjsjann123
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-30T09:36:22Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T16:47:18Z

Review comments:
diff_hunk: @@ -254,12 +326,12 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
 
   scheduler_utils::clearMemorySpace(fusion);
 
-  auto ref_tv = getReferenceTensor(fusion);
-  NVF_ERROR(ref_tv != nullptr);
-
   scheduler_utils::cacheInputs(fusion, true);
   scheduler_utils::cacheAndForkOutputs(fusion, true);
 
+  auto ref_tv = getReferenceTensor(fusion);
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: This change is not strictly required, but just to pick a reference from fusion outputs. `cacheAndForkOutputs` may change output tensors.
created_at: 2025-04-16T20:13:56Z
author_association: COLLABORATOR
diff_hunk: @@ -419,6 +491,10 @@ void ResizeScheduler::schedule(Fusion* fusion, const HeuristicParams* params) {
   // [..., I0/bdimx/max_gdimx, max_gdimx(BIDx), bdimx(TIDx), vec_factor] or
   // [..., I0/bdimx(BIDx), bdimx(TIDx), vec_factor]
 
+  // Before propagating the reference transformations, make sure all
+  // tensors are ready to get transformations from the ref
+  prepareForBackwardTransformPropagation(ref_tv);
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: > // This may not work if there's a tensor that is reshaped
  // from the reference tensor, but that should not be the case as the
  // reference is picked by the same routine used for the pointwise
  // scheduler.

nitpick, the comment in the new line 500 - 503 reads a bit strange in the context of this PR.
Sounds like it is indeed happening and we are patching that now. Should we remove that comment?
created_at: 2025-04-17T17:17:41Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,78 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Before propagating loop transformations of the ref tensor, all
+// tensors need their loop domains to be reachable from the ref by
+// just a backward traversal in the exact graph. This is to make
+// tranformations to be consistently propagated through resize-induced
+// cycles in the exact graph.
+//
+// There may be tensors that have further transformations that do not
+// exist in the reference tensor, e.g., a reshape of the
+// reference. Here, we transform those tensors so that they become
+// ready to get the reference transformations propagated to.
+void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
+  Fusion* fusion = ref_tv->fusion();
+
+  // Tensors that appear before the ref should not need anything here
+  // as they should already be ready to get propagated from the ref
+  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
+  const std::unordered_set<Statement*> all_dep_stmt_set{
+      all_dep_stmts.begin(), all_dep_stmts.end()};
+
+  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
+        return all_dep_stmt_set.contains(tv);
+      })) {
+    return;
+  }
+
+  IdModel id_model(fusion);
+  const auto& graph = id_model.buildBroadcastGraph();
+
+  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());
+
+  std::vector<TensorView*> tvs_with_extra_transforms;
+
+  // Look for tensor whose logical domain is not rechable from ref by
+  // just a backward traversal
+  for (auto tv : fusion->allTvs()) {
+    if (all_dep_stmt_set.contains(tv)) {
+      continue;
+    }
+
+    const auto tv_logical =
+        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));
+
+    auto all_visited = ValGraphBFS::getExprGroupsBetween(
+                           graph,
+                           ref_logical,
+                           tv_logical,
+                           /*require_all_to_visited=*/false,
+                           Direction::Backward)
+                           .second;
+
+    if (all_visited) {
+      continue;
+    }
+
+    // This tensor is unreachable by just traversing backward
+    tvs_with_extra_transforms.push_back(tv);
+  }
+
+  if (tvs_with_extra_transforms.empty()) {
+    return;
+  }
+
+  // Make their loop domains look like the ref so that loop
+  // transformations of the ref can be propagated
+  scheduler_tools::scheduleLoopDomainsLike(
+      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: Out of curiosity. How do we know it's a safe to overwrite the loop domain here?

i.e. if the overwrite loop domain from ref_tv is *smaller* than the actual tv, we are effectively applying a false predicate here, which might lead to silent wrong results?

I'm assuming the resize scheduler has a check somewhere that ensures the coverage from ref to outputs?
created_at: 2025-04-17T17:26:04Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,78 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Before propagating loop transformations of the ref tensor, all
+// tensors need their loop domains to be reachable from the ref by
+// just a backward traversal in the exact graph. This is to make
+// tranformations to be consistently propagated through resize-induced
+// cycles in the exact graph.
+//
+// There may be tensors that have further transformations that do not
+// exist in the reference tensor, e.g., a reshape of the
+// reference. Here, we transform those tensors so that they become
+// ready to get the reference transformations propagated to.
+void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
+  Fusion* fusion = ref_tv->fusion();
+
+  // Tensors that appear before the ref should not need anything here
+  // as they should already be ready to get propagated from the ref
+  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
+  const std::unordered_set<Statement*> all_dep_stmt_set{
+      all_dep_stmts.begin(), all_dep_stmts.end()};
+
+  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
+        return all_dep_stmt_set.contains(tv);
+      })) {
+    return;
+  }
+
+  IdModel id_model(fusion);
+  const auto& graph = id_model.buildBroadcastGraph();
+
+  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());
+
+  std::vector<TensorView*> tvs_with_extra_transforms;
+
+  // Look for tensor whose logical domain is not rechable from ref by
+  // just a backward traversal
+  for (auto tv : fusion->allTvs()) {
+    if (all_dep_stmt_set.contains(tv)) {
+      continue;
+    }
+
+    const auto tv_logical =
+        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));
+
+    auto all_visited = ValGraphBFS::getExprGroupsBetween(
+                           graph,
+                           ref_logical,
+                           tv_logical,
+                           /*require_all_to_visited=*/false,
+                           Direction::Backward)
+                           .second;
+
+    if (all_visited) {
+      continue;
+    }
+
+    // This tensor is unreachable by just traversing backward
+    tvs_with_extra_transforms.push_back(tv);
+  }
+
+  if (tvs_with_extra_transforms.empty()) {
+    return;
+  }
+
+  // Make their loop domains look like the ref so that loop
+  // transformations of the ref can be propagated
+  scheduler_tools::scheduleLoopDomainsLike(
+      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: That's a good question. Since this is only for pointwise ops with no reduction, as long as the loop domain of the reference tensor is kept intact, all tensors should be also covered. If not, that seems to indicate the reference is not really a right reference. We found a couple of issues with finding right references, but as far as I can see, a reference tensor chosen by the current logic should guarantee all other tensors can be transformed to the same shape as the reference without losing any dimension.

That won't be the case with reductions, so that's one big difference when we extend this scheduling scheme to handle reductions.
created_at: 2025-04-30T01:02:02Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,78 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Before propagating loop transformations of the ref tensor, all
+// tensors need their loop domains to be reachable from the ref by
+// just a backward traversal in the exact graph. This is to make
+// tranformations to be consistently propagated through resize-induced
+// cycles in the exact graph.
+//
+// There may be tensors that have further transformations that do not
+// exist in the reference tensor, e.g., a reshape of the
+// reference. Here, we transform those tensors so that they become
+// ready to get the reference transformations propagated to.
+void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
+  Fusion* fusion = ref_tv->fusion();
+
+  // Tensors that appear before the ref should not need anything here
+  // as they should already be ready to get propagated from the ref
+  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
+  const std::unordered_set<Statement*> all_dep_stmt_set{
+      all_dep_stmts.begin(), all_dep_stmts.end()};
+
+  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
+        return all_dep_stmt_set.contains(tv);
+      })) {
+    return;
+  }
+
+  IdModel id_model(fusion);
+  const auto& graph = id_model.buildBroadcastGraph();
+
+  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());
+
+  std::vector<TensorView*> tvs_with_extra_transforms;
+
+  // Look for tensor whose logical domain is not rechable from ref by
+  // just a backward traversal
+  for (auto tv : fusion->allTvs()) {
+    if (all_dep_stmt_set.contains(tv)) {
+      continue;
+    }
+
+    const auto tv_logical =
+        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));
+
+    auto all_visited = ValGraphBFS::getExprGroupsBetween(
+                           graph,
+                           ref_logical,
+                           tv_logical,
+                           /*require_all_to_visited=*/false,
+                           Direction::Backward)
+                           .second;
+
+    if (all_visited) {
+      continue;
+    }
+
+    // This tensor is unreachable by just traversing backward
+    tvs_with_extra_transforms.push_back(tv);
+  }
+
+  if (tvs_with_extra_transforms.empty()) {
+    return;
+  }
+
+  // Make their loop domains look like the ref so that loop
+  // transformations of the ref can be propagated
+  scheduler_tools::scheduleLoopDomainsLike(
+      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
path: csrc/scheduler/resize.cpp
user: login: jjsjann123
body: nitpick: I wonder if we should add a check looking for reduction op and the reasoning as a comment in this function?
Not pushing strongly for that though.
created_at: 2025-04-30T09:36:15Z
author_association: COLLABORATOR
diff_hunk: @@ -245,6 +245,78 @@ std::unique_ptr<HeuristicParams> ResizeScheduler::computeHeuristics(
   return params;
 }
 
+namespace {
+
+// Before propagating loop transformations of the ref tensor, all
+// tensors need their loop domains to be reachable from the ref by
+// just a backward traversal in the exact graph. This is to make
+// tranformations to be consistently propagated through resize-induced
+// cycles in the exact graph.
+//
+// There may be tensors that have further transformations that do not
+// exist in the reference tensor, e.g., a reshape of the
+// reference. Here, we transform those tensors so that they become
+// ready to get the reference transformations propagated to.
+void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
+  Fusion* fusion = ref_tv->fusion();
+
+  // Tensors that appear before the ref should not need anything here
+  // as they should already be ready to get propagated from the ref
+  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
+  const std::unordered_set<Statement*> all_dep_stmt_set{
+      all_dep_stmts.begin(), all_dep_stmts.end()};
+
+  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
+        return all_dep_stmt_set.contains(tv);
+      })) {
+    return;
+  }
+
+  IdModel id_model(fusion);
+  const auto& graph = id_model.buildBroadcastGraph();
+
+  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());
+
+  std::vector<TensorView*> tvs_with_extra_transforms;
+
+  // Look for tensor whose logical domain is not rechable from ref by
+  // just a backward traversal
+  for (auto tv : fusion->allTvs()) {
+    if (all_dep_stmt_set.contains(tv)) {
+      continue;
+    }
+
+    const auto tv_logical =
+        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));
+
+    auto all_visited = ValGraphBFS::getExprGroupsBetween(
+                           graph,
+                           ref_logical,
+                           tv_logical,
+                           /*require_all_to_visited=*/false,
+                           Direction::Backward)
+                           .second;
+
+    if (all_visited) {
+      continue;
+    }
+
+    // This tensor is unreachable by just traversing backward
+    tvs_with_extra_transforms.push_back(tv);
+  }
+
+  if (tvs_with_extra_transforms.empty()) {
+    return;
+  }
+
+  // Make their loop domains look like the ref so that loop
+  // transformations of the ref can be propagated
+  scheduler_tools::scheduleLoopDomainsLike(
+      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
path: csrc/scheduler/resize.cpp
user: login: naoyam
body: Thanks for the suggestion. Added a comment (not here, but the first place where potentially shrinking propagations happen).
created_at: 2025-04-30T16:47:18Z
author_association: COLLABORATOR

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

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


### **Description**
- Added function to prepare tensors for backward transform propagation

- Updated schedule function to call prepareForBackwardTransformPropagation

- Added test case for reshape after reference tensor


___



### **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 Resize Scheduler for Reshape Handling</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/resize.cpp

<li>Added <code>prepareForBackwardTransformPropagation</code> function to handle <br>tensors with extra transformations<br> <li> Updated <code>schedule</code> function to call <br><code>prepareForBackwardTransformPropagation</code> before propagating <br>transformations<br> <li> Modified loop to check for resize-based operations instead of specific <br>types


</details>


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

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_resize.cpp</strong><dd><code>Add Test Case for Reshape After Reference</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

tests/cpp/test_resize.cpp

- Added test case `ReshapeAfterRef` to reproduce issue #4250


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4259/files#diff-98e2e8f5cb3849584331987b87d7cc114cb88cfddddd2f739ca8b452d92f4ece">+46/-0</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/4259/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427aR260-R316'><strong>Logic Verification</strong></a>

Ensure that the logic in `prepareForBackwardTransformPropagation` correctly identifies and handles tensors with extra transformations, especially those involving reshapes.</summary>

```c++
void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
  Fusion* fusion = ref_tv->fusion();

  // Tensors that appear before the ref should not need anything here
  // as they should already be ready to get propagated from the ref
  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
  const std::unordered_set<Statement*> all_dep_stmt_set{
      all_dep_stmts.begin(), all_dep_stmts.end()};

  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
        return all_dep_stmt_set.contains(tv);
      })) {
    return;
  }

  IdModel id_model(fusion);
  const auto& graph = id_model.buildBroadcastGraph();

  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());

  std::vector<TensorView*> tvs_with_extra_transforms;

  // Look for tensor whose logical domain is not rechable from ref by
  // just a backward traversal
  for (auto tv : fusion->allTvs()) {
    if (all_dep_stmt_set.contains(tv)) {
      continue;
    }

    const auto tv_logical =
        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));

    auto all_visited = ValGraphBFS::getExprGroupsBetween(
                           graph,
                           ref_logical,
                           tv_logical,
                           /*require_all_to_visited=*/false,
                           Direction::Backward)
                           .second;

    if (all_visited) {
      continue;
    }

    // This tensor is unreachable by just traversing backward
    tvs_with_extra_transforms.push_back(tv);
  }

  if (tvs_with_extra_transforms.empty()) {
    return;
  }

  // Make their loop domains look like the ref so that loop
  // transformations of the ref can be propagated
  scheduler_tools::scheduleLoopDomainsLike(
      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4259/files#diff-a8e9f807f9091349106134e438b7e81e3543c83c196ac75a73e1095d057b427aR260-R316'><strong>Performance Impact</strong></a>

Evaluate the performance impact of the changes introduced in `prepareForBackwardTransformPropagation` to ensure that the performance goals are met.</summary>

```c++
void prepareForBackwardTransformPropagation(TensorView* ref_tv) {
  Fusion* fusion = ref_tv->fusion();

  // Tensors that appear before the ref should not need anything here
  // as they should already be ready to get propagated from the ref
  const auto all_dep_stmts = StmtSort::getStmtsTo({ref_tv});
  const std::unordered_set<Statement*> all_dep_stmt_set{
      all_dep_stmts.begin(), all_dep_stmts.end()};

  if (std::ranges::all_of(fusion->allTvs(), [&](TensorView* tv) {
        return all_dep_stmt_set.contains(tv);
      })) {
    return;
  }

  IdModel id_model(fusion);
  const auto& graph = id_model.buildBroadcastGraph();

  const auto ref_logical = graph.toGroups(ref_tv->getLogicalDomain());

  std::vector<TensorView*> tvs_with_extra_transforms;

  // Look for tensor whose logical domain is not rechable from ref by
  // just a backward traversal
  for (auto tv : fusion->allTvs()) {
    if (all_dep_stmt_set.contains(tv)) {
      continue;
    }

    const auto tv_logical =
        graph.toGroups(TensorDomain::noBroadcasts(tv->getLogicalDomain()));

    auto all_visited = ValGraphBFS::getExprGroupsBetween(
                           graph,
                           ref_logical,
                           tv_logical,
                           /*require_all_to_visited=*/false,
                           Direction::Backward)
                           .second;

    if (all_visited) {
      continue;
    }

    // This tensor is unreachable by just traversing backward
    tvs_with_extra_transforms.push_back(tv);
  }

  if (tvs_with_extra_transforms.empty()) {
    return;
  }

  // Make their loop domains look like the ref so that loop
  // transformations of the ref can be propagated
  scheduler_tools::scheduleLoopDomainsLike(
      tvs_with_extra_transforms, ref_tv->getLogicalDomain());
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4259/files#diff-98e2e8f5cb3849584331987b87d7cc114cb88cfddddd2f739ca8b452d92f4eceR6132-R6176'><strong>Test Coverage</strong></a>

Verify that the added test case `ReshapeAfterRef` adequately covers the scenario described in issue #4250 and that it effectively tests the changes made in the PR.</summary>

```c++
// Repro of issue #4250
TEST_F(ResizeTest, ReshapeAfterRef) {
  auto fusion_ptr = std::make_unique<Fusion>();
  auto& fusion = *fusion_ptr;
  FusionGuard fg(fusion_ptr.get());

  std::vector<int64_t> shape({2, 16, 100});

  EnableOptionsGuard enable_options_guard;
  EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});

  auto tv0 = makeConcreteTensor(shape);
  fusion.addInput(tv0);
  auto tv1 = sin(tv0);
  auto tv2 = slice(
      tv1,
      {{fusion.zeroVal(), tv1->getLogicalDomain().at(0)->extent()},
       {fusion.zeroVal(), tv1->getLogicalDomain().at(1)->extent()},
       {fusion.zeroVal(), IrBuilder::create<Val>(shape[2] / 2)}});
  auto tv3 = sin(tv0);
  auto tv4 = slice(
      tv3,
      {{fusion.zeroVal(), tv3->getLogicalDomain().at(0)->extent()},
       {fusion.zeroVal(), tv3->getLogicalDomain().at(1)->extent()},
       {IrBuilder::create<Val>(shape[2] / 2),
        IrBuilder::create<Val>(shape[2])}});
  auto tv5 = cat({tv4, tv2}, 2);
  auto tv6 = add(tv0, tv5);
  fusion.addOutput(tv6);

  auto tv7 = transpose(tv6, 0, 1);
  auto tv8 = reshape(tv7, {IrBuilder::create<Val>(-1)});
  fusion.addOutput(tv8);

  // tv6 will be picked as the reference. The resize scheduler will
  // try to update tv7 and t8 by propagating transformations from
  // tv6. Before doing so, the loop domain of tv8 needs to be updated
  // to match with tv6.

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

  auto outputs = scheduleAndRun(&fusion, SchedulerType::Resize, {t0});
  testValidate(&fusion, outputs.outputs, {t0}, __LINE__, __FILE__);
}
```

</details>

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

author_association: NONE
user: login: naoyam
created_at: 2025-04-16T20:12:24Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-30T01:15:37Z
body: !test
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-30T16:47:27Z
body: !build
author_association: COLLABORATOR