validateDomain in TransformPropagator (csarofeen#1796)

shmsong · Jul 1, 2022 · 3f2c263 · 3f2c263
1 parent c077085
commit 3f2c263
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/torch/csrc/jit/codegen/cuda/transform_replay.cpp b/torch/csrc/jit/codegen/cuda/transform_replay.cpp
@@ -835,6 +835,19 @@ bool TransformReplay::fullSelfMatching(
   return true;
 }
 
+namespace {
+
+// Make sure if tv is set to new_td it doesn't violate set compute at and max
+// produce at positions.
+bool validateDomain(TensorView* tv, TensorDomain* new_td) {
+  auto first_mismatch =
+      BestEffortReplay::findFirstMismatchedID(tv->domain(), new_td);
+  return first_mismatch >= (int)tv->getMaxProducerPosition() &&
+      first_mismatch >= (int)tv->getComputeAtPosition();
+}
+
+} // namespace
+
 void TransformPropagator::propagateTvPasC(TensorView* from, TensorView* to) {
   int pos = replayed_pos_.at(from);
   // Note: [Using multiple TransformPropagators]
@@ -849,6 +862,13 @@ void TransformPropagator::propagateTvPasC(TensorView* from, TensorView* to) {
       TransformReplay::getMatchedLeafPosWithoutReplayPasC(to, from, pos);
   if (new_pos < 0) {
     auto replay = TransformReplay::replayPasC(to, from, pos);
+    TORCH_INTERNAL_ASSERT(
+        validateDomain(to, replay.first),
+        "Tried to set the domain of ",
+        to,
+        " to ",
+        replay.first,
+        " but that would invalidate previously compute at position or max producer position.");
     to->setDomain(replay.first);
     new_pos = replay.second;
   }
@@ -862,6 +882,13 @@ void TransformPropagator::propagateTvCasP(TensorView* from, TensorView* to) {
       TransformReplay::getMatchedLeafPosWithoutReplayCasP(to, from, pos);
   if (new_pos < 0) {
     auto replay = TransformReplay::replayCasP(to, from, pos);
+    TORCH_INTERNAL_ASSERT(
+        validateDomain(to, replay.first),
+        "Tried to set the domain of ",
+        to,
+        " to ",
+        replay.first,
+        " but that would invalidate previously compute at position or max producer position.");
     to->setDomain(replay.first);
     new_pos = replay.second;
   }
@@ -873,6 +900,13 @@ void TransformPropagator::propagateTvSibling(TensorView* from, TensorView* to) {
   // See note [Using multiple TransformPropagators]
   if (!TransformReplay::fullSelfMatching(to, from)) {
     auto replay = TransformReplay::fullSelfReplay(to->domain(), from->domain());
+    TORCH_INTERNAL_ASSERT(
+        validateDomain(to, replay),
+        "Tried to set the domain of ",
+        to,
+        " to ",
+        replay,
+        " but that would invalidate previously compute at position or max producer position.");
     to->setDomain(replay);
   }
   replayed_pos_[to] = pos;