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

nvfuser/include/nvfuser/kernel.h

@@ -0,0 +1,299 @@
+// 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/circular_buffer.h>
+#include <device_lower/analysis/sync_information.h>
+#include <device_lower/pass/warp_reduce.h>
+#include <fusion.h>
+#include <ir/base_nodes.h>
+#include <ir/builder.h>
+#include <parallel_dimension_map.h>
+#include <utils.h>
+#include <vectorization_info.h>
+#include <visibility.h>
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace nvfuser {
+namespace kir {
+
+//! Summary of interesting facts about the kernel
+// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+struct KernelSummary {
+  //! Count of WAR (write-after-read) hazard barriers
+  int64_t war_hazard_syncs_count = 0;
+
+  //! List of global buffers
+  std::vector<const kir::Allocate*> global_allocations;
+
+  //! List of dynamic shared memory buffers
+  std::vector<const kir::Allocate*> dynamic_smem_allocations;
+
+  //! List of static shared memory buffers
+  std::vector<const kir::Allocate*> static_smem_allocations;
+
+  //! Indicate the need to generate random numbers
+  bool has_philox_op = false;
+
+  //! Do we have any block reductions?
+  bool has_block_reductions = false;
+
+  //! Number of static grid reductions
+  bool has_grid_reductions = false;
+
+  //! Do we have any grid reduction in a loop, or grid reductions dependent on
+  //! grid reductions
+  bool has_cooperative_grid_reduction = false;
+
+  //! Do we have any block broadcasts?
+  bool has_block_broadcasts = false;
+
+  //! Do we have any grid broadcasts?
+  bool has_grid_broadcasts = false;
+
+  //! Do we have any welford op?
+  bool has_welford = false;
+
+  //! Do we have any welford op?
+  bool has_block_welford = false;
+
+  //! Do we have any welford op?
+  bool has_grid_welford = false;
+
+  //! Do we have any iter grouped outer block reduction op?
+  bool has_iter_grouped_reductions = false;
+
+  //! number of grouped iters for grouped outer block reduction
+  int64_t num_grouped_iterations = 1;
+
+  //! Do we have any outer grouped grid welford op?
+  bool has_outer_grouped_grid_welford = false;
+
+  //! Largest shared memory buffer size of outer grouped grid welford
+  int64_t outer_grouped_grid_welford_largest_smem_size = 0;
+
+  //! Largest shared memory buffer base type
+  DataType largest_smem_data_type = DataType::Null;
+
+  //! Do we have allocations of dynamic local memory?
+  bool has_dynamic_local_memory_allocations = false;
+
+  //! List of dynamic local memory buffers.
+  //! Only used for debugging.
+  std::vector<const kir::Allocate*> dynamic_lmem_allocations;
+
+  //! Validations needed and information about them. 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;
+
+  //! Effective ParallelTypes of broadcast ops
+  std::unordered_map<const BroadcastOp*, ParallelTypeBitmap>
+      broadcast_parallel_types;
+
+  //! Track which tensor views are inputs or outputs of a vectorized operation
+  //! and their maximum vectorized access size
+  std::unordered_map<TensorView*, int64_t> vectorized_accesses;
+
+  // Sync map is needed to figure out if global memory buffers need to be marked
+  // as volatile because they're used for communication.
+  std::shared_ptr<const SyncMap> sync_map;
+
+  // Parallel dimension map needed to set the correct properties of grid buffers
+  // (is a dim inactive)
+  ParallelDimensionMap parallel_dimension_map;
+
+  //! Track information on vectorized set operations for runtime validation
+  std::vector<VectorizedSetInfo> vectorized_set_info;
+
+  //! Minimum compute capability of device that can execute this kernel
+  std::pair<int64_t, int64_t> min_device_version;
+
+  //! Plain text description of why min_device_version_ is required
+  std::string min_device_version_reason;
+
+  //! Track Circular Buffer TensorViews
+  CircularBufferInfo circular_buffer_info;
+
+  //! Track if there are ElectSync predicates in this Kernel.
+  //! Reason: At runtime, we check that at least a single warp along TIDx axis
+  //! exists.
+  bool has_elect_sync_predicate = false;
+};
+
+class KernelPerformanceProfile {
+ public:
+  //! Register an expression to profile
+  void registerExpr(const Expr* expr);
+
+  //! Query if an expression is profiled
+  bool isProfiled(const Expr* expr) const;
+
+  //! Get the number of profiled expressions
+  int64_t getNumberOfProfileEntries() const {
+    return num_profile_entries_;
+  }
+
+  //! Set the backing buffer of profile.
+  void setBuffer(TensorView* buffer) {
+    buffer_ = buffer;
+  }
+
+  //! Get the backing buffer
+  TensorView* getBuffer() const {
+    return buffer_;
+  }
+
+  //! Get the indices of the profile of an expression in the backing buffer
+  std::array<int64_t, 2> getIndicesInProfileBuffer(const Expr* expr) const;
+
+  std::string toString(const at::Tensor& buffer) const;
+
+ private:
+  //! Get the new profile index
+  int64_t getNewIndex();
+
+  //! Get the profile index
+  std::optional<int64_t> getIndex(const Expr* expr) const;
+
+ private:
+  int64_t num_profile_entries_ = 0;
+
+  //! Backing buffer of Nx2 integer tensor, where N is the number of profiled
+  //! regions. Each region has two integer values, one representing
+  //! the cycles spent, and another the count.
+  TensorView* buffer_ = nullptr;
+
+  //! Map profiled expressions to profile entry offsets
+  std::unordered_map<const Expr*, int64_t> expr_entry_map_;
+
+  // TODO: Allow profiling of ForLoops
+  //! Map profiled ForLoop to profile entry offsets
+  // std::unordered_map<const ForLoop*, int64_t> loop_entry_map_;
+};
+
+class KernelInternalProxy;
+
+//! Container for a lowered Kernel IR
+//!
+// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+class NVF_API Kernel final : public Fusion {
+  friend KernelInternalProxy;
+
+ public:
+  // Kernel starts by grabbing all the nodes from the provided fusion.
+  // Kernel is not SSA, if a definition is not set, we should update it, but
+  // not remove previous definition if it is set. This is primarily because when
+  // we do something like generate an initialization statement for a reduction
+  // TV, we may want to continue to do fusion like analysis on the original
+  // expression.
+  Kernel(Fusion* fusion, PrimDataType index_type = PrimDataType::Int);
+  Kernel() = delete;
+
+  // No move or copy semantics
+  Kernel(const Kernel&) = delete;
+  Kernel& operator=(const Kernel&) = delete;
+
+  //! Finalize a kernel definition
+  //!
+  //! At this point we have a complete kernel definition and we can
+  //! run analysis passes to build a KernelSummary.
+  void finalize(std::vector<Expr*> top_level_exprs);
+
+  const std::vector<Expr*>& topLevelExprs() const {
+    return top_level_exprs_;
+  }
+
+  const KernelSummary& summary() const {
+    return summary_;
+  }
+
+  PrimDataType indexType() const {
+    return index_type_;
+  }
+
+  //! Checks if parallel type is padded
+  bool isParallelTypePadded(ParallelType ptype) const {
+    return ptype == ParallelType::TIDx &&
+        warp_padded_parallel_info_.is_tidx_padded;
+  }
+
+  const WarpPaddedParallelInfo& getWarpPaddedParallelInfo() const {
+    return warp_padded_parallel_info_;
+  }
+
+  const KernelPerformanceProfile& profile() const {
+    return profile_;
+  }
+
+  //! Debug dump of the Kernel IR
+  void print() const;
+
+  const std::vector<Val*>& parameters() const {
+    return parameters_;
+  }
+
+ protected:
+  using IrContainer::registerExpr;
+  using IrContainer::registerVal;
+
+  //! Register the Val with this fusion
+  void registerVal(Val* val) override;
+
+  //! Register expr with this fusion.
+  //! When we register an expression, we want to update the dependency tracking
+  //! of Vals. We add expr to our general expr_set_,
+  void registerExpr(Expr* expr) override;
+
+ private:
+  // Analyze the kernel IR and caches the summary of interesting data
+  void analyze();
+
+  // Top level statements
+  std::vector<Expr*> top_level_exprs_;
+
+  // Summary of interesting kernel data
+  KernelSummary summary_;
+
+  // Is this kernel being compiled with int32 or int64 indexing. This
+  // information is required to resolve DataType::Index
+  PrimDataType index_type_ = PrimDataType::Int;
+
+  WarpPaddedParallelInfo warp_padded_parallel_info_;
+
+  KernelPerformanceProfile profile_;
+
+  // Parameters of the kernel. The parameters contain the inputs and outputs of
+  // the kernel, intermediate buffers, and special items such as RNG state and
+  // tensor map for TMA support, etc. The parameters are not required to have no
+  // definition. If a parameter has a definition, its definition will be
+  // evaluated before the kernel is executed.
+  std::vector<Val*> parameters_;
+};
+
+//! A special debugging proxy for Kernel.
+//!
+//! Should not be used for other than testing and debugging.
+class NVF_API KernelInternalProxy {
+ public:
+  KernelInternalProxy(Kernel* kernel) : kernel_(kernel) {}
+
+  std::vector<Expr*>& topLevelExprs();
+
+ private:
+  Kernel* kernel_ = nullptr;
+};
+
+} // namespace kir
+} // namespace nvfuser

