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

nvfuser/include/nvfuser/device_lower/analysis/thread_predicate.h

@@ -0,0 +1,158 @@
+// 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 <visibility.h>
+
+#include <device_lower/utils.h>
+#include <ir/all_nodes.h>
+#include <parallel_type_bitmap.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+namespace nvfuser {
+
+//! Maps TensorViews to a { ParallelTypeBitmap, SourceMap } pair
+//!
+//! Map from TensorView to bit set represnting <BIDx, BIDy, BIDz, TIDx, TIDy,
+//! TIDz> If any dependency of TV had a parallelized reduction, we will track
+//! it here. This will be used for predicate generation to prevent
+//! parallelization on that axis. This is important if we have a reduction on
+//! for example TIDx, as the reduced value is only valid on threadIdx.x == 0
+//! therefore if we use that value later in the kernel we have that predicate.
+//! If we follow a reduction parallelized on TIDx with a broadcast on TIDx we
+//! no longer need the predicate and can reset the bit accordingly
+//!
+//! In addition, if a parallel thread type is not used, it is
+//! redundant to use all threads/blocks. That isn't a problem
+//! generally although it can be inefficient, but when an aliased smem
+//! buffer is used as an output, redundant writes can be invalid (see issue
+//! #1110). PredicateInfo::redundant_types track which parallel types
+//! are redundant for each tensor and is used to let only one
+//! thread/block of a redundant type execute the expression for a
+//! tensor.
+class ThreadPredicateMap {
+ public:
+  using SourceMap =
+      std::unordered_map<ParallelType, std::unordered_set<const TensorView*>>;
+
+  //! Thread predicate information for each tensor
+  struct PredicateInfo {
+    // Parallel types where only one thread/block is valid.
+    ParallelTypeBitmap limited_types;
+    // Parallel types where only one thread/block is enough.
+    ParallelTypeBitmap redundant_types;
+
+    // when a loop domain of a Tensor stored in global memory
+    // is merged from concretized broadcast logical domain, the broadcasted
+    // logical domains should be skipped when writing to global memory.
+    // broadcast_ld_indices_map maps a parallel type to a list of indices
+    // of the broadcasted logical domains. The write to global memory is needed
+    // only when the index equals to 0.
+    std::unordered_map<ParallelType, std::vector<Val*>>
+        broadcast_ld_indices_map;
+
+    // Tracking use chain of redundant writes:
+    //  [Redundant use chain]
+    //  a parallel type is a `redundant_consumer_type` only
+    //    if all of its propagation use chains terminate with
+    //    a redundant write of this type.
+    //  A propagation use chain is currently either a reg-to-reg
+    //   chain for a shared mem tv, or a reg/smem-to-reg/smem chain
+    //   for a global tv.
+    // This is complementary information to `redundant_types`.
+    //  If a tensor view is redundantly written and not redundantly
+    //  used by all consumers, see FusionRedundantPredSync3,
+    //  a RAW sync will need to be inserted before reading
+    //  this redundantly written tensor.
+    ParallelTypeBitmap redundant_use_types;
+    bool operator==(const PredicateInfo& other) const {
+      return limited_types == other.limited_types &&
+          redundant_types == other.redundant_types &&
+          redundant_use_types == other.redundant_use_types;
+    }
+  };
+
+  using MapType = std::unordered_map<const TensorView*, PredicateInfo>;
+
+  using const_iterator = MapType::const_iterator;
+
+  //! Build a map from each tensor to PredicateInfo.
+  void build(Fusion* fusion);
+
+  //! Get a PredicateInfo for a given tensor. If it's an output of
+  //! a parallel broadcast, unmask the limited_types_ bit of the
+  //! corresponding parallel type since it must join the broadcast
+  //! operation although the valid input is only available at one of
+  //! the threads/blocks.
+  NVF_API PredicateInfo getPredicateInfo(const TensorView* tv) const;
+
+  //! Returns a flag set that indicates which parallel types should be
+  //! predicated.
+  ParallelTypeBitmap getPredicatedParallelTypes(const TensorView* tv) const;
+
+  //! Returns a Bool predicate for a given TensorView.
+  Val* getPredicate(
+      const TensorView* tv,
+      ParallelTypeBitmap mask = ParallelTypeBitmap().setAll()) const;
+
+  //! Returns a ParallelTypeBitmap representing which domain needs
+  //! blockBroadcast.
+  //!
+  //! Even when a domain is broadcast and parallelized, it does not need
+  //! blockBroadcast unless it is predicated by limited_types_
+  ParallelTypeBitmap getParallelBroadcastDomains(const TensorView* tv) const;
+
+  //! Mark tv as updated so that rebuilding the map should recompute
+  //! its predicates and those of its dependents.
+  void markAsUpdated(const TensorView* tv);
+
+  void print() const;
+
+  //! Generate a Bool value from PredicateInfo.
+  static Val* getPredicateFromPredicateInfo(
+      const ThreadPredicateMap::PredicateInfo& pred_info,
+      const ParallelTypeBitmap& mask);
+
+  //! Get the redundant use types of the given expr, see [Redundant use chain]
+  ParallelTypeBitmap getRedundantConsumerType(Expr* expr) const;
+
+ private:
+  // Update the thread_predicates bitset based on provided Expr
+  void updateBitSet(const Expr*);
+  void avoidConcretizedBroadcastRedundantWrite(const TensorView* out_tv);
+  const_iterator find(const TensorView* tv) const;
+  const_iterator end() const;
+
+  const PredicateInfo& at(const TensorView* tv) const;
+  PredicateInfo& at(const TensorView* tv);
+
+  //! Update a mapping
+  bool update(
+      const TensorView* tv,
+      const ParallelTypeBitmap& limited_types,
+      const ParallelTypeBitmap& redundant_types);
+
+  //! Update a mapping
+  bool update(const TensorView* tv, const PredicateInfo& pred_and_src);
+
+  //! Backward populate redundant use chain info once the redundant
+  //!  parallel writes have been identified.
+  void populateRedundantUseMap(Fusion* fusion);
+
+ private:
+  MapType thread_predicates_;
+  //! Keep track of updated tensors that need predicates to be computed
+  std::unordered_set<const TensorView*> updated_tvs_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/tma.h

@@ -0,0 +1,93 @@
+// 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 <ostream>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+#include <ir/all_nodes.h>
+#include <val_graph.h>
+
+namespace nvfuser {
+
+// See doc/dev/tma.md for design
+
+// All ValGroups are in the traversal graph of tensor indexer
+
+struct TMADim {
+  ValGroup partitioned;
+  ValGroup box;
+  ValGroup tile;
+  ValGroup stride;
+  Val* gmem_stride_bytes;
+
+  Val* tensorSize() const {
+    return partitioned->front()->as<IterDomain>()->extent();
+  }
+  Val* boxSize() const {
+    return box ? box->front()->as<IterDomain>()->extent()
+               : gmem_stride_bytes->fusion()->oneVal();
+  }
+  Val* tileSize() const {
+    return tile ? tile->front()->as<IterDomain>()->extent()
+                : gmem_stride_bytes->fusion()->oneVal();
+  }
+  Val* elementStride() const {
+    return stride ? stride->front()->as<IterDomain>()->extent()
+                  : gmem_stride_bytes->fusion()->oneVal();
+  }
+};
+
+std::ostream& operator<<(std::ostream& os, const TMADim& d);
+
+class TMAInfo {
+  std::vector<TMADim> dims_;
+  MmaInputSmemSwizzle swizzle_;
+  TensorView* gmem_tv_;
+
+ public:
+  TMAInfo(
+      std::vector<TMADim> dims,
+      MmaInputSmemSwizzle swizzle,
+      TensorView* gmem_tv)
+      : dims_(std::move(dims)), swizzle_(swizzle), gmem_tv_(gmem_tv) {}
+
+  const std::vector<TMADim>& dims() const {
+    return dims_;
+  }
+
+  std::vector<ValGroup> getTMADomain() const {
+    std::vector<ValGroup> result;
+    std::transform(
+        dims_.begin(),
+        dims_.end(),
+        std::back_inserter(result),
+        [](const auto& d) { return d.partitioned; });
+    return result;
+  }
+
+  Val* tileSizeBytes() const {
+    int64_t itemsize = dataTypeSize(gmem_tv_->dtype());
+    Val* size = IrBuilder::create<Val>(itemsize, DataType::Index);
+    for (const auto& d : dims_) {
+      size = SimplifyingIrBuilder::mulExpr(size, d.tileSize());
+    }
+    return size;
+  }
+
+  Val* tensorMap() const;
+};
+
+std::unordered_map<TensorView*, const TMAInfo> getConsumerToTMAInfoMap(
+    Fusion* fusion);
+
+MmaInputSmemSwizzle getSwizzle(TensorView* tv);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/analysis/trivial_broadcast.h

@@ -0,0 +1,75 @@
+// 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 <logical_domain_map.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+//! Traverse and collect all concretized broadcast domains.
+//!
+//! The traversal first initializes the origin map with broadcast
+//! domains in input tensors. Then, a new entry is added to the origin
+//! map when a broadcast op is encountered during a forward traversal
+//! of the given fusion. For non-broadcast ops, mappings are just
+//! propagated forward using PairwiseLogicalDomainMap.
+//!
+//! When the mapped consumer domain is not broadcast, it means the
+//! producer broadcast domain is concretized, and its origin broadcast
+//! domains are marked as concretized.
+class NVF_API ConcretizedBroadcastDomains : private IterVisitor {
+ public:
+  ConcretizedBroadcastDomains() = delete;
+  ConcretizedBroadcastDomains(Fusion* fusion);
+
+  //! Is a domain concretized?
+  bool isConcretized(IterDomain* id) const;
+
+  //! Is a domain concretized to a unique concrete domain?
+  bool isUniquelyConcretized(IterDomain* id) const;
+
+  //! Is a domain concretized to multiple concrete domains?
+  bool maybeNonUniquelyConcretized(IterDomain* id) const;
+
+  //! Return all domains id is concretized to, if concretized
+  std::unordered_set<IterDomain*> allConcretizedDomains(IterDomain* id) const;
+
+ private:
+  using IterVisitor::handle;
+
+  void handle(TensorView* tv) final;
+
+  void handle(BroadcastOp* bop) final;
+
+  void dispatch(Expr* expr) final;
+
+  void markAsConcretized(
+      IterDomain* broadcast_root_domain,
+      IterDomain* concrete_root_domain);
+
+  bool insertRootDomainToConcreteDomainSet(
+      IterDomain* new_root_id,
+      std::unordered_set<IterDomain*>& id_set);
+
+ private:
+  //! Maps each root broadcast domain to its original root broadcast
+  //! domains. Their can be multiple original domains due to, e.g.,
+  //! binary ops with broadcast domains in both inputs.
+  std::unordered_map<IterDomain*, std::unordered_set<IterDomain*>>
+      broadcast_origin_map_;
+  //! Map all broadcast domains to concrete root domains
+  std::unordered_map<IterDomain*, std::unordered_set<IterDomain*>>
+      broadcast_to_concrete_map_;
+
+  std::unique_ptr<ExactLogicalDomainMap> exact_map_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/id_model_options.h

@@ -0,0 +1,135 @@
+// 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 <id_model/utils.h>
+
+#include <sstream>
+
+namespace nvfuser {
+
+class IdModelOptions {
+ public:
+  IdModelOptions()
+      : build_id_model_(isOptionEnabled(EnableOption::IdModel)),
+        consumer_index_(
+            isIdModelOptionEnabled(IdModelEnableOption::ConsumerIndex)),
+        producer_index_(
+            isIdModelOptionEnabled(IdModelEnableOption::ProducerIndex)),
+        inline_predicate_(
+            isIdModelOptionEnabled(IdModelEnableOption::InlinePredicate)),
+        unswitch_predicate_(
+            isIdModelOptionEnabled(IdModelEnableOption::UnswitchPredicate)),
+        loop_(isIdModelOptionEnabled(IdModelEnableOption::Loop)) {
+    ensureConsistency();
+  }
+
+  bool buildIdModel() const {
+    return build_id_model_;
+  }
+
+  void setBuildIdModel(bool b) {
+    build_id_model_ = b;
+    ensureConsistency();
+  }
+
+  bool buildTensorIndexer() const {
+    return build_tensor_indexer_;
+  }
+
+  void setBuildTensorIndexer(bool b) {
+    build_tensor_indexer_ = b;
+    ensureConsistency();
+  }
+
+  bool consumerIndex() const {
+    return consumer_index_;
+  }
+
+  void setConsumerIndex(bool b) {
+    consumer_index_ = b;
+    ensureConsistency();
+  }
+
+  bool producerIndex() const {
+    return producer_index_;
+  }
+
+  void setProducerIndex(bool b) {
+    producer_index_ = b;
+    ensureConsistency();
+  }
+
+  bool inlinePredicate() const {
+    return inline_predicate_;
+  }
+
+  void setInlinePredicate(bool b) {
+    inline_predicate_ = b;
+    ensureConsistency();
+  }
+
+  bool unswitchPredicate() const {
+    return unswitch_predicate_;
+  }
+
+  void setUnswitchPredicate(bool b) {
+    unswitch_predicate_ = b;
+    ensureConsistency();
+  }
+
+  bool loop() const {
+    return loop_;
+  }
+
+  void setLoop(bool b) {
+    loop_ = b;
+    ensureConsistency();
+  }
+
+  std::string toString() const {
+    auto bool2str = [](bool b) { return b ? "true" : "false"; };
+
+    std::stringstream ss;
+    ss << "build_id_model=" << bool2str(build_id_model_)
+       << ", build_tensor_indexer=" << bool2str(build_tensor_indexer_)
+       << ", consumer_index=" << bool2str(consumer_index_)
+       << ", producer_index=" << bool2str(producer_index_)
+       << ", inline_predicate=" << bool2str(inline_predicate_)
+       << ", unswitch_predicate=" << bool2str(unswitch_predicate_)
+       << ", loop=" << bool2str(loop_);
+    return ss.str();
+  }
+
+ private:
+  void ensureConsistency() {
+    // TensorIndexer is required if these options are enabled
+    build_tensor_indexer_ = build_tensor_indexer_ || consumer_index_ ||
+        producer_index_ || inline_predicate_ || unswitch_predicate_ || loop_;
+    // Similarly, IdModel needs to be built if TensorIndexer is used
+    build_id_model_ = build_id_model_ || build_tensor_indexer_;
+  }
+
+ private:
+  // Build IdModel
+  bool build_id_model_ = false;
+  // Build TensorIndexer
+  bool build_tensor_indexer_ = false;
+  // Globally enables consumer indexing.
+  bool consumer_index_ = false;
+  // Globally enables producer indexing.
+  bool producer_index_ = false;
+  // Globally enables inline predicate
+  bool inline_predicate_ = false;
+  // Globally enables unswitch predicate
+  bool unswitch_predicate_ = false;
+  // Generate loops using IdModel
+  bool loop_ = false;
+};
+
+} // namespace nvfuser