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

nvfuser/include/nvfuser/scheduler/reduction_heuristic.h

@@ -0,0 +1,339 @@
+// 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 <sstream>
+
+namespace nvfuser {
+
+class TensorView;
+
+// Parameters of the reduction heuristic to describe the optimal schedule.
+// Warning: equal operator is intended for use in caching the kernel associated
+// with these reduction parameters. It does not check if the launch parameters
+// are equivelent!
+class ReductionParams : public HeuristicParams {
+ public:
+  // Note that heuristictype can be different from SchedulerType::Reduction
+  // since ReductionParams is also used by, e.g., normalization schedulers.
+  ReductionParams(SchedulerType scheduler_type = SchedulerType::Reduction)
+      : HeuristicParams(scheduler_type) {};
+  // Reducing inner most dimension?
+  bool fastest_dim = false;
+
+  // Store input in shared memory or registers to reduce global memory reads
+  bool persistent_kernel = false;
+
+  // Project persistent buffers back to inputs to reduce persistent buffer size
+  bool project_persistent_buffers = false;
+
+  // Are we treating the scheduling as 3 dimensional, can be useful for patterns
+  // like [reduction, iteration, reduction].
+  bool schedule_3D = false;
+
+  // For outer reductions we may want to swap the gdimx and gdimy bindings to
+  // amortize the cost of the final cleanup in grid reductions.
+  bool flip_grid = false;
+
+  // Inner Reduction Domain:
+
+  // Reduce across the block?
+  bool cross_block_inner_reduction = false;
+  // Reduce across the grid?
+  bool cross_grid_inner_reduction = false;
+  // Unrolling/Vectorization factor for inner reduction dimension
+  int64_t unroll_factor_inner_reduction = 1;
+
+  // Extra unroll on top of vectorization
+  int64_t unroll_factor_top_of_vectorization = 1;
+
+  // vectorize instead of unroll
+  bool vectorize_inner_reduction = false;
+  // Split grid dim for iteration axis in case it's too large for cuda
+  bool split_grid_dim_inner_reduction = false;
+  // Pad inner dimension to nearest warp
+  bool pad_inner_reduction_to_warp = false;
+  // Register persistent buffer size in inner dimension
+  int64_t batches_per_block_inner_reduction = 1;
+
+  // Which block parallel dimension should be used for the inner reduction.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType block_dim_inner_reduction = ParallelType::Serial;
+  // Which grid parallel dimension should be used for the inner reduction.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType grid_dim_inner_reduction = ParallelType::Serial;
+
+  // Iteration Domain:
+
+  // Perform multiple reductions per block?
+  bool multiple_reds_per_blk = false;
+  // Unrolling/Vectorization factor for iteration dimension
+  int64_t unroll_factor_iter_dom = 1;
+  // vectorize instead of unroll
+  bool vectorize_iter_dom = false;
+  // Inner split grid dim for iteration axis in case it's too large for cuda
+  bool split_grid_dim_iter_dom_inner = false;
+  // Outer split grid dim for iteration axis in case it's too large for cuda
+  bool split_grid_dim_iter_dom_outer = false;
+
+  // Which block parallel dimension should be used for the iter domain.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType block_dim_iter_dom = ParallelType::Serial;
+  // Which grid parallel dimension should be used for the iter domain.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType grid_dim_iter_dom = ParallelType::Serial;
+
+  // Outer Reduction Domain if 3D Scheduled:
+
+  // Reduce across the block?
+  bool cross_block_outer_reduction = false;
+  // Reduce across the grid?
+  bool cross_grid_outer_reduction = false;
+  // Split grid dim for iteration axis in case it's too large for cuda
+  bool split_grid_dim_outer_reduction = false;
+  // Register persistent buffer size in outer dimension
+  int64_t batches_per_block_outer_reduction = 1;
+  // Unrolling/Vectorization factor for outer reduction factor
+  int64_t unroll_factor_outer_reduction = 1;
+
+  // Which block parallel dimension should be used for the outer reduction.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType block_dim_outer_reduction = ParallelType::Serial;
+  // Which grid parallel dimension should be used for the outer reduction.
+  // !!WARNING!! Convenience method, this be unique based on non-parallel type
+  // parameters, not used for equivalence/hashing.
+  ParallelType grid_dim_outer_reduction = ParallelType::Serial;
+
+  // Use computeWith to persistent buffers
+  bool compute_persistent_buffer_with_first_consumer = false;
+
+  bool static_bdimx = false;
+  bool static_bdimy = false;
+
+  bool isUnrolled() const {
+    return unroll_factor_inner_reduction > 1 || unroll_factor_iter_dom > 1 ||
+        unroll_factor_outer_reduction > 1;
+  }
+
+  // specific to combined inner and outer reduction
+  bool combined_inner_outer = false;
+  // use TIDx for out reduction axis
+  bool tidx_for_outer_reduction = false;
+  // pad outer reduction to warp
+  bool pad_outer_reduction_to_warp = false;
+  // in outer reduction part of inner-outer persistent scheduler, may further
+  // split inner dim by grid
+  bool combined_split_grid_inner_dim = false;
+  // partial result of outer reduction is written to gmem then read back in a
+  // different parallel pattern set the vectorization factor of its read and
+  // write
+  int64_t vectorization_factor_outer = 1;
+  int64_t vectorization_factor_tmp_gmem_write = 1;
+  // inner reduction axis is parallelized by block_dim_inner_reduction (usually
+  // TIDx) the remaining part is further parallelized by
+  // block_dim_inner_reduction_extra (usually TIDy)
+  ParallelType block_dim_inner_reduction_extra = ParallelType::Serial;
+
+  // vector stores buffer should be moved to shared memory.
+  // TODO: For innerOuterPersistentHeuristic, only the persistent tensors in the
+  // original fusion definition may be moved to shared memory, the intermediate
+  // persistent tensors which are creased by the scheduler to store the partial
+  // outer reduction results are always stored in registers because they are
+  // frequently accessed by both read and write. The code can be extended to
+  // allow the move of these intermediate persistent tensors to shared memory
+  // when the shared memory is much larger than the register file.
+  std::vector<TensorView*> smem_persistent_buffers;
+
+ public:
+  using HeuristicParams::HeuristicParams;
+
+  // Warning: Does not check launch parameters!
+  bool sameAs(const HeuristicParams* other_base) const override {
+    auto other = dynamic_cast<const ReductionParams*>(other_base);
+    if (other == nullptr) {
+      return false;
+    }
+
+    bool attr_equal = other->cparams == cparams &&
+        other->fastest_dim == fastest_dim &&
+        other->persistent_kernel == persistent_kernel &&
+        other->project_persistent_buffers == project_persistent_buffers &&
+        other->schedule_3D == schedule_3D && other->flip_grid == flip_grid &&
+        other->cross_block_inner_reduction == cross_block_inner_reduction &&
+        other->cross_grid_inner_reduction == cross_grid_inner_reduction &&
+        other->unroll_factor_inner_reduction == unroll_factor_inner_reduction &&
+        other->vectorize_inner_reduction == vectorize_inner_reduction &&
+        other->split_grid_dim_inner_reduction ==
+            split_grid_dim_inner_reduction &&
+        other->pad_inner_reduction_to_warp == pad_inner_reduction_to_warp &&
+        other->batches_per_block_inner_reduction ==
+            batches_per_block_inner_reduction &&
+        other->multiple_reds_per_blk == multiple_reds_per_blk &&
+        other->unroll_factor_iter_dom == unroll_factor_iter_dom &&
+        other->vectorize_iter_dom == vectorize_iter_dom &&
+        other->split_grid_dim_iter_dom_inner == split_grid_dim_iter_dom_inner &&
+        other->split_grid_dim_iter_dom_outer == split_grid_dim_iter_dom_outer &&
+        other->cross_block_outer_reduction == cross_block_outer_reduction &&
+        other->cross_grid_outer_reduction == cross_grid_outer_reduction &&
+        other->unroll_factor_outer_reduction == unroll_factor_outer_reduction &&
+        other->split_grid_dim_outer_reduction ==
+            split_grid_dim_outer_reduction &&
+        other->batches_per_block_outer_reduction ==
+            batches_per_block_outer_reduction &&
+        other->compute_persistent_buffer_with_first_consumer ==
+            compute_persistent_buffer_with_first_consumer &&
+        other->combined_inner_outer == combined_inner_outer &&
+        other->tidx_for_outer_reduction == tidx_for_outer_reduction &&
+        other->pad_outer_reduction_to_warp == pad_outer_reduction_to_warp &&
+        other->vectorization_factor_outer == vectorization_factor_outer &&
+        other->combined_split_grid_inner_dim == combined_split_grid_inner_dim &&
+        other->unroll_factor_top_of_vectorization ==
+            unroll_factor_top_of_vectorization &&
+        other->vectorization_factor_tmp_gmem_write ==
+            vectorization_factor_tmp_gmem_write;
+
+    if (other->static_bdimy || static_bdimy) {
+      attr_equal = attr_equal && other->lparams.bdimy() == lparams.bdimy();
+    }
+    if (other->static_bdimx || static_bdimx) {
+      attr_equal = attr_equal && other->lparams.bdimx() == lparams.bdimx();
+    }
+    return attr_equal;
+  }
+
+  std::string toString() const override {
+    std::stringstream ss;
+    ss << "\n===== Reduction Parameters ========\n"
+       << (tag.empty() ? "" : "Tag: ") << tag << "\n"
+       << (fastest_dim ? "Red On Fastest Dim\n" : "Red On Slow Dim\n")
+       << (persistent_kernel ? "Persistent Kernel\n" : "")
+       << (project_persistent_buffers ? "Project Persistent Buffers\n" : "");
+    if (batches_per_block_inner_reduction > 1 || persistent_kernel) {
+      ss << "Batches per block: " << batches_per_block_inner_reduction << "\n";
+    }
+
+    if (schedule_3D) {
+      ss << "3D Schedule\n"
+         << "Outer Reduction: ";
+      if (cross_block_outer_reduction) {
+        ss << "cross block - " << block_dim_outer_reduction << " / ";
+      }
+      if (cross_grid_outer_reduction) {
+        ss << "cross grid - " << grid_dim_outer_reduction << " / ";
+        ss << (split_grid_dim_outer_reduction ? "split grid dim / " : "");
+      }
+
+      ss << (unroll_factor_outer_reduction > 1 ? "unroll / " : "");
+      if (unroll_factor_outer_reduction > 1) {
+        ss << "factor " << unroll_factor_outer_reduction << " ";
+      }
+
+      if (batches_per_block_outer_reduction > 1 || persistent_kernel) {
+        ss << "persistent batch - " << batches_per_block_outer_reduction;
+      }
+    }
+
+    ss << "\nIteration Domain: ";
+
+    if (grid_dim_iter_dom != ParallelType::Serial) {
+      ss << grid_dim_iter_dom << " / ";
+      if (split_grid_dim_iter_dom_outer) {
+        ss << "split grid dimension outer / ";
+      } else if (split_grid_dim_iter_dom_inner) {
+        ss << "split grid dimension inner / ";
+      }
+    }
+    if (block_dim_iter_dom != ParallelType::Serial) {
+      ss << block_dim_iter_dom << " / ";
+    }
+    ss << (multiple_reds_per_blk ? "multiple reductions per block / " : "")
+       << (vectorize_iter_dom ? "vectorize / " : "")
+       << (unroll_factor_iter_dom > 1 && !vectorize_iter_dom ? "unroll / "
+                                                             : "");
+    if (unroll_factor_iter_dom > 1) {
+      ss << "factor " << unroll_factor_iter_dom;
+    }
+
+    ss << "\nInner Reduction Domain: ";
+
+    if (cross_block_inner_reduction) {
+      ss << "cross block - " << block_dim_inner_reduction << " / ";
+      ss << (pad_inner_reduction_to_warp ? " pad to warp / " : "");
+    }
+    if (cross_grid_inner_reduction) {
+      ss << "cross grid - " << grid_dim_inner_reduction << " / ";
+      ss << (split_grid_dim_inner_reduction ? "split grid dim / " : "");
+    }
+    if (batches_per_block_inner_reduction > 1 || persistent_kernel) {
+      ss << "persistent batch - " << batches_per_block_inner_reduction << " / ";
+    }
+    ss << (cross_grid_inner_reduction && split_grid_dim_inner_reduction
+               ? "split grid dimension / "
+               : "")
+       << (vectorize_inner_reduction ? "vectorize / " : "")
+       << (unroll_factor_inner_reduction > 1 && !vectorize_inner_reduction
+               ? "unroll / "
+               : "");
+    if (unroll_factor_inner_reduction > 1) {
+      ss << "factor " << unroll_factor_inner_reduction;
+    }
+
+    if (compute_persistent_buffer_with_first_consumer) {
+      ss << "\ncomputeWith persistent buffers";
+    }
+
+    ss << "\n" << lparams.toString();
+    ss << cparams.toString() << "\n";
+    ss << "====================================\n";
+    return ss.str();
+  }
+
+  // Warning: Hash is not based on launch parameters!
+  size_t hash() const override {
+    constexpr size_t bits = sizeof(std::size_t) * 8;
+    size_t attr_hash = static_cast<size_t>(fastest_dim) << (bits - 1) ^
+        static_cast<size_t>(persistent_kernel) << (bits - 2) ^
+        static_cast<size_t>(project_persistent_buffers) << (bits - 3) ^
+        static_cast<size_t>(schedule_3D) << (bits - 4) ^
+        static_cast<size_t>(flip_grid) << (bits - 5) ^
+        static_cast<size_t>(cross_block_inner_reduction) << (bits - 6) ^
+        static_cast<size_t>(cross_grid_inner_reduction) << (bits - 7) ^
+        static_cast<size_t>(unroll_factor_inner_reduction) << (bits - 8) ^
+        static_cast<size_t>(vectorize_inner_reduction) << (bits - 9) ^
+        static_cast<size_t>(split_grid_dim_inner_reduction) << (bits - 10) ^
+        static_cast<size_t>(pad_inner_reduction_to_warp) << (bits - 11) ^
+        static_cast<size_t>(batches_per_block_inner_reduction) << (bits - 12) ^
+        static_cast<size_t>(multiple_reds_per_blk) << (bits - 13) ^
+        static_cast<size_t>(unroll_factor_iter_dom) << (bits - 14) ^
+        static_cast<size_t>(vectorize_iter_dom) << (bits - 15) ^
+        static_cast<size_t>(split_grid_dim_iter_dom_outer) << (bits - 16) ^
+        static_cast<size_t>(split_grid_dim_iter_dom_inner) << (bits - 17) ^
+        static_cast<size_t>(cross_block_outer_reduction) << (bits - 18) ^
+        static_cast<size_t>(cross_grid_outer_reduction) << (bits - 19) ^
+        static_cast<size_t>(split_grid_dim_outer_reduction) << (bits - 20) ^
+        static_cast<size_t>(batches_per_block_outer_reduction) << (bits - 21) ^
+        static_cast<size_t>(unroll_factor_outer_reduction) << (bits - 22) ^
+        static_cast<size_t>(compute_persistent_buffer_with_first_consumer)
+            << (bits - 23) ^
+        static_cast<size_t>(unroll_factor_top_of_vectorization) << (bits - 24);
+    return attr_hash;
+  }
+
+  std::unique_ptr<HeuristicParams> clone() const override {
+    return std::make_unique<ReductionParams>(*this);
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/reduction_utils.h

@@ -0,0 +1,159 @@
+// 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.h>
+#include <ir/all_nodes.h>
+#include <scheduler/reduction_heuristic.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+namespace reduction_scheduler_utils {
+
+// Consistent parallelization based on provided reduction parameters. Provided
+// tensor is expected to be reduced by canonicalDimReduction before sending
+// here. reduction_tv should be provided as the tensorview to reduce.
+// RFactor of reduction_tv will be returned if applicable otherwise reduction_tv
+// is returned
+TensorView* scheduleReductionTV(
+    const ReductionParams* rparams,
+    TensorView* reduction_tv,
+    bool has_iter_axis);
+
+// Inlining function intended for single or multi reduction fusions.
+void multiReductionInliner(
+    Fusion* fusion,
+    TensorView* reduction_tv,
+    TensorView* reference_tv,
+    const bool unroll,
+    const bool vectorize,
+    const bool use_grouped_reduction,
+    const int64_t vectorizatoin_factor,
+    std::vector<TensorView*> reduction_tvs,
+    std::vector<TensorView*> cached_inputs,
+    std::vector<std::pair<TensorView*, TensorView*>> cached_outputs,
+    std::vector<TensorView*> smem_persistent_buffer_consumers = {},
+    std::vector<TensorView*> dummy_outputs = {});
+
+// Propagate transformations with internal cutoff boundary at boundaryNodesSet
+// in P2C forward propagate, disable propagation to TensorView in
+// boundaryNodesSet in C2P backward propagate, disable propagation from
+// TensorView in boundaryNodesSet
+NVF_API void propagateTransformation(
+    TensorView* reference_tv,
+    const std::unordered_set<TensorView*>& boundaryNodesSet =
+        std::unordered_set<TensorView*>());
+
+// Propagate RFactor from first reduction TensorView to others
+void propagateRFactor(
+    TensorView* reference_tv,
+    TensorView* reduction_tv,
+    const std::vector<TensorView*>& reduction_tvs);
+
+// Get all cached input/output and shared memory TensorViews that are
+// vectorizable and unrollable.
+//
+// Parameters:
+//   reference_tv: TensorView created during RFactor, used to find vectorizable
+//                 TensorViews.
+//   is_vectorize: Indicates if vectorization is applied in the scheduler.
+//   cached_inputs: Inputs cached in registers or shared memory.
+//   cached_outputs: Outputs cached in registers.
+NVF_API std::unordered_set<TensorView*> getCachedTvsToUnrollOrVectorize(
+    TensorView* reference_tv,
+    bool is_vectorize,
+    const std::vector<TensorView*>& cached_inputs,
+    const std::vector<std::pair<TensorView*, TensorView*>>& cached_outputs);
+
+// Propagate parallelization from the reference TensorView to other TensorViews.
+// Unroll, Vectorize, and MisalignedVectorize types are explicitly handled for
+// TensorViews in unroll_vectorizable_cached_tvs. Clears unroll parallelization
+// for reduction_tv and reference_tv if they shouldn't be unrolled.
+//
+// Parameters:
+//   reduction_tv: The reduction TensorView being scheduled and parallelized.
+//                 Needs to clear its vectorization or convert to grouped
+//                 reduction.
+//
+//   reference_tv: The reference TensorView being scheduled and parallelized,
+//                 propagates parallelization to other selected TensorViews.
+//
+//   is_unroll_or_vectorization: Indicates if unroll or vectorization is used in
+//                               the scheduler.
+//
+//   reduction_tvs: All reduction TensorViews in the fusion. May add grouped
+//                  parallelization.
+//
+//   unroll_vectorizable_cached_tvs: Cached TensorViews that are unrollable
+//                                   or vectorizable.
+//
+//   selected_tvs: TensorViews selected for parallelization, default is all Tvs.
+NVF_API void propagateParallelization(
+    TensorView* reduction_tv,
+    TensorView* reference_tv,
+    const bool is_unroll_or_vectorization,
+    const bool use_grouped_reduction,
+    const std::vector<TensorView*>& reduction_tvs,
+    const std::unordered_set<TensorView*>& unroll_vectorizable_cached_tvs,
+    const std::vector<TensorView*>& selected_tvs = {});
+
+// Sort and rfactor the reference tv in a consistent way for reduction inliner.
+// Order of the sort is:
+//
+// [i-device dims, i-block dims, i-thread dims, i-constant sized, i-non-constant
+// sized
+//  r-block dims, r-thread dims, r-non-constant sized, r-constant sized,
+//  i/r-unswitched, i/r-unroll/vectorized, broadcasted dims]
+//
+// Rfactored axes are reductions bound to grid or blocks. If no axes are bound
+// to a grid or block dimension it will rfactor the r-unswitch dimension.
+// Reduction inliner expects an rfactored domain.
+NVF_API TensorView* sortAndRFactor(TensorView* reference_tv);
+
+// If project_to_inputs is true, take all projectable persistent buffers,
+// and move them to the inputs. Otherwise, try to project to their immediate
+// producers if these producers are persistent buffers.
+// This function create dummy outputs which should be used in later stages of
+// the scheduling.
+NVF_API std::vector<TensorView*> projectPersistentBuffers(
+    Fusion* fusion,
+    const bool project_to_inputs);
+
+//! Get reduction types based on the given fusion or reduction tvs.
+//! If there are no reduction tvs, return None.
+//! If there are only inner reduction tvs, return Inner.
+//! If there are only outer reduction tvs, return Outer.
+//! If there are both inner and outer reduction tvs, return InnerOuter.
+enum class ReductionType { Inner, Outer, InnerOuter, None };
+std::ostream& operator<<(std::ostream& os, ReductionType reduction_type);
+std::string toString(ReductionType reduction_type);
+ReductionType getReductionType(Fusion* fusion);
+ReductionType getReductionType(const std::vector<TensorView*>& reduction_tvs);
+
+/**
+ * @brief Vectorize shared memory consumers
+ *
+ * Applies vectorization to shared memory consumers.
+ * If extent of the last dim multiples vectorization factor exceeds hardware
+ * limitations, additional split is added.
+ *
+ * @param smem_consumers Vector of TensorView pointers representing shared
+ * memory consumers
+ * @param io_vectorization_factor Vectorization factor set for fusion inputs and
+ * outputs
+ * @note TODO: Optimize writing to shared memory and address bank conflicts for
+ * float32 with innermost extent of 8
+ */
+void sharedMemoryConsumerVectorization(
+    std::vector<TensorView*>& smem_consumers,
+    const int64_t io_vectorization_factor);
+
+} // namespace reduction_scheduler_utils
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/registry.h

@@ -0,0 +1,97 @@
+// 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 <expr_evaluator.h>
+#include <fusion.h>
+#include <scheduler/compile_time_info.h>
+#include <scheduler/utils.h>
+
+namespace nvfuser {
+
+class HeuristicDataCache;
+class HeuristicParams;
+class SchedulerRuntimeInfo;
+
+//! Virtual base class for schedule heuristics
+//!   heuristic implementations derive from this
+//!   class and implement a schedule(Fusion*)
+//!   and a bool canSchedule(Fusion*) interface
+class SchedulerEntry {
+ public:
+  NVF_API virtual ~SchedulerEntry() = default;
+
+  //! Fusion runtime facing API,
+  //!   schedule the given fusion with heuristics owned
+  //!   by this entry, for actual heuristics to override
+  NVF_API virtual void schedule(
+      Fusion* fusion,
+      const HeuristicParams* params) = 0;
+
+  virtual std::unique_ptr<HeuristicParams> computeHeuristics(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) = 0;
+
+  // Compile check that the scheduler maybe able to schedule the fusion
+  virtual bool canScheduleCompileTime(Fusion* fusion) = 0;
+
+  // Runtime check that the scheduler can take the fusion. Scheduler must be
+  // able to schedule the fusion if canScheduleCompileTime && this returns True.
+  virtual bool canScheduleRunTime(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) = 0;
+
+  // Dispatch heuristic type to the right derived class of scheduler entry.
+  // Scheduler entries are stateless so it's a lightweight class to dispatch to
+  // the virtual functions in this abstract class.
+  NVF_API static std::unique_ptr<SchedulerEntry> makeSchedulerInstance(
+      SchedulerType scheduler_type);
+
+  // Checks the provided scheduler type can schedule the fusion with the
+  // provided inputs. Schedules the fusion according to the heuristics provided
+  // by the scheduler. Returns the heuristics. This is simply a convenience
+  // function for a common testing pattern. If validate_scheduler is set to
+  // false canSchedule will not be checked.
+  NVF_API static std::unique_ptr<HeuristicParams> scheduleWith(
+      Fusion* fusion,
+      SchedulerType scheduler_type,
+      const at::ArrayRef<c10::IValue>& runtime_inputs,
+      bool validate_scheduler = true);
+
+  //! Heuristic comparison
+  NVF_API bool sameAs(const SchedulerEntry* other);
+
+  NVF_API const HeuristicParams* params() const {
+    return params_.get();
+  }
+
+  std::unique_ptr<HeuristicParams> params_ = nullptr;
+};
+
+namespace Schedule {
+
+//! External access for canSchedule utilities through SchedulerEntry
+//!  to avoid exposing a single function to the namespace
+bool canSchedule(
+    SchedulerType sh,
+    Fusion* fusion,
+    SchedulerRuntimeInfo& runtime_info,
+    HeuristicDataCache* data_cache = nullptr,
+    bool skip_compile_time_checks = false);
+
+//! Fusion segmenter facing API,
+//!   returns a schedule that applies in the given fusion, returns
+//!   SchedulerType::None if no schedule in the registry can handle.
+SchedulerType proposeHeuristics(
+    Fusion* fusion,
+    SchedulerRuntimeInfo& runtime_info);
+} // namespace Schedule
+
+} // namespace nvfuser