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

nvfuser/include/nvfuser/scheduler/registry_utils.h

@@ -0,0 +1,111 @@
+// 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/all_schedulers.h>
+
+namespace nvfuser {
+
+class TensorView;
+class ComputeAtLogicalDomainMap;
+class ComputeAtMap;
+class ExpressionEvaluator;
+class KernelArgumentHolder;
+
+namespace registry_utils {
+
+bool checkPatternEquivalence(
+    TensorView* out_tv0,
+    TensorView* out_tv1,
+    const ComputeAtLogicalDomainMap& logical_map);
+
+// Reusing some code from lowering specifically in lower_trivial_broadcast.cpp
+// ConcretizedBroadcastDomains::maybeNonUniquelyConcretized this checks if
+// there's a broadcast iteration domain that's being broadcasted to seemingly
+// different extents, meaning we don't know in the kernel if the dimension is
+// being broadcasted to one size multiple times or different sizes. This is a
+// hard to optimize problem and likely indicates we shouldn't be fusing.
+bool hasNonUniqueBcast(Fusion* fusion);
+
+// TODO: remove this requirement entirely
+bool rejectScheduleForMemoryPromotion(
+    Fusion* fusion,
+    SchedulerType scheduler_type);
+
+bool isConnectedFusionGraph(Fusion* fusion);
+
+// Returns if a fusion cannot transformed into a consistent format since we
+// can't transform forward through view operations, for exmaple:
+//
+// tv0[I0, I1, I2]
+// tv1[I0*I1, I2] = view(tv0)
+// tv2[I0, I1*I2] = view(tv0)
+//
+// If we start transform propagation at either tv1 or tv2, it would require
+// "replaying forward" through the other. If we started at tv1 we'd have to be
+// able to take tv2[I0, I1*I2] and transform it to [I0*I1, I2], however this
+// would "undo" the view transformation which we do not support today.
+//
+// Returns true if a scenario like above is found in the fusion.
+bool requiresForwardViewReplay(Fusion* fusion, ComputeAtMap& ca_map);
+
+// Returns if view interferes with how we want to treat the reference, being
+// at least a 2D reduction schedule but maybe a 3D reduction schedule.
+bool reductionInterferingView(
+    Fusion* fusion,
+    const ComputeAtMap& ca_map,
+    TensorView* reduction_reference);
+
+// Check inputs, outputs and intermediates
+// Intermediates are contiguous, so strides are not necessary
+// Strides are required for inputs and also maybe for outputs as
+// they may be non-contiguous. However, in our current interface,
+// output strides are not available, so if there's any outputs that
+// are non contiguous, need to fall back to 64-bit indexing
+PrimDataType getIndexTypeOfKernel(
+    Fusion* fusion,
+    const std::vector<TensorView*>& all_tvs,
+    const KernelArgumentHolder& inputs,
+    ExpressionEvaluator& ee);
+
+class SchedulerTopologyChecker {
+ public:
+  // Checks if any broadcasts are resolved after a reduction that don't follow
+  // the normalization pattern
+  static bool hasNonNormalizePostReductionBCast(Fusion* fusion);
+
+  // Checks if any broadcasts are resolved after a reduction, this shouldn't
+  // be accepted in the single reduction or multi-reduction scheduler
+  static bool hasPostReductionBCast(Fusion* fusion);
+
+  // Checks if there's any unsupported operations post reduction. If outer
+  // reduction we can fuse some pointwise ops if they don't require
+  // broadcasting (checked in hasPostReductionBCast). For inner reductions we
+  // cannot fuse any binary like operation (includes operations like shift
+  // that we're not fusing right now) involving "new" inputs (not going
+  // through a reduction).
+  static bool supportedPostReductionFusion(
+      Fusion* fusion,
+      std::vector<TensorView*> reduction_tvs);
+
+  // Checks if there's any gather-like ops that result in non-resolved
+  // broadcast domains and then get squeezed before reaching reduction
+  // TVs. The reduction scheduler uses reduction TVs as a scheduling
+  // reference, so that won't be able to schedule the broadcast ID if
+  // squeezed and its corresponding index-accessed producer ID, and
+  // any IDs that the producer ID depends on.
+  //
+  // This analysis has some similarity as DomainMap. Can be
+  // consolidated?
+  static bool hasGatherToBroadcastBeforeReduction(
+      Fusion* fusion,
+      const std::vector<TensorView*>& reduction_tvs);
+};
+
+} // namespace registry_utils
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/resize.h

@@ -0,0 +1,41 @@
+// 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;
+
+class ResizeScheduler : public SchedulerEntry {
+ public:
+  bool canScheduleCompileTime(Fusion* fusion) override;
+  bool canScheduleRunTime(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) 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::Resize;
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/resize_heuristic.h

@@ -0,0 +1,67 @@
+// 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 <c10/util/hash.h>
+#include <ir/interface_nodes.h>
+#include <scheduler/heuristic.h>
+#include <utils.h>
+
+#include <sstream>
+
+namespace nvfuser {
+
+class ResizeParams : public HeuristicParams {
+ public:
+  ResizeParams() : HeuristicParams(SchedulerType::Resize) {};
+
+  // Split grid x dimension
+  bool split_grid_x_dim = false;
+
+  int64_t largest_input = -1;
+
+  int64_t vectorization_factor = 1;
+
+  static constexpr int64_t max_gdimx = (1L << 31) - 1L;
+
+  using HeuristicParams::HeuristicParams;
+
+  // Warning: Does not check launch parameters!
+  bool sameAs(const HeuristicParams* other_base) const override {
+    auto other = dynamic_cast<const ResizeParams*>(other_base);
+    if (other == nullptr) {
+      return false;
+    }
+    bool attr_equal = other->cparams == cparams &&
+        other->split_grid_x_dim == split_grid_x_dim &&
+        other->largest_input == largest_input &&
+        other->vectorization_factor == vectorization_factor;
+    return attr_equal;
+  }
+
+  std::string toString() const override {
+    std::stringstream ss;
+    ss << "\n===== Resize Parameters ========\n"
+       << (tag.empty() ? "" : "Tag: ") << tag << " Resize Characteristics:\n"
+       << " split grid x dim: " << split_grid_x_dim << "\n"
+       << " index of largest input: " << largest_input << "\n"
+       << " vectorization factor: " << vectorization_factor << "\n";
+    ss << "====================================\n";
+    return ss.str();
+  }
+
+  size_t hash() const override {
+    return c10::get_hash(split_grid_x_dim);
+  }
+
+  std::unique_ptr<HeuristicParams> clone() const override {
+    return std::make_unique<ResizeParams>(*this);
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/runtime_info.h

@@ -0,0 +1,166 @@
+// 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 <cstddef>
+#include <cstdint>
+
+#include <expr_evaluator.h>
+#include <fusion.h>
+#include <runtime/executor_kernel_arg.h>
+#include <utils.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+class ExpressionEvaluator;
+
+//!  SchedulerRuntimeInfo is the abstraction introduced in
+//! this PR for passing runtime input dependent information
+//! to the schedulers and kernel caches.
+//!
+//! Note:
+//!  if any additional info needed,  or maybe just the inputs themselves it
+//!    could just be added to this class, and they will be distributed to the
+//!    segmenter and schedulers.
+//!  It is important that input id encoding should be up to date with any change
+//!   of this class to avoid launching compiled kernels with illegal inputs.
+
+class SchedulerRuntimeInfo : public NonCopyable {
+ public:
+  // Max vector size we will consider, in bytes,
+  //  currently set to 16B = 128b
+  static constexpr int64_t max_alignment_size_in_byte = 16;
+
+  //! Create runtime info for given fusion and input. Creating and binding
+  //! evaluator is optional. The evaluator is used to manage intermediate
+  //! integers in the fusion. We need them for segmenter and schedulers,
+  //! but we don't need them when we are just using this class to provide
+  //! additional encoding for kernel cache lookup.
+  //!
+  //! The index type of forced_index_type is used if given, no matter
+  //! how large the actual arguments and fusion tensors
+  //! are. CORRECTNESS IS NOT GUARANTEED.
+  SchedulerRuntimeInfo(
+      Fusion* complete_fusion,
+      KernelArgumentHolder args,
+      PrecomputedValues* precomputed_values = nullptr,
+      const std::vector<TensorView*>& all_tvs = {},
+      std::optional<PrimDataType> forced_index_type = std::nullopt);
+
+  NVF_API SchedulerRuntimeInfo(
+      Fusion* complete_fusion,
+      const at::ArrayRef<c10::IValue>& aten_inputs);
+
+  //! Lookup for the alignment sizes of the given tv. Currently only returns
+  //!  actual alignment info for input tensors to the complete fusion,
+  //!  and for other intermediate/fuser-allocated tensors will
+  //!  return max_alignment_size_in_byte.
+  size_t getAlignmentSize(TensorView* tv);
+
+  //! Returns sizes of tensor dimensions in same order as allocation domain,
+  //! ignoring any IterType::Reduction domains in the allocation domain. This
+  //! only works for complete Fusion inputs whose allocation domain is a
+  //! permutation of their root domain and will raise an exception otherwise.
+  const std::vector<int64_t>& getInputAllocationSizes(TensorView* tv) const {
+    NVF_ERROR(
+        isInputTv(tv),
+        "TensorView ",
+        tv->toString(),
+        " is not an input or its logical domain is not a permutation of its ",
+        "allocation domain");
+    auto sizes_it = input_sizes_.find(tv);
+    NVF_ERROR(sizes_it != input_sizes_.end());
+    return sizes_it->second;
+  }
+
+  //! Returns strides of tensor in same order as allocation domain, in elements
+  //! instead of bytes. Only works for complete Fusion inputs whose allocation
+  //! domain is a permutation of their root domain and will raise an exception
+  //! otherwise.
+  const std::vector<int64_t>& getInputAllocationStrides(TensorView* tv) const {
+    NVF_ERROR(
+        isInputTv(tv),
+        "TensorView ",
+        tv->toString(),
+        " is not an input or its logical domain is not a permutation of its ",
+        "allocation domain");
+    auto strides_it = input_strides_elements_.find(tv);
+    NVF_ERROR(strides_it != input_strides_elements_.end());
+    return strides_it->second;
+  }
+
+  // Computes alignment size in bytes for provided ptr address
+  static size_t computeAlignmentSize(size_t ptr_address);
+
+  // Return the runtime pointer value for provided tensor view
+  size_t ptrOf(TensorView* tv) const;
+
+  PrimDataType getIndexType() const {
+    return index_type_;
+  }
+
+  Fusion* fusion() {
+    return complete_fusion_;
+  }
+
+  ExpressionEvaluator& expressionEvaluator() {
+    NVF_ERROR(expression_evaluator_ != nullptr);
+    return *expression_evaluator_;
+  }
+
+ private:
+  // Build and bind full fusion inputs to an expression evaluator
+  std::unique_ptr<ExpressionEvaluator> getExpressionEvaluator(
+      const KernelArgumentHolder& inputs,
+      PrecomputedValues* precomputed_values);
+
+  bool isInputTv(TensorView* tv) const {
+    return std::find(
+               complete_fusion_->inputs().begin(),
+               complete_fusion_->inputs().end(),
+               tv) != complete_fusion_->inputs().end();
+  }
+
+ private:
+  // Returns the offset of tv in the inputs ignoring non tensor views. Used to
+  // access input_sizes, input_strides, input_ptr
+  int offsetTensorPos(TensorView* tv);
+
+  // Expression evaluator used to probe sizes in the fusion IR
+  std::unique_ptr<ExpressionEvaluator> expression_evaluator_ = nullptr;
+
+  // Fusion reference that this runtime info is associated with
+  Fusion* complete_fusion_ = nullptr;
+
+  // Copy of aten input pointer addresses
+  // TODO: Support output tensor pointers
+  std::unordered_map<Val*, size_t> input_ptrs_;
+
+  // Copy of aten input tensor sizes ordered like the TensorView's allocation
+  // domain
+  std::unordered_map<Val*, std::vector<int64_t>> input_sizes_;
+
+  // Copy of aten input tensor strides (in elements) ordered like the
+  // TensorView's allocation domain
+  std::unordered_map<Val*, std::vector<int64_t>> input_strides_elements_;
+
+  // Copy of aten input tensor strides (in bytes) for only discontiguous
+  // dimensions
+  std::unordered_map<Val*, std::vector<size_t>> input_discontig_strides_;
+
+  // Cache for getAlignmentSize
+  std::unordered_map<TensorView*, size_t> alignment_map_;
+
+  // Found index mode kernel needs to be run in
+  PrimDataType index_type_ = PrimDataType::Int;
+
+  // TODO: Remove
+  std::unordered_map<TensorView*, size_t> vectorword_map_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/scheduler_types.h

@@ -0,0 +1,80 @@
+// 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 <visibility.h>
+#include <array>
+#include <ostream>
+#include <string>
+
+namespace nvfuser {
+
+//! Each SchedulerType maps to a scheduler in distinct CPP files.
+//! For instance, SchedulerType::PointWise maps to PointWiseScheduler in
+//! pointwise.cpp.
+//!
+//!    Each of the scheduler needs to provide 3 interface functions:
+//!
+//!      1. canScheduleCompileTime(Fusion* fusion) :
+//!
+//!        This function contains compiled-time checks on the graph itself
+//!        without runtime input information. Only `fusion` is given in the
+//!        argument to make sure only compile-time available info is needed in
+//!        the check.
+//!
+//!        This function is to be called exactly once on each segmented group
+//!        created in a segmented fusion so this part will not contribute to
+//!        dynamic shape latency.
+//!
+//!     2. canScheduleRunTime(
+//!            Fusion* fusion,
+//!            SchedulerRuntimeInfo& runtime_info,
+//!           HeuristicDataCache* data_cache = nullptr):
+//!        This function contains all canSchedule checks that will have to
+//!        involve runtime input information, and will be run both by the
+//!        segmenter and the kernel cache. The latency of this function will
+//!        contribute to dynamic shape latency so `data_cache` should be used as
+//!        much as possible to save re-computation.
+//!
+//!     3. schedule(fusion):
+//!
+//!        This function will be called when compiling a kernel. It should apply
+//!        scheduling to the given fusion
+
+enum class SchedulerType {
+  None,
+  NoOp,
+  PointWise,
+  Matmul,
+  Reduction,
+  InnerPersistent,
+  InnerOuterPersistent,
+  OuterPersistent,
+  Transpose,
+  ExprEval,
+  Resize
+};
+
+//! Define a schedule table to loop over all the heuristics in priority order.
+constexpr std::array<SchedulerType, 10> all_heuristics_in_priority_order = {
+    SchedulerType::ExprEval,
+    SchedulerType::NoOp,
+    SchedulerType::Matmul,
+    SchedulerType::Reduction,
+    SchedulerType::Resize,
+    SchedulerType::Transpose,
+    SchedulerType::PointWise,
+    SchedulerType::InnerPersistent,
+    SchedulerType::OuterPersistent,
+    SchedulerType::InnerOuterPersistent};
+
+std::string toString(SchedulerType sh);
+
+NVF_API std::ostream& operator<<(std::ostream& os, SchedulerType sh);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/transpose.h

@@ -0,0 +1,114 @@
+// 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 <ATen/core/ivalue.h>
+#include <exceptions.h>
+#include <fusion.h>
+#include <scheduler/registry.h>
+#include <scheduler/transpose_heuristic.h>
+#include <visibility.h>
+
+#define SUPPORT_SPLITTING_INNERMOST_DIM 0
+
+namespace nvfuser {
+
+// Note [Transpose scheduling]
+//
+// The target of transpose scheduling is to get coalesced global memory access
+// to as much input and output tensors as possible. For a DAG with only pure
+// pointwise operators, the scheduling is very simple because the inner most
+// dimension of all input and output tensors are all mapped together in the
+// ComputeAtMap, i.e., there is essentially only one inner most dimension. In
+// such case, we just vectorize that inner most dimension and bind it to
+// threadIdx.x identically for all input and output tensors. In the case where
+// transposes are present in the DAG, the inner most dimensions of different
+// inputs and outputs might not match. And there is no fixed pattern on which
+// input/output tensors should share the same inner most dimension with which.
+// Consider the following example DAGs ([T] represents transpose, all tensors
+// are 2D):
+//
+//   t0    t1      t0    t1      t0    t1        t0    t1         t0
+//    \    |        \    /        \    |          \    |          |
+//     \  [T]       [T] [T]        \  [T]          t2 [T]        [T]
+//      \ /           \ /           \ / \         / \ / \         |
+//      t2             t2           t2   t3      t3  t4 t5       [T]
+//                                                                |
+//                                                                t1
+//
+// In order to support all these cases in a general way, the following
+// perspective is very important: What we are looking for is to bind threadIdx.x
+// differently for different inputs and outputs, so there has to be some tensor
+// somewhere in the DAG that we write and read with different threadIdx.x
+// bindings. The tensor of binding swap can be any tensor on the path that
+// connects inputs/outputs with different inner most dimension, especially, it
+// does not necessarily have to be the tensor of the transpose operator. In
+// other words, thanks to our indexing system who is already taking care of the
+// correctness of transpose, the scheduler can freely choose where to realize
+// these transposes as different threadIdx.x bindings. This observation greatly
+// simplifies our scheduling.
+//
+// Our scheduling strategy is as follows: We first split the inputs and outputs
+// of the fusion into two groups according to their inner most dimension. The
+// inner most dimensions of tensors in the same group are mapped to each other,
+// and they are not mapped to the inner most dimesion of tensors in a different
+// group. Depending on the transpose pattern, there can be more than two groups,
+// if this is the case, we only consider the two largest groups, and the tensors
+// in the remaining groups will just be accessed unvectorized and uncoalesced.
+// We call the largest group as `group1` and the second largest group as
+// `group2`. When we have the groups, we will make a 2D tiling [I1, I2] ->
+// [I1/tile1, tile1, I2/tile2, tile2] on the inner most dimensions of group1 and
+// group2. If I1 and I2 are too small to make a 32x32 tile, such as in the
+// fusion of tanspose(T1[1024, 2, 1024, 2], {1, 3}), we merge in other
+// dimensions to make a virtual I1 and I2. The details of how we create virtual
+// I1 and I2 are described in note [Supporting small transpose dimensions].
+//
+// Each tile [tile1, tile2] will be handled by a block, and the tensors that
+// have mismatched threadIdx.x bindings will use shared memory. The outer IDs of
+// the tiling split will be merged with non-tiled IDs and then binded to
+// blockIdx.x for the entire DAG, regardless of which group a tensor belongs to.
+// For the inner tile IDs [tile1, tile2], we need to transform and parallelize
+// group 1 and group 2 differently. The intermediate tensors can be transformed
+// and parallelized consistently either with group 1 or group 2. Here, since
+// group 1 is larger than group 2, we decide to only transform and parallelize
+// the cached inputs of group 2 together with group 2, and keep the rest of the
+// DAG consistent with group 1.
+//
+// If you would like to see an example of how to manually schedule a complicated
+// DAG using this idea, refer to:
+//   FusionManualScheduleTransposeComplexDAG1_CUDA
+
+class SchedulerRuntimeInfo;
+class HeuristicDataCache;
+
+//! Utility for canSchedule interface to check if this fusion has at least two
+//! groups, each with a fully broadcasted reference tensor.
+NVF_API bool hasAtLeastTwoValidGroups(Fusion* fusion);
+
+class TransposeScheduler : public SchedulerEntry {
+ public:
+  bool canScheduleCompileTime(Fusion* fusion) override;
+
+  bool canScheduleRunTime(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) override;
+
+  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::Transpose;
+  }
+};
+
+} // namespace nvfuser