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

nvfuser/include/nvfuser/mma_type.h

@@ -0,0 +1,257 @@
+// 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 <macros.h>
+
+#include <exceptions.h>
+#include <type.h>
+#include <visibility.h>
+
+#include <cstring>
+#include <ostream>
+
+#include <cstdint>
+
+namespace nvfuser {
+
+constexpr std::string_view MATMUL_LOG_PREFIX = "[MATMUL DEBUG] ";
+
+//! Named descriptors of domains in matmul
+enum class MatmulDimRole { M = 0, N, K, Batch };
+
+std::string toString(MatmulDimRole role);
+
+//! Named descriptors of TensorView roles in fusion
+//!  OPERAND_A - an input to the fusion that is a producer of a matmul "A" input
+//!  OPERAND_B - an input to the fusion that is a producer of a matmul "B" input
+//!  OUTPUT - fusion outputs that have the matmul as a dependency
+//!  EPILOGUE_INPUT - an input to the fusion that is a producer of an
+//!    OUTPUT, but not of an MMA input
+//!
+//!  Note: bias vector tensors will be assigned to the EPILOGUE_INPUT role.
+enum class MatmulTensorRole {
+  OPERAND_A = 0,
+  OPERAND_B,
+  OUTPUT,
+  EPILOGUE_INPUT
+};
+
+//! The expected number of occurances of core TensorView roles in fusion
+static constexpr size_t MATMUL_CORE_ROLES_EXPECTED_COUNT = 1;
+
+//! Utility data structure for recording gemm tiles
+struct GemmTile {
+  int64_t m, n, k;
+  GemmTile(int64_t m_, int64_t n_, int64_t k_) : m(m_), n(n_), k(k_) {}
+
+  bool operator==(const GemmTile& other) const {
+    return m == other.m && n == other.n && k == other.k;
+  }
+
+  GemmTile operator/(const GemmTile& other) const {
+    return GemmTile(m / other.m, n / other.n, k / other.k);
+  }
+
+  std::vector<int64_t> toVector() const {
+    return {m, n, k};
+  }
+};
+
+//! Utility data structure for recording gemm tiles
+struct MatMulTileOptions {
+  GemmTile cta_tile = GemmTile(128, 128, 32);
+  GemmTile warp_tile = GemmTile(64, 64, 32);
+
+  MatMulTileOptions() = default;
+  MatMulTileOptions(GemmTile cta_tile_, GemmTile warp_tile_)
+      : cta_tile(cta_tile_), warp_tile(warp_tile_) {}
+
+  bool operator==(const MatMulTileOptions& other) const {
+    return cta_tile == other.cta_tile && warp_tile == other.warp_tile;
+  }
+};
+
+enum class MmaMacro : uint64_t;
+
+struct MmaMacroEncode {
+  enum class Arch : uint16_t { NoMma, Volta, Turing, Ampere, Hopper } arch;
+  uint16_t m;
+  uint16_t n;
+  uint16_t k;
+
+  constexpr operator uint64_t() {
+    return (uint64_t)arch << 48 | (uint64_t)m << 32 | (uint64_t)n << 16 |
+        (uint64_t)k;
+  }
+
+  constexpr operator MmaMacro() {
+    return static_cast<MmaMacro>(static_cast<uint64_t>(*this));
+  }
+
+  constexpr MmaMacroEncode(MmaMacro macro)
+      : arch(Arch(toUnderlying(macro) >> 48)),
+        m((toUnderlying(macro) >> 32) & 0xFFFF),
+        n((toUnderlying(macro) >> 16) & 0xFFFF),
+        k(toUnderlying(macro) & 0xFFFF) {}
+
+  constexpr MmaMacroEncode(Arch arch_, uint16_t m_, uint16_t n_, uint16_t k_)
+      : arch(arch_), m(m_), n(n_), k(k_) {}
+};
+
+static_assert(sizeof(MmaMacroEncode) == sizeof(uint64_t));
+
+//! Type of mma instrinsic macro to use
+//!  This will translate to which mma intrinsic from runtime string
+//!    to be generated to implement the mma op. The current plan
+//!    is to have exactly one macro for each
+//!  (arch, datatype, operand layout) triple, though there
+//!  exists multiple possibilities for some cases, e.g. for Turing and fp16
+//!  one can use 16_8_8 or 16_8_16.
+//! Will consider adding more choices that the scheduler can pick from
+//!  when our perf target becomes more fine grained, which is more likely in
+//!  latency bound kernels.
+
+#define MACRO(arch, m, n, k) \
+  arch##_##m##_##n##_##k = MmaMacroEncode(MmaMacroEncode::Arch::arch, m, n, k)
+
+enum class MmaMacro : uint64_t {
+  NoMMA = 0,
+
+  MACRO(Turing, 16, 8, 8),
+  MACRO(Turing, 16, 8, 16),
+  MACRO(Turing, 16, 16, 16),
+
+  MACRO(Ampere, 16, 8, 16),
+  MACRO(Ampere, 16, 16, 16),
+
+  MACRO(Hopper, 64, 8, 16),
+  MACRO(Hopper, 64, 16, 16),
+  MACRO(Hopper, 64, 24, 16),
+  MACRO(Hopper, 64, 32, 16),
+  MACRO(Hopper, 64, 40, 16),
+  MACRO(Hopper, 64, 48, 16),
+  MACRO(Hopper, 64, 56, 16),
+  MACRO(Hopper, 64, 64, 16),
+  MACRO(Hopper, 64, 72, 16),
+  MACRO(Hopper, 64, 80, 16),
+  MACRO(Hopper, 64, 88, 16),
+  MACRO(Hopper, 64, 96, 16),
+  MACRO(Hopper, 64, 104, 16),
+  MACRO(Hopper, 64, 112, 16),
+  MACRO(Hopper, 64, 120, 16),
+  MACRO(Hopper, 64, 128, 16),
+  MACRO(Hopper, 64, 136, 16),
+  MACRO(Hopper, 64, 144, 16),
+  MACRO(Hopper, 64, 152, 16),
+  MACRO(Hopper, 64, 160, 16),
+  MACRO(Hopper, 64, 168, 16),
+  MACRO(Hopper, 64, 176, 16),
+  MACRO(Hopper, 64, 184, 16),
+  MACRO(Hopper, 64, 192, 16),
+  MACRO(Hopper, 64, 200, 16),
+  MACRO(Hopper, 64, 208, 16),
+  MACRO(Hopper, 64, 216, 16),
+  MACRO(Hopper, 64, 224, 16),
+  MACRO(Hopper, 64, 232, 16),
+  MACRO(Hopper, 64, 240, 16),
+  MACRO(Hopper, 64, 248, 16),
+  MACRO(Hopper, 64, 256, 16),
+};
+
+#undef MACRO
+
+//! [Operand Layout Convention]
+//! Operand layout, T=transposed/row_major, N=normal/col_major
+//! Ordered by position of K
+//! NT : K,M x K,N -> M,N
+//! TT : M,K X K,N -> M,N
+//! TN : M,K X N,K -> M,N
+//! NN : K,M X N,K -> M,N
+enum class MmaLayout { NT = 0, TT, TN, NN };
+
+//! Indicates which dimension is innermost in the allocation domain of an
+//! operand
+enum class UnitDim { K, M_or_N };
+
+//! Utility to annotate which input of mma this option struct describes
+enum class MmaOperand { A, B };
+
+//! GPU arch check for macro type
+inline bool isTuring(MmaMacro macro) {
+  return MmaMacroEncode(macro).arch == MmaMacroEncode::Arch::Turing;
+}
+
+inline bool isAmpere(MmaMacro macro) {
+  return MmaMacroEncode(macro).arch == MmaMacroEncode::Arch::Ampere;
+}
+
+inline bool isHopper(MmaMacro macro) {
+  return MmaMacroEncode(macro).arch == MmaMacroEncode::Arch::Hopper;
+}
+
+//! Get the m size from macro type
+inline int64_t getM(MmaMacro macro) {
+  return MmaMacroEncode(macro).m;
+}
+
+//! Get the n size from macro type
+inline int64_t getN(MmaMacro macro) {
+  return MmaMacroEncode(macro).n;
+}
+
+//! Get the k size from macro type
+inline int64_t getK(MmaMacro macro) {
+  return MmaMacroEncode(macro).k;
+}
+
+// Unpacked constants from macro type:
+//   exact numbers are defined by each individual instruction.
+int getOutputRegisterSize(MmaMacro macro);
+int getInputARegisterSize(MmaMacro macro);
+int getInputBRegisterSize(MmaMacro macro);
+
+// Unpack MMA op shape
+GemmTile getMmaOpShape(MmaMacro macro);
+
+// Warning: The values of the enum class must match the matrix descriptor as
+// specified in:
+// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor
+// Do not edit the values of the enum class unless you know what you are doing.
+enum class MmaInputSmemSwizzle {
+  None = 0,
+  B128 = 1,
+  B64 = 2,
+  B32 = 3,
+};
+
+constexpr int64_t core_matrix_width_bytes = 16;
+
+int64_t getBytesFromSwizzle(MmaInputSmemSwizzle swizzle);
+MmaInputSmemSwizzle getSwizzleFromBytes(int64_t bytes);
+
+// MMA stringify utils
+NVF_API std::string toString(MmaLayout input_layout);
+std::string toString(const GemmTile& tile);
+NVF_API std::string toString(const MatMulTileOptions& opts);
+NVF_API std::string toString(MmaMacro macro);
+NVF_API std::string toString(MmaInputSmemSwizzle swizzle);
+inline std::ostream& operator<<(
+    std::ostream& os,
+    MmaInputSmemSwizzle input_layout) {
+  os << toString(input_layout);
+  return os;
+}
+
+// MMA hash utils
+NVF_API size_t hash(MmaMacro macro);
+size_t hash(MmaLayout input_layout);
+size_t hash(const GemmTile& tile);
+NVF_API size_t hash(const MatMulTileOptions& opts);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/multidevice/c10d_mock.h

@@ -0,0 +1,175 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <ATen/core/TensorBody.h>
+#include <ATen/core/ivalue.h>
+#include <c10/util/intrusive_ptr.h>
+
+namespace c10d {
+
+inline void setDebugLevelFromEnvironment() {}
+
+class Work : public torch::CustomClassHolder {
+ public:
+  void wait() {}
+};
+
+struct ReduceOp : torch::CustomClassHolder {
+  enum RedOpType {
+    SUM,
+    AVG,
+    PRODUCT,
+    MIN,
+    MAX,
+    BAND,
+    BOR,
+    BXOR,
+    UNUSED,
+  };
+
+  ReduceOp() = default;
+  ReduceOp(RedOpType op) : op_(op) {}
+
+  RedOpType op_ = UNUSED;
+};
+
+struct ReduceScatterOptions {
+  ReduceOp reduceOp = ReduceOp::UNUSED;
+};
+
+struct ScatterOptions {
+  int64_t rootRank = 0;
+};
+
+struct AllgatherOptions {};
+
+struct GatherOptions {
+  int64_t rootRank = 0;
+};
+
+struct BroadcastOptions {
+  int64_t rootRank = 0;
+};
+
+struct AllreduceOptions {
+  ReduceOp reduceOp = ReduceOp::UNUSED;
+};
+
+struct ReduceOptions {
+  ReduceOp reduceOp = ReduceOp::UNUSED;
+  int64_t rootRank = 0;
+};
+
+struct BarrierOptions {
+  std::vector<int64_t> device_ids;
+};
+
+class Backend : public torch::CustomClassHolder {
+ public:
+  void startCoalescing() {}
+
+  c10::intrusive_ptr<Work> endCoalescing() {
+    return c10::make_intrusive<Work>();
+  }
+
+  const std::string getBackendName() const {
+    return "";
+  };
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& tensors,
+      const BroadcastOptions& opts = BroadcastOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) {
+    return c10::make_intrusive<Work>();
+  }
+
+  int getSize() const {
+    return 0;
+  }
+};
+
+struct TCPStoreOptions {
+  static constexpr uint16_t kDefaultPort = 0;
+};
+
+class TCPStore : public torch::CustomClassHolder {};
+
+} // namespace c10d

