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

nvfuser/include/nvfuser/device_lower/lower2device.h

@@ -0,0 +1,391 @@
+// 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 <compute_at_map.h>
+#include <device_lower/analysis/circular_buffer.h>
+#include <device_lower/analysis/fused_reduction.h>
+#include <device_lower/analysis/predicate_elimination.h>
+#include <device_lower/analysis/sync_information.h>
+#include <device_lower/analysis/tensor_memory.h>
+#include <device_lower/analysis/thread_predicate.h>
+#include <device_lower/analysis/tma.h>
+#include <device_lower/analysis/trivial_broadcast.h>
+#include <device_lower/id_model_options.h>
+#include <device_lower/pass/allocation.h>
+#include <device_lower/pass/circular_buffer.h>
+#include <device_lower/pass/predicate.h>
+#include <device_lower/pass/scalar_hoist.h>
+#include <device_lower/pass/warp_reduce.h>
+#include <exceptions.h>
+#include <expr_simplifier.h>
+#include <id_model/id_model.h>
+#include <id_model/indexing.h>
+#include <ir/all_nodes.h>
+#include <kernel.h>
+#include <kernel_ir.h>
+#include <logical_domain_map.h>
+#include <non_divisible_split.h>
+#include <options.h>
+#include <parallel_dimension_map.h>
+#include <runtime/executor_params.h>
+#include <vectorization_info.h>
+#include <visibility.h>
+
+#include <functional>
+#include <memory>
+#include <ostream>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace nvfuser {
+
+// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+class GpuLower : public NonCopyable {
+  class KernelIrMapper;
+
+ public:
+  GpuLower() = delete;
+
+  using Pass = std::pair<
+      std::string, // name of the pass
+      std::function<std::vector<Expr*>(const std::vector<Expr*>&)>>;
+
+  // GpuLower lowers the provided fusion into a kernel which can be translated
+  // into cuda code. index_type allows to compile the kernel based on int32
+  // indexing instead of int64 for additional performance.
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  NVF_API explicit GpuLower(
+      Fusion* fusion,
+      const CompileParams& cparams = CompileParams());
+
+  NVF_API kir::Kernel* kernel() const;
+
+  //! Returns the currently active lowering object.
+  //! It's an error if no lowering is in progress.
+  static GpuLower* current();
+
+  //! Query if lowering is in progress
+  static bool hasCurrent();
+
+  //! Actually run the lowering by executing the passes in the order given by
+  //! passes_
+  NVF_API kir::Kernel* run();
+
+  const PrimDataType& indexType() const {
+    return cparams_.index_type.value();
+  }
+
+  const auto& minDeviceVersion() const {
+    return min_device_version_;
+  }
+
+  const std::string& minDeviceVersionReason() const {
+    return min_device_version_reason_;
+  }
+
+  std::shared_ptr<const ConcretizedBroadcastDomains>
+  concretizedBroadcastDomains() {
+    return concretized_broadcast_domains_;
+  }
+
+  const ThreadPredicateMap& threadPredMap() const {
+    return thread_pred_map_;
+  }
+
+  // Returns non-const reference. Necessary to reset a predicate flag
+  // when a broadcast expression is fused into a reduction.
+  ThreadPredicateMap& threadPredMap() {
+    return thread_pred_map_;
+  }
+
+  std::shared_ptr<const ComputeAtMap> caMap() const {
+    return std::const_pointer_cast<const ComputeAtMap>(compute_at_map_);
+  }
+
+  bool hasIdModel() const {
+    return id_model_.get() != nullptr;
+  }
+
+  IdModel& idModel() {
+    NVF_ERROR(id_model_.get());
+    return *id_model_;
+  }
+
+  const IdModel& idModel() const {
+    NVF_ERROR(id_model_.get());
+    return *id_model_;
+  }
+
+  bool isTensorIndexerEnabled() const {
+    return tensor_indexer_.get() != nullptr;
+  }
+
+  TensorIndexer& tensorIndexer() {
+    NVF_ERROR(tensor_indexer_.get());
+    return *tensor_indexer_;
+  }
+
+  const TensorIndexer& tensorIndexer() const {
+    NVF_ERROR(tensor_indexer_.get());
+    return *tensor_indexer_;
+  }
+
+  const ParallelDimensionMap& parallelDimensionMap() const {
+    return parallel_dimension_map_;
+  }
+
+  ParallelDimensionMap& parallelDimensionMap() {
+    return parallel_dimension_map_;
+  }
+
+  PredicateElimination& predicateElimination() {
+    NVF_ERROR(pred_elimination_.get() != nullptr);
+    return *pred_elimination_;
+  }
+
+  const PredicateElimination& predicateElimination() const {
+    NVF_ERROR(pred_elimination_.get() != nullptr);
+    return *pred_elimination_;
+  }
+
+  LocalAllocationInfoMap& localAllocationInfoMap() {
+    return local_allocation_info_map_;
+  }
+
+  const WarpPaddedParallelInfo& getWarpPaddedParallelInfo() const {
+    return warp_pad_info_;
+  }
+
+  auto& nonDivisibleSplitInfo() {
+    return non_divisible_split_info_;
+  }
+
+  const auto& nonDivisibleSplitInfo() const {
+    return non_divisible_split_info_;
+  }
+
+  const auto& divisibleSplitSet() const {
+    return divisible_splits_;
+  }
+
+  CircularBufferInfo& circularBufferInfo() {
+    return circular_buffer_info_;
+  }
+
+  TmaCircularBufferInfo& tmaCircularBufferInfo() {
+    return tma_circular_buffer_info_;
+  }
+
+  CommonScalarMap& commonScalarMap() {
+    return common_scalar_map_;
+  }
+
+  const auto& vectorizedAccesses() const {
+    return vectorized_accesses_;
+  }
+
+  auto& vectorizedAccesses() {
+    return vectorized_accesses_;
+  }
+
+  const auto& vectorizedSetInfo() const {
+    return vectorized_set_info_;
+  }
+
+  auto& vectorizedSetInfo() {
+    return vectorized_set_info_;
+  }
+
+  FusedReductionInfo& fusedReductionInfo() {
+    return fused_reduction_info_;
+  }
+
+  std::shared_ptr<const SyncMap> syncMap() const {
+    return sync_map_;
+  }
+
+  kir::KernelPerformanceProfile& profile() {
+    return profile_;
+  }
+
+  std::unordered_map<const Expr*, TensorView*>& ldstMBarrierMap() {
+    return ldst_mbarrier_map_;
+  }
+
+  const std::unordered_map<const Expr*, TensorView*>& ldstMBarrierMap() const {
+    return ldst_mbarrier_map_;
+  }
+
+  bool isNvFuserZeroEnabled() {
+    if (isOptionDisabled(DisableOption::MagicZero)) {
+      return false;
+    }
+    return cparams_.enable_magic_zero;
+  }
+
+  // This is an interface to propagate information after expression
+  //  replacement on the kernel IR. E.g.:
+  //    for ...
+  //       c = a + b   (expr 0)
+  //  after any pass that does replacement:
+  //    for ...
+  //       c1 = a1 + b1 (expr1)
+  //  The previous analysis that was performed on expr0 might still
+  //    be valid on expr1 but that info would be lost after replacement.
+  //  This function provides an interface to manually update the info
+  //    in any pass that performs replacement.
+  void propagateExprInfo(const Expr* old_expr, const Expr* new_expr);
+
+  std::vector<Val*>& allKnownVals() {
+    return all_known_vals_;
+  }
+
+  const std::vector<Val*>& allKnownVals() const {
+    return all_known_vals_;
+  }
+
+  const std::vector<Pass>& passes() const {
+    return passes_;
+  }
+
+  std::vector<Pass>& passes() {
+    return passes_;
+  }
+
+  std::unordered_map<TensorView*, const TMAInfo>& consumerToTMAInfo() {
+    return consumer_to_tma_info_;
+  }
+
+  const std::unordered_map<TensorView*, const TMAInfo>& consumerToTMAInfo()
+      const {
+    return consumer_to_tma_info_;
+  }
+
+  const TensorMemoryInfo& tmemInfo() const {
+    return tmem_info_;
+  }
+
+  TensorMemoryInfo& tmemInfo() {
+    return tmem_info_;
+  }
+
+  // Register a boolean Val as a predicate to validate at the run time. Optional
+  // validation error messages can be given as args.
+  template <typename... Args>
+  void validate(Val* validation_condition, Args... args) {
+    auto sv = simplifyExpr(validation_condition);
+    if (sv->isTrue()) {
+      // If validation_condition is simplified to true, we know that the
+      // condition is always true regardless of the runtime values of the
+      // inputs. We can skip the validation. For example, we are not interested
+      // in validating that 3 < 4 or i % 8 < 8 every time we run the kernel.
+      return;
+    }
+    std::string message = to_str(args...);
+    NVF_ERROR(!sv->isFalse(), message);
+    validations_.emplace_back(sv, message);
+  }
+
+  const std::vector<std::pair<const Val*, std::string>>& validations() const {
+    return validations_;
+  }
+
+  std::vector<std::pair<const Val*, std::string>>& validations() {
+    return validations_;
+  }
+
+  // Get the index variable assigned for a given loop ID. Currently
+  //  it's a wrapper around ComputeAtMap::getIndexVariable or
+  // IdModel::getLoopIndexVariable if IdModelEnableOption::Loop is
+  //  enabled.
+  Val* getLoopIndexVariable(
+      IterDomain* id,
+      CircularBufferLoopStage stage =
+          CircularBufferLoopStage::NotApplicable) const;
+
+  const IdModelOptions idModelOptions() const {
+    return id_model_options_;
+  }
+
+ private:
+  void analysis(Fusion* fusion);
+
+  // Goes through the parallelized iterdomains of the used TVs and find
+  //  the parallel dimensions that need to be padded to a multiples of
+  //  warp size.
+  void collectPaddedParallelDims();
+
+  bool resolveComputeWith(Fusion* fusion);
+
+ private:
+  // Lowered Kernel IR
+  std::unique_ptr<kir::Kernel> kernel_;
+
+  // Passes to lower kernel, in order
+  std::vector<Pass> passes_;
+
+  // Some stateful information during lowering
+  // TODO: A lot of this information uses a define class then call build. It
+  // would be safer to wrap all of these in unique pointers and remove the build
+  // interface and default constructor. That way they couldn't be accessed
+  // without being initialized.
+  std::pair<int64_t, int64_t> min_device_version_;
+  std::string min_device_version_reason_;
+  std::shared_ptr<const ConcretizedBroadcastDomains>
+      concretized_broadcast_domains_;
+  ThreadPredicateMap thread_pred_map_;
+  std::unique_ptr<PredicateElimination> pred_elimination_;
+  std::shared_ptr<ComputeAtMap> compute_at_map_;
+  LocalAllocationInfoMap local_allocation_info_map_;
+  WarpPaddedParallelInfo warp_pad_info_;
+  ParallelDimensionMap parallel_dimension_map_;
+  NonDivisibleSplitInfo non_divisible_split_info_;
+  CircularBufferInfo circular_buffer_info_;
+  TmaCircularBufferInfo tma_circular_buffer_info_;
+  CommonScalarMap common_scalar_map_;
+  FusedReductionInfo fused_reduction_info_;
+  std::shared_ptr<const SyncMap> sync_map_;
+  kir::KernelPerformanceProfile profile_;
+  std::unordered_set<Split*> divisible_splits_;
+  CompileParams cparams_;
+  std::unique_ptr<IdModel> id_model_;
+  std::unique_ptr<TensorIndexer> tensor_indexer_;
+  std::unordered_map<TensorView*, const TMAInfo> consumer_to_tma_info_;
+
+  // Track which tensor views are inputs or outputs of a vectorized operation
+  // and their maximum vectorized access size
+  // std::unordered_map<TensorView*, VectorizationInfo> vectorized_accesses_;
+  std::unordered_map<TensorView*, int64_t> vectorized_accesses_;
+  // Info on each vectorized set op
+  std::vector<VectorizedSetInfo> vectorized_set_info_;
+
+  // All vals that are known to the kernel, including fusion inputs and
+  // precomputed values
+  std::vector<Val*> all_known_vals_;
+
+  // Keep track of the mbarrier used for each load/store operation
+  std::unordered_map<const Expr*, TensorView*> ldst_mbarrier_map_;
+
+  // Information about tensor memory usage
+  TensorMemoryInfo tmem_info_;
+
+  // Keep track of validations needed at runtime. For example, a pair of
+  //! "extent mod split_factor == 0" and an error message for divisibility check
+  //! for vectorization.
+  std::vector<std::pair<const Val*, std::string>> validations_;
+
+  Fusion* fusion_ = nullptr;
+
+  // A temporary option set to selectively enable IdModel usage
+  IdModelOptions id_model_options_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/alias_memory.h

