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

nvfuser/include/nvfuser/evaluator_common.h

@@ -0,0 +1,295 @@
+// 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/lower2device.h>
+#include <exceptions.h>
+#include <fusion.h>
+#include <ir/all_nodes.h>
+#include <polymorphic_value.h>
+#include <runtime/executor_params.h>
+#include <utils.h>
+#include <visibility.h>
+
+#include <c10/core/DeviceType.h>
+
+namespace nvfuser {
+
+class PrecomputedValues;
+class KernelArgumentHolder;
+struct TensorArgAbstract;
+
+//! NaiveValueMachine:
+//!  This is an un-optimized runtime for evaluating a
+//!   set of values in one run. The runtime contains
+//!   a vector of instructions inferred from IR at compile-time
+//!   and it currently must be associated with an instance of
+//!   PrecomputedValues that will provide the workspace
+//!   containing the concrete values for the values.
+class NaiveValueMachine {
+  //! The generic types of instructions supported for this machine.
+  enum class InstructionType { UNARY_OP, BINARY_OP, TERNARY_OP, SET_OP };
+
+ public:
+  //! Constructor lowers all the expr IR nodes stored in precomputed_values
+  //!  and stores them in the private state.
+  NaiveValueMachine(PrecomputedValues& precomputed_values);
+
+  //! Copy all values other than `precomputed_values_` from other
+  //! This would be better implemented as a copy constructor, except that would
+  //! also presumably bind precomputed_values_ which we could not then rebind,
+  //! as we need to during cloning.
+  void copyFrom(const NaiveValueMachine& other);
+
+  //! Runs all the instructions and write results to the associated
+  //!  precomputed_values.
+  void run();
+
+ private:
+  //! Convert an unary IR expr to an instruction
+  void makeUnaryOp(UnaryOp* uop);
+
+  //! Convert an binary IR expr to an instruction
+  void makeBinaryOp(BinaryOp* bop);
+
+  //! Convert an ternary IR expr to an instruction
+  void makeTernaryOp(TernaryOp* bop);
+
+  //! Create an empty instruction with all default values
+  //!  and place it at the end of the instruction buffer.
+  int makeInstructionEntry();
+
+  //! Run a single instruction at the given index of
+  //!  the instruction buffer. Decodes and dispatches
+  //!  to the corresponding instruction handle functions.
+  void runInstruction(int index);
+
+  //! Runs a unary operation at given index of instruction buffer
+  void runUnaryOp(int index);
+
+  //! Runs a binary operation at given index of instruction buffer
+  void runBinaryOp(int index);
+
+  //! Runs a ternary operation at given index of instruction buffer
+  void runTernaryOp(int index);
+
+ private:
+  friend PrecomputedValues;
+
+  //! Reference to the PrecomputedValues workspace associated with
+  //!   this runtime. All the instructions will read and write the
+  //!   values in this workspace.
+  PrecomputedValues& precomputed_values_;
+
+  //! Instruction buffer. All states are in separate vectors and
+  //!  the entry of each vector at the same index correspond to
+  //!  the same instruction.
+
+  //! Total number of instructions
+  int num_of_instructions_ = 0;
+
+  //! Machine instruction type for each instruction i.e.
+  //!  unary or binary
+  std::vector<InstructionType> inst_type_;
+
+  //! Unary operator type if applicable, contains a default
+  //!  value at each index corresponding to a binary op.
+  std::vector<UnaryOpType> uop_type_;
+
+  //! Data type for unary op of type UnaryOpType::Cast, contains a default
+  //!  value at each index corresponding other ops.
+  std::vector<DataType> data_type_;
+
+  //! Binary operator type if applicable, contains a default
+  //!  value at each index corresponding to a binary op.
+  std::vector<BinaryOpType> bop_type_;
+
+  //! Ternary operator type if applicable, contains a default
+  //!  value at each index corresponding to a ternary op.
+  std::vector<TernaryOpType> top_type_;
+
+  //! Indexes of operands and destination of each instruction.
+  //!  The indexes corresponds to positions in the workspace
+  //!  where concrete values are hosted.
+
+  //! Operand 0 of each instruction.
+  std::vector<int> src0_;
+
+  //! Operand 1 of each instruction, a default value at
+  //!  each index corresponding to a unary op.
+  std::vector<int> src1_;
+
+  //! Operand 2 of each instruction, a default value at
+  //!  each index corresponding to a unary or binary op.
+  std::vector<int> src2_;
+
+  //! Destination of each instruction.
+  std::vector<int> dest_;
+};
+
+//! PrecomputedValues:
+//!  A class to support optimized evaluation of values
+//!  at runtime.
+//!    At compile time all necessary values are collected
+//!  from given IR nodes and a runtime and a workspace containing
+//!  the concrete values is created and pre-allocated.
+//!    At runtime the value vm is used to evaluate all the
+//!  values and store them in the workspace ahead of time.
+class PrecomputedValues {
+ public:
+  PrecomputedValues() = delete;
+  NVF_API PrecomputedValues(PrecomputedValues&&) = default;
+
+  NVF_API explicit PrecomputedValues(Fusion* fusion);
+
+  NVF_API ~PrecomputedValues();
+
+  //! Bind a list of concrete values to the fusion's runtime inputs.
+  void bindInputs(const KernelArgumentHolder& args);
+
+  //! Bind a list of concrete values to a list of fusion values. The two lists
+  //! must have the same size.
+  void bindValues(
+      const std::vector<Val*>& values,
+      const KernelArgumentHolder& args);
+
+  using ParallelExtentMap =
+      std::unordered_map<ParallelType, std::vector<const Val*>>;
+
+  //! Bind concrete values from launch constraints
+  void bindParallelExtents(
+      const ParallelExtentMap& parallel_extents,
+      const LaunchParams& launch_constraint);
+
+  //! Bind the NamedScalars corresponding to the
+  //!  concrete parallel dimension sizes after the
+  //!  actual value has been resolved.
+  void bindConcreteParallelTypeValue(ParallelType pt, PolymorphicValue value);
+
+  //! Returns if the workspace contains evaluated results.
+  bool hasValidValues() {
+    return has_valid_values_;
+  }
+
+  //! Runs the internal value machine that will compute
+  //!  the values allocated in the workspace.
+  void evaluate();
+
+  //! Returns value for the given IR node if it's stored
+  //!  in the workspace and has been evaluated.
+  const PolymorphicValue& getMaybeValueFor(const Val* val) const;
+
+  //! Debugging helper, prints all the currently known values
+  void print() const;
+
+  PrecomputedValues clone(IrCloner& ir_cloner) const;
+
+ protected:
+  // Fusion IR associated with the precomputed values. Can be kir::Kernel or
+  // Fusion.
+  Fusion* fusion_ = nullptr;
+
+  //! Contains all the named scalars correspond
+  //!  to thread size of each parallel type.
+  std::unordered_map<ParallelType, std::unique_ptr<std::vector<int>>>
+      thread_dim_value_indices_;
+
+  //! Initialize the workspace before first use.
+  //!  Assume the given value list IR nodes have
+  //!  been topologically sorted.
+  void initializeValueList(const std::vector<Val*>& sorted_value_list);
+
+  //! Bind concrete value to the given index
+  //!  if the index is valid.
+  void bindValue_(int index, const PolymorphicValue& value) {
+    if (index < 0 || is_constant_[index]) {
+      return;
+    }
+    defined_[index] = true;
+    values_[index] = value;
+    binding_log_.emplace_back(index, value);
+  }
+  template <typename T>
+  void bindValue(int index, const T& value) {
+    bindValue_(index, PolymorphicValue(value));
+  }
+
+  //! Invalidate all computed values in the workspace.
+  void invalidate();
+
+  //! Interface for subclasses to access symbols_
+  void loadSymbols(std::vector<Val*> symbols) {
+    symbols_ = std::move(symbols);
+  }
+
+  //! Interface for subclasses to access symbols_
+  std::vector<Val*>& symbols() {
+    return symbols_;
+  }
+
+  //! Initialize the value runtime that will
+  //!  infer instructions from the workspace.
+  void initializeIntegerMachine() {
+    value_machine_ = std::make_unique<NaiveValueMachine>(*this);
+  }
+
+  //! Iterate through all the named scalars corresponding
+  //!  to thread sizes and pre-group them by their parallel
+  //!  types.
+  void initializeNamedScalars();
+
+ private:
+  //! Post evaluation check, throws if any computed value
+  //!  is inconsistent with its bound value
+  void validate();
+
+  //! Returns true if workspace has a computed or constant
+  //!  value for given index.
+  bool hasValue(int index) {
+    NVF_ERROR(index > 0);
+    return defined_[index] || is_constant_[index];
+  }
+
+  void bindTensorMetaData(TensorView* tv, const at::Tensor& tensor);
+
+ private:
+  friend NaiveValueMachine;
+
+  //! Marks if an evaluation has finished
+  bool has_valid_values_ = false;
+
+  //! The size of workspace
+  int num_of_values_ = -1;
+
+  //! Marks if a value has been bound or
+  //!  computed at each index.
+  std::vector<bool> defined_;
+
+  //! Marks if a value is compile-time constant
+  //!  at each index.
+  std::vector<bool> is_constant_;
+
+  //! Stores the concrete values at each index.
+  std::vector<PolymorphicValue> values_;
+
+  //! Use a single monostate to represent null, instead of creating a new
+  //! PolymorphicValue for each null.
+  PolymorphicValue null_ = std::monostate{};
+
+  //! Stores the IR nodes corresponding to each index.
+  std::vector<Val*> symbols_;
+
+  //! An internal log to keep track of all the bindings
+  //!  used in each evaluation cycle. To be used for
+  //!  consistency check.
+  std::vector<std::pair<int, PolymorphicValue>> binding_log_;
+
+  //! Integer runtime for realizing the values computations.
+  std::unique_ptr<NaiveValueMachine> value_machine_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/exceptions.h

@@ -0,0 +1,283 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+// This is a refactor of the NVF_ERROR and NVF_CHECK macros
+// from PyTorch for implementing NVFuser specific macros.
+
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <deque>
+#include <exception>
+#include <iosfwd>
+#include <optional>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include <visibility.h>
+
+namespace nvfuser {
+
+// This function will demangle the mangled function name into a more human
+// readable format, e.g. _Z1gv -> g().
+// More information:
+// https://github.com/gcc-mirror/gcc/blob/master/libstdc%2B%2B-v3/libsupc%2B%2B/cxxabi.h
+// NOTE: `__cxa_demangle` returns a malloc'd string that we have to free
+// ourselves.
+std::string demangle(const char* name);
+
+std::string _get_backtrace(
+    size_t frames_to_skip = 0,
+    size_t maximum_number_of_frames = 64,
+    bool skip_python_frames = true);
+
+struct SourceLocation {
+  const char* function;
+  const char* file;
+  uint32_t line;
+};
+
+std::ostream& operator<<(std::ostream& out, const SourceLocation& loc);
+
+struct CompileTimeEmptyString {
+  operator const std::string&() const {
+    static const std::string empty_string_literal;
+    return empty_string_literal;
+  }
+  operator const char*() const {
+    return "";
+  }
+};
+
+template <typename T>
+struct CanonicalizeStrTypes {
+  using type = const T&;
+};
+
+template <size_t N>
+struct CanonicalizeStrTypes<std::array<char, N>> {
+  using type = const char*;
+};
+
+inline std::ostream& _to_str(std::ostream& ss) {
+  return ss;
+}
+
+template <typename T>
+inline std::ostream& _to_str(std::ostream& ss, const T& t) {
+  // NOLINTNEXTLINE(clang-analyzer-core.CallAndMessage)
+  ss << t;
+  return ss;
+}
+
+template <>
+inline std::ostream& _to_str<CompileTimeEmptyString>(
+    std::ostream& ss,
+    const CompileTimeEmptyString&) {
+  return ss;
+}
+
+template <typename T, typename... Args>
+inline std::ostream& _to_str(
+    std::ostream& ss,
+    const T& t,
+    const Args&... args) {
+  return _to_str(_to_str(ss, t), args...);
+}
+
+template <typename... Args>
+struct _str_wrapper final {
+  static std::string call(const Args&... args) {
+    std::ostringstream ss;
+    _to_str(ss, args...);
+    return ss.str();
+  }
+};
+
+// Specializations for already-a-string types.
+template <>
+struct _str_wrapper<std::string> final {
+  // return by reference to avoid the binary size of a string copy
+  static const std::string& call(const std::string& str) {
+    return str;
+  }
+};
+
+template <>
+struct _str_wrapper<const char*> final {
+  static const char* call(const char* str) {
+    return str;
+  }
+};
+
+// For nvfuser::to_str() with an empty argument list (which is common in our
+// assert macros), we don't want to pay the binary size for constructing and
+// destructing a stringstream or even constructing a string.
+template <>
+struct _str_wrapper<> final {
+  static CompileTimeEmptyString call() {
+    return CompileTimeEmptyString();
+  }
+};
+
+// Convert a list of string-like arguments into a single string.
+template <typename... Args>
+inline decltype(auto) to_str(const Args&... args) {
+  return _str_wrapper<typename CanonicalizeStrTypes<Args>::type...>::call(
+      args...);
+}
+
+class NVF_API nvfError : public std::exception {
+  // The actual error message.
+  std::string msg_;
+  // Context for the message (in order of decreasing specificity).  Context will
+  // be automatically formatted appropriately, so it is not necessary to add
+  // extra leading/trailing newlines to strings inside this vector
+  std::vector<std::string> context_;
+  // The C++ backtrace at the point when this exception was raised.  This
+  // may be empty if there is no valid backtrace.  (We don't use optional
+  // here to reduce the dependencies this file has.)
+  std::string backtrace_;
+  // These two are derived fields from msg_stack_ and backtrace_, but we need
+  // fields for the strings so that we can return a const char* (as the
+  // signature of std::exception requires).  Currently, the invariant
+  // is that these fields are ALWAYS populated consistently with respect
+  // to msg_stack_ and backtrace_.
+  std::string what_;
+  std::string what_without_backtrace_;
+  // This is a little debugging trick: you can stash a relevant pointer
+  // in caller, and then when you catch the exception, you can compare
+  // against pointers you have on hand to get more information about
+  // where the exception came from.  In Caffe2, this is used to figure
+  // out which operator raised an exception.
+  const void* caller_;
+
+ public:
+  // PyTorch-style Error constructor.
+  nvfError(SourceLocation source_location, std::string msg);
+  // Base constructor
+  nvfError(
+      std::string msg,
+      std::string backtrace,
+      const void* caller = nullptr);
+  // Add some new context to the message stack.  The last added context
+  // will be formatted at the end of the context list upon printing.
+  // WARNING: This method is O(n) in the size of the stack, so don't go
+  // wild adding a ridiculous amount of context to error messages.
+  void add_context(std::string msg);
+  const std::string& msg() const {
+    return msg_;
+  }
+  const std::vector<std::string>& context() const {
+    return context_;
+  }
+  const std::string& backtrace() const {
+    return backtrace_;
+  }
+  /// Returns the complete error message, including the source location.
+  /// The returned pointer is invalidated if you call add_context() on
+  /// this object.
+  const char* what() const noexcept override {
+    return what_.c_str();
+  }
+  const void* caller() const noexcept {
+    return caller_;
+  }
+  /// Returns only the error message string, without source location.
+  /// The returned pointer is invalidated if you call add_context() on
+  /// this object.
+  const char* what_without_backtrace() const noexcept {
+    return what_without_backtrace_.c_str();
+  }
+
+ private:
+  void refresh_what();
+  std::string compute_what(bool include_backtrace) const;
+};
+
+[[noreturn]] NVF_API void nvfCheckFail(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const std::string& msg);
+
+[[noreturn]] NVF_API void nvfCheckFail(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* msg);
+
+[[noreturn]] void nvfErrorFail(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* condMsg,
+    const char* userMsg);
+
+[[noreturn]] inline void nvfErrorFail(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* condMsg,
+    CompileTimeEmptyString /*userMsg*/) {
+  nvfCheckFail(func, file, line, condMsg);
+}
+
+[[noreturn]] NVF_API void nvfErrorFail(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* condMsg,
+    const std::string& userMsg);
+
+template <typename... Args>
+decltype(auto) nvfCheckMsgImpl(const char* /*msg*/, const Args&... args) {
+  return to_str(args...);
+}
+inline const char* nvfCheckMsgImpl(const char* msg) {
+  return msg;
+}
+// If there is just 1 user-provided C-string argument, use it.
+inline const char* nvfCheckMsgImpl(const char* /*msg*/, const char* args) {
+  return args;
+}
+
+} // namespace nvfuser
+
+#define STRINGIZE_IMPL(x) #x
+#define STRINGIZE(x) STRINGIZE_IMPL(x)
+
+#define NVF_THROW(...) \
+  nvfuser::nvfErrorFail(                                    \
+        __FUNCTION__,                                       \
+        __FILE__,                                           \
+        static_cast<uint32_t>(__LINE__),                    \
+        " INTERNAL ASSERT FAILED at "                       \
+        STRINGIZE(__FILE__) ":" STRINGIZE(__LINE__)         \
+        ", please report a bug with repro script to NVFuser at " \
+        "https://github.com/NVIDIA/Fuser/issues. ",         \
+        nvfuser::to_str(__VA_ARGS__));
+
+#define NVF_ERROR(cond, ...) \
+  if ((!(cond))) {           \
+    NVF_THROW(__VA_ARGS__)   \
+  }
+
+#define NVF_CHECK_MSG(cond, type, ...) \
+  (nvfuser::nvfCheckMsgImpl(           \
+      "Expected " #cond " to be true, but got false.  ", ##__VA_ARGS__))
+
+#define NVF_CHECK(cond, ...)                     \
+  if ((!(cond))) {                               \
+    nvfuser::nvfCheckFail(                       \
+        __func__,                                \
+        __FILE__,                                \
+        static_cast<uint32_t>(__LINE__),         \
+        NVF_CHECK_MSG(cond, "", ##__VA_ARGS__)); \
+  }

nvfuser/include/nvfuser/expr_evaluator.h

@@ -0,0 +1,125 @@
+// 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 <evaluator_common.h>
+#include <exceptions.h>
+#include <ir/cloner.h>
+#include <ir/interface_nodes.h>
+#include <iter_visitor.h>
+#include <logical_domain_map.h>
+#include <polymorphic_value.h>
+#include <visibility.h>
+
+#include <string>
+#include <unordered_map>
+
+namespace nvfuser {
+
+class PrecomputedValues;
+
+//! Calculate Fusion IR expressions
+class ExpressionEvaluator {
+ public:
+  //! Bind a concrete value to an IR variable
+  //! If evaluate_validate is true, and value is evaluatable with the
+  //! information already known, then evaluate and validate the value with the
+  //! concrete value.
+  void bind(
+      const Val* value,
+      PolymorphicValue concrete_value,
+      bool evaluate_validate = false) {
+    bind_(value, std::move(concrete_value), evaluate_validate);
+  }
+
+  //! Bind a concrete value to a named scalar
+  void bind(const std::string& name, PolymorphicValue concrete_value) {
+    bind_(name, std::move(concrete_value));
+  }
+
+  //! Set a concrete value for a parallel dimension
+  void bind(ParallelType pt, PolymorphicValue concrete_value);
+
+  //! Try to evaluate a Fusion IR value
+  NVF_API const PolymorphicValue& evaluate(const Val* value);
+
+  //! Try to evaluate a parallel dimension
+  const PolymorphicValue& evaluate(ParallelType pt);
+
+  //! Evaluates a value through a const evaluator reference.
+  //! Initializes a known_values map to store intermediate values in lieu of
+  //! known_values_.
+  NVF_API PolymorphicValue evaluate(const Val* value) const;
+
+  //! Base evaluate method called by other overloads and Expr::evaluate.
+  const PolymorphicValue& evaluate(
+      const Val* value,
+      std::unordered_map<const Val*, PolymorphicValue>& known_values) const;
+
+  bool isKnown(const Val* value) const {
+    return known_values_.count(value) > 0;
+  }
+
+  void invalidate(const Val* value) {
+    known_values_.erase(value);
+  }
+
+  //! Debugging helper, prints all the currently known values
+  void print() const;
+
+  void bindPrecomputedValues(PrecomputedValues* precomputed_values) {
+    precomputed_values_ = precomputed_values;
+  }
+
+  auto& precomputedValues() {
+    return precomputed_values_;
+  }
+
+  //! Augment the evaluator with the exact root-domain map such that
+  //! if the extent of a root ID is known, the extents of all other
+  //! root IDs that are exactly mapped also get bound to the same
+  //! value. This is currently just done with ExactLogicalDomainMap, but
+  //! can be similarly done with the Exact CA map as well.
+  void propagateBoundValuesThroughExactMaps(
+      Fusion* fusion,
+      ExactLogicalDomainMap* exact_map = nullptr);
+
+  ExpressionEvaluator clone(IrCloner& ir_cloner) const;
+
+ private:
+  void bind_(
+      const Val* value,
+      PolymorphicValue concrete_value,
+      bool evaluate_validate);
+
+  void bind_(const std::string& name, PolymorphicValue concrete_value);
+
+  void bindTensorDomain(
+      const TensorView* tv,
+      const at::Tensor& t,
+      bool evaluate_validate);
+
+  const PolymorphicValue& getValue(
+      const Val* value,
+      const std::unordered_map<const Val*, PolymorphicValue>&
+          additional_known_values) const;
+
+ private:
+  // TODO: Consider make this const. It can't be const as bind() of
+  // this class calls
+  // PrecomputedValuess::bindConcreteParallelTypeValue, but it's
+  // unclear why the precompute values cannot be kept constant and
+  // binding a value to ExpressionEvaluator just updates
+  // known_named_scalars_.
+  PrecomputedValues* precomputed_values_ = nullptr;
+  std::unordered_map<const Val*, PolymorphicValue> known_values_;
+  std::unordered_map<std::string, PolymorphicValue> known_named_scalars_;
+  PolymorphicValue null_ = std::monostate{};
+};
+
+} // namespace nvfuser