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

nvfuser/include/nvfuser/device_lower/pass/grid_serialization.h

@@ -0,0 +1,26 @@
+// 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/macros/Export.h>
+#include <exceptions.h>
+
+#include <ir/all_nodes.h>
+#include <kernel_ir.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! Detect ReductionOps that have serialGridReductionRequested() == true. When
+//! found, confirm that no conflicting operations exist, then place sync nodes
+//! before and after outer-most non-parallelized loop.
+std::vector<Expr*> insertGridSerializationSyncs(
+    const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/index.h

@@ -0,0 +1,200 @@
+// 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 <instrumentation.h>
+#include <kernel_ir.h>
+#include <kernel_ir_dispatch.h>
+#include <logical_domain_map.h>
+
+#include <unordered_set>
+#include <vector>
+
+namespace nvfuser {
+
+// TODO: Replace with mutator as IndexLowering is replacing expr's with
+// versions that are doing indexing
+class IndexLowering : private OptOutConstDispatch {
+ public:
+  static std::vector<Expr*> getIndexedExprs(std::vector<Expr*> incoming_exprs);
+
+ private:
+  IndexLowering() = default;
+
+  void pushBack(Expr*);
+
+  // Return the most recently inserted
+  //  expression in the current active
+  //  scope or global scope.
+  Expr* back() const;
+
+  // Insert an expression before the current top-level expression.
+  void insertAtTopLevel(Expr* expr);
+
+  void handle(const FullOp*) final;
+  void handle(const IotaOp*) final;
+  void handle(const EyeOp*) final;
+  void handle(const ViewAsScalar*) final;
+  void handle(const UnaryOp*) final;
+  void handle(const BinaryOp*) final;
+  void handle(const TernaryOp*) final;
+  void handle(const ArrayConstruct*) final;
+  void handle(const StructConstruct*) final;
+  void handle(const GetAttr*) final;
+  void handle(const GetItem*) final;
+  void handle(const GetMetaData*) final;
+  void handle(const TensorConstruct*) final;
+  void handle(const SelectOp*) final;
+  void handle(const IndexSelectOp*) final;
+  void handle(const TorchGatherOp*) final;
+  void handle(const ScatterOp*) final;
+  void handle(const RNGOp*) final;
+  void handle(const ReductionOp*) final;
+  void handle(const GroupedReductionOp*) final;
+  void handle(const WelfordOp*) final;
+  void handle(const GroupedWelfordOp*) final;
+  void handle(const LoadStoreOp*) final;
+  void handle(const MmaOp*) final;
+  void handle(const BroadcastOp*) final;
+  void handle(const PadOp*) final;
+  void handle(const SliceOp*) final;
+  void handle(const CatOp*) final;
+
+  void handle(const kir::Asm*) final;
+  void handle(const ForLoop*) final;
+  void handle(const kir::IfThenElse*) final;
+  void handle(const kir::Allocate*) final;
+  void handle(const kir::BlockSync*) final;
+  void handle(const kir::GridSync*) final;
+  void handle(const kir::FenceAsyncProxy*) final;
+  void handle(const kir::WgMmaFence*) final;
+  void handle(const kir::SetMaxNReg*) final;
+  void handle(const kir::Return*) final;
+  void handle(const kir::MBarrierInit*) final;
+  void handle(const kir::MBarrierInvalidate*) final;
+  void handle(const kir::MBarrierArrive*) final;
+  void handle(const kir::MBarrierArriveExpectTx*) final;
+  void handle(const kir::MBarrierWait*) final;
+  void handle(const kir::MBarrierWaitParity*) final;
+  void handle(const kir::AsyncWait*) final;
+  void handle(const kir::AsyncCommit*) final;
+  void handle(const kir::BlockSerializeWait*) final;
+  void handle(const kir::BlockSerializeRelease*) final;
+
+  void generate(const std::vector<Expr*>& exprs);
+
+  // Get the loop in which the currently visiting expr is a rotated expr.
+  const std::unordered_set<ForLoop*>& getRotatedLoop() const {
+    return rotated_loop_;
+  }
+
+  // lower index for producer. The `override_index` is a mapping `id->index`,
+  // where `id` must be an IterDomain in the rFactor domain of the producer.
+  // This is can used to manually set the index for the given rFactor ID.
+  // Currently, this `override_index` is only used by indexing ops like
+  // select/index_select.
+  // The argument `generate_pointer` specifies whether to generate pointer for
+  // the tensor. If global tensor, then generate T1.data. If shared memory
+  // tensor, then use `cvta` ptx to convert shared memory address to unsigned
+  // int for indexing. Search `toSmem` in the codebase for additional
+  // information. This argument is effective only if the indexed tensor is a
+  // shared memory or global tensor. On other memory type, this argument will
+  // cause an error.
+  Val* lowerSrcIndex(
+      Val* val,
+      Val* dst,
+      const std::unordered_map<IterDomain*, Val*>& override_index = {},
+      bool generate_pointer = false,
+      DataType as_type = DataType::Null) const;
+
+  Val* lowerDstIndex(
+      Val* dst,
+      const std::unordered_map<int, Val*>& override_index = {},
+      bool generate_pointer = false,
+      DataType as_type = DataType::Null) const;
+
+  void handleCpAsyncBulkLoad(const LoadStoreOp* ldst);
+  void handleCpAsyncBulkStore(const LoadStoreOp* ldst);
+
+  void handleBlockReduction(const ReductionOp* rop, Val* out, Val* in);
+  void handleGridReduction(const ReductionOp* rop, Val* out, Val* in);
+  //! Called by handleGridReduction, this returns true if rop is lowered as a
+  //! serial grid reduction.
+  void handleSerialGridReduction(const ReductionOp* rop, Val* out, Val* in);
+
+  void handleBlockReduction(
+      const GroupedReductionOp* rop,
+      const std::vector<Val*>& outputs,
+      const std::vector<Val*>& inputs);
+  void handleGridReduction(
+      const GroupedReductionOp* rop,
+      const std::vector<Val*>& outputs,
+      const std::vector<Val*>& inputs);
+
+  void handleGridWelford(WelfordOp* new_wop);
+
+  void handleGroupedBlockWelford(
+      const GroupedWelfordOp* wop,
+      const std::vector<WelfordTriplet>& output_vals,
+      const std::vector<WelfordTriplet>& input_vals,
+      const std::vector<WelfordTriplet>& init_vals);
+  void handleGroupedGridWelford(
+      const GroupedWelfordOp* wop,
+      const std::vector<WelfordTriplet>& output_vals,
+      const std::vector<WelfordTriplet>& input_vals,
+      const std::vector<WelfordTriplet>& init_vals);
+
+  // Allocate a unique buffer for grid reductions and broadcast. A
+  // buffer is uniquely allocated for each output tensor of an
+  // expression.
+  kir::Allocate* allocateUniqueBuffer(
+      Val* buffer_size,
+      DataType dtype,
+      bool zero_init,
+      TensorView* out_tv,
+      std::unordered_map<TensorView*, kir::Allocate*>& alloc_map);
+
+  std::vector<kir::Allocate*> allocateWelfordWorkBuffer(
+      const std::vector<WelfordTriplet>& triplets,
+      WelfordTriplet::ValName name,
+      Val* buffer_size);
+
+  // Allocate a fused reduction object uniquely for a given
+  // TensorView. Parameter expr is the expression corresponding to the
+  // fused reduction.
+  void allocateUniqueFusedReduction(Expr* expr, TensorView* out_tv);
+
+ private:
+  std::vector<Expr*> lowered_exprs_;
+
+  // This is a slight work around as scope has a couple definitions, we have the
+  // Scope that's in ForLoop/IfThenElse which is really just a wrapper around
+  // std::vector<Expr*> and then we have the actual ForLoop/IfThenElse. We want
+  // to be able to carry both around because when we push back to a scope it
+  // could be either the body or else body of the IfThenElse. However, we want
+  // to understand the nesting of IfThenElse/ForLoop nodes.
+  Scope* active_scope_ = nullptr;
+
+  // Track for loops to send to indexing. Similar to what's done in
+  // kir::IrVisitor
+  std::vector<ForLoop*> for_loops_;
+
+  // Keep track of the loop in which the currently visiting expr is a rotated.
+  std::unordered_set<ForLoop*> rotated_loop_;
+
+  // Maps to keep track of allocated buffers and objects that must be
+  // allocated only once
+  std::unordered_map<TensorView*, kir::Allocate*> sync_buffer_map_;
+  std::unordered_map<TensorView*, kir::Allocate*> work_buffer_map_;
+  std::unordered_map<TensorView*, kir::AllocateFusedReduction*>
+      fused_reduction_map_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/inline_ptx.h

@@ -0,0 +1,16 @@
+// 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>
+
+namespace nvfuser {
+
+std::vector<Expr*> lowerToInlinePtx(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/insert_syncs.h

@@ -0,0 +1,39 @@
+// 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 {
+
+//! Insert sync at end of for-loops to prevent write-after-read race condition.
+//!
+//! WAR race condition occurs when the next iteration of the loop overwrites
+//! shared memory value before a previous operation has finished reading it.
+std::vector<Expr*> insertWarThreadSynchronization(
+    const std::vector<Expr*>& exprs);
+
+//! Insert syncs between writing to shared memory and then reading it.
+//! RAW pass is run before indexing, unrolling (loop duplication), memory
+//! aliasing, and index (grid/block bcast/reduction)
+std::vector<Expr*> insertRawThreadSynchronization(
+    const std::vector<Expr*>& exprs);
+
+//! Insert wait expressions such as wgmma.wait_group.sync.aligned at end of
+//! for-loops for async expressions to prevent write-after-read race condition.
+//!
+//! WAR race condition occurs when the next iteration of the loop overwrites
+//! the buffer before a previous async expression has finished reading it.
+std::vector<Expr*> insertWarAsyncWait(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/instrument.h

@@ -0,0 +1,24 @@
+// 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>
+
+namespace nvfuser {
+
+//! Set up KernelPerformanceProfile of GpuLower when enabled, which
+//! keeps track of expressions to profile. A new TensorView is added
+//! for storing profiling results. The expression list is prepended
+//! with an kir::Allocate node to allocate the TensorView profile
+//! buffer. Note that any expression added after this pass will not be
+//! profiled, so this pass should be called after all expressions are
+//! lowered. KernelPerformanceProfile is copied to Kernel after
+//! lowering.
+std::vector<Expr*> instrumentKernel(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/loop_rotation.h

@@ -0,0 +1,150 @@
+// 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>
+
+namespace nvfuser {
+
+// Note [Loop Rotation]
+// Loop Rotation is an optimization pass to improve instruction scheduling. For
+// a given loop, for example:
+//   for (int i = 0; i < n; i++) {
+//     line1(i);
+//     line2(i);
+//     line3(i);
+//     line4(i);
+//   }
+// If we rotate one line up, then we get
+//   if (0 < n) {
+//     line1(0);
+//   }
+//   for (int i = 0; i < n; i++) {
+//     line2(i);
+//     line3(i);
+//     line4(i);
+//     if (i + 1 < n) {
+//       line1(i + 1);
+//     }
+//   }
+// Similarly, if we rotate two lines up, then we get
+//   if (0 < n) {
+//     line1(0);
+//     line2(0);
+//   }
+//   for (int i = 0; i < n; i++) {
+//     line3(i);
+//     line4(i);
+//     if (i + 1 < n) {
+//       line1(i + 1);
+//       line2(i + 1);
+//     }
+//   }
+// In order to take advantage of this pass, the scheduler needs to specify which
+// loop to rotate and the consumers whose allocation and computation will be
+// rotated, and pass this information as compilation parameter. For example, if
+// I have a fusion that will create the following loop structure:
+//   for (int i = 0; i < id1.extent(); i++) {
+//     float T1[5];
+//     for (int j = 0; j < 5; j++) {
+//       if (i < T0.size[0]) {
+//         T1[j] = sin(T0[i, j]);
+//       }
+//     }
+//     float T2[5];
+//     for (int j = 0; j < 5; j++) {
+//       T2[j] = cos(T1[j]);
+//     }
+//     float T3[5];
+//     for (int j = 0; j < 5; j++) {
+//       T3[j] = exp(T2[j]);
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       if (i < T4.size[0]) {
+//         T4[i, j] = log(T3[j]);
+//       }
+//     }
+//   }
+// Then the scheduler could make a compilation parameter {id1, {T1, T2}} to the
+// fusion, and this pass will transform the code as
+//   float T1[5];
+//   float T2[5];
+//   if (0 < id1.extent()) {
+//     for (int j = 0; j < 5; j++) {
+//       if (0 < T0.size[0]) {
+//         T1[j] = sin(T0[0, j]);
+//       }
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       T2[j] = cos(T1[j]);
+//     }
+//   }
+//   for (int i = 0; i < id1.extent(); i++) {
+//     float T3[5];
+//     for (int j = 0; j < 5; j++) {
+//       T3[j] = exp(T2[j]);
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       if (i < T4.size[0]) {
+//         T4[i, j] = log(T3[j]);
+//       }
+//     }
+//     if (i + 1 < id1.extent()) {
+//       for (int j = 0; j < 5; j++) {
+//         if (i + 1 < T0.size[0]) {
+//           T1[j] = sin(T0[i + 1, j]);
+//         }
+//       }
+//       for (int j = 0; j < 5; j++) {
+//         T2[j] = cos(T1[j]);
+//       }
+//     }
+//   }
+// Currently, because all our existing predicates should already cover
+// out-of-bound access, so we are omitting the predicates to get a
+// better-looking code:
+//   float T1[5];
+//   float T2[5];
+//   for (int j = 0; j < 5; j++) {
+//     if (0 < T0.size[0]) {
+//       T1[j] = sin(T0[0, j]);
+//     }
+//   }
+//   for (int j = 0; j < 5; j++) {
+//     T2[j] = cos(T1[j]);
+//   }
+//   for (int i = 0; i < id1.extent(); i++) {
+//     float T3[5];
+//     for (int j = 0; j < 5; j++) {
+//       T3[j] = exp(T2[j]);
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       if (i < T4.size[0]) {
+//         T4[i, j] = log(T3[j]);
+//       }
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       if (i + 1 < T0.size[0]) {
+//         T1[j] = sin(T0[i + 1, j]);
+//       }
+//     }
+//     for (int j = 0; j < 5; j++) {
+//       T2[j] = cos(T1[j]);
+//     }
+//   }
+
+// vector of (tv, dim, selection)
+// For each entry in the vector, the selected tv/expr in loop tv->axis(dim)
+// will be rotated
+using LoopRotationParam = std::vector<
+    std::tuple<TensorView*, int64_t, std::unordered_set<Statement*>>>;
+
+std::vector<Expr*> rotateLoops(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/loops.h

@@ -0,0 +1,68 @@
+// 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/thread_predicate.h>
+#include <instrumentation.h>
+#include <ir/all_nodes.h>
+#include <kernel_ir.h>
+
+namespace nvfuser {
+
+//! Loop nest generator pass will get IR that looks something like:
+//! T0[I0o{ceil(I0/4)}, I1o{ceil(I1/128)}, I0iU{4}, I1i{128}] = ...
+
+//  and will generate the loop nest structure for these exprs like:
+//!
+//! for( i : I0o{ceil(I0/4)} ) {
+//!   for( j : I1o{ceil(I1/128)} ) {
+//!     for( k : I0i{4} )
+//!       for( l : I1i{128} )
+//!         T0[I0o{ceil(I0/4)}, I1o{ceil(I1/128)}, I0iU{4}, I1i{128}] = ...
+//!
+//! It does not generate predicates, but it will generate allocations, and loop
+//! nests to initialize reduction buffers.
+class LoopNestGenerator {
+ public:
+  static std::vector<Expr*> loweredExprs(const std::vector<Expr*>& exprs);
+
+ private:
+  LoopNestGenerator(const std::vector<Expr*>& exprs);
+
+  // Open a new inner most for loop, track which TV it was constructed from
+  // according to the computeAt chain.
+  void openFor(IterDomain*);
+
+  // Close the inner most for loop
+  void closeFor();
+
+  // Appends an expression to the current scope
+  void pushFront(Expr* expr);
+
+  void handle(Expr* expr);
+
+  // Run the pass and accumulate output in lowered_exprs_
+  void generate(const std::vector<Expr*>& exprs);
+
+ private:
+  // Lowered exprs to return
+  std::vector<Expr*> lowered_exprs_;
+
+  // Keep all for loops conveniently to make unrolling easier, basically just a
+  // stack of the active for_loops
+  std::vector<ForLoop*> for_loops_;
+
+  // Loop structure of each expression
+  std::unordered_map<TensorView*, std::vector<IterDomain*>> loop_structures_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/magic_zero.h

@@ -0,0 +1,86 @@
+// 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 <visibility.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+struct IndexFromIdGraph;
+
+//! Insert magic zero definition at the begining of the kernel. Insert magic
+//! zero update after every (outer most) loop nest with a compile time extent.
+//!
+//! This will make sure nvrtc does not aggressively save predicate and indices.
+std::vector<Expr*> insertMagicZero(const std::vector<Expr*>& exprs);
+
+//! Check if val is a reference to the magic zero variable
+NVF_API bool isMagicZero(const Val* val);
+
+//! Check if val is protected with magic zero.
+//!
+//! Specifically, this returns true if val is defined as "x + magic_zero".
+bool isProtectedWithMagicZero(const Val* val);
+
+//! Check if val is protected with magic zero, if yes, unwrap it.
+//! maybeUnwrapMagicZero(i1) -> i1
+//! maybeUnwrapMagicZero(i2 + magic_zero) -> i2
+Val* maybeUnwrapMagicZero(Val* val);
+
+// Determine if we may run into over reuse of predicates or registers in the
+// compiler. If the loop can be unrolled and the index and domain are not
+// "simple" we likely want the loop protected.
+//
+// Magic zero protection should only be done for global memory and predicates.
+// We should avoid use on registers. Shared memory does not require it, but
+// likely wouldn't hurt.
+bool needsMagicZero(
+    ForLoop* loop,
+    IterDomain* reference_domain = nullptr,
+    Val* ind = nullptr);
+
+struct IndexMagicZeroInfo {
+  //! Index that may be updated with magic zero
+  Val* index = nullptr;
+  //! Loop index that is protected by magic zero. nullptr if no loop
+  //! is protected
+  Val* original_loop_index = nullptr;
+  //! Protected loop index. nullptr if no loop is protected
+  Val* protected_loop_index = nullptr;
+  //! Protected loop. nullptr if no loop is protected
+  IterDomain* loop_id = nullptr;
+};
+
+//! Protect an index val of an IterDomain with magic zero
+//!
+//! This should be only used for predicate indexing.
+//!
+//! No protection is done if none of the loops is determined to require
+//! protection by needsMagicZero.
+IndexMagicZeroInfo protectPredicateIndexWithMagicZero(
+    Val* index,
+    const IndexFromIdGraph& id_graph,
+    const std::vector<ForLoop*>& loops);
+
+//! Protect an index val of a tensor with magic zero
+//!
+//! This should be only used for non-predicate indexing.
+//!
+//! No protection is done if none of the loops is determined to require
+//! protection by needsMagicZero.
+void protectNonPredicateIndexWithMagicZero(
+    const std::vector<ForLoop*>& loops,
+    const std::vector<IterDomain*>& loop_domains,
+    std::unordered_map<IterDomain*, Val*>& concrete_loop_idx_map);
+
+} // namespace nvfuser