add SchedulerTopologyChecker::hasIncompatibleReshape

Author: liqiangxl

csrc/scheduler/pointwise.cpp

@@ -879,6 +879,9 @@ void schedulePointwise(Fusion* fusion, const PointwiseParams* pparams) {
 
   if (!ir_utils::getReshapeOps(fusion).empty()) {
     ComputeAtMap ca_map(fusion);
+    std::cout << "fusion before propagateReshapeTransforms" << std::endl;
+    fusion->print();
+    std::cout << ca_map.toString() << std::endl;
     // Propagate reshape transforms through the graph, expecially the reference.
     scheduler_utils::propagateReshapeTransforms(fusion, ca_map);
 

csrc/scheduler/registry.cpp

@@ -97,6 +97,13 @@ bool checkCanSchedule(Fusion* fusion, SchedulerType scheduler_type) {
     return false;
   }
 
+  if (registry_utils::SchedulerTopologyChecker::hasIncompatibleReshape(
+          fusion)) {
+    scheduler_debug_utils::canScheduleRejectReason(
+        scheduler_type, "Fusion has cyclic reshapes.");
+    return false;
+  }
+
   if (registry_utils::SchedulerTopologyChecker::
           rejectScheduleFusionGlobalBufferRequirement(fusion, scheduler_type)) {
     scheduler_debug_utils::canScheduleRejectReason(

csrc/scheduler/registry_utils.cpp

@@ -5,6 +5,7 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 // clang-format on
+#include <ir/builder.h>
 #include <ir/utils.h>
 #include <logical_domain_map.h>
 #include <runtime/executor_kernel_arg.h>
@@ -1170,6 +1171,69 @@ bool SchedulerTopologyChecker::hasCyclicReshape(Fusion* fusion) {
   return false;
 }
 
+namespace {
+
+// Two sequences of constant outer-split factors are considered compatible
+// if they share the same leading prefix up to the length of the shorter one.
+// That is, either they are identical (same length) or the shorter sequence
+// matches the first N elements of the longer one.
+bool areConstExtentSequencesCompatible(
+    const std::vector<int64_t>& a,
+    const std::vector<int64_t>& b) {
+  if (a.size() <= b.size()) {
+    return a == std::vector<int64_t>(b.begin(), b.begin() + a.size());
+  } else {
+    return b == std::vector<int64_t>(a.begin(), a.begin() + b.size());
+  }
+}
+
+} // namespace
+
+// Reshape is expressed as merges and outer splits by constant factors.
+// Merges are replay-safe and do not introduce conflicts during transform
+// propagation. Compatibility therefore hinges on the constant outer-split
+// factors: if the leading prefix of factors matches, the splits can be
+// replayed across reshapes.
+bool SchedulerTopologyChecker::hasIncompatibleReshape(Fusion* fusion) {
+  const auto reshape_ops = ir_utils::getOpsOfType<ReshapeOp>(fusion);
+  if (reshape_ops.size() < 2) {
+    return false;
+  }
+
+  // Collect constant outer-split factors of the reshape output, ordered
+  // from outer-most to inner-most.
+  auto collect_extent_prefix = [](ReshapeOp* reshape) {
+    auto reshape_out = reshape->out()->as<TensorView>();
+    std::vector<int64_t> outer_split_factors;
+    for (auto* id : reshape_out->getLogicalDomain()) {
+      if (auto* def = id->definition()) {
+        if (auto* split = dynamic_cast<Split*>(def)) {
+          if (!split->innerSplit() && split->outer() == id &&
+              split->factor()->isConstInt()) {
+            outer_split_factors.push_back(
+                split->factor()->evaluate().as<int64_t>());
+          }
+        }
+      }
+    }
+
+    return outer_split_factors;
+  };
+
+  for (const auto i : arange(std::ssize(reshape_ops) - 1)) {
+    const auto& out_ids_i_extents = collect_extent_prefix(reshape_ops[i]);
+    for (const auto j : arange(i + 1, std::ssize(reshape_ops))) {
+      const auto& out_ids_j_extents = collect_extent_prefix(reshape_ops[j]);
+      if (!areConstExtentSequencesCompatible(
+              out_ids_i_extents, out_ids_j_extents)) {
+        return true;
+      }
+    }
+  }
+
+  return false;
+}
+
 } // namespace registry_utils
 
 } // namespace nvfuser

csrc/scheduler/registry_utils.h

@@ -123,6 +123,8 @@ class SchedulerTopologyChecker {
   // propagateReshapeTransforms won't work as it won't find any
   // terminating reshape IDs.
   static bool hasCyclicReshape(Fusion* fusion);
+
+  static bool hasIncompatibleReshape(Fusion* fusion);
 };
 
 } // namespace registry_utils

csrc/scheduler/utils.cpp

@@ -2366,6 +2366,12 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
       }
     }
   }
+  std::cout << "transformed_disjoint_sets: " << transformed_disjoint_sets.size()
+            << std::endl;
+  for (auto set : transformed_disjoint_sets) {
+    std::cout << "  transformed_disjoint_sets: " << set->toString()
+              << std::endl;
+  }
 
   std::unordered_set<IterDomain*> terminating_reshape_dims;
   for (const auto& disjoint_set_shared_ptr :
@@ -2390,7 +2396,11 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
       terminating_reshape_dims.emplace(id);
     }
   }
-
+  std::cout << "terminating_reshape_dims: " << terminating_reshape_dims.size()
+            << std::endl;
+  for (auto set : terminating_reshape_dims) {
+    std::cout << "  terminating_reshape_dims: " << set->toString() << std::endl;
+  }
   // If iter domains are involved in any transformation from root domains to
   // logical domains they should be considered "contaminated".
   for (auto tv : fusion->allTvs()) {
@@ -2455,8 +2465,11 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
     }
 
     // Propagate the view transformations
+    std::cout << "before reorder " << tv->toString() << std::endl;
     tv->reorder(old2new);
     //! Propagate current transformations on from_tv to all graphs
+    std::cout << "transformPropagateToAllFrom " << tv->toString() << " @ "
+              << (int64_t)old2new.size() << std::endl;
     transformPropagateToAllFrom(tv, (int64_t)old2new.size());
 
     // Propgating the transforms will not replay the DIDx parallelization, so we

csrc/transform_replay.cpp

@@ -1201,6 +1201,8 @@ void TransformPropagator::propagateP2C(TensorView* from, TensorView* to) {
     debug() << "TransformPropagator::propagateP2C" << std::endl;
     debug() << "  from: " << from << " @ " << pos << std::endl;
     debug() << "  to: " << to << std::endl;
+    std::cout << "  from transforms: " << std::endl;
+    from->printTransforms();
   }
   if (new_pos < 0) {
     auto replay = TransformReplay::replayCasP(
@@ -1255,6 +1257,9 @@ void TransformPropagator::propagateSibling(TensorView* from, TensorView* to) {
 
 TransformPropagator::TransformPropagator(TensorView* from, int64_t pos) {
   replayed_pos_[from] = wrapDim(pos, from->nDims() + 1);
+  std::cout << "TransformPropagator::TransformPropagator" << std::endl;
+  std::cout << "  from: " << from->toString() << " @ " << replayed_pos_[from]
+            << std::endl;
 }
 
 void MostInlinedTransformPropagator::propagateC2P(