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

nvfuser/include/nvfuser/transform_replay.h

@@ -0,0 +1,297 @@
+// 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 <ir/internal_nodes.h>
+#include <scheduler/tools/maxinfo_propagator.h>
+#include <visibility.h>
+
+#include <algorithm>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace nvfuser {
+
+/*
+ * compute_at is a relative property between two TensorViews which marks at what
+ * iteration domain we're going to generate a tensor to be consumed by another.
+ * For example if we have: T2[I, J, K] = T1[I, J, K] * 2.0 and then we call
+ * T2.split(axis = 0, factor = ...): T2[Io, Ii, J, K] = T1[I, J, K] * 2.0 where
+ * Io is the outer axes from the split, and Ii is the inner axes from the split.
+ * then we call T1.compute_at(T2, axis=1) we would expect to have:
+ * T2[Io, Ii, J, K] = T1[Io, Ii, J, K] * 2.0
+ * which would produce the following loop nest structure:
+ *
+ * for(io : Io)
+ *  for(ii : Ii)
+ *   for(j : J)
+ *    for(k : K)
+ *     //produce T1:
+ *     T1[io, ii, j, k] = ...
+ *  for(ii : Ii)
+ *   for(j : J)
+ *    for(k : K)
+ *     //consume T1, produce T2
+ *     T2[io, ii, j, k] = T1[io, ii, j, k] * 2.0
+ *
+ * This file provides the replay function that allows us to construct T1's
+ * domain from T2 at a desired level (compute_at_axis) without modifying any
+ * unnecessary parts of the domain.
+ *
+ * EXAMPLES:
+ *
+ * ANOTHER ITER EXAMPLE:
+ *   T2[I, J, K] = T1[I, J, K] * 2.0
+ * T2.split(axis = 0, factor = ...)
+ *   T2[Io, Ii, J, K] = T1[I, J, K] * 2.0
+ * T2.split(axis = 2, factor = ...)
+ *   T2[Io, Ii, Jo, Ji, K] = T1[I, J, K] * 2.0
+ * T1.compute_at(T2, axis=1)
+ *   T2[Io, Ii, Jo, Ji, K] = T1[Io, Ii, J, K] * 2.0
+ *
+ * Note: compute_at axis:
+ * T2[ 0 Io, 1 Ii, 2 Jo, 3 Ji, 4 K 5 ] //5 is inline, 0 is at "root" which means
+ * completely separate loop nests.
+ *
+ * for(io : Io)
+ *  for(ii : Ii)
+ *   for(j : J)
+ *    for(k : K)
+ *     //produce T1, this is the view that replay generates:
+ *     T1[io, ii, j, k] = ...
+ *  for(ii : Ii)
+ *   for(jo : Jo)
+ *     for(ji : Ji)
+ *      for(k : K)
+ *       //consume T1, produce T2
+ *       T2[io, ii, jo, ji, k] = T1[io, ii, jo, ji, k] * 2.0
+ *       //consumer view on T1 will be produced at a later stage.
+ *
+ *
+ * SIMPLE REDUCTION EXAMPLE:
+ *   T1[I, J, K] = ...
+ *   T2[I, R, K] = T1[I, J, K] //.sum(axis = 1), we reduce on R/J to produce
+ * T2[I, K] T2.split(axis = 0, factor = ...) T2[Io, Ii, R, K] = T1[I, J, K]
+ * T1.compute_at(T2, axis=3)
+ *   T2[Io, Ii, R, K] = T1[Io, Ii, J, K]
+ *
+ * for(io : Io)
+ *  for(ii : Ii)
+ *   for(k : K)
+ *    T2[io, ii, k] = init
+ *   for(r : R)
+ *    for(k : K)
+ *     //produce T1:
+ *     T1[io, ii, r, k] = ...
+ *     //consume T1 produce T2:
+ *     T2[io, ii, k] += T1[io, ii, r, k]
+ *
+ *
+ * REDUCTION EXAMPLE RESULTING IN AN ERROR:
+ *   T1[I, R, K] = ... //R is reduction domain, we reduce on R to produce T1[I,
+ * K] T2[I, K] = T1[I, K]
+ *
+ * for(i : I)
+ *   for(k : K)
+ *     T1[i, k] = init
+ *   for(r : R)
+ *     for(k : K)
+ *       T1[i, k] += ...[i, r, k]
+ * for(i : I)
+ *   for(k : K)
+ *     T2[i, k] = T1[i, k]
+ *
+ * T1.compute_at(T2, axis=2)
+ * This should be an error, or a warning and changed to:
+ * T1.compute_at(T2, axis=1)
+ * The error is because the kernel would have to be:
+ *
+ * for(i : I)
+ *   T1[i, k] = init
+ *   for(r : R)
+ *     for(k : K)
+ *       T1[i, k] += ...[i, r, k]
+ *   for(k : K)
+ *     T2[i, k] = T1[i, k]
+ *
+ * Otherwise we would produce incorrect results.
+ *
+ */
+
+class TensorDomain;
+class TensorView;
+class LogicalDomainMap;
+
+struct TransformReplayOptions {
+  // In theory, it makes more sense to have skip_target_swizzle = true by
+  // default because this is how we index into the producer and how we propagate
+  // transformations. However, we are in a very funny situation that:
+  // BestEffortReplay for swizzle is broken. For example, if we have a
+  // producer <=> consumer pair like:
+  //       I1             I0
+  //      /  \           /  |
+  //     I1o I1i        I0o I0i
+  //      |   |          |   |
+  // swizzle I1i    swizzle I0i     <=>     I3  I2
+  //      |   |          |   |
+  //    I1o' I1i       I0o' I0i
+  //      \  /           \  /
+  //       I1'            I0'
+  // where I1o', I0o' = swizzle(I1o, I0o), we never really skipped swizzle to
+  // map I1' with I3 and I0' with I2. But even with this error, our swizzle
+  // indexing worked due to luck. So effectively we were doing
+  // skip_target_swizzle = false. But today, we can not make this `true` for
+  // vectorization validation and indexing, because of another bug in
+  // BestEffortReplay: swizzle skip should happen in an all-or-nothing fashion.
+  // We can not just skip X but not skip Y, but we are not implementing this
+  // skip like that. If we make it `true`, this will trigger some error in some
+  // schedule. So here, in order to avoid exposing one bug, we are more
+  // explicitly using a wrong behavior that we have been using because this
+  // wrong behavior has a better luck.
+  // For more info, see https://github.com/NVIDIA/Fuser/issues/554
+  bool skip_target_swizzle = false;
+  bool replay_swizzle = false;
+  bool replay_resize = false;
+  bool replay_allocation = false;
+
+  TransformReplayOptions& skipTargetSwizzle(bool value = true) {
+    skip_target_swizzle = value;
+    return *this;
+  }
+
+  TransformReplayOptions& replaySwizzle(bool value = true) {
+    replay_swizzle = value;
+    return *this;
+  }
+
+  TransformReplayOptions& replayResize(bool value = true) {
+    replay_resize = value;
+    return *this;
+  }
+
+  TransformReplayOptions& replayAllocation(bool value = true) {
+    replay_allocation = value;
+    return *this;
+  }
+};
+
+class NVF_API TransformReplay {
+ public:
+  // Replay producer as consumer, returns {producer, producer_compute_at_axis}.
+  //
+  // replay_resize indicates whether resize should be replayed or
+  // ignored. It is only replayed when replaying a producer for
+  // indexing.
+  // replay_allocation indicates whether to replace the producer's allocation
+  // domain with corresponding consumer's allocation domain. By default, we
+  // should preserve producer's current allocation domain, and if that
+  // allocation domain is inconsistent with the replay, an error will be raised.
+  // This option is used in cacheBefore, cacheAfter, and cacheFork
+  static std::pair<TensorDomain*, int64_t> replayPasC(
+      const TensorView* producer,
+      const TensorView* consumer,
+      int64_t consumer_compute_at_axis,
+      TransformReplayOptions opt = {});
+  static std::pair<TensorDomain*, int64_t> replayPasC(
+      const TensorView* producer,
+      const TensorView* consumer,
+      int64_t consumer_compute_at_axis,
+      const LogicalDomainMap& logical_map,
+      TransformReplayOptions opt = {});
+
+  // Replay producer as consumer, returns {replayed_consumer_domain,
+  // consumer_compute_at_axis}.
+  //
+  // Unlike replayPasC, it always ignores resize.
+  static std::pair<TensorDomain*, int64_t> replayCasP(
+      const TensorView* consumer,
+      const TensorView* producer,
+      int64_t producer_compute_at_axis,
+      TransformReplayOptions opt = {});
+  static std::pair<TensorDomain*, int64_t> replayCasP(
+      const TensorView* consumer,
+      const TensorView* producer,
+      int64_t producer_compute_at_axis,
+      const LogicalDomainMap& logical_map,
+      TransformReplayOptions opt = {});
+
+  // Self replay.
+  static TensorDomain* fullSelfReplay(
+      const TensorDomain* new_self_root,
+      const TensorDomain* self);
+
+  // Returns the loop position in producer that matches with `consumer_pos` in
+  // consumer. Returns -1 if matching is impossible. This function can be used
+  // to test if replay is needed for getting matching outer dims. This function
+  // should be consistent with `replayPasC`: if you pass the tensors just
+  // replayed by replayPasC as inputs, you should return exactly the same
+  // position as `replayPasC`. However, this function is more tolerant than
+  // fully matching `replayPasC`: if in the consumer, there are unmappable
+  // dimensions, these dimensions are just ignored.
+  //
+  // When skip_resize is true, mapping is done more permissively by
+  // skipping resize ops. For example, that is done when this is used
+  // by TransformPropagator, whereas it isn't when used for
+  // determining the inlining position by MaxPosCalculator as inlining
+  // isn't allowed with different extents.
+  static int64_t getMatchedLeafPosWithoutReplayPasC(
+      const TensorView* producer,
+      const TensorView* consumer,
+      int64_t consumer_pos,
+      bool skip_resize = false);
+
+  // Returns the loop position in consumer that matches with `producer_pos` in
+  // producer. Behavior similar to getMatchedLeafPosWithoutReplayPasC, except
+  // that we are also ignoring reductions in the producer.
+  //
+  // When skip_resize is true, mapping is done more permissively by
+  // skipping resize ops. For example, that is done when this is used
+  // by TransformPropagator, whereas it isn't when used for
+  // determining the inlining position by MaxPosCalculator as inlining
+  // isn't allowed with different extents.
+  static int64_t getMatchedLeafPosWithoutReplayCasP(
+      const TensorView* consumer,
+      const TensorView* producer,
+      int64_t producer_pos,
+      bool skip_resize = false);
+
+  // tests if two tensors has fully matching transformations
+  static bool fullSelfMatching(
+      const TensorView* replay,
+      const TensorView* target);
+};
+
+class NVF_API TransformPropagator
+    : public MaxLogicalDomainInfoSpanningTree::Propagator {
+ protected:
+  std::unordered_map<TensorView*, int64_t> replayed_pos_;
+
+ public:
+  void propagateC2P(TensorView* from, TensorView* to) override;
+  void propagateP2C(TensorView* from, TensorView* to) override;
+  void propagateSibling(TensorView* from, TensorView* to) override;
+  TransformPropagator(TensorView* from, int64_t pos = -1);
+};
+
+struct MostInlinedTransformPropagator
+    : public MaxLogicalDomainInfoSpanningTree::Propagator {
+  void propagateC2P(TensorView* from, TensorView* to) override;
+  void propagateP2C(TensorView* from, TensorView* to) override;
+  void propagateSibling(TensorView* from, TensorView* to) override;
+};
+
+// Replays an `Expr` with the new input, `new_in`. This function currently has
+// the following limitations:
+// 1. It requires `e` to be a unary op, and therefore takes a single new input.
+// 2. It requires `e` to be a TensorView op, which takes and produces only
+// TensorViews.
+Expr* replayExprWithNewInput(Expr* e, Val* new_in);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/transform_rfactor.h

