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

nvfuser/include/nvfuser/polymorphic_value.h

@@ -0,0 +1,432 @@
+// 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 <any>
+#include <complex>
+#include <cstddef>
+#include <functional>
+#include <numeric>
+#include <ostream>
+#include <unordered_map>
+
+#include <ATen/ATen.h>
+
+#ifndef DYNAMIC_TYPE_CHECK
+#define DYNAMIC_TYPE_CHECK NVF_ERROR
+#endif
+
+#include <dynamic_type/dynamic_type.h>
+#include <macros.h>
+#include <opaque_type.h>
+
+namespace nvfuser {
+
+struct DataType;
+
+// Use a single pointer type to represent all pointers, otherwise we would need
+// exponential compilation time for all pointer types in PolymorphicValue.
+class Pointer {
+  std::byte* ptr_;
+  int64_t size_;
+
+ public:
+  template <typename T>
+  Pointer(T* ptr) : ptr_(reinterpret_cast<std::byte*>(ptr)), size_(sizeof(T)) {}
+
+  inline Pointer(void* ptr, DataType dtype);
+
+  Pointer() : ptr_(nullptr), size_(-1) {}
+
+  int64_t size() const {
+    return size_;
+  }
+
+  template <typename T>
+  explicit operator T*() const {
+    return reinterpret_cast<T*>(ptr_);
+  }
+
+  Pointer& operator+=(int64_t offset) {
+    ptr_ += offset * size_;
+    return *this;
+  }
+
+  Pointer& operator-=(int64_t offset) {
+    ptr_ -= offset * size_;
+    return *this;
+  }
+
+  Pointer& operator++() {
+    ptr_ += size_;
+    return *this;
+  }
+
+  Pointer& operator--() {
+    ptr_ -= size_;
+    return *this;
+  }
+
+  Pointer operator++(int) {
+    Pointer tmp = *this;
+    ++*this;
+    return tmp;
+  }
+
+  Pointer operator--(int) {
+    Pointer tmp = *this;
+    --*this;
+    return tmp;
+  }
+
+  Pointer operator+(int64_t offset) const {
+    Pointer tmp = *this;
+    tmp += offset;
+    return tmp;
+  }
+
+  Pointer operator-(int64_t offset) const {
+    Pointer tmp = *this;
+    tmp -= offset;
+    return tmp;
+  }
+
+  int64_t operator-(const Pointer& other) const {
+    NVF_ERROR(size_ == other.size_);
+    return (ptr_ - other.ptr_) / (int64_t)size_;
+  }
+
+  bool operator==(const Pointer& other) const {
+    return ptr_ == other.ptr_;
+  }
+
+  bool operator==(std::nullptr_t) const {
+    return ptr_ == nullptr;
+  }
+
+  bool operator!=(const Pointer& other) const {
+    return ptr_ != other.ptr_;
+  }
+
+  bool operator!=(std::nullptr_t) const {
+    return ptr_ != nullptr;
+  }
+
+  bool operator<(const Pointer& other) const {
+    return ptr_ < other.ptr_;
+  }
+
+  bool operator>(const Pointer& other) const {
+    return ptr_ > other.ptr_;
+  }
+
+  bool operator<=(const Pointer& other) const {
+    return ptr_ <= other.ptr_;
+  }
+
+  bool operator>=(const Pointer& other) const {
+    return ptr_ >= other.ptr_;
+  }
+
+  bool operator!() const {
+    return !ptr_;
+  }
+
+  explicit operator bool() const {
+    return ptr_;
+  }
+
+  explicit operator int64_t() const {
+    return reinterpret_cast<int64_t>(ptr_);
+  }
+
+  explicit operator unsigned() const {
+    return (unsigned)(int64_t)(*this);
+  }
+
+  explicit operator size_t() const {
+    return reinterpret_cast<size_t>(ptr_);
+  }
+};
+
+inline Pointer operator+(int64_t offset, const Pointer& ptr) {
+  return ptr + offset;
+}
+
+inline std::ostream& operator<<(std::ostream& os, const Pointer& ptr) {
+  os << (void*)ptr;
+  return os;
+}
+
+struct Struct;
+class Accessor;
+struct StructType;
+
+// See Note [Struct Support in PolymorphicValue] for documentation.
+class StructHandle {
+  std::shared_ptr<Struct> struct_ptr_;
+
+ public:
+  StructHandle(std::shared_ptr<Struct> struct_ptr)
+      : struct_ptr_(std::move(struct_ptr)) {}
+  StructHandle& operator=(std::shared_ptr<Struct> struct_ptr) {
+    struct_ptr_ = std::move(struct_ptr);
+    return *this;
+  }
+
+  StructHandle(const StructHandle& other) = default;
+  StructHandle(StructHandle&& other) = default;
+  StructHandle& operator=(const StructHandle& other) = default;
+  StructHandle& operator=(StructHandle&& other) = default;
+
+  bool operator==(const StructHandle& other) const;
+
+  template <typename T>
+  bool is() const {
+    return std::dynamic_pointer_cast<T>(struct_ptr_) != nullptr;
+  }
+
+  template <typename T>
+  inline T& as() const {
+    return *std::dynamic_pointer_cast<T>(struct_ptr_);
+  }
+
+  inline StructType type() const;
+
+  template <typename Ret, typename Class>
+  inline std::enable_if_t<std::is_base_of_v<Struct, Class>, Ret&> operator->*(
+      Ret Class::* member) const {
+    return as<Class>().*member;
+  }
+
+  inline Accessor operator->*(const std::string& key) const;
+};
+
+using PolymorphicValue = dynamic_type::DynamicType<
+    dynamic_type::Containers<std::vector>,
+    StructHandle,
+    Pointer,
+    Opaque,
+    at::Tensor,
+    std::complex<double>,
+    double,
+    int64_t,
+    bool>;
+
+namespace PolymorphicValue_functions {
+
+NVF_API std::string toString(const PolymorphicValue& v);
+
+template <typename T>
+inline bool isNan(const T& a) {
+  return std::isnan(a);
+}
+
+// For example, `nan+i` and `nan-i` are treated equal because both are NaNs.
+// This is consistent with pytorch's implementation:
+// https://github.com/pytorch/pytorch/blob/6d8e0c4b5a3be8201cab731dfd1e6513162cf25c/c10/util/complex_utils.h#L43.
+template <typename T>
+inline bool isNan(const std::complex<T>& a) {
+  return std::isnan(a.real()) || std::isnan(a.imag());
+}
+
+// NaNs are treated equal.
+template <typename T>
+inline bool isSameNanSensitive(const T& a, const T& b) {
+  if (isNan(a) && isNan(b)) {
+    return true;
+  }
+  return a == b;
+}
+
+inline bool isSame(const PolymorphicValue& a, const PolymorphicValue& b) {
+  if (a.type() != b.type()) {
+    return false;
+  }
+  if (a.is<at::Tensor>()) {
+    return (a.as<at::Tensor>().is_same(b.as<at::Tensor>()));
+  }
+  if (a.is<double>()) {
+    return isSameNanSensitive(a.as<double>(), b.as<double>());
+  }
+  if (a.is<std::complex<double>>()) {
+    return isSameNanSensitive(
+        a.as<std::complex<double>>(), b.as<std::complex<double>>());
+  }
+  return a == b;
+}
+
+inline PolymorphicValue signbit(const PolymorphicValue& a) {
+  if (a.is<int64_t>()) {
+    return PolymorphicValue(std::signbit(a.as<int64_t>()));
+  }
+  if (a.is<double>()) {
+    return PolymorphicValue(std::signbit(a.as<double>()));
+  }
+  if (a.is<at::Tensor>()) {
+    return PolymorphicValue(a.as<at::Tensor>().signbit());
+  }
+  NVF_THROW("PolymorphicValue signbit not implemented for ", a.type().name());
+}
+
+inline PolymorphicValue fmod(
+    const PolymorphicValue& a,
+    const PolymorphicValue& b) {
+  // TODO: relax the type check
+  NVF_ERROR(
+      a.is<at::Tensor>() || a.type() == b.type(),
+      "fmod is not implemented for mismatch dtypes");
+  if (a.is<int64_t>()) {
+    if (b.is<int64_t>()) {
+      return PolymorphicValue(std::fmod(a.as<int64_t>(), b.as<int64_t>()));
+    }
+    if (b.is<double>()) {
+      return PolymorphicValue(std::fmod(a.as<int64_t>(), b.as<double>()));
+    }
+  }
+  if (a.is<double>()) {
+    if (b.is<int64_t>()) {
+      return PolymorphicValue(std::fmod(a.as<double>(), b.as<int64_t>()));
+    }
+    if (b.is<double>()) {
+      return PolymorphicValue(std::fmod(a.as<double>(), b.as<double>()));
+    }
+  }
+  if (a.is<at::Tensor>()) {
+    if (b.is<int64_t>()) {
+      return PolymorphicValue(a.as<at::Tensor>().fmod(b.as<int64_t>()));
+    }
+    if (b.is<double>()) {
+      return PolymorphicValue(a.as<at::Tensor>().fmod(b.as<double>()));
+    }
+    if (b.is<at::Tensor>()) {
+      return PolymorphicValue(a.as<at::Tensor>().fmod(b.as<at::Tensor>()));
+    }
+  }
+  NVF_THROW(
+      "PolymorphicValue fmod not implemented for ",
+      a.type().name(),
+      " , ",
+      b.type().name());
+}
+
+inline PolymorphicValue ceildiv(
+    const PolymorphicValue& a,
+    const PolymorphicValue& b) {
+  if (a.is<int64_t>() && b.is<int64_t>()) {
+    auto aa = a.as<int64_t>();
+    auto bb = b.as<int64_t>();
+    if (bb > 0) {
+      return PolymorphicValue((aa + bb - 1) / bb);
+    } else {
+      return PolymorphicValue((aa + bb + 1) / bb);
+    }
+  }
+  return PolymorphicValue(std::ceil((a / b).as<double>()));
+}
+
+inline PolymorphicValue max(
+    const PolymorphicValue& a,
+    const PolymorphicValue& b) {
+  return PolymorphicValue(a > b ? a : b);
+}
+
+inline PolymorphicValue min(
+    const PolymorphicValue& a,
+    const PolymorphicValue& b) {
+  return PolymorphicValue(a < b ? a : b);
+}
+
+inline PolymorphicValue gcd(
+    const PolymorphicValue& a,
+    const PolymorphicValue& b) {
+  return PolymorphicValue(std::gcd(a.as<int64_t>(), b.as<int64_t>()));
+}
+
+inline PolymorphicValue abs(const PolymorphicValue& a) {
+  if (a.is<int64_t>()) {
+    return PolymorphicValue(std::abs(a.as<int64_t>()));
+  }
+  if (a.is<double>()) {
+    return PolymorphicValue(std::abs(a.as<double>()));
+  }
+  if (a.is<bool>()) {
+    return a;
+  }
+  if (a.is<std::complex<double>>()) {
+    return std::abs(a.as<std::complex<double>>());
+  }
+  if (a.is<at::Tensor>()) {
+    return a.as<at::Tensor>().abs();
+  }
+  NVF_THROW("PolymorphicValue abs not implemented for ", a.type().name());
+}
+
+inline PolymorphicValue erf(const PolymorphicValue& a) {
+  if (a.is<at::Tensor>()) {
+    return PolymorphicValue(a.as<at::Tensor>().erf());
+  }
+  NVF_THROW("PolymorphicValue erf not implemented for ", a.type().name());
+}
+
+// Convert scalars, vector of scalars, vector of vector of scalars, etc., into
+// an at::Tensor. device argument allows for the creation of CPU Scalars.
+inline PolymorphicValue toTensor(
+    const PolymorphicValue& x,
+    at::DeviceType device_type = at::kCUDA,
+    int8_t device_index = 0) {
+  if (x.is<at::Tensor>()) {
+    return x;
+  }
+  auto options = at::TensorOptions().device(device_type, device_index);
+  if (x.is<int64_t>()) {
+    return PolymorphicValue(
+        at::tensor(x.as<int64_t>(), options.dtype(at::kLong)).squeeze());
+  }
+  if (x.is<double>()) {
+    return PolymorphicValue(
+        at::tensor(x.as<double>(), options.dtype(at::kDouble)).squeeze());
+  }
+  if (x.is<bool>()) {
+    return PolymorphicValue(
+        at::tensor(x.as<bool>(), options.dtype(at::kBool)).squeeze());
+  }
+  if (x.is<std::complex<double>>()) {
+    return PolymorphicValue(
+        at::tensor(
+            (c10::complex<double>)x.as<std::complex<double>>(),
+            options.dtype(at::kComplexDouble))
+            .squeeze());
+  }
+  if (x.is<std::vector>()) {
+    auto vec = x.as<std::vector>();
+    std::vector<at::Tensor> tensors;
+    tensors.reserve(vec.size());
+    for (const auto& elem : vec) {
+      tensors.push_back(toTensor(elem).as<at::Tensor>());
+    }
+    return PolymorphicValue(at::stack(tensors));
+  }
+  NVF_THROW("PolymorphicValue toTensor not implemented for ", x.type().name());
+}
+
+// Convert PolymorphicValue to c10::Scalar.
+inline c10::Scalar toScalar(const PolymorphicValue& x) {
+  if (x.is<std::complex<double>>()) {
+    return (c10::complex<double>)x.as<std::complex<double>>();
+  } else {
+    return (c10::Scalar)x;
+  }
+}
+
+} // namespace PolymorphicValue_functions
+
+} // namespace nvfuser
+
+#include <struct.inl>

