WIP

NVIDIA · Nov 27, 2024 · bd177e3 · bd177e3
1 parent 1427f1b
commit bd177e3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/csrc/multidevice/utils.cpp b/csrc/multidevice/utils.cpp
@@ -344,6 +344,13 @@ bool haveDifferentShardings(
     if (!is_mapped_in_id_model(p_loop_id, c_loop_id, id_model)) {
       return true;
     }
+    if (p_parallel_type_to_id.count(parallel_type)) {
+      IterDomain* p_id = p_parallel_type_to_id.at(parallel_type);
+      IterDomain* c_id = c_parallel_type_to_id.at(parallel_type);
+      if (!exact_graph.disjointValSets().strictAreMapped(p_id, c_id)) {
+        return true;
+      }
+    }
   }
 
   return false;

diff --git a/csrc/preseg_passes/insert_reshardings.cpp b/csrc/preseg_passes/insert_reshardings.cpp
@@ -30,6 +30,8 @@ bool shouldReshardAfter(Expr* expr) {
 }
 
 void insertReshardingsBefore(Fusion* fusion) {
+  IdModel id_model(fusion, false, false, true);
+  id_model.buildPermissiveGraph();
   // Remove this after we refactor this as a pre-segmenter pass.
   FusionGuard fg(fusion);
   for (Expr* expr : fusion->exprs()) {
@@ -77,6 +79,9 @@ void insertReshardingsBefore(Fusion* fusion) {
 }
 
 void insertReshardingsAfter(Fusion* fusion) {
+  IdModel id_model(fusion, false, false, true);
+  id_model.buildPermissiveGraph();
+
   // Remove this after we refactor this as a pre-segmenter pass.
   FusionGuard fg(fusion);
   // Iterate backwards over fusion expressions. Reshard after will