@@ -0,0 +1,33 @@
+// 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 <ir/all_nodes.h>
+#include <transform_iter.h>
+
+#include <algorithm>
+#include <vector>
+
+namespace nvfuser {
+
+// TODO: Only replay dispatch is really borrowed from TransformIter, we should
+// reevaluate the reuse of dispatch for classes that inherit TransformIter.
+class TransformRFactor {
+ public:
+  // Transform the provided tensor domain to two domains, a producer and
+  // consumer domain. These domains are created by taking axes and reducing them
+  // in the producer domain, and taking the remaining reduction axes and
+  // reducing them in the consumer domain.
+  static std::pair<TensorDomain*, TensorDomain*> runReplay(
+      TensorDomain*,
+      std::vector<int64_t> axes);
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/transform_view.h

@@ -0,0 +1,136 @@
+// 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 <visibility.h>
+
+#include <ir/all_nodes.h>
+
+#include <memory>
+#include <vector>
+
+namespace nvfuser {
+
+class ViewTransform;
+
+//!
+//! The goal of analyzeView is to find the minimum number of transformations
+//! to convert from the original size to the new size. A naive view algorithm
+//! would merge all axis together and then split according to the new sizes.
+//!
+//! This implementation will keep the original domains, if the domains are the
+//! same size in the original and new shapes. If an original domain is not
+//! evenly divisible by the new domain, we will merge the minimum number of
+//! adjacent original domains.
+//!
+//! The view transformations are processed in the following order:
+//! 1. Squeeze - Removes size-1 broadcast dimensions
+//! 2. Keep, Merge, Split - Used to create new logical domain
+//! 3. Broadcast - Inserts size-1 dimensions
+//!
+//! Broadcast is handled last because size-1 dimension can be inserted anywhere
+//! in the new shape.
+//!
+
+struct AnalyzeViewResult {
+  std::vector<bool> broadcast_axes;
+  std::vector<bool> squeeze_axes;
+  std::vector<std::shared_ptr<ViewTransform>> transforms;
+
+  std::string toString() const;
+
+  bool operator==(const AnalyzeViewResult& other) const;
+
+  bool operator!=(const AnalyzeViewResult& other) const {
+    return !(*this == other);
+  }
+
+  size_t hash() const;
+};
+
+struct AnalyzeViewConstraint {
+  // 1 if size 1 dimension, otherwise 0;
+  std::vector<int64_t> original_constraint;
+  std::vector<int64_t> new_constraint;
+  // Just the positions of true in AnalyzeViewResult::squeeze_axes
+  std::vector<int64_t> squeeze_string;
+  // Just the positions of true in AnalyzeViewResult:broadcast_axes
+  std::vector<int64_t> broadcast_string;
+  // A stringified version of the transformations:
+  std::vector<int64_t> split_merge_string;
+
+  std::vector<int64_t> conglomerateString() const {
+    // Don't think this is necessary but just being safe. Using
+    // -3 as a dilimeter between value groups.
+    std::vector<int64_t> conglomerate = {
+        (int64_t)original_constraint.size(),
+        (int64_t)new_constraint.size(),
+        -3};
+    auto add_vec = [&conglomerate](const std::vector<int64_t>& vec) {
+      conglomerate.insert(conglomerate.end(), vec.begin(), vec.end());
+      conglomerate.push_back(-3);
+    };
+    add_vec(original_constraint);
+    add_vec(new_constraint);
+    add_vec(squeeze_string);
+    add_vec(broadcast_string);
+    add_vec(split_merge_string);
+    return conglomerate;
+  }
+
+  bool operator==(const AnalyzeViewConstraint& other) const {
+    return other.conglomerateString() == this->conglomerateString();
+  }
+
+  // Naive hashing function, likely has a lot of collisions, but may not matter
+  // too much if we don't expact many types of views.
+  size_t hash() const {
+    size_t hash_value = 0;
+    for (auto val : conglomerateString()) {
+      if (val == std::numeric_limits<int64_t>::max()) {
+        continue;
+      }
+      hash_value += val;
+    }
+    return hash_value;
+  }
+};
+
+//! Infer -1 value in new view std::vector<int64_t> based on original view
+//! std::vector<int64_t>. This shouldn't generally be used directly but is
+//! useful for testing.
+NVF_API std::pair<std::vector<int64_t>, std::vector<int64_t>> inferViewShapes(
+    const std::vector<int64_t>& original_sizes,
+    const std::vector<int64_t>& new_sizes);
+
+// Find the transformations necessary to convert TensorView
+// from original size to new size.
+AnalyzeViewResult analyzeView(
+    const TensorView* tv,
+    const std::vector<int64_t>& original_sizes,
+    const std::vector<int64_t>& new_sizes);
+
+// Find the constraints derived from the view transformations
+NVF_API AnalyzeViewConstraint analyzeViewConstraint(
+    const std::vector<int64_t>& original_sizes,
+    const std::vector<int64_t>& new_sizes);
+
+// Generate a new TensorDomain from the given view transformations.
+// The original root domain is kept in the new TensorDomain,
+// but a new logical domain is created from the view transformations.
+TensorDomain* transformView(
+    TensorDomain* original_domain,
+    const AnalyzeViewResult& view_analysis);
+
+//! Apply the reshape transformations of view_analysis to inp_tv
+NVF_API TensorView* reshape(
+    TensorView* inp_tv,
+    const AnalyzeViewResult& view_analysis);
+
+} // namespace nvfuser