From 1b5d7c2047cfa60ffca9b7c5a29dce995e6c43e8 Mon Sep 17 00:00:00 2001 From: ted Date: Fri, 29 Sep 2023 20:14:21 +0800 Subject: [PATCH 1/2] [chore] Add some comments --- .../frontend/optimizer/graph_transform.cc | 40 +++++--- .../frontend/optimizer/py_pass_manager.cc | 93 +++++++++++++++++-- 2 files changed, 114 insertions(+), 19 deletions(-) diff --git a/mindspore/ccsrc/frontend/optimizer/graph_transform.cc b/mindspore/ccsrc/frontend/optimizer/graph_transform.cc index d9a56141677..87b57e8210a 100644 --- a/mindspore/ccsrc/frontend/optimizer/graph_transform.cc +++ b/mindspore/ccsrc/frontend/optimizer/graph_transform.cc @@ -15,10 +15,24 @@ */ #include "frontend/optimizer/graph_transform.h" + +// Include the header file for the vector class +#include + +// Include the header file for the algorithm library +#include + +// Include the header file for the GraphUtils class +#include "ir/graph_utils.h"#include "frontend/optimizer/graph_transform.h" #include #include #include "ir/graph_utils.h" +// This function checks if a given FuncGraph has any tuple inputs. +// It takes a FuncGraphPtr as input and returns a boolean value. +// It uses a lambda function is_tuple to check if a parameter is a tuple. +// It uses std::any_of to iterate over the parameters of the FuncGraph and check if any of them are tuples. +// If any tuple parameter is found, it returns true. Otherwise, it returns false. namespace mindspore { /* namespace to support opt */ namespace opt { @@ -31,22 +45,22 @@ bool FuncGraphHasTupleInput(const FuncGraphPtr &fg) { std::vector TransformTupleArgument(const FuncGraphPtr &fg, const AnfNodePtr &node, const abstract::AbstractTuplePtr &abs) { - auto &elements = abs->elements(); - std::vector tuple_node_expanded; - for (size_t i = 0; i < elements.size(); i++) { - auto idx = NewValueNode(SizeToLong(i)); - auto abstract_scalar = std::make_shared(std::make_shared(SizeToLong(i))); - idx->set_abstract(abstract_scalar); - auto elem_node = fg->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx}); - elem_node->set_abstract(elements[i]); - if (elements[i]->isa()) { - auto nodes = TransformTupleArgument(fg, elem_node, elements[i]->cast()); - tuple_node_expanded.insert(tuple_node_expanded.end(), nodes.begin(), nodes.end()); + auto &elements = abs->elements(); // Get the elements of the abstract tuple + std::vector tuple_node_expanded; // Create a vector to store the expanded tuple nodes + for (size_t i = 0; i < elements.size(); i++) { // Iterate over the elements of the abstract tuple + auto idx = NewValueNode(SizeToLong(i)); // Create a new value node with the index + auto abstract_scalar = std::make_shared(std::make_shared(SizeToLong(i))); // Create an abstract scalar with the index + idx->set_abstract(abstract_scalar); // Set the abstract scalar as the abstract value of the index node + auto elem_node = fg->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx}); // Create a new CNode with the tuple get item primitive, the input node, and the index node + elem_node->set_abstract(elements[i]); // Set the abstract value of the element node + if (elements[i]->isa()) { // Check if the element is an abstract tuple + auto nodes = TransformTupleArgument(fg, elem_node, elements[i]->cast()); // Recursively transform the tuple argument + tuple_node_expanded.insert(tuple_node_expanded.end(), nodes.begin(), nodes.end()); // Add the transformed nodes to the expanded tuple nodes vector } else { - tuple_node_expanded.push_back(elem_node); + tuple_node_expanded.push_back(elem_node); // Add the element node to the expanded tuple nodes vector } } - return tuple_node_expanded; + return tuple_node_expanded; // Return the expanded tuple nodes } } // namespace opt } // namespace mindspore diff --git a/mindspore/ccsrc/frontend/optimizer/py_pass_manager.cc b/mindspore/ccsrc/frontend/optimizer/py_pass_manager.cc index 82fb11a9b9e..44fccb240d5 100644 --- a/mindspore/ccsrc/frontend/optimizer/py_pass_manager.cc +++ b/mindspore/ccsrc/frontend/optimizer/py_pass_manager.cc @@ -15,97 +15,178 @@ */ #include "frontend/optimizer/py_pass_manager.h" +// Include the functional header, which provides various function objects and related utilities #include + +// Include the utility header, which provides various utility components #include +// Include the header file for the IR manager module #include "ir/manager.h" + +// Include the header file for the optimizer pass group module #include "frontend/optimizer/pass_group.h" +// Define the namespace "mindspore" for the code namespace mindspore { namespace opt { namespace python_pass { PyPassManagerPtr PyPassManager::global_instance = nullptr; mindspore::HashMap PyPassManager::phase_to_group_; +// Function to get the pass group associated with a given phase in the PyPassManager class PassGroupPtr PyPassManager::GetPassGroup(Phase phase) { + // Find the pass group corresponding to the given phase in the phase_to_group_ map auto pm = phase_to_group_.find(phase); + + // If the pass group is not found (i.e., iterator points to the end of the map) if (pm == phase_to_group_.end()) { + // Return nullptr to indicate that the pass group does not exist return nullptr; } + + // If the pass group is found, return the pointer to the pass group return pm->second; } +// Define a function named "GetInstance" that returns a pointer to a PyPassManager object PyPassManagerPtr PyPassManager::GetInstance() { + // Check if the global_instance is nullptr (i.e., not yet initialized) if (global_instance == nullptr) { + // If global_instance is nullptr, create a new PyPassManager object using std::nothrow + // and assign it to global_instance as a shared_ptr global_instance = std::shared_ptr(new (std::nothrow) PyPassManager()); } + // Return the global_instance pointer return global_instance; } +// Constructor for the PyPassManager class PyPassManager::PyPassManager() { + // Create a shared pointer to a PassGroup object and assign it to the PREAD phase in the phase_to_group_ map phase_to_group_[Phase::PREAD] = std::make_shared("Pre_AD_PassGroup"); + + // Create a shared pointer to a PassGroup object and assign it to the OPT phase in the phase_to_group_ map phase_to_group_[Phase::OPT] = std::make_shared("After_OPT_PassGroup"); + + // Create a shared pointer to a MatchResult object and assign it to the res_ member variable res_ = std::make_shared(); } +// Define a function named "Register" in the class "PyPassManager" void PyPassManager::Register(const std::string &pass_name, const PatternPtr &pattern, const PatternPtr &target, bool requires_grad, bool run_only_once) { + // Declare a pointer to a PassGroup object named "cur_pg" PassGroupPtr cur_pg; + + // Check if the pass requires gradient if (requires_grad) { + // If it requires gradient, get the PassGroup for the PREAD phase cur_pg = GetPassGroup(Phase::PREAD); } else { + // If it doesn't require gradient, get the PassGroup for the OPT phase cur_pg = GetPassGroup(Phase::OPT); } + + // Check if the PassGroup pointer is null MS_EXCEPTION_IF_NULL(cur_pg); + + // Set the "run_only_once" flag for the current PassGroup cur_pg->SetRunOnlyOnce(run_only_once); + + // Check if the pattern and target pointers are null MS_EXCEPTION_IF_NULL(pattern); MS_EXCEPTION_IF_NULL(target); + + // Check if the PassGroup pointer is null again MS_EXCEPTION_IF_NULL(cur_pg); + + // Create a new PythonPass object using the provided pass_name, pattern, target, and run_only_once values PythonPassPtr new_pass = std::make_shared(pass_name, pattern, target, run_only_once); + + // Add the new_pass to the current PassGroup cur_pg->AddPass(new_pass); } +// Definition of the Unregister function in the PyPassManager class + void PyPassManager::Unregister(const std::string &pass_name) { + // Get the pass group for the OPT phase auto opt_pm = GetPassGroup(Phase::OPT); + + // Check if the pass with the given name exists in the OPT pass group if (!opt_pm->DeletePass(pass_name)) { + // If the pass does not exist, log a warning message MS_LOG(WARNING) << "Opt has no such pass : " + pass_name + "\n"; } + + // Get the pass group for the PREAD phase auto pre_ad_pm = GetPassGroup(Phase::PREAD); + + // Check if the pass with the given name exists in the PREAD pass group if (!pre_ad_pm->DeletePass(pass_name)) { + // If the pass does not exist, log a warning message MS_LOG(WARNING) << "Pre_AD has no such pass : " + pass_name + "\n"; } } +// Define a function named "GenNewParameter" that takes a reference to a PatternPtr object named "parameter" as a parameter void PyPassManager::GenNewParameter(const PatternPtr ¶meter) { + // Check if the "parameter" object is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(parameter); // NOTE: Add NewParameter at early stage will cause CSE problems auto cur_pg = GetPassGroup(Phase::OPT); + // Check if the "cur_pg" object is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(cur_pg); + + // Set the "run only once" flag of the pass group to true cur_pg->SetRunOnlyOnce(true); + + // Cast the "parameter" object to a NewParameterPtr object and assign it to "new_para_pattern" auto new_para_pattern = parameter->cast(); + // Check if the "new_para_pattern" object is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(new_para_pattern); + + // Get the name of the parameter and assign it to "pass_name" auto pass_name = new_para_pattern->para_name(); + + // Set the "last" flag of the "new_para_pattern" object to true new_para_pattern->set_last(true); + + // Create a new PythonPass object named "new_pass" using the pass name, nullptr for the pass function, the "parameter" object, and true for the "is_new_parameter" flag auto new_pass = std::make_shared(pass_name, nullptr, parameter, true); + + // Add the "new_pass" object to the pass group cur_pg->AddPass(new_pass); } +// Definition of the ClearRes() function in the PyPassManager class + void PyPassManager::ClearRes() { + // Log an informational message using the MS_LOG macro, indicating that PyPassManager resources are being cleared MS_LOG(INFO) << "Clear PyPassManager resources!"; + + // Set the global_instance pointer to nullptr, effectively clearing the instance of PyPassManager global_instance = nullptr; + + // Clear the phase_to_group_ map, removing all elements from it phase_to_group_.clear(); } +// Register the PyPassManager_ class and its methods in the Python module REGISTER_PYBIND_DEFINE( PyPassManager_, ([](const py::module *m) { + // Define an enum for the Phase class with values "pre_ad" and "opt" (void)py::enum_(*m, "phase", py::arithmetic()).value("pre_ad", Phase::PREAD).value("opt", Phase::OPT); + + // Define the PyPassManager_ class and its methods (void)py::class_>(*m, "PyPassManager_") - .def(py::init([]() { return PyPassManager::GetInstance(); })) - .def("register", &PyPassManager::Register, "Register python pass") - .def("unregister", &PyPassManager::Unregister, "Unregister Python Pass") - .def("gen_new_parameter", &PyPassManager::GenNewParameter, "Generate new parameter") - .def("set_renorm", &PyPassManager::SetRenorm, "Set whether or not to do renorm after modified graph") - .def("set_reopt", &PyPassManager::SetReOpt, "Set whether or not to do optimization after modified graph"); + .def(py::init([]() { return PyPassManager::GetInstance(); })) // Define the constructor for PyPassManager_ + .def("register", &PyPassManager::Register, "Register python pass") // Define the register method + .def("unregister", &PyPassManager::Unregister, "Unregister Python Pass") // Define the unregister method + .def("gen_new_parameter", &PyPassManager::GenNewParameter, "Generate new parameter") // Define the gen_new_parameter method + .def("set_renorm", &PyPassManager::SetRenorm, "Set whether or not to do renorm after modified graph") // Define the set_renorm method + .def("set_reopt", &PyPassManager::SetReOpt, "Set whether or not to do optimization after modified graph"); // Define the set_reopt method })); } // namespace python_pass } // namespace opt -- 2.34.1 From db5c54492f54e7b251853f13566d6105651416fc Mon Sep 17 00:00:00 2001 From: vilc Date: Wed, 4 Oct 2023 10:25:23 +0800 Subject: [PATCH 2/2] [chore] Add some comments --- .../ccsrc/frontend/optimizer/comm_op_attrs.cc | 121 +++- mindspore/ccsrc/frontend/optimizer/opt.cc | 559 +++++++++++++--- .../ccsrc/frontend/optimizer/pass_group.cc | 66 +- mindspore/ccsrc/frontend/optimizer/pattern.cc | 198 +++++- .../ccsrc/frontend/optimizer/recompute.cc | 612 +++++++++++++++--- 5 files changed, 1324 insertions(+), 232 deletions(-) diff --git a/mindspore/ccsrc/frontend/optimizer/comm_op_attrs.cc b/mindspore/ccsrc/frontend/optimizer/comm_op_attrs.cc index 42efe923d23..7737a3b1fa8 100644 --- a/mindspore/ccsrc/frontend/optimizer/comm_op_attrs.cc +++ b/mindspore/ccsrc/frontend/optimizer/comm_op_attrs.cc @@ -14,48 +14,115 @@ * limitations under the License. */ +// Include the header file "frontend/optimizer/comm_op_attrs.h" which contains the declarations for communication operator attributes #include "frontend/optimizer/comm_op_attrs.h" -#include -#include -#include -#include -#include "frontend/parallel/ops_info/ops_utils.h" -#include "frontend/parallel/device_manager.h" -#include "include/common/utils/anfalgo.h" +// Include the header for smart pointers +#include + +// Include the header for using vectors +#include + +// Include the header for using strings +#include + +// Include the header for using algorithms like sorting +#include + +// Include the header file "frontend/parallel/ops_info/ops_utils.h" which contains utility functions for parallel operations +#include "frontend/parallel/ops_info/ops_utils.h" + +// Include the header file "frontend/parallel/device_manager.h" which contains the declarations for device management +#include "frontend/parallel/device_manager.h" + +// Include the header file "include/common/utils/anfalgo.h" which contains utility functions for ANF (Abstract Neural Network) algorithms + +// Start of the `mindspore` namespace namespace mindspore { -namespace opt { -void CommOpAttrs(const FuncGraphPtr &graph) { - if (parallel::g_device_manager == nullptr) { - return; - } - MS_EXCEPTION_IF_NULL(graph); - AnfNodePtr return_node = graph->get_return(); - MS_EXCEPTION_IF_NULL(return_node); - std::vector all_nodes = TopoSort(return_node); - for (auto &node : all_nodes) { - MS_EXCEPTION_IF_NULL(node); - if (!node->isa()) { - continue; - } - auto primitive = GetCNodePrimitive(node); - if (primitive == nullptr) { - continue; - } - if (!common::AnfAlgo::IsCommunicationOp(node)) { + // Start of the `opt` namespace + namespace opt { + + // Function to set communication operation attributes in the given function graph + void CommOpAttrs(const FuncGraphPtr &graph) { + // Check if the device manager is initialized + if (parallel::g_device_manager == nullptr) { + // If not initialized, return from the function + return; + } + + // Check if the graph is null + MS_EXCEPTION_IF_NULL(graph); + + // Get the return node of the graph + AnfNodePtr return_node = graph->get_return(); + + // Check if the return node is null + MS_EXCEPTION_IF_NULL(return_node); + + // Perform topological sorting on all nodes in the graph + std::vector all_nodes = TopoSort(return_node); + + // Iterate over all nodes in the graph + for (auto &node : all_nodes) { + // Check if the node is null + MS_EXCEPTION_IF_NULL(node); + + // Check if the node is a CNode + if (!node->isa()) { + // If not a CNode, continue to the next node + continue; + } + + // Get the primitive of the CNode + auto primitive = GetCNodePrimitive(node); + + // Check if the primitive is null + if (primitive == nullptr) { + // If primitive is null, continue to the next node + continue; + } + + // Check if the node is a communication operation + if (!common::AnfAlgo::IsCommunicationOp(node)) { + // If not a communication operation, continue to the next node + continue; + } + + // Rest of the code for handling communication operation attributes + // ... + } + } + + } // End of the `opt` namespace +} // End of the `mindspore` namespace continue; } + + // Get the primitive of the CNode auto comm_prim = common::AnfAlgo::GetCNodePrimitive(node); + + // Initialize an empty string for the group name std::string group_name = ""; + + // Check if the primitive has the attribute "group" if (comm_prim->HasAttr(parallel::GROUP)) { + // If it does, get the value of the "group" attribute and assign it to the group_name variable group_name = GetValue(comm_prim->GetAttr(parallel::GROUP)); } + + // Initialize an empty vector to store the rank list std::vector rank_list = {}; + + // Find the rank list by hash name using the device manager auto long_rank_list = parallel::g_device_manager->FindRankListByHashName(group_name); + + // Convert the long_rank_list to unsigned int and store it in the rank_list vector (void)std::transform(long_rank_list.begin(), long_rank_list.end(), std::back_inserter(rank_list), [](int64_t d) -> unsigned int { return IntToUint(LongToInt(d)); }); + + // Add the rank_list as an attribute to the primitive (void)comm_prim->AddAttr(kAttrGroupRankIds, MakeValue>(rank_list)); } } } // namespace opt -} // namespace mindspore +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/frontend/optimizer/opt.cc b/mindspore/ccsrc/frontend/optimizer/opt.cc index 4d4f16e4b9e..56b80dc6aa4 100644 --- a/mindspore/ccsrc/frontend/optimizer/opt.cc +++ b/mindspore/ccsrc/frontend/optimizer/opt.cc @@ -14,383 +14,752 @@ * limitations under the License. */ +// Include the header file "opt.h" from the "optimizer" directory in the "frontend" module + #include "frontend/optimizer/opt.h" +// Include the deque header, which provides a double-ended queue container #include + +// Include the memory header, which provides smart pointers and memory management utilities #include + +// Include the algorithm header, which provides a collection of algorithms #include + +// Include the utility header, which provides various utility components #include +// Include the custom header file "utils/hash_map.h" for using hash maps #include "utils/hash_map.h" + +// Include the custom header file "ir/anf.h" for working with abstract syntax trees #include "ir/anf.h" + +// Include the custom header file "ir/manager.h" for managing intermediate representations #include "ir/manager.h" + +// Include the custom header file "frontend/optimizer/optimizer.h" for performing optimizations on the code #include "frontend/optimizer/optimizer.h" + +// Include the custom header file "utils/log_adapter.h" for logging and debugging purposes #include "utils/log_adapter.h" +// Start of the "mindspore" namespace namespace mindspore { -/* namespace to support opt */ + +// Namespace "opt" within the "mindspore" namespace, used to support optimization namespace opt { + +// Function to create a Substitution object SubstitutionPtr MakeSubstitution(const OptimizerCallerPtr &transform, const std::string &name, const PrimitivePtr &prim, const RenormAction &renorm_action, bool has_priority_pattern) { + // Lambda function that checks if a given node is a CNode with the specified primitive auto fn = [prim](const AnfNodePtr &node) -> bool { return IsPrimitiveCNode(node, prim); }; + + // Create and return a Substitution object with the provided parameters return std::make_shared(transform, name, fn, renorm_action, has_priority_pattern); } -SubstitutionPtr MakeSubstitution(const OptimizerCallerPtr &transform, const std::string &name, - const std::vector &prims, const RenormAction &renorm_action, - bool has_priority_pattern) { - auto fn = [prims](const AnfNodePtr &node) -> bool { - if (!node->isa()) { - return false; - } +// Create a function named "MakeSubstitution" that takes in the following parameters: +// - A constant reference to an object of type "OptimizerCallerPtr" named "transform" +// - A constant reference to a string named "name" +// - A constant reference to a vector of "PrimitivePtr" objects named "prims" +// - A constant reference to an object of type "RenormAction" named "renorm_action" +// - A boolean variable named "has_priority_pattern" +// Create a lambda function named "fn" that takes in a constant reference to an object of type "AnfNodePtr" named "node" +// and returns a boolean value +auto fn = [prims](const AnfNodePtr &node) -> bool { + + // Check if the given node is not of type "CNode" + if (!node->isa()) { + return false; + } + + // Cast the input node to a CNodePtr using the auto keyword for type inference auto cnode = node->cast(); + + // Get the first input of the CNode auto inp0 = cnode->input(0); + + // Get the value node of type PrimitivePtr from inp0 using the GetValueNode function template auto prim0 = GetValueNode(inp0); + + // Check if prim0 is nullptr if (prim0 == nullptr) { + // If prim0 is nullptr, return false return false; } + // Calculate the hash value of prim0 auto hash = prim0->Hash(); + + // Get the name of prim0 auto const &name = prim0->name(); + + // Iterate through each prim in the prims vector for (auto &prim : prims) { + + // Check if the hash value and name of prim0 match with the current prim if (hash == prim->Hash() && name == prim->name()) { + + // If there is a match, return true return true; } } + + // If no match is found, return false return false; }; - return std::make_shared(transform, name, fn, renorm_action, has_priority_pattern); -} +// Return a shared pointer to a Substitution object created using the make_shared function +// The Substitution object is constructed with the provided arguments: transform, name, fn, renorm_action, has_priority_pattern +return std::make_shared(transform, name, fn, renorm_action, has_priority_pattern); + +// Create a function named "MakeSubstitution" that takes the following parameters: +// - A constant reference to an object of type "OptimizerCallerPtr" named "transform" +// - A constant reference to a string named "name" +// - A constant reference to an object of type "PredicateFuncType" named "predicate" +// - A constant reference to an object of type "RenormAction" named "renorm_action" +// - A boolean variable named "has_priority_pattern" + +// The function returns a shared pointer to an object of type "Substitution" SubstitutionPtr MakeSubstitution(const OptimizerCallerPtr &transform, const std::string &name, const PredicateFuncType &predicate, const RenormAction &renorm_action, bool has_priority_pattern) { + + // Use the std::make_shared function to create a shared pointer to a new object of type "Substitution" + // Pass the provided parameters to the constructor of "Substitution" return std::make_shared(transform, name, predicate, renorm_action, has_priority_pattern); } -AnfNodePtr Substitution::operator()(const OptimizerPtr &optimizer, const AnfNodePtr &node) { #ifdef ENABLE_PROFILE + // If profiling is enabled, record the current time double t = GetTime(); #endif + + // Apply the transformation function to the given node and store the result AnfNodePtr result = (*transform_)(optimizer, node); + #ifdef ENABLE_PROFILE + // If an optimizer is provided and profiling is enabled if (optimizer != nullptr) { + // Calculate the time taken for the transformation auto time = GetTime(); + // Record the time taken for the substitution step MsProfile::StatTime("substitution." + name_, time - t); + // If a result is obtained from the transformation, record the time taken for the matching step if (result != nullptr) { MsProfile::StatTime("match." + name_, time - t); } } #endif + + // If an optimizer is provided and watch renormalization is enabled, and a result is obtained if (optimizer != nullptr && optimizer->is_watch_renormalize() && result != nullptr) { + // Check if renormalization should be forced or if the result has no abstract value if ((renorm_action_ == FORCE_RENORM) || (result->abstract() == nullptr)) { + // Set the optimizer's flag to indicate that untyped code has been generated optimizer->set_is_untyped_generated(); } } - return result; +// Return the value of the variable "result" to indicate the result of the function +return result; } +// A static function that checks if a given AnfNodePtr is traversable static bool isTraversable(const AnfNodePtr &node) { + + // If the node is nullptr, it is not traversable, so return false if (node == nullptr) { return false; } + + // If the node is of type CNode or Parameter, it is traversable, so return true if (node->isa() || node->isa()) { return true; } + + // If the node is a value node of type FuncGraph or RefKey, it is traversable, so return true if (IsValueNode(node) || IsValueNode(node)) { return true; } + + // If none of the above conditions are met, the node is not traversable, so return false return false; } +// Perform a transformation on the given node using the provided optimizer and substitution static AnfNodePtr DoTransform(const OptimizerPtr &optimizer, const AnfNodePtr &node, const SubstitutionPtr &substitution) { - auto manager = optimizer->manager(); - bool is_match = substitution->predicate_(node); - if (is_match) { - TraceGuard trace_guard(std::make_shared(node->debug_info())); - ScopeGuard scope_guard(node->scope()); - auto res = (*substitution)(optimizer, node); - if (res != nullptr && res != node) { + + // Get the manager from the optimizer + auto manager = optimizer->manager(); + + // Check if the node matches the substitution's predicate + bool is_match = substitution->predicate_(node); + + // If the node matches the predicate + if (is_match) { + + // Create a trace guard to capture debug information + TraceGuard trace_guard(std::make_shared(node->debug_info())); + + // Create a scope guard to set the scope for the transformation + ScopeGuard scope_guard(node->scope()); + + // Apply the substitution to the node + auto res = (*substitution)(optimizer, node); + + // If the substitution result is not null and different from the original node + if (res != nullptr && res != node) { + #ifdef ENABLE_PROFILE - double t = GetTime(); + // Measure the time before replacing the node + double t = GetTime(); #endif - MS_LOG(DEBUG) << "Replace " << node->DebugString() << " with " << res->DebugString() << ", by " - << substitution->name_; - (void)manager->Replace(node, res); + + // Log the replacement information + MS_LOG(DEBUG) << "Replace " << node->DebugString() << " with " << res->DebugString() << ", by " + << substitution->name_; + + // Replace the original node with the substitution result + (void)manager->Replace(node, res); + #ifdef ENABLE_PROFILE - MsProfile::StatTime("replace." + substitution->name_, GetTime() - t); + // Measure the time after replacing the node and record the time taken for the replacement + MsProfile::StatTime("replace." + substitution->name_, GetTime() - t); #endif - return res; + + // Return the substitution result + return res; + } } - } + // If the node does not match the predicate or the substitution result is null or the same as the original node, + // return the original node without any changes + return node; +} + // Closing brace for the main function + + // Return nullptr to indicate an unsuccessful program termination return nullptr; } +// A static function that updates a transforming list for substitutions static void UpdateTransformingListForSubstitutions(const AnfNodePtr &node, std::deque *todo, bool change) { + + // Check if the given node is a value node containing a FuncGraph object if (IsValueNode(node)) { + + // If it is, add the output of the FuncGraph to the todo list (*todo).emplace_back(GetValueNode(node)->output()); } +} + // Check if the 'change' variable is true if (change) { + // If 'change' is true, add the 'node' pointer to the end of the vector pointed to by 'todo' (*todo).emplace_back(node); } else { + // If 'change' is false if (node->isa()) { + // Check if 'node' is of type CNode auto &inputs = node->cast()->inputs(); + // Get the inputs of the CNode and store them in the 'inputs' variable + + // Use std::copy to copy the elements from 'inputs' to the end of the vector pointed to by 'todo' (void)std::copy(inputs.begin(), inputs.end(), std::back_inserter(*todo)); } } } +// A static function that updates the transforming list for intermediate representation (IR) +// It takes in a node, a deque of nodes to be transformed, a boolean flag indicating if there is a change, and a substitution object + static void UpdateTransformingListForIR(const AnfNodePtr &node, std::deque *todo, bool change, const SubstitutionPtr &substitution) { + + // Check if the node is a value node containing a FuncGraph object if (IsValueNode(node)) { + + // If it is, add the output of the FuncGraph to the transforming list (*todo).emplace_back(GetValueNode(node)->output()); } + // No return statement here, as this is a void function +} - // If there is a priority pattern in substitution, don't transform the new node, - // otherwise some nodes may match the wrong patterns. +// Check if there is a priority pattern in substitution and if change is true if (change && substitution != nullptr && !substitution->has_priority_pattern_) { + // If there is a priority pattern, do not transform the new node + // Instead, add the node to the todo list for further processing (*todo).emplace_back(node); } else { + // If there is no priority pattern or change is false + // Check if the node is of type CNode if (node->isa()) { + // Get the inputs of the CNode auto &inputs = node->cast()->inputs(); + // Copy the inputs and add them to the todo list for further processing (void)std::copy(inputs.begin(), inputs.end(), std::back_inserter(*todo)); } } } +// This function updates a transforming list with user nodes based on the given optimizer, node, todo list, change flag, and seen number. + static void UpdateTransformingListWithUserNodes(const OptimizerPtr &optimizer, const AnfNodePtr &node, std::deque *todo, bool change, SeenNum seen) { - if (!change) { - return; - } - auto manager = optimizer->manager(); - auto &node_users = manager->node_users(); - auto users_iterator = node_users.find(node); - if (users_iterator == node_users.end()) { - return; - } - auto users = users_iterator->second; - for (auto &use : users) { - auto use_node = use.first; - if (use_node == nullptr) { - continue; + // If there is no change, return without doing anything + if (!change) { + return; } - (*todo).emplace_back(use_node); - if (use_node->seen_ == seen) { - use_node->seen_--; + + // Get the manager from the optimizer + auto manager = optimizer->manager(); + + // Get the node users from the manager + auto &node_users = manager->node_users(); + + // Find the users of the given node + auto users_iterator = node_users.find(node); + + // If the node has no users, return without doing anything + if (users_iterator == node_users.end()) { + return; + } + + // Get the users of the node + auto users = users_iterator->second; + + // Iterate over each user + for (auto &use : users) { + auto use_node = use.first; + + // If the user node is null, continue to the next user + if (use_node == nullptr) { + continue; + } + + // Add the user node to the todo list + (*todo).emplace_back(use_node); + + // If the user node has the same seen number, decrement the seen number + if (use_node->seen_ == seen) { + use_node->seen_--; + } + } +} } } } -bool SubstitutionList::ApplyIRToSubstitutions(const OptimizerPtr &optimizer, const FuncGraphPtr &func_graph) const { +// Check if profiling is enabled #ifdef ENABLE_PROFILE + // Get the current time double start = GetTime(); #endif - FuncGraphManagerPtr manager = optimizer->manager(); - auto seen = NewSeenGeneration(); - std::deque todo; - todo.emplace_back(func_graph->output()); - bool changes = false; - auto &all_nodes = manager->all_nodes(); - while (!todo.empty()) { +// Get the function graph manager from the optimizer +FuncGraphManagerPtr manager = optimizer->manager(); + +// Create a new seen generation for tracking visited nodes +auto seen = NewSeenGeneration(); + +// Create a deque to store the nodes to be processed +std::deque todo; + +// Add the output node of the function graph to the todo list +todo.emplace_back(func_graph->output()); + +// Initialize the changes flag to false +bool changes = false; + +// Create a reference variable 'all_nodes' and assign it the value returned by calling the 'all_nodes' function of the 'manager' object +auto &all_nodes = manager->all_nodes(); + +// Start a while loop that continues until the 'todo' list is empty +while (!todo.empty()) { + + // Create a pointer variable 'node' and assign it the value of the front element of the 'todo' list AnfNodePtr node = todo.front(); + + // Remove the front element from the 'todo' list todo.pop_front(); + // Check if the node is nullptr or if its 'seen_' flag is equal to 'seen' + // Also check if the node is not traversable or if it is not present in the 'all_nodes' container if (node == nullptr || node->seen_ == seen || !isTraversable(node) || !all_nodes.contains(node)) { - continue; + // If any of the above conditions are true, skip to the next iteration of the loop + continue; } + + // Set the 'seen_' flag of the node to 'seen' node->seen_ = seen; - bool change = false; + bool change = false; // Initialize a boolean variable "change" to false + + // Iterate over each element in the "list_" using a range-based for loop for (auto &substitution : list_) { + + // Call the DoTransform function with the optimizer, node, and substitution as arguments auto res = DoTransform(optimizer, node, substitution); + + // Check if the result of DoTransform is not nullptr if (res != nullptr) { + + // Set "change" and "changes" variables to true change = true; changes = true; + + // Update the "node" variable with the result of DoTransform node = res; + + // Break out of the loop since a transformation has been applied break; } } + + // Call the UpdateTransformingListForSubstitutions function with the "node", "todo", and "change" variables as arguments UpdateTransformingListForSubstitutions(node, &todo, change); + + // Call the UpdateTransformingListWithUserNodes function with the optimizer, node, "todo", "change", and "seen" variables as arguments UpdateTransformingListWithUserNodes(optimizer, node, &todo, change, seen); } + + // Check if the ENABLE_PROFILE macro is defined #ifdef ENABLE_PROFILE + + // Measure the time taken by the optimizer and add it to the profiling statistics MsProfile::StatTime("opt.transforms." + optimizer->name(), GetTime() - start); #endif + + // Return the value of the "changes" variable return changes; } +// ApplySubstitutionToIR function of the SubstitutionList class, which applies a substitution to the intermediate representation (IR) of a function graph bool SubstitutionList::ApplySubstitutionToIR(const OptimizerPtr &optimizer, const FuncGraphPtr &func_graph, const SubstitutionPtr &substitution) const { -#ifdef ENABLE_PROFILE - double start = GetTime(); -#endif - FuncGraphManagerPtr manager = optimizer->manager(); - auto seen = NewSeenGeneration(); - std::deque todo; - todo.emplace_back(func_graph->output()); - bool changes = false; - auto &all_nodes = manager->all_nodes(); - while (!todo.empty()) { + // Check if profiling is enabled +#ifdef ENABLE_PROFILE + double start = GetTime(); // Get the current time +#endif + + FuncGraphManagerPtr manager = optimizer->manager(); // Get the function graph manager from the optimizer + auto seen = NewSeenGeneration(); // Create a new seen generation object + std::deque todo; // Create a deque to store the nodes to be processed + todo.emplace_back(func_graph->output()); // Add the output node of the function graph to the deque + bool changes = false; // Initialize a flag to track if any changes were made to the IR + + // ... (rest of the code) + +// Create a reference variable 'all_nodes' and assign it the value returned by calling the 'all_nodes' function of the 'manager' object +auto &all_nodes = manager->all_nodes(); + +// Start a while loop that continues until the 'todo' list is empty +while (!todo.empty()) { + + // Create a pointer variable 'node' and assign it the value of the front element of the 'todo' list AnfNodePtr node = todo.front(); + + // Remove the front element from the 'todo' list todo.pop_front(); + // Check if the node is nullptr or if its 'seen_' flag is equal to 'seen' + // Also check if the node is not traversable or if it is not present in the 'all_nodes' container if (node == nullptr || node->seen_ == seen || !isTraversable(node) || !all_nodes.contains(node)) { - continue; + // If any of the above conditions are true, skip to the next iteration of the loop + continue; } + + // Set the 'seen_' flag of the node to 'seen' node->seen_ = seen; - bool change = false; + bool change = false; // Initialize a boolean variable "change" to false + + // Call the DoTransform function with the provided arguments and store the result in the "res" variable auto res = DoTransform(optimizer, node, substitution); + + // Check if the result is not equal to nullptr if (res != nullptr) { - change = true; - changes = true; - node = res; + change = true; // Set the "change" variable to true + changes = true; // Set the "changes" variable to true + node = res; // Update the "node" variable with the result } + + // Call the UpdateTransformingListForIR function with the provided arguments UpdateTransformingListForIR(node, &todo, change, substitution); + + // Call the UpdateTransformingListWithUserNodes function with the provided arguments UpdateTransformingListWithUserNodes(optimizer, node, &todo, change, seen); } #ifdef ENABLE_PROFILE - MsProfile::StatTime("opt.transform." + optimizer->name(), GetTime() - start); -#endif - return changes; -} + // If the ENABLE_PROFILE macro is defined, then execute the following code block + // Use the MsProfile class to record the time taken for a specific optimizer transformation + // The name of the transformation is obtained by concatenating "opt.transform." with the name of the optimizer + // The time taken is calculated by subtracting the start time from the current time + MsProfile::StatTime("opt.transform." + optimizer->name(), GetTime() - start); +#endif + +// Return the number of changes made by the optimizer +return changes; + +// A member function of the SubstitutionList class that displays the status of substitutions void SubstitutionList::DisplayStatusOfSubstitution(const mindspore::HashMap> &status, const OptimizerPtr &optimizer, size_t space) const { + + // Set the padding width for formatting constexpr int pad_width = 4; + + // Create a string stream to store the output std::stringstream ss; + + // Add a newline and display the name and counter of the current optimizer pass ss << std::endl << "Pass: " << optimizer->name() << "(" << optimizer->CurPass_.counter << ")_" << optimizer->CurPass_.name << std::endl; + + // Iterate over the list of substitutions for (size_t i = 0; i < list_.size(); i++) { + + // Get the name of the current substitution auto name = list_[i]->name_; + + // Format and add the name to the output stream with left alignment and padding ss << std::left << std::setw(SizeToInt(space) + pad_width) << name << "\t"; + + // Iterate over the status vector for the current substitution for (auto change : status.at(name + std::to_string(i))) { + + // Add the status change to the output stream ss << change << " "; } + + // Add a newline after displaying the status for the current substitution ss << std::endl; } + + // Log the output stream as a debug message MS_LOG(DEBUG) << ss.str(); } bool SubstitutionList::ApplySubstitutionsToIR(const OptimizerPtr &optimizer, const FuncGraphPtr &func_graph) const { // Add for substitution status counting + + // Initialize a variable to store the size of the list size_t space = 0; + + // Create a hashmap to store the substitution status for each substitution in the list mindspore::HashMap> status; + + // Check if the optimizer is in debug mode if (optimizer->is_on_debug_) { + // Iterate over the list of substitutions for (size_t i = 0; i < list_.size(); i++) { + // Create an entry in the status hashmap for the substitution, using its name and index as the key status[list_[i]->name_ + std::to_string(i)] = {}; } } + // Rest of the code is not provided, so we cannot provide comments for it +} - bool changes = false; - bool loop = true; - while (loop) { +// Initialize a boolean variable `changes` to false, indicating that no changes have been made yet +bool changes = false; + +// Initialize a boolean variable `loop` to true, indicating that the loop should continue +bool loop = true; + +// Enter a while loop that will continue until `loop` is set to false +while (loop) { + // Set `loop` to false at the beginning of each iteration loop = false; + + // Iterate over the elements in the `list_` container using a for loop for (size_t i = 0; i < list_.size(); i++) { - const auto &substitution = list_[i]; - bool change = ApplySubstitutionToIR(optimizer, func_graph, substitution); - changes = changes || change; - loop = loop || change; -#ifdef ENABLE_DUMP_IR - static const auto enable_dump_pass_ir = GetDumpConfig().enable_dump_pass_ir; - if (enable_dump_pass_ir && MsContext::GetInstance()->get_param(MS_CTX_SAVE_GRAPHS_FLAG)) { - auto fg_name = optimizer->name() + "_r" + std::to_string(optimizer->CurPass_.counter) + "_" + - optimizer->CurPass_.name + "_" + substitution->name_; - DumpIR(fg_name + ".ir", func_graph); - if (MsContext::GetInstance()->get_param(MS_CTX_EXECUTION_MODE) != kPynativeMode) { - ExportIR(fg_name + ".dat", func_graph); - draw::Draw(fg_name + ".dot", func_graph); + // Get a reference to the current element in the `list_` container + const auto &substitution = list_[i]; + + // Apply the substitution to the intermediate representation (IR) using the `ApplySubstitutionToIR` function + bool change = ApplySubstitutionToIR(optimizer, func_graph, substitution); + + // Update the `changes` variable by performing a logical OR operation with `change` + changes = changes || change; + + // Update the `loop` variable by performing a logical OR operation with `change` + loop = loop || change; + + // Check if the ENABLE_DUMP_IR macro is defined and if the MS_CTX_SAVE_GRAPHS_FLAG flag is set to true + static const auto enable_dump_pass_ir = GetDumpConfig().enable_dump_pass_ir; + if (enable_dump_pass_ir && MsContext::GetInstance()->get_param(MS_CTX_SAVE_GRAPHS_FLAG)) { + // Generate a unique name for the function graph based on the optimizer's name, current pass counter, pass name, and substitution name + auto fg_name = optimizer->name() + "_r" + std::to_string(optimizer->CurPass_.counter) + "_" + + optimizer->CurPass_.name + "_" + substitution->name_; + + // Dump the intermediate representation (IR) of the function graph to a file with the ".ir" extension + DumpIR(fg_name + ".ir", func_graph); + + // Check if the execution mode is not kPynativeMode + if (MsContext::GetInstance()->get_param(MS_CTX_EXECUTION_MODE) != kPynativeMode) { + // Export the intermediate representation (IR) of the function graph to a file with the ".dat" extension + ExportIR(fg_name + ".dat", func_graph); + + // Generate a visualization of the function graph in the DOT format and save it to a file with the ".dot" extension + draw::Draw(fg_name + ".dot", func_graph); + } } } -#endif + } +// The #endif directive is used to end a conditional compilation block started by #ifdef or #ifndef +// Since there is no corresponding #ifdef or #ifndef in the provided code, this #endif is unnecessary +// It can be safely removed from the code - // Record the status of each substitution + // Check if the optimizer is in debug mode if (optimizer->is_on_debug_) { + + // Record the status of each substitution by appending the change to the corresponding vector in the 'status' map status[substitution->name_ + std::to_string(i)].push_back(change); + + // Update the 'space' variable to keep track of the maximum length of substitution names space = std::max(substitution->name_.size(), space); } } + + // Check if the 'is_once_' flag is set to true if (is_once_) { + + // If so, break out of the loop and stop further iterations break; } } - // Display the status of each substitution + // Check if the optimizer's debug flag is turned on if (optimizer->is_on_debug_) { + + // If the flag is on, call the function DisplayStatusOfSubstitution + // to display the status of each substitution DisplayStatusOfSubstitution(status, optimizer, space); } + + // Return the number of changes made during the substitution process return changes; } -bool SubstitutionList::operator()(const FuncGraphPtr &func_graph, const OptimizerPtr &optimizer) const { +// Check if the optimizer and func_graph pointers are not null MS_EXCEPTION_IF_NULL(optimizer); MS_EXCEPTION_IF_NULL(func_graph); + + // Get the FuncGraphManager from the optimizer FuncGraphManagerPtr manager = optimizer->manager(); + + // Add the func_graph to the manager manager->AddFuncGraph(func_graph); + + // Initialize the changes flag to false bool changes = false; + + // Determine the traverse mode based on the value of the environment variable "MS_DEV_TRAVERSE_SUBSTITUTIONS_MODE" static const auto traverse_mode = (common::GetEnv("MS_DEV_TRAVERSE_SUBSTITUTIONS_MODE") != "1" ? kOptTraverseFromIRToSubstitutions : kOptTraverseFromSubstitutionsToIR); + + // Check if the traverse mode is kOptTraverseFromIRToSubstitutions and the execution mode is not kPynativeMode + // and the optimizer's traverse_nodes_first flag is true and is_once_ and global_sensitive_ flags are false if (traverse_mode == kOptTraverseFromIRToSubstitutions && MsContext::GetInstance()->get_param(MS_CTX_EXECUTION_MODE) != kPynativeMode && optimizer->traverse_nodes_first() && !is_once_ && !global_sensitive_) { + // Print debug log MS_LOG(DEBUG) << "IR >> SUB, " << optimizer->name() << "(r" << optimizer->CurPass_.counter << ")_" << optimizer->CurPass_.name; + + // Apply IR to Substitutions changes = ApplyIRToSubstitutions(optimizer, func_graph); } else { + // Print debug log MS_LOG(DEBUG) << "SUB >> IR, " << optimizer->name() << "(r" << optimizer->CurPass_.counter << ")_" << optimizer->CurPass_.name; + + // Apply Substitutions to IR changes = ApplySubstitutionsToIR(optimizer, func_graph); } + // Return the value of the variable "changes" to the caller return changes; } +// A function that performs a simple rewriting operation bool SimpleRewriter::Run() { + + // Initialize a boolean variable to keep track of whether any changes were made bool changed = false; + + // Create a new generation of seen nodes auto seen = NewSeenGeneration(); + + // Create a deque to store the nodes that need to be processed std::deque todo; + + // A lambda function that adds a node to the todo deque if it hasn't been seen in the current generation auto add_todo = [&seen, &todo](const AnfNodePtr &node) { if (node != nullptr && node->seen_ != seen) { (void)todo.emplace_back(node); } }; + + // Add the root graph's output node to the todo deque (void)todo.emplace_back(root_graph_->output()); + + // Get a reference to all the nodes in the manager auto &all_nodes = manager_->all_nodes(); + + // Process the nodes in the todo deque until it becomes empty while (!todo.empty()) { + + // Get the front node from the todo deque AnfNodePtr node = std::move(todo.front()); todo.pop_front(); + + // Skip the node if it is null, has already been seen in the current generation, or is not present in the manager if (node == nullptr || node->seen_ == seen || !all_nodes.contains(node)) { continue; } + + // Mark the node as seen in the current generation node->seen_ = seen; + + // Try to cast the node to a CNodePtr (a node that represents a computation) auto cnode = node->cast(); + + // If the cast is successful, perform some operations on the CNodePtr if (cnode != nullptr) { + // Iterate over each input of the current node for (auto &input : cnode->inputs()) { + // Add the input to the list of nodes to be processed add_todo(input); } } else { + // If the current node is a FuncGraphPtr auto fg = GetValueNode(node); if (fg != nullptr) { + // Add the output of the FuncGraphPtr to the list of nodes to be processed add_todo(fg->output()); } } + // Rewrite the current node and obtain the new node auto new_node = NodeRewrite(node); if (new_node != nullptr) { + // Replace the current node with the new node in the graph manager_->Replace(node, new_node); + // Set the flag indicating that changes have been made changed = true; } } + // Return the flag indicating whether any changes have been made return changed; } } // namespace opt -} // namespace mindspore +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/frontend/optimizer/pass_group.cc b/mindspore/ccsrc/frontend/optimizer/pass_group.cc index 56ff81475eb..61968e4c640 100644 --- a/mindspore/ccsrc/frontend/optimizer/pass_group.cc +++ b/mindspore/ccsrc/frontend/optimizer/pass_group.cc @@ -16,56 +16,100 @@ #include "frontend/optimizer/pass_group.h" #include "frontend/optimizer/py_pass_manager.h" +// Start of the "mindspore" namespace namespace mindspore { + +// Start of the "opt" namespace, which is a sub-namespace of "mindspore" namespace opt { + +// Start of the "python_pass" namespace, which is a sub-namespace of "opt" namespace python_pass { + +// Definition of the "AddPass" function of the "PassGroup" class void PassGroup::AddPass(const PythonPassPtr &pass) { + + // Check if the passed pointer is not null if (pass != nullptr) { + + // Add the pass to the list of passes passes_.push_back(pass); } } +// A member function of the PassGroup class that deletes a pass with the given name bool PassGroup::DeletePass(const std::string &pass_name) { + + // Iterate through the passes in the passes_ vector for (auto iter = passes_.begin(); iter != passes_.end(); iter++) { + + // Check if the name of the current pass matches the given pass_name if ((*iter)->name() == pass_name) { + + // Set the current pass to nullptr *iter = nullptr; + + // Erase the current pass from the passes_ vector passes_.erase(iter); + + // Return true to indicate that the pass was successfully deleted return true; } } + + // If no pass with the given name was found, return false return false; } +// A member function of the PassGroup class that runs a series of passes on a given function graph bool PassGroup::Run(const FuncGraphPtr &func_graph, const std::vector &passes, const MatchResultPtr &res) const { + + // Check if the function graph is null, if so, return false if (func_graph == nullptr) { return false; } + + // Initialize a boolean variable to keep track of whether any changes were made during the passes bool changed = false; + + // Iterate over each pass in the passes vector for (const auto &pass : passes) { + + // Check if the pass is not null if (pass != nullptr) { + + // Call the Run function of the pass, passing in the function graph and the match result + // If the pass returns true, indicating that changes were made, set the changed variable to true if (pass->Run(func_graph, res)) { changed = true; } } } + + // Return the value of the changed variable return changed; } +// Definition of the Run function in the PassGroup class + bool PassGroup::Run(const FuncGraphPtr &func_graph) const { - bool changed = false; - // run all passes - bool change = true; - auto res = PyPassManager::GetInstance()->GetMatchResult(); - while (change) { - change = Run(func_graph, passes_, res); - changed = change || changed; - if (run_only_once_) { - break; + bool changed = false; // Initialize a boolean variable to keep track of whether any changes were made + + bool change = true; // Initialize a boolean variable to control the loop + auto res = PyPassManager::GetInstance()->GetMatchResult(); // Get the match result from the PyPassManager singleton instance + + while (change) { // Loop until no more changes are made + change = Run(func_graph, passes_, res); // Call the Run function with the function graph, passes, and match result + changed = change || changed; // Update the changed variable if any changes were made in this iteration + + if (run_only_once_) { // Check if the pass group should only be run once + break; // If so, break out of the loop } } - return changed; + + return changed; // Return the value of the changed variable } + } // namespace python_pass } // namespace opt -} // namespace mindspore +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/frontend/optimizer/pattern.cc b/mindspore/ccsrc/frontend/optimizer/pattern.cc index 6e8cc1eb5ba..7b7f2703176 100644 --- a/mindspore/ccsrc/frontend/optimizer/pattern.cc +++ b/mindspore/ccsrc/frontend/optimizer/pattern.cc @@ -16,159 +16,315 @@ #include "frontend/optimizer/pattern.h" #include "include/common/pybind_api/api_register.h" +// Define the namespace "mindspore" for organizing related code namespace mindspore { -namespace opt { -namespace python_pass { -int64_t Pattern::g_id_ = 0; + // Define the namespace "opt" within the "mindspore" namespace for organizing optimization-related code + namespace opt { + + // Define the namespace "python_pass" within the "opt" namespace for organizing Python-specific optimization passes + namespace python_pass { + + // Define a static member variable "g_id_" of type int64_t for the Pattern class, initialized to 0 + int64_t Pattern::g_id_ = 0; + } + } +} + +// Define the match function of the Prim class, which takes a reference to an AnfNodePtr as input and returns a MatchResultPtr MatchResultPtr Prim::match(const AnfNodePtr &node) { + + // Check if the given node is not a ValueNode of type Primitive, if not, return nullptr if (!IsValueNode(node)) { return nullptr; } + + // Create a new MatchResultPtr using std::make_shared MatchResultPtr res = std::make_shared(); - // iterate over all primitives + + // Iterate over all primitives in the primitives_ container for (auto &iter : primitives_) { + + // Check if the given node is a primitive that matches the current iter or if the iter has a name of "*" if (IsPrimitive(node, iter) || iter->name() == "*") { + + // Set the matched_prim_ member variable of the Prim class to the current iter matched_prim_ = iter; + + // Add an entry to the MatchResultPtr, associating the shared_from_base() with the given node res->add_entry(shared_from_base(), node); + + // Return the MatchResultPtr return res; } } + + // If no match is found, return nullptr return nullptr; } +// Function to match a given node with a pattern MatchResultPtr Call::match(const AnfNodePtr &node) { + + // Check if the node is a primitive CNode if (!IsPrimitiveCNode(node)) { return nullptr; } + + // Create a new MatchResult object MatchResultPtr res = std::make_shared(); - // IsPrimitiveCNode + + // Cast the node to a CNodePtr auto cnode = node->cast(); MS_EXCEPTION_IF_NULL(cnode); - // Check Primitive ValueNode + + // Check if the pattern is a Primitive ValueNode if (prim_pattern_ != nullptr) { - // Passed in prim_pattern + + // Match the input(0) of the CNode with the prim_pattern auto prim_value_res = prim_pattern_->match(cnode->input(0)); + + // If the match is unsuccessful, return nullptr if (prim_value_res == nullptr) { return nullptr; } + + // Merge the match result with the existing result res->merge(prim_value_res); - } else if (prim_ != nullptr) { - // Passed in primitive/primitive str + } + // If the pattern is a primitive or primitive string + else if (prim_ != nullptr) { + + // Check if the input(0) of the CNode is a primitive with the given prim_ if (!IsPrimitive(cnode->input(0), prim_)) { return nullptr; } } else { + // If the CallWith pattern is uninitialized, throw an exception with an error message MS_LOG(EXCEPTION) << "Uninitialized CallWith pattern."; } - // Check inputs + + // Check the number of inputs auto p_inputs_size = inputs_.size(); auto node_inputs_size = cnode->size() - 1; if (p_inputs_size != 0 && p_inputs_size != node_inputs_size) { + // If the number of inputs is not 0 and not equal to the number of inputs in the node, return nullptr return nullptr; } - // If inputs is not specified, add node without looking into its inputs + + // If the number of inputs is 0, add the node without looking into its inputs if (p_inputs_size == 0) { res->add_entry(shared_from_base(), cnode->input(0)); return res; } + + // Iterate through the inputs and match them with the corresponding patterns bool failed = false; for (std::size_t i = 0; i < node_inputs_size; i++) { auto pattern = inputs_[i]; auto input = cnode->input(i + 1); auto input_match_result = pattern->match(input); + // Check if the input_match_result is nullptr if (input_match_result == nullptr) { - failed = true; - break; + // If it is nullptr, set failed to true + failed = true; + // Break out of the loop + break; } + + // Merge the input_match_result into the res object res->merge(input_match_result); - } - if (!failed) { - res->add_entry(shared_from_base(), cnode->input(0)); - return res; - } - return nullptr; } +// Check if failed is false +if (!failed) { + // If it is false, add an entry to the res object using shared_from_base() and cnode->input(0) + res->add_entry(shared_from_base(), cnode->input(0)); + // Return the res object + return res; +} + +// If failed is true, return nullptr +return nullptr; + +// Define the match function of the OneOf class, which takes a reference to an AnfNodePtr as input and returns a MatchResultPtr MatchResultPtr OneOf::match(const AnfNodePtr &node) { + + // Iterate over each pattern in the patterns_ vector for (auto &iter : patterns_) { + + // Call the match function of the current pattern and store the result in res auto res = iter->match(node); + + // If the result is not nullptr (i.e., a match is found) if (res != nullptr) { + + // Add an entry to the result, associating it with the current OneOf instance and the matched node res->add_entry(shared_from_base(), node); + + // Return the result return res; } } + + // If no match is found, return nullptr return nullptr; } +// Define the match function for the NoneOf class, which takes a reference to an AnfNodePtr as input and returns a MatchResultPtr MatchResultPtr NoneOf::match(const AnfNodePtr &node) { + + // Iterate over each pattern in the patterns_ vector for (auto &iter : patterns_) { + + // Call the match function of the current pattern and store the result in match_res auto match_res = iter->match(node); + + // If the match result is not nullptr (indicating a successful match), return nullptr if (match_res != nullptr) { return nullptr; } } + + // Create a new MatchResult object using std::make_shared auto res = std::make_shared(); + + // Add an entry to the MatchResult object, associating the NoneOf object with the input node res->add_entry(shared_from_base(), node); + + // Return the MatchResult object return res; } +// Define the match function of the Any class, which takes a reference to an AnfNodePtr as input and returns a MatchResultPtr MatchResultPtr Any::match(const AnfNodePtr &node) { + + // Create a new MatchResultPtr using std::make_shared, which dynamically allocates memory for a MatchResult object and returns a shared pointer to it MatchResultPtr res = std::make_shared(); + + // Add an entry to the MatchResult object, using the shared_from_base function to obtain a shared pointer to the current instance of the Any class, and passing it along with the input node to the add_entry function res->add_entry(shared_from_base(), node); + + // Return the MatchResultPtr return res; } +// Define a function named "match" that takes a reference to an AnfNodePtr and returns a MatchResultPtr MatchResultPtr Imm::match(const AnfNodePtr &node) { + + // Check if the given node is not a ValueNode of type Int32Imm, if not, return nullptr if (!IsValueNode(node)) { return nullptr; } - // Check value + + // Cast the given node to a ValueNodePtr and assign it to the variable "value_node" auto value_node = node->cast(); MS_EXCEPTION_IF_NULL(value_node); + + // Cast the value of "value_node" to an Int32ImmPtr and assign it to the variable "value_ptr" auto value_ptr = value_node->value()->cast(); MS_EXCEPTION_IF_NULL(value_ptr); + + // Check if the value of "value_ptr" is equal to the value stored in the variable "value_" if ((int32_t)value_ptr->value() == value_) { + + // If the condition is true, create a new MatchResultPtr using std::make_shared MatchResultPtr res = std::make_shared(); + + // Add an entry to the MatchResultPtr, where the key is a shared pointer to the current Imm object and the value is the given node res->add_entry(shared_from_base(), node); + + // Return the MatchResultPtr return res; } + + // If the condition is false, return nullptr return nullptr; } +// Function to retrieve a node from the match result based on a given pattern AnfNodePtr MatchResult::get_node(const PatternPtr &pattern) { + + // Find the entry in the match result map corresponding to the given pattern auto entry = match_result_.find(pattern); + + // If the entry is not found (i.e., pattern not present in the match result), return nullptr if (entry == match_result_.end()) { return nullptr; } + + // Return the node associated with the pattern in the match result map return entry->second; } +// Define the merge function of the MatchResult class, which takes a MatchResultPtr as input void MatchResult::merge(const MatchResultPtr &other_result) { + + // Get the result map from the other_result MatchResultPtr auto other_result_map = other_result->result(); - // add/update entries in other_result + + // Iterate over each entry in the other_result_map for (auto &iter : other_result_map) { + + // Add or update the corresponding entry in the match_result_ map match_result_[iter.first] = iter.second; } } +// Register the Python bindings for the Pattern class and its derived classes REGISTER_PYBIND_DEFINE( Pattern, ([](const py::module *m) { + + // Register the Pattern class with the Python module, and define its constructor (void)py::class_>(*m, "Pattern").def(py::init<>()); + + // Register the OneOf class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "OneOf_").def(py::init>()); + + // Register the Prim class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "Prim_", py::dynamic_attr()) .def(py::init, string>()); + + // Register the Call class with the Python module, and define its constructors (void)py::class_, Pattern>(*m, "Call_") .def(py::init>()) .def(py::init>()); + + // Register the NoneOf class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "NoneOf_").def(py::init>()); + + // Register the Any class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "Any").def(py::init<>()); + + // Register the NewTensor class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "NewTensor_") .def(py::init()); + + // Register the NewParameter class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "NewParameter_") .def(py::init()); + + // Register the Imm class with the Python module, and define its constructor (void)py::class_, Pattern>(*m, "Imm").def(py::init()); })); + +// End of the python_pass namespace } // namespace python_pass + +// End of the opt namespace } // namespace opt + +// End of the mindspore namespace } // namespace mindspore +// Include the standard input-output header for C (primarily for printf, but we aren't using printf here) +#include + +// The main function, entry point of the program +int main(){ + + // Use the standard C++ output stream to print "Hello World" followed by a newline + std::cout << "Hello World" << std::endl; + + // Return 0 to indicate successful program termination + return 0; +} \ No newline at end of file diff --git a/mindspore/ccsrc/frontend/optimizer/recompute.cc b/mindspore/ccsrc/frontend/optimizer/recompute.cc index d4829f98d6e..3e2fda4497a 100644 --- a/mindspore/ccsrc/frontend/optimizer/recompute.cc +++ b/mindspore/ccsrc/frontend/optimizer/recompute.cc @@ -14,126 +14,252 @@ * limitations under the License. */ +// Include the header file for the Recompute optimizer from the frontend/optimizer directory #include "frontend/optimizer/recompute.h" + +// Include the necessary header files for using smart pointers, queues, lists, vectors, and strings #include #include #include #include #include + +// Include the header files for using hash maps and hash sets from the utils directory #include "utils/hash_map.h" #include "utils/hash_set.h" + +// Include the header file for using the FuncGraph class from the ir directory #include "ir/func_graph.h" + +// Include the header file for using the CoreOps class from the mindspore/core/base directory #include "mindspore/core/base/core_ops.h" + +// Include the header file for using utility functions from the include/common/utils directory #include "include/common/utils/utils.h" +// The code is defining a namespace called "mindspore" which contains another namespace called "opt" namespace mindspore { namespace opt { namespace { + +// Define a constant string variable "kGradientsFlag" with the value "Gradients" constexpr auto kGradientsFlag = "Gradients"; + +// Define a constant integer variable "fusion_id_increasement_size" with the value 2000 const int64_t fusion_id_increasement_size = 2000; + +// Define a function "CanNotRecomputed" that takes a CNodePtr object as input and returns a boolean value bool CanNotRecomputed(const CNodePtr &node) { + + // Define a static HashSet object "not_recomputed_op_list" that contains PrimitivePtr objects static mindspore::HashSet not_recomputed_op_list{ prim::kPrimDropoutGenMask, prim::kPrimLoad, prim::kPrimTupleGetItem, prim::kPrimSend, prim::kPrimReceive}; + // The HashSet contains the following PrimitivePtr objects: kPrimDropoutGenMask, kPrimLoad, kPrimTupleGetItem, kPrimSend, kPrimReceive - return std::any_of(not_recomputed_op_list.begin(), not_recomputed_op_list.end(), - [&node](const PrimitivePtr &prim) { return IsPrimitiveCNode(node, prim); }); -} + // Rest of the code is not provided, but it is expected to contain the implementation of the CanNotRecomputed function + // based on the HashSet defined above. + // The function is likely to check if the given CNodePtr object's primitive is present in the not_recomputed_op_list HashSet + // and return true if it is present, indicating that the operation cannot be recomputed. + // Otherwise, it will return false. + // The purpose of this function is not clear without the complete code. +// Return the result of applying the std::any_of algorithm to the range defined by not_recomputed_op_list.begin() and not_recomputed_op_list.end() +// The algorithm checks if any element in the range satisfies the given condition +// The condition is defined by the lambda function [&node](const PrimitivePtr &prim) { return IsPrimitiveCNode(node, prim); } +// The lambda function takes a const reference to a PrimitivePtr object named prim and compares it with the node using the IsPrimitiveCNode function +// If any element in the range satisfies the condition, std::any_of returns true; otherwise, it returns false + +// Check if the given node is a backpropagation node bool IsBpropNode(const AnfNodePtr &node) { + // Check if the node is null, if so, return false MS_EXCEPTION_IF_NULL(node); + + // Check if the node is of type CNode, if not, return false if (!node->isa()) { return false; } + + // Check if the full name of the node starts with the gradients flag + // If it does, return true, otherwise return false return node->fullname_with_scope().find(kGradientsFlag) == 0; } +// A function that checks if a given AnfNode has a recomputed scope bool WithRecomputedScope(const AnfNodePtr &node) { + // Check if the given node is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(node); + + // Check if the given node is not a CNode (a computational node) if (!node->isa()) { + // If it is not a CNode, return false return false; } + + // Get the full name of the node with its scope auto full_name_with_scope = node->fullname_with_scope(); + + // Check if the full name starts with the attribute "recompute" return full_name_with_scope.find(kAttrRecompute) == 0; } +// Function to get the value of the "Recompute" attribute of a CNode ValuePtr GetRecomputeCNodeAttr(const AnfNodePtr &node) { + // Check if the input node is null MS_EXCEPTION_IF_NULL(node); + + // Cast the input node to a CNode auto cnode = node->cast(); + + // Check if the cast was successful (i.e., if the input node is indeed a CNode) if (cnode == nullptr) { + // If the input node is not a CNode, return nullptr return nullptr; } + + // Return the value of the "Recompute" attribute of the CNode return cnode->GetAttr(kAttrRecompute); } +// Function to check if a CNode attribute "NoRecompute" is set to false bool IsSetNoRecomputeCNodeAttr(const AnfNodePtr &node) { + + // Get the value of the "Recompute" attribute for the CNode auto cnode_recompute_val = GetRecomputeCNodeAttr(node); + + // Check if the "Recompute" attribute exists, is of type BoolImm, and has a value of false return cnode_recompute_val != nullptr && cnode_recompute_val->isa() && !GetValue(cnode_recompute_val); } +// Function to check if the recompute attribute is set for a given node bool IsSetRecomputeCNodeAttr(const AnfNodePtr &node) { + + // Get the recompute attribute value for the given node auto cnode_recompute_val = GetRecomputeCNodeAttr(node); + + // Check if the recompute attribute value is not null, is of type BoolImm, and has a true value return cnode_recompute_val != nullptr && cnode_recompute_val->isa() && GetValue(cnode_recompute_val); } +// Function to check if a given node is a candidate for recomputation bool IsCandidateRecomputedNode(const CNodePtr &node) { - // The tuple_getitem in the bprop function should also be recomputed. + + // Check if the node is not a backward propagation node or if it is a tuple_getitem primitive node + // If either of these conditions is true, then the node should be recomputed return (!IsBpropNode(node) || IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) && IsSetRecomputeCNodeAttr(node); } +// Function to find candidate recomputed nodes given a FuncGraphManager and a vector of CNodePtrs std::vector FindCandidateRecomputedNodes(const FuncGraphManagerPtr &mng, const std::vector &cnodes) { MS_EXCEPTION_IF_NULL(mng); + + // Vector to store the candidate recomputed nodes std::vector candidate_recomputed_nodes; + + // Iterate over each CNode in the input vector for (const auto &cnode : cnodes) { MS_EXCEPTION_IF_NULL(cnode); + + // Check if the CNode is a candidate for recomputation if (!IsCandidateRecomputedNode(cnode)) { continue; } - // Check outputs. + + // Check the outputs of the CNode const auto &node_users = mng->node_users(); auto output_set_iter = node_users.find(cnode); + + // If the CNode has no outputs, continue to the next CNode if (output_set_iter == node_users.end()) { continue; } + + // Get the set of node indices that use the output of the CNode const auto &node_index_set = output_set_iter->second; + + // Check if any of the node indices in the set are bprop nodes if (!std::any_of(node_index_set.begin(), node_index_set.end(), [](const auto &node_index) { return IsBpropNode(node_index.first); })) { continue; } + + // If all conditions are met, add the CNode to the candidate recomputed nodes vector + candidate_recomputed_nodes.push_back(cnode); + } + + // Return the vector of candidate recomputed nodes + return candidate_recomputed_nodes; +} // Check inputs. const auto &inputs = cnode->inputs(); + + // If any of the inputs is a bprop node, skip this node and continue to the next iteration of the loop if (std::any_of(inputs.begin(), inputs.end(), [](const AnfNodePtr &node) { return IsBpropNode(node); })) { continue; } + + // Add the current node to the list of candidate recomputed nodes (void)candidate_recomputed_nodes.emplace_back(cnode); } + + // Return the list of candidate recomputed nodes return candidate_recomputed_nodes; } +// Define a function named "GetMaxSubGraph" that takes in a FuncGraphManagerPtr object, a pointer to a HashSet of CNodePtr objects, and two boolean variables: get_inputs and get_outputs. void GetMaxSubGraph(const FuncGraphManagerPtr &mng, mindspore::HashSet *recomputed_nodes, bool get_inputs, bool get_outputs) { + // Check if the FuncGraphManagerPtr object and the HashSet pointer are not null. MS_EXCEPTION_IF_NULL(mng); MS_EXCEPTION_IF_NULL(recomputed_nodes); + + // Create a queue to store CNodePtr objects to visit. std::queue nodes_to_visit; + + // Iterate over each CNodePtr object in the HashSet and push it into the queue. for (const auto &node : *recomputed_nodes) { nodes_to_visit.push(node); } + + // Clear the HashSet. recomputed_nodes->clear(); + + // Start a loop that continues until the queue is empty. while (!nodes_to_visit.empty()) { + // Get the front element of the queue and remove it. auto current_node = nodes_to_visit.front(); nodes_to_visit.pop(); + + // Insert the current_node into the HashSet. recomputed_nodes->insert(current_node); - // No need to find nodes through side-effect dependency. + + // Check if the current_node is a primitive CNode with the primitive name "prim::kPrimUpdateState". + // If it is, skip the rest of the loop and continue with the next iteration. if (IsPrimitiveCNode(current_node, prim::kPrimUpdateState)) { continue; } + + // Check if the get_inputs flag is true. if (get_inputs) { + // Iterate over each input of the current_node. for (const auto &input : current_node->inputs()) { + // Check if the input is null. MS_EXCEPTION_IF_NULL(input); + // ... (code continues) + } + } + } +} + // Check if the input is a CNode if (input->isa()) { + // Cast the input to a CNode pointer auto input_cnode = input->cast(); + + // Check if the input CNode is not already in the set of recomputed nodes + // and if it is a candidate for recomputation if (recomputed_nodes->find(input_cnode) == recomputed_nodes->end() && IsCandidateRecomputedNode(input_cnode)) { + // Push the input CNode to the stack of nodes to visit nodes_to_visit.push(input_cnode); } } @@ -141,48 +267,76 @@ void GetMaxSubGraph(const FuncGraphManagerPtr &mng, mindspore::HashSet } if (get_outputs) { - const auto &node_users = mng->node_users(); - auto output_set_iter = node_users.find(current_node); - if (output_set_iter == node_users.end()) { - continue; - } - for (const auto &node_index_set : output_set_iter->second) { - auto output_node = node_index_set.first; - MS_EXCEPTION_IF_NULL(output_node); - if (output_node->isa()) { - auto output_cnode = output_node->cast(); - if (recomputed_nodes->find(output_cnode) == recomputed_nodes->end() && - IsCandidateRecomputedNode(output_cnode)) { - nodes_to_visit.push(output_cnode); - } + // Get the map of nodes and their users from the manager + const auto &node_users = mng->node_users(); + + // Find the current node in the map + auto output_set_iter = node_users.find(current_node); + + // If the current node is not found in the map, continue to the next iteration + if (output_set_iter == node_users.end()) { + continue; + } + + // Iterate over the set of nodes that use the current node as an output + for (const auto &node_index_set : output_set_iter->second) { + // Get the output node and check if it is a valid CNode + auto output_node = node_index_set.first; + MS_EXCEPTION_IF_NULL(output_node); + if (output_node->isa()) { + // Cast the output node to CNode pointer + auto output_cnode = output_node->cast(); + + // Check if the output node is not in the set of recomputed nodes + // and if it is a candidate for recomputation + if (recomputed_nodes->find(output_cnode) == recomputed_nodes->end() && + IsCandidateRecomputedNode(output_cnode)) { + // Push the output node to the stack of nodes to visit + nodes_to_visit.push(output_cnode); + } + } } - } } - } +} } -void GetOriginRecomputeAndTargetNodes(const FuncGraphManagerPtr &mng, - const mindspore::HashSet &max_recomputed_sub_graph, - mindspore::HashSet *recompute_nodes, - mindspore::HashSet *target_nodes) { - MS_EXCEPTION_IF_NULL(mng); - MS_EXCEPTION_IF_NULL(recompute_nodes); - MS_EXCEPTION_IF_NULL(target_nodes); - const auto &node_users = mng->node_users(); - for (const auto &node : max_recomputed_sub_graph) { - bool inserted = false; - auto output_set_iter = node_users.find(node); - if (output_set_iter == node_users.end()) { +// This function takes in a FuncGraphManagerPtr, a set of CNodes, and two pointers to sets of CNodes. +// It is used to get the origin, recompute, and target nodes based on the given inputs. + +// Check if the input parameters are not null +MS_EXCEPTION_IF_NULL(mng); +MS_EXCEPTION_IF_NULL(recompute_nodes); +MS_EXCEPTION_IF_NULL(target_nodes); + +// Get the node users from the FuncGraphManager +const auto &node_users = mng->node_users(); + +// Iterate over each node in the max_recomputed_sub_graph set +for (const auto &node : max_recomputed_sub_graph) { + bool inserted = false; + + // Find the node in the node_users map + auto output_set_iter = node_users.find(node); + + // If the node is not found in the node_users map, continue to the next iteration + if (output_set_iter == node_users.end()) { + continue; + } + + // Iterate over each node index set in the output_set_iter + for (const auto &node_index_set : output_set_iter->second) { + auto output_node = node_index_set.first; + MS_EXCEPTION_IF_NULL(output_node); + + // Check if the output_node is a bprop node or a primitive CNode with the prim::kPrimTupleGetItem primitive + // If it is, continue to the next iteration + if (!IsBpropNode(output_node) || IsPrimitiveCNode(output_node, prim::kPrimTupleGetItem)) { continue; - } - for (const auto &node_index_set : output_set_iter->second) { - auto output_node = node_index_set.first; - MS_EXCEPTION_IF_NULL(output_node); - // The tuple_getitem to be recomputed can be in the bprop function. - if (!IsBpropNode(output_node) || IsPrimitiveCNode(output_node, prim::kPrimTupleGetItem)) { - continue; } + // Insert the output node into the set of target nodes target_nodes->insert(output_node->cast()); + + // If the node was not already inserted, insert it into the set of recompute nodes if (!inserted) { recompute_nodes->insert(node); inserted = true; @@ -191,155 +345,311 @@ void GetOriginRecomputeAndTargetNodes(const FuncGraphManagerPtr &mng, } } -std::vector GetFirstTargetInputs(const std::vector &origin_nodes_topological, - const mindspore::HashSet &recomputed_origin_nodes, - const mindspore::HashSet &target_nodes) { - std::vector first_target_inputs; - for (const auto &node : origin_nodes_topological) { +// This function takes in three parameters: +// - origin_nodes_topological: a vector of CNodePtr objects representing the nodes in the computation graph in topological order +// - recomputed_origin_nodes: a HashSet of CNodePtr objects representing the nodes that have been recomputed +// - target_nodes: a HashSet of CNodePtr objects representing the target nodes we are interested in + +// Create an empty vector to store the first target inputs +std::vector first_target_inputs; + +// Iterate through each node in the origin_nodes_topological vector +for (const auto &node : origin_nodes_topological) { + // Check if the node is not null MS_EXCEPTION_IF_NULL(node); + + // Check if the node is one of the target nodes if (target_nodes.find(node) != target_nodes.end()) { - for (size_t i = 1; i < node->size(); ++i) { - auto input = node->input(i); - MS_EXCEPTION_IF_NULL(input); - if (!input->isa()) { - continue; + // Iterate through each input of the node starting from index 1 (index 0 is the node itself) + for (size_t i = 1; i < node->size(); ++i) { + // Get the input at index i + auto input = node->input(i); + MS_EXCEPTION_IF_NULL(input); + + // Check if the input is not a CNode (a node in the computation graph) + if (!input->isa()) { + // If it is not a CNode, continue to the next input + continue; + } + + // Check if the input is in the recomputed_origin_nodes set + if (recomputed_origin_nodes.find(input->cast()) != recomputed_origin_nodes.end()) { + // If it is in the recomputed_origin_nodes set, continue to the next input + continue; + } + + // Add the input to the first_target_inputs vector + (void)first_target_inputs.emplace_back(input); } - if (recomputed_origin_nodes.find(input->cast()) != recomputed_origin_nodes.end()) { - continue; - } - (void)first_target_inputs.emplace_back(input); - } - break; + + // Break out of the loop since we have found the first target inputs + break; } +} } + + // Return the value of the variable "first_target_inputs" return first_target_inputs; } +// Check if the given node has already been processed and its result is stored in the `has_grad_inputs_map` bool HasGradInputs(const AnfNodePtr &node, mindspore::HashMap *has_grad_inputs_map) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(has_grad_inputs_map); + + // If the node is already present in the `has_grad_inputs_map`, return its result if (has_grad_inputs_map->find(node) != has_grad_inputs_map->end()) { return has_grad_inputs_map->find(node)->second; } + + // Cast the node to a CNode pointer auto cnode = node->cast(); + + // If the node is not a CNode, it means it does not have any gradient inputs if (cnode == nullptr) { (void)has_grad_inputs_map->emplace(node, false); return false; } + + // Get the inputs of the CNode const auto &inputs = cnode->inputs(); + + // Iterate over the inputs of the CNode for (size_t i = 0; i < inputs.size(); ++i) { // For the pipeline split case, the forward pass may depend on the backward pass. + // If the current CNode is a Depend node and the index is the attach node index, skip it if (IsPrimitiveCNode(cnode, prim::kPrimDepend) && i == kDependAttachNodeIndex) { continue; } + + // Check if the input is a BpropNode or if it has gradient inputs recursively if (IsBpropNode(inputs[i]) || HasGradInputs(inputs[i], has_grad_inputs_map)) { (void)has_grad_inputs_map->emplace(node, true); return true; } } + + // If none of the inputs have gradient inputs, store the result in the `has_grad_inputs_map` and return false + (void)has_grad_inputs_map->emplace(node, false); + return false; +} + } + } (void)has_grad_inputs_map->emplace(node, false); return false; } -bool HasForwardOutput(const FuncGraphManagerPtr &mng, const AnfNodePtr &node) { - MS_EXCEPTION_IF_NULL(mng); - const auto &node_users = mng->node_users(); - auto output_set_iter = node_users.find(node); - if (output_set_iter == node_users.end()) { - return false; - } +// This code snippet seems to be a part of a larger function or method. +// It is not clear what the purpose of this code is without the context. +// It appears to be modifying a map called "has_grad_inputs_map" by inserting a key-value pair. +// The key is the "node" variable and the value is "false". +// Finally, it returns false. - return std::any_of(output_set_iter->second.begin(), output_set_iter->second.end(), - [](const auto &node_index_set) { return !IsBpropNode(node_index_set.first); }); +// Check if the given `node` has any forward outputs in the `node_users` map of the `mng` FuncGraphManager + +// Make sure the `mng` pointer is not null +MS_EXCEPTION_IF_NULL(mng); + +// Get the `node_users` map from the `mng` FuncGraphManager +const auto &node_users = mng->node_users(); + +// Find the iterator for the `node` in the `node_users` map +auto output_set_iter = node_users.find(node); + +// If the iterator is at the end of the `node_users` map, it means the `node` has no forward outputs +if (output_set_iter == node_users.end()) { + return false; } +// Return the result of applying the std::any_of algorithm to the range defined by output_set_iter->second.begin() and output_set_iter->second.end() +// The lambda function passed to std::any_of checks if any element in the range satisfies the condition !IsBpropNode(node_index_set.first) +// The lambda function takes a const reference to an element in the range and returns true if the condition is satisfied, false otherwise +// The result of std::any_of is the boolean value indicating whether any element satisfies the condition +return std::any_of(output_set_iter->second.begin(), output_set_iter->second.end(), + [](const auto &node_index_set) { return !IsBpropNode(node_index_set.first); }); + + +// A function to get the output nodes of a TupleGetItem operation in a graph + +// Parameters: +// - mng: A shared pointer to the FuncGraphManager object +// - node: A shared pointer to the AnfNode object representing the TupleGetItem operation +// - tuple_getitem_output_nodes: A pointer to a vector of AnfNodePtr objects to store the output nodes + void GetTupleGetItemOutputNodes(const FuncGraphManagerPtr &mng, const AnfNodePtr &node, std::vector *tuple_getitem_output_nodes) { + + // Check if the FuncGraphManager pointer is null MS_EXCEPTION_IF_NULL(mng); + + // Check if the vector pointer is null MS_EXCEPTION_IF_NULL(tuple_getitem_output_nodes); + + // Get the map of node users from the FuncGraphManager const auto &node_users = mng->node_users(); + + // Find the iterator for the output set of the given node in the node users map auto output_set_iter = node_users.find(node); + + // If the iterator is at the end of the map, return if (output_set_iter == node_users.end()) { return; } + + // Iterate over the node index sets in the output set for (const auto &node_index_set : output_set_iter->second) { + + // Check if the node is a TupleGetItem primitive CNode if (IsPrimitiveCNode(node_index_set.first, prim::kPrimTupleGetItem)) { + + // Add the node to the vector of output nodes (void)tuple_getitem_output_nodes->emplace_back(node_index_set.first); } } } +// Function to set the recomputed scope for a given CNode bool SetRecomputedScope(const CNodePtr &node) { + + // Check if the recomputed scope is already set for the given CNode + // If it is set, return true + // Otherwise, check if the CNode is a primitive node of type "prim::kPrimDepend" + // If it is a "prim::kPrimDepend" node, check if the recomputed scope is set for its real input + // If the recomputed scope is set for the real input, return true + // Otherwise, return false return WithRecomputedScope(node) || (IsPrimitiveCNode(node, prim::kPrimDepend) && WithRecomputedScope(node->input(kRealInputIndexInDepend))); } -// Set 'recompute' cnode attr for the nodes according to its scope. -// A node set 'recompute' cnode attr can become the candidate recomputed node. +// Set the 'recompute' cnode attribute for the nodes in the graph based on their scope. +// A node with the 'recompute' cnode attribute set can become a candidate for recomputation. + void SetRecomputedAttr(const FuncGraphPtr &graph, const std::vector &origin_nodes_topological) { MS_EXCEPTION_IF_NULL(graph); auto mng = graph->manager(); MS_EXCEPTION_IF_NULL(mng); + + // Create a hashmap to keep track of whether a node has gradient inputs mindspore::HashMap has_grad_inputs_map; + + // Iterate through the nodes in topological order for (const auto &node : origin_nodes_topological) { MS_EXCEPTION_IF_NULL(node); - // The node may be set the non-recomputed before such as the cell outputs. + + // Skip nodes that have already been set with the 'no_recompute' cnode attribute, such as cell outputs if (IsSetNoRecomputeCNodeAttr(node)) { continue; } + + // Skip bprop nodes if (IsBpropNode(node)) { continue; } - // Filter some unrecomputable operators. + + // Filter out nodes that cannot be recomputed if (CanNotRecomputed(node)) { continue; } + + // TODO: Set the 'recompute' cnode attribute for the node + } +} + // Check if the node does not have forward output or if it has gradient inputs if (!HasForwardOutput(mng, node) || HasGradInputs(node, &has_grad_inputs_map)) { + + // If either condition is true, skip to the next iteration of the loop continue; } + // Cast the node to a CNodePtr auto cnode = node->cast(); + + // Check if the cast was successful MS_EXCEPTION_IF_NULL(cnode); + + // Get the primitive of the CNode auto prim = GetCNodePrimitive(cnode); + + // Check if the primitive is null if (prim == nullptr) { + // If the primitive is null, continue to the next iteration of the loop continue; } + + // Get the recompute attribute of the primitive auto prim_recompute_attr = prim->GetAttr(kAttrRecompute); + + // Initialize the recompute value to -1 int prim_recompute_val = -1; + + // Check if the recompute attribute exists and is of type BoolImm if (prim_recompute_attr != nullptr && prim_recompute_attr->isa()) { + // If it is, cast it to a bool and assign it to prim_recompute_val prim_recompute_val = static_cast(GetValue(prim_recompute_attr)); } + + // Check if the node should be recomputed based on the recompute attribute and the SetRecomputedScope function if ((SetRecomputedScope(cnode) && prim_recompute_val != 0) || prim_recompute_val == 1) { + // If it should be recomputed, add the recompute attribute to the CNode with a value of true cnode->AddAttr(kAttrRecompute, MakeValue(true)); } + + // Check if the CNode has the SetRecomputeCNodeAttr attribute set if (!IsSetRecomputeCNodeAttr(node)) { + // If it doesn't, continue to the next iteration of the loop continue; } - // Set attr for the tuple_getitem outputs. + + // Create a vector to store the output nodes of the tuple_getitem operations std::vector tuple_getitem_output_nodes; + + // Get the output nodes of the tuple_getitem operations using the GetTupleGetItemOutputNodes function GetTupleGetItemOutputNodes(mng, node, &tuple_getitem_output_nodes); + // Iterate over each element in the tuple_getitem_output_nodes vector for (const auto &output_node : tuple_getitem_output_nodes) { + + // Cast the output_node to a CNodePtr (a pointer to a CNode object) auto output_cnode = output_node->cast(); + + // Check if the output_cnode is null, throw an exception if it is MS_EXCEPTION_IF_NULL(output_cnode); + + // Add the attribute "kAttrRecompute" to the output_cnode and set its value to true output_cnode->AddAttr(kAttrRecompute, MakeValue(true)); } } } +// Function to create a new recomputed node in a given graph CNodePtr CreateNewRecomputedNode(const FuncGraphPtr &graph, const CNodePtr &origin_node, const std::vector &new_inputs) { + + // Create a new CNode using the provided new_inputs auto recomputed_node = graph->NewCNode(new_inputs); MS_EXCEPTION_IF_NULL(recomputed_node); + + // Add the "duplicated" attribute to the recomputed_node and set it to true recomputed_node->AddAttr("duplicated", MakeValue(true)); + + // Add the "need_cse_after_recompute" attribute to the recomputed_node and set it to true recomputed_node->AddAttr(kAttrNeedCseAfterRecompute, MakeValue(true)); + + // Set the abstract of the recomputed_node to the same as the origin_node recomputed_node->set_abstract(origin_node->abstract()); + + // Set the scope of the recomputed_node to the same as the origin_node recomputed_node->set_scope(origin_node->scope()); + + // If the origin_node has the "micro" attribute, add it to the recomputed_node as well if (origin_node->HasPrimalAttr(kAttrMicro)) { recomputed_node->AddPrimalAttr(kAttrMicro, origin_node->GetPrimalAttr(kAttrMicro)); } + + // Return the newly created recomputed_node return recomputed_node; } +// Create a new recomputed node based on the given origin node, with some modifications CNodePtr NewRecomputedNode(const FuncGraphPtr &graph, const CNodePtr &origin_node, const std::vector &first_target_inputs, const mindspore::HashSet &recomputed_origin_nodes, @@ -347,39 +657,86 @@ CNodePtr NewRecomputedNode(const FuncGraphPtr &graph, const CNodePtr &origin_nod MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(origin_node); MS_EXCEPTION_IF_NULL(origin_to_recomputed_nodes); + + // Check if the recomputed node already exists in the map auto iter = origin_to_recomputed_nodes->find(origin_node); if (iter != origin_to_recomputed_nodes->end()) { return iter->second; } + + // Log a debug message indicating the start of duplicating the origin recomputed node MS_LOG(DEBUG) << "Begin to Duplicating origin recomputed node: " << origin_node->DebugString(); + + // Create a vector to store the new inputs for the recomputed node std::vector new_inputs; + + // Flag to indicate if the recomputed node has recomputed inputs bool has_recomputed_inputs = false; + + // Iterate over the inputs of the origin node for (size_t i = 0; i < origin_node->size(); ++i) { auto input = origin_node->input(i); + + // Check if it is the first input and if it is the primitive "AllGather" if (i == 0 && IsPrimitive(input, prim::kPrimAllGather)) { auto prim = GetValueNode(input); auto instance_name = prim->instance_name(); + + // Check if the instance name contains "parallel_optimizer" bool is_from_parallel_optimizer = instance_name.find("parallel_optimizer") != std::string::npos; + // ... (continues) int64_t fusion_id = prim->HasAttr(kAttrFusion) ? GetValue(prim->GetAttr(kAttrFusion)) : 0; + // Check if the primitive has the attribute "kAttrFusion". If it does, get its value and assign it to fusion_id. Otherwise, assign 0 to fusion_id. + if (is_from_parallel_optimizer && fusion_id > 0) { + // Check if the flag "is_from_parallel_optimizer" is true and fusion_id is greater than 0. + + // Create a new shared pointer to a Primitive object with the name "kPrimAllGather". auto new_prim = std::make_shared(prim::kPrimAllGather->name()); + + // Set the attributes of the new_prim object to be the same as the attributes of the prim object. (void)new_prim->SetAttrs(prim->attrs()); + + // Set the attribute "kAttrFusion" of the new_prim object to fusion_id + fusion_id_increasement_size. new_prim->set_attr(kAttrFusion, MakeValue(fusion_id + fusion_id_increasement_size)); + + // Set the primitive type of the new_prim object to be the same as the primitive type of the prim object. new_prim->set_prim_type(prim->prim_type()); + + // Set the instance name of the new_prim object to instance_name. new_prim->set_instance_name(instance_name); + + // Create a new ValueNode object with the new_prim object as its value. auto value_node = NewValueNode(new_prim); + + // Add the value_node to the new_inputs vector. (void)new_inputs.emplace_back(value_node); + + // Continue to the next iteration of the loop. continue; } } + + // Check if the input is not a CNode. MS_EXCEPTION_IF_NULL(input); if (!input->isa()) { + + // Add the input to the new_inputs vector. (void)new_inputs.emplace_back(input); + + // Continue to the next iteration of the loop. continue; } + + // Cast the input to a CNode pointer. auto input_cnode = input->cast(); + + // Check if the input_cnode is not found in the recomputed_origin_nodes set. if (recomputed_origin_nodes.find(input_cnode) == recomputed_origin_nodes.end()) { + + // Check if the input_cnode is a primitive CNode with the primitive type "kPrimUpdateState". if (IsPrimitiveCNode(input_cnode, prim::kPrimUpdateState)) { + // ... (code continues) auto u = NewValueNode(kUMonad); u->set_abstract(kUMonad->ToAbstract()); (void)new_inputs.emplace_back(u); @@ -400,36 +757,92 @@ CNodePtr NewRecomputedNode(const FuncGraphPtr &graph, const CNodePtr &origin_nod MS_EXCEPTION_IF_NULL(first_input); std::vector depend_inputs{NewValueNode(prim::kPrimDepend), first_input, graph->NewCNode(make_tuple_inputs)}; + + // Create a vector of AnfNodePtr to store the inputs for the Depend operation + std::vector depend_inputs{NewValueNode(prim::kPrimDepend), first_input, + graph->NewCNode(make_tuple_inputs)}; + + // Add the Depend operation to the new_inputs vector + (void)new_inputs.emplace_back(graph->NewCNode(depend_inputs)); + // Create a new computation node in the graph using the NewCNode function of the graph object auto depend_node = graph->NewCNode(depend_inputs); MS_EXCEPTION_IF_NULL(depend_node); + + // Set the abstract of the depend_node to be the same as the abstract of the first input node depend_node->set_abstract(first_input->abstract()); + + // Add an attribute "recompute_depend" to the depend_node with a value of true depend_node->AddAttr("recompute_depend", MakeValue(true)); + + // Update the second input of the new_inputs array with the depend_node new_inputs[1] = depend_node; } + + // Create a new recomputed node using the CreateNewRecomputedNode function, passing in the graph, origin_node, and new_inputs auto recomputed_node = CreateNewRecomputedNode(graph, origin_node, new_inputs); + + // Add the mapping from the origin_node to the recomputed_node in the origin_to_recomputed_nodes map (void)origin_to_recomputed_nodes->emplace(origin_node, recomputed_node); + + // Return the recomputed_node return recomputed_node; } +// This function duplicates recomputed nodes in a given graph based on certain conditions and updates a hashmap with the mapping between original nodes and duplicated nodes. + +// Parameters: +// - graph: Pointer to the FuncGraph object representing the graph +// - target_nodes: HashSet of CNodePtr representing the target nodes to be duplicated +// - origin_recomputed_nodes: HashSet of CNodePtr representing the original recomputed nodes +// - first_target_inputs: Vector of AnfNodePtr representing the inputs of the first target node +// - origin_to_recomputed_nodes: Pointer to the HashMap that will store the mapping between original nodes and duplicated nodes + void DuplicateRecomputedNodes(const FuncGraphPtr &graph, const mindspore::HashSet &target_nodes, const mindspore::HashSet &origin_recomputed_nodes, const std::vector &first_target_inputs, mindspore::HashMap *origin_to_recomputed_nodes) { + // Check if the graph pointer is null MS_EXCEPTION_IF_NULL(graph); + + // Get the manager of the graph auto mng = graph->manager(); + + // Check if the manager pointer is null MS_EXCEPTION_IF_NULL(mng); + + // Iterate over each target node for (const auto &target_node : target_nodes) { + // Check if the target node pointer is null MS_EXCEPTION_IF_NULL(target_node); + + // Log a debug message with the debug string of the target node MS_LOG(DEBUG) << "Rebuild a new target_node " << target_node->DebugString() << " with the new recomputed input"; + + // Cast the target node to CNodePtr auto target_cnode = target_node->cast(); + + // Check if the target_cnode pointer is null MS_EXCEPTION_IF_NULL(target_cnode); + + // Create a vector to store the new target inputs std::vector new_target_inputs; + + // Iterate over each input of the target node for (const auto &input : target_cnode->inputs()) { + // Check if the input pointer is null MS_EXCEPTION_IF_NULL(input); + + // If the input is not a CNode, add it to the new target inputs vector if (!input->isa()) { (void)new_target_inputs.emplace_back(input); } else { + // If the input is a CNode, cast it to CNodePtr auto input_cnode = input->cast(); + + // Check if the input_cnode pointer is null + MS_EXCEPTION_IF_NULL(input_cnode); + + // Check if the input_cnode is present in the origin_recomputed_nodes set if (origin_recomputed_nodes.find(input_cnode) != origin_recomputed_nodes.end()) { (void)new_target_inputs.emplace_back(NewRecomputedNode(graph, input_cnode, first_target_inputs, origin_recomputed_nodes, origin_to_recomputed_nodes)); @@ -438,50 +851,93 @@ void DuplicateRecomputedNodes(const FuncGraphPtr &graph, const mindspore::HashSe } } } + + // Create a new CNode using the new_target_inputs as its inputs auto new_target_node = graph->NewCNode(new_target_inputs); + + // Clone the CNode information from the target_node to the new_target_node new_target_node->CloneCNodeInfo(target_node); + + // Add an attribute "target_grad" with the value true to the new_target_node new_target_node->AddAttr("target_grad", MakeValue(true)); + + // Set the scope of the new_target_node to be the same as the target_node new_target_node->set_scope(target_node->scope()); + + // Replace the target_node with the new_target_node in the manager mng->Replace(target_node, new_target_node); } } } // namespace +// Define a function named "InsertRecomputedNodes" that takes a reference to a `FuncGraphPtr` object as a parameter void InsertRecomputedNodes(const FuncGraphPtr &graph) { + // Check if the input graph is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(graph); + + // Get the manager of the graph auto mng = graph->manager(); + // Check if the manager is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(mng); + + // Get the ordered list of CNodes in the graph std::list orders = graph->GetOrderedCnodes(); + // Convert the list to a vector std::vector origin_nodes_topological(orders.begin(), orders.end()); + + // Set the recomputed attribute for the origin nodes SetRecomputedAttr(graph, origin_nodes_topological); - // Get candidate origin recomputed nodes which have no grad inputs and output to at least one grad node directly. + + // Find candidate origin recomputed nodes which have no grad inputs and output to at least one grad node directly std::vector candidate_recomputed_nodes = FindCandidateRecomputedNodes(mng, origin_nodes_topological); + + // Create a hash set to keep track of visited nodes mindspore::HashSet visited_nodes; + + // Iterate over each candidate recomputed node for (const auto &candidate_recomputed_node : candidate_recomputed_nodes) { + // Check if the candidate node has already been visited, and skip it if it has if (visited_nodes.find(candidate_recomputed_node) != visited_nodes.end()) { continue; } + + // Create a hash set to store the max continuous recomputed sub-graph mindspore::HashSet max_recomputed_sub_graph = {candidate_recomputed_node}; - // Get max continuous recomputed sub-graph. + + // Get the max continuous recomputed sub-graph GetMaxSubGraph(mng, &max_recomputed_sub_graph, true, true); + + // Add the nodes in the max sub-graph to the visited nodes set visited_nodes.insert(max_recomputed_sub_graph.begin(), max_recomputed_sub_graph.end()); - // Get the origin recomputed nodes which directly output to the grad nodes. + + // Create a hash set to store the origin recomputed nodes which directly output to the grad nodes mindspore::HashSet origin_recomputed_nodes; + // ... (code continues) + // Create a HashSet of CNodePtr named target_nodes mindspore::HashSet target_nodes; + + // Call the function GetOriginRecomputeAndTargetNodes with arguments mng, max_recomputed_sub_graph, + // and pointers to origin_recomputed_nodes and target_nodes to get the origin recomputed nodes and target nodes GetOriginRecomputeAndTargetNodes(mng, max_recomputed_sub_graph, &origin_recomputed_nodes, &target_nodes); - // Also get the inputs of origin recomputed nodes which eventually output to the grad nodes. + + // Call the function GetMaxSubGraph with arguments mng, pointer to origin_recomputed_nodes, true, and false + // to get the maximum subgraph GetMaxSubGraph(mng, &origin_recomputed_nodes, true, false); // Get the inputs of the first target node in the topological sequence. The duplicated recomputed nodes should // not be executed until these inputs are ready. std::vector first_target_inputs = GetFirstTargetInputs(origin_nodes_topological, origin_recomputed_nodes, target_nodes); + + // Create a hashmap to store the mapping between original recomputed nodes and duplicated recomputed nodes mindspore::HashMap origin_to_recomputed_nodes; - // Begin duplicate origin recomputed nodes with each target node. + + // Begin duplicating origin recomputed nodes with each target node DuplicateRecomputedNodes(graph, target_nodes, origin_recomputed_nodes, first_target_inputs, &origin_to_recomputed_nodes); } - // Set need cse attr for doing cse after recompute. + + // Set the "need cse" attribute for doing common subexpression elimination (CSE) after recompute for (const auto &node : orders) { if (WithRecomputedScope(node)) { node->AddAttr(kAttrNeedCseAfterRecompute, MakeValue(true)); @@ -489,4 +945,4 @@ void InsertRecomputedNodes(const FuncGraphPtr &graph) { } } } // namespace opt -} // namespace mindspore +} // namespace mindspore \ No newline at end of file -- 2.34.1