nvfuser-cu121-torch25 0.2.25.dev20250201__cp310-cp310-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/device_lower/analysis/divisible_split.h

@@ -0,0 +1,28 @@
+// 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 <compute_at_map.h>
+#include <fusion.h>
+#include <ir/all_nodes.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+// Looks through all transformations assocaited with view, or enforced divisible
+// vectorization splits and gathers all splits that provably don't have a
+// remainder, therefore the extents of the associated IterDomains do not require
+// a ceilDiv expressions.
+NVF_API std::unordered_set<Split*> getAllDivisibleSplits(Fusion* fusion);
+
+// Same as above but will use provided ComputeAtMap instead of building its own.
+NVF_API std::unordered_set<Split*> getAllDivisibleSplits(
+    Fusion* fusion,
+    const ComputeAtMap* ca_map);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/fused_reduction.h

@@ -0,0 +1,36 @@
+// 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 <exceptions.h>
+#include <ir/all_nodes.h>
+
+namespace nvfuser {
+
+//! Keep track of certain patterns of reductions.
+//!
+//! - Allreduce IterDomain: reduced and broadcast domain.
+class FusedReductionInfo {
+ public:
+  void markAsAllreduce(IterDomain* id);
+
+  bool isAllreduce(IterDomain* id) const;
+
+ private:
+  // Reduction IterDomains that are also broadcast
+  std::unordered_set<IterDomain*> allreduce_ids_;
+};
+
+//! Detect reductions and broadcasts that are eligible for the fused
+//! reduction kernel. When found, the predicate flags of the broadcast
+//! is unset, which effectively makes the broadcast just a unary set
+//! op.
+//! TODO: Consider moving the warp-based fused reduction here.
+void fuseReductionsAndBroadcasts(Fusion*);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/index_compute.h

@@ -0,0 +1,322 @@
+// 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 <exceptions.h>
+#include <fusion.h>
+#include <index_compute.h>
+
+namespace nvfuser {
+
+// Struct to hold useful information from an index pass on iterdomain graph.
+// Used to return the IndexCompute structure back to the indexing calls in
+// index_compute.cpp. Other structurs are required to resolve the actual
+// indexing math there.
+struct IndexFromIdGraph {
+  IndexCompute index;
+  IndexCompute concrete_index;
+  std::unordered_map<IterDomain*, Val*> initial_concrete_index_map;
+  std::vector<IterDomain*> resolved_loop_domains;
+
+  explicit IndexFromIdGraph(
+      IndexCompute index,
+      IndexCompute concrete_index,
+      std::unordered_map<IterDomain*, Val*> initial_concrete_index_map,
+      std::vector<IterDomain*> loop_domains);
+};
+
+//! Indexing interface, returns IndexFromIdGraph which the IndexCompute object
+//! can be queried from directly for the produced indexing. If producer_tv !=
+//! nullptr producer will be indexed, if producer_tv == nullptr consumer will be
+//! indexed. If is_global global indexing will be done, else shared memory or
+//! local indexing will be performed.
+IndexFromIdGraph getTensorIndexFromIdGraph(
+    const std::vector<ForLoop*>& loops,
+    const std::unordered_set<ForLoop*>& rotated_loops,
+    const TensorView* consumer_tv,
+    const TensorView* producer_tv = nullptr,
+    bool is_global = true,
+    const std::unordered_map<IterDomain*, IterDomain*>& c2p_map = {});
+
+//! Indexing interface for calculating predicate index returns IndexFromIdGraph
+//! which the IndexCompute object can be queried from directly for the produced
+//! indexing If is_start_predicate, will produce indexing math for the start
+//! predicates.
+IndexFromIdGraph getPredicateIndexingFromIdGraph(
+    const std::vector<ForLoop*>& loops,
+    const std::unordered_set<ForLoop*>& rotated_loops,
+    TensorView* consumer_tv,
+    ForLoop* unswitch_or_vec_loop,
+    IterDomain* circular_buffer_axis,
+    bool is_start_predicate);
+
+//! getTensorIndexFromIdGraph is the function that index_compute will call very
+//! straightforwardly. However, for implementing the new indexing logic that
+//! starts to abstract some of the indexing away from index_compute we need to
+//! move quite a bit of the intertwined indexing logic away from the
+//! index_compute file and the index_reference_replay file. This is because we
+//! want to separate out what has to be done on the fly, from what analysis we
+//! can do early on with the iter domain graph and associated properties.
+//!
+//! getTensorIndexFromIdGraph places this analysis internally in
+//! LoopIndexingAnalysis. LoopIndexingAnalysis though has to communicate to:
+//!   1) index_compute.cpp::IndexCompute to tell IndexCompute which expressions
+//!   it needs to traverse to compute the indexing math.
+//!
+//! LoopIndexing is nothing but a mechanism for this communication.
+//!
+//! Holds information needed to produce indexing math. In the current version of
+//! indexing pass, the iter domains combined with the loop nests are the source
+//! of truth in terms of resolving the actual integer indexing math from the
+//! sequence of iterdomain transforms.
+//!
+//! This information is crtiical in resolving indexing associated with complex
+//! broadcast patterns. Check FusionComplexBCast* test cases as well as
+//! Indexing* tests for examples where resolving indices from IterDomain
+//! transformations can be challenging.
+//!
+//! The source of this challenge is due to inling patterns where the IterDomains
+//! responsible for control flow are not local to a particular TensorView.
+//! Broadcast, operations like view/reshape, and gather/shift can make indexing
+//! local buffers complex because of the complex effects inlining into other
+//! TensorViews produce.
+//!
+//! TODO:
+//!   The first iteration tries to match the semantics of reference
+//!     replay without any new logic. In a follow up iteration will
+//!     need to revisit a few further pathological patterns.
+//!
+//! Note:
+//!   The current implementation of loop indexing pass works on
+//! equivalent classes defined by ComputeAt exact map. The
+//! list of expressions stored in this class form a "reference", graph of
+//! iterdomain expressions when all of their inputs and outputs are replaced
+//! with their exact concrete mapped id's.
+//!
+//!   Here an invariant in a graph of iterdomain expressions is that
+//! each iterdomain is produced exactly once and is either a loop domain
+//! or has been consumed exactly once by another expression. This makes sure
+//! that a well defined indexing can be generated for each of the concrete ids
+//! whenever we either forward or backward traverse the graph.
+class LoopIndexing {
+ public:
+  //! Returns the original loop nest.
+  const auto& loops() const {
+    return loops_;
+  }
+
+  //! Returns the vector of Iterdomains
+  //!  that match the original loop pattern.
+  const auto& loopDomains() const {
+    return loop_domains_;
+  }
+
+  const auto& loopRootDomains() const {
+    return loop_root_;
+  }
+
+  //! Returns the consumer tv that the view info
+  //!  was derived from.
+  auto consumerTv() const {
+    return consumer_tv_;
+  }
+
+  //! Returns the set of Iterdomain transforms that
+  //!  define the correct indexing path, in forward
+  //!  topological order.
+  std::vector<Expr*> getForwardExprList() const;
+
+  //! Returns the set of Iterdomain transforms that
+  //!  define the correct indexing path, in backward
+  //!  topological order.
+  std::vector<Expr*> getBackwardExprList() const;
+
+  //! Returns the set of out of line expressions in
+  //!  reverse topological order.
+  const std::vector<Expr*>& getBackwardOutOfLineExprList() const {
+    return out_of_line_exprs_;
+  }
+
+  //! Returns all exact concrete id's that were produced
+  //!  or consumed in the selected indexing expressions
+  std::unordered_set<IterDomain*> getAllExactConcreteIdSet() const;
+
+ private:
+  friend class LoopIndexingAnalysis;
+
+  //! The loop nest that this loop indexing is derived from.
+  std::vector<ForLoop*> loops_;
+
+  //! Consumer tv, where the view related info was derived from.
+  const TensorView* consumer_tv_ = nullptr;
+
+  //! The source iterdomains that all the Iterdomain transforms
+  //!   in this loop nest originated from.
+  std::vector<IterDomain*> loop_root_;
+
+  //! The loop iterdomains that the original loop nests correspond
+  //!  to. May be longer than loops_ with the dangling iterdomains
+  //!  appended towards the end.
+  std::vector<IterDomain*> loop_domains_;
+
+  //! The selected sequence of expressions that should represent
+  //!  the correct indexing math from the given loop nest.
+  std::vector<Expr*> index_exprs_;
+
+  //! The subset of sequence of expressions that can be resolved
+  //!  with only the iterdomains on the right of consumer tv's ca
+  //!  axis.
+  //! Expressions are ordered in reverse topological order.
+  std::vector<Expr*> out_of_line_exprs_;
+};
+
+class LoopIndexingAnalysis {
+ public:
+  static LoopIndexing fromLoopAndConsumer(
+      const std::vector<ForLoop*>& loops,
+      const TensorView* consumer_tv);
+
+  //! Return all concrete IDs that can be reachable from a given list
+  //! of consumer loop IDs. Reachability is defined as the existence
+  //! an indexing path from the the loop IDs
+  static VectorOfUniqueEntries<IterDomain*> getReplayableConcreteIDs(
+      const std::vector<IterDomain*>& consumer_loop_ids,
+      const TensorView* consumer_tv);
+
+ private:
+  explicit LoopIndexingAnalysis(
+      const std::vector<ForLoop*>& loops,
+      const TensorView* consumer_tv);
+
+  explicit LoopIndexingAnalysis(
+      const std::vector<IterDomain*>& consumer_loop_ids,
+      const TensorView* consumer_tv);
+
+  void run();
+
+  //! Populate derived information into a LoopIndexing
+  //!  data structure.
+  LoopIndexing getLoopIndexing(const std::vector<ForLoop*>& loops) {
+    LoopIndexing indexing;
+    indexing.loops_ = loops;
+    indexing.consumer_tv_ = consumer_tv_;
+    indexing.loop_root_ = loop_root_domains_;
+    indexing.loop_domains_ = loop_domains_.vector();
+    indexing.index_exprs_ = replayed_exprs_;
+    indexing.out_of_line_exprs_ = out_of_line_exprs_;
+    return indexing;
+  }
+
+  //! Validates that the current loop structure is well formed, in the sense
+  //! that ca_map would not map any two loops in the loop nest together.
+  void validateLoopStructure(const std::vector<ForLoop*>& loops);
+
+  //! Start at the loop iter domains, and traverse back into history on the
+  //! concrete IDs in the exact map calling "visitExpr" expressions through the
+  //! history.
+  void traverseFromDomainVals();
+
+  //! Concretize the given iterdomain and record the visit (in deterministic
+  //! order) in terms of the exact mapped concrete id. Marks the mapping of the
+  //! id to the concrete id in "concrete_to_original_id_" and returns the
+  //! concrete id.
+  IterDomain* concretizeAndVisitId(IterDomain* id);
+
+  //! If an equivalent expression has already been processed this function
+  //! simply returns. Otherwise puts the exact concrete IDs of inputs in
+  //! consumed_concrete_, and concrete IDs of outputs in produced_concrete_.
+  //! Then adds the expression to replayed_exprs_.
+  void visitExpr(Expr* expr);
+
+  //! Iterates through provided vals, calls concretizeAndVisitId on them, and
+  //! returns if any of the returned vals are in existing_ids. This is used to
+  //! check if inputs or outputs of ID expressions have already been
+  //! produced/consumed in the traversal. Indexing only needs to consume/produce
+  //! one IterDomain per exact disjoint set.
+  bool visitIdsAndCheckDuplication(
+      const std::vector<Val*>& vals,
+      const std::unordered_set<IterDomain*>& existing_ids);
+
+  //! Fills loop_domains_ with the corresponding replayed_concrete_id mapping to
+  //! the provided loops. Must be done after the exact iterdomain "replay"
+  //! (traverseFromDomainVals). loop_domains_ are the original_id not the
+  //! concrete_id (translated with concrete_to_original_id). These iter domains
+  //! are used to grab the history that will be replayed in IndexCompute. We're
+  //! looking for "new" root domains and subsequent transformations, filling in
+  //! any missing "outputs" (or inputs for backward traversal). Then fills
+  //! loop_domains_ with all of these iter domains.
+  void constructLoopDomains();
+
+  //! Fills out_of_line_exprs_ by traversing the selected list of
+  //!  expressions in reverse topological order and collect iterdomains
+  //!  on the indexing paths that only involves loop id's on the right
+  //!  of consumer's ca axis.
+  void collectOutOfLineExprs();
+
+ private:
+  //! Original consumer tv to derive view info from.
+  const TensorView* consumer_tv_ = nullptr;
+
+  // Exact concrete domains that has been used
+  //  in the traversal connection.
+  std::unordered_set<IterDomain*> produced_concrete_;
+  std::unordered_set<IterDomain*> consumed_concrete_;
+
+  //! Iterdomains that the corresponding loops are generated from.
+  std::vector<IterDomain*> initial_loop_domain_ids_;
+
+  //! All Id's in consumer's transform history
+  std::vector<Val*> all_consumer_id_vals_;
+
+  //! Concrete iterdomains visited in the domain traversal,
+  //!  in the order they are visited in traverseFromDomainVals.
+  VectorOfUniqueEntries<IterDomain*> replayed_concrete_ids_;
+
+  //! Keeping track of the original visited id's before they
+  //!  were concretized.
+  std::unordered_map<IterDomain*, IterDomain*> concrete_to_original_id_;
+
+  //! Map from concrete id to its single consumer on the selected
+  //!  iterdomain expression list.
+  std::unordered_map<IterDomain*, Expr*> concrete_id_to_consumer_;
+
+  //! Source domains that all the Iterdomain transforms
+  //!  in the loop nest originated from.
+  std::vector<IterDomain*> loop_root_domains_;
+
+  //! Leaf domains representing the original loop structure
+  VectorOfUniqueEntries<IterDomain*> loop_domains_;
+
+  //! Selected list of exprs that will produce and consume each
+  //!  of the exact concrete ids from the loop nest exactly once.
+  std::vector<Expr*> replayed_exprs_;
+
+  //! Set of expressions from the selected list that can be
+  //!  resolved from axes on the right of ca axes.
+  std::vector<Expr*> out_of_line_exprs_;
+};
+
+// When indexing there are sometimes an option to propagate an index down
+// multiple paths. This will return the IterDomains in the history of the
+// reference domain and mark which paths should be taken (if there's a
+// preference) to reach the roots provided in preferred_roots.
+std::unordered_set<IterDomain*> buildLoopIndexingPreferredPath(
+    const TensorView* original_tv,
+    const LoopIndexing& loop_indexing,
+    bool use_replay_map = false,
+    std::unordered_map<IterDomain*, IterDomain*> p2c_map = {});
+
+// Get an logical IterDomain that is mapped with an IterDomain. If
+// multiple such IDs exist, select one whose input IDs are mapped with
+// the consumer IDs. This is to ensure the path from the loop
+// IterDomains to the root matches with the consumer tensor.
+IterDomain* getLogicalIDToTraverse(
+    IterDomain* id,
+    const std::vector<Val*>& consumer_all_ids);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/predicate_elimination.h

@@ -0,0 +1,71 @@
+// 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 <exceptions.h>
+
+#include <ir/all_nodes.h>
+#include <kernel_ir.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+class PredicateElimination : public IterVisitor {
+ public:
+  PredicateElimination(Fusion* fusion);
+
+  //! True if expr does not need a predicate
+  //!
+  //! \param expr Tensor expression
+  bool canOmitPredicate(const Expr* expr) const;
+
+  bool needsSharedMemoryPredicate(const Expr* expr) const;
+
+  //! Value to initialize out-of-bound regions
+  Val* getInitValue(TensorView* tv) const;
+
+  //! Dump to string for debugging
+  std::string toString() const;
+
+  // A utility to set removal info of `to` the same as `from`.
+  //  See issue #1641
+  // We build predicate info before lowering but more expressions
+  //  are created during lowering that this class also need to
+  //  keep track of to make sure correct predicate removal is
+  //  applied.
+  // This utility is a quick patch for the missing information
+  //  since it might be better just to recompute predicate info
+  //  if all expressions were mutated, but that'd take much more
+  //  global info to reliably track.
+  void propagateRemovalInfo(const Expr* from, const Expr* to);
+
+  const std::unordered_set<const Expr*>& getNonPredicatedExprs() const {
+    return non_predicated_exprs_;
+  }
+
+ private:
+  using IterVisitor::handle;
+
+  void dispatch(Expr* expr) final;
+
+  //! Set a value to initialize out-of-bound regions
+  bool setDefaultInitValue(TensorView* tv);
+  //! Set a value to initialize out-of-bound regions of reduction tensors
+  bool setReductionInitValue(TensorView* tv, Val* reduction_init);
+
+  //! Check if expr needs to be predicated
+  bool needsPredicate(Expr* expr) const;
+
+ private:
+  //! Expressions that are found to be safe without predicates
+  std::unordered_set<const Expr*> non_predicated_exprs_;
+  //! Tensors and their initialization values
+  std::unordered_map<TensorView*, Val*> init_value_map_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/sync_information.h

@@ -0,0 +1,47 @@
+// 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 <exceptions.h>
+#include <ir/all_nodes.h>
+#include <parallel_type_bitmap.h>
+#include <visibility.h>
+
+#include <unordered_map>
+
+namespace nvfuser {
+
+class SyncMap {
+ public:
+  //! Validates all tensors are consistently parallelized. Basically,
+  //! when a producer axis is threaded, either with threadIdx or
+  //! blockIdx, there must be a mapped consumer axis with the
+  //! same ParallelType with some exceptions.
+  //!
+  //! ComputeAtMap is already built as they are used to validate consistency.
+  //!
+  //! Fills needs_raw_sync with output TVs if they need a raw sync if on smem or
+  //! gmem. The second entry in this map is the parallel dimensions being
+  //! communicated across.
+  NVF_API SyncMap(Fusion* fusion);
+
+  std::string toString() const;
+
+  ParallelTypeBitmap needsRawSync(TensorView* tv) const {
+    auto it = needs_raw_sync_.find(tv);
+    if (it != needs_raw_sync_.end()) {
+      return it->second;
+    }
+    return ParallelTypeBitmap();
+  }
+
+ private:
+  std::unordered_map<TensorView*, ParallelTypeBitmap> needs_raw_sync_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/tensor_memory.h

@@ -0,0 +1,65 @@
+// 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
+
+namespace nvfuser {
+
+class Fusion;
+
+// Information used to lower tensor memory. So far, there is no information
+// needed, the computeTMemInfo just check that there is only one tensor on TMem
+// in the fusion. This limitation is described in the note below, and it is only
+// for incremental development. This limitation will be removed soon in the
+// future.
+struct TensorMemoryInfo;
+TensorMemoryInfo computeTMemInfo(Fusion* fusion);
+
+// Note: [Tensor Memory Allocation]
+//
+// Tensor memory is a very special memory, so its allocation is also very
+// different from other memory types.
+//
+// It is highly recommended to read the PTX documentation for tensor memory
+// if you are not alreay familiar with it:
+// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#tensor-memory
+//
+// The first thing to note is, TMem does not have virtualization. This means:
+// We can not just allocate starting from address 0 like how we allocate shared
+// memory, and rely on page table to translate the same virtual address of
+// different CTA to different physical address. There is no virtual TMem
+// address. All addresses are physical addresses.
+//
+// Because multiple CTAs can execute on the same SM simultaneously, there must
+// be some handshaking mechanism for each CTA to know the region of TMem that it
+// can use. This is done by using the PTX instruction tcgen05.alloc. To ensure
+// safety, there is a mutex "I have the right to allocate TMem" in the
+// hardware. At the beginning of each CTA, the CTA will try to acquire the mutex
+// automatically. If it fails, the CTA will be blocked until the mutex is free.
+// This means, only one CTA can allocate TMem at a time. Once the CTA has
+// finished allocating TMem, it should release the mutex to relinquish the right
+// to allocate. After the right to allocate is relinquished, this CTA can not
+// allocate new TMem any more, but it can still access the TMem that it has
+// allocated, and it can also free the TMem that it has allocated. Once one CTA
+// relinquishes the right to allocate, the next CTA that is blocked will be
+// unblocked and can acquire the mutex to allocate TMem.
+//
+// Currently, the TMem allocation is not supported in nvFuser. We currently only
+// allow one TensorView to be on TMem, and because we never relinquish the right
+// to allocate TMem, CTA will be serialized on SM. A new CTA can be scheduled on
+// an SM only after the previous CTA on that SM has completely finished
+// executing. Thanks to this serialization, we can just skip allocating and
+// think that our only TMem TensorView own the entire TMem, because we are sure
+// that there will not be another CTA using that address. As a result, we could
+// just provide address 0 to our instructions that access TMem. In principle, it
+// is clearly wrong to write to an address that is not allocated, but because we
+// are sure that it will in practice work for the specific unit test that we are
+// targeting, we just do it so we have incremental development.
+
+struct TensorMemoryInfo {};
+
+} // namespace nvfuser