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

nvfuser/include/nvfuser/kernel_ir_dispatch.h

@@ -0,0 +1,147 @@
+// 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 <exceptions.h>
+
+namespace nvfuser {
+
+class Expr;
+class ForLoop;
+class Scope;
+
+namespace kir {
+class Predicate;
+class TensorIndex;
+class IfThenElse;
+
+// Base visitor class that visits all nodes in provided vector<Expr*>.
+//
+// Includes visiting through scopes like IfThenElse and ForLoop, and tracks
+// them in scopes_ and for_loops_.
+//
+// Makes a copy of exprs at exprs_ which could be used to modify and return.
+//
+// When traversing through ITE/FLs it will use a copy
+// of the provided expressions to make it safe to insert/delete nodes.
+//
+// Provides a simple base class to inherit from for typical lowering passes on
+// Expr list
+class NVF_API IrVisitor : public OptOutDispatch {
+ public:
+  std::vector<Expr*> handle(const std::vector<Expr*>& exprs);
+
+ protected:
+  using OptOutDispatch::handle;
+
+  void handle(ForLoop*) override;
+  void handle(IfThenElse*) override;
+
+ protected:
+  std::vector<ForLoop*> for_loops_;
+  std::vector<Scope*> scope_;
+  std::vector<Expr*> scope_exprs_;
+  std::vector<Expr*> exprs_;
+};
+
+// Const version of IrVisitor
+class ConstIrVisitor : public OptOutConstDispatch {
+ public:
+  std::vector<const Expr*> handle(const std::vector<const Expr*>& exprs);
+
+ protected:
+  using OptOutConstDispatch::handle;
+
+  void handle(const ForLoop*) override;
+  void handle(const IfThenElse*) override;
+
+ protected:
+  std::vector<const ForLoop*> for_loops_;
+  std::vector<const Scope*> scope_;
+  std::vector<const Expr*> scope_exprs_;
+  std::vector<const Expr*> exprs_;
+};
+
+// Base Expr Mutator class that visits all nodes with IrVisitor, and then
+// inserts new expressions, replaces expressions based on insertion/replace
+// maps provided or removes existing expressions. These replacement
+// maps are expected to accumulate during an initial traversal, then
+// runs an insertion based on them after the overloaded traversal.
+//
+// Order of mutations may be important, mutations are ordered according to the
+// following rules:
+//   Before/After insertions are ordered as registered when reverse_order ==
+//   false,
+//
+//   Before/After insertions are in reverse order as registered when
+//   reverse_order == true,
+//
+//   Before/After insertions are done before Expr replacements, so reference for
+//   insertions must be on pre-replaced Exprs
+//
+//   Removal of expressions is done after replacements.
+//
+// To place in a scope that is empty, simply provide a nullptr reference
+// Since insertions are done in order, it's possible to insert an expression in
+// an empty scope, and then use that inserted scope as a reference for
+// subsequent mutations.
+class ExprMutator : public IrVisitor {
+ public:
+  std::vector<Expr*> traverseAndInsert(
+      const std::vector<Expr*>& exprs,
+      bool reverse_order = false);
+
+ protected:
+  std::vector<Expr*> mutate(bool reverse_order = false);
+
+  using IrVisitor::handle;
+  // Registration function which *don't* need to be called "in place" during
+  // visiting.
+  void registerInsertBefore(Expr* reference, Expr* new_expr, Scope* scope);
+  void registerInsertAfter(Expr* reference, Expr* new_expr, Scope* scope);
+  void registerReplace(Expr* reference, Expr* new_expr, Scope* scope);
+  void registerRemove(Expr* expr_to_remove, Scope* scope);
+
+  // Registration function which need to be called "in place" during visiting.
+  // I.E.
+  // if you want to insert before/after or replace an Expr, you must register
+  // when in handle(Expr*) of that expr.
+  void registerInsertBefore(Expr* reference, Expr* new_expr);
+  void registerInsertAfter(Expr* reference, Expr* new_expr);
+  void registerReplace(Expr* reference, Expr* new_expr);
+  void registerRemove(Expr* expr_to_remove);
+
+ private:
+  enum class MutationMode { BEFORE, AFTER, REPLACE, REMOVE };
+
+  void registerMutation(
+      Expr* ref,
+      Expr* new_expr,
+      Scope* scope,
+      MutationMode mode);
+
+  struct MutationInformation {
+    Expr* reference = nullptr;
+    Expr* new_expr = nullptr;
+    Scope* scope = nullptr;
+    MutationMode mode = MutationMode::BEFORE;
+  };
+
+  // Track insertions as they're registered
+  std::vector<MutationInformation> insertions_;
+
+  // Track replacements as they're registered
+  std::vector<MutationInformation> replacements_;
+
+  // Track removal as they're registered
+  std::vector<MutationInformation> removal_;
+};
+
+} // namespace kir
+} // namespace nvfuser