nvfuser/include/nvfuser/predicate_compute.h

@@ -0,0 +1,213 @@
+// 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/thread_predicate.h>
+#include <device_lower/utils.h>
+#include <exceptions.h>
+#include <index_compute.h>
+#include <kernel_ir.h>
+#include <logical_domain_map.h>
+
+namespace nvfuser {
+
+class PredicateCompute {
+ public:
+  // ignore_internal_syncthread_ops will prevent creation of predicates on
+  // block/grid broadcast/reduce as these have syncthread calls within them
+  // so all threads need to execute the function.
+  static Val* getInlinePredicate(
+      const Expr* expr,
+      const std::vector<ForLoop*>& loops,
+      const std::unordered_set<ForLoop*>& rotated_loops,
+      Val* thread_pred,
+      PredicateType pred_type);
+};
+
+//! Parallelized domains may need to be predicated with threading
+//! indices and IterDomain extents. For example, if a domain is
+//! parallelized by TIDx, when TIDx is not exact, i.e., it can be
+//! larger than the extents of domains parallelized by TIDx,
+//! threadIdx.x may be larger than the IterDomain extent. This can be
+//! harmless for Local tensors, however, for it would
+//! result in out-of-bounds access for Shared tensors as they are
+//! allocated based on tensor shapes rather than threading
+//! dimensions.
+class ParallelizedDomainPredicate {
+ public:
+  //! Predicate information for parallelized domains
+  class PredicateInfo {
+   public:
+    explicit PredicateInfo(ParallelType pt) : pt_(pt) {}
+
+    //! Adds a domain that is parallized by the same parallel type
+    bool addDomain(IterDomain* id);
+
+    const std::vector<IterDomain*>& ids() const {
+      return ids_;
+    }
+
+    //! Generates a predicate Val from predicate information
+    Val* getPredicate() const;
+
+   private:
+    ParallelType pt_;
+    //! Domains parallelized by the same parallel type
+    std::vector<IterDomain*> ids_;
+  };
+
+  //! Returns a predicate Val for parallelied domains of an expression.
+  static Val* getPredicate(
+      const Expr* expr,
+      const std::vector<ForLoop*>& loops);
+
+  //! Returns predicate information for parallelied domains of an
+  //! expression.
+  static std::unordered_map<ParallelType, PredicateInfo> getPredicateMap(
+      const Expr* expr,
+      const std::vector<ForLoop*>& loops,
+      ForLoop* unswitched_loop = nullptr);
+};
+
+//! Keys to identify unique unswitch predicates. Just consists of a
+//! predicated concrete domain if not parallelized. If parallelized,
+//! pick one for each different parallelization. When the same
+//! parallel type is used for different concrete domains, they are
+//! considered different predicates and are included in the unswitch
+//! condition lists.
+class UnswitchPredicateKey {
+ public:
+  UnswitchPredicateKey();
+
+  // Parameter loop_ids represents the loop domains used for the
+  // predicated domain
+  UnswitchPredicateKey(
+      IterDomain* predicated_consumer_id,
+      TensorView* consumer_tv,
+      IterDomain* predicated_concrete_id,
+      std::unordered_set<IterDomain*> loop_ids);
+
+  bool operator==(const UnswitchPredicateKey& other) const {
+    return predicated_concrete_id_ == other.predicated_concrete_id_ &&
+        parallel_concrete_ids_ == other.parallel_concrete_ids_;
+  }
+
+  const auto& predicatedId() const {
+    return predicated_concrete_id_;
+  }
+
+  const auto& parallelConcreteIds() const {
+    return parallel_concrete_ids_;
+  }
+
+  IterDomain* parallelId(ParallelType pt) const {
+    auto it = parallelConcreteIds().find(pt);
+    if (it == parallelConcreteIds().end()) {
+      return nullptr;
+    } else {
+      return it->second;
+    }
+  }
+
+  std::string toString() const;
+
+ private:
+  //! Predicated concrete domain
+  IterDomain* predicated_concrete_id_ = nullptr;
+  //! Dependent loop domains
+  std::unordered_set<IterDomain*> loop_ids_;
+  //! Store parallelized concrete domains
+  std::unordered_map<ParallelType, IterDomain*> parallel_concrete_ids_;
+};
+
+struct UnswitchPredicateKeyHash {
+  std::size_t operator()(const UnswitchPredicateKey& key) const;
+};
+
+// Generate predicates for loops that are unswitched, unrolled or
+// vectorized loops
+class UnswitchPredicate {
+ public:
+  // Get a predicate for a loop that is unswitched, unrolled or
+  // vectorized. The outer_loops parameter represents the outer loops
+  // of the unswitched/unrolled/vectorized loop.
+  static Val* get(
+      const std::vector<ForLoop*>& outer_loops,
+      ForLoop* unrolled_loop);
+
+ private:
+  //! Predicate information for each UnswitchPredicateKey.
+  struct MergedPredicates {
+    //! Predicate information for the start and stop predicates.
+    struct Info {
+      //! Most restrictive static predicate. Nullptr if no static
+      //! predicate found.
+      Val* static_pred = nullptr;
+      //! The offset value of static_pred
+      PolymorphicValue static_offset = 0L;
+      //! List of dynamic predicates.
+      std::vector<Val*> dynamic_preds;
+      //! Circular buffer loop stage if applicable. The predicate
+      //! generated in the main loop where no epilogue is generated
+      //! needs to be used.
+      CircularBufferLoopStage loop_stage =
+          CircularBufferLoopStage::NotApplicable;
+    };
+    UnswitchPredicateKey predicate_key;
+    Info start;
+    Info stop;
+  };
+
+  UnswitchPredicate(std::vector<ForLoop*> outer_loops, ForLoop* unrolled_loop);
+
+  void predicateOn(Expr*);
+
+  void openLoop(ForLoop*);
+
+  void openIte(kir::IfThenElse*);
+
+  //! Generates the final predicates from the predicated_keys map
+  void finalize();
+
+  //! Merge predicates as much as possible. If a predicate offset is
+  //! static, only pick the most restrictive one, e.g., the one with the
+  //! minimum offset for the start predication.
+  void mergeUnswitchPredicates(
+      Val* predicate,
+      Val* offset,
+      CircularBufferLoopStage loop_stage,
+      MergedPredicates::Info& merged_predicate_info,
+      bool is_start);
+
+  //! Adds new predicates for parallelized domains
+  void addParallelizedDomainPredicates(Expr*);
+
+ private:
+  //! Track which iter domains have been predicated
+  std::unordered_set<UnswitchPredicateKey, UnswitchPredicateKeyHash>
+      predicated_keys_;
+
+  //! The predicates that have been recorded but not yet finalized
+  std::vector<MergedPredicates> pending_predicates_;
+
+  //! Track which parallelized domains have been predicated
+  std::unordered_map<ParallelType, ParallelizedDomainPredicate::PredicateInfo>
+      parallelized_dom_predicates_;
+
+  //! The predicates that have been generated.
+  std::vector<Val*> predicates_;
+
+  std::vector<ForLoop*> for_loops_;
+
+  // Keep track of the loop in which the currently visiting expr is a rotated.
+  std::unordered_set<ForLoop*> rotated_loop_;
+
+  ForLoop* unrolled_loop_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/python_frontend/distributed_tensor.h

@@ -0,0 +1,50 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+
+#pragma once
+
+#include <ATen/core/TensorBody.h>
+
+#include <multidevice/device_mesh.h>
+#include <type.h>
+
+namespace nvfuser::python_frontend {
+
+// A class that represents a distributed tensor. It wraps a local tensor, a
+// mesh, and a mapping from mesh axes to tensor axes. If the mesh is empty,
+// it degenerates into a non-distributed tensor.
+class DistributedTensor {
+ public:
+  explicit DistributedTensor(
+      at::Tensor local_tensor,
+      DeviceMesh mesh = DeviceMesh())
+      : local_(std::move(local_tensor)), mesh_(std::move(mesh)) {}
+  DistributedTensor(const DistributedTensor&) = delete;
+  DistributedTensor& operator=(const DistributedTensor&) = delete;
+  DistributedTensor(DistributedTensor&&) = default;
+  DistributedTensor& operator=(DistributedTensor&&) = default;
+
+  const DeviceMesh& mesh() const {
+    return mesh_;
+  }
+
+  at::Tensor local() const {
+    return local_;
+  }
+
+  void setAxisIsShardedOn(int64_t axis, ParallelType parallel_type);
+
+  int64_t axisShardedOn(ParallelType parallel_type) const;
+
+ private:
+  at::Tensor local_;
+  DeviceMesh mesh_;
+  std::unordered_map<ParallelType, int64_t> axis_sharded_on_;
+};
+
+} // namespace nvfuser::python_frontend