nvfuser/include/nvfuser/multidevice/communication.h

@@ -0,0 +1,232 @@
+// 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 <ir/base_nodes.h>
+#include <ir/builder.h>
+#include <ir/interface_nodes.h>
+#include <multidevice/communicator.h>
+#include <multidevice/device_mesh.h>
+#include <multidevice/multidevice.h>
+#ifdef NVFUSER_DISTRIBUTED
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#else
+#include <multidevice/c10d_mock.h>
+#endif
+#include <type.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+enum class CommunicationType {
+  Gather,
+  Allgather,
+  Scatter,
+  Reduce,
+  Allreduce,
+  ReduceScatter,
+  Broadcast,
+  SendRecv
+};
+
+std::ostream& operator<<(std::ostream& os, const CommunicationType& type);
+
+using RedOpType = c10d::ReduceOp::RedOpType;
+
+// The class "Communication" represents a MPI-style communication
+// communication operation to be executed on the network. The base class
+// Communication should not be used directly but through its derived classes:
+// Broadcast, Gather, Scatter, Allgather, and SendRecv. Other collectives will
+// be added later.
+class Communication : public Expr {
+ public:
+  using Expr::Expr;
+  // Only specify `root` for types that have root.
+  // Only specify `red_op` for reduction types.
+  // Only specify `scattered_axis` for ReduceScatter.
+  Communication(
+      IrBuilderPasskey passkey,
+      CommunicationType type,
+      TensorView* out,
+      TensorView* in,
+      Team team, // All devices involved in this communication. It must include
+                 // `root`. It can be a subset of `root`+`mesh` in case of 2D
+                 // sharding.
+      DeviceIdxType root = -1,
+      RedOpType red_op = RedOpType::UNUSED,
+      int64_t scattered_axis = -1);
+
+  Communication(const Communication& other) = delete;
+  Communication& operator=(const Communication& other) = delete;
+  Communication(Communication&& other) = delete;
+  Communication& operator=(Communication&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "Communication";
+  }
+
+  CommunicationType type() const {
+    return attribute<CommunicationType>(0);
+  }
+
+  TensorView* out() const {
+    return output(0)->as<TensorView>();
+  }
+
+  TensorView* in() const {
+    return input(0)->as<TensorView>();
+  }
+
+  const Team& team() const {
+    return attribute<Team>(1);
+  }
+
+  // A convenience helper so the user doesn't need to convert size_t to int64_t.
+  int64_t team_size() const {
+    return static_cast<int64_t>(team().size());
+  }
+
+  DeviceIdxType root() const {
+    return attribute<DeviceIdxType>(2);
+  }
+
+  RedOpType reduceOp() const {
+    return attribute<RedOpType>(3);
+  }
+
+  int64_t scatteredAxis() const {
+    return attribute<int64_t>(4);
+  }
+
+  // PyTorch's process group expects the root to be specified
+  // as an integer between 0 and world_size-1. We choose it to be
+  // the device's relative index within the team
+  int64_t getRootRelativeIndex();
+
+ private:
+  void validate();
+};
+
+enum class P2PCommunicationType { SEND, RECV };
+
+std::ostream& operator<<(std::ostream& os, const P2PCommunicationType& type);
+
+class P2PCommunication : public Expr {
+ public:
+  using Expr::Expr;
+
+  P2PCommunication(
+      IrBuilderPasskey passkey,
+      P2PCommunicationType type,
+      TensorView* buffer,
+      Val* peer);
+
+  P2PCommunication(const P2PCommunication& other) = delete;
+  P2PCommunication& operator=(const P2PCommunication& other) = delete;
+  P2PCommunication(P2PCommunication&& other) = delete;
+  P2PCommunication& operator=(P2PCommunication&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "P2PCommunication";
+  }
+
+  P2PCommunicationType type() const {
+    return attribute<P2PCommunicationType>(0);
+  }
+
+  TensorView* buffer() const {
+    return input(0)->as<TensorView>();
+  }
+
+  Val* peer() const {
+    return attributeVal(1);
+  }
+};
+
+// The method "post" triggers the execution of the communication. This call is
+// non-blocking. The communication can be posted multiple times.
+// It is assumed that the current device_index (given by
+// communicator.deviceId()) belongs to the team of the communication,
+// otherwise an error is thrown.
+//
+// NOTE: pytorch's NCCL process group API needs <team_size> buffers on root for
+// scatter/gather operation.
+// (*) Broadcast
+// Copies the root's src buffer to each device's dst buffer
+// Requirements:
+//   - the root is set and belongs to the team
+//   - the root has one src buffer, and no or one dst buffer
+//   - non-roots have no src buffer and one dst buffer
+//   - all buffers have the same size
+// (*) Gather
+// Copies each device's source buffer to the root's respective src
+// buffer. The order of the sender devices matches the order of the
+// root's buffers.
+// Requirements:
+//   - the root is set and belongs to the team
+//   - the root has one src buffer and <team_size> dst buffers
+//   - non-roots have one src buffer and no dst buffer
+//   - all buffers have the same size
+// (*) Allgather
+// Copies each device's src buffer to each device's respective src
+// buffer. The order of the devices matches the order of the
+// buffers
+// Requirements:
+//   - all device have one src buffer and <team_size> dst buffers
+//   - all buffers have the same size
+// (*) Scatter
+// Copies each root's src buffer to each device's dst buffer.
+// The order of the buffers matches the order of the receiver devices
+// Requirements:
+//   - the root is set and belongs to the team
+//   - the root has <team_size> src buffers and one dst buffer
+//   - non-roots have no src buffer and one dst buffer
+//   - all buffers have the same size
+// (*) Reduce
+// Reduce the src buffers to the root's dst buffer.
+// Requirements:
+//   - the root is set and belongs to the team
+//   - the root has one src buffers and one dst buffer
+//   - non-roots have one src buffer and no dst buffer
+//   - all buffers have the same size
+// (*) Allreduce
+// Reduce the src buffers to the dst buffer.
+// Requirements:
+//   - all devices have one src buffer and one dst buffer
+//   - all buffers have the same size
+// (*) ReduceScatter
+// Reduce all the src buffers and shard the result to the dst buffers.
+// Requirements:
+//   - all devices have <team_size> src buffer and one dst buffer
+//   - all buffers have the same size
+// (*) SendRecv
+// Copies the sender's src buffers to the receiver's dst buffer
+// It is equivalent to a Broadcast with a team of size == 2
+c10::intrusive_ptr<c10d::Work> postSingleCommunication(
+    Communication* communication,
+    DeviceIdxType my_device_index,
+    c10d::Backend* backend,
+    at::Tensor input_tensor,
+    at::Tensor output_tensor);
+
+c10::intrusive_ptr<c10d::Work> postSingleCommunication(
+    P2PCommunication* communication,
+    DeviceIdxType my_device_index,
+    DeviceIdxType peer,
+    c10d::Backend* backend,
+    at::Tensor buffer);
+
+} // namespace nvfuser