@@ -0,0 +1,37 @@
+// 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 <dispatch.h>
+#include <ir/all_nodes.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! Reuse Allocation nodes via pointer aliasing
+//!
+//! First pass finds candidate TensorViews
+//! A candidate TensorView is anything in shared memory OR
+//! in local memory with a static size larger than register_size_threshold
+//!
+//! Second pass finds appropriate input Allocate Node
+//! among candidate TensorViews
+//!
+//! Alias Criteria:
+//! If input is a candidate TensorView,
+//!          input allocation has the same size as output allocation,
+//!          thread bindings match,
+//!          is not used after this op:
+//! then alias output Allocate to input Allocate.
+//!
+std::vector<Expr*> reuseMemoryAllocations(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/allocation.h

@@ -0,0 +1,32 @@
+// 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 <kernel_ir.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! Buffer allocation information to store in GPU lower to avoid
+//!  logic duplication
+struct LocalAllocationInfo {
+  kir::Allocate* alloc_expr = nullptr;
+  std::vector<IterDomain*> alloc_domains;
+};
+
+using LocalAllocationInfoMap = std::
+    unordered_map<const kir::Allocate*, std::unique_ptr<LocalAllocationInfo>>;
+
+//! Insert buffer allocations
+std::vector<Expr*> insertAllocations(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/circular_buffer.h

@@ -0,0 +1,191 @@
+// 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 <kernel_ir.h>
+#include <kernel_ir_dispatch.h>
+
+// Double buffering a tensor doubles its allocation size and uses two
+// buffers to facilitate computation and memory access
+// overlapping. The basic form of code looks like as follows:
+//
+// Before:
+// for i
+//   x[S]; // allocation
+//   for j:
+//     x[j] = y[i, j]
+//   for j:
+//     ... = x[j]
+//
+// After:
+// X[S * 2]; // allocation
+// for i in 0 to 1: // Prologue
+//   for j:
+//     x[j] = y[i, j]
+//
+// for i in 0 to N-1: // Main
+//   for j:
+//     x[j + (1 - i % 2) * S] = y[i + 1, j]
+//   for j:
+//     ... = x[j + (i % 2) * S]
+//
+// for i in N-1 to N: // Epilogue
+//   for j:
+//     ... = x[j + (i % 2) * S]
+//
+// Here, S is the original size of tensor x.
+//
+// The i loop is the double buffer loop of tensor x, where double
+// buffering is applied to the tensor. The first step of lowering is
+// to find the double buffering axis for each double buffered
+// tensor. It must not be parallelized as it isn't possible to double
+// buffer parallelized loops. Also, an unrolled axis expands the
+// allocation and is intended to make the loop completely unrolled,
+// which also conflicts with double buffering. So, basically, the double
+// buffering axis is the inner-most axis within the axes left
+// of the CA position. However, when it is parallelized or unrolled, a
+// further left axis is picked.
+//
+// Once the double buffer axis is determined, the main task is to
+// replicate the corresponding double buffer loop as illustrated
+// above. The Prologue loop is to just fetch the first element to
+// populate the buffer. The main loop is mostly the same as the
+// original loop, except for the indexing change to switch the two
+// buffers. When used as a consumer, an offset of (1 - i % 2) * S is
+// added, whereas (i % 2) * S is added when used as a producer. Here,
+// i is the index of the double buffer loop. The Epilogue loop is just
+// for the last iteration of the loop. Since the main loop reads one
+// element ahead of the producer of the double buffered tensor, it
+// would require an additional guard to prevent buffer overruns with
+// the producer if the main loop were also used for the last
+// iteration. However, the value loaded by the invalid load would not
+// be used, so instead of adding the additional predicate, the Epilogue
+// loop is replicated from the original loop, except for the load
+// expression since it's not used. Note that this overrun does not
+// happen when the producer is on gmem, so in that case, this
+// additional replication is not done.
+//
+// When creating those three types of loops, additional care must be
+// taken when multiple tensors are double buffered. When multiple
+// tensors use the same loop as their double buffer loop, one pass of
+// replication takes care of them at once, meaning the same Prologue,
+// Main, Epilogue loops are used for the multiple tensors.
+//
+// Other tasks to do for a double buffer tensor include:
+// - Move allocation to outside of the double buffer loop
+// - Double the allocation size
+// - Omit the RAW sync in the Main and Epilogue loops
+
+// [Cicular buffer] An generalization of double buffering.
+// On sm80+ hardware there is asynchronous copy infrastructure that
+//  motivates a circular buffering generalization of double buffering.
+// Almost all analyses previously done for double buffering are exactly
+//  the same with circular buffering, except for the introduction of
+//  new concept: `stage depth`.
+//
+// The `stage depth` is defined as the multiplier of extra buffering
+//  space used. In the case of double buffering, the stage depth would
+//  be 2.
+//
+// A circular buffered loop structure would look like follows, which
+//  exactly parallels the case of double buffered loop structure, since
+//  it is a exact generalization to the same purpose.
+//
+// Here S is the original allocation size as above,
+//  D is the stage depth. With D=2, the below loop structure becomes
+//  exactly the same as the case in double buffering.
+//
+// allocate X[S*D] // allocation
+// for i in 0..D-1: // prolog
+//   for j in ...
+//     if pred:
+//       x[i*S+j] = y[i, j];
+//
+// for i in 0..N: // main loop
+//   for j in ...
+//     if pred:
+//       x[((i+D-1)%D)*S+j] = y[i+D-1, j];
+//   for j in ...
+//     .. = x[(i%D)*S+j]
+//
+// (Epilog omitted since this only makes sense in using
+// cp.async, where producer will be in global mem and consumer will
+// be in shared mem).
+//
+// The profitability of this optimization comes from extra tolerance
+//  of global memory pipeline latency, as on the expression `.. = x[(i%D)*S+j]`
+//  we only need to make sure the data for the current iteration is
+//  completed while the remaining D-2 load iterations could still be in progress
+//  and overlap with the computes of the current loop.
+//
+// To express this pattern on sm80+ hardware we can group the loads
+//  in each iteration of the circular buffered loop as one "transaction",
+//  and specify how many transactions we want to ensure completion when
+//  we insert the async barriers.
+//
+// allocate X[S*D] // allocation
+// for i in 0..D-1: // prolog
+//   for j in ...
+//     if pred:
+//       x[i*S+j] = y[i, j];
+//   cp.async.commit; // mark the transaction boundary
+//
+// # At this point we have D-1 transactions on the fly.
+//   and for the first iteration of the main loop we need
+//   one transaction completed, so we leave D-2 transactions
+//   on the fly, which would be the input to the barrier instruction.
+//
+// cp.async.wait D-2 // ensure all but the last D-2 transactions complete.
+//
+// for i in 0..N: // main loop
+//   # At this point we always have D-2 transactions on the fly.
+//      and one completed.
+//   for j in ...
+//     if pred:
+//       x[((i+D-1)%D)*S+j] = y[i+D-1, j];
+//   for j in ...
+//     .. = x[(i%D)*S+j]
+//   cp.async.commit; // mark the transaction boundary for the
+//                       load issued in this iteration.
+//   # At this point we have D-1 transactions on the fly,
+//       and none completed.
+//   cp.async.wait D-2; // Ensure all but the last D-2 transactions complete.
+//   __syncthreads(); // Need to syncthreads because each thread will only
+//                      ensure completion of its own async copies so
+//                      would need to sync to this point to ensure
+//                      completion of the whole tile.
+
+namespace nvfuser {
+
+class TmaCircularBufferInfo {
+ public:
+  // Map cpAsyncBulk to its tensor index
+  void recordTensorIndex(const Expr* expr, kir::TensorIndex* index);
+
+  // Check if tensor index exists for expression
+  bool existsTensorIndex(const Expr* expr) const;
+
+  // Get tensor index for expression
+  kir::TensorIndex* getTensorIndex(const Expr* expr);
+
+ private:
+  // Track mbarrier used for cpAsyncBulk load operation. Required by indexing
+  // pass.
+  std::unordered_map<const Expr*, kir::TensorIndex*> ldst_mbarrier_index_map_;
+};
+
+class CircularBufferPass {
+ public:
+  //! Apply circular buffering transformations
+  static std::vector<Expr*> run(const std::vector<Expr*>& exprs);
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/expr_sort.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 <exceptions.h>
+#include <ir/base_nodes.h>
+
+namespace nvfuser {
+
+std::vector<Expr*> reorderExprsForComputeAt();
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/fusion_simplifier.h

@@ -0,0 +1,21 @@
+// 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 <dispatch.h>
+#include <fusion.h>
+#include <ir/all_nodes.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+// Transpose, Shift, Gather, and View Ops with LoadStoreOps
+std::vector<Expr*> loadStoreOpInserter(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser