WIP: IterDomain Graphs

csrc/device_lower/lower2device.cpp

@@ -390,7 +390,7 @@ void GpuLower::analysis(Fusion* fusion) {
   // functionality should be affected. New IterDomains may be created,
   // so it is expected that generated code may use diffrent variable
   // names
-  if (isOptionEnabled(EnableOption::IdModel)) {
+  if (true || isOptionEnabled(EnableOption::IdModel)) {
     IdModel id_model(fusion_);
   }
 

csrc/disjoint_set.h

@@ -81,6 +81,15 @@ class VectorOfUniqueEntries {
     return false;
   }
 
+  // Returns if a node was actually added
+  bool pushFront(T entry) {
+    if (set_.emplace(entry).second) {
+      vector_.insert(vector_.begin(), entry);
+      return true;
+    }
+    return false;
+  }
+
   // Returns true if any node was added
   bool pushBack(const VectorOfUniqueEntries<T, Hash>& other) {
     return pushBack(other.vector());
@@ -170,6 +179,14 @@ class VectorOfUniqueEntries {
     return v;
   }
 
+  // Remove and returns the last element in vector
+  T popFront() {
+    T v = vector_.front();
+    set_.erase(v);
+    vector_.erase(vector_.begin());
+    return v;
+  }
+
   // Returns if this container is empty
   bool empty() const {
     return vector_.empty();
@@ -394,7 +411,9 @@ class DisjointSets {
         entry_it != disjointSetMap().end(),
         "Strict mapping failed on element: ",
         abstractToString(entry0),
-        " either an error occurred, or non strict mapping should have been used.");
+        " either an error occurred, or non strict mapping should have been used.",
+        " ",
+        entry0->name());
     return entry_it->second->has(entry1);
   }
 

csrc/id_model/id_model.cpp

@@ -9,6 +9,7 @@
 #include <id_model/loop_promotion.h>
 #include <id_model/to_string.h>
 #include <id_model/transform_replay.h>
+#include <id_model/utils.h>
 #include <id_model/validation_utils.h>
 
 #include <device_lower/analysis/trivial_broadcast.h>
@@ -233,6 +234,203 @@ std::string IdModel::toString() const {
   return ss.str();
 }
 
+// Generate a new expr with the IterDomain inputs/outputs replaced based on map.
+// Replaced inputs/outputs should almost exact match with provided expr.
+Expr* IdModel::addExprWithReplacement(
+    const std::unordered_map<IterDomain*, IterDomain*>& old_2_new_ids,
+    Expr* old_expr) {
+  // Figure out which graphs are already initialized to make sure we add the new
+  // expression to them.
+  std::vector<IdMappingMode> initialized_modes;
+  for (auto mode : kIdMappingModes) {
+    auto graph_it = id_graphs_.find(mode);
+    if (graph_it == id_graphs_.end()) {
+      continue;
+    }
+
+    auto& graph = graph_it->second;
+    if (graph.disjointValSets().disjointSetMap().empty()) {
+      continue;
+    }
+
+    initialized_modes.push_back(mode);
+  }
+
+  // We will fill this map for every IterDomain in input and output.
+  std::unordered_map<IterDomain*, IterDomain*> replacement_map = old_2_new_ids;
+
+  // Validate replacement map. Make sure the keys are an input or output
+  for (auto replacement_entry : replacement_map) {
+    NVF_ERROR(
+        std::find(
+            old_expr->inputs().begin(),
+            old_expr->inputs().end(),
+            replacement_entry.first) != old_expr->inputs().end() ||
+            std::find(
+                old_expr->outputs().begin(),
+                old_expr->outputs().end(),
+                replacement_entry.first) != old_expr->outputs().end(),
+        "Wanted to replace ",
+        replacement_entry.first->toString(),
+        " however the is not an input or output of:\n",
+        old_expr->toString());
+  }
+
+  // If all inputs and or all output were replaced
+  bool all_inps_replaced = true;
+  bool all_outs_replaced = true;
+  {
+    for (auto inp_id : ir_utils::filterByType<IterDomain>(old_expr->inputs())) {
+      if (replacement_map.find(inp_id) == replacement_map.end()) {
+        all_inps_replaced = false;
+        replacement_map[inp_id] = inp_id->cloneWithoutRFactor();
+      }
+    }
+
+    for (auto out_id :
+         ir_utils::filterByType<IterDomain>(old_expr->outputs())) {
+      if (replacement_map.find(out_id) == replacement_map.end()) {
+        all_outs_replaced = false;
+        replacement_map[out_id] = out_id->cloneWithoutRFactor();
+      }
+    }
+
+    NVF_ERROR(
+        (all_inps_replaced || all_outs_replaced),
+        "Either all the inputs or all the outputs need to be replaced when using this function.");
+
+    for (auto mode : initialized_modes) {
+      for (auto inp_or_out_id : all_inps_replaced
+               ? ir_utils::filterByType<IterDomain>(old_expr->inputs())
+               : ir_utils::filterByType<IterDomain>(old_expr->outputs())) {
+        NVF_ERROR(
+            idGraph(mode).hasGroup(inp_or_out_id),
+            "Expected ",
+            inp_or_out_id->toString(),
+            " to be initialized in graph mode: ",
+            mode);
+      }
+    }
+  }
+
+  // Create the new expression with provided outputs
+  auto replay = ReplacementTransformCloner::clone(replacement_map, old_expr);
+
+  // Add new output iter domains to id_definitions_/id_uses_ of IdModel
+  for (auto out_id : ir_utils::filterByType<IterDomain>(replay->outputs())) {
+    id_definitions_[out_id].pushBack(replay);
+    id_uses_[out_id];
+  }
+
+  // Add new input iter domains to id_definitions_/id_uses_ of IdModel
+  for (auto inp_id : ir_utils::filterByType<IterDomain>(replay->inputs())) {
+    id_definitions_[inp_id];
+    id_uses_[inp_id].pushBack(replay);
+  }
+
+  // Update all the initialized graph mappings
+  for (auto mode : initialized_modes) {
+    auto& graph = idGraph(mode);
+
+    graph.registerExpr(replay);
+    auto replay_group = graph.toGroup(replay);
+
+    // Initialize any non-existent input ids, update existing ones
+    for (auto inp_id : ir_utils::filterByType<IterDomain>(replay->inputs())) {
+      if (!graph.disjointValSets().mappingExists(inp_id)) {
+        // inp_id is not initialized in the map, initialize it
+        graph.initializeVal(inp_id, {}, {replay});
+      } else {
+        // Update unique uses of existing input ids
+        auto inp_group = graph.toGroup(inp_id);
+        graph.addUniqueUses(inp_group, replay_group);
+      }
+    }
+
+    // Initialize any non-existent output ids, update existing ones
+    for (auto out_id : ir_utils::filterByType<IterDomain>(replay->outputs())) {
+      if (!graph.disjointValSets().mappingExists(out_id)) {
+        // out_id is not initialized in the map, initialize it
+        graph.initializeVal(out_id, {replay}, {});
+      } else {
+        // out_id is already initialized, add the replay as a unique definition
+        // of its group
+        auto out_group = graph.toGroup(out_id);
+        graph.addUniqueDefinitions(out_group, replay_group);
+      }
+    }
+
+    // If the inputs were replaced we want to map through forward the newly
+    // added expression. If the outputs were replaced we want to map through
+    // backwards the newly added expression.
+
+    // Forward
+    VectorOfUniqueEntries<Expr*> representative_uses;
+    for (auto in : ir_utils::filterByType<IterDomain>(replay->inputs())) {
+      for (const ExprGroup& use_group : graph.getUses(graph.toGroup(in))) {
+        if (use_group == replay_group) {
+          continue;
+        }
+        representative_uses.pushBack(use_group->front());
+      }
+    }
+
+    for (auto rep_use : representative_uses) {
+      graph.maybeMapThroughExprs(rep_use, replay, true);
+    }
+
+    // Backwards
+    VectorOfUniqueEntries<Expr*> representative_defs;
+    for (auto out : ir_utils::filterByType<IterDomain>(replay->outputs())) {
+      for (const ExprGroup& def_group :
+           graph.getDefinitions(graph.toGroup(out))) {
+        if (def_group == replay_group) {
+          continue;
+        }
+        representative_defs.pushBack(def_group->front());
+      }
+    }
+
+    for (auto rep_def : representative_defs) {
+      graph.maybeMapThroughExprs(rep_def, replay, false);
+    }
+  }
+  return replay;
+}
+
+// Clone provided iter domain and return the new copy. Map that copy in relevant
+// maps.
+IterDomain* IdModel::cloneIterDomain(IterDomain* id) {
+  // Figure out which graphs are already initialized to make sure we add the new
+  // expression to them.
+  std::vector<IdMappingMode> initialized_modes;
+  for (auto mode : kIdMappingModes) {
+    auto graph_it = id_graphs_.find(mode);
+    if (graph_it == id_graphs_.end()) {
+      continue;
+    }
+
+    auto& graph = graph_it->second;
+    if (graph.disjointValSets().disjointSetMap().empty()) {
+      continue;
+    }
+
+    initialized_modes.push_back(mode);
+  }
+
+  auto id_copy = id->cloneWithoutRFactor();
+
+  id_uses_[id_copy] = {};
+  id_definitions_[id_copy] = {};
+
+  for (auto mode : initialized_modes) {
+    idGraph(mode).initializeVal(id_copy, {}, {});
+    idGraph(mode).mapVals(id, id_copy);
+  }
+
+  return id_copy;
+}
+
 ValGraph IdModel::initializeIdGraph(bool propagate_through_exprs) const {
   ValGraph id_graph(propagate_through_exprs);
 
@@ -603,13 +801,25 @@ void IdModel::buildLoopGraph() {
   maybeBuildGraph(IdMappingMode::EXACT);
   maybeBuildGraph(IdMappingMode::PERMISSIVE);
 
+  if (!tv_exprs_.empty()) {
+    std::stringstream ss;
+    tv_exprs_.at(0)->fusion()->print(ss);
+    VERBOSE() << ss.str();
+  }
+
   const StatefulInliningInfo inlining_info =
       buildStatefulInliningInfo(tv_exprs_, idGraph(IdMappingMode::PERMISSIVE));
 
   initializeLoopGraph(inlining_info);
 
   validateLoopGraphHasNoSelfMappedLeafDomains();
 
+  VERBOSE() << "Initial loop graph:\n";
+  for (const auto& group :
+       idGraph(IdMappingMode::LOOP).disjointValSets().disjointSets()) {
+    VERBOSE() << nvfuser::toString(group) << std::endl;
+  }
+
   loop_promotion_map_ = LoopPromotionMapBuilder::get(
       *this, inlining_info, loop_promotion_map_builder_callback_);
 
@@ -620,7 +830,55 @@ void IdModel::buildLoopGraph() {
   idGraph(IdMappingMode::LOOP).validateConsistency();
 }
 
+// TODO: Reenable after reenabling parallel propagation.
+//        propagateLoopPTypes
+void IdModel::validatePTypes(const std::vector<TensorView*>& all_tvs) const {
+  // VectorOfUniqueEntries<IterDomain*> leaf_ids;
+  // for (auto tv : all_tvs) {
+  //   leaf_ids.pushBack(tv->domain()->leaf());
+  // }
+
+  // for (const auto& disjoint_set :
+  //      idGraph(IdMappingMode::EXACT).disjointValSets().disjointSets()) {
+  //   for (auto id : disjoint_set->vector()) {
+  //     auto id_ptype = id->getParallelType();
+
+  //     NVF_ERROR(
+  //         leaf_ids.has(id) || id_ptype == ParallelType::Serial,
+  //         "Invalid parallelization of non leaf iter domain: ",
+  //         id->toString());
+  //   }
+  // }
+}
+
+void IdModel::propagateLoopPTypes() const {
+  for (const auto& loop_disjoint_set :
+       idGraph(IdMappingMode::LOOP).disjointValSets().disjointSets()) {
+    ParallelType common_ptype = ParallelType::Serial;
+    for (auto id : loop_disjoint_set->vector()) {
+      auto id_ptype = id->as<IterDomain>()->getParallelType();
+
+      NVF_ERROR(
+          id_ptype == common_ptype || id_ptype == ParallelType::Serial ||
+              common_ptype == ParallelType::Serial,
+          "Issue validating parallel type disjoint ptype is, ",
+          common_ptype,
+          " but found in the set the id: ",
+          id->toString());
+
+      common_ptype =
+          common_ptype == ParallelType::Serial ? id_ptype : common_ptype;
+    }
+
+    for (auto id : loop_disjoint_set->vector()) {
+      id->as<IterDomain>()->parallelize(common_ptype);
+    }
+  }
+}
+
 void IdModel::buildAllGraphs() {
+  VERBOSE() << "*** Building all graphs ***";
+
   if (tvs_.empty()) {
     return;
   }
@@ -663,6 +921,11 @@ void IdModel::buildAllGraphs() {
         idGraph(IdMappingMode::PERMISSIVE));
   }
 
+  // Permissive graph needs the trivial exprs from the almost exact graph to
+  // build correctly. Once built though we can remove the trivial expressions
+  // from the almost exact graph.
+  idGraph(IdMappingMode::ALMOSTEXACT).removeTrivialExprs();
+
   buildLoopGraph();
 }
 
@@ -870,4 +1133,12 @@ std::unordered_map<ValGroup, IterDomain*> updateValGroupIdMap(
   return new_map;
 }
 
+std::unordered_map<IterDomain*, IterDomain*> IdModel::buildIndexGraph(
+    const std::vector<Expr*>& exprs,
+    const std::vector<TensorView*>& all_tvs,
+    StatefulInliningInfo& info,
+    std::unordered_map<ValGroup, IterDomain*> stale_promotion_map) {
+  NVF_ERROR(false, "Not implemented yet.");
+}
+
 } // namespace nvfuser