nvfuser/include/nvfuser/linked_hash_map.h

@@ -0,0 +1,97 @@
+// 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 <list>
+#include <unordered_map>
+#include <utility>
+
+#include <exceptions.h>
+
+namespace nvfuser {
+
+// This mimics
+// https://docs.oracle.com/javase/8/docs/api/java/util/LinkedHashMap.html. The
+// implementation is a combination of a linked list of <K,V> and a hash map from
+// K to a position in the linked list.
+template <typename K, typename V>
+class LinkedHashMap {
+ public:
+  using value_type = std::pair<K, V>;
+  using const_iterator = typename std::list<value_type>::const_iterator;
+  using iterator = typename std::list<value_type>::iterator;
+
+  LinkedHashMap() = default;
+  LinkedHashMap(const LinkedHashMap&) = delete;
+  LinkedHashMap& operator=(const LinkedHashMap&) = delete;
+  LinkedHashMap(LinkedHashMap&&) = default;
+  LinkedHashMap& operator=(LinkedHashMap&&) = default;
+
+  // Returns the value associated with `key` and the list iterator following the
+  // removed element.
+  std::pair<V, iterator> erase(const K& key);
+
+  void insert(const_iterator i, const K& key, const V& value);
+  void insert(const_iterator i, const K& key, V&& value);
+
+  void pushBack(const K& key, const V& value);
+  void pushBack(const K& key, V&& value);
+
+  const_iterator begin() const {
+    return order_.begin();
+  }
+  const_iterator end() const {
+    return order_.end();
+  }
+
+ private:
+  std::list<value_type> order_;
+  std::unordered_map<K, iterator> key_to_index_;
+};
+
+template <typename K, typename V>
+std::pair<V, typename LinkedHashMap<K, V>::iterator> LinkedHashMap<K, V>::erase(
+    const K& key) {
+  iterator index = key_to_index_.at(key);
+  key_to_index_.erase(key);
+  V value = std::move(index->second);
+  return {std::move(value), order_.erase(index)};
+}
+
+template <typename K, typename V>
+void LinkedHashMap<K, V>::insert(
+    LinkedHashMap<K, V>::const_iterator i,
+    const K& key,
+    const V& value) {
+  auto j = order_.emplace(i, key, value);
+  bool inserted = key_to_index_.emplace(key, j).second;
+  NVF_CHECK(inserted, "Key already existed");
+}
+
+template <typename K, typename V>
+void LinkedHashMap<K, V>::insert(
+    LinkedHashMap<K, V>::const_iterator i,
+    const K& key,
+    V&& value) {
+  auto j = order_.emplace(i, key, std::move(value));
+  bool inserted = key_to_index_.emplace(key, j).second;
+  NVF_CHECK(inserted, "Key already existed");
+}
+
+template <typename K, typename V>
+void LinkedHashMap<K, V>::pushBack(const K& key, const V& value) {
+  insert(order_.end(), key, value);
+}
+
+template <typename K, typename V>
+void LinkedHashMap<K, V>::pushBack(const K& key, V&& value) {
+  insert(order_.end(), key, std::move(value));
+}
+
+} // namespace nvfuser