save

torch/csrc/jit/codegen/cuda/test/test_gpu.cpp

@@ -23599,8 +23599,6 @@ TEST_F(NVFuserTest, TestTransformPropagatorSibling_CUDA) {
 
   TransformPropagator::from(tvs2.var_sum);
 
-  fusion.print();
-
   // check that the resulting tensors in tvs2 are identical
   auto checkSiblingConsistency = [](TensorView* replay, TensorView* target) {
     auto replay_root = replay->getRootDomain();

torch/csrc/jit/codegen/cuda/transform_replay.cpp

@@ -692,18 +692,17 @@ std::deque<TensorView*> producersFor(TensorView* tv) {
 bool TransformPropagator::replayPasC(
     TensorView* producer_tv,
     TensorView* consumer_tv) {
-  if (producer_tv == starting_tv) {
-    return false;
-  }
-
+  unsigned int pos;
   auto consumer_pos_it = replayed_pos.find(consumer_tv);
   if (consumer_pos_it == replayed_pos.end()) {
-    return false;
+    pos = consumer_tv->nDims();
+  } else {
+    pos = consumer_pos_it->second;
   }
 
   auto pairwiseMap = PairwiseRootDomainMap(producer_tv, consumer_tv);
   auto replayed_producer = TransformReplay::replayPasC(
-      producer_tv, consumer_tv, consumer_pos_it->second, pairwiseMap);
+      producer_tv, consumer_tv, pos, pairwiseMap);
 
   auto producer_root = producer_tv->getMaybeRFactorDomain();
   auto replayed_domain = replayed_producer.first->domain();
@@ -741,18 +740,17 @@ bool TransformPropagator::replayPasC(
 bool TransformPropagator::replayCasP(
     TensorView* consumer_tv,
     TensorView* producer_tv) {
-  if (consumer_tv == starting_tv) {
-    return false;
-  }
-
+  unsigned int pos;
   auto producer_pos_it = replayed_pos.find(producer_tv);
   if (producer_pos_it == replayed_pos.end()) {
-    return false;
+    pos = producer_tv->nDims();
+  } else {
+    pos = producer_pos_it->second;
   }
 
   auto pairwiseMap = PairwiseRootDomainMap(producer_tv, consumer_tv);
   auto replayed_consumer = TransformReplay::replayCasP(
-      consumer_tv, producer_tv, producer_pos_it->second, pairwiseMap);
+      consumer_tv, producer_tv, pos, pairwiseMap);
 
   auto consumer_root = consumer_tv->getRootDomain();
   auto replayed_domain = replayed_consumer.first->domain();
@@ -790,9 +788,6 @@ bool TransformPropagator::replayCasP(
 TransformPropagator::TransformPropagator(TensorView* from) : starting_tv(from) {
   VectorOfUniqueEntries<TensorView*> propagation{starting_tv};
 
-  // Seed position with local tv
-  replayed_pos[from] = from->nDims();
-
   // While tensor views are being replayed, if they're modified, make sure we
   // propagate back to all producers as well as consumers. This is definitely
   // not the most efficient implementation as what we do is any time a tv is
@@ -802,17 +797,21 @@ TransformPropagator::TransformPropagator(TensorView* from) : starting_tv(from) {
 
     // Replay tv forward to its consumers.
     for (auto consumer_tv : consumersOf(tv)) {
+      std::cout << "replayCasP: C=" << consumer_tv << ", P=" << tv << std::endl;
       auto replayed = replayCasP(consumer_tv, tv);
       // If consumer has changed, mark we should propagate
       if (replayed) {
+        std::cout << "Replayed: " << consumer_tv << std::endl;
         propagation.pushBack(consumer_tv);
       }
     }
 
     for (auto producer_tv : producersFor(tv)) {
       // If producer has changed, mark we should propagate
+      std::cout << "replayPasC: P=" << producer_tv << ", C=" << tv << std::endl;
       auto replayed = replayPasC(producer_tv, tv);
       if (replayed) {
+        std::cout << "Replayed: " << producer_tv << std::endl;
         propagation.pushBack(producer_tv);
       }
     }