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

nvfuser/include/nvfuser/device_lower/pass/misaligned_vectorization.h

@@ -0,0 +1,118 @@
+// 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 <vector>
+
+namespace nvfuser {
+
+//! Transform for-loop structure to handle misaligned addresses
+//!
+//! Sections of misaligned addresses are handled sequentially
+//! while aligned addresses use vectorized memory accesses.
+//!
+//! ---------------------------------------------------------------------------
+//! Before Misaligned Vectorization:
+//!
+//! Inputs: T0
+//! Outputs: T3
+//!
+//! for(...) {
+//!   T1[vector_size];
+//!   for( i : vector_size ) {
+//!     T1[i] = T0[...]
+//!   }
+//!
+//!   T2[vector_size];
+//!   for( i : vector_size ) {
+//!     T2[i] = unaryOp(T1[i])
+//!   }
+//!
+//!   for( i : vector_size ) {
+//!     T3[...] = T2[i]
+//!   }
+//! }
+//!
+//! ---------------------------------------------------------------------------
+//! After Misaligned Vectorization:
+//!
+//! Inputs: T0
+//! Outputs: T3
+//!
+//! for(...) {
+//!   T1[vector_size];
+//!   T2[vector_size];
+//!
+//!   if (inline_predicate_except_last_root_domain) {
+//!     index_except_last_root_domain = ...
+//!     address = (int64_t) &T1[index_except_last_root_domain]
+//!
+//!     offset_size = (address % vector_size_bytes) / data_type_size_bytes
+//!     shift_init = vector_size - offset_size
+//!     shift = (shift_init == vector_size) ? 0 : shift_init
+//!
+//!     // size of the last root domain
+//!     extent = ...
+//!     remainder = (extent - shift) % vector_size
+//!
+//!     last_root_domain_index = ...
+//!
+//!     // Vectorize Section
+//!     if ( (last_root_domain_index + shift) < (extent - remainder) ) {
+//!       T1[0] = vectorize_load( T0[index + shift] );
+//!
+//!       for( i : vector_size ) {
+//!         T2[i] = unaryOp(T1[i])
+//!       }
+//!
+//!       T3[index + shift] = vectorize_store( T2[0] );
+//!     }
+//!
+//!     // Initial Section
+//!     if ( last_root_domain_index == 0 ) {
+//!       for( i : shift ) {
+//!         T1[i] = T0[...]
+//!       }
+//!
+//!       for( i : shift ) {
+//!         T2[i] = unaryOp(T1[i])
+//!       }
+//!
+//!       for( i : shift ) {
+//!         T3[...] = T2[i]
+//!       }
+//!     }
+//!
+//!     // Remainder Section
+//!     if ( (last_root_domain_index + shift) >= (extent - remainder) &&
+//!          (last_root_domain_index + shift) < extent) {
+//!
+//!       for( i : remainder ) {
+//!         T1[i] = T0[index + shift]
+//!       }
+//!
+//!       for( i : remainder ) {
+//!         T2[i] = unaryOp(T1[i])
+//!       }
+//!
+//!       for( i : remainder ) {
+//!         T3[index + shift] = T2[i]
+//!       }
+//!     }
+//!   }
+//! }
+//!
+std::vector<Expr*> processMisalignedVectorization(
+    const std::vector<Expr*>& exprs);
+
+bool containsAnyDirectChildMisalignedVectorize(const ForLoop* fl);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/predicate.h

@@ -0,0 +1,23 @@
+// 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 {
+
+//! Update predicates with valid bool conditionals
+//!
+std::vector<Expr*> generateConditionalFromPredicate(
+    const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/replace_size.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 <dispatch.h>
+#include <fusion.h>
+#include <ir/all_nodes.h>
+
+namespace nvfuser {
+
+// TensorViews are all based on symbolic sizes. When we first initialize them
+// we don't know if they're inputs or outputs which would mean that they have
+// runtime shapes. Intermediate tensors (those not going to global memory) do
+// not have this information. Since we need to have the correct information in
+// the kernel being fetched for shapes, we want to replace input and output
+// tensors to reference the runtime structure containing sizes.
+void replaceSymbolicSizes(Fusion*);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/scalar_hoist.h

@@ -0,0 +1,115 @@
+// 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 <list>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace nvfuser {
+
+// Hoisting common subexpressions for scalar expressions, including indices,
+// predicates, tensor factories, etc.
+//
+// Class CommonScalarMap is updated during the lowering as new scalar
+// are inserted.
+//
+// Once all scalars are inserted to CommonScalarMap, allocations of the
+// the hoisted scalars are inserted by allocateCommonScalars. Note
+// that this assumes that the CUDA code generator does not inline a
+// scalar Val with allocation (PR #1434).
+
+class CommonScalarMap {
+ public:
+  //! For the given scalar, insert the subexpressions in its definition to the
+  //! loop that has minimum amount of computation. For example, if I have a loop
+  //! FOR i1
+  //!   FOR i2
+  //!     FOR i3
+  //!       FOR i4
+  //!         index = ((i1*1 + i2*2) + i3*3) + i4*4
+  //! and I want to hoist `index`. Then this function will try insert i1*1 to
+  //! common_scalar_map_[FOR i1], try insert i1*1 + i2*2 to
+  //! common_scalar_map_[FOR i2], try insert ((i1*1 + i2*2) + i3*3) to
+  //! common_scalar_map_[FOR i3], try insert ((i1*1 + i2*2) + i3*3) + i4*4 to
+  //! common_scalar_map_[FOR i4], Before insertion, this function recursively
+  //! uses reuseScalarIfAlreadyComputed to find existing
+  //! expressions/subexpressions in common_scalar_map_ that can be reused. If a
+  //! reuse oppportunity is found, then this function will modify the definition
+  //! of `value` to use the existing subexpression. This function returns the
+  //! modified value whose definition reuses other expressions in the list.
+  Val* hoistScalar(Val* value, const std::vector<ForLoop*>& loops);
+
+  //! common_scalar_map_ stores all seen indices in a given loop, however, we
+  //! don't want to create separate allocation for all of them. We are only
+  //! interested in allocating for the indices that is actually hoisted, or used
+  //! more than once. This method returns the Vals that will get its separate
+  //! allocation.
+  std::vector<Val*> getHoistedScalars(ForLoop* loop) const;
+
+  //! Initialize the common_scalar_map_ with lowered exprs. If some scalar is
+  //! already computed in these lowered exprs and is recomputed in indexing or
+  //! predicate math, then we should reuse these existing computation.
+  void initialize(const std::vector<Expr*> exprs);
+
+ private:
+  //! This is the underlying implementation of the public hoistScalar, with some
+  //! additional arguments and return values.
+  //! Returns (hoisted value, has tensor index dependency)
+  std::pair<Val*, bool> hoistScalarImpl(
+      Val* value,
+      const std::vector<ForLoop*>& loops,
+      std::vector<Val*>&
+          seen_subexprs, // Stores the subexpressions that has already been seen
+                         // during the recursion. This is used to detect
+                         // self-reuse. For example, if I have
+                         // i3 = i1 * i2 + i1 * i2
+                         // when visiting the second i1 * i2, I will have the
+                         // first i1 * i2 in this vector, so that we know we can
+                         // reuse that i1 * i2.
+      int64_t position, // if `value` is given to `hoistScalar` (i.e., is_give
+                        // == true), then this is the position of the outer-most
+                        // loop nest that contains all the dependencies of
+                        // `value`. if `value` is a subexpression of the value
+                        // given to `hoistScalar`, then this is the position of
+                        // the outer-most loop nest that contains all the
+                        // dependencies of its parent.
+      bool is_given = false // true for the given scalar from the public
+                            // `hoistScalar`, false otherwise.
+  );
+
+  //! If there is already an expression in common_scalar_map_[loop] which is
+  //! sameAs `value`, then just return that expression. Otherwise, if there is a
+  //! subexpression of an existing expression sameAs `value`, then that
+  //! subexpression will be split out as a separate item in the mapped list, and
+  //! that subexpression will be returned. If nothing sameAs `value`, then
+  //! return nullptr.
+  Val* reuseScalarIfAlreadyComputed(Val* value, ForLoop* loop);
+
+ private:
+  //! Map to hold hoisted common indices. The order matters and indicates data
+  //! dependency. For example, my list might have [i1*4, i1*4+2, i1*4/16]
+  std::unordered_map<ForLoop*, std::list<Val*>> common_scalar_map_;
+
+  //! A set to identify that if a val is hoisted (an expression used in the
+  //! inner loop, but its value only depend on outer loop variables, so the
+  //! computation of this expression is hoisted to an outer loop) or reused (one
+  //! expression is used in multiple indices/predicates).
+  std::unordered_set<Val*> hoisted_or_reused_;
+};
+
+//! Insert allocations of hoisted indices. Must be called after
+//! collecting all common indices.
+std::vector<Expr*> allocateCommonScalars(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/unroll.h

@@ -0,0 +1,98 @@
+// 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 <device_lower/analysis/thread_predicate.h>
+#include <device_lower/utils.h>
+#include <kernel_ir.h>
+#include <kernel_ir_dispatch.h>
+#include <logical_domain_map.h>
+
+#include <bitset>
+#include <unordered_map>
+
+namespace nvfuser {
+
+//! Unroll pass
+//!
+//! A bit deceptively: UnrollPass adds all predicates, so it needs to be run
+//! even if we don't unroll any loops.
+//!
+//! Unrolling pass will get IR that looks something 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}] = ...
+//!
+//! And it will return the following:
+//! for( i : I0o{ceil(I0/4)} ) {
+//!   for( j : I1o{ceil(I1/128)} ) {
+//!
+//!     if( i * 4 + 3 < I && j * 128 + 127 < J ){
+//!       for( k : I0i{4} )
+//!         for( l : I1i{128} )
+//!           T0[ ( i * 4 + k ) * J + j * 128 + l ] = ...
+//!     } else {
+//!       for( k : I0i{4} )
+//!         for( l : I1i{128} )
+//!           if( i * 4 + k < I && j * 128 + l < J)
+//!              T0[ ( i * 4 + k ) * J + j * 128 + l ] = ...
+//!     }
+//!
+//!   }
+//! }
+//!
+//! As can be seen it generates two sets of loops for I0i{4} and I1i{128}. The
+//! first set is protected by a predicate that makes sure there's a full
+//! internal tile we can iterate over. This way we remove the predicate nested
+//! in the inner most loop. There's of course a second set of loops, which has a
+//! predicate still in the inner most loop, making sure that we cover edges and
+//! corners.
+//!
+class UnrollPass : kir::ExprMutator {
+ public:
+  // Take the incoming exprs and run loop unrolling, returning the new IR
+  static std::vector<Expr*> runPass(const std::vector<Expr*>& exprs);
+
+  static bool canOmitElseClause(ForLoop* fl);
+
+ private:
+  void registerReplace(Expr* reference, Expr* new_expr);
+
+  // Generate the for Expr replacement map
+  UnrollPass(const std::vector<Expr*>& exprs);
+
+  const std::unordered_map<Expr*, Expr*>& replacementMap() const {
+    return expr_replacement_map_;
+  }
+
+  using kir::ExprMutator::handle;
+
+  void handle(ForLoop* fl) final;
+
+  void dispatch(Expr* expr) final;
+
+ private:
+  // We will track which loops in the incoming IR will be replaced and by what
+  std::unordered_map<Expr*, Expr*> expr_replacement_map_;
+
+  // keep track if we're within an unrolled loop
+  bool look_for_unroll_ = true;
+
+  // Indicates if the currently visited expression is inside a
+  // unswitched path
+  bool unswitched_loop_ = false;
+
+  // As we generate inline predicates check if we actually generated a
+  // non-trivial one.
+  bool non_trivial_pred_found_ = false;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/vectorize_welford.h

@@ -0,0 +1,45 @@
+// 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 <vector>
+
+namespace nvfuser {
+
+class Expr;
+
+// Apply loop-invariant code hoisting to serial WelfordOps. For
+// example, when the innermost loop looks like:
+//
+// for () {
+//   welfordCombine(...);
+/// }
+//
+// The count input should be invariant when the loop is not a
+// reduction loop, and then this can be transformed as:
+//
+// After:
+// nvfuser_index_t new_count = outN()[0] + 1;
+// float reciprocal = 1 / new_count;
+// for () {
+//   welfordVectorized(..., new_count, reciprocal);
+// }
+//
+// Here, welfordVectorized does not need to compute the division. This
+// transformation can be applied when the innermost loop is a
+// non-reduction domain and there's no predicate depending on the loop
+// index of the innermost loop. A common case is when the read of a
+// fusion input is vectorized and that input is fed to an outer
+// welford reduction. In this case, the innermost domain is a
+// non-reduction domain and is vectorized, so the prediacte should not
+// have any dependency with the loop index, which enables the code
+// moition as the above.
+std::vector<Expr*> vectorizeWelford(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/device_lower/pass/warp_reduce.h

@@ -0,0 +1,23 @@
+// 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 <kernel_ir.h>
+
+namespace nvfuser {
+
+struct WarpPaddedParallelInfo {
+  bool is_tidx_padded = false;
+  bool is_tidx_single_warp = false;
+  bool has_warp_reduction = false;
+};
+
+std::vector<Expr*> fuseWarpReduce(const std::vector<Expr*> exprs);
+
+} // namespace nvfuser