nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/id_model/indexing.h

@@ -0,0 +1,208 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <device_lower/analysis/trivial_broadcast.h>
+#include <id_model/id_model.h>
+#include <ir/base_nodes.h>
+#include <ir/interface_nodes.h>
+#include <options.h>
+#include <type.h>
+
+// Just for PredicateInfo. Should be moved to its own header file
+#include <index_compute.h>
+
+#include <unordered_map>
+
+namespace nvfuser {
+
+struct IndexingInfo {
+  std::vector<IterDomain*> loop_domains;
+  // Indexing traversal path from loop domains
+  ExprPath<ExprGroup> traversal_path;
+  // Index mappings of ID groups along the traversal path
+  std::unordered_map<ValGroup, Val*> index_map;
+  // Mappings from ID groups to dependent loop groups
+  std::unordered_map<ValGroup, ValGroups> loop_group_dependencies;
+};
+
+struct IndexingAllocationInfo {
+  std::vector<IterDomain*> domains;
+  std::vector<Val*> strides;
+  std::vector<bool> contiguity;
+};
+
+// The basic algorithm of indexing is:
+//
+// 1. Find the loop domains
+// 2. Find the allocation domains
+// 3. Find the path from the loop domains to the allocation domains
+// 4. Set the initial index vals for the loop domains
+// 5. Propagate the initial indices of the loop domains to the allocation
+// domains
+//
+// The indexing traversal is done on the AlmostExact graph augmented
+// with the loop promotion map since both the loop and allocations
+// domains may be promoted.
+class TensorIndexer {
+ public:
+  // Using non-const references of IdModel because traversalGraph() returns a
+  // non-const reference
+  TensorIndexer(IdModel& id_model);
+
+  bool isContigIndexingEnabled() const {
+    return !isOptionDisabled(DisableOption::ContigIndexing);
+  }
+
+  // Get a linear index of a given tensor appearing in a given expr, either
+  // as a consumer or a producer. The predicate indexing will have a
+  // separate interface.
+  //
+  // The actual for-loops are required for handling circular buffering
+  Val* getLinearIndex(
+      TensorView* tv,
+      const Expr* expr,
+      const std::vector<ForLoop*>& loops) const;
+
+  // Get the index of a loop domain.
+  Val* getLoopIndex(IterDomain* loop_id, const std::vector<ForLoop*>& for_loops)
+      const;
+
+  // Get the index of the given ID groups
+  std::vector<Val*> getIndexFor(
+      const Expr* expr,
+      bool as_consumer,
+      const std::vector<IterDomain*>& index_ids,
+      const std::vector<ForLoop*>& loops) const;
+
+  // Get the contig indices of the given ID groups with their strides
+  std::pair<std::vector<Val*>, std::vector<Val*>> getContigIndexFor(
+      const Expr* expr,
+      bool as_consumer,
+      const IndexingAllocationInfo& alloc_info,
+      const std::vector<ForLoop*>& loops) const;
+
+  // The AlmostExact graph is used since size-1 splits and merges
+  // should not affect actual index exprs.
+  // Returns non-const reference because indexing may create new domains and
+  // need to update the graph.
+
+  static IdMappingMode traversalGraphType() {
+    return IdMappingMode::ALMOSTEXACT;
+  }
+
+  ValGraph& traversalGraph() const {
+    return id_model_.idGraph(traversalGraphType());
+  }
+
+  // Traverse exprs and set allocation info for each tensor
+  void setupAllocationDomains(const std::vector<Expr*>& exprs);
+
+  // Get the list of predicates of a given tensor appearing in a given
+  // expr as a consumer. Each predicate corresponds to a domain of the
+  // tensor, which is by default one of the logical domains but can be
+  // an intermediate domain with contiguous indexing.
+  //
+  // An optional ForLoop parameter specifies a loop that is either
+  // unswitched/unrolled or vectorized, both of which are handled by
+  // UnswitchPredicate. For normal inline predicates, the parameter
+  // should be nullptr.
+  std::vector<PredicateInfo> getPredicates(
+      TensorView* tv,
+      const Expr* expr,
+      const std::vector<ForLoop*>& for_loops,
+      ForLoop* unswitched_loop = nullptr) const;
+
+  // Get the indexing traversal path for indexing a given list of IDs
+  // for a given expr
+  ExprPath<ExprGroup> getIndexingPath(
+      const Expr* expr,
+      const std::vector<IterDomain*>& index_ids) const;
+
+ private:
+  // Build a map of loop groups to their index Vals. See the comment
+  // on loop_index_map_.
+  void buildLoopIndexMap();
+
+  const IndexingAllocationInfo& getIndexingAllocationInfo(
+      TensorView* tv) const {
+    auto it = alloc_info_.find(tv);
+    NVF_ERROR(
+        it != alloc_info_.end(),
+        "No allocation info found for ",
+        tv->toString());
+    return it->second;
+  }
+
+  // Returns the index map as well as its traversal path of given
+  // index domains appearing in a given expr. Used by
+  // getIndexFor.
+  IndexingInfo computeIndex(
+      const Expr* expr,
+      const std::vector<IterDomain*>& index_ids,
+      const std::vector<ForLoop*>& for_loops) const;
+
+  // Propagate the loop indices of a given list of loop domains to the
+  // traversal graph (i.e., the AlmostExact graph). Uses the loop
+  // index map, which is built for the Loop graph.
+  std::unordered_map<ValGroup, Val*> getInitialIndexMap(
+      const std::vector<IterDomain*>& loop_domains,
+      const std::vector<ForLoop*>& for_loops) const;
+
+  // Get the loop domains of a given expr. Currently, they're always
+  // the loop domains of a consumer tensor, but in the future this
+  // function may return the loop domains of a producer for
+  // producer-based indexing.
+  std::vector<IterDomain*> getLoopDomains(const Expr* expr) const;
+
+  // For a given indexng traversal path toward allocation_domains,
+  // return the contiguous domains and their strides that can provide
+  // equivalent indexing results.
+  //
+  // Currently, only backward traversal is supported.
+  std::pair<std::vector<ValGroup>, std::vector<Val*>> getContigDomainsAndStrides(
+      const IndexingAllocationInfo& alloc_info,
+      const ExprPath<ExprGroup>& traversal_path) const;
+
+  // Get a replace map for tensor indexing. Examples include replacing
+  // an index of a vectorized loop with zero.
+  //
+  // This replacement map is used to replace a tensor index after an
+  // index map is generated. Since replacment is only done for loop
+  // domains, this could be done as part of getInitialIndexMap. One
+  // reason that we might want to first generate an index and do some
+  // replacements, rather than using final index vals to build the
+  // index map, is that one index map could be used for multiple
+  // indices. For normal tensor indexing, this may not matter, but for
+  // predicate indexing, it needs to generate both start and stop
+  // predicates, and one index map would be sufficient for both
+  // indices by using different replacement maps.
+  std::unordered_map<Val*, Val*> getIndexReplacementMap(
+      const Expr* expr,
+      bool as_consumer,
+      const std::vector<IterDomain*>& loop_domains,
+      const std::vector<ForLoop*>& for_loops,
+      const std::unordered_map<ValGroup, Val*>& index_map) const;
+
+ private:
+  // Using non-const references of IdModel because traversalGraph() returns a
+  // non-const reference
+  IdModel& id_model_;
+
+  // Mappings from loop groups to their indices. Serial loops will
+  // be mapped a unique loop index Val. Parallel loops will be mapped
+  // to NamedScalar such as "threadIdx.x". This map needs to be built
+  // once and can be reused for different tensors.
+  std::unordered_map<ValGroup, Val*> loop_index_map_;
+
+  // Allocation info for each tensor. Must be filled before computing
+  // the index of each tensor
+  std::unordered_map<TensorView*, IndexingAllocationInfo> alloc_info_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/indexing_traversal.h

@@ -0,0 +1,72 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <device_lower/utils.h>
+#include <val_graph_visitor.h>
+
+namespace nvfuser {
+
+// BFS traversal for indexing. The only difference with the default
+// ValGraphBFS is that for indexing there must be a special care taken
+// when resize is involved since there can be multiple paths and
+// there's only one correct path. Specifically, any resize expr group
+// node must appear in the root-logical path of the consumer
+// tensor. Otherwise, resize nodes should be ignored. See
+// IndexingTest.ResizePath for a concrete example.
+class IndexingTraversal : public ValGraphBFS {
+ public:
+  IndexingTraversal(
+      const Expr* expr,
+      const ValGraph& graph,
+      std::vector<NodeType> from_groups,
+      std::vector<NodeType> to_groups,
+      bool require_all_to_visited = true);
+
+  ~IndexingTraversal() override = default;
+
+  static ExprPath getExprsBetween(
+      const Expr* expr,
+      const ValGraph& graph,
+      const std::vector<IterDomain*>& from_domains,
+      const std::vector<IterDomain*>& to_domains);
+
+  static std::optional<ExprPath> getExprsBetweenForResize(
+      const Expr* expr,
+      const ValGraph& graph,
+      const std::vector<IterDomain*>& from_domains,
+      const std::vector<IterDomain*>& to_domains);
+
+  using ValGraphBFS::isVisited;
+
+  bool excludeFromTraversal(const NodeType& group) const override {
+    if (const ExprGroup* eg = std::get_if<ExprGroup>(&group)) {
+      if ((*eg)->empty()) {
+        return false;
+      }
+      auto resize = dynamic_cast<Resize*>((*eg)->front());
+      if (resize == nullptr) {
+        return false;
+      }
+      if (std::none_of((*eg)->begin(), (*eg)->end(), [&](Expr* expr) -> bool {
+            return resize_paths_.find(expr->as<Resize>()) !=
+                resize_paths_.end();
+          })) {
+        // This resize node should never be traversed for indexing of
+        // the given expr
+        return true;
+      }
+    }
+    return false;
+  }
+
+ private:
+  std::unordered_set<Resize*> resize_paths_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/indexing_utils.h

@@ -0,0 +1,62 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <device_lower/analysis/index_compute.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <id_model/id_model.h>
+#include <id_model/to_string.h>
+
+namespace nvfuser {
+namespace indexing_utils {
+
+// Get a matching ForLoop for a given loop iter domain. There may not
+// be such a loop if this loop-nest is for initializing a reduction
+// buffer.
+inline ForLoop* getForLoop(
+    IterDomain* loop_id,
+    const std::vector<ForLoop*>& for_loops,
+    const ValGraph& loop_graph) {
+  auto it = std::find_if(
+      for_loops.begin(), for_loops.end(), [&](ForLoop* for_loop) -> bool {
+        IterDomain* for_loop_id = for_loop->iter_domain();
+        return loop_graph.disjointValSets().strictAreMapped(
+            loop_id, for_loop_id);
+      });
+  if (it != for_loops.end()) {
+    return *it;
+  } else {
+    return nullptr;
+  }
+}
+
+// Check if unswitching a given for-loop actually matters. For example,
+// if a loop is parallelized, unswitching doesn't mean anything as we
+// don't unswitch threading dimensions, e.g., "threadIdx.x + ... < N"
+// is generated rather than "blockDim.x + ... < N".
+inline bool isEffectiveUnswitchLoop(ForLoop* fl) {
+  // Threaded domain is not unswitched
+  if (fl->iter_domain()->isThread() || fl->iter_domain()->isDeviceDim()) {
+    return false;
+  }
+
+  // If it's vectorized, it must be true that any of the iteration
+  // values can be used to generate the predicates of the tensor, so
+  // unswitching has no effect. Same for loops that are known to be
+  // safe to just predicate at the end.
+  if (fl->iter_domain()->getParallelType() == ParallelType::Vectorize ||
+      lower_utils::predicateAtEnd(fl)) {
+    return false;
+  }
+
+  return true;
+}
+
+} // namespace indexing_utils
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/loop_promotion.h

@@ -0,0 +1,180 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <val_graph.h>
+
+namespace nvfuser {
+
+class IdModel;
+struct StatefulInliningInfo;
+
+// Callback interface for LoopPromotionMapBuilder. Allow exposing the
+// temporary maps for testing and debugging
+class LoopPromotionMapBuilderCallback {
+ public:
+  virtual ~LoopPromotionMapBuilderCallback() = default;
+
+  // Called after Step 1 with the root resolution map and the
+  // corresponding IEL graph
+  virtual void postStep1(
+      const std::unordered_map<ValGroup, IterDomain*>& iel_root_resolution_map,
+      const ValGraph& iel_graph) {}
+  // Called after Step 2 with the IEL promotion map and the
+  // corresponding IEL graph
+  virtual void postStep2(
+      const std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
+      const ValGraph& iel_graph) {}
+  // Called after Step 3 with the loop promotion map
+  virtual void postStep3(
+      const std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map) {}
+  // Called after Step 4 with the IEL promotion map and the
+  // corresponding IEL graph
+  virtual void postStep4(
+      const std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
+      const ValGraph& iel_graph) {}
+  // Called after Step 3 with the final loop promotion map
+  virtual void postStep5(
+      const std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map) {}
+};
+
+class LoopPromotionMapBuilder {
+ public:
+  // Build a map of loop groups to IterDomains that represent actual
+  // loops. The map is built based on the broadcast resolution with
+  // root domains between inlined producer and consumer tensors.
+  //
+  // (For debugging only) When force_full_loop_promotion_analysis is
+  // true, it always performs the full loop promotion analysis even
+  // when it's possible to take a quicker shortcut.
+  static std::unordered_map<ValGroup, IterDomain*> get(
+      IdModel& id_model,
+      const StatefulInliningInfo& inlining_info,
+      LoopPromotionMapBuilderCallback* callback = nullptr,
+      bool force_full_loop_promotion_analysis = false);
+
+ private:
+  LoopPromotionMapBuilder(
+      IdModel& id_model,
+      const StatefulInliningInfo& inlining_info,
+      LoopPromotionMapBuilderCallback* callback = nullptr,
+      bool force_full_loop_promotion_analysis = false);
+
+  std::unordered_map<ValGroup, IterDomain*> build();
+
+  // Shortcut to build a map of promotion IDs without doing the full
+  // loop promotion analysis. Can only be used when the full analysis
+  // is not requierd.
+  std::unordered_map<ValGroup, IterDomain*> buildWithNoBroadcast();
+
+  ValGraph& idGraph(IdMappingMode mode);
+  const ValGraph& idGraph(IdMappingMode mode) const;
+
+  std::unordered_map<ValGroup, IterDomain*> buildInlineRootResolutionMap(
+      const ValGraph& iel_graph,
+      const StatefulInliningInfo& info) const;
+
+  // Helper function for building loop promotion map.
+  //
+  // Propagate promotion mappings from root IEL groups to intermediate
+  // and loop IEL groups by traversing IEL exprs. For each expr, if an
+  // input is promoted, the output needs to be promoted too. If
+  // there's already an equivalent expr that uses the promoted inputs,
+  // create a mapping from the outputs of the IEL expr to the outputs
+  // of the equivalent expr. We only consider exprs that are mapped
+  // in the loop graph as we are looking for domains that represent
+  // the actual loops of the input and output domains of the IEL
+  // expr. If no such expr is found, the IEL expr is replayed with the
+  // promoted inputs.
+  //
+  // This is used twice when building the promotion map. The first time
+  // it is used there's no loop graph promotion yet, so only the IEL
+  // promotions are propagated. In that case, loop_graph_promotion_map
+  // should be just empty.
+  //
+  // Propagation uses iel_promotion_map and
+  // loop_graph_promotion_map. If both are available for an IEL group,
+  // the former has the precedence. This is because when this function
+  // is used for step 4, the given iel_promotion_map starts as an
+  // empty map and gets populated during this propagation, so any
+  // mapping in the map is guaranteed to be the correct final mapping,
+  // whereas the loop graph may have invalid mappings for partially
+  // inlined domains.
+  void propagatePromotionsInIELGraph(
+      const ValGraph& iel_graph,
+      std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
+      const ValGraph& loop_graph,
+      const std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map);
+
+  // Same as the other propagatePromotionsInIELGraph but without loop
+  // graph map. This is used for step 2, where there's no loop
+  // graph map yet.
+  void propagatePromotionsInIELGraph(
+      const ValGraph& iel_graph,
+      std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map);
+
+  // Given an IEL promotion map, identify the mapping of each loop
+  // group. The promotion must represent all the domains in each loop
+  // group. If a valid representative promotion is not found for a
+  // loop group, no mapping is added for the group.
+  std::unordered_map<ValGroup, IterDomain*> projectIELPromotionToLoopGraph(
+      const ValGraph& iel_graph,
+      const std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
+      const ValGraph& loop_graph,
+      const StatefulInliningInfo& inlining_info) const;
+
+  // Find a promoted iter domain of a given loop group that covers all
+  // the exact groups representative of the resolved transformations
+  // within the loop group. Specifically, we examine each IEL group of
+  // the loop group, and if an IEL group has a promotion, we consider it as a
+  // candidate of the promotion of this loop group. If not, we include a
+  // domain of the IEL group as a candidate too. Once all candidates are
+  // obtained, we pick one that covers all the exact domains (cf. concrete
+  // domains in ComputeAtMap)
+  IterDomain* findPromotionOfLoopGroup(
+      const ValGroup& loop_group,
+      const ValGraph& iel_graph,
+      const std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
+      const std::unordered_map<ValGroup, ValGroups>& exact_covered_ids,
+      const VectorOfUniqueEntries<IterDomain*>& terminal_loop_ids) const;
+
+  // Terminal loop ids are iteration domains in each loop group that:
+  // 1) Don't have an entry in p2c_ca_permissive_maps, which would mean a
+  //    consumer TV's iter domain maps to this domain in a way that that domain
+  //    is also in the same loop group
+  // 2) Don't have a direct IterDomain consumer within the group
+  VectorOfUniqueEntries<IterDomain*> computeTerminalLoopIds(
+      const StatefulInliningInfo& info) const;
+
+  // Given the Step-3 promotion results, returns only promotions of
+  // groups that are producers to partially inlined groups. Those
+  // partially inlined groups may not have correct promotions as of
+  // Step 3 and need another propagation pass.
+  std::unordered_map<ValGroup, IterDomain*>
+  getProducerPromotionsOfPartiallyInlinedGroups(
+      const std::unordered_map<ValGroup, IterDomain*>&
+          initial_loop_promotion_map,
+      const ValGraph& loop_graph) const;
+
+  // Basic consistency check of the given loop promotion map
+  void sanityCheckLoopPromotionMap(
+      const std::unordered_map<ValGroup, IterDomain*>& loop_promotion_map)
+      const;
+
+ private:
+  IdModel& id_model_;
+  const StatefulInliningInfo& inlining_info_;
+  LoopPromotionMapBuilderCallback* callback_ = nullptr;
+
+  // (For debugging only) When force_full_loop_promotion_analysis_ is
+  // true, it always performs the full loop promotion analysis even
+  // when it's possible to take a quicker shortcut.
+  bool force_full_loop_promotion_analysis_ = false;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/predicate_indexing.h

@@ -0,0 +1,104 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <id_model/id_model.h>
+
+namespace nvfuser {
+
+// Get the domains to predicate for a given tensor used as a consumer
+// of a given expr.
+std::vector<IterDomain*> getPredicateDomains(
+    TensorView* consumer_tv,
+    const Expr* expr);
+
+// Get a replace map for predicate indexing of a given tensor appearing
+// in a given loop-nest.
+//
+// The unswitched_loop parameter is an optional ForLoop that is used
+// when this predicate is for an unswitched, unrolled or vectorized
+// loop.
+std::unordered_map<Val*, Val*> getPredicateIndexReplacementMap(
+    TensorView* tv,
+    const std::vector<ForLoop*>& for_loops,
+    const std::unordered_map<ValGroup, Val*>& index_map,
+    const ValGraph& traversal_graph,
+    const ExprPath<ExprGroup>& traversal_path,
+    const IdModel& id_model,
+    bool is_start_predicate,
+    ForLoop* unswitched_loop = nullptr);
+
+// Check if a given ExprGroup is a split that needs an additional
+// predicate due to its non-divisibility.
+inline bool isNonDivisibleSplit(const ExprGroup& expr_group) {
+  if (!expr_group->front()->isA<Split>()) {
+    return false;
+  }
+
+  const auto& non_divisible_split_info =
+      GpuLower::current()->nonDivisibleSplitInfo();
+
+  std::vector<PredicateDomainInfo> pred_info_vec;
+
+  // The splitsToPredicate map is for each tensor. Here, it's assumed
+  // that if any tensor has a non-divisible split that's mapped with
+  // expr_group, it should be considered a non-divisible split. This
+  // may result in duplicate predicates, which should be removed by
+  // the expression simplifier.
+  //
+  // For example, suppose tv0 is a 1D tensor of size 16:
+  //
+  // auto tv1 = reshape(tv0, {16}, {2, 8});
+  // tv1->split(1, 3);
+  //
+  // propagate_transformation(to: tv0, from: tv1)
+  //
+  // Here, the split by 3 of tv1 is not included in its non-divisible
+  // split list even though it is indeed non-divisible. The reason is
+  // that the input to the non-divisible split, the inner logical
+  // domain of extent 8, is predicated anyway since it's a logical
+  // domain. Specifically, its predicate should consist of:
+  //
+  // - Predicate for the outer logical domain
+  // - Predicate for the inner logical domain
+  //
+  // However, for tv0, it is indeed included in the non-divisible
+  // split list since the domain of extent 8 is not part
+  // of its logical domain.
+  //
+  // - Predicate for the sole logical domain
+  // - Predicate for the non-divisible split
+  //
+  // This would mean that when generating a predicate for tv1, since
+  // the below check would find a mapping with the non-divisible split
+  // for tv0, the tv1 predicate would be:
+  //
+  // - Predicate for the outer logical domain
+  // - Predicate for the inner logical domain
+  // - Predicate for the non-divisible split
+  //
+  // Here, the last two predicates are redundant since both of them
+  // guard the index with respect to the domain of extent 8, which is
+  // redundant. This is a bit annonying but should have no actual
+  // impact as the redundancy should be removed by the expression
+  // simplifier.
+  for (const auto& [tv, splits] :
+       non_divisible_split_info.splitsToPredicate()) {
+    if (std::find_if(splits.begin(), splits.end(), [&](Split* split) {
+          return expr_group->has(split);
+        }) != splits.end()) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/schedule.h

@@ -0,0 +1,54 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <val_graph.h>
+
+namespace nvfuser {
+
+// Choose an IterDomain from the ValGroup that is amenable to transforms.
+// Specifically we prefer, in descending order:
+//   1. Iteration domains
+//   2. Broadcast domains
+//   3. Reduction domains
+IterDomain* representativeId(const ValGroup& vg);
+
+// Given a ValGraph and two ValGroups g0 and g1 in this graph, if there is
+// already a merge of g0 with g1 in graph, return the output ValGroup of that
+// merge. Otherwise create an new ValGroup that is a merge of g0 and g1 in
+// graph, and a new ExprGroup that is the definition of the new ValGroup.
+// After the merge, g0 and g1 will remain valid pointers.
+ValGroup merge(ValGraph* graph, const ValGroup& g0, const ValGroup& g1);
+
+// Given a ValGraph and a ValGroup g in this graph, if there is already a split
+// of g in graph with the same factor, then return the output ValGroups of that
+// split. Otherwise create two new ValGroups that are a split of g in
+// graph, and a new ExprGroup that is the definition of the new ValGroups.
+// After the split, g will remain valid pointers.
+std::pair<ValGroup, ValGroup> split(
+    ValGraph* graph,
+    const ValGroup& g,
+    Val* factor,
+    bool inner_split = true);
+std::pair<ValGroup, ValGroup> split(
+    ValGraph* graph,
+    const ValGroup& g,
+    int64_t factor,
+    bool inner_split = true);
+
+// Given a ValGraph and two ValGroups g0 and g1 in this graph, if there is
+// already a swizzle of g0 with g1 in graph, return the output ValGroups of that
+// swizzle. Otherwise create two new ValGroups that are a swizzle of g0 and g1
+// in graph, and a new ExprGroup that is the definition of the new ValGroups.
+std::pair<ValGroup, ValGroup> swizzle(
+    ValGraph* graph,
+    SwizzleType swizzle_type,
+    const ValGroup& g0,
+    const ValGroup& g1);
+
+} // namespace nvfuser