Make loop the same as allocation to avoid a change in transform_replay.

NVIDIA · Nov 21, 2024 · cf126e6 · cf126e6
1 parent 44696e4
commit cf126e6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/csrc/fusion_segmenter.cpp b/csrc/fusion_segmenter.cpp
@@ -1889,7 +1889,7 @@ void eraseInputDistinctRootDomains(Fusion* fusion) {
           /*root_domain=*/std::vector<IterDomain*>(),
           new_logical_domain,
           new_alloc,
-          /*loop_domain=*/new_logical_domain,
+          /*loop_domain=*/new_alloc,
           tv->domain()->contiguity());
     } else {
       new_td = IrBuilder::create<TensorDomain>(

diff --git a/csrc/transform_replay.cpp b/csrc/transform_replay.cpp
@@ -757,19 +757,19 @@ std::pair<TensorDomain*, int64_t> TransformReplay::replayCasP(
       consumer->domain()->contiguity());
 
   if (producer->hasAllocation()) {
-    std::vector<IterDomain*> producer_allocation =
-        TensorDomain::noReductions(producer->getAllocationDomain());
-    ReplayTransformations replay_CasP(
-        producer_allocation,
+    auto replay_CasP = BestEffortReplay(
+        new_IDs,
+        producer->getLoopDomain(),
         logical_map.mapProducerToConsumer(producer->domain(), replayed));
     const auto& p2c_map = replay_CasP.getReplay();
 
+    auto producer_rank = producer->getAllocationDomain().size();
     std::vector<IterDomain*> new_allocation_domain;
-    new_allocation_domain.reserve(producer_allocation.size());
+    new_allocation_domain.reserve(producer_rank);
     std::vector<std::optional<bool>> new_contiguity;
-    new_contiguity.reserve(producer_allocation.size());
+    new_contiguity.reserve(producer_rank);
 
-    for (auto i : c10::irange(producer->getAllocationDomain().size())) {
+    for (auto i : c10::irange(producer_rank)) {
       IterDomain* id = producer->getAllocationDomain()[i];
       // We won't find reduction IterDomains in the map. See
       // AllocationDomainTest.CacheBefore.