Use linked hash map to update logical domain in rfactor (#5100)

This approach should be faster and simplifies the update logic.

For PR #5090

Author: Priya2698

csrc/transform_rfactor.cpp

@@ -13,6 +13,7 @@
 #include <ir/iostream.h>
 #include <ir/utils.h>
 #include <iter_visitor.h>
+#include <linked_hash_map.h>
 #include <ops/arith.h>
 
 namespace nvfuser {
@@ -55,40 +56,26 @@ namespace {
 // in this replay.
 class ReplayRFactor : public ReplayTransformations {
  private:
-  // Perform the update of the logical domain by replacing "replace0" with
-  // "with0" and if not nullptr "with1", also removes "replace1" if not nullptr.
-  void updateRFactorDomain(
-      IterDomain* replace0,
-      IterDomain* replace1,
-      IterDomain* with0,
-      IterDomain* with1) {
-    NVF_ERROR(
-        with0 != nullptr,
-        "The first provided IterDomain should be a real pointer,",
-        " the second iter domain provided can be a nullptr.");
-    auto pos =
-        std::find(logical_domain_.begin(), logical_domain_.end(), replace0);
-    NVF_ERROR(
-        pos != logical_domain_.end(),
-        "Could not find iter domain: ",
-        replace0->toString(),
-        " in the logical domain to replace.");
-    logical_domain_.insert(pos, with0);
-    if (with1 != nullptr) {
-      pos = std::find(logical_domain_.begin(), logical_domain_.end(), replace0);
-      logical_domain_.insert(pos, with1);
+  void splitId(Split* split) {
+    auto it = logical_domain_.erase(split->in()).second;
+    logical_domain_.insert(it, split->outer(), std::monostate());
+    logical_domain_.insert(it, split->inner(), std::monostate());
+  }
+
+  void mergeId(Merge* merge) {
+    auto outer_it = logical_domain_.erase(merge->outer()).second;
+    logical_domain_.insert(outer_it, merge->out(), std::monostate());
+    logical_domain_.erase(merge->inner());
+  }
+
+  void updateRFactorDomain(Expr* expr) {
+    if (expr->isA<Split>()) {
+      splitId(expr->as<Split>());
     }
-    pos = std::find(logical_domain_.begin(), logical_domain_.end(), replace0);
-    logical_domain_.erase(pos);
-    if (replace1 != nullptr) {
-      pos = std::find(logical_domain_.begin(), logical_domain_.end(), replace1);
-      NVF_ERROR(
-          pos != logical_domain_.end(),
-          "Wanted to replace ",
-          replace1->toString(),
-          " but it's not in the logical domain.");
-      logical_domain_.erase(pos);
+    if (expr->isA<Merge>()) {
+      mergeId(expr->as<Merge>());
     }
+    NVF_ERROR("Unrecognized expression: ", expr->toString());
   }
 
   // Took a good bit of this from ReplayTransformations::handle(Split...)
@@ -153,7 +140,7 @@ class ReplayRFactor : public ReplayTransformations {
     id_map_[s->inner()] = idi;
 
     if (static_logical_ids_.count(s->in())) {
-      updateRFactorDomain(s->in(), nullptr, s->outer(), s->inner());
+      updateRFactorDomain(s);
     }
   }
 
@@ -211,7 +198,7 @@ class ReplayRFactor : public ReplayTransformations {
               static_logical_ids_.count(m->outer()),
           "If one input to a merge is a static logical id, the other must be "
           "as well.");
-      updateRFactorDomain(m->outer(), m->inner(), m->out(), nullptr);
+      updateRFactorDomain(m);
     }
   }
 
@@ -241,7 +228,7 @@ class ReplayRFactor : public ReplayTransformations {
   // The updated domain matching the producer's logical domain. This rfactor
   // domain is relative to the iter domains in the origianl_domain and must be
   // updated to grab the mapped id's later.
-  std::vector<IterDomain*> logical_domain_;
+  LinkedHashMap<IterDomain*, std::monostate> logical_domain_;
 
   ReplayRFactor(
       // Original domain the rfactor is in reference to.
@@ -257,8 +244,7 @@ class ReplayRFactor : public ReplayTransformations {
       std::unordered_set<IterDomain*> static_logical_ids)
       : ReplayTransformations(original_domain->loop(), std::move(id_map)),
         rfactor_axes_(std::move(rfactor_axes)),
-        static_logical_ids_(std::move(static_logical_ids)),
-        logical_domain_(original_domain->logical()) {
+        static_logical_ids_(std::move(static_logical_ids)) {
     const auto all_dep_vals = DependencyCheck::getAllValsBetween(
         {original_domain->maybeRoot().begin(),
          original_domain->maybeRoot().end()},
@@ -267,8 +253,19 @@ class ReplayRFactor : public ReplayTransformations {
     auto all_dep_ids = ir_utils::filterByType<IterDomain>(all_dep_vals);
     rfactor_dep_ids_.insert(all_dep_ids.begin(), all_dep_ids.end());
 
+    for (IterDomain* id : original_domain->logical()) {
+      logical_domain_.pushBack(id, std::monostate());
+    }
+
     setErrorOnFailure(false);
   }
+
+  std::vector<IterDomain*> logical() const {
+    auto logical_ids = std::views::keys(logical_domain_);
+    std::vector<IterDomain*> transformed_logical(
+        logical_ids.begin(), logical_ids.end());
+    return transformed_logical;
+  }
 };
 
 } // namespace
@@ -421,10 +418,11 @@ std::pair<TensorDomain*, TensorDomain*> TransformRFactor::runReplay(
   // Specify the logical domain of the producer which will match the consumer
   // root domain.
   std::vector<IterDomain*> new_producer_logical_domain;
-  new_producer_logical_domain.reserve(replay_rfactor.logical_domain_.size());
+  std::vector<IterDomain*> transformed_original_logical =
+      replay_rfactor.logical();
   std::transform(
-      replay_rfactor.logical_domain_.begin(),
-      replay_rfactor.logical_domain_.end(),
+      transformed_original_logical.begin(),
+      transformed_original_logical.end(),
       std::back_inserter(new_producer_logical_domain),
       [&](IterDomain* id) {
         auto replayed_id_it = original_to_producer_id_map.find(id);