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

nvfuser/include/nvfuser/scheduler/transpose_heuristic.h

@@ -0,0 +1,164 @@
+// 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 <scheduler/heuristic.h>
+#include <utils.h>
+
+#include <sstream>
+
+namespace nvfuser {
+
+// Parameters of the transpose heuristic to describe the optimial 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 TransposeParams : public HeuristicParams {
+ public:
+  TransposeParams() : HeuristicParams(SchedulerType::Transpose) {};
+  static constexpr int64_t getMaxThreadsPerBlock() {
+    return 128;
+  }
+
+  static constexpr int64_t getDefaultTileSize() {
+    return 32;
+  }
+
+  // See note [Supporting small transpose dimensions], all dims are positions in
+  // reference1
+  std::vector<std::pair<int64_t, int64_t>> split_before_tiling = {};
+  std::vector<int64_t> dims_merged_with_1 = {};
+  std::vector<int64_t> dims_merged_with_2 = {};
+
+  // Vectorization factor for tensors in the first group
+  int64_t vectorize_factor1 = 1;
+
+  // Vectorization factor for tensors in the second group
+  int64_t vectorize_factor2 = 1;
+
+  // TODO: support symbolic tile size
+  // https://github.com/csarofeen/pytorch/pull/1854#discussion_r928143729
+
+  // Tile size for the inner most dim of tensors in the first group
+  int64_t tile_size1 = getDefaultTileSize();
+
+  // Tile size for the inner most dim of tensors in the second group
+  int64_t tile_size2 = getDefaultTileSize();
+
+  using HeuristicParams::HeuristicParams;
+
+  // Warning: Does not check launch parameters!
+  bool sameAs(const HeuristicParams* other_base) const override {
+    auto other = dynamic_cast<const TransposeParams*>(other_base);
+    if (other == nullptr) {
+      return false;
+    }
+    bool attr_equal = other->cparams == cparams &&
+        other->split_before_tiling == split_before_tiling &&
+        other->dims_merged_with_1 == dims_merged_with_1 &&
+        other->dims_merged_with_2 == dims_merged_with_2 &&
+        other->vectorize_factor1 == vectorize_factor1 &&
+        other->vectorize_factor2 == vectorize_factor2 &&
+        other->tile_size1 == tile_size1 && other->tile_size2 == tile_size2;
+    return attr_equal;
+  }
+
+  std::string toString() const override {
+    std::stringstream ss;
+    ss << "\n===== Transpose Parameters ========\n"
+       << (tag.empty() ? "" : "Tag: ") << tag << " Transpose Characteristics:\n"
+       << " BlckX: " << lparams.bdimx() << "\n";
+    ss << " input tile size: " << tile_size1 << "\n";
+    ss << " output tile size: " << tile_size2 << "\n";
+    int64_t elements_per_tile = tile_size1 * tile_size2;
+    ss << " elements per tile: " << elements_per_tile << "\n";
+    int64_t elements_per_thread = elements_per_tile / lparams.bdimx();
+    ss << " elements per thread: " << elements_per_thread << "\n";
+    if (vectorize_factor1 > 1) {
+      ss << "Vectorize group 1, Factor: " << vectorize_factor1 << "\n";
+    }
+    int64_t unroll_factor1 = elements_per_thread / vectorize_factor1;
+    if (unroll_factor1 > 1) {
+      ss << "Unroll group 1, Factor: " << unroll_factor1 << "\n";
+    }
+    if (vectorize_factor2 > 1) {
+      ss << "Vectorize group 2, Factor: " << vectorize_factor2 << "\n";
+    }
+    int64_t unroll_factor2 = elements_per_thread / vectorize_factor2;
+    if (unroll_factor2 > 1) {
+      ss << "Unroll group 2, Factor: " << unroll_factor2 << "\n";
+    }
+    if (!split_before_tiling.empty() || !dims_merged_with_1.empty() ||
+        !dims_merged_with_2.empty()) {
+      ss << "Virtual inner-most dim:\n";
+      if (!split_before_tiling.empty()) {
+        ss << "  ";
+        bool first = true;
+        for (auto pair : split_before_tiling) {
+          if (!first) {
+            ss << ", ";
+          }
+          first = false;
+          ss << "split(" << pair.first << ", " << pair.second << ")";
+        }
+        ss << "\n";
+      }
+      if (!dims_merged_with_1.empty()) {
+        ss << "  merge ";
+        bool first = true;
+        for (auto dim : dims_merged_with_1) {
+          if (!first) {
+            ss << ", ";
+          }
+          first = false;
+          ss << dim;
+        }
+        ss << " with innermost1\n";
+      }
+      if (!dims_merged_with_2.empty()) {
+        ss << "  merge ";
+        bool first = true;
+        for (auto dim : dims_merged_with_2) {
+          if (!first) {
+            ss << ", ";
+          }
+          first = false;
+          ss << dim;
+        }
+        ss << " with innermost2\n";
+      }
+    }
+    ss << "====================================\n";
+    return ss.str();
+  }
+
+  size_t hash() const override {
+    return c10::get_hash(
+        split_before_tiling,
+        dims_merged_with_1,
+        dims_merged_with_2,
+        vectorize_factor1,
+        vectorize_factor2,
+        tile_size1,
+        tile_size2);
+  }
+
+  std::unique_ptr<HeuristicParams> clone() const override {
+    return std::make_unique<TransposeParams>(*this);
+  }
+
+  int64_t getThreadsPerBlock() const {
+    int64_t tile_vectors1 = ceilDiv(tile_size1 * tile_size2, vectorize_factor1);
+    int64_t tile_vectors2 = ceilDiv(tile_size1 * tile_size2, vectorize_factor2);
+    int64_t tile_vectors = std::min(tile_vectors1, tile_vectors2);
+    return std::min(getMaxThreadsPerBlock(), tile_vectors);
+  }
+};
+
+} // namespace nvfuser