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

nvfuser/include/nvfuser/id_model/contiguity.h

@@ -0,0 +1,166 @@
+// 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 <contiguity.h>
+#include <id_model/id_model.h>
+
+namespace nvfuser {
+
+// Minimal adaptation of OrderedIdInformation for IdModel. Note that
+// the analysis only propagates forward for now.
+class OrderedIdGroupInformation : public OrderedIdInformation {
+ public:
+  // Run the ordering analysis from given allocation domains through
+  // a given traversal path
+  static OrderedIdGroupInformation get(
+      const std::vector<IterDomain*>& alloc_domain,
+      const ExprPath<ExprGroup>& path_from_alloc,
+      const ValGraph& graph) {
+    OrderedIdGroupInformation info(alloc_domain, graph);
+    info.traverse(path_from_alloc);
+    return info;
+  }
+
+  // Traversal is based on the AlmostExact graph, so matching of iter
+  // domains also needs to be done with the same graph
+  bool isConsistentlyOrdered(IterDomain* id) const override {
+    return std::find_if(
+               consistently_ordered_ids_.begin(),
+               consistently_ordered_ids_.end(),
+               [&](IterDomain* consistent_id) -> bool {
+                 return graph_.disjointValSets().strictAreMapped(
+                     consistent_id, id);
+               }) != consistently_ordered_ids_.end();
+  }
+
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<IterDomain*>>::
+      const_iterator
+      findAllocIDs(IterDomain* id) const override {
+    // This is a little ugly workaround. id_to_alloc_ids_ is a map of
+    // iter domains. If it were a map from ValGroup, this lookup
+    // should have been O(1)
+    return std::find_if(
+        id_to_alloc_ids_.begin(),
+        id_to_alloc_ids_.end(),
+        [&](const auto& kv) -> bool {
+          return graph_.disjointValSets().strictAreMapped(kv.first, id);
+        });
+  }
+
+ protected:
+  OrderedIdGroupInformation(
+      const std::vector<IterDomain*>& alloc_domain,
+      const ValGraph& graph)
+      : OrderedIdInformation(alloc_domain), graph_(graph) {
+    using_id_graph_ = true;
+  }
+
+  // Currently only forward propagation is supported
+  void traverse(const ExprPath<ExprGroup>& path_from_alloc) {
+    for (const auto& [eg, direction] : path_from_alloc) {
+      if (direction == Direction::Backward) {
+        // TODO: support Backward prop
+        continue;
+      }
+      dispatch(eg->front());
+    }
+  }
+
+  std::vector<IterDomain*>::const_iterator findActiveId(
+      IterDomain* id) const override {
+    NVF_ERROR(id != nullptr);
+    auto it = std::find_if(
+        active_ids_.begin(),
+        active_ids_.end(),
+        [&](IterDomain* active_id) -> bool {
+          return active_id != nullptr &&
+              graph_.disjointValSets().strictAreMapped(active_id, id);
+        });
+    return it;
+  }
+
+ private:
+  const ValGraph& graph_;
+};
+
+// Adapted from ContigIDs
+class ContigIDGroups {
+ public:
+  ContigIDGroups(
+      const std::vector<IterDomain*>& alloc_domains,
+      std::vector<bool> contiguity,
+      const ExprPath<ExprGroup>& path_from_alloc,
+      const ValGraph& graph,
+      bool is_predicate_pass);
+
+  void dispatch(const ExprGroup& eg, Direction direction) {
+    NVF_ERROR(!eg->empty());
+    Expr* expr = eg->front();
+
+    // Currently not propagating any contiguity information with
+    // swizzles as contiguity is generally not preserved after swizzles.
+    // But in follow ups we could gradually add back a few special
+    // cases, depending on specific swizzle type and axes.
+
+    if (auto merge = dynamic_cast<Merge*>(expr)) {
+      handle(merge, direction);
+    } else if (auto split = dynamic_cast<Split*>(expr)) {
+      handle(split, direction);
+    } else if (auto resize = dynamic_cast<Resize*>(expr)) {
+      handle(resize, direction);
+    }
+  }
+
+  void handle(Merge* merge, Direction direction);
+
+  void handle(Split* split, Direction direction);
+
+  void handle(Resize* resize, Direction direction);
+
+  const std::unordered_set<ValGroup>& contigIDs() const {
+    return contig_ids_;
+  }
+
+  const std::unordered_map<IterDomain*, ValGroup>& allocToContigIDs() const {
+    return alloc_to_contig_ids_;
+  }
+
+ private:
+  // Indexing traversal graph.
+  const ValGraph& graph_;
+  // Domains to analyze contiguity. They are typically allocation
+  // domains but if this is a predicate indexing pass, they are
+  // likely logical domains.
+  const std::vector<IterDomain*> alloc_domains_;
+  // Contiguity of alloc_domains_
+  const std::vector<bool> alloc_contiguity_;
+  const bool is_predicate_pass_;
+  std::unique_ptr<const OrderedIdGroupInformation> consistent_transform_info_;
+
+  // Contig domain groups
+  std::unordered_set<ValGroup> contig_ids_;
+  // Mapping of allocation domains to contig groups
+  std::unordered_map<IterDomain*, ValGroup> alloc_to_contig_ids_;
+  // All domains that have dependencies with resize ops
+  std::unordered_set<ValGroup> resize_deps_;
+  // All domains that have dependencies with non-divisible split ops
+  std::unordered_set<ValGroup> non_divisible_deps_;
+};
+
+// Get a contiguous indexing domain for a given allocation domain. If
+// no such domain is found, just the allocation domain itself is
+// returned.
+std::unordered_map<IterDomain*, ValGroup> getContigDomains(
+    const std::vector<IterDomain*>& alloc_domains,
+    const std::vector<bool>& alloc_contiguity,
+    const ExprPath<ExprGroup>& path_from_alloc,
+    const ValGraph& graph,
+    bool is_predicate_pass);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/id_model.h

@@ -0,0 +1,359 @@
+// 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 <disjoint_set.h>
+#include <fusion.h>
+#include <ir/all_nodes.h>
+#include <val_graph.h>
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace nvfuser {
+
+class ValGraph;
+class LoopPromotionMapBuilderCallback;
+
+struct StatefulInliningInfo {
+  // All producer ids within (including dependencies of) inlined loop domains,
+  // used for deterministic order
+  VectorOfUniqueEntries<IterDomain*> ordered_p_ca_ids;
+
+  // p2c mappings through the fusion within (including dependencies of) inlined
+  // loop domains.
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<Val*>>
+      p2c_ca_permissive_maps;
+
+  // Broadcast resolution map for root domains, including non-inlined
+  // root domains
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<IterDomain*>>
+      p2c_root_broadcast_resolution_map;
+
+  // All IDs of all first siblings
+  VectorOfUniqueEntries<IterDomain*> ordered_sibling_ids;
+
+  // Mappings to other sibling IDs from ordered_sibling_ids
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<Val*>> sibling_maps;
+};
+
+StatefulInliningInfo buildStatefulInliningInfo(
+    const std::vector<Expr*>& exprs,
+    const ValGraph& exact_graph,
+    const ValGraph& permissive_graph);
+
+// A collection of ValGraphs that are built from a fusion or series of
+// expressions. These graphs are related, but have some distinct features based
+// on the IdMappingMode.
+//
+// EXACT/PERMISSIVE mode:
+//
+// consumer[i0, b1] = producer[i0]
+// consumer->merge(0) (consumer will now be [i0 * b1])
+//
+// When producer is replayed as consumer (the direction we use for mapping)
+// with forwarding from ForwardingInfo the producer to consumer map will have
+// both a mapping of consumer(i0) to producer(i0) as well as consumer(i0*b1) to
+// producer(i0). This latter mapping is important for loop nest mappings as the
+// consumer will generate a loop based on i0*b1 and the producer may be
+// computeAt inside this loop nest. However, for indexing we do not want these
+// two iter domains mapped as producer may be indexed as i0*i1 depending on the
+// loop nest structure and how it was built.
+//
+// Exact mode is if the iter domain relationships from producer to consumer are
+// considered the exact same size operating on matching dimensions from the root
+// domain mapping.
+//
+// LOOP mode is important to resolve inlined broadcassts. If we have something
+// like: consumer[i0o, threadIdx.x{i0i}] = producer[i0o,
+// threadIdx.y{i0i}](computeAt = 1) which can easily happen when using shared
+// memory. Loop is actually defined for all iteration domains, and resembles
+// groups of iter domains that are effectively inlined with each other.
+// Therefore iter domain's that are a common dependency of inlined loop domains
+// may be loop mapped together.
+//
+// Loop promotion is a mechanism by which to capture inlined resolved
+// broadcasts. If a consumer resolves a broadcast of a producer, and the
+// producer's broadcast is inlined (in total or partially). Then the producer's
+// iter domain will be "promoted" to the size of the consumers iter domain.
+//
+// IdMappingMode::EXACT
+//   Don't map any broadcast axes to non-broadcast axes
+//   Do not forward through any broadcast IDs
+// IdMappingMode::BROADCAST
+//   Map any broadcast axes to non-broadcast axes
+//   Do not forward through any broadcast IDs
+// IdMappingMode::PERMISSIVE
+//   Forward broadcast axes in replay
+//   Map all iteration domains
+//   Always contain root mappings (otherwise they could have been forwarded in
+//   broadcast)
+// IdMappingMode::ALMOSTEXACT
+//   Forward through broadcast axes, but not through to a non-broadcast axis
+//     i.e. id{b1*i0}, id{i0} are mapped
+//          id{i1*i0}, id{i0} are not mapped (this part is the difference from
+//          PERMISSIVE)
+//   Forward through split one axes, i.e. id{ceilDiv(i0, 1)}, id{i0} are mapped
+// IdMappingMode::LOOP
+//   Subgraph of the permissive graph. Maps only CA and their
+//   dependent domains.
+class IdModel : public PolymorphicBase {
+ public:
+  // Sometimes fusion inputs or outputs are disconnected from expressions, in
+  // those cases we still may want to send in some additional tensor views from
+  // the Fusion that don't have expressions associated with them.
+  //
+  // All graphs are built by default. It can be disabled with
+  // build_graphs=false.
+  IdModel(
+      const std::vector<Expr*>& exprs,
+      const std::vector<TensorView*>& additional_tvs = {},
+      bool build_graphs = true,
+      bool allow_self_mapping = false,
+      LoopPromotionMapBuilderCallback* loop_promotion_map_builder_callback =
+          nullptr);
+
+  // Same as the above constructor with fusion->exprs() excpet fusion may have
+  // some dangling inputs/outputs that are expected to have IterDomain entries
+  // even though there's no possible connections from them.
+  //
+  // The validate parameter is a temporary option during the
+  // transition from the current ComputeAtMap.
+  IdModel(
+      Fusion* fusion,
+      bool build_graphs = true,
+      bool allow_self_mapping = false,
+      bool validate = false,
+      LoopPromotionMapBuilderCallback* loop_promotion_map_builder_callback =
+          nullptr);
+
+  bool hasIdGraph(IdMappingMode mode) const {
+    return id_graphs_.find(mode) != id_graphs_.end();
+  }
+
+  // Returns iter domain graph of provided mode. The graph must have
+  // been already built.
+  const ValGraph& idGraph(IdMappingMode mode) const;
+  ValGraph& idGraph(IdMappingMode mode);
+
+  const std::unordered_map<IterDomain*, VectorOfUniqueEntries<Expr*>>& idUses()
+      const {
+    return id_uses_;
+  }
+
+  const std::unordered_map<IterDomain*, VectorOfUniqueEntries<Expr*>>&
+  idDefinitions() const {
+    return id_definitions_;
+  }
+
+  // TODO: Seems a bit unfortunate that this isn't IterDomain local information.
+  const std::unordered_set<IterDomain*>& viewRfactorIds() const {
+    return view_rfactor_ids_;
+  }
+
+  std::string toString() const;
+
+  bool empty() const {
+    return tvs_.empty();
+  }
+
+  const std::vector<TensorView*>& tvs() const {
+    return tvs_;
+  }
+
+  Fusion* fusion() const {
+    return fusion_;
+  }
+
+  // Build all graphs, i.e., Exact, AlmostExact, Permissive and
+  // LOOP. This is by default called from the constructor
+  void buildAllGraphs();
+
+  // Fills disjoint_ids_[IdMappingMode::EXACT] for relationships between inputs
+  // and first output of expr
+  ValGraph& buildExactGraph();
+
+  // Fills disjoint_ids_[IdMappingMode::ALMOSTEXACT]. Initialize AlmostExact as
+  // Exact entries, then map anything that's either merged with a size-1 or
+  // split by a size-1 dimension.
+  ValGraph& buildAlmostExactGraph();
+
+  // Fills disjoint_ids_[IdMappingMode::BROADCAST]. Initialize it as
+  // Exact entries, then map through broadcasts. Build the Exact graph
+  // as well if not yet done.
+  ValGraph& buildBroadcastGraph();
+
+  // Fills disjoint_ids_[IdMappingMode::PERMISSIVE]. Initialize it as
+  // BROADCAST entries, then map through forwarded domains. Build the
+  // BROADCAST graph as well if not yet done.
+  ValGraph& buildPermissiveGraph();
+
+  // Fills disjoint_ids_[IdMappingMode::LOOP]. Map only inlined
+  // domains that are mapped in the permissive graph. Build the Exact
+  // and Permissive graphs as well if not yet done.
+  //
+  // (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.
+  ValGraph& buildLoopGraph(bool force_full_loop_promotion_analysis = false);
+
+  // Build a graph. Dependent graphs are also built if not yet done.
+  void buildGraph(IdMappingMode mode);
+
+  // Build a graph if not already built
+  void maybeBuildGraph(IdMappingMode mode);
+
+  // Iterates over all IterDomains in id_definitions_ and calls initializeVal on
+  // a new ValGraph and returns it.
+  ValGraph initializeIdGraph(bool propagate_through_exprs = true) const;
+
+  // Returns an IdGraph with all Id's mapped that are mapped both in graph0 and
+  // graph1.
+  ValGraph buildIntersection(
+      const ValGraph& graph0,
+      const ValGraph& graph1,
+      bool propagate_exprs = true) const;
+
+  const std::unordered_map<ValGroup, IterDomain*>& loopPromotionMap() const {
+    return loop_promotion_map_;
+  }
+
+  // Replay Expr but with the inputs provided. ValGraphs will be updated
+  // for all maps that have entries, adding the output iter domains of the
+  // replayed expression and adding potential mappings through the expression.
+  Expr* addReplayAs(std::vector<IterDomain*> new_inputs, Expr* expr);
+
+  //! Run through disjoint sets in the LOOP graph, make sure there's only one
+  //! non-serial parallel type in each disjoint set, set the parallel type of
+  //! all IterDomains in the disjoint set to that PType.
+  void validateAndPropagatePType();
+
+  //! (Copied from ComputeAtMap::allocateIndexVariables)
+  //!  Run through disjoint sets in the LOOP map and allocate the index
+  //!  variable for the associated for loop that will be generated
+  //!  for each disjoint sets in the loop map. This pre-allocation makes
+  //!  2 key assumptions about computeAt map that would very likely be
+  //!  long term invariant:
+  //!    1. All kir::forloop created in the lowering pass should belong
+  //!  to one of the disjoint sets in loop map.
+  //!    2. The lowering pass will *never* create a loop nest with 2
+  //!  different nesting levels mapped together, i.e. the case below
+  //!  never occurs:
+  //!   for i in IterDomain1
+  //!    for j in IterDomain2
+  //!     ...
+  //!   With loop_map.areMapped(IterDomain1, IterDomain2) == true.
+  //! Under this condition, we can pre-allocate all required index
+  //!  variable integers before creating any kir::forloop, and this
+  //!  would help optimizing the generated integer math for indexing.
+  void allocateLoopIndexVariables();
+
+  // Get the index variable assigned for a given loop ID
+  Val* getLoopIndexVariable(
+      IterDomain* id,
+      CircularBufferLoopStage circular_buffer_loop_stage =
+          CircularBufferLoopStage::NotApplicable) const;
+
+  // Get the index variable assigned for a given loop group
+  Val* getLoopIndexVariable(
+      const ValGroup& loop_group,
+      CircularBufferLoopStage circular_buffer_loop_stage =
+          CircularBufferLoopStage::NotApplicable) const;
+
+ protected:
+  // Fills id_uses_ and id_definitions_ for all IterDomains active in the
+  // fusion.
+  void buildIterDomainDefinitionsAndUses();
+
+  // Start loop map by grouping inlined iter domains
+  void initializeLoopGraph(const StatefulInliningInfo& info);
+
+  // 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.
+  std::unordered_map<ValGroup, IterDomain*> buildLoopPromotionMap(
+      const StatefulInliningInfo& info);
+
+  // Errors if self mapping occurs
+  void assertNoSelfMapping();
+
+  // Loop graph represents the loop structure of the given fusion, so
+  // there must not be any mapping between the loop domains of each
+  // tensor.
+  void validateLoopGraphHasNoSelfMappedLeafDomains() const;
+
+ protected:
+  // Fusion where iter domains belong
+  Fusion* fusion_ = nullptr;
+
+  // All tensor expressions that this model analyzes
+  std::vector<Expr*> tv_exprs_;
+
+  // All tensors that this model analyzes
+  std::vector<TensorView*> tvs_;
+
+  // Tensors should not have domains that are mapped with another
+  // domains of the same tensor. This flag disables the check
+  bool allow_self_mapping_ = false;
+
+  // If true, validate graphs by comparing them with ComputeAtMap
+  bool validate_ = false;
+
+  // Optional callback for the loop promotion map builder for
+  // debugging and testing
+  LoopPromotionMapBuilderCallback* loop_promotion_map_builder_callback_ =
+      nullptr;
+
+  // By default, the permissive graph should map compliment domains as
+  // well. See the design doc for more details
+  bool permissive_graph_map_compliment_ids_ = true;
+
+  // Keeps ValGraphs containing all IterDomains for all mapping mode types.
+  //
+  // Using an array here might be nice, but it seems hard to use an enum as an
+  // array key
+  // https://stackoverflow.com/questions/2102582/how-can-i-count-the-items-in-an-enum
+  std::unordered_map<IdMappingMode, ValGraph> id_graphs_;
+
+  // If multiple transformations occur IterDomains could have multiple uses,
+  // however only one should be active in the given Fusion. When we resolve loop
+  // promotions during lowering, we can generate new iter domains from existing
+  // ones, so there can be multiple uses generated. Tracks all the active iter
+  // domain uses.
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<Expr*>> id_uses_;
+
+  // Make sure we don't blindly use definitions as we don't want to grab
+  // transformations before a tensor view's root domain. There can be
+  // multiple definitions due to replays.
+  std::unordered_map<IterDomain*, VectorOfUniqueEntries<Expr*>> id_definitions_;
+
+  std::unordered_set<IterDomain*> view_rfactor_ids_;
+
+  // Promotion domain for each loop group
+  std::unordered_map<ValGroup, IterDomain*> loop_promotion_map_;
+
+  // Allocated Loop index variable through the LOOP graph
+  std::unordered_map<ValGroup, Val*> loop_index_variable_map_;
+
+  // Allocated loop indices for circular buffer loops
+  std::unordered_map<
+      ValGroup,
+      std::unique_ptr<std::unordered_map<CircularBufferLoopStage, Val*>>>
+      circular_buffered_loop_index_variable_map_;
+};
+
+// A utility function to update a map of ValGroups to ID from an old
+// Valgraph to a new ValGraph. The new graph must be a superset of the
+// old graph.
+std::unordered_map<ValGroup, IterDomain*> updateValGroupIdMap(
+    const std::unordered_map<ValGroup, IterDomain*>& stale_map,
+    ValGraph& new_graph);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/id_model_index_compute.h

@@ -0,0 +1,81 @@
+// 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 <dispatch.h>
+#include <id_model/id_model.h>
+
+namespace nvfuser {
+
+// Similar to IndexCompute but adapted for the graph-based indexing
+class IdGraphIndexCompute : public OptOutDispatch {
+ public:
+  IdGraphIndexCompute(
+      const ValGraph& traversal_graph,
+      std::unordered_map<ValGroup, Val*> initial_index_map)
+      : traversal_graph_(traversal_graph),
+        index_map_(std::move(initial_index_map)) {}
+
+  // Propagate the index map through a given expr of a specified
+  // direction.
+  void propagate(const ExprGroup& expr_group, Direction direction) {
+    NVF_ERROR(!expr_group->empty());
+    // This looks a little ugly but the dispatch interface doesn't
+    // have a way to pass arguments
+    current_direction_ = direction;
+    dispatch(expr_group->front());
+    current_direction_ = Direction::Undefined;
+  }
+
+  const std::unordered_map<ValGroup, Val*> indexMap() const {
+    return index_map_;
+  }
+
+ private:
+  using OptOutDispatch::handle;
+
+  void handle(Split* split) override;
+
+  void handle(Merge* merge) override;
+
+  void handle(Swizzle* swizzle) override;
+
+  void handle(Resize* resize) override;
+
+  bool isForward(Expr* expr) const {
+    return current_direction_ == Direction::Forward;
+  }
+
+  bool hasIndex(IterDomain* id) const {
+    return indexMap().find(toGroup(id)) != indexMap().end();
+  }
+
+  Val* getIndex(IterDomain* id) const {
+    auto it = index_map_.find(toGroup(id));
+    NVF_ERROR(it != index_map_.end(), "Index not found: ", id->toString());
+    return it->second;
+  }
+
+  void setIndex(IterDomain* id, Val* idx) {
+    // May overwrite index. When the graph is cyclic due to, e.g.,
+    // resize, the index obtained by traversing most through the
+    // indexing path should be used (see also PR #3454)
+    index_map_[toGroup(id)] = idx;
+  }
+
+  const ValGroup& toGroup(IterDomain* id) const {
+    return traversal_graph_.toGroup(id);
+  }
+
+ private:
+  const ValGraph& traversal_graph_;
+  std::unordered_map<ValGroup, Val*> index_map_;
+  Direction current_direction_ = Direction::Undefined;
+};
+
+} // namespace nvfuser