From 38bddf2d257d8d10b2e9107f32e46e6e8dc0ba35 Mon Sep 17 00:00:00 2001 From: Michael Davis Date: Fri, 30 Jan 2026 18:37:11 -0800 Subject: [PATCH] rename storage.h/cpp -> container.h/cpp --- CMakeLists.txt | 3 +- csrc/fusion.h | 2 +- csrc/ir/builder.h | 2 +- csrc/ir/{storage.cpp => container.cpp} | 0 csrc/ir/container.h | 191 ++++++++++++++++++++++- csrc/ir/storage.h | 203 ------------------------- 6 files changed, 193 insertions(+), 208 deletions(-) rename csrc/ir/{storage.cpp => container.cpp} (100%) delete mode 100644 csrc/ir/storage.h diff --git a/CMakeLists.txt b/CMakeLists.txt index d6289a29bb7..d21425d2e9e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -215,8 +215,7 @@ list(APPEND NVFUSER_SRCS ${NVFUSER_SRCS_DIR}/ir/base_nodes.cpp ${NVFUSER_SRCS_DIR}/ir/builder.cpp ${NVFUSER_SRCS_DIR}/ir/cloner.cpp - #${NVFUSER_SRCS_DIR}/ir/container.cpp - ${NVFUSER_SRCS_DIR}/ir/storage.cpp + ${NVFUSER_SRCS_DIR}/ir/container.cpp ${NVFUSER_SRCS_DIR}/ir/graphviz.cpp ${NVFUSER_SRCS_DIR}/ir/iostream.cpp ${NVFUSER_SRCS_DIR}/ir/internal_base_nodes.cpp diff --git a/csrc/fusion.h b/csrc/fusion.h index 4e7ce658574..7873b852d63 100644 --- a/csrc/fusion.h +++ b/csrc/fusion.h @@ -21,7 +21,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/csrc/ir/builder.h b/csrc/ir/builder.h index 80e7828067c..af0fea66d32 100644 --- a/csrc/ir/builder.h +++ b/csrc/ir/builder.h @@ -11,7 +11,7 @@ #include "exceptions.h" #include "fusion_guard.h" #include "ir/builder_passkey.h" -#include "ir/storage.h" +#include "ir/container.h" #include "visibility.h" namespace nvfuser { diff --git a/csrc/ir/storage.cpp b/csrc/ir/container.cpp similarity index 100% rename from csrc/ir/storage.cpp rename to csrc/ir/container.cpp diff --git a/csrc/ir/container.h b/csrc/ir/container.h index d6a53fb1e3a..e361b8743ee 100644 --- a/csrc/ir/container.h +++ b/csrc/ir/container.h @@ -7,8 +7,197 @@ // clang-format on #pragma once +#include +#include +#include + +#include "base.h" +#include "exceptions.h" +#include "ir/base_nodes.h" +#include "visibility.h" + namespace nvfuser { -// Empty for now... +// Passkey for container to register names with statements +class IrContainerPasskey { + friend class IrContainer; + + private: + explicit IrContainerPasskey() = default; +}; + +class NamedScalar; + +class IrContainer { + public: + NVF_API IrContainer(); + + // Copy/Move Constructors and Operators are deleted. IrContainer is managed + // through a smart pointer in IrContainer. Semantic operations for Fusion + // types are handled directly through copy and swap functions. + IrContainer(const IrContainer& other) = delete; + IrContainer(IrContainer&& other) noexcept = delete; + + IrContainer& operator=(const IrContainer& other) = delete; + IrContainer& operator=(IrContainer&& other) noexcept = delete; + + ~IrContainer(); + + bool inContainer(const Statement* stmt) const; + + void assertInContainer(const Statement* stmt, const std::string& msg) const { + NVF_CHECK( + inContainer(stmt), msg, " it was not found in the active container."); + } + + //! Return values in insertion order + const std::deque deterministic_vals() const noexcept; + + //! Return expression in insertion order + const std::deque deterministic_exprs() const noexcept; + + //! Return mapping from value to integer id + const std::unordered_map deterministic_vals_map() + const noexcept; + + //! Return mapping from expression to integer id + const std::unordered_map deterministic_exprs_map() + const noexcept; + + //! Return the set of Exprs registered with this fusion. Warning: This will + //! return exprs outside inputs/outputs, so can be unsafe for use with + //! segmented fusions. + const std::unordered_set& unordered_exprs() const noexcept { + return exprs_; + } + + //! Return the set of Vals registered with this fusion + const std::unordered_set& vals() const noexcept { + return vals_; + } + + int64_t numExprs() const noexcept { + return std::ssize(exprs_); + } + + // When include_shortcuts is true, it will count the shortcuts like true_val_. + int64_t numVals(bool include_shortcuts) const noexcept { + return include_shortcuts ? std::ssize(vals_) : std::ssize(vals_up_); + } + + // Shortcuts for frequently used vals + NVF_API Val* zeroVal(); + NVF_API Val* oneVal(); + Val* falseVal(); + Val* trueVal(); + NamedScalar* magicZeroVal(); + NVF_API Val* zeroVal(DataType dtype); + NVF_API Val* oneVal(DataType dtype); + Val* metadataOf(Val*); + + // Axioms about CUDA programming, for example: threadIdx.x < blockDim.x + const std::vector& axioms() { + lazyInitAxioms(); + return *axioms_; + } + + void assumePositive(Val* val); + void assumeNonNegative(Val* val); + + protected: + static IrCloner copy(const IrContainer* from, IrContainer* to); + + static void swap(IrContainer& a, IrContainer& b) noexcept; + + // Let Fusion access IrContainer::clear() + friend class Fusion; + + void removeExpr(Expr* expr); + + //! Completely remove val from the fusion, break all dependencies associated + //! with it + void removeVal(Val* val); + + //! Register the Val with this container + NVF_API void registerVal(Val* val); + + //! Register expr with this container. + NVF_API void registerExpr(Expr* expr); + + StmtNameType getValName(ValType vtype) { + if (val_type_name_map_.find(vtype) == val_type_name_map_.end()) { + val_type_name_map_[vtype] = 0; + } + return val_type_name_map_[vtype]++; + } + + StmtNameType getExprName() { + return expr_name_counter_++; + } + + void clear() noexcept; + + void lazyInitAxioms(); + + friend class StatementGuard; + + // A simple garbage collection mechanism to remove all Exprs and Vals that + // were created after a certain point. This is useful for analysis that + // creates new Exprs and Vals in the container and wants to clean up after + // itself. + // + // Used by StatementGuard only. + void removeStatementsCreatedAfter( + int64_t prev_num_exprs, + int64_t prev_num_vals); + + // Deque of unique pointer is the memory owning data structure + std::deque> vals_up_; + + // A convenient set to return when we just need an unordered set to do + // something like check if a Val is in this container + std::unordered_set vals_; + + // Deque of unique pointer is the memory owning data structure + std::deque> exprs_up_; + + // A convenient set to return when we just need an unordered set to do + // something like check if an Expr is in this container + std::unordered_set exprs_; + + // Values names counters + std::unordered_map val_type_name_map_; + + // Expression names counter + StmtNameType expr_name_counter_ = 0; + + // Manually store some persistent, frequently used nodes. It's very + // challenging to do this anything but manually as detecting when a container + // may or may not have one of these vals is tricky. Specifically because if + // the container doesn't own it, it's hard to understand from the outside if + // the node may have been removed then re-registered. It could also be tricky + // to know when we're using a different container as in FusionCopy_test + // demonstrates deleting then creating containers can result in the same + // pointer for the container. + std::unique_ptr true_val_; + std::unique_ptr false_val_; + std::unique_ptr one_val_; + std::unique_ptr zero_val_; + std::unique_ptr magic_zero_val_; + std::unique_ptr> axioms_; + std::unordered_map> metadata_; + + public: + Fusion* parent() const { + NVF_ERROR( + parent_ != nullptr, "Call to IrContainer::parent() holds nullptr.") + return parent_; + } + + private: + // Parent Fusion that owns this container (for pure composition pattern) + // Used by Statement::fusion() to navigate back to owning Fusion + Fusion* parent_ = nullptr; +}; } // namespace nvfuser diff --git a/csrc/ir/storage.h b/csrc/ir/storage.h deleted file mode 100644 index e361b8743ee..00000000000 --- a/csrc/ir/storage.h +++ /dev/null @@ -1,203 +0,0 @@ -// 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 -#include -#include - -#include "base.h" -#include "exceptions.h" -#include "ir/base_nodes.h" -#include "visibility.h" - -namespace nvfuser { - -// Passkey for container to register names with statements -class IrContainerPasskey { - friend class IrContainer; - - private: - explicit IrContainerPasskey() = default; -}; - -class NamedScalar; - -class IrContainer { - public: - NVF_API IrContainer(); - - // Copy/Move Constructors and Operators are deleted. IrContainer is managed - // through a smart pointer in IrContainer. Semantic operations for Fusion - // types are handled directly through copy and swap functions. - IrContainer(const IrContainer& other) = delete; - IrContainer(IrContainer&& other) noexcept = delete; - - IrContainer& operator=(const IrContainer& other) = delete; - IrContainer& operator=(IrContainer&& other) noexcept = delete; - - ~IrContainer(); - - bool inContainer(const Statement* stmt) const; - - void assertInContainer(const Statement* stmt, const std::string& msg) const { - NVF_CHECK( - inContainer(stmt), msg, " it was not found in the active container."); - } - - //! Return values in insertion order - const std::deque deterministic_vals() const noexcept; - - //! Return expression in insertion order - const std::deque deterministic_exprs() const noexcept; - - //! Return mapping from value to integer id - const std::unordered_map deterministic_vals_map() - const noexcept; - - //! Return mapping from expression to integer id - const std::unordered_map deterministic_exprs_map() - const noexcept; - - //! Return the set of Exprs registered with this fusion. Warning: This will - //! return exprs outside inputs/outputs, so can be unsafe for use with - //! segmented fusions. - const std::unordered_set& unordered_exprs() const noexcept { - return exprs_; - } - - //! Return the set of Vals registered with this fusion - const std::unordered_set& vals() const noexcept { - return vals_; - } - - int64_t numExprs() const noexcept { - return std::ssize(exprs_); - } - - // When include_shortcuts is true, it will count the shortcuts like true_val_. - int64_t numVals(bool include_shortcuts) const noexcept { - return include_shortcuts ? std::ssize(vals_) : std::ssize(vals_up_); - } - - // Shortcuts for frequently used vals - NVF_API Val* zeroVal(); - NVF_API Val* oneVal(); - Val* falseVal(); - Val* trueVal(); - NamedScalar* magicZeroVal(); - NVF_API Val* zeroVal(DataType dtype); - NVF_API Val* oneVal(DataType dtype); - Val* metadataOf(Val*); - - // Axioms about CUDA programming, for example: threadIdx.x < blockDim.x - const std::vector& axioms() { - lazyInitAxioms(); - return *axioms_; - } - - void assumePositive(Val* val); - void assumeNonNegative(Val* val); - - protected: - static IrCloner copy(const IrContainer* from, IrContainer* to); - - static void swap(IrContainer& a, IrContainer& b) noexcept; - - // Let Fusion access IrContainer::clear() - friend class Fusion; - - void removeExpr(Expr* expr); - - //! Completely remove val from the fusion, break all dependencies associated - //! with it - void removeVal(Val* val); - - //! Register the Val with this container - NVF_API void registerVal(Val* val); - - //! Register expr with this container. - NVF_API void registerExpr(Expr* expr); - - StmtNameType getValName(ValType vtype) { - if (val_type_name_map_.find(vtype) == val_type_name_map_.end()) { - val_type_name_map_[vtype] = 0; - } - return val_type_name_map_[vtype]++; - } - - StmtNameType getExprName() { - return expr_name_counter_++; - } - - void clear() noexcept; - - void lazyInitAxioms(); - - friend class StatementGuard; - - // A simple garbage collection mechanism to remove all Exprs and Vals that - // were created after a certain point. This is useful for analysis that - // creates new Exprs and Vals in the container and wants to clean up after - // itself. - // - // Used by StatementGuard only. - void removeStatementsCreatedAfter( - int64_t prev_num_exprs, - int64_t prev_num_vals); - - // Deque of unique pointer is the memory owning data structure - std::deque> vals_up_; - - // A convenient set to return when we just need an unordered set to do - // something like check if a Val is in this container - std::unordered_set vals_; - - // Deque of unique pointer is the memory owning data structure - std::deque> exprs_up_; - - // A convenient set to return when we just need an unordered set to do - // something like check if an Expr is in this container - std::unordered_set exprs_; - - // Values names counters - std::unordered_map val_type_name_map_; - - // Expression names counter - StmtNameType expr_name_counter_ = 0; - - // Manually store some persistent, frequently used nodes. It's very - // challenging to do this anything but manually as detecting when a container - // may or may not have one of these vals is tricky. Specifically because if - // the container doesn't own it, it's hard to understand from the outside if - // the node may have been removed then re-registered. It could also be tricky - // to know when we're using a different container as in FusionCopy_test - // demonstrates deleting then creating containers can result in the same - // pointer for the container. - std::unique_ptr true_val_; - std::unique_ptr false_val_; - std::unique_ptr one_val_; - std::unique_ptr zero_val_; - std::unique_ptr magic_zero_val_; - std::unique_ptr> axioms_; - std::unordered_map> metadata_; - - public: - Fusion* parent() const { - NVF_ERROR( - parent_ != nullptr, "Call to IrContainer::parent() holds nullptr.") - return parent_; - } - - private: - // Parent Fusion that owns this container (for pure composition pattern) - // Used by Statement::fusion() to navigate back to owning Fusion - Fusion* parent_ = nullptr; -}; - -} // namespace nvfuser