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

nvfuser/include/nvfuser/ir/builder.h

@@ -0,0 +1,215 @@
+// 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_guard.h>
+#include <ir/builder_passkey.h>
+#include <ir/container.h>
+#include <utils.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+namespace kir {
+class Kernel;
+}
+
+class ArrayConstruct;
+class IrCloner;
+class NamedScalar;
+class StructConstruct;
+class TensorView;
+class Val;
+
+//! IR builder interface
+class IrBuilder {
+ public:
+  //! Allocate a new IR node, forwarding the arguments to the appropriate
+  //! constructor and registering with the container
+  template <class T, class... Args>
+  static T* create(Args&&... args) {
+    Fusion* fusion = FusionGuard::getCurFusion();
+    return createInContainer<T>(fusion, std::forward<Args>(args)...);
+  }
+
+  //! Allocate a new IR node, forwarding the arguments to the appropriate
+  //! constructor and registering with the container
+  template <class T, class... Args>
+  static T* createInContainer(IrContainer* container, Args&&... args) {
+    NVF_ERROR(container != nullptr, "Need an active container to build IR.");
+    T* node = new T(IrBuilderPasskey(container), std::forward<Args>(args)...);
+
+    container->registerStmt(IrBuilderPasskey(container), node);
+
+    return node;
+  }
+
+  //! Clone an IR node, forwarding the arguments to the IrCloner constructor.
+  //! Register clones with IrCloner's target container.
+  template <class T>
+  static T* clone(const T* src, IrCloner* ir_cloner);
+
+  // Unary operations
+  NVF_API static Val* derefExpr(Val* val);
+  NVF_API static Val* negExpr(Val* val);
+  NVF_API static Val* logicalNotExpr(Val* val);
+  static Val* bitwiseNotExpr(Val* val);
+  NVF_API static Val* absExpr(Val* val);
+  static Val* setExpr(Val* val);
+  static Val* maybeCastExpr(DataType dtype, Val* val);
+  static Val* maybeRefCastExpr(DataType dtype, Val* val);
+  static Val* addressExpr(Val* val);
+  static NamedScalar* setExprNamedScalar(const std::string& name, Val* val);
+  static NamedScalar* addressExprNamedScalar(const std::string& name, Val* val);
+
+  // Binary operations
+  NVF_API static Val* logicalAndExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* logicalOrExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* bitwiseAndExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* bitwiseOrExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* bitwiseXorExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* lShiftExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* rShiftExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* eqExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* neExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* gtExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* ltExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* leExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* geExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* addExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* subExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* mulExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* divExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* ceilDivExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* modExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* maxExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* minExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* gcdExpr(Val* lhs, Val* rhs);
+  NVF_API static Val* isDivisibleExpr(Val* dividend, Val* divisor);
+
+  // Ternary operations
+  NVF_API static Val* whereExpr(Val* pred, Val* lhs, Val* rhs);
+
+  // Array and struct access
+  NVF_API static Val* getItemExpr(Val* array, Val* index);
+  NVF_API static Val* getItemExpr(Val* array, PolymorphicValue index);
+  NVF_API static Val* getAttrExpr(Val* struct_, std::string attr);
+  NVF_API static Val* reverseArrayExpr(Val* array);
+
+  // Get tensor metadata
+  NVF_API static Val* metadataExpr(TensorView* tv);
+
+  // Get tensor base address, for gmem tensor, it is something like
+  // `T1.data`. For smem tensor, it is something like `toSmem(T1)`.
+  static Val* baseAddressExpr(TensorView* tv);
+
+  // Construct an array of values, or nested arrays of values.
+  template <typename T>
+  static Val* arrayExpr(std::vector<T> members) {
+    if constexpr (std::is_same_v<T, Val*>) {
+      NVF_ERROR(!members.empty(), "Cannot create an array with no members.");
+      auto in_dtype = members.at(0)->dtype();
+      auto out_dtype =
+          ArrayType{std::make_shared<DataType>(in_dtype), members.size()};
+      auto out = create<Val>(out_dtype);
+      create<ArrayConstruct>(out, members);
+      return out;
+    } else {
+      static_assert(
+          is_std_vector_v<T>,
+          "Argument for function array must be vector of value or nested vector");
+      std::vector<Val*> array_members;
+      std::transform(
+          members.begin(),
+          members.end(),
+          std::back_inserter(array_members),
+          [](const T& member) { return arrayExpr(member); });
+      return arrayExpr(array_members);
+    }
+  }
+
+  template <typename T = NotImplementedStruct>
+  static Val* structExpr(
+      const std::vector<std::pair<std::string, Val*>>& fields,
+      std::string name = "") {
+    std::vector<StructType::FieldInfo> field_infos;
+    field_infos.reserve(fields.size());
+    for (auto& field : fields) {
+      field_infos.emplace_back(StructType::FieldInfo{
+          field.first,
+          std::make_shared<DataType>(field.second->dtype()),
+          true});
+    }
+    DataType dtype =
+        StructType::make<T>(std::move(field_infos), std::move(name));
+    auto out = create<Val>(dtype);
+    create<StructConstruct>(out, fields);
+    return out;
+  }
+
+ private:
+  static Val* newArithmeticExpr(BinaryOpType op_type, Val* lhs, Val* rhs);
+  static Val* newLogicExpr(BinaryOpType op_type, Val* lhs, Val* rhs);
+};
+
+//! A wrapper builder with static expression simplification
+//!
+//! Example:
+//! - addExpr(new Scalar(DataType::Int, 1), new Scalar(DataType::Int, 2)) ->
+//! Scalar(DataType::Int, 3)
+//! - addExpr(new Scalar(DataType::Int, 0), new NamedScalar("foo")) ->
+//! NamedScalar("foo")
+//!
+//! Designed to be used to simplify predicate and index expressions in
+//! generated code. Also, the shift validation may fail without
+//! this simplification.
+class SimplifyingIrBuilder : public IrBuilder {
+ public:
+  static Val* negExpr(Val* val);
+  static Val* logicalNotExpr(Val* val);
+  static Val* bitwiseNotExpr(Val* val);
+  static Val* maybeCastExpr(DataType dtype, Val* val);
+
+  static Val* addExpr(
+      Val* lhs,
+      PolymorphicValue rhs,
+      DataType rhs_dtype = DataType::Null);
+  static Val* addExpr(Val* lhs, Val* rhs);
+
+  static Val* subExpr(Val* lhs, Val* rhs);
+
+  NVF_API static Val* mulExpr(
+      Val* lhs,
+      PolymorphicValue rhs,
+      DataType rhs_dtype = DataType::Null);
+  static Val* mulExpr(Val* lhs, Val* rhs);
+  static Val* divExpr(Val* lhs, Val* rhs);
+
+  static Val* ceilDivExpr(Val* lhs, Val* rhs);
+
+  static Val* modExpr(Val* lhs, Val* rhs);
+  static Val* logicalAndExpr(Val* lhs, Val* rhs);
+  static Val* logicalOrExpr(Val* lhs, Val* rhs);
+  static Val* bitwiseAndExpr(Val* lhs, Val* rhs);
+  static Val* bitwiseOrExpr(Val* lhs, Val* rhs);
+  static Val* maxExpr(Val* lhs, Val* rhs);
+  static Val* minExpr(Val* lhs, Val* rhs);
+  static Val* gcdExpr(Val* lhs, Val* rhs);
+
+  static Val* ltExpr(Val* lhs, Val* rhs);
+  static Val* leExpr(Val* lhs, Val* rhs);
+  static Val* eqExpr(Val* lhs, Val* rhs);
+  static Val* neExpr(Val* lhs, Val* rhs);
+  static Val* geExpr(Val* lhs, Val* rhs);
+  static Val* gtExpr(Val* lhs, Val* rhs);
+
+  static Val* whereExpr(Val* pred, Val* lhs, Val* rhs);
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ir/builder_passkey.h

@@ -0,0 +1,29 @@
+// 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 IrContainer;
+
+// Passkey for builder to register properties with statements, and to call
+// functions in IrContainer
+class IrBuilderPasskey {
+  friend class IrBuilder;
+
+ public:
+  // TODO: Collapse ir_container and Kernel once Kernel inherits from
+  // IrContainer
+  IrContainer* const ir_container_ = nullptr;
+
+ private:
+  explicit IrBuilderPasskey(IrContainer* ir_container)
+      : ir_container_(ir_container) {}
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ir/cloner.h

@@ -0,0 +1,185 @@
+// 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 <dispatch.h>
+#include <exceptions.h>
+#include <ir/builder.h>
+#include <visibility.h>
+
+#include <tuple>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace nvfuser {
+
+class IrContainer;
+
+//! Clones nodes from an exiting Fusion
+//!
+//! \warning IrCloner machinery is a specialized helper for implementing
+//!   Fusion copy operations and the and limited scope of RecomputeTv below.
+//!   It is not intended for any other uses.
+//!
+class IrCloner {
+  friend class Statement;
+  friend class IrBuilder;
+
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  explicit IrCloner(IrContainer* container);
+  virtual ~IrCloner() = default;
+
+  NVF_API Statement* clone(const Statement* statement);
+
+  int64_t clone(int64_t x) {
+    return x;
+  }
+
+  template <class T>
+  T* clone(const T* node) {
+    return node ? clone(node->template as<Statement>())->template as<T>()
+                : nullptr;
+  }
+
+  template <class T>
+  std::vector<T> clone(const std::vector<T>& container) {
+    // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
+    std::vector<T> copy;
+    copy.reserve(container.size());
+    for (auto p : container) {
+      copy.push_back(clone(p));
+    }
+    return copy;
+  }
+
+  template <class T>
+  std::unordered_set<T> clone(const std::unordered_set<T>& container) {
+    std::unordered_set<T> copy;
+    copy.reserve(container.size());
+    for (auto p : container) {
+      copy.insert(clone(p));
+    }
+    return copy;
+  }
+
+  template <typename... Ts>
+  std::tuple<Ts...> clone(const std::tuple<Ts...>& tup) {
+    return std::apply(
+        [this](auto&... x) {
+          return std::make_tuple<Ts...>(this->clone(x)...);
+        },
+        tup);
+  }
+
+  template <typename T, typename U>
+  std::pair<T, U> clone(const std::pair<T, U>& p) {
+    return std::make_pair<T, U>(this->clone(p.first), this->clone(p.second));
+  }
+
+  template <class T, class U>
+  std::unordered_map<T, U> clone(const std::unordered_map<T, U>& container) {
+    std::unordered_map<T, U> copy;
+    copy.reserve(container.size());
+    for (const auto& [k, v] : container) {
+      copy.emplace(clone(k), clone(v));
+    }
+    return copy;
+  }
+
+  template <typename T, typename Hash = std::hash<T>>
+  DisjointSets<T, Hash> clone(const DisjointSets<T, Hash>& disjoint_sets) {
+    DisjointSets<T, Hash> cloned_disjoint_sets;
+    for (const auto& original_set : disjoint_sets.disjointSets()) {
+      NVF_ERROR(!original_set->empty());
+      bool first = true;
+      typename DisjointSets<T, Hash>::DisjointSet new_set;
+      for (const auto& val : *original_set) {
+        auto clone_of_val = clone(val);
+        if (first) {
+          auto it = cloned_disjoint_sets.initializeSet(clone_of_val).first;
+          new_set = it->second;
+          first = false;
+        } else {
+          cloned_disjoint_sets.appendToSet(clone_of_val, new_set);
+        }
+      }
+    }
+
+    return cloned_disjoint_sets;
+  }
+
+  IrContainer* container() const {
+    return ir_container_;
+  }
+
+ protected:
+  NVF_API void registerClone(const Statement* src, Statement* clone);
+  virtual Statement* handle(const Statement* s);
+
+ protected:
+  // We keep track of the original -> clone map so we don't
+  // duplicate clones of the same object if referenced multiple times
+  std::unordered_map<const Statement*, Statement*> clones_map_;
+
+ private:
+  // The destination Fusion container
+  IrContainer* ir_container_ = nullptr;
+
+  // Builder to make all the new nodes
+  IrBuilder builder_;
+};
+
+// Replicates all expressions used to generate the provided TensorView. Does not
+// replicate inputs. Does not replicate scalar values. In other words the value
+// provided will be recomputed from the inputs of the fusion.
+class RecomputeTv : private IrCloner {
+ public:
+  // Replicates expressions and values in provided expressions.
+  NVF_API static TensorView* recompute(
+      TensorView* tv,
+      const std::vector<Val*>& from = {});
+
+ private:
+  RecomputeTv(Fusion* fusion);
+  Statement* handle(const Statement* s) override;
+  Statement* handle(const TensorDomain*);
+};
+
+//! Clone an IR node, forwarding the arguments to the IrCloner constructor.
+template <class T>
+T* IrBuilder::clone(const T* src, IrCloner* ir_cloner) {
+  NVF_ERROR(
+      ir_cloner != nullptr,
+      "Cannot use create when a cloner object is set. Use clone.");
+
+  NVF_ERROR(
+      ir_cloner->container() != nullptr,
+      "Cloner doesn't have a valid container to store cloned object.");
+
+  T* dest = new T(src, ir_cloner);
+  const Statement* src_stmt = dynamic_cast<const Statement*>(src);
+  Statement* dest_stmt = dynamic_cast<Statement*>(dest);
+
+  auto dest_container = ir_cloner->container();
+  auto src_container = src_stmt->container();
+
+  dest_container->registerStmt(IrBuilderPasskey(dest_container), dest_stmt);
+
+  if (src_container != dest_container) {
+    dest_stmt->setName(IrBuilderPasskey(dest_container), src_stmt->name());
+  }
+
+  ir_cloner->registerClone(src_stmt, dest_stmt);
+
+  return dest;
+}
+
+} // namespace nvfuser

nvfuser/include/nvfuser/ir/container.h

@@ -0,0 +1,226 @@
+// 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 <ir/base_nodes.h>
+#include <utils.h>
+
+#include <deque>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace nvfuser {
+
+class IrBuilderPasskey;
+class ExprPasskey;
+class OptOutMutator;
+
+class NamedScalar;
+
+// Passkey for container to register names with statements
+class IrContainerPasskey {
+  friend class IrContainer;
+
+ private:
+  explicit IrContainerPasskey() = default;
+};
+
+class IrContainer : public PolymorphicBase {
+ public:
+  NVF_API IrContainer();
+
+  IrContainer(const IrContainer& other);
+  IrContainer(IrContainer&& other) noexcept;
+
+  IrContainer& operator=(const IrContainer& other);
+  IrContainer& operator=(IrContainer&& other) noexcept;
+
+  ~IrContainer() override;
+
+  bool inContainer(const Statement* stmt) const;
+
+  void assertInContainer(const Statement* stmt, const std::string& msg) const {
+    NVF_CHECK(
+        inContainer(stmt), msg, " it was not found in the active container.");
+  }
+
+  //! Return values in insertion order
+  const std::deque<Val*> deterministic_vals() const noexcept {
+    std::deque<Val*> vals_deque;
+    std::transform(
+        vals_up_.begin(),
+        vals_up_.end(),
+        std::back_inserter(vals_deque),
+        [](const std::unique_ptr<Val>& val_up) { return val_up.get(); });
+    return vals_deque;
+  }
+
+  //! Return expression in insertion order
+  const std::deque<Expr*> deterministic_exprs() const noexcept {
+    std::deque<Expr*> exprs_deque;
+    std::transform(
+        exprs_up_.begin(),
+        exprs_up_.end(),
+        std::back_inserter(exprs_deque),
+        [](const std::unique_ptr<Expr>& expr_up) { return expr_up.get(); });
+    return exprs_deque;
+  }
+
+  //! Return mapping from value to integer id
+  const std::unordered_map<Val*, int64_t> deterministic_vals_map()
+      const noexcept {
+    std::unordered_map<Val*, int64_t> vals_map;
+    int64_t count = 0;
+    std::transform(
+        vals_up_.begin(),
+        vals_up_.end(),
+        std::inserter(vals_map, vals_map.end()),
+        [&count](const std::unique_ptr<Val>& val_up) {
+          return std::make_pair(val_up.get(), count++);
+        });
+    return vals_map;
+  }
+
+  //! Return mapping from expression to integer id
+  const std::unordered_map<Expr*, int64_t> deterministic_exprs_map()
+      const noexcept {
+    std::unordered_map<Expr*, int64_t> exprs_map;
+    int64_t count = 0;
+    std::transform(
+        exprs_up_.begin(),
+        exprs_up_.end(),
+        std::inserter(exprs_map, exprs_map.end()),
+        [&count](const std::unique_ptr<Expr>& expr_up) {
+          return std::make_pair(expr_up.get(), count++);
+        });
+    return exprs_map;
+  }
+
+  //! Register the Statement with this container
+  NVF_API virtual void registerStmt(IrBuilderPasskey, Statement* stmt);
+
+  //! Register the Val with this container
+  NVF_API virtual void registerVal(IrBuilderPasskey, Val* val);
+
+  //! Register expr with this container.
+  NVF_API virtual void registerExpr(IrBuilderPasskey, Expr* expr);
+
+  //! Return the set of Exprs registered with this fusion. Warning: This will
+  //! return exprs outside inputs/outputs, so can be unsafe for use with
+  //! segmented fusions.
+  const std::unordered_set<Expr*>& unordered_exprs() const noexcept {
+    return exprs_;
+  }
+
+  //! Return the set of Vals registered with this fusion
+  const std::unordered_set<Val*>& vals() const noexcept {
+    return vals_;
+  }
+
+  // Shortcuts for frequently used vals
+  NVF_API Val* zeroVal();
+  NVF_API Val* oneVal();
+  Val* falseVal();
+  Val* trueVal();
+  NamedScalar* magicZeroVal();
+  NVF_API Val* zeroVal(DataType dtype);
+  NVF_API Val* oneVal(DataType dtype);
+  Val* metadataOf(Val*);
+
+  // Axioms about CUDA programming, for example: threadIdx.x < blockDim.x
+  const std::vector<Val*>& axioms() {
+    lazyInitAxioms();
+    return *axioms_;
+  }
+
+  void assumePositive(Val* val);
+  void assumeNonNegative(Val* val);
+
+ protected:
+  static IrCloner copy(const IrContainer* from, IrContainer* to);
+
+  friend void swap(IrContainer& a, IrContainer& b) noexcept;
+
+  // Let mutator remove Exprs.
+  friend OptOutMutator;
+
+  virtual void removeExpr(Expr* expr);
+
+  //! Completely remove val from the fusion, break all dependencies associated
+  //! with it
+  virtual void removeVal(Val* val);
+
+  //! Register the Val with this container
+  virtual void registerVal(Val* val);
+
+  //! Register expr with this container.
+  virtual void registerExpr(Expr* expr);
+
+  StmtNameType getValName(ValType vtype) {
+    if (val_type_name_map_.find(vtype) == val_type_name_map_.end()) {
+      val_type_name_map_[vtype] = 0;
+    }
+    return val_type_name_map_[vtype]++;
+  }
+
+  StmtNameType getExprName() {
+    return expr_name_counter_++;
+  }
+
+  void clear() noexcept;
+
+  void lazyInitAxioms();
+
+  // Deque of unique pointer is the memory owning data structure
+  std::deque<std::unique_ptr<Val>> vals_up_;
+
+  // A convenient set to return when we just need an unordered set to do
+  // something like check if a Val is in this container
+  std::unordered_set<Val*> vals_;
+
+  // Deque of unique pointer is the memory owning data structure
+  std::deque<std::unique_ptr<Expr>> exprs_up_;
+
+  // A convenient set to return when we just need an unordered set to do
+  // something like check if an Expr is in this container
+  std::unordered_set<Expr*> exprs_;
+
+  // Used to implement a generic "inContainer" that can be passed an invalid
+  // pointer. Specifically a pointer to a Statement owned by another container
+  // that has been freed. We can't check normally with the unordered_sets we
+  // already have because it would require a const_cast from a constant
+  // expr/val, or a dynamic cast from a Statement.
+  std::unordered_set<void*> raw_ptrs_;
+
+  // Values names counters
+  std::unordered_map<ValType, StmtNameType> val_type_name_map_;
+
+  // Expression names counter
+  StmtNameType expr_name_counter_ = 0;
+
+  // Manually store some persistent, frequently used nodes. It's very
+  // challenging to do this anything but manually as detecting when a container
+  // may or may not have one of these vals is tricky. Specifically because if
+  // the container doesn't own it, it's hard to understand from the outside if
+  // the node may have been removed then re-registered. It could also be tricky
+  // to know when we're using a different container as in FusionCopy_test
+  // demonstrates deleting then creating containers can result in the same
+  // pointer for the container.
+  std::unique_ptr<Val> true_val_;
+  std::unique_ptr<Val> false_val_;
+  std::unique_ptr<Val> one_val_;
+  std::unique_ptr<Val> zero_val_;
+  std::unique_ptr<NamedScalar> magic_zero_val_;
+  std::unique_ptr<std::vector<Val*>> axioms_;
+  std::unordered_map<Val*, std::pair<Val*, Expr*>> metadata_;
+};
+
+} // namespace nvfuser