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

nvfuser/include/nvfuser/scheduler/pointwise.h

@@ -0,0 +1,183 @@
+// 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/pointwise_heuristic.h>
+#include <scheduler/registry.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+/*
+ * The 2D pointwise scheduling logic is a bit interesting. We'll start by giving
+ * motivation for what the scheduling is attempting to do. What we're going to
+ * do with the scheduling is attempt to make it two dimensional in a way that
+ * minimizes the refetching of broadcasted dimensions. If we think of the
+ * trivial case:
+ * T0[i0, b1]
+ * T1[b0, i1]
+ * T2[i0, i1] = T0 + T1
+ * If we scheduled T2 as 1-dimensional we would do something along the lines of
+ * merging i0 and i1 then splitting out a block and thread dimension. If i1 is
+ * greater than the thread dimension, then all threads would pull the same value
+ * from T0. However, they would all be pulling different values from T1. In this
+ * case we have perfect reuse of the broadcast dimension T0 but potentially no
+ * reuse of the broadcast dimension of T1. "Potentially" because if i1 isn't too
+ * big it should be efficiently cached in L2. If i1 is big, then by the time we
+ * increment the i0 dimension the i1 dimension will be pushed out of cache.
+ *
+ * Instead what we do is we map this to a two dimensional problem. Instead of
+ * having the schedule that merges the two dimensions, we'll actually leave the
+ * dimensions separate and we'll take i0, split it to BIDy, TIDy, and take i1
+ * and split it to BIDx and TIDx. Therefore we'll have a parallelization on T2
+ * like [BIDy, TIDy | BIDx, TIDx], where | denotes the separation of the
+ * original i0 and i1. This helps because all threads in the TIDx dimension will
+ * reuse the same value in the i0 dimension (holding BIDy and TIDy constant),
+ * all the threads in the TIDy dimension (holding BIDx, and TIDx constant) will
+ * reuse the same value in the i1 dimension. This reuse of values reduces the
+ * number of redundant values pulled from T0 and T1. The same thing can be said
+ * for when incrementing BIDy, but since BIDy is strided on BIDx there's no
+ * effective increment of BIDy without incrementing BIDx. Since all threads are
+ * executed within a block we can effectively consider the block incrementing
+ * TIDx BDIMx times while holding TIDy constant and incrementing TIDy BDIMy
+ * times while holding TIDx constant. Since multiple BIDx's are running at the
+ * same time on the device we can consider a wave on the GPU of incrementing
+ * BIDx (wave number of times), while holding TIDy constant BDIMy * wave number
+ * of times.
+ *
+ * If instead we have a situation like:
+ * T0[i0, i1, b2]
+ * T1[i0, b1, i2]
+ * T2[i0, i1, i2] = T0 + T1
+ * It makes sense that the break point would be in position 2, between i1 and
+ * i2. This is because when we map [i0, i1 | i2] to [BIDy, TIDy| BIDx, TIDx]
+ * BIDx, and TIDx will access the same elements of T0 on b2, and TIDy will
+ * likely access the same elements of T1 (as long as i1 > BDIMy). Even if i1 on
+ * the order of BDIMy we'll only access ~two unique elements per increment of
+ * BIDx or TIDx. This means we'll still reuse many of the same values and limit
+ * the amount we need to read duplicate values in T0 and T1.
+ *
+ * If instead we have:
+ * T0[i0, b1, i2]
+ * T1[b0, i1, i2]
+ * T2[i0, i1, i2] = T0 + T1
+ * The analysis gets a bit more complicated. First if i2 is very large and i0
+ * and i1 are relatively small it would make sense to have [i0, i1 | i2]. If b0
+ * is very small it's unlikely beneficial to have [i0 | i1, i2] as there would
+ * be small reuse on b0, and potentially no reuse on b1. If i2 is very small it
+ * may be worthwhile to have [i0 | i1, i2]. If i1 and i2 are not small, and
+ * their product is relatively large (i.e. you can't fit T2[i, :, :] in L2) then
+ * it's unlikely we'll get any significant reuse across i0.
+ *
+ * What we should (but don't due to complexity) assume then, is that we will get
+ * strong reuse across TIDx and TIDy for dimensions that are on the inner
+ * portion of the 2D tile.
+ *
+ * For example if we have:
+ * T0[i0, b1, i2]
+ * T1[b0, b1, i2]
+ * T2[b0, i1, i2]
+ * T3[i0, i1, i2] = T0 + T1 + T2
+ * We may want to break point at position 1 or position 2 (i.e. [i0 | i1, i2] or
+ * [i0, i1 | i2]). We can't immediately tell from the structure.
+ *
+ * If we choose [i0, i1 | i2] then we'll get:
+ * Strong reuse of T0 on TIDy (b1 dim)
+ * Perfect reuse across T1 on TIDy (b0 and b1)
+ * If BIDx is bound to the LHS of the tile we'll get:
+ * Maybe strong reuse of T0 on BIDx (b1 dim if it's large)
+ * Perfect reuse across T1 on BIDx
+ * Potentially no reuse on T2 if i1 is very large
+ *
+ * If we pick [i0 | i1, i2], then we'll get:
+ * We'll perfect reuse across TIDy on T1 and T2 on b0
+ * Some reuse on T0 and T1 on b1 across BIDx if i2 is relatively small and BIDx
+ * is bound to the RHS of the 2D schedule Perfect reuse on T1 and T2 on b0
+ * across BIDx if BIDx is bound to the LHS of the 2D schedule
+ *
+ * Materializing these benefits is dependent on the decisions the scheduler
+ * makes when parallelizing the problem. The heuristics logic at the moment is
+ * fairly simplistic where it assumes that there's only reuse across the break
+ * points for tensors that have no iteration domain on the entire side of the
+ * breakpoint. This is not optimal but for the time being it seems sufficient.
+ * We would ideally take into consideration the parallelization scheme and
+ * partial broadcasting on the lhs or rhs.
+ *
+ * An example of how this analysis is done is given the DAG:
+ * T0[i0, i1, b2] float
+ * T1[i0, b1, i2] half
+ * T2[i0, b1, i2] = cast(T1, float)
+ * T4[i0, i1, i2] float = T0 + T2
+ * With values of 10, 100, 1000 as [i0, i1, i2]
+ * Our break point analysis for positions 0, 1, 2, 3 will be:
+ *
+ * 0: 10*10 * 100*10 * 1000*10 = 1e9
+ * 1: 10*10 * 100*10 * 1000*10 = 1e9
+ * 2: 10*10 * 100*10 * 1000*6  = 6e8
+ * 3: 10*10 * 100*10 * 1000*10 = 1e9
+ *
+ * Where for each computation the LHS of the * pairs is the number of elements
+ * in that dimension on the reference and the RHS of the * pairs is the
+ * broadcast multiple where any tensor that has all broadcasts on the rhs or lhs
+ * of the break point doesn't contribute to the broadcast multiple of the rhs or
+ * lhs.
+ *
+ * So we'll pick position 2 since we're confident we can get broadcast reuse on
+ * the rhs of tensor 0. As already mentioned this is a pretty big
+ * simplification/assumption and in reality it may be harder/easier to take
+ * advantage of broadcast on the inner or outer dimension. This is a reasonable
+ * way to make relative decisions on break points, however, this computation is
+ * ont doing an effective estimate of actual DRAM transfers which it should be
+ * modified to do so.
+ *
+ * For view schedules there can be some incoherent break points for example:
+ * T1[i0, i1*i2] = view(T0[i0, i1, i2])
+ * would make the position 2 "incoherent". In otherwords we cannot replay
+ * through the view a schedule that tries to merge i0 and i1, without i2. So for
+ * positions that are incoherent we won't consider break point positions there.
+ *
+ * See FusionBroadcastViewMultiples_CUDA for what we expect with view handling.
+ * Shortly any dimensions that are inputs or outputs of view transformations are
+ * considered together, since it's hard to account for partial dimensions that
+ * are being broadcasted. So for view it's primarily an all or nothing situation
+ * when it comes to the 2D pointwise scheduler.
+ *
+ * DID axes, which are not allocated, are ignored in the analysis.
+ * Specifically, two fusions that only differ by DID axes result in
+ * the same scheduling decisions.
+ */
+
+class SchedulerRuntimeInfo;
+class HeuristicDataCache;
+
+class PointWiseScheduler : 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::PointWise;
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/pointwise_heuristic.h

@@ -0,0 +1,118 @@
+// 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 {
+
+// Parameters of the pointwise heuristic to describe the optimial schedule.
+// Warning: equal operator is intended for use in caching the kernel associated
+// with these pointwise parameters. It does not check if the launch parameters
+// are equivelent!
+class PointwiseParams : public HeuristicParams {
+ public:
+  PointwiseParams() : HeuristicParams(SchedulerType::PointWise) {};
+
+  // Treat pointwise operation as 2-Dimensional, this is the location where we
+  // split from left side of the domain to right. i.e. 0 means problem is
+  // treated as 1-D, 1 of 3 would mean we treat the first dimension as the outer
+  // dimension, and all the others as an inner dimension.
+  int64_t break_point = 0;
+
+  // Split block across left and right dimension
+  bool split_block = false;
+
+  // Split grid y dimension, if otherwise it would be too large
+  bool split_grid_y_dim = false;
+
+  // For many instances having BIDx on the inner most dimension is the most
+  // performant parallel binding. However, if we're broadcasting the outer
+  // dimension with a large inner dimension, it can be more performant to bind
+  // BIDy on the inner most dimension.
+  bool flip_grid_binding = false;
+
+  // vectorization factor
+  int64_t vectorization_factor = 1;
+
+  // Unroll on top of vectorization
+  // In the 2D scheduler, unroll the outer dimension to reuse loaded data across
+  // rows, reducing loaded bytes by the unroll factor.
+  // Always equals 1 for 1D scheduler.
+  int64_t unroll_factor_outer = 1;
+
+  // In the 2D scheduler, unroll the inner dimension to reuse loaded data across
+  // cols, reducing loaded bytes by the unroll factor.
+  // Also used in 1D scheduler.
+  int64_t unroll_factor_inner = 1;
+
+  using HeuristicParams::HeuristicParams;
+
+  // Warning: Does not check launch parameters!
+  bool sameAs(const HeuristicParams* other_base) const override {
+    auto other = dynamic_cast<const PointwiseParams*>(other_base);
+    if (other == nullptr) {
+      return false;
+    }
+    bool attr_equal = other->cparams == cparams &&
+        other->vectorization_factor == vectorization_factor &&
+        other->break_point == break_point &&
+        other->split_block == split_block &&
+        other->split_grid_y_dim == split_grid_y_dim &&
+        other->unroll_factor_outer == unroll_factor_outer &&
+        other->unroll_factor_inner == unroll_factor_inner &&
+        other->flip_grid_binding == flip_grid_binding;
+    return attr_equal;
+  }
+
+  std::string toString() const override {
+    std::stringstream ss;
+    ss << "\n===== Pointwise Parameters ========\n"
+       << (tag.empty() ? "" : "Tag: ") << tag << " Pointwise Characteristics:\n"
+       << " Gridx: " << lparams.gdimx() << " BlckY: " << lparams.bdimy()
+       << " BlckX: " << lparams.bdimx() << "\n";
+    if (break_point) {
+      ss << "2D Schedule\n"
+         << "  Bcast break point: " << break_point << "\n";
+      if (split_block) {
+        ss << "Split block into y-dim\n";
+      }
+      if (split_grid_y_dim) {
+        ss << "  Split y grid dim\n";
+      }
+    }
+    ss << "vectorization_factor: " << vectorization_factor << "\n";
+    ss << "unroll_factor_outer: " << unroll_factor_outer << "\n";
+    ss << "unroll_factor_inner: " << unroll_factor_inner << "\n";
+    if (flip_grid_binding) {
+      ss << "Flip BIDx/BIDy bindings\n";
+    }
+    ss << "====================================\n";
+    return ss.str();
+  }
+
+  // Warning: Hash is not based on launch parameters!
+  size_t hash() const override {
+    size_t attr_hash = static_cast<size_t>(vectorization_factor) ^
+        static_cast<size_t>(break_point) << 4 ^
+        static_cast<size_t>(split_block) << 5 ^
+        static_cast<size_t>(split_grid_y_dim) << 6 ^
+        static_cast<size_t>(unroll_factor_outer) << 7 ^
+        static_cast<size_t>(unroll_factor_inner) << 9 ^
+        static_cast<size_t>(flip_grid_binding) << 10;
+    return attr_hash;
+  }
+
+  std::unique_ptr<HeuristicParams> clone() const override {
+    return std::make_unique<PointwiseParams>(*this);
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/pointwise_utils.h

@@ -0,0 +1,24 @@
+// 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 <compute_at_map.h>
+#include <exceptions.h>
+#include <ir/all_nodes.h>
+#include <ir/utils.h>
+#include <scheduler/tools/domain_map.h>
+#include <scheduler/utils.h>
+
+namespace nvfuser {
+namespace pointwise_utils {
+
+// Return reference tensor view.
+TensorView* getReferenceTensor(Fusion* fusion);
+
+} // namespace pointwise_utils
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/reduction.h

@@ -0,0 +1,43 @@
+// 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/reduction_heuristic.h>
+#include <scheduler/registry.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+class SchedulerRuntimeInfo;
+class HeuristicDataCache;
+
+class ReductionScheduler : 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::Reduction;
+  }
+};
+
+} // namespace nvfuser