nvfuser/include/nvfuser/kernel_db/kernel_db.h

@@ -0,0 +1,109 @@
+// 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
+
+#if defined(__cplusplus) && (__cplusplus >= 201703L)
+#include <filesystem>
+namespace fs = std::filesystem;
+#elif defined(__cplusplus) && (__cplusplus >= 201402L)
+#include <experimental/filesystem>
+namespace fs = std::experimental::filesystem;
+#else
+#error "C++14 or Higher is required for filesystem library!"
+#endif
+
+#include <unordered_map>
+#include <vector>
+
+#include <visibility.h>
+
+namespace nvfuser {
+
+//! KernelDbEntry captures information to be printed per fusion in a csv file
+//! that is used to restore a hash map
+struct KernelDbEntry {
+  //! Cuda kernel function signature that is required to load the Cubin
+  std::string kernel_signature;
+  //! Compilation args supplied to NVRTC -- register usage and compute
+  //! capability can be specific to a kernel instance
+  std::string compile_args;
+  //! Full file path to Cuda Kernel
+  std::string kernel_code_file;
+  //! Full file path to cubin
+  std::string cubin_file;
+};
+
+//! KernelDb class is a singleton structure that is used to open, query, and
+//! write to the the database that is held in a hash map.  The kernel code is
+//! used as string key to the hash map.
+class KernelDb {
+  KernelDb(bool _disabled);
+
+  //! Open is private because this method should only be called once by the
+  //! singleton upon creation to create a new db or restore an existing one.
+  bool open(
+      const std::string& kernel_db_dir,
+      const std::string& kernel_db_file,
+      bool use_temp_dir);
+
+ public:
+  // clang-tidy - deleted member function should be public
+  KernelDb(const KernelDb&) = delete;
+  KernelDb& operator=(const KernelDb&) = delete;
+
+  //! Thread-Safe method to get the Meyer's singleton -- Interface
+  static KernelDb& get();
+  //! Thread-Safe method to get the Meyer's singleton -- For testing
+  NVF_API static KernelDb& get(
+      const std::string& kernel_db_dir,
+      const std::string& kernel_db_file,
+      bool use_temp_dir = true,
+      bool disabled = false,
+      bool reset = false);
+
+  //! Enable is derived from two booleans
+  bool enabled() const {
+    return !disabled_ && initialized_;
+  }
+  //! Returns the number entries in the db
+  size_t size() const {
+    return kernel_map_.size();
+  }
+
+  //! Query uses the string of the kernel code to lookup whether a cubin already
+  //! exists for the given kernel.  Additionally, the compile args are also
+  //! matched.
+  NVF_API bool query(
+      const std::string& kernel_code,
+      const std::string& compile_args,
+      std::string& kernel_signature,
+      std::vector<char>& cubin) const;
+  //! Write is used to write a new entry to the db upon compilation of a
+  //! new fusion
+  NVF_API bool write(
+      const std::string& kernel_code,
+      const std::string& compile_args,
+      const std::string& kernel_signature,
+      const std::vector<char>& cubin);
+
+ private:
+  //! Disablement is specified by the user and can also be set by a
+  //! failure to open the db
+  bool disabled_ = true;
+  //! Db is only initialized after it is successfully open
+  bool initialized_ = false;
+  //! Hash Map of kernel_string -> db_entry
+  std::unordered_map<std::string, KernelDbEntry> kernel_map_;
+
+  //! Full path to the db directory
+  fs::path kernel_db_path_;
+  //! Full path to csv file used to record and restore the db
+  fs::path kernel_db_txt_file_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/kernel_db/utils.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 <string>
+#include <vector>
+
+#include <visibility.h>
+
+namespace nvfuser {
+
+//! Helper methods to faciliate moving data between data buffers and files based
+//! on what type of data is being copied.
+
+NVF_API bool append_to_text_file(
+    const std::string& file_path,
+    const std::string& src);
+
+NVF_API bool copy_from_binary_file(
+    const std::string& file_path,
+    std::vector<char>& dst);
+NVF_API bool copy_from_text_file(
+    const std::string& file_path,
+    std::string& dst);
+
+bool copy_to_binary_file(
+    const std::string& file_path,
+    const std::vector<char>& dst);
+NVF_API bool copy_to_text_file(
+    const std::string& file_path,
+    const std::string& src);
+
+} // namespace nvfuser