Add slice tests to demonstrate manual scheduling

CMakeLists.txt

@@ -141,6 +141,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/id_model/id_model.cpp
   ${NVFUSER_SRCS_DIR}/id_model/id_model_index_compute.cpp
   ${NVFUSER_SRCS_DIR}/id_model/indexing.cpp
+  ${NVFUSER_SRCS_DIR}/id_model/indexing_traversal.cpp
   ${NVFUSER_SRCS_DIR}/id_model/loop_promotion.cpp
   ${NVFUSER_SRCS_DIR}/id_model/predicate_indexing.cpp
   ${NVFUSER_SRCS_DIR}/id_model/schedule.cpp

csrc/id_model/indexing_traversal.cpp

@@ -0,0 +1,46 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <id_model/indexing_traversal.h>
+
+namespace nvfuser {
+
+IndexingTraversal::IndexingTraversal(
+    const Expr* expr,
+    const ValGraph& graph,
+    std::vector<NodeType> from_groups,
+    std::vector<NodeType> to_groups)
+    : ValGraphBFS(graph, from_groups, to_groups) {
+  auto consumer_tv = ir_utils::getTvOutput(expr);
+  NVF_ERROR(consumer_tv != nullptr);
+  // Remember the resize exprs appearing in the consumer
+  // tensor. These resize exprs are the only ones that should be
+  // valid to visit when indexing the inputs and outputs of the
+  // expr.
+  //
+  // This is a WAR for cases like
+  // ResizeTest.SliceScheduledLikeProducer. Alternatively, we could
+  // have a separate graph for indexing that does not map producer and
+  // consumers in non-unary ops. See PR #2897.
+  auto all_ids = consumer_tv->domain()->allIDs();
+  std::unordered_set<IterDomain*> all_id_set(all_ids.begin(), all_ids.end());
+  for (auto id : all_ids) {
+    auto resize = dynamic_cast<Resize*>(id->definition());
+    if (resize == nullptr) {
+      continue;
+    }
+    auto resize_in = resize->in();
+    if (all_id_set.find(resize_in) == all_id_set.end()) {
+      // ths resize must not be part of the exprs involved for
+      // the domains of consumer_tv
+      continue;
+    }
+    resize_paths_.insert(resize);
+  }
+}
+
+} // namespace nvfuser

csrc/id_model/indexing_traversal.h

@@ -25,26 +25,7 @@ class IndexingTraversal : public ValGraphBFS {
       const Expr* expr,
       const ValGraph& graph,
       std::vector<NodeType> from_groups,
-      std::vector<NodeType> to_groups)
-      : ValGraphBFS(graph, from_groups, to_groups) {
-    auto consumer_tv = ir_utils::getTvOutput(expr);
-    NVF_ERROR(consumer_tv != nullptr);
-    if (consumer_tv->hasRoot()) {
-      // Remember the resize exprs appearing in the consumer
-      // tensor. These resize exprs are the only ones that should be
-      // valid to visit when indexing the inputs and outputs of the expr
-      auto root_to_logical_exprs = StmtSort::getExprsBetween(
-          {consumer_tv->getRootDomain().begin(),
-           consumer_tv->getRootDomain().end()},
-          {consumer_tv->getLogicalDomain().begin(),
-           consumer_tv->getLogicalDomain().end()});
-      for (Expr* root_to_logical_expr : root_to_logical_exprs) {
-        if (auto resize = dynamic_cast<Resize*>(root_to_logical_expr)) {
-          resize_paths_.insert(resize);
-        }
-      }
-    }
-  }
+      std::vector<NodeType> to_groups);
 
   ~IndexingTraversal() override = default;