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

nvfuser/include/nvfuser/python_frontend/translation_utils.h

@@ -0,0 +1,308 @@
+// 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/all_nodes.h>
+#include <ops/all_ops.h>
+
+namespace nvfuser::python_frontend {
+
+// Get std::function for UnaryOp
+template <typename ResultType, typename... ArgTypes>
+std::function<ResultType(ArgTypes...)> getFunction(const UnaryOp* uop) {
+  auto wrap_function = [](ResultType (*fn)(ArgTypes...)) { return fn; };
+
+  switch (uop->getUnaryOpType()) {
+    case UnaryOpType::Abs:
+      return wrap_function(abs);
+    case UnaryOpType::Acos:
+      return wrap_function(acos);
+    case UnaryOpType::Acosh:
+      return wrap_function(acosh);
+    case UnaryOpType::Asin:
+      return wrap_function(asin);
+    case UnaryOpType::Asinh:
+      return wrap_function(asinh);
+    case UnaryOpType::Atan:
+      return wrap_function(atan);
+    case UnaryOpType::Atanh:
+      return wrap_function(atanh);
+    case UnaryOpType::Ceil:
+      return wrap_function(ceil);
+    case UnaryOpType::Cos:
+      return wrap_function(cos);
+    case UnaryOpType::Cosh:
+      return wrap_function(cosh);
+    case UnaryOpType::Exp:
+      return wrap_function(exp);
+    case UnaryOpType::Exp2:
+      return wrap_function(exp2);
+    case UnaryOpType::Expm1:
+      return wrap_function(expm1);
+    case UnaryOpType::Erf:
+      return wrap_function(erf);
+    case UnaryOpType::Erfc:
+      return wrap_function(erfc);
+    case UnaryOpType::Erfinv:
+      return wrap_function(erfinv);
+    case UnaryOpType::Erfcinv:
+      return wrap_function(erfcinv);
+    case UnaryOpType::Floor:
+      return wrap_function(floor);
+    case UnaryOpType::Frac:
+      return wrap_function(frac);
+    case UnaryOpType::Lgamma:
+      return wrap_function(lgamma);
+    case UnaryOpType::Log:
+      return wrap_function(log);
+    case UnaryOpType::Log10:
+      return wrap_function(log10);
+    case UnaryOpType::Log1p:
+      return wrap_function(log1p);
+    case UnaryOpType::Log2:
+      return wrap_function(log2);
+    case UnaryOpType::Neg:
+      return wrap_function(neg);
+    case UnaryOpType::LogicalNot:
+      return wrap_function(logical_not);
+    case UnaryOpType::BitwiseNot:
+      return wrap_function(bitwise_not);
+    case UnaryOpType::Reciprocal:
+      return wrap_function(reciprocal);
+    case UnaryOpType::Relu:
+      return wrap_function(relu);
+    case UnaryOpType::Rsqrt:
+      return wrap_function(rsqrt);
+    case UnaryOpType::Round:
+      return wrap_function(round);
+    case UnaryOpType::Sigmoid:
+      return wrap_function(sigmoid);
+    case UnaryOpType::Signbit:
+      return wrap_function(signbit);
+    case UnaryOpType::Silu:
+      return wrap_function(silu);
+    case UnaryOpType::Sin:
+      return wrap_function(sin);
+    case UnaryOpType::Sinh:
+      return wrap_function(sinh);
+    case UnaryOpType::Sqrt:
+      return wrap_function(sqrt);
+    case UnaryOpType::Tan:
+      return wrap_function(tan);
+    case UnaryOpType::Tanh:
+      return wrap_function(tanh);
+    case UnaryOpType::Trunc:
+      return wrap_function(trunc);
+    case UnaryOpType::IsFinite:
+      return wrap_function(isfinite);
+    case UnaryOpType::IsInf:
+      return wrap_function(isinf);
+    case UnaryOpType::IsNan:
+      return wrap_function(isnan);
+    case UnaryOpType::IsNegInf:
+      return wrap_function(isneginf);
+    case UnaryOpType::IsPosInf:
+      return wrap_function(isposinf);
+    case UnaryOpType::IsReal:
+      return wrap_function(isreal);
+    case UnaryOpType::Real:
+      return wrap_function(real);
+    case UnaryOpType::Imag:
+      return wrap_function(imag);
+    default:
+      NVF_CHECK(
+          false,
+          "Unexpected operator type: ",
+          uop->getUnaryOpType(),
+          " in ",
+          uop->toString());
+  }
+}
+
+// Get std::function for BinaryOp
+template <typename ResultType, typename... ArgTypes>
+std::function<ResultType(ArgTypes...)> getFunction(const BinaryOp* bop) {
+  auto wrap_function = [](ResultType (*fn)(ArgTypes...)) { return fn; };
+
+  switch (bop->getBinaryOpType()) {
+    case BinaryOpType::Add:
+      return wrap_function(add);
+      break;
+    case BinaryOpType::Atan2:
+      return wrap_function(atan2);
+      break;
+    case BinaryOpType::Div:
+      return wrap_function(div);
+      break;
+    case BinaryOpType::Fmod:
+      return wrap_function(fmod);
+      break;
+    case BinaryOpType::Mul:
+      return wrap_function(mul);
+      break;
+    case BinaryOpType::Nextafter:
+      return wrap_function(nextafter);
+      break;
+    case BinaryOpType::Pow:
+      return wrap_function(pow);
+      break;
+    case BinaryOpType::Remainder:
+      return wrap_function(remainder);
+      break;
+    case BinaryOpType::Sub:
+      return wrap_function(sub);
+      break;
+    case BinaryOpType::Mod:
+      return wrap_function(mod);
+      break;
+    case BinaryOpType::Eq:
+      return wrap_function(eq);
+      break;
+    case BinaryOpType::NE:
+      return wrap_function(ne);
+      break;
+    case BinaryOpType::GT:
+      return wrap_function(gt);
+      break;
+    case BinaryOpType::GE:
+      return wrap_function(ge);
+      break;
+    case BinaryOpType::LT:
+      return wrap_function(lt);
+      break;
+    case BinaryOpType::LE:
+      return wrap_function(le);
+      break;
+    case BinaryOpType::BitwiseAnd:
+      return wrap_function(bitwise_and);
+      break;
+    case BinaryOpType::BitwiseOr:
+      return wrap_function(bitwise_or);
+      break;
+    case BinaryOpType::BitwiseXor:
+      return wrap_function(bitwise_xor);
+      break;
+    case BinaryOpType::LogicalAnd:
+      return wrap_function(logical_and);
+      break;
+    case BinaryOpType::LogicalOr:
+      return wrap_function(logical_or);
+      break;
+    case BinaryOpType::Lshift:
+      return wrap_function(bitwise_left_shift);
+      break;
+    case BinaryOpType::Rshift:
+      return wrap_function(bitwise_right_shift);
+      break;
+    case BinaryOpType::Gcd:
+      return wrap_function(gcd);
+      break;
+    case BinaryOpType::Min:
+      return wrap_function(minimum);
+      break;
+    case BinaryOpType::Max:
+      return wrap_function(maximum);
+      break;
+    case BinaryOpType::CeilDiv:
+      return wrap_function(ceilDiv);
+      break;
+    default:
+      NVF_CHECK(
+          false,
+          "Unexpected operator type: ",
+          bop->getBinaryOpType(),
+          " in ",
+          bop->toString());
+  }
+}
+
+// Get std::function for TernaryOp
+template <typename ResultType, typename... ArgTypes>
+std::function<ResultType(ArgTypes...)> getFunction(const TernaryOp* top) {
+  auto wrap_function = [](ResultType (*fn)(ArgTypes...)) { return fn; };
+
+  // clamp and threshold define a subset of TernaryOp configurations, so they
+  // are handled in a separate template specialization.
+  switch (top->getTernaryOpType()) {
+    case TernaryOpType::Lerp:
+      return wrap_function(lerp);
+      break;
+    case TernaryOpType::Where:
+      return wrap_function(where);
+      break;
+    case TernaryOpType::Threshold:
+    case TernaryOpType::Clamp:
+      NVF_CHECK(
+          false,
+          "Invalid function arguments for operator type",
+          top->getTernaryOpType(),
+          " in ",
+          top->toString());
+    default:
+      NVF_CHECK(
+          false,
+          "Unexpected operator type: ",
+          top->getTernaryOpType(),
+          " in ",
+          top->toString());
+  }
+}
+
+// Fully specialized template functions to create std::function for TernaryOp.
+template <>
+std::function<TensorView*(TensorView*, Val*, Val*)> getFunction<
+    TensorView*,
+    TensorView*,
+    Val*,
+    Val*>(const TernaryOp* top);
+
+template <>
+std::function<Val*(Val*, Val*, Val*)> getFunction<Val*, Val*, Val*, Val*>(
+    const TernaryOp* top);
+
+// Get std::function for ReductionOp
+template <typename ResultType, typename... ArgTypes>
+std::function<ResultType(ArgTypes...)> getFunction(const ReductionOp* rop) {
+  switch (rop->getReductionOpType()) {
+    case BinaryOpType::Add:
+      return sum;
+      break;
+    case BinaryOpType::Mul:
+      return prod;
+      break;
+    case BinaryOpType::Max:
+      return max;
+      break;
+    case BinaryOpType::Min:
+      return min;
+      break;
+    default:
+      NVF_CHECK(
+          false,
+          "Unexpected reduction operator type: ",
+          rop->getReductionOpType(),
+          " in ",
+          rop->toString());
+  }
+}
+
+// Get string name for UnaryOp
+std::string getString(const UnaryOp* uop);
+
+// Get string name for BinaryOp
+std::string getString(const BinaryOp* bop);
+
+// Get string name for TernaryOp
+std::string getString(const TernaryOp* bop);
+
+// Get string name for ReductionOp
+std::string getString(const ReductionOp* rop);
+
+// Get serde record type for ReductionOp
+serde::RecordType getSerdeType(const ReductionOp* rop);
+
+} // namespace nvfuser::python_frontend

