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

nvfuser/include/nvfuser/scheduler/compile_time_info.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 <fusion.h>
+#include <scheduler/pointwise_utils.h>
+#include <scheduler/scheduler_types.h>
+#include <scheduler/tools/domain_map.h>
+#include <scheduler/utils.h>
+#include <scheduler/vectorize_helper.h>
+
+namespace nvfuser {
+
+//! namespace for hosting catalog of possible compile time
+//!  info that can be cached. Each possible entry type has
+//!  a value in `CompileTimeEntryType` and an entry type class
+//!  definition like `VectorizableInputsAndOutputs`. The corresponnding
+//!  classes contain their entry type, data type and maybe more
+//!  later depending on use cases.
+namespace HeuristicCompileTime {
+
+//! Each entry type under this category represent some information
+//!  that can be inferred compile-time, i.e. without any runtime input
+//!  meta data. They will be stored in `HeuristicDataCache` and will
+//!  be re-used each time the same fusion is visited.
+
+//! Enum for all possible types of cached entries of compile-time info.
+enum class CompileTimeEntryType {
+  DOMAIN_MAP,
+  TRANSPOSE_DOMAIN_MAP,
+  REFERENCE_TENSORS,
+  REFERENCE_TENSORS_FOR_GROUPS,
+  VECTORIZABLE_INPUTS_AND_OUTPUTS,
+  INPUTS_AND_OUTPUTS_INNER_DIM_GROUPS,
+  TV_TO_CONTIG_INNER_SIZE_MAPS,
+  UNROLLABLE_INPUTS_AND_OUTPUTS,
+  REDUCTION_TVS,
+  PERSISTENT_BUFFER_INFO,
+  SCOPE_PERSISTENT_FACTOR_INFO,
+  BROADCAST_BYTE_MULTIPLES,
+  INNER_MOST_DIMS_INFO,
+  CAN_SCHEDULE_TRANSPOSE,
+  CAN_SCHEDULE_MUL_SUM_AS_MMA,
+  LOGICAL_REORDER_MAP,
+  VECTORIZATION_BREAK_POINT_OF_RED_PROD,
+  SCHEDULE_HYPERPARAMETERS
+};
+
+//! Entry type definition class for `DOMAIN_MAP`,
+//!  stores the domain map of a fusion.
+class DomainMap {
+ public:
+  using DataType = scheduler_tools::DomainMap;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::DOMAIN_MAP;
+};
+
+//! Entry type definition class for `DOMAIN_MAP`,
+//!  stores the domain map of a fusion, used by transpose scheduler.
+class TransposeDomainMap {
+ public:
+  using DataType = scheduler_tools::DomainMap;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::TRANSPOSE_DOMAIN_MAP;
+};
+
+//! Entry type definition class for `REFERENCE_TENSORS`,
+//!  stores the the reference TensorViews used to schedule a fusion.
+class ReferenceTensors {
+ public:
+  using DataType = std::vector<TensorView*>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::REFERENCE_TENSORS;
+};
+
+//! Entry type definition class for `REFERENCE_TENSORS`,
+//!  stores the the reference TensorViews used to schedule a fusion, used by
+//!  transpose scheduler.
+class ReferenceTensorsForGroups {
+ public:
+  using DataType = std::vector<TensorView*>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::REFERENCE_TENSORS_FOR_GROUPS;
+};
+
+//! Entry type definition class for `VECTORIZABLE_INPUTS_AND_OUTPUTS`,
+//!  stores the vectorizable TensorViews on a fusion's inputs and outputs.
+class VectorizableInputsAndOutputs {
+ public:
+  using DataType = std::vector<TensorView*>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::VECTORIZABLE_INPUTS_AND_OUTPUTS;
+};
+
+//! Entry type definition class for `TV_TO_CONTIG_INNER_SIZE_MAPS`,
+//!  stores the vectorizable TensorViews on a fusion's inputs and outputs.
+class TvToContigInnerSizeMaps {
+ public:
+  using DataType = std::vector<std::unordered_map<TensorView*, Val*>>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::TV_TO_CONTIG_INNER_SIZE_MAPS;
+};
+
+//! Entry type definition class for `INPUTS_AND_OUTPUTS_INNER_DIM_GROUPS`,
+//!  stores the fusion's inputs and outputs grouped by inner most dimension.
+class InputsOutputsInnerDimGroups {
+ public:
+  using DataType = std::vector<std::vector<TensorView*>>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::INPUTS_AND_OUTPUTS_INNER_DIM_GROUPS;
+};
+
+//! Entry type definition class for `UNROLLABLE_INPUTS_AND_OUTPUTS`,
+//!  stores the unrollable TensorViews on a fusion's inputs and outputs.
+class UnrollableInputsAndOutputs {
+ public:
+  using DataType = std::vector<TensorView*>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::UNROLLABLE_INPUTS_AND_OUTPUTS;
+};
+
+//! Entry type definition class for `REDUCTION_TVS`,
+//!  stores the all tvs with reduction axes in a fusion.
+class ReductionTVs {
+ public:
+  using DataType = std::vector<TensorView*>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::REDUCTION_TVS;
+};
+
+//! Entry type definition class for `PERSISTENT_BUFFER_INFO`,
+//!  stores persistent buffers inferred from topology and scheduling of fusion.
+class PersistentBufferInfo {
+ public:
+  using DataType = scheduler_utils::PersistentBufferInfo;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::PERSISTENT_BUFFER_INFO;
+};
+
+//! Entry type definition class for `INNER_MOST_DIMS_INFO`,
+//!  Used in the transpose scheduler to store inner most IterDomains and their
+//!  position in reference1 of group 1 and group 2
+//!  Note, negative value indicates mapping failure
+class InnerMostDimInfo {
+ public:
+  using DataType = std::vector<int64_t>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::INNER_MOST_DIMS_INFO;
+};
+
+//! Auxiliary data types for `SCOPE_PERSISTENT_FACTOR_INFO` entry type.
+using ScopedPersistenceBufferMap = std::unordered_map<Val*, std::vector<bool>>;
+
+//! Entry type definition class for `SCOPE_PERSISTENT_FACTOR_INFO`,
+// Tracks which buffers are active at a given Val*, order of bool vector is
+// based on persistence buffer order from persistence buffer info, this is then
+// appended by the projectable persistent buffers' inputs. True in the bool
+// vector means the persistent buffer is active at the generation of the key.
+class ScopePersistentFactorInfo {
+ public:
+  using DataType = ScopedPersistenceBufferMap;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::SCOPE_PERSISTENT_FACTOR_INFO;
+};
+
+//! Entry type definition class for `BROADCAST_BYTE_MULTIPLES`,
+//!  stores "byte multiples" information. This information can be used to figure
+//!  out if using a 2D scheduler how many bytes have to be transferred with
+//!  varying split locations. See BroadcastMultiple definition for more
+//!  information.
+class BroadcastMultiples {
+ public:
+  using DataType = scheduler_utils::BroadcastMultipleInformation;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::BROADCAST_BYTE_MULTIPLES;
+};
+
+//! Entry type definition class for `CAN_SCHEDULE_TRANSPOSE`,
+//!  stores if the transpose scheduler can scheduler this fusion
+class CanScheduleTranspose {
+ public:
+  using DataType = bool;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::CAN_SCHEDULE_TRANSPOSE;
+};
+
+//! Entry type definition class for `LOGICAL_REORDER_MAP`,
+//!  stores the domain map of a fusion.
+class LogicalReorderMap {
+ public:
+  using DataType = std::unordered_map<int64_t, int64_t>;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::LOGICAL_REORDER_MAP;
+};
+
+class VectorizationBreakPointOfReductionProducer {
+ public:
+  using DataType = int64_t;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::VECTORIZATION_BREAK_POINT_OF_RED_PROD;
+};
+
+//! Entry type definition class for `SCHEDULE_HYPERPARAMETERS`,
+//!  stores hyperparameters for SchedulerEntry::computeHeuristics
+class SchedulerHyperParameters {
+ public:
+  using DataType = scheduler_utils::SchedulerHyperParameters;
+  static const CompileTimeEntryType EntryType =
+      CompileTimeEntryType::SCHEDULE_HYPERPARAMETERS;
+};
+
+//! Base abstract class for unified storage in `HeuristicDataCache`,
+//!  each entry in `HeuristicDataCache` will be a subclass.
+class CompileTimeInfoBase : public PolymorphicBase {
+ public:
+  CompileTimeInfoBase(CompileTimeEntryType entry_type)
+      : entry_type_(entry_type) {}
+  CompileTimeEntryType type() {
+    return entry_type_;
+  }
+
+ private:
+  CompileTimeEntryType entry_type_;
+};
+
+} // namespace HeuristicCompileTime
+
+//! Note: Do NOT export this class. MSVC issue with exported class that contains
+//! std::vector<unique_ptr<xxx>>: https://godbolt.org/z/3E4e8T1P1
+//! Compile-time information cache for `canSchedule` and `getHeuristics`
+//! interfaces. Each cache instance stores information that could be inferred at
+//! compile time in a fusion and therefore corresponds to an instance of
+//! KernelExecutor.
+class HeuristicDataCache {
+  using EntryOwningPtr =
+      std::unique_ptr<HeuristicCompileTime::CompileTimeInfoBase>;
+  using EntryPtr = HeuristicCompileTime::CompileTimeInfoBase*;
+  using EntryType = HeuristicCompileTime::CompileTimeEntryType;
+
+ public:
+  bool hasEntry(EntryType entry_type) {
+    return entry_type_map_.find(entry_type) != entry_type_map_.end();
+  }
+
+  void insert(EntryOwningPtr new_entry);
+
+  EntryPtr at(EntryType entry_type) {
+    return entry_type_map_.at(entry_type);
+  }
+
+ private:
+  std::vector<EntryOwningPtr> entries_;
+  std::unordered_map<EntryType, EntryPtr> entry_type_map_;
+};
+
+//! A utility class to facilitate accessing HeuristicDataCache.
+//!  This utility is needed because the information to be stored
+//!    in HeuristicDataCache is used in several different scenarios
+//!    and we want to support all these use cases in one interface.
+//!  The current use examples are:
+//!   1. During fusion segmentation process, all the fusions
+//!     given to canSchedule are temporary and therefore the
+//!     compile time info do not need to be cached, and in fact
+//!     a cache wouldn't be instantiated by that time.
+//!
+//!   2. When a kernel is created for the first time, entries will be
+//!     missing in the cache and all the computed information will be
+//!     captured and written into the cache.
+//!
+//!   3. When we check a compiled fusion for heuristic hit, we want to
+//!     use the cached info to save runtime latency.
+//!
+//! The designed interface is used as:
+//!   auto entry = HeuristicDataCacheEntry<EntryClass>(data_cache, maker_fn);
+//!   auto& data = entry.get();
+//!
+//!  `maker_fn` will be called to compute the information when no cached data
+//!   exists and `entry` will own the computed data when no data cache is
+//!   supplied.
+template <typename EntryClass>
+class HeuristicDataCacheEntry {
+  using EntryDataType = typename EntryClass::DataType;
+  using EntryDataTypeOwnPtr = std::unique_ptr<EntryDataType>;
+  using MakerFnType = std::function<EntryDataTypeOwnPtr()>;
+
+ public:
+  //! Creates a data entry with type defined in EntryClass,
+  //!  eg. EntryClass = VectorizableInputsAndOutputs;
+  //!
+  //! @param data_cache, a pointer to an instantiated compile-time
+  //!  info cache. The info data will be
+  //!    1. read from data cache if data cache is not recording.
+  //!    2. written into  data cache if data cache is recording.
+  //!    3. managed by owned_data_ if data cache is nullptr
+  //! @param fn:
+  //!   The factory function that needs to return a owning pointer
+  //!  i.e. std::unique_ptr<EntryClass::DataType>. It will only
+  //!  be called either when data cache is recording or when no data
+  //!  cache is given.
+  HeuristicDataCacheEntry(HeuristicDataCache* data_cache, MakerFnType fn);
+
+  //! Unified interface to get actual data, either from cache
+  //!  or from factory function.
+  EntryDataType& get() {
+    return *data_ptr_;
+  }
+
+ private:
+  //! Internal data owing pointer that will manage the computed
+  //!  data where there is no data cache.
+  EntryDataTypeOwnPtr owned_data_ = nullptr;
+
+  //! Pointer to the valid data entry that could be accessed.
+  EntryDataType* data_ptr_ = nullptr;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/debug_utils.h

@@ -0,0 +1,68 @@
+// 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 <debug.h>
+#include <options.h>
+#include <utils.h>
+
+#include <iostream>
+
+namespace nvfuser {
+
+namespace scheduler_debug_utils {
+
+// Basic logging utility for any messages in scheduler or segmenter
+template <typename... Args>
+void canScheduleMessage(const Args&... args) {
+  // Using builtin expect to reduce the overhead slightly,
+  //  alternatively may want to allow this message in debug
+  //  build only but that'd be inconvenient for user support.
+  if (C10_UNLIKELY(
+          isDebugDumpEnabled(DebugDumpOption::FusionSegmenterLog) ||
+          isDebugDumpEnabled(DebugDumpOption::SchedulerVerbose))) {
+    debug() << c10::str(args...) << "\n";
+  }
+}
+
+// Short-cut message for flagging why shedulers cannot schedule fusions,
+//  assuming first argument is heuristic type (not actively checked).
+template <typename SchedulerType, typename... Args>
+void canScheduleRejectReason(
+    SchedulerType scheduler_type,
+    const Args&... args) {
+  canScheduleMessage(
+      "Scheduler _", scheduler_type, "_ ***rejected*** because : ", args...);
+}
+
+// Based on
+// https://learn.microsoft.com/en-us/cpp/cpp/ellipses-and-variadic-templates?view=msvc-170#example
+inline void log() {
+  if (isDebugDumpEnabled(DebugDumpOption::SchedulerVerbose)) {
+    debug() << std::endl;
+  }
+}
+
+template <typename T>
+void log(const T& t) {
+  if (isDebugDumpEnabled(DebugDumpOption::SchedulerVerbose)) {
+    debug() << t << std::endl;
+  }
+}
+
+template <typename First, typename... Rest>
+void log(const First& first, const Rest&... rest) {
+  if (isDebugDumpEnabled(DebugDumpOption::SchedulerVerbose)) {
+    debug() << first;
+    log(rest...);
+  }
+}
+
+} // namespace scheduler_debug_utils
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/expr_eval_sched.h

@@ -0,0 +1,45 @@
+// 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 <scheduler/heuristic.h>
+#include <scheduler/registry.h>
+
+namespace nvfuser {
+
+class Fusion;
+class SchedulerRuntimeInfo;
+class HeuristicDataCache;
+
+// ExprEval scheduler represents the case where we allocate outputs directly
+// using EE. No code is generated.
+class ExprEvalScheduler : public SchedulerEntry {
+ public:
+  // This scheduler only accepts MatmulOp.
+  bool canScheduleCompileTime(Fusion* fusion) override;
+
+  bool canScheduleRunTime(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache) override {
+    return true;
+  }
+
+  std::unique_ptr<HeuristicParams> computeHeuristics(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache) override;
+
+  void schedule(Fusion* fusion, const HeuristicParams* params) override;
+
+  constexpr static SchedulerType schedulerType() {
+    return SchedulerType::ExprEval;
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/heuristic.h

@@ -0,0 +1,113 @@
+// 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 <runtime/executor_params.h>
+#include <scheduler/scheduler_types.h>
+#include <utils.h>
+
+#include <string>
+
+namespace nvfuser {
+
+class SchedulerRuntimeInfo;
+class HeuristicDataCache;
+
+// Top-level class representing heuristic parameters. Most schedulers
+// have their own subclasses to have their specific parameters, except
+// for ExprEval schedulers.
+class HeuristicParams : public PolymorphicBase {
+ public:
+  std::string tag = "";
+
+  LaunchParams lparams;
+  CompileParams cparams;
+  const SchedulerType scheduler_type;
+
+  virtual std::string toString() const {
+    std::stringstream ss;
+    ss << "Heuristic Params (" << scheduler_type << ")";
+    return ss.str();
+  }
+
+  virtual size_t hash() const {
+    return 0;
+  };
+
+  virtual bool sameAs(const HeuristicParams* other) const {
+    if (!other->isStrictlyA<HeuristicParams>()) {
+      return false;
+    }
+    if (other->scheduler_type != scheduler_type) {
+      return false;
+    }
+    return other->cparams == cparams;
+  }
+
+  HeuristicParams() = delete;
+  explicit HeuristicParams(SchedulerType _scheduler_type)
+      : scheduler_type(_scheduler_type) {};
+
+  virtual std::unique_ptr<HeuristicParams> clone() const {
+    return std::make_unique<HeuristicParams>(*this);
+  }
+};
+
+//! Auxiliary class for storing heuristics. The managed data is either
+//!  a single heursitic for complete fusion, or a vector of heuristics used for
+//!  a segmented fusion.
+class HeuristicParamsList {
+ public:
+  //! Constructor for segmented fusion case. Created with empty list and
+  //!  uses emplaceBack for inserting heuristics in order
+  explicit HeuristicParamsList() = default;
+
+  //! Constructor fills heuristics_ with nullptr, which allows us to create
+  //! SchedulerEntries out of order.
+  explicit HeuristicParamsList(size_t num_heuristics) {
+    heuristics_.reserve(num_heuristics);
+    std::fill_n(std::back_inserter(heuristics_), num_heuristics, nullptr);
+  }
+
+  //! Constructor for complete fusion case, generates the scheduler entry
+  //!  for the fusion owning the given expression
+  explicit HeuristicParamsList(
+      SchedulerType scheduler_type,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr);
+
+  HeuristicParamsList(const HeuristicParamsList&) = delete;
+  HeuristicParamsList& operator=(const HeuristicParamsList&) = delete;
+
+  std::unique_ptr<HeuristicParams>& at(int index) {
+    return heuristics_.at(index);
+  }
+
+  //! Place a heuristics on the list. Applies to segmented fusion only.
+  void emplaceBack(std::unique_ptr<HeuristicParams>&& pt) {
+    NVF_ERROR(is_segmented_);
+    heuristics_.emplace_back(std::move(pt));
+  }
+
+  //! Returns list of heuristics for a segmneted fusion.
+  const std::vector<std::unique_ptr<HeuristicParams>>& heuristicsList() const {
+    return heuristics_;
+  }
+
+  //! Returns the single heuristics for a complete fusion.
+  HeuristicParams* singleKernelHeuristics() const {
+    NVF_ERROR(!is_segmented_);
+    return heuristics_.begin()->get();
+  }
+
+ private:
+  std::vector<std::unique_ptr<HeuristicParams>> heuristics_;
+  bool is_segmented_ = true;
+};
+
+} // namespace nvfuser