From 0c5629725a0cba851d1a57cf198bf8f54f79cbe4 Mon Sep 17 00:00:00 2001
From: Naoya Maruyama <nmaruyama@nvidia.com>
Date: Thu, 12 Dec 2024 14:08:20 -0800
Subject: [PATCH] Fix and a repro

---
 csrc/ir/nodes.cpp       |  6 ++++--
 tests/cpp/test_gpu3.cpp | 45 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp
index 3dae5c4f42a..9c033a66c2c 100644
--- a/csrc/ir/nodes.cpp
+++ b/csrc/ir/nodes.cpp
@@ -3763,8 +3763,10 @@ std::vector<IterDomain*> TensorDomain::allIDs() const {
   while (!ids_to_be_sorted.empty()) {
     auto it = ids_to_be_sorted.begin();
     while (it != ids_to_be_sorted.end()) {
-      auto in_it = out2in.find(*it);
-      if (in_it == out2in.end() || sorted_ids.has(in_it->second)) {
+      auto range = out2in.equal_range(*it);
+      if (std::all_of(range.first, range.second, [&](const auto& kv) {
+            return sorted_ids.has(kv.second);
+          })) {
         sorted_ids.pushBack(*it);
         it = ids_to_be_sorted.erase(it);
       } else {
diff --git a/tests/cpp/test_gpu3.cpp b/tests/cpp/test_gpu3.cpp
index abb5ae471a9..8fa235357b1 100644
--- a/tests/cpp/test_gpu3.cpp
+++ b/tests/cpp/test_gpu3.cpp
@@ -9225,6 +9225,51 @@ TEST_F(NVFuserTest, ParallelDimensionsInAllocation) {
   ASSERT_TRUE(tidx_dim != nullptr);
 }
 
+// Check the topological ordering of TensorDomain::allIDs(). Repro of
+// issue #3583
+TEST_F(NVFuserTest, AllIdsMultipleDependencies) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeConcreteTensor({10, 20});
+  fusion.addInput(tv0);
+
+  auto tv1 = slice(
+      tv0,
+      {{fusion.zeroVal(), IrBuilder::create<Val>(2)},
+       {fusion.zeroVal(), tv0->getLogicalDomain().at(1)->extent()}});
+
+  fusion.addOutput(tv1);
+
+  tv1->merge(0);
+  tv1->split(0, 4);
+  tv1->split(0, 8);
+
+  fusion.print();
+
+  auto all_ids = tv1->domain()->allIDs();
+
+  auto split2 = tv1->axis(0)->definition()->as<Split>();
+  auto split1 = split2->input(0)->definition()->as<Split>();
+  auto merge = split1->input(0)->definition()->as<Merge>();
+  auto resize = merge->input(0)->definition()->as<Resize>();
+
+  std::vector<Expr*> exprs{resize, merge, split1, split2};
+
+  for (auto expr : exprs) {
+    for (auto inp : ir_utils::filterByType<IterDomain>(expr->inputs())) {
+      auto inp_it = std::find(all_ids.begin(), all_ids.end(), inp);
+      for (auto out : ir_utils::filterByType<IterDomain>(expr->outputs())) {
+        auto out_it = std::find(all_ids.begin(), all_ids.end(), out);
+        EXPECT_LT(inp_it, out_it)
+            << "Invalid ordering: " << out->toString() << " detected before "
+            << inp->toString() << ". All IDs: " << toDelimitedString(all_ids)
+            << "\n";
+      }
+    }
+  }
+}
+
 // Test file size should be up to 10K LoC. Create a new file for more tests.
 
 } // namespace nvfuser