nvfuser/include/nvfuser/scheduler/all_schedulers.h

@@ -0,0 +1,17 @@
+// 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/expr_eval_sched.h>
+#include <scheduler/matmul.h>
+#include <scheduler/no_op.h>
+#include <scheduler/normalization_inner.h>
+#include <scheduler/normalization_inner_outer.h>
+#include <scheduler/normalization_outer.h>
+#include <scheduler/pointwise.h>
+#include <scheduler/reduction.h>
+#include <scheduler/transpose.h>

nvfuser/include/nvfuser/scheduler/ampere_multi_matmul.h

@@ -0,0 +1,206 @@
+// 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/cuda/CUDAContext.h>
+#include <scheduler/multi_matmul.h>
+
+namespace nvfuser {
+
+// MmaOps in the scheduled tensor. Each one outputs a TensorView* which we call
+// an mma_result. Each MmaOp will also have two input TensorViews which we call
+// "ab" and "bb" since they are the immediate A and B operands and they contain
+// broadcast dimensions. Again there can be multiple abs and multiple bbs in
+// one fusion. These TensorViews are loaded from global memory tensors that we
+// call "a" and "b" into shared memory tensors called acw_smem and bcw_smem.
+// They are loaded from shared memory to register buffers we call "acr" and
+// "bcr" ("cr" meaning "cache read" in this context).
+//
+// Putting this all together we have the following order for a simple matmul
+//
+//   a -> acw_smem -> acr -> ... -> ab
+//                                    \                                      .
+//                                      mma_result ->  ... -> dc -> d
+//                                    /
+//   b -> bcw_smem -> bcr -> ... -> bb
+//
+// The ... indicate that there might be other tensors involved in a prologue or
+// epilogue section at that location.
+//
+// In this example there are two matmuls both using the same "a" operand:
+//
+//   b1 -> bcw_smem1 -> bcr1 -> ... -> bb1
+//                                        \                                  .
+//                                          mma_result1
+//                                        /             \                    .
+//       a -> acw_smem -> acr -> ... -> ab                ... -> dc -> d
+//                                        \             /
+//                                          mma_result2
+//                                        /
+//   b2 -> bcw_smem2 -> bcr2 -> ... -> bb2
+//
+// Note that there can be more than one output d and each one will have its own
+// register cache dc.
+//
+// Split-K and smem epilogue unswizzling add two additional tensors for each
+// mma in the fusion: splitk_sum and smem_epilogue.
+//
+//   // No split-K, no smem epilogue unswizzling:
+//     mma_result ->  ... -> dc -> d
+//   // split-K, no smem epilogue unswizzling:
+//     mma_result -> splitk_sum -> ... -> dc -> d
+//   // smem epilogue unswizzling, no split-K:
+//     mma_result -> smem_epilogue -> ... -> dc -> d
+//   // split-K and smem epilogue unswizzling:
+//     mma_result -> smem_epilogue -> splitk_sum -> ... -> dc -> d
+//
+// These additional tensors are added to each mma_result in the fusion.
+//
+// Each of the named tensors above is scheduled differently. We schedule them
+// by building AbstractTensors for each tensor category; these are held in
+// AmpereMultipleMatmulScheduler::schedules_.
+// TODO: Inherit from SchedulerEntry
+class AmpereMultipleMatmulScheduler : public MultipleMatmulScheduler {
+ public:
+  AmpereMultipleMatmulScheduler(Fusion* fusion, const MatmulParams* params)
+      : MultipleMatmulScheduler(fusion, params) {
+    const auto device_prop = at::cuda::getCurrentDeviceProperties();
+    const int cc = device_prop->major * 10 + device_prop->minor;
+    NVF_ERROR(
+        cc >= 75 && cc < 90,
+        "This matmul scheduler is restricted to Ampere and Turing.");
+  }
+
+  void run() final;
+
+ private:
+  void cacheInputsAndOutputs();
+
+  // Including current tensor naming convention for reference,
+  //  this is very temporary and will change over time and
+  //  in fact the whole body of this function will
+  //  eventually be a set of utility functions for different
+  //  sections of matmul(fusion) kernels, with
+  //  each having its own build out to do.
+  //
+  // Current naming convention is based on the following formula:
+  //
+  //  d = alpha * (a x b) + beta * c
+  //
+  // and is defined in the following way:
+  //
+  //  operands assumed in global memory : a, b, c
+  //
+  //  registers staging global load : ar, br (short for a/b read)
+  //
+  //  shared mem cache of operands : acw_smem, bcw_smem (short for a/b
+  //  cache_write smem)
+  //
+  //  registers at shared memory load output : acr, bcr (short for a/b cache
+  //  read)
+  //
+  //  register tensor input to the actual mma op: ab, bb (short for a/b
+  //  broadcasted)
+  //
+  //  accumulator register: mma_result
+  //   - mma_result is MmaOp output if there is epilogue
+  //   - mma_result is dc (short for d cache) if there is no epilogue
+  //
+  //  result in global memory: d
+
+  // Currently the support is for a, b, c and d as fusion inputs/outputs
+  //  aka. no prolog fusion yet.
+  void defineOperandCaches();
+
+  void cacheOperandsToSmem(
+      const std::vector<TensorView*>& operands,
+      std::vector<TensorView*>& smem_operands,
+      int64_t vec_size);
+
+  // We add two LoadStore operators to the inputs of our fusions. The first
+  // one is for a read from global memory and the second one (below) is for a
+  // cache read. As an optimizaton, we avoid adding an operator if there's an
+  // existing LoadStoreOp present. Please note that for the second LoadStore
+  // we don't propagate the allocation domain, since the scheduler sets the
+  // allocation domain in the registers.
+  void cacheOperandsToRegisters(
+      const std::vector<TensorView*>& tv_smems,
+      std::vector<TensorView*>& tv_rs);
+
+  //! Swizzle the M and N outer dimensions after makeTile has been called.
+  //! This updates outer_dim_roles if we introduce a new dimension, which can
+  //! happen if tv is missing a merged axis, in which case we skip merging after
+  //! the split. This is analogous to forwarding during transform propagation.
+  void swizzleBlockTiles(
+      TensorView* tv,
+      std::vector<MatmulDimRole>& outer_dim_roles);
+
+  //! This calls orig->cacheAfter() and also updates the broadcast graph to
+  //! reflect the new IterDomain mappings
+  TensorView* cacheAfter(
+      TensorView* orig,
+      LoadStoreOpType op_type = LoadStoreOpType::Set,
+      CacheOp cache_op = CacheOp::AllLevels,
+      bool propagate_allocation_domain = false);
+
+  //! Do block tiling for a collection of TensorViews. The tensors should be
+  //! unscheduled before this method is called.
+  //!   1) Axes will be ordered according to canonicalDimOrdering, and then axes
+  //! with the same role will be merged.
+  //!   2) After that, we perform splits according to
+  //!   params_->tile_sizes.cta_tile, e.g. [M, K] -> [Mo, Ko, Mi, Ki].
+  //!   3) Depending on the value of params_->grid_swizzle_factor, if the TV has
+  //! both M and N dimensions, we perform a 2D swizzle of the outer dimensions
+  //! Mo and No.
+  //!   4) Finally, we do a split-K split if the splitk_factor is not 1
+  std::vector<std::vector<MatmulDimRole>> blockTileTensors(
+      const std::vector<TensorView*>& tvs);
+
+  //! Schedule the loads of all operands from global memory to shared memory.
+  //! Starting from the basic tiled schedule, we swizzle the operand memory.
+  //! Note that the cache op and LoadStoreOpType are already set during
+  //! defineOperandCaches().
+  void scheduleOperandSmemStores();
+
+  void scheduleMmaOperands(
+      std::vector<TensorView*>& tvs,
+      const std::optional<MmaOperand> operand_type);
+
+  // MmaOperand contains only A and B. If tvs are outputs (i.e. not operands),
+  // then operand_type should be std::nullopt.
+  void scheduleMmaResults();
+
+  void schedulePrologues();
+
+  void scheduleOutputTensor(TensorView* c);
+
+  void scheduleEpilogue();
+
+  //! Propagates transformations from fusion output to fusion tv inputs that are
+  //!  producers in the epilogue. Transformations' propagation aims at input tvs
+  //!  which are not assigned to core roles, that is, are not MMA inputs.
+  void scheduleFusionInputsForEpilogue();
+
+  void scheduleSplitKSum();
+
+  void setUpInlining();
+
+  // NOTE: this should be called after acw_smem, acr, ..., ab, and mma_result
+  // transforms have been applied and inlining
+  void setUpCircularBuffering();
+
+ private:
+  std::vector<std::pair<TensorView*, TensorView*>> cached_outputs_;
+
+  std::vector<ValGroup> canonical_dim_ordering_;
+
+  std::vector<TensorView*> acw_smems_, bcw_smems_, acrs_, bcrs_, abs_, bbs_,
+      splitk_sums_, smem_epilogues_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/cache_policy_refiner.h

@@ -0,0 +1,19 @@
+// 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 <fusion.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+// Visits all global-to-local vector loads in `fusion` and refines their cache
+// policies.
+NVF_API void refineCachePolicy(Fusion* fusion);
+
+} // namespace nvfuser