Some further cleanup for the new computeAt interface (#1793)

Revert MaxProducerPosUpdater to old algo.

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

@@ -165,7 +165,7 @@ void ComputeAt::runAt(
     TensorView* consumer,
     unsigned int consumer_position,
     ComputeAtMode mode) {
-  FUSER_PERF_SCOPE("ComputeAt::run");
+  FUSER_PERF_SCOPE("ComputeAt::runAt");
 
   // Make sure the correct fusion is setup between this and consumer.
   TORCH_CHECK(
@@ -175,6 +175,10 @@ void ComputeAt::runAt(
       consumer,
       " are not in the same fusion.");
 
+  if (mode == ComputeAtMode::MostInlined) {
+    consumer_position = consumer->nDims();
+  }
+
   FusionGuard fg(producer->fusion());
 
   auto selected = getPropagationSubgraph(producer, consumer);
@@ -206,6 +210,10 @@ void ComputeAt::runWith(
       consumer,
       " are not in the same fusion.");
 
+  if (mode == ComputeAtMode::MostInlined) {
+    producer_position = producer->nDims();
+  }
+
   FusionGuard fg(producer->fusion());
 
   auto selected = getPropagationSubgraph(producer, consumer);

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

@@ -309,9 +309,15 @@ void InlinePropagator::propagateTvPasC(TensorView* from, TensorView* to) {
   int pos = getReplayPosPasC(to, from);
   auto to_pos =
       TransformReplay::getMatchedLeafPosWithoutReplayPasC(to, from, pos);
-  // TODO: Can we make TransformPropagator do the transformation, and
-  // InlinePropagator only set the CA positions?
-  //   TORCH_CHECK(to_pos >= 0);
+  if (mode_ != ComputeAtMode::MostInlined) {
+    TORCH_CHECK(
+        to_pos >= 0,
+        "Unable to propagate CA position from consumer ",
+        from,
+        " to producer ",
+        to,
+        " because this would require replay.");
+  }
   if (to_pos < 0) {
     auto replay = TransformReplay::replayPasC(to, from, pos);
     TORCH_INTERNAL_ASSERT(
@@ -335,9 +341,15 @@ void InlinePropagator::propagateTvCasP(TensorView* from, TensorView* to) {
   int pos = getReplayPosCasP(to, from);
   auto to_pos =
       TransformReplay::getMatchedLeafPosWithoutReplayCasP(to, from, pos);
-  // TODO: Can we make TransformPropagator do the transformation, and
-  // InlinePropagator only set the CA positions?
-  //   TORCH_CHECK(to_pos >= 0);
+  if (mode_ != ComputeAtMode::MostInlined) {
+    TORCH_CHECK(
+        to_pos >= 0,
+        "Unable to propagate CA position from producer ",
+        from,
+        " to consumer ",
+        to,
+        " because this would require replay.");
+  }
   if (to_pos < 0) {
     auto replay = TransformReplay::replayCasP(to, from, pos);
     TORCH_INTERNAL_ASSERT(
@@ -373,23 +385,71 @@ void InlinePropagator::propagateTvSibling(TensorView* from, TensorView* to) {
   recordReplayedPos(to, from_pos);
 }
 
+namespace {
+
 // Try to find the aligned position on consumer's domain corresponding to the
-// compute at position of producer domain.
-void MaxProducerPosUpdater::handle(TensorView* consumer) {
+//  compute at position of producer domain. Used in computeAt pass only. No
+//  checking on actual producer-consumer relationship.
+unsigned int getConsumerPosAlignedToProducerCA(
+    TensorView* consumer,
+    TensorView* producer) {
+  // Locate consumer's position that aligns with
+  //  the producer's new compute at axis. We need broadcast axes forwarded so we
+  //  need to replay PasC as CasP will not forward braodcast dims. For example
+  //  if we have:
+  // T2[ iS22{( 3 * 1 )} ] ca_pos( 1 ) = broadcast( T1[ iS1{3} ] ca_pos( 1 )
+  // produce_pos( 1) ) CasP will have the mapping iS1{3} -> iS2{3} and PasC will
+  // have the mapping iS22{( 3 * 1 )} <- iS1{3} We need the latter. Refer to
+  // NVFuserTest.FusionComplexBCast1_CUDA
+
+  auto c2p_map =
+      BestEffortReplay::replayPasC(
+          producer,
+          consumer,
+          -1,
+          // Compute at root domain may not be valid here, as all
+          // producers don't have to be able to map into consumer at
+          // max producer position. Since computeAt should be valid
+          // and this mechanism is only intended to lower produce
+          // position of consumer, we can simply use the pairwise map.
+          PairwiseRootDomainMap(producer, consumer))
+          .getReplay();
+
+  // Find the innermost position of consumer that has
+  //  been mapped within the producer ca axis.
   unsigned int consumer_pos = consumer->nDims();
   while (consumer_pos > 0) {
-    for (auto producer : ir_utils::producerTvsOf(consumer)) {
-      auto producer_pos = TransformReplay::getMatchedLeafPosWithoutReplayPasC(
-          producer, consumer, consumer_pos);
-      if (producer_pos >= 0 &&
-          producer_pos <= producer->getComputeAtPosition()) {
-        goto finished;
-      }
+    auto consumer_id = consumer->axis((int)consumer_pos - 1);
+    auto p_dom = producer->domain()->domain();
+    if (std::any_of(
+            p_dom.begin(),
+            p_dom.begin() + producer->getComputeAtPosition(),
+            [&consumer_id, &c2p_map](IterDomain* p_id) {
+              auto c_id_it = c2p_map.find(consumer_id);
+              if (c_id_it != c2p_map.end()) {
+                return c_id_it->second == p_id;
+              }
+              return false;
+            })) {
+      break;
     }
     consumer_pos--;
   }
-finished:
-  consumer->setMaxProducer(consumer_pos, true);
+
+  return consumer_pos;
+}
+
+} // namespace
+
+// Try to find the aligned position on consumer's domain corresponding to the
+// compute at position of producer domain.
+void MaxProducerPosUpdater::handle(TensorView* consumer) {
+  unsigned int consumer_pos = 0;
+  for (auto producer : ir_utils::producerTvsOf(consumer)) {
+    consumer_pos = std::max(
+        consumer_pos, getConsumerPosAlignedToProducerCA(consumer, producer));
+  }
+  consumer->setMaxProducer(consumer_pos);
 }
 
 void MaxProducerPosUpdater::propagateTvPasC(TensorView* from, TensorView* to) {