diff --git a/mindspore/ccsrc/runtime/collective/collective_comm_lib_loader.cc b/mindspore/ccsrc/runtime/collective/collective_comm_lib_loader.cc index 5705a456ba4..5886170423b 100644 --- a/mindspore/ccsrc/runtime/collective/collective_comm_lib_loader.cc +++ b/mindspore/ccsrc/runtime/collective/collective_comm_lib_loader.cc @@ -14,38 +14,65 @@ * limitations under the License. */ +// Include the header file "runtime/collective/collective_comm_lib_loader.h" which contains the declarations for the collective communication library loader. #include "runtime/collective/collective_comm_lib_loader.h" +// Start of the "mindspore" namespace namespace mindspore { -namespace device { -bool CollectiveCommLibLoader::Initialize() { - std::string err_msg = ""; -#ifndef _WIN32 - collective_comm_lib_ptr_ = dlopen(comm_lib_name_.c_str(), RTLD_LAZY); - err_msg = GetDlErrorMsg(); -#else - collective_comm_lib_ptr_ = LoadLibrary(comm_lib_name_.c_str()); - err_msg = std::to_string(GetLastError()); -#endif - if (collective_comm_lib_ptr_ == nullptr) { - MS_LOG(EXCEPTION) << "Loading " + comm_lib_name_ + " failed. Error: " + err_msg; - } - return true; + // Start of the "device" namespace + namespace device { + // Definition of the "Initialize" function of the "CollectiveCommLibLoader" class + bool CollectiveCommLibLoader::Initialize() { + // Initialize an empty error message string + std::string err_msg = ""; + + // Check if the platform is not Windows + #ifndef _WIN32 + // Load the collective communication library using dlopen and store the result in collective_comm_lib_ptr_ + collective_comm_lib_ptr_ = dlopen(comm_lib_name_.c_str(), RTLD_LAZY); + // Get the error message from dlopen and store it in err_msg + err_msg = GetDlErrorMsg(); + // If the platform is Windows + #else + // Load the collective communication library using LoadLibrary and store the result in collective_comm_lib_ptr_ + collective_comm_lib_ptr_ = LoadLibrary(comm_lib_name_.c_str()); + // Get the last error code as a string and store it in err_msg + err_msg = std::to_string(GetLastError()); + // End of the platform check + #endif + + // Check if the collective_comm_lib_ptr_ is nullptr, indicating a failure in loading the library + if (collective_comm_lib_ptr_ == nullptr) { + // Log an exception with the error message + MS_LOG(EXCEPTION) << "Loading " + comm_lib_name_ + " failed. Error: " + err_msg; + } + + // Return true to indicate successful initialization + return true; + } + } // End of the "device" namespace +} // End of the "mindspore" namespace + +// Define the function `Finalize` for the `CollectiveCommLibLoader` class +bool CollectiveCommLibLoader::Finalize() { + // Check if the `collective_comm_lib_ptr_` is null, and throw an exception if it is + MS_EXCEPTION_IF_NULL(collective_comm_lib_ptr_); } -bool CollectiveCommLibLoader::Finalize() { - MS_EXCEPTION_IF_NULL(collective_comm_lib_ptr_); - #ifndef _WIN32 + // If the operating system is not Windows, use dlclose to close the dynamic library handle if (dlclose(collective_comm_lib_ptr_) != 0) { + // If dlclose fails, throw an exception with an error message MS_LOG(EXCEPTION) << "Closing " + comm_lib_name_ + " handle failed. Error: " + GetDlErrorMsg(); } #else + // If the operating system is Windows, use FreeLibrary to close the dynamic library handle if (!FreeLibrary(reinterpret_cast(collective_comm_lib_ptr_))) { + // If FreeLibrary fails, throw an exception with an error message MS_LOG(EXCEPTION) << "Closing " + comm_lib_name_ + " handle failed. Error: " + std::to_string(GetLastError()); } #endif - return true; -} -} // namespace device -} // namespace mindspore +// Return true to indicate successful closing of the library handle +return true; +} // End of namespace device +} // End of namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/collective/collective_communication_lib.cc b/mindspore/ccsrc/runtime/collective/collective_communication_lib.cc index 6411b96af99..7f451a2ad30 100644 --- a/mindspore/ccsrc/runtime/collective/collective_communication_lib.cc +++ b/mindspore/ccsrc/runtime/collective/collective_communication_lib.cc @@ -14,67 +14,119 @@ * limitations under the License. */ +// Include the header file "collective_communication_lib.h" from the "runtime/collective" directory. #include "runtime/collective/collective_communication_lib.h" +// Start of the namespace "mindspore" namespace mindspore { + // Start of the namespace "device" namespace device { +// Definition of the function "Finalize" in the class "CollectiveCommunicationLib" bool CollectiveCommunicationLib::Finalize() { +// Check if the library is not initialized or already finalized if (!initialized_ || finalized_.load()) { return true; } + // Iterate over each element in the 'groups_' map using a range-based for loop for (const auto &group : groups_) { + // Check if the value of the current element is null CHECK_IF_NULL(group.second); + // Call the 'Finalize' function of the current group and check if it returns false if (!group.second->Finalize()) { + // If 'Finalize' returns false, return false from the main function return false; } } + // Clear the 'groups_' map groups_.clear(); + // Set the 'initialized_' flag to false initialized_ = false; + // Set the 'finalized_' flag to true finalized_ = true; + // Return true to indicate successful completion of the main function return true; -} +// Function to destroy a communication group based on its name bool CollectiveCommunicationLib::DestroyCommunicationGroup(const std::string &group_name) { + // Check if the group exists in the map of groups if (groups_.count(group_name) == 0) { - return false; + return false; // Return false if the group does not exist } - auto group = groups_[group_name]; - CHECK_IF_NULL(group); - if (!group->Finalize()) { - return false; + auto group = groups_[group_name]; // Get the group object from the map + CHECK_IF_NULL(group); // Check if the group object is null + if (!group->Finalize()) { // Call the Finalize function of the group object and check if it returns false + return false; // Return false if the Finalize function fails } - (void)groups_.erase(group_name); - return true; + (void)groups_.erase(group_name); // Remove the group from the map + return true; // Return true to indicate successful destruction of the group } +// Function to get the rank ID of a group in the CollectiveCommunicationLib class uint32_t CollectiveCommunicationLib::GetRankId(const std::string &group_name) { + // Check if the group exists in the groups_ map CHECK_RET(groups_.count(group_name) != 0, true, "The group " + group_name + " does not exist."); + + // Get the group object from the groups_ map auto group = groups_[group_name]; + + // Check if the group object is not null CHECK_IF_NULL(group); + + // Return the rank ID of the group using the GetGroupRank() function of the group object return group->GetGroupRank(global_rank_id_); } +// Function to get the size of a group in the CollectiveCommunicationLib class uint32_t CollectiveCommunicationLib::GetGroupSize(const std::string &group_name) { + // Check if the group exists in the groups_ map CHECK_RET(groups_.count(group_name) != 0, true, "The group " + group_name + " does not exist."); + + // Get the group object from the groups_ map auto group = groups_[group_name]; + + // Check if the group object is null CHECK_IF_NULL(group); + + // Return the size of the group return group->group_size(); } +// Function to get a communication group based on its name CommunicationGroupPtr CollectiveCommunicationLib::GetGroup(const std::string &group_name) { + // Check if the group exists in the map of groups if (groups_.count(group_name) == 0) { + // If the group does not exist, return a nullptr return nullptr; } + // If the group exists, return a pointer to the group return groups_[group_name]; } -const std::string &CollectiveCommunicationLib::global_group_name() const { return global_group_name_; } +// Define a member function named "global_group_name" in the "CollectiveCommunicationLib" class +// The function returns a constant reference to a std::string object +const std::string &CollectiveCommunicationLib::global_group_name() const { + // Return the value of the member variable "global_group_name_" + return global_group_name_; +} -uint32_t CollectiveCommunicationLib::global_rank_id() const { return global_rank_id_; } +// Define a member function named "global_rank_id" in the "CollectiveCommunicationLib" class +uint32_t CollectiveCommunicationLib::global_rank_id() const { + // Return the value of the member variable "global_rank_id_" + return global_rank_id_; +} -uint32_t CollectiveCommunicationLib::local_rank_id() const { return local_rank_id_; } +// This is a member function of the CollectiveCommunicationLib class +// It returns the value of the private member variable local_rank_id_ +// The return type is uint32_t +uint32_t CollectiveCommunicationLib::local_rank_id() const { + return local_rank_id_; +} -uint32_t CollectiveCommunicationLib::global_rank_size() const { return global_rank_size_; } -} // namespace device -} // namespace mindspore +// Define the function `global_rank_size` in the `CollectiveCommunicationLib` namespace, returning a `uint32_t` value +uint32_t CollectiveCommunicationLib::global_rank_size() const { + return global_rank_size_; // Return the value of the member variable `global_rank_size_` +} + +} // End of the `device` namespace +} // End of the `mindspore` namespace \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/collective/communication_group.cc b/mindspore/ccsrc/runtime/collective/communication_group.cc index ea0a0184bbf..54d60a58bda 100644 --- a/mindspore/ccsrc/runtime/collective/communication_group.cc +++ b/mindspore/ccsrc/runtime/collective/communication_group.cc @@ -16,42 +16,80 @@ #include "runtime/collective/communication_group.h" +// Start of the "mindspore" namespace namespace mindspore { -namespace device { -CommunicationGroup::CommunicationGroup(const std::string &name, const std::vector &group_ranks, - uint32_t global_rank) - : initialized_(false), - global_rank_(global_rank), - size_(group_ranks.size()), - name_(name), - group_ranks_(group_ranks) { - uint32_t group_rank = 0; - // The input group_ranks contains the global ranks of the processes in this group. - (void)std::for_each(group_ranks.begin(), group_ranks.end(), [&](const uint32_t &global_rank) { - global_to_group_ranks_[global_rank] = group_rank; - group_to_global_ranks_[group_rank] = global_rank; - group_rank++; - }); -} + // Start of the "device" namespace + namespace device { + // Constructor for the CommunicationGroup class + CommunicationGroup::CommunicationGroup(const std::string &name, const std::vector &group_ranks, + uint32_t global_rank) + : initialized_(false), // Initialize the "initialized_" member variable to false + global_rank_(global_rank), // Initialize the "global_rank_" member variable with the provided value + size_(group_ranks.size()), // Initialize the "size_" member variable with the size of the "group_ranks" vector + name_(name), // Initialize the "name_" member variable with the provided value + group_ranks_(group_ranks) { // Initialize the "group_ranks_" member variable with the provided vector + uint32_t group_rank = 0; // Initialize the "group_rank" variable to 0 + // Iterate over each element in the "group_ranks" vector using a lambda function + (void)std::for_each(group_ranks.begin(), group_ranks.end(), [&](const uint32_t &global_rank) { + global_to_group_ranks_[global_rank] = group_rank; // Map the global rank to the group rank + group_to_global_ranks_[group_rank] = global_rank; // Map the group rank to the global rank + group_rank++; // Increment the group rank + }); + } + } // End of the "device" namespace +} // End of the "mindspore" namespace + +// This function is a member function of the CommunicationGroup class +// It takes a uint32_t parameter named global_rank and returns a uint32_t value uint32_t CommunicationGroup::GetGroupRank(uint32_t global_rank) { + // Check if the global_rank exists in the global_to_group_ranks_ map + // If it doesn't exist, throw an exception with an error message CHECK_RET((global_to_group_ranks_.count(global_rank) != 0), true, "Group " + name_ + " doesn't contain the global rank " + std::to_string(global_rank)); + + // If the global_rank exists in the map, return the corresponding group rank return global_to_group_ranks_[global_rank]; } +// This function is a member function of the CommunicationGroup class +// It takes a group_rank as input and returns the corresponding global rank + uint32_t CommunicationGroup::GetGlobalRank(uint32_t group_rank) { + // Check if the group_rank exists in the group_to_global_ranks_ map CHECK_RET((group_to_global_ranks_.count(group_rank) != 0), true, "Group " + name_ + " doesn't contain the group rank " + std::to_string(group_rank)); + + // If the group_rank exists, return the corresponding global rank from the map return group_to_global_ranks_[group_rank]; } -uint32_t CommunicationGroup::group_size() const { return size_; } +// Define a member function named "group_size" for the class "CommunicationGroup" +// The function returns a value of type "uint32_t" +uint32_t CommunicationGroup::group_size() const { + // Return the value of the member variable "size_" + return size_; +} -const std::vector &CommunicationGroup::group_ranks() const { return group_ranks_; } +// Define a member function named "group_ranks" for the class "CommunicationGroup" +// The function returns a constant reference to a vector of unsigned 32-bit integers +const std::vector &CommunicationGroup::group_ranks() const { + // Return the private member variable "group_ranks_" + return group_ranks_; +} -const std::map &CommunicationGroup::global_to_group_ranks() const { return global_to_group_ranks_; } +// Define a constant member function named global_to_group_ranks() that returns a reference to a constant std::map +const std::map &CommunicationGroup::global_to_group_ranks() const { + // Return the member variable global_to_group_ranks_ + return global_to_group_ranks_; +} + +// Define a constant member function named "group_to_global_ranks" in the "CommunicationGroup" class +// This function returns a reference to a constant std::map object with key type uint32_t and value type uint32_t +const std::map &CommunicationGroup::group_to_global_ranks() const { + // Return the private member variable "group_to_global_ranks_" + return group_to_global_ranks_; +} -const std::map &CommunicationGroup::group_to_global_ranks() const { return group_to_global_ranks_; } } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/device/kernel_runtime.cc b/mindspore/ccsrc/runtime/device/kernel_runtime.cc index 21c185df051..71a3000ac2f 100644 --- a/mindspore/ccsrc/runtime/device/kernel_runtime.cc +++ b/mindspore/ccsrc/runtime/device/kernel_runtime.cc @@ -14,712 +14,1433 @@ * limitations under the License. */ +// Include the header file for the kernel runtime of the device #include "runtime/device/kernel_runtime.h" + +// Include the functional header for using std::function #include + +// Include the utility header for using std::pair and std::vector #include #include + +// Include the set header for using std::set #include + +// Include the helper header for optimizer functions #include "backend/common/optimizer/helper.h" + +// Include the ANF runtime algorithm header for working with ANF runtime #include "backend/common/session/anf_runtime_algorithm.h" + +// Include the dynamic shape helper header for working with dynamic shape optimization #include "backend/common/optimizer/dynamic_shape/dynamic_shape_helper.h" + +// Include the ANF algorithm header for working with ANF nodes #include "include/common/utils/anfalgo.h" + +// Include the kernel graph header for working with kernel graphs #include "backend/common/session/kernel_graph.h" + +// Include the device shape transfer header for transferring device shapes #include "runtime/device/ms_device_shape_transfer.h" + +// Include the op runtime info header for op runtime information #include "runtime/pynative/op_runtime_info.h" + +// Include the kernel runtime manager header for managing kernel runtimes #include "runtime/device/kernel_runtime_manager.h" + +// Include the dump JSON parser header for parsing dump JSON files #include "debug/data_dump/dump_json_parser.h" + +// Include the operator header for working with operators #include "frontend/operator/ops.h" + +// Include the value header for working with IR values #include "ir/value.h" + +// Include the MS context header for working with MindSpore context #include "utils/ms_context.h" + +// Include the MS utils header for working with MindSpore utilities #include "utils/ms_utils.h" + +// Include the shape utils header for working with shape utilities #include "utils/shape_utils.h" + +// Include the utils header for general utilities #include "include/common/utils/utils.h" +// Include the header file for the parallel context utility #include "include/common/utils/parallel_context.h" + +// Include the header file for the environment configuration parser #include "include/common/debug/env_config_parser.h" + +// Include the header file for the Ascend device address in the Ascend device HAL #include "plugin/device/ascend/hal/device/ascend_device_address.h" + +// Check if the ENABLE_CPU macro is defined and the _WIN32 macro is not defined #if ((defined ENABLE_CPU) && (!defined _WIN32)) -#include "ps/ps_cache/ps_cache_manager.h" + + // Include the header file for the PS cache manager + #include "ps/ps_cache/ps_cache_manager.h" + #endif +// Import the Address and AddressPtr classes from the mindspore::kernel namespace using mindspore::kernel::Address; using mindspore::kernel::AddressPtr; +// Define the namespace "mindspore" namespace mindspore { -namespace device { -constexpr size_t kAtomicCleanInputSize = 2; -namespace { -std::vector GetGraphInputs(const session::KernelGraph &graph) { - auto graph_inputs = graph.inputs(); - std::vector result(graph_inputs.begin(), graph_inputs.end()); - std::set inputs_set(graph_inputs.begin(), graph_inputs.end()); - auto kernels = graph.execution_order(); - for (auto &kernel : kernels) { - MS_EXCEPTION_IF_NULL(kernel); - auto input_num = common::AnfAlgo::GetInputTensorNum(kernel); - for (size_t i = 0; i < input_num; ++i) { - auto input_node = kernel->input(i + 1); - auto input_real_node = common::AnfAlgo::VisitKernelWithReturnType(input_node, 0).first; - MS_EXCEPTION_IF_NULL(input_real_node); - if (input_real_node->isa() && inputs_set.find(input_real_node) == inputs_set.end()) { - (void)inputs_set.insert(input_real_node); - (void)result.emplace_back(input_real_node); + + // Define the namespace "device" within the "mindspore" namespace + namespace device { + + // Define a constant variable "kAtomicCleanInputSize" with a value of 2 + constexpr size_t kAtomicCleanInputSize = 2; + + // Define an anonymous namespace to limit the scope of the following functions + namespace { + + // Function to get the inputs of a given KernelGraph + std::vector GetGraphInputs(const session::KernelGraph &graph) { + + // Get the inputs of the graph + auto graph_inputs = graph.inputs(); + + // Create a vector to store the result + std::vector result(graph_inputs.begin(), graph_inputs.end()); + + // Create a set to keep track of the input nodes + std::set inputs_set(graph_inputs.begin(), graph_inputs.end()); + + // Get the execution order of the kernels in the graph + auto kernels = graph.execution_order(); + + // Iterate over the kernels + for (auto &kernel : kernels) { + MS_EXCEPTION_IF_NULL(kernel); + + // Get the number of input tensors for the current kernel + auto input_num = common::AnfAlgo::GetInputTensorNum(kernel); + + // Iterate over the input tensors + for (size_t i = 0; i < input_num; ++i) { + + // Get the input node + auto input_node = kernel->input(i + 1); + + // Get the real input node (ignoring any virtual nodes) + auto input_real_node = common::AnfAlgo::VisitKernelWithReturnType(input_node, 0).first; + MS_EXCEPTION_IF_NULL(input_real_node); + + // Check if the input node is a Parameter and not already in the inputs set + if (input_real_node->isa() && inputs_set.find(input_real_node) == inputs_set.end()) { + + // Add the input node to the inputs set and the result vector + (void)inputs_set.insert(input_real_node); + (void)result.emplace_back(input_real_node); + } + } + } + + // Return the result vector + return result; } + + } // End of anonymous namespace + + } // End of namespace device + +} // End of namespace mindspore } } return result; } } // namespace + +// Define a constant variable for the minimum input size constexpr size_t kMinInputSize = 2; + +// Destructor for the KernelRuntime class KernelRuntime::~KernelRuntime() { + // Set the stream pointers to nullptr to release any resources stream_ = nullptr; independent_stream_ = nullptr; communication_stream_ = nullptr; } +// Define a member function named LockRuntime in the class KernelRuntime that returns a std::lock_guard object std::lock_guard KernelRuntime::LockRuntime() { + + // Define a static std::mutex object named mutex, which will be shared among all instances of KernelRuntime static std::mutex mutex; + + // Create a std::lock_guard object using the mutex and return it return std::lock_guard(mutex); } +// Implementation of the Load function in the KernelRuntime class + bool KernelRuntime::Load(const session::KernelGraph &, bool) { + + // Log an informational message using the MS_LOG macro MS_LOG(INFO) << "Call default load."; + + // Return true to indicate successful loading return true; } +// Implementation of the LoadData function in the KernelRuntime class + bool KernelRuntime::LoadData(const session::KernelGraph &) { + + // Log an informational message using the MS_LOG macro MS_LOG(INFO) << "Call default load data."; + + // Return false to indicate that the data loading was not successful return false; } +// Check if the output device address exists for a given kernel and index bool KernelRuntime::NodeOutputDeviceAddressExist(const AnfNodePtr &kernel, size_t index) { MS_EXCEPTION_IF_NULL(kernel); + + // Check if the output address exists for the given kernel and index if (AnfAlgo::OutputAddrExist(kernel, index)) { + // Get the output address for the given kernel and index const auto &address = AnfAlgo::GetOutputAddr(kernel, index); MS_EXCEPTION_IF_NULL(address); + + // Check if the device type of the output address matches the target device address type return address->DeviceType() == GetTargetDeviceAddressType(); } + + // Return false if the output address does not exist return false; } +// A function to assign memory to the nodes in a kernel graph void KernelRuntime::AssignMemory(const session::KernelGraph &graph) { + + // Get the instance of the MsContext class auto context_ptr = MsContext::GetInstance(); + + // Throw an exception if the context pointer is null MS_EXCEPTION_IF_NULL(context_ptr); + + // Check if memory scheduler is enabled if (UseMemScheduler()) { + + // Assign static memory to value nodes in the graph AssignStaticMemoryValueNode(graph); + + // Reset the addresses of the nodes in the graph ResetNodeAddress(graph); + + // Assign communication memory to the nodes in the graph AssignCommunicationMem(graph); + } else { + + // Throw an exception if the memory manager pointer is null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Reset the dynamic memory managed by the memory manager mem_manager_->ResetDynamicMemory(); + + // Assign static memory to the nodes in the graph AssignStaticMemory(graph); + + // Assign dynamic memory to the nodes in the graph AssignDynamicMemory(graph); } + + // Update the output memory of reference nodes in the graph UpdateRefNodeOutputMem(graph); } -void KernelRuntime::GetCommunicationInputInfo(const AnfNodePtr &node, size_t *total_size, - DeviceAddressPtrList *address_list, - std::vector *align_size_list) const { - MS_EXCEPTION_IF_NULL(node); - MS_EXCEPTION_IF_NULL(total_size); - MS_EXCEPTION_IF_NULL(address_list); - MS_EXCEPTION_IF_NULL(align_size_list); - size_t input_num = common::AnfAlgo::GetInputTensorNum(node); - for (size_t i = 0; i < input_num; ++i) { +// This function is a member function of the `KernelRuntime` class. +// It is used to get information about the input tensors of a given `node`. +// The information includes the total size of the input tensors, a list of device addresses for each input tensor, +// and a list of alignment sizes for each input tensor. + +// Check if the `node` pointer is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(node); + +// Check if the `total_size` pointer is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(total_size); + +// Check if the `address_list` pointer is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(address_list); + +// Check if the `align_size_list` pointer is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(align_size_list); + +// Get the number of input tensors for the given `node` +size_t input_num = common::AnfAlgo::GetInputTensorNum(node); + +// Iterate over each input tensor +for (size_t i = 0; i < input_num; ++i) { + // Get the previous node and output index for the current input tensor auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(node, i, true); auto input_node = input_node_with_index.first; + + // Check if the input node pointer is null, and throw an exception if it is null MS_EXCEPTION_IF_NULL(input_node); + + // Create a null device address pointer DeviceAddressPtr address = nullptr; + + // Check if the output address exists for the input node and output index if (AnfAlgo::OutputAddrExist(input_node, input_node_with_index.second)) { - address = AnfAlgo::GetMutableOutputAddr(input_node, input_node_with_index.second); + // If the output address exists, get the mutable output address for the input node and output index + address = AnfAlgo::GetMutableOutputAddr(input_node, input_node_with_index.second); } else { - address = PreAssignCNodeMemory(input_node, input_node_with_index.second); + // If the output address does not exist, pre-assign memory for the input node and output index + address = PreAssignCNodeMemory(input_node, input_node_with_index.second); } + + // Check if the address pointer is null, and throw an exception if it is null MS_EXCEPTION_IF_NULL(address); + + // Get the alignment size for the address size auto align_size = MemoryManager::GetCommonAlignSize(address->size()); + + // Add the address and align size to the respective lists + address_list->push_back(address); + align_size_list->push_back(align_size); +} + // Add the align_size to the value pointed by total_size *total_size += align_size; + + // Add the address to the end of the address_list vector address_list->emplace_back(address); + + // Add the align_size to the end of the align_size_list vector align_size_list->emplace_back(align_size); } } +// Define a function named "AssignCommunicationInputFromMemoryPool" that takes a reference to an AnfNodePtr as a parameter void KernelRuntime::AssignCommunicationInputFromMemoryPool(const AnfNodePtr &node) const { + + // Check if the given node is a communication operation using the IsCommunicationOp function from the AnfAlgo class if (!common::AnfAlgo::IsCommunicationOp(node)) { - return; + return; // If it is not a communication operation, return from the function } + + // Throw an exception if the given node is null MS_EXCEPTION_IF_NULL(node); + + // Throw an exception if the memory manager is null MS_EXCEPTION_IF_NULL(mem_manager_); - - size_t total_size = 0; - DeviceAddressPtrList address_list; - std::vector align_size_list; - GetCommunicationInputInfo(node, &total_size, &address_list, &align_size_list); - if (align_size_list.empty()) { - MS_LOG(WARNING) << "No inputs for " << node->fullname_with_scope(); - return; - } - - if (!mem_manager_->MallocContinuousMemFromMemPool(address_list, total_size, align_size_list)) { - MS_LOG(EXCEPTION) << "Allocate continuous memory failed, totol_size:" << total_size; - } } +// Declare a variable to store the total size +size_t total_size = 0; + +// Declare a list to store device addresses +DeviceAddressPtrList address_list; + +// Declare a vector to store alignment sizes +std::vector align_size_list; + +// Call the function GetCommunicationInputInfo and pass the node, total_size, address_list, and align_size_list as arguments +GetCommunicationInputInfo(node, &total_size, &address_list, &align_size_list); + +// Check if the align_size_list is empty +if (align_size_list.empty()) { + // If it is empty, print a warning message indicating that there are no inputs for the node's fullname_with_scope + MS_LOG(WARNING) << "No inputs for " << node->fullname_with_scope(); + + // Return from the function + return; +} + +// Check if the memory manager is unable to allocate continuous memory from the memory pool +if (!mem_manager_->MallocContinuousMemFromMemPool(address_list, total_size, align_size_list)) { + + // If allocation fails, throw an exception with an error message indicating the total size of memory requested + MS_LOG(EXCEPTION) << "Allocate continuous memory failed, total_size:" << total_size; +} +// End of the if statement block + +// Define the function `GetCommunicationOutputInfo` of the class `KernelRuntime` void KernelRuntime::GetCommunicationOutputInfo(const AnfNodePtr &node, size_t *total_size, DeviceAddressPtrList *address_list, std::vector *align_size_list) const { + // Check if the input arguments are null, and throw exceptions if they are MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(total_size); MS_EXCEPTION_IF_NULL(align_size_list); MS_EXCEPTION_IF_NULL(address_list); - - const auto kernel_mod = AnfAlgo::GetKernelMod(node); - MS_EXCEPTION_IF_NULL(kernel_mod); - const auto output_size_list = kernel_mod->GetOutputSizeList(); - for (size_t i = 0; i < output_size_list.size(); ++i) { - DeviceAddressPtr address = nullptr; - if (AnfAlgo::OutputAddrExist(node, i)) { - address = AnfAlgo::GetMutableOutputAddr(node, i); - } else { - const std::string output_format = AnfAlgo::GetOutputFormat(node, i); - const auto output_type = AnfAlgo::GetOutputDeviceDataType(node, i); - const auto tensor_size = AnfAlgo::GetOutputTensorMemSize(node, i); - address = CreateDeviceAddress(nullptr, tensor_size, output_format, output_type, {node, i}); - AnfAlgo::SetOutputAddr(address, i, node.get()); - } - MS_EXCEPTION_IF_NULL(address); - auto align_size = MemoryManager::GetCommonAlignSize(address->size()); - *total_size += align_size; - align_size_list->emplace_back(align_size); - address_list->emplace_back(address); - } + // ... } +// Get the kernel module associated with the given node +const auto kernel_mod = AnfAlgo::GetKernelMod(node); +MS_EXCEPTION_IF_NULL(kernel_mod); + +// Get the list of output sizes from the kernel module +const auto output_size_list = kernel_mod->GetOutputSizeList(); + +// Iterate over the output size list +for (size_t i = 0; i < output_size_list.size(); ++i) { + // Create a null pointer for the device address + DeviceAddressPtr address = nullptr; + + // Check if the output address exists for the given node and index + if (AnfAlgo::OutputAddrExist(node, i)) { + // If the output address exists, get the mutable output address + address = AnfAlgo::GetMutableOutputAddr(node, i); + } else { + // If the output address does not exist, create a new device address + const std::string output_format = AnfAlgo::GetOutputFormat(node, i); + const auto output_type = AnfAlgo::GetOutputDeviceDataType(node, i); + const auto tensor_size = AnfAlgo::GetOutputTensorMemSize(node, i); + address = CreateDeviceAddress(nullptr, tensor_size, output_format, output_type, {node, i}); + AnfAlgo::SetOutputAddr(address, i, node.get()); + } + + // Check if the device address is null + MS_EXCEPTION_IF_NULL(address); + + // Get the aligned size for the device address + auto align_size = MemoryManager::GetCommonAlignSize(address->size()); + + // Add the aligned size to the total size + *total_size += align_size; + + // Add the aligned size to the align size list + align_size_list->emplace_back(align_size); + + // Add the device address to the address list + address_list->emplace_back(address); +} +// Closing brace to end the main function +} + +// Define a function named "AssignCommunicationOutputFromMemoryPool" that takes a reference to an AnfNodePtr as a parameter void KernelRuntime::AssignCommunicationOutputFromMemoryPool(const AnfNodePtr &node) const { + + // Check if the given node is a communication operation, if not, return from the function if (!common::AnfAlgo::IsCommunicationOp(node)) { return; } + + // Check if the given node is null, if so, throw an exception MS_EXCEPTION_IF_NULL(node); + + // Check if the memory manager is null, if so, throw an exception MS_EXCEPTION_IF_NULL(mem_manager_); - - size_t total_size = 0; - std::vector align_size_list; - std::vector address_list; - GetCommunicationOutputInfo(node, &total_size, &address_list, &align_size_list); - if (align_size_list.empty()) { - MS_LOG(WARNING) << "No output for " << node->fullname_with_scope(); - return; - } - - if (!mem_manager_->MallocContinuousMemFromMemPool(address_list, total_size, align_size_list)) { - MS_LOG(EXCEPTION) << "Allocate continuous memory failed, totol_size:" << total_size; - } } -void KernelRuntime::RunOpMallocPre(const session::KernelGraph &graph, - const std::vector &input_tensors) { - const auto &nodes = graph.execution_order(); - // Malloc for Node output - for (const auto &node : nodes) { - auto output_num = common::AnfAlgo::GetOutputTensorNum(node); - for (size_t i = 0; i < output_num; ++i) { - MS_EXCEPTION_IF_NULL(node); - auto runtime_info = node->user_data(); - MS_EXCEPTION_IF_NULL(runtime_info); - auto const &output_format = runtime_info->output_format(i); - auto output_type = runtime_info->output_type(i); - auto tensor_size = runtime_info->output_tensor_size(i); - // Create DeviceAddress without ptr. - // Get real device ptr after KernelBuild finish. - auto device_address = CreateDeviceAddress(nullptr, tensor_size, output_format, output_type); - device_address->set_host_shape(trans::GetRuntimePaddingShape(node, i)); - AnfAlgo::SetOutputAddr(device_address, i, node.get()); - } - } +// Declare a variable to store the total size +size_t total_size = 0; - // Malloc for graph input +// Declare two vectors to store the alignment size and device addresses +std::vector align_size_list; +std::vector address_list; + +// Call the function GetCommunicationOutputInfo and pass the node, total_size, address_list, and align_size_list as arguments +GetCommunicationOutputInfo(node, &total_size, &address_list, &align_size_list); + +// Check if the align_size_list is empty +if (align_size_list.empty()) { + // If it is empty, print a warning message indicating that there is no output for the given node's fullname with scope + MS_LOG(WARNING) << "No output for " << node->fullname_with_scope(); + + // Return from the function + return; +} + +// Check if the memory manager is unable to allocate continuous memory from the memory pool +if (!mem_manager_->MallocContinuousMemFromMemPool(address_list, total_size, align_size_list)) { + + // If allocation fails, throw an exception with an error message indicating the total size of memory requested + MS_LOG(EXCEPTION) << "Allocate continuous memory failed, total_size:" << total_size; +} +// End of the if statement block + +// This function is called before running the kernel to allocate memory for the output tensors of each node in the graph. +// It takes in the graph and a vector of input tensors as parameters. + +// Get the execution order of the nodes in the graph +const auto &nodes = graph.execution_order(); + +// Iterate over each node in the execution order +for (const auto &node : nodes) { + // Get the number of output tensors for the current node + auto output_num = common::AnfAlgo::GetOutputTensorNum(node); + + // Iterate over each output tensor of the current node + for (size_t i = 0; i < output_num; ++i) { + // Check if the node is null + MS_EXCEPTION_IF_NULL(node); + + // Get the runtime information for the current node + auto runtime_info = node->user_data(); + MS_EXCEPTION_IF_NULL(runtime_info); + + // Get the output format, output type, and tensor size for the current output tensor + auto const &output_format = runtime_info->output_format(i); + auto output_type = runtime_info->output_type(i); + auto tensor_size = runtime_info->output_tensor_size(i); + + // Create a device address without a pointer + // The real device pointer will be obtained after the KernelBuild finishes + auto device_address = CreateDeviceAddress(nullptr, tensor_size, output_format, output_type); + + // Set the host shape of the device address using the runtime padding shape of the node and the output index + device_address->set_host_shape(trans::GetRuntimePaddingShape(node, i)); + + // Set the output address of the node to the device address + AnfAlgo::SetOutputAddr(device_address, i, node.get()); + } +} + + // Check if the number of input tensors matches the number of graph input parameters if (input_tensors.size() != graph.inputs().size()) { + // If they don't match, throw an exception with an error message MS_LOG(EXCEPTION) << "Input tensors size " << input_tensors.size() << " should be equal to graph input parameter size " << graph.inputs().size(); } + + // Iterate over each input parameter of the graph for (size_t input_index = 0; input_index < graph.inputs().size(); ++input_index) { + // Get the current input parameter auto item = graph.inputs()[input_index]; MS_EXCEPTION_IF_NULL(item); + + // Check if the current item is a Parameter node if (!item->isa()) { + // If it's not a Parameter node, skip to the next iteration continue; } + + // Get the number of output tensors for the current input parameter auto output_size = common::AnfAlgo::GetOutputTensorNum(item); + + // Iterate over each output tensor of the current input parameter for (size_t index = 0; index < output_size; index++) { + // Get the current input tensor auto current_tensor = input_tensors[input_index]; MS_EXCEPTION_IF_NULL(current_tensor); + + // Get the device address of the current input tensor auto output_address = std::dynamic_pointer_cast(current_tensor->device_address()); + + // Check if the output address is not null and the device type matches the target device address type if (output_address != nullptr && output_address->DeviceType() == GetTargetDeviceAddressType()) { + // Set the output address for the current output tensor of the current input parameter AnfAlgo::SetOutputAddr(output_address, index, item.get()); + // Continue to the next iteration continue; } + // Get the runtime information of the current operation from the user data of the item auto op_runtime_info = item->user_data(); MS_EXCEPTION_IF_NULL(op_runtime_info); + + // Get the output type ID of the current operation at the given index TypeId output_type_id = op_runtime_info->output_type(index); + + // Get the size of the output tensor at the given index auto output_tensor_size = op_runtime_info->output_tensor_size(index); + + // Get the output format of the current operation at the given index auto output_format = op_runtime_info->output_format(index); - auto device_address = - CreateDeviceAddress(nullptr, output_tensor_size, output_format, output_type_id, {item, index}); + + // Create a device address for the output tensor using the obtained information + auto device_address = CreateDeviceAddress(nullptr, output_tensor_size, output_format, output_type_id, {item, index}); + + // Set the device address to be from persistent memory if the current tensor is a parameter device_address->set_from_persistent_mem(current_tensor->is_parameter()); + + // Set the device address as the output address for the current operation at the given index AnfAlgo::SetOutputAddr(device_address, index, item.get()); + + // Set the device address for the current tensor current_tensor->set_device_address(device_address); + + // Set the synchronization status of the current tensor to indicate that it needs to be synchronized from host to device current_tensor->set_sync_status(kNeedSyncHostToDevice); } } } +// Reset the device address of each node in the kernel graph void KernelRuntime::ResetNodeAddress(const session::KernelGraph &kernel_graph) { + // Get the execution order of the kernels in the kernel graph auto kernels = kernel_graph.execution_order(); + + // Iterate over each kernel in the execution order for (auto &kernel : kernels) { + // Get the kernel module associated with the kernel auto kernel_mod = AnfAlgo::GetKernelMod(kernel); MS_EXCEPTION_IF_NULL(kernel_mod); + + // Get the number of input tensors for the kernel size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel); + + // Iterate over each input tensor of the kernel for (size_t j = 0; j < input_num; ++j) { + // Get the real input index of the tensor auto input_index = AnfAlgo::GetRealInputIndex(kernel, j); + + // Get the previous node's output tensor connected to the current input tensor KernelWithIndex kernel_with_index = common::AnfAlgo::GetPrevNodeOutput(kernel, input_index, true); + + // Get the index of the output tensor in the previous node auto index = kernel_with_index.second; + + // Get the previous node itself auto &input_node = kernel_with_index.first; + + // Check if the device address of the output tensor already exists if (NodeOutputDeviceAddressExist(input_node, index)) { continue; } + + // Get the data type of the output tensor TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(input_node, index); + + // Check if the output tensor has unknown data type if (output_type_id == kTypeUnknown) { MS_LOG(WARNING) << "It is not suggested to use a lonely weight parameter as the output of graph"; continue; } - auto tensor_size = AnfAlgo::GetOutputTensorMemSize(input_node, index); - auto device_address = CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(input_node, index), - output_type_id, {input_node, index}); - AnfAlgo::SetOutputAddr(device_address, index, input_node.get()); - } - auto output_sizes = kernel_mod->GetOutputSizeList(); - for (size_t i = 0; i < output_sizes.size(); ++i) { - auto output_format = AnfAlgo::GetOutputFormat(kernel, i); - auto output_type = AnfAlgo::GetOutputDeviceDataType(kernel, i); - AnfAlgo::SetOutputAddr(CreateDeviceAddress(nullptr, output_sizes[i], output_format, output_type), i, - kernel.get()); - } - auto workspace_sizes = kernel_mod->GetWorkspaceSizeList(); - for (size_t i = 0; i < workspace_sizes.size(); ++i) { - AnfAlgo::SetWorkspaceAddr(CreateDeviceAddress(nullptr, workspace_sizes[i], kOpFormat_DEFAULT, kNumberTypeFloat32), - i, kernel.get()); + // Get the memory size of the output tensor + auto tensor_size = AnfAlgo::GetOutputTensorMemSize(input_node, index); + + // Rest of the code is not provided, but it is assumed to perform some operations on the tensor size + // and the device address of the output tensor } } } + // Create a device address using the CreateDeviceAddress function, passing nullptr as the device pointer, tensor_size as the size of the tensor, + // AnfAlgo::GetOutputFormat(input_node, index) as the output format, output_type_id as the output type ID, and {input_node, index} as the input node and index. + auto device_address = CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(input_node, index), output_type_id, {input_node, index}); + + // Set the output address of the input node using the SetOutputAddr function, passing device_address as the device address, index as the output index, and input_node.get() as the input node. + AnfAlgo::SetOutputAddr(device_address, index, input_node.get()); + } + + // Get the list of output sizes from the kernel module + auto output_sizes = kernel_mod->GetOutputSizeList(); + + // Iterate over the output sizes + for (size_t i = 0; i < output_sizes.size(); ++i) { + // Get the output format and device data type for the current output + auto output_format = AnfAlgo::GetOutputFormat(kernel, i); + auto output_type = AnfAlgo::GetOutputDeviceDataType(kernel, i); + + // Set the output address for the current output using the CreateDeviceAddress function + AnfAlgo::SetOutputAddr(CreateDeviceAddress(nullptr, output_sizes[i], output_format, output_type), i, kernel.get()); + } + + // Get the list of workspace sizes from the kernel module + auto workspace_sizes = kernel_mod->GetWorkspaceSizeList(); + + // Iterate over the workspace sizes + for (size_t i = 0; i < workspace_sizes.size(); ++i) { + // Set the workspace address for the current workspace using the CreateDeviceAddress function + AnfAlgo::SetWorkspaceAddr(CreateDeviceAddress(nullptr, workspace_sizes[i], kOpFormat_DEFAULT, kNumberTypeFloat32), i, kernel.get()); + } + } +} + +// Definition of the function `RunOpAssignMemory` belonging to the class `KernelRuntime` void KernelRuntime::RunOpAssignMemory(const std::vector &input_tensors, const session::KernelGraph &graph, bool is_gradient_out, const std::map &tensor_to_node) { - MS_EXCEPTION_IF_NULL(mem_manager_); - mem_manager_->ResetDynamicMemory(); + // Check if the memory manager is null, and throw an exception if it is + MS_EXCEPTION_IF_NULL(mem_manager_); - for (const auto &node : graph.execution_order()) { - AssignCommunicationOutputFromMemoryPool(node); - AssignCommunicationInputFromMemoryPool(node); - } - - RunOpAssignInputMemory(input_tensors, graph); - AssignStaticMemoryValueNode(graph); - for (const auto &node : graph.execution_order()) { - RunOpAssignOutputMemory(node, tensor_to_node, is_gradient_out); - RunOpAssignWorkSpaceMemory(node); - } - UpdateRefNodeOutputMem(graph); + // Reset the dynamic memory managed by the memory manager + mem_manager_->ResetDynamicMemory(); } +// Iterate over each node in the execution order of the graph +for (const auto &node : graph.execution_order()) { + + // Assign the communication output from the memory pool for the current node + AssignCommunicationOutputFromMemoryPool(node); + + // Assign the communication input from the memory pool for the current node + AssignCommunicationInputFromMemoryPool(node); +} + +// Call the function to assign input memory to the input tensors of the graph +RunOpAssignInputMemory(input_tensors, graph); + +// Assign static memory value to the nodes in the graph +AssignStaticMemoryValueNode(graph); + +// Iterate over the nodes in the graph in the order of execution +for (const auto &node : graph.execution_order()) { + + // Call the function to assign output memory to the node's outputs + RunOpAssignOutputMemory(node, tensor_to_node, is_gradient_out); + + // Call the function to assign workspace memory to the node + RunOpAssignWorkSpaceMemory(node); +} + +// Update the reference node's output memory in the graph +UpdateRefNodeOutputMem(graph); + void KernelRuntime::RunOpClearMemory(const session::KernelGraph &graph) const { // clear input parameter memory resource for (const auto &input_node : graph.inputs()) { MS_EXCEPTION_IF_NULL(input_node); AnfAlgo::SetOutputAddr(nullptr, 0, input_node.get()); } + // clear input value node memory resource for (const auto &value_node : graph.graph_value_nodes()) { MS_EXCEPTION_IF_NULL(value_node); AnfAlgo::SetOutputAddr(nullptr, 0, value_node.get()); } + + // clear output and workspace memory resource for each node in the execution order for (const auto &cnode : graph.execution_order()) { MS_EXCEPTION_IF_NULL(cnode); + // clear output memory resource size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode); for (size_t index = 0; index < output_num; ++index) { AnfAlgo::SetOutputAddr(nullptr, index, cnode.get()); } + // clear workspace memory resource auto kernel_mod = AnfAlgo::GetKernelMod(cnode); - MS_EXCEPTION_IF_NULL(kernel_mod); - auto workspace_lists = kernel_mod->GetWorkspaceSizeList(); - for (size_t index = 0; index < workspace_lists.size(); ++index) { - AnfAlgo::SetWorkspaceAddr(nullptr, index, cnode.get()); - } + // ... } } +// Check if the pointer `kernel_mod` is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(kernel_mod); + +// Get the workspace size list from the `kernel_mod` object and assign it to the `workspace_lists` variable +auto workspace_lists = kernel_mod->GetWorkspaceSizeList(); + +// Iterate over the elements of the `workspace_lists` vector using an index variable `index` +for (size_t index = 0; index < workspace_lists.size(); ++index) { + + // Set the workspace address to nullptr for the given `index` and `cnode` object using the `AnfAlgo::SetWorkspaceAddr` function + AnfAlgo::SetWorkspaceAddr(nullptr, index, cnode.get()); +} + +// Close the innermost block of code +} #ifdef ENABLE_DEBUGGER + +// Function to check if data dumping is enabled bool KernelRuntime::DumpDataEnabled() { - // Returns true if e2e dump is enabled. + + // Get an instance of the DumpJsonParser class auto &dump_json_parser = DumpJsonParser::GetInstance(); + + // Return true if end-to-end (e2e) dump is enabled return dump_json_parser.e2e_dump_enabled(); } -bool KernelRuntime::DumpDataEnabledIteration() { - // Returns true if e2e dump is enabled and current iteration must be dumped. - auto &dump_json_parser = DumpJsonParser::GetInstance(); - if (!dump_json_parser.e2e_dump_enabled()) { - return false; - } +#endif // ENABLE_DEBUGGER - auto cur_iter = dump_json_parser.cur_dump_iter(); - if (dump_json_parser.IsDumpIter(cur_iter)) { - return true; +// Function to check if data dumping is enabled for the current iteration +bool KernelRuntime::DumpDataEnabledIteration() { + + // Get an instance of the DumpJsonParser class + auto &dump_json_parser = DumpJsonParser::GetInstance(); + + // Check if end-to-end (e2e) data dumping is enabled + if (!dump_json_parser.e2e_dump_enabled()) { + return false; // If e2e data dumping is not enabled, return false } - return false; + // If e2e data dumping is enabled, continue with the rest of the code + // ... } + +// Assign the value of the current dump iterator to the variable cur_iter using the auto keyword +auto cur_iter = dump_json_parser.cur_dump_iter(); + +// Check if the current iterator is a valid dump iterator using the IsDumpIter function of dump_json_parser +if (dump_json_parser.IsDumpIter(cur_iter)) { + // If it is a valid dump iterator, return true + return true; +} + +// If it is not a valid dump iterator, return false +return false; + +// End of the function #endif +// This function is a member function of the `KernelRuntime` class void KernelRuntime::AssignStaticMemory(const session::KernelGraph &graph) { + + // Call the `AssignStaticMemoryInput` function to assign static memory for input nodes in the given `KernelGraph` AssignStaticMemoryInput(graph); + + // Call the `AssignStaticMemoryValueNode` function to assign static memory for value nodes in the given `KernelGraph` AssignStaticMemoryValueNode(graph); + + // Call the `AssignStaticMemoryOutput` function to assign static memory for output nodes in the given `KernelGraph` AssignStaticMemoryOutput(graph); } +// Define the function `RunOpAssignInputMemory` which takes in a vector of `TensorPtr` objects as `input_tensors` and a `KernelGraph` object as `graph` void KernelRuntime::RunOpAssignInputMemory(const std::vector &input_tensors, const session::KernelGraph &graph) { + // Check if the memory manager is null, if so, throw an exception MS_EXCEPTION_IF_NULL(mem_manager_); + + // Check if the size of `input_tensors` is not equal to the size of `graph.inputs()`, if so, throw an exception if (input_tensors.size() != graph.inputs().size()) { MS_LOG(EXCEPTION) << "Input tensors size " << input_tensors.size() << " should be equal to graph input parameter size " << graph.inputs().size(); } +} for (size_t input_index = 0; input_index < graph.inputs().size(); ++input_index) { auto item = graph.inputs()[input_index]; MS_EXCEPTION_IF_NULL(item); + + // Check if the item is a Parameter if (!item->isa()) { continue; } + + // Get the number of output tensors for the current item auto output_size = common::AnfAlgo::GetOutputTensorNum(item); + + // Iterate over each output tensor for (size_t index = 0; index < output_size; index++) { auto current_tensor = input_tensors[input_index]; MS_EXCEPTION_IF_NULL(current_tensor); + + // Get the device address of the current tensor auto output_address = std::dynamic_pointer_cast(current_tensor->device_address()); - // Device address have already create + + // Check if the device address has already been created and if it matches the target device address type if (output_address != nullptr && output_address->DeviceType() == GetTargetDeviceAddressType()) { + + // Check if the device address pointer is null if (output_address->ptr_ == nullptr) { + + // Allocate memory from the memory pool for the output address if (!mem_manager_->MallocMemFromMemPool(output_address, output_address->size())) { MS_LOG(EXCEPTION) << "Allocate memory failed, size:" << output_address->size(); } } + } + // Set the output address using the provided output_address, index, and item AnfAlgo::SetOutputAddr(output_address, index, item.get()); continue; } + + // Get the output device data type for the item at the given index TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(item, index); + + // If the output type is unknown, get the output infer data type for the item at the given index if (output_type_id == kTypeUnknown) { output_type_id = common::AnfAlgo::GetOutputInferDataType(item, index); } + + // Get the size of the output tensor memory for the item at the given index auto tensor_size = AnfAlgo::GetOutputTensorMemSize(item, index); - // Device address new create - auto device_address = - CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(item, index), output_type_id, {item, index}); + + // Create a new device address with the given tensor size, output format, output type, and item-index pair + auto device_address = CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(item, index), output_type_id, {item, index}); + + // Check if the device address is null MS_EXCEPTION_IF_NULL(device_address); + + // Check if the memory manager is null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Set the device address to be from persistent memory device_address->set_from_persistent_mem(true); + + // Allocate memory from the memory pool using the memory manager and the device address auto ret = mem_manager_->MallocMemFromMemPool(device_address, tensor_size); + + // If memory allocation failed, throw an exception if (!ret) { MS_LOG(EXCEPTION) << "Device memory isn't enough and alloc failed, alloc size:" << tensor_size; } + + // Set the output address using the device address, index, and item AnfAlgo::SetOutputAddr(device_address, index, item.get()); } +// Closing brace for the main function } +// Closing brace for the namespace } +// Define the function `RunOpAssignOutputMemory` which takes in a kernel node, a map of tensors to nodes, and a boolean flag indicating if it is a gradient output void KernelRuntime::RunOpAssignOutputMemory(const AnfNodePtr &kernel, const std::map &tensor_to_node, bool is_gradient_out) { + // Check if the kernel node is null MS_EXCEPTION_IF_NULL(kernel); + + // Check if the memory manager is null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Get the kernel module associated with the kernel node auto kernel_mod = AnfAlgo::GetKernelMod(kernel); + + // Check if the kernel module is null MS_EXCEPTION_IF_NULL(kernel_mod); + + // Get the list of output sizes from the kernel module auto output_sizes = kernel_mod->GetOutputSizeList(); + + // If the output sizes list is empty, return if (output_sizes.empty()) { return; } - // Use device_address Allocated in RunOpMallocPre. - for (auto &iter : tensor_to_node) { - auto device_address = iter.first->device_address(); - AnfAlgo::SetOutputAddr(std::dynamic_pointer_cast(device_address), iter.second.second, - iter.second.first.get()); - } +// Iterate over each element in the `tensor_to_node` map using a range-based for loop +for (auto &iter : tensor_to_node) { - for (size_t i = 0; i < output_sizes.size(); ++i) { + // Get the device address of the current element + auto device_address = iter.first->device_address(); + + // Set the output address of the current element using the `SetOutputAddr` function from the `AnfAlgo` class + // The output address is obtained by casting the `device_address` to a `DeviceAddress` pointer using `std::dynamic_pointer_cast` + // The output address is set for the second element of the `iter` pair (iter.second.second) and the first element of the `iter` pair (iter.second.first.get()) + AnfAlgo::SetOutputAddr(std::dynamic_pointer_cast(device_address), iter.second.second, iter.second.first.get()); +} + +// Iterate over the output sizes vector using a for loop +for (size_t i = 0; i < output_sizes.size(); ++i) { + + // Check if the output address exists for the current kernel and output index if (AnfAlgo::OutputAddrExist(kernel, i, false)) { - auto address = AnfAlgo::GetMutableOutputAddr(kernel, i, false); - MS_EXCEPTION_IF_NULL(address); - if (address->ptr() == nullptr) { - MS_EXCEPTION_IF_NULL(mem_manager_); - if (!mem_manager_->MallocMemFromMemPool(address, address->size())) { - MS_LOG(EXCEPTION) << "Allocate memory failed, size:" << address->size(); + + // Get the mutable output address for the current kernel and output index + auto address = AnfAlgo::GetMutableOutputAddr(kernel, i, false); + MS_EXCEPTION_IF_NULL(address); + + // Check if the address pointer is null + if (address->ptr() == nullptr) { + + // If the address pointer is null, allocate memory from the memory pool + MS_EXCEPTION_IF_NULL(mem_manager_); + if (!mem_manager_->MallocMemFromMemPool(address, address->size())) { + MS_LOG(EXCEPTION) << "Allocate memory failed, size:" << address->size(); + } } - } - continue; + + // Continue to the next iteration of the loop + continue; } + + // Check if the current kernel is of type "ApplyMomentumOp" if (common::AnfAlgo::GetCNodeName(kernel) == kApplyMomentumOpName) { - auto device_address = AnfAlgo::GetPrevNodeMutableOutputAddr(kernel, i); - AnfAlgo::SetOutputAddr(device_address, i, kernel.get()); - continue; + + // Get the mutable output address of the previous node for the current kernel and output index + auto device_address = AnfAlgo::GetPrevNodeMutableOutputAddr(kernel, i); + + // Set the output address of the current kernel to the device address + AnfAlgo::SetOutputAddr(device_address, i, kernel.get()); + + // Continue to the next iteration of the loop + continue; } + + // Get the output format and output device data type for the current kernel and output index std::string output_format = AnfAlgo::GetOutputFormat(kernel, i); auto output_type = AnfAlgo::GetOutputDeviceDataType(kernel, i); + + // Create a device address with null pointer, output size, output format, output type, and kernel and output index pair auto device_address = CreateDeviceAddress(nullptr, output_sizes[i], output_format, output_type, {kernel, i}); + + // ... (code continues) + // Check if the device address is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(device_address); + + // Set the host shape of the device address using the GetRuntimePaddingShape function from the trans namespace device_address->set_host_shape(trans::GetRuntimePaddingShape(kernel, i)); + + // Check if the output is a gradient, and if so, set the "from_persistent_mem" flag of the device address to true if (is_gradient_out) { device_address->set_from_persistent_mem(true); } + + // Allocate memory from the memory pool using the MallocMemFromMemPool function of the mem_manager_ auto ret = mem_manager_->MallocMemFromMemPool(device_address, output_sizes[i]); + + // Check if the memory allocation was successful, and if not, throw an exception with the allocation size if (!ret) { MS_LOG(EXCEPTION) << "Device memory isn't enough and alloc failed, alloc size:" << output_sizes[i]; } + + // Set the output address of the kernel using the SetOutputAddr function of the AnfAlgo namespace AnfAlgo::SetOutputAddr(device_address, i, kernel.get()); } } +// This function is a member function of the `KernelRuntime` class and is responsible for assigning workspace memory to a given kernel. + void KernelRuntime::RunOpAssignWorkSpaceMemory(const AnfNodePtr &kernel) { + // Check if the kernel is null MS_EXCEPTION_IF_NULL(kernel); + + // Check if the memory manager is null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Check if the kernel is a CNode (a node in the computation graph) if (kernel->isa()) { + // Get the kernel module associated with the kernel auto kernel_mod = AnfAlgo::GetKernelMod(kernel); MS_EXCEPTION_IF_NULL(kernel_mod); + + // Get the list of workspace sizes required by the kernel auto workspace_lists = kernel_mod->GetWorkspaceSizeList(); + + // Iterate over the workspace sizes for (size_t i = 0; i < workspace_lists.size(); ++i) { + // Create a device address for the workspace memory auto device_address = CreateDeviceAddress(nullptr, workspace_lists[i], "", kTypeUnknown); MS_EXCEPTION_IF_NULL(device_address); + + // Allocate memory from the memory pool for the workspace auto ret = mem_manager_->MallocMemFromMemPool(device_address, workspace_lists[i]); + + // Check if the allocation was successful if (!ret) { + // If allocation failed, log an exception with the size of the failed allocation MS_LOG(EXCEPTION) << "Device memory isn't enough and alloc failed, alloc size:" << workspace_lists[i]; } + + // Set the workspace address in the kernel node AnfAlgo::SetWorkspaceAddr(device_address, i, kernel.get()); } } } +// This function is responsible for assigning memory to the output nodes of a kernel in a runtime environment. +// It takes in the pre_output_value, which is the value of the output node before memory assignment, and the graph containing the output nodes. void KernelRuntime::RunOpAssignOutputNodeMemory(const ValuePtr &pre_output_value, const session::KernelGraph &graph) { + // If the pre_output_value is nullptr, there is nothing to assign memory to, so return. if (pre_output_value == nullptr) { return; } + + // Create a vector to store the pre_output_tensors std::vector pre_output_tensors; + // Convert the pre_output_value to tensors and store them in the pre_output_tensors vector TensorValueToTensor(pre_output_value, &pre_output_tensors); + + // Get the output nodes of the graph auto output_nodes = graph.outputs(); + + // Check if the number of pre_output_tensors is equal to the number of output nodes in the graph if (pre_output_tensors.size() != output_nodes.size()) { + // If they are not equal, throw an exception with an error message indicating the mismatch MS_LOG(EXCEPTION) << "The size of pre output tensors [" << pre_output_tensors.size() << "] is not equal to the size of output nodes of graph [" << output_nodes.size() << "]"; } - // share output address with pre output tensors + + // Iterate over the output nodes and assign memory to them for (size_t i = 0; i < output_nodes.size(); ++i) { + // Get the output node and its index auto output_node_with_index = common::AnfAlgo::VisitKernel(output_nodes[i], 0); auto output_node = output_node_with_index.first; MS_EXCEPTION_IF_NULL(output_node); + + // Check if the output node is a CNode if (!output_node->isa()) { + // If it is not a CNode, check if it is a Parameter if (output_node->isa()) { auto param = output_node->cast(); if (param != nullptr && !param->has_default()) { + // If it is a Parameter and does not have a default value, assign memory to it + // ... + } + } + } + } +} + // Check if the output parameter is a real parameter, if not, throw an exception MS_LOG(EXCEPTION) << "The output parameter should be real parameter!"; } } continue; } + // Cast the output node to a CNode pointer auto real_output_cnode = output_node->cast(); MS_EXCEPTION_IF_NULL(real_output_cnode); MS_EXCEPTION_IF_NULL(pre_output_tensors[i]); + // Check if the device address of the pre output tensor is null, if so, log a message and continue to the next iteration if (pre_output_tensors[i]->device_address() == nullptr) { MS_LOG(INFO) << "The address of pre output tensor [" << i << "] is a nullptr!"; continue; } + // Check if the real output node is a NOP node if (common::AnfAlgo::IsNopNode(real_output_cnode)) { + // Check if the input size of the real output node is less than the minimum input size required if (real_output_cnode->inputs().size() < kMinInputSize) { MS_LOG(EXCEPTION) << "The input size of output node: " << real_output_cnode->DebugString() - << " should large than one!"; + << " should be larger than one!"; } + // Set the output address of the pre output tensor to the second input of the real output node AnfAlgo::SetOutputAddr(std::dynamic_pointer_cast(pre_output_tensors[i]->device_address()), output_node_with_index.second, real_output_cnode->input(1).get()); } else { + // Call the SetOutputAddr function of the AnfAlgo namespace + // Pass the device address of the i-th pre_output_tensor as a shared pointer to DeviceAddress + // Pass the output_node_with_index.second as the second argument + // Pass the output_node_with_index.first.get() as the third argument AnfAlgo::SetOutputAddr(std::dynamic_pointer_cast(pre_output_tensors[i]->device_address()), output_node_with_index.second, output_node_with_index.first.get()); } } } -void KernelRuntime::AssignStaticMemoryInput(const session::KernelGraph &graph) { - MS_EXCEPTION_IF_NULL(mem_manager_); - auto graph_id = graph.graph_id(); - MS_LOG(INFO) << "AssignStaticMemoryInput start for graph " << graph_id; - auto graph_inputs = GetGraphInputs(graph); - auto graph_valid_input = graph.valid_inputs(); - graph_inputs.insert(graph_inputs.end(), graph.child_graph_result().begin(), graph.child_graph_result().end()); - std::vector need_alloc_nodes; - auto add_need_alloc_nodes = [&need_alloc_nodes, graph_id, this](const AnfNodePtr &node) { - MS_EXCEPTION_IF_NULL(node); - if (!node->isa()) { +// The function `AssignStaticMemoryInput` is a member function of the class `KernelRuntime`. +// It is responsible for assigning static memory input to the given `graph`. + +// Check if the memory manager is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(mem_manager_); + +// Get the graph ID of the given graph +auto graph_id = graph.graph_id(); + +// Log an informational message indicating the start of the static memory input assignment for the graph +MS_LOG(INFO) << "AssignStaticMemoryInput start for graph " << graph_id; + +// Get the inputs of the graph +auto graph_inputs = GetGraphInputs(graph); + +// Get the valid inputs of the graph +auto graph_valid_input = graph.valid_inputs(); + +// Append the child graph results to the graph inputs +graph_inputs.insert(graph_inputs.end(), graph.child_graph_result().begin(), graph.child_graph_result().end()); + +// Create an empty vector to store the nodes that need memory allocation +std::vector need_alloc_nodes; + +// Define a lambda function `add_need_alloc_nodes` that adds a node to the `need_alloc_nodes` vector if it meets certain conditions +auto add_need_alloc_nodes = [&need_alloc_nodes, graph_id, this](const AnfNodePtr &node) { + MS_EXCEPTION_IF_NULL(node); + + // Check if the node is not a parameter, and return if it is not + if (!node->isa()) { + return; + } + + // Check if the output device address of the node exists and is not null, and return if it does + if (NodeOutputDeviceAddressExist(node, 0)) { + const auto &address = AnfAlgo::GetOutputAddr(node, 0); + MS_EXCEPTION_IF_NULL(address); + if (address->GetPtr() != nullptr) { return; } - if (NodeOutputDeviceAddressExist(node, 0)) { - const auto &address = AnfAlgo::GetOutputAddr(node, 0); - MS_EXCEPTION_IF_NULL(address); - if (address->GetPtr() != nullptr) { - return; - } - } + } + // Continue with the rest of the function +} + // Declare a variable "input_param" and assign it the value of the "node" pointer casted to a ParameterPtr auto input_param = node->cast(); + + // Check if "input_param" is not a nullptr and if it is not used by a real kernel in the graph with the given "graph_id" if (input_param != nullptr && !input_param->IsUsedByRealKernelInGraph(graph_id)) { + // If the above condition is true, return from the current function return; } + + // If the above condition is false, add the "node" pointer to the "need_alloc_nodes" vector need_alloc_nodes.push_back(node); }; +// Iterate over the graph_inputs vector using a for loop for (size_t i = 0; i < graph_inputs.size(); ++i) { auto input_node = graph_inputs[i]; MS_EXCEPTION_IF_NULL(input_node); + + // Check if the current input node is valid and should be processed if (i < graph_valid_input.size() && !graph_valid_input[i]) { continue; } + + // Check if the input node is of type prim::kPrimMakeTuple if (common::AnfAlgo::CheckPrimitiveType(input_node, prim::kPrimMakeTuple)) { auto outs = common::AnfAlgo::GetAllOutput(input_node); + + // Iterate over all the outputs of the input node for (auto &out : outs) { MS_EXCEPTION_IF_NULL(out); + + // Call the add_need_alloc_nodes function with the current output node add_need_alloc_nodes(out); } } + + // Call the add_need_alloc_nodes function with the current input node add_need_alloc_nodes(input_node); } + + // Check if ENABLE_CPU is defined and _WIN32 is not defined #if ((defined ENABLE_CPU) && (!defined _WIN32)) bool ps_cache_check = false; #endif + + // Create an empty map to store the shadow_backend_node_map std::map shadow_backend_node_map; + + // Call the GetShadowBackendNodeMap function with the graph and the shadow_backend_node_map GetShadowBackendNodeMap(graph, &shadow_backend_node_map); for (auto &item : need_alloc_nodes) { + // Check if the item is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(item); + + // Get the number of output tensors for the current item auto output_size = common::AnfAlgo::GetOutputTensorNum(item); + + // Iterate over each output tensor of the current item for (size_t index = 0; index < output_size; index++) { + // Get the data type of the current output tensor TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(item, index); - // if graph output is a weight and doesn't link to any cnode, it's data type will be unknown + + // Check if the output type is unknown if (output_type_id == kTypeUnknown) { + // Log a warning message indicating that using a lonely weight parameter as the output of the graph is not suggested MS_LOG(WARNING) << "It is not suggested to use a lonely weight parameter as the output of graph"; + + // Continue to the next iteration of the loop continue; } + + // Get the device address of the current item DeviceAddressPtr device_address = GetInternalDeviceAddress(graph, item); -#if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__)) + + // Check if the current parameter is a hash table in parameter server training mode const std::string ¶m_name = item->fullname_with_scope(); if (ps::ps_cache_instance.IsHashTable(param_name)) { + // Log an info message indicating that the parameter enables the embeddingLookup cache in parameter server training mode MS_LOG(INFO) << "Parameter(" << param_name << ")" << " enables the embeddingLookup cache in parameter server training mode."; - // PS embeddingLookup cache check. + + // Check if the PS embeddingLookup cache has been checked if (!ps_cache_check) { + // Check if the current graph supports PS embeddingLookup cache CheckIfSupportPSEmbeddingCache(graph); + + // Set the PS embeddingLookup cache check flag to true ps_cache_check = true; } + } + // Query the hash table address for the given parameter name const auto &address = ps::ps_cache_instance.QueryHashTableAddr(param_name); + // Check if the address is null MS_EXCEPTION_IF_NULL(address.addr); + // Create a device address using the queried address, size, output format, output type id, and item index device_address = CreateDeviceAddress(address.addr, address.size, AnfAlgo::GetOutputFormat(item, index), output_type_id, {item, index}); + // Set the host shape of the device address using the runtime padding shape of the item and index device_address->set_host_shape(trans::GetRuntimePaddingShape(item, index)); + // Set the output address of the item at the given index AnfAlgo::SetOutputAddr(device_address, index, item.get()); + // Continue to the next iteration of the loop continue; } #endif + // Get the device address for the item at the given index using the shadow backend node map and graph ID GetDeviceAddress(item, shadow_backend_node_map, index, graph.graph_id(), &device_address); + // Set the output address of the item at the given index AnfAlgo::SetOutputAddr(device_address, index, item.get()); } } + // Log the end of assigning static memory input MS_LOG(INFO) << "AssignStaticMemoryInput end"; } +// This function is a member function of the `KernelRuntime` class. +// It is used to get the device address of an `AnfNode` based on the provided parameters. + +// Parameters: +// - `item`: The `AnfNode` for which the device address is to be obtained. +// - `shadow_backend_node_map`: A map that maps `AnfNode` pointers to their corresponding shadow nodes. +// - `index`: The index of the output of the `item` node for which the device address is to be obtained. +// - `graph_id`: The ID of the graph to which the `item` node belongs. +// - `device_address`: A pointer to a `DeviceAddressPtr` object that will store the obtained device address. + void KernelRuntime::GetDeviceAddress(const AnfNodePtr &item, const std::map shadow_backend_node_map, size_t index, uint32_t graph_id, DeviceAddressPtr *device_address) { + + // Initialize a pointer to the shadow node corresponding to the `item` node AnfNodePtr shadow_node = nullptr; + + // Find the shadow node in the `shadow_backend_node_map` using the `item` node as the key auto iter = shadow_backend_node_map.find(item); + + // If the shadow node is found, assign it to the `shadow_node` pointer if (iter != shadow_backend_node_map.end()) { shadow_node = iter->second; } + + // If the `device_address` pointer is null and the `shadow_node` pointer is not null if (*device_address == nullptr && shadow_node != nullptr) { + + // Get the mutable output address of the `shadow_node` at the specified `index` auto conj_device_address = AnfAlgo::GetMutableOutputAddr(shadow_node, index); + + // If the `conj_device_address` is not null and its device type is Ascend if (conj_device_address != nullptr && conj_device_address->DeviceType() == DeviceAddressType::kAscend) { + + // Assign the `conj_device_address` to the `device_address` pointer *device_address = conj_device_address; } - } else if (*device_address == nullptr) { - auto tensor_size = AnfAlgo::GetOutputTensorMemSize(item, index); - TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(item, index); - *device_address = - CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(item, index), output_type_id, {item, index}); } - if (*device_address != nullptr && (*device_address)->GetPtr() == nullptr) { + // If the `device_address` pointer is null + else if (*device_address == nullptr) { + + // Get the size of the output tensor of the `item` node at the specified `index` auto tensor_size = AnfAlgo::GetOutputTensorMemSize(item, index); - (*device_address)->set_host_shape(trans::GetRuntimePaddingShape(item, index)); - MS_LOG(INFO) << "Assign Static Memory for Input node, size:" << tensor_size - << " node:" << item->fullname_with_scope() << " index: " << index; - if (mem_manager_->MallocMem(kStaticMem, tensor_size, *device_address, graph_id) == nullptr) { - MS_LOG(EXCEPTION) << "Cannot alloc address when flag is: " << kStaticMem << ", tensor size is: " << tensor_size; + + // Get the output device data type of the `item` node at the specified `index` + TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(item, index); + + // Create a device address object using the obtained tensor size, output format, output type ID, and the `item` node and index as the source + *device_address = CreateDeviceAddress(nullptr, tensor_size, AnfAlgo::GetOutputFormat(item, index), output_type_id, {item, index}); + } + + // If the `device_address` pointer is not null and its pointer is null + if (*device_address != nullptr && (*device_address)->GetPtr() == nullptr) { + // ... + } +} + auto tensor_size = AnfAlgo::GetOutputTensorMemSize(item, index); // Get the size of the output tensor memory + (*device_address)->set_host_shape(trans::GetRuntimePaddingShape(item, index)); // Set the host shape of the device address + MS_LOG(INFO) << "Assign Static Memory for Input node, size:" << tensor_size // Log the size of the tensor memory + << " node:" << item->fullname_with_scope() << " index: " << index; // Log the node name and index + if (mem_manager_->MallocMem(kStaticMem, tensor_size, *device_address, graph_id) == nullptr) { // Allocate memory for the device address + MS_LOG(EXCEPTION) << "Cannot alloc address when flag is: " << kStaticMem << ", tensor size is: " << tensor_size; // Log an exception if memory allocation fails } } } +// This function is used to assign static memory output for a given graph. void KernelRuntime::AssignStaticMemoryOutput(const session::KernelGraph &graph) { + // Log the start of the function along with the graph ID MS_LOG(INFO) << "AssignStaticMemoryOutput start for graph " << graph.graph_id(); + + // Get all the output nodes of the graph that are of type prim::kPrimTupleGetItem auto nodes = common::AnfAlgo::GetAllOutput(graph.output(), {prim::kPrimTupleGetItem}); + + // Create a vector to store non-communication operations std::vector non_communication_op; - // Assign Communicate Op Memory firstly. + + // Assign memory for communication operations first for (const auto &node : nodes) { - // Assign output address to nop node that the attribute of "skip_nop_op_addr" is false; + // Check if the node is a NOP node or if the "skip_nop_op_addr" attribute is true auto is_skip = !common::AnfAlgo::IsNopNode(node) || common::AnfAlgo::IsNeedSkipNopOpAddr(node); + + // Get the kernel with index for the node auto kernel_with_index = common::AnfAlgo::VisitKernelWithReturnType(node, 0, is_skip); MS_EXCEPTION_IF_NULL(kernel_with_index.first); + + // Check if the kernel is a CNode and a real kernel if (!kernel_with_index.first->isa() || !AnfUtils::IsRealKernel(kernel_with_index.first)) { continue; } + + // Check if the kernel is a communication operation if (common::AnfAlgo::IsCommunicationOp(kernel_with_index.first)) { + // Assign memory for the communication node AssignCommunicationNodeMem(kStaticMem, kernel_with_index.first); } else { + // Add the kernel to the non-communication operations vector non_communication_op.emplace_back(kernel_with_index); } } +} + // Iterate over each item in the non_communication_op container using a range-based for loop for (const auto &item_with_index : non_communication_op) { + // Check if the first element of the item is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(item_with_index.first); + + // Log a debug message indicating the AssignNodeOutputMem operation for the current item MS_LOG(DEBUG) << "AssignNodeOutputMem for " << item_with_index.first->fullname_with_scope(); + + // Call the AssignNodeOutputMem function with the appropriate arguments AssignNodeOutputMem(kStaticMem, item_with_index.first, SizeToInt(item_with_index.second)); } + + // Log an info message indicating the end of the AssignStaticMemoryOutput function MS_LOG(INFO) << "AssignStaticMemoryOutput end"; } +// This function updates the output memory addresses of reference nodes in a given kernel graph. void KernelRuntime::UpdateRefNodeOutputMem(const session::KernelGraph &graph) { + // Get the execution order of kernels in the graph auto &kernels = graph.execution_order(); + + // Iterate over each kernel in the execution order for (auto &kernel : kernels) { + // Check if the kernel is null MS_EXCEPTION_IF_NULL(kernel); + + // Get the number of output tensors for the current kernel auto output_num = common::AnfAlgo::GetOutputTensorNum(kernel); + + // If the kernel has no output tensors, skip to the next kernel if (output_num == 0) { MS_LOG(DEBUG) << "This kernel has no output size."; continue; } + + // Iterate over each output tensor of the current kernel for (size_t i = 0; i < output_num; ++i) { + // Create a pair of the current kernel and the output index session::AnfWithOutIndex out_pair(kernel, i); + + // Check if the output tensor is in the reference output map of the graph if (graph.IsInRefOutputMap(out_pair)) { + // Get the corresponding original pair from the reference output map auto origin_pair = graph.GetRefCorrespondOutput(out_pair); + + // Check if the original pair is null MS_EXCEPTION_IF_NULL(origin_pair.first); + + // Get the mutable output address of the original node auto origin_node_output_addr = AnfAlgo::GetMutableOutputAddr(origin_pair.first, origin_pair.second); + + // Check if the origin node output address is null MS_EXCEPTION_IF_NULL(origin_node_output_addr); + + // Get the mutable output address of the current node auto cur_node_output_addr = AnfAlgo::GetMutableOutputAddr(kernel, i); + + // Check if the current node output address is not the same as the origin node output address if (origin_node_output_addr.get() != cur_node_output_addr.get()) { MS_LOG(DEBUG) << "REF address is not same, ref node output need address update"; - MS_LOG(DEBUG) << "REF origin op is " << origin_pair.first->DebugString() << ", output index is " + MS_LOG(DEBUG) << "REF origin op is " << origin_pair.first->DebugString() << ", output index is " << i; + // TODO: Update the output address of the reference node + } + } + } + } +} << origin_pair.second << ", cur op is " << kernel->DebugString() << ", out index is " << i; + // Check if the current node's output address has a non-empty host shape if (!cur_node_output_addr->host_shape().empty()) { + // Set the host shape of the origin node's output address to the host shape of the current node's output address origin_node_output_addr->set_host_shape(cur_node_output_addr->host_shape()); } + // Set the output address of the origin node at index i to the kernel's address AnfAlgo::SetOutputAddr(origin_node_output_addr, i, kernel.get()); } } @@ -727,126 +1448,252 @@ void KernelRuntime::UpdateRefNodeOutputMem(const session::KernelGraph &graph) { } } +// A function to assign memory for a communication node based on its type and whether there is any reusable memory available + void KernelRuntime::AssignCommunicationNodeMem(MemType type, const AnfNodePtr &node) { + + // Check if there is any reusable memory available if (!reuse_communication_address_.empty()) { + + // If there is reusable memory available, assign dynamic memory type type = kDynamicMem; } + + // Assign memory for input of the communication node AssignCommunicationNodeInputMem(type, node); + + // Assign memory for output of the communication node AssignCommunicationNodeOutputMem(type, node); + + // Assign memory for workspace of the communication node AssignWorkSpaceMem(type, node); } +// Function to assign memory for the output of a communication node void KernelRuntime::AssignCommunicationNodeOutputMem(MemType type, const AnfNodePtr &node) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(mem_manager_); + + // Get the kernel module associated with the node auto kernel_mod = AnfAlgo::GetKernelMod(node); MS_EXCEPTION_IF_NULL(kernel_mod); + + // Get the list of output sizes from the kernel module auto output_sizes = kernel_mod->GetOutputSizeList(); + + // If the output sizes list is empty, log a message and return if (output_sizes.empty()) { MS_LOG(INFO) << "This kernel[" << node->DebugString() << "] has no output size."; return; } + + // Get the current context auto context_ptr = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(context_ptr); + + // Initialize the total size and output index size_t total_size = 0; size_t output_index = 0; + + // Create a vector to store the aligned sizes std::vector align_size_list; + + // Iterate over the output sizes for (uint64_t mem_size : output_sizes) { + // Check if the output address already exists for the current output index if (AnfAlgo::OutputAddrExist(node, output_index++)) { MS_LOG(INFO) << "Communication op " << node->fullname_with_scope() << " has output device address"; return; } + // ... (rest of the code) + // Check if the "MS_CTX_ENABLE_HCCL" parameter is set to true in the context if (context_ptr->get_param(MS_CTX_ENABLE_HCCL)) { - mem_size = MemoryManager::GetCommonAlignSize(mem_size); + // If it is true, get the common aligned size of mem_size using the MemoryManager class + mem_size = MemoryManager::GetCommonAlignSize(mem_size); } - total_size += mem_size; - align_size_list.emplace_back(mem_size); - } + // Add mem_size to the total_size + total_size += mem_size; + + // Add mem_size to the align_size_list + align_size_list.emplace_back(mem_size); + + // Check if the align_size_list is empty if (align_size_list.empty()) { + // If it is empty, return from the function return; } + // Check if the type is kSomasReuseDynamicMem if (type == kSomasReuseDynamicMem) { + + // Call the KernelMemNotReuse function to check if memory reuse is not possible for the given node bool not_reuse = KernelMemNotReuse(node); + + // If memory reuse is not possible if (not_reuse) { + + // Set the type to kDynamicMem type = kDynamicMem; + + // Print an informational message indicating that memory reuse is disabled for the output of the node MS_LOG(INFO) << "Disable Memory Reuse for " << node->fullname_with_scope() << "'s output."; } } - uint8_t *output_ptr = nullptr; - int64_t valid_reuse_index = -1; - auto cnode = node->cast(); - MS_EXCEPTION_IF_NULL(cnode); - if (common::AnfAlgo::HasNodeAttr(kAttrReuseCommunication, cnode)) { - auto reuse_index = common::AnfAlgo::GetNodeAttr(cnode, kAttrReuseCommunication); - auto it = reuse_communication_address_.find(reuse_index); - if (it != reuse_communication_address_.end()) { - valid_reuse_index = reuse_index; - output_ptr = it->second.second; - } +// Declare a pointer to an 8-bit unsigned integer and initialize it to nullptr +uint8_t *output_ptr = nullptr; + +// Declare a variable of type int64_t and initialize it to -1 +int64_t valid_reuse_index = -1; + +// Cast the input node to a CNode pointer +auto cnode = node->cast(); + +// Throw an exception if the casted CNode pointer is null +MS_EXCEPTION_IF_NULL(cnode); + +// Check if the CNode has the attribute "kAttrReuseCommunication" +if (common::AnfAlgo::HasNodeAttr(kAttrReuseCommunication, cnode)) { + + // Get the value of the attribute "kAttrReuseCommunication" from the CNode + auto reuse_index = common::AnfAlgo::GetNodeAttr(cnode, kAttrReuseCommunication); + + // Find the reuse communication address in the map using the reuse index + auto it = reuse_communication_address_.find(reuse_index); + + // Check if the reuse communication address is found in the map + if (it != reuse_communication_address_.end()) { + + // Set the valid reuse index to the found reuse index + valid_reuse_index = reuse_index; + + // Set the output pointer to the second element of the found reuse communication address + output_ptr = it->second.second; } +} for (size_t j = 0; j < align_size_list.size(); ++j) { + // Get the output format for the j-th output of the node std::string output_format = AnfAlgo::GetOutputFormat(node, j); + + // Get the output device data type for the j-th output of the node auto output_type = AnfAlgo::GetOutputDeviceDataType(node, j); + + // Create a device address for the j-th output of the node auto address = CreateDeviceAddress(nullptr, output_sizes[j], output_format, output_type, {node, j}); MS_EXCEPTION_IF_NULL(address); + + // If output_ptr is nullptr, allocate memory for the output and set output_ptr to the allocated memory if (output_ptr == nullptr) { output_ptr = mem_manager_->MallocOutputMem(node, 0, type, total_size, address, true); MS_EXCEPTION_IF_NULL(output_ptr); + + // If valid_reuse_index is not -1, update the communication address for reuse if (valid_reuse_index != -1) { auto &it = reuse_communication_address_[valid_reuse_index]; it.second = output_ptr; } } else { + // Set the pointer of the device address to output_ptr address->set_ptr(output_ptr); } + + // Set the host shape of the device address address->set_host_shape(trans::GetRuntimePaddingShape(node, j)); + + // Set the output address of the node for the j-th output AnfAlgo::SetOutputAddr(address, j, node.get()); + + // Increment output_ptr by the size of the j-th output output_ptr += align_size_list[j]; } } +// A function named "KernelMemNotReuse" that takes an AnfNodePtr as input and returns a boolean value bool KernelRuntime::KernelMemNotReuse(const AnfNodePtr &node) { - MS_EXCEPTION_IF_NULL(node); - return false; + // Check if the input node is null, and throw an exception if it is + MS_EXCEPTION_IF_NULL(node); + + // Return false, indicating that kernel memory reuse is not allowed + return false; } +// Function to pre-assign memory for a CNode in the kernel runtime DeviceAddressPtr KernelRuntime::PreAssignCNodeMemory(const AnfNodePtr &anf_node, size_t index) const { MS_EXCEPTION_IF_NULL(anf_node); + + // Check if the given ANF node is a NOP node (no operation) if (common::AnfAlgo::IsNopNode(anf_node)) { + // Get the previous node and its output index connected to the NOP node auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(anf_node, index); + + // Recursively call PreAssignCNodeMemory for the previous node return PreAssignCNodeMemory(input_node_with_index.first, input_node_with_index.second); } + // If the given ANF node is not a NOP node, return nullptr + return nullptr; +} + // Get the size of the output tensor memory using the AnfAlgo::GetOutputTensorMemSize function auto output_size = AnfAlgo::GetOutputTensorMemSize(anf_node, index); + + // Get the output format of the tensor using the AnfAlgo::GetOutputFormat function std::string output_format = AnfAlgo::GetOutputFormat(anf_node, index); + + // Get the output device data type using the AnfAlgo::GetOutputDeviceDataType function auto output_type = AnfAlgo::GetOutputDeviceDataType(anf_node, index); + + // Create a device address using the CreateDeviceAddress function, passing nullptr for the device pointer, + // the output size, output format, output type, and the node and index as a vector auto address = CreateDeviceAddress(nullptr, output_size, output_format, output_type, {anf_node, index}); + + // Set the output address using the AnfAlgo::SetOutputAddr function, passing the address, index, and the node pointer AnfAlgo::SetOutputAddr(address, index, anf_node.get()); + + // Return the created address return address; } -void KernelRuntime::AssignCommunicationNodeInputMem(MemType type, const AnfNodePtr &node) { - auto context_ptr = MsContext::GetInstance(); - MS_EXCEPTION_IF_NULL(context_ptr); - MS_EXCEPTION_IF_NULL(node); - MS_EXCEPTION_IF_NULL(mem_manager_); - size_t total_size = 0; - std::vector> addr_size; - size_t input_num = common::AnfAlgo::GetInputTensorNum(node); - for (size_t i = 0; i < input_num; ++i) { - auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(node, i, true); - auto input_node = input_node_with_index.first; - MS_EXCEPTION_IF_NULL(input_node); - if (AnfAlgo::OutputAddrExist(input_node, input_node_with_index.second)) { - MS_LOG(INFO) << "Communication op " << input_node->fullname_with_scope() << " has input device address"; - return; - } - DeviceAddressPtr address = nullptr; +// This function is used to assign memory for the input of a communication node. +// It takes in the type of memory to be assigned and the communication node as parameters. - address = PreAssignCNodeMemory(input_node, input_node_with_index.second); +// Get the instance of the global context +auto context_ptr = MsContext::GetInstance(); +MS_EXCEPTION_IF_NULL(context_ptr); + +// Check if the node is null +MS_EXCEPTION_IF_NULL(node); + +// Check if the memory manager is null +MS_EXCEPTION_IF_NULL(mem_manager_); + +// Initialize the total size of the memory to be assigned to 0 +size_t total_size = 0; + +// Create a vector to store pairs of device addresses and their corresponding sizes +std::vector> addr_size; + +// Get the number of input tensors for the communication node +size_t input_num = common::AnfAlgo::GetInputTensorNum(node); + +// Iterate over each input tensor +for (size_t i = 0; i < input_num; ++i) { + // Get the previous node and output index for the current input tensor + auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(node, i, true); + auto input_node = input_node_with_index.first; + MS_EXCEPTION_IF_NULL(input_node); + + // Check if the output address exists for the input node and output index + if (AnfAlgo::OutputAddrExist(input_node, input_node_with_index.second)) { + // Log an information message indicating that the communication op already has an input device address + MS_LOG(INFO) << "Communication op " << input_node->fullname_with_scope() << " has input device address"; + return; + } + + // Initialize the device address to nullptr + DeviceAddressPtr address = nullptr; + +// Assign the result of calling the function PreAssignCNodeMemory with the arguments input_node and input_node_with_index.second to the variable address. MS_EXCEPTION_IF_NULL(address); auto mem_size = MemoryManager::GetCommonAlignSize(address->size()); @@ -860,151 +1707,292 @@ void KernelRuntime::AssignCommunicationNodeInputMem(MemType type, const AnfNodeP bool not_reuse = KernelMemNotReuse(node); if (not_reuse) { type = kDynamicMem; + // Print an informational message indicating that memory reuse is disabled for the given node MS_LOG(INFO) << "Disable Memory Reuse for " << node->fullname_with_scope() << "'s input."; } } auto cnode = node->cast(); MS_EXCEPTION_IF_NULL(cnode); if (cnode->inputs().size() < kMinInputSize) { - // communication node's input should contain itself and at least on input + // Print an error message indicating that the given CNode does not have enough inputs + // Communication node's input should contain itself and at least one input MS_LOG(ERROR) << "No inputs for " << cnode->fullname_with_scope(); - return; - } +// Return statement without a value, indicating an early termination of the program +return; - int64_t valid_reuse_index = -1; - uint8_t *input_ptr = nullptr; - if (common::AnfAlgo::HasNodeAttr(kAttrReuseCommunication, cnode)) { +// Declare a variable of type int64_t named valid_reuse_index and initialize it with -1 +int64_t valid_reuse_index = -1; + +// Declare a pointer of type uint8_t named input_ptr and initialize it with nullptr +uint8_t *input_ptr = nullptr; + +// Check if the given node has the attribute kAttrReuseCommunication using the common::AnfAlgo::HasNodeAttr function +if (common::AnfAlgo::HasNodeAttr(kAttrReuseCommunication, cnode)) { + + // If the attribute is present, get the value of the attribute as an int64_t using the common::AnfAlgo::GetNodeAttr function auto reuse_index = common::AnfAlgo::GetNodeAttr(cnode, kAttrReuseCommunication); - auto it = reuse_communication_address_.find(reuse_index); - if (it != reuse_communication_address_.end()) { - valid_reuse_index = reuse_index; - input_ptr = it->second.first; - } - } + // Search for the reuse_index in the reuse_communication_address_ map using the find function + auto it = reuse_communication_address_.find(reuse_index); + + // If the reuse_index is found in the map + if (it != reuse_communication_address_.end()) { + + // Update the valid_reuse_index with the found reuse_index + valid_reuse_index = reuse_index; + + // Update the input_ptr with the first element of the value associated with the found reuse_index in the map + input_ptr = it->second.first; + } +} + + // Check if the input pointer is nullptr if (input_ptr == nullptr) { + // Get the first input node of the current CNode auto first_input_node = cnode->input(1); + + // Visit the kernel with the first input node and get the prenode index auto prenode_index = common::AnfAlgo::VisitKernelWithReturnType(first_input_node, 0, true); + + // Allocate memory for the output of the prenode using the memory manager input_ptr = mem_manager_->MallocOutputMem(prenode_index.first, prenode_index.second, type, total_size, addr_size[0].first, true); + + // Check if there is a valid reuse index if (valid_reuse_index != -1) { + // Update the communication address for the valid reuse index with the new input pointer auto &it = reuse_communication_address_[valid_reuse_index]; it.first = input_ptr; } } - for (const auto &iter : addr_size) { +// Iterate over each element in the addr_size container using a range-based for loop +for (const auto &iter : addr_size) { + + // Check if the first element of the current pair is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(iter.first); + + // Set the pointer of the first element to the input_ptr iter.first->set_ptr(input_ptr); + + // Increment the input_ptr by the value of the second element in the current pair input_ptr += iter.second; - } } +// End of the loop + +// Function to assign memory to the output of a node in the kernel runtime void KernelRuntime::AssignNodeOutputMem(MemType type, const AnfNodePtr &node, int index) { + // Check if the node is null, throw an exception if it is MS_EXCEPTION_IF_NULL(node); + + // Check if the memory manager is null, throw an exception if it is MS_EXCEPTION_IF_NULL(mem_manager_); + // Rest of the function implementation goes here... +} + + // Check if the type is kSomasReuseDynamicMem if (type == kSomasReuseDynamicMem) { + + // Call the KernelMemNotReuse function to check if memory reuse is not possible for the given node bool not_reuse = KernelMemNotReuse(node); + + // If memory reuse is not possible if (not_reuse) { + + // Set the type to kDynamicMem type = kDynamicMem; + + // Print an informational message indicating that memory reuse is disabled for the output of the node MS_LOG(INFO) << "Disable Memory Reuse for " << node->fullname_with_scope() << "'s output."; } } - auto kernel_mod = AnfAlgo::GetKernelMod(node); - MS_EXCEPTION_IF_NULL(kernel_mod); - auto output_sizes = kernel_mod->GetOutputSizeList(); - if (output_sizes.empty()) { - return; +// Get the kernel module associated with the given node +auto kernel_mod = AnfAlgo::GetKernelMod(node); +MS_EXCEPTION_IF_NULL(kernel_mod); + +// Get the list of output sizes from the kernel module +auto output_sizes = kernel_mod->GetOutputSizeList(); + +// If the output_sizes list is empty, return from the function +if (output_sizes.empty()) { + return; +} + +// Iterate over the output_sizes list +for (size_t i = 0; i < output_sizes.size(); ++i) { + // If the index is not kGetAllOuts and not equal to the current index, continue to the next iteration + if ((kGetAllOuts != index) && (SizeToInt(i) != index)) { + continue; } - for (size_t i = 0; i < output_sizes.size(); ++i) { - if ((kGetAllOuts != index) && (SizeToInt(i) != index)) { - continue; - } - if (NodeOutputDeviceAddressExist(node, i)) { - MS_LOG(DEBUG) << "Already malloc index:" << i; - continue; - } - MS_LOG(DEBUG) << "Assign Node:" << node->fullname_with_scope() << " output memory size:" << output_sizes[i]; - if (type == kStaticMem) { - MS_LOG(INFO) << "Assign Static Memory for Output node, size:" << output_sizes[i] - << " node:" << node->fullname_with_scope(); - } - std::string output_format = AnfAlgo::GetOutputFormat(node, i); + + // If the device address for the output node and index already exists, log a debug message and continue to the next iteration + if (NodeOutputDeviceAddressExist(node, i)) { + MS_LOG(DEBUG) << "Already malloc index:" << i; + continue; + } + + // Log a debug message indicating the assignment of memory for the output node + MS_LOG(DEBUG) << "Assign Node:" << node->fullname_with_scope() << " output memory size:" << output_sizes[i]; + + // If the type is kStaticMem, log an info message indicating the assignment of static memory for the output node + if (type == kStaticMem) { + MS_LOG(INFO) << "Assign Static Memory for Output node, size:" << output_sizes[i] + << " node:" << node->fullname_with_scope(); + } + + // Get the output format for the node and index + std::string output_format = AnfAlgo::GetOutputFormat(node, i); + // ... (code continues) + // Get the output device data type for the given node and index auto output_type = AnfAlgo::GetOutputDeviceDataType(node, i); + + // Create a device address with nullptr as the device pointer, output sizes, output format, output type, and node and index information auto device_address = CreateDeviceAddress(nullptr, output_sizes[i], output_format, output_type, {node, i}); + + // Check if the device address is null, throw an exception if it is MS_EXCEPTION_IF_NULL(device_address); + + // Allocate memory for the output using the memory manager, passing the node, index, type, output size, device address, and false for not being a dynamic output uint8_t *ptr = mem_manager_->MallocOutputMem(node, i, type, output_sizes[i], device_address, false); + + // Check if the pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(ptr); + + // Set the host shape of the device address using the runtime padding shape of the node and index device_address->set_host_shape(trans::GetRuntimePaddingShape(node, i)); + + // Set the output address of the device address for the given index and node AnfAlgo::SetOutputAddr(device_address, i, node.get()); } } +// Function to assign extra static memory to a tensor for a given node and index DeviceAddressPtr KernelRuntime::AssignExtraStaticMem(const TensorPtr &tensor, const AnfNodePtr &node, size_t index) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(mem_manager_); + + // Get the device address of the tensor auto tensor_address = std::dynamic_pointer_cast(tensor->device_address()); + + // Log the assignment of static memory for the output node MS_LOG(DEBUG) << "Assign Node:" << node->fullname_with_scope() << "Assign Static Memory for Output node, size:" << tensor_address->size(); + + // Create a new device address with the same properties as the tensor's device address auto device_address = CreateDeviceAddress(nullptr, tensor_address->size(), tensor_address->format(), tensor_address->type_id(), {node, index}); MS_EXCEPTION_IF_NULL(device_address); + + // Allocate memory for the device address using the memory manager uint8_t *ptr = mem_manager_->MallocOutputMem(node, index, kStaticMem, tensor_address->size(), device_address, false); MS_EXCEPTION_IF_NULL(ptr); + + // Return the device address return device_address; } +// Function to assign a tensor value to a value node in the kernel runtime + void KernelRuntime::AssignValueNodeTensor(const ValueNodePtr &value_node, const ValuePtr &node_value, size_t output_idx) { + // Check if the value node and node value are not null MS_EXCEPTION_IF_NULL(value_node); MS_EXCEPTION_IF_NULL(node_value); + + // Check if the memory manager is not null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Get the instance of the MsContext auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); + + // Create a vector to store the tensors std::vector tensors; + + // Convert the node value to tensors TensorValueToTensor(node_value, &tensors); - // Graph id should be passed to record static memory if profiling is enabled. + + // Get the kernel info from the value node auto kernel_info = dynamic_cast(value_node->kernel_info()); MS_EXCEPTION_IF_NULL(kernel_info); + + // Get the graph id from the kernel info uint32_t graph_id = kernel_info->graph_id(); + + // Iterate through each tensor in the vector for (const auto &tensor : tensors) { + // Check if the tensor is null if (tensor == nullptr) { MS_LOG(WARNING) << "Tensor is null"; return; } + + // Get the device address of the tensor auto output_address = std::dynamic_pointer_cast(tensor->device_address()); + + // Check if the output address is not null and the device type matches the target device address type if (output_address != nullptr && output_address->DeviceType() == GetTargetDeviceAddressType()) { + // Continue with the assignment of the tensor to the value node + // ... + } + } +} + // Set the output address of the tensor using the AnfAlgo::SetOutputAddr function + // The address is obtained by casting the device address of the tensor to a shared pointer of DeviceAddress + // The output index is incremented by 1 each time this function is called + // The value_node is passed as a parameter to identify the node to which the output belongs AnfAlgo::SetOutputAddr(std::dynamic_pointer_cast(tensor->device_address()), output_idx++, value_node.get()); continue; } + // Calculate the size of the tensor in bytes size_t tensor_size = LongToSize(tensor->data().nbytes()); + // Get the size of the node's output tensor memory auto node_size = AnfAlgo::GetOutputTensorMemSize(value_node, output_idx); + // Get the data type of the output tensor TypeId output_type_id = AnfAlgo::GetOutputDeviceDataType(value_node, output_idx); + // If the output type is unknown, get the inferred data type if (output_type_id == kTypeUnknown) { output_type_id = common::AnfAlgo::GetOutputInferDataType(value_node, output_idx); } + // Get the output format of the tensor auto output_format = AnfAlgo::GetOutputFormat(value_node, output_idx); + // Create a device address with the specified parameters DeviceAddressPtr address = CreateDeviceAddress(nullptr, node_size, output_format, output_type_id, {value_node, output_idx}); + // Set the host shape of the address using the trans::GetRuntimePaddingShape function address->set_host_shape(trans::GetRuntimePaddingShape(value_node, output_idx)); + // Set the from_persistent_mem flag of the address to true address->set_from_persistent_mem(true); + // Check if the address is null MS_EXCEPTION_IF_NULL(address); + // Check if the MS_CTX_ENABLE_PYNATIVE_INFER flag is enabled and there is enough device memory if (ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_INFER) && !mem_manager_->MallocMemFromMemPool(address, node_size)) { + // If there is not enough memory, throw an exception with the allocation size MS_LOG(EXCEPTION) << "Device memory isn't enough and alloc failed, alloc size:" << node_size; } else { + // Log an information message indicating that static memory is being assigned for a value node MS_LOG(INFO) << "Assign Static Memory for Value node, size:" << node_size << " node:" << value_node->fullname_with_scope(); + + // Allocate memory for the value node using the memory manager if (mem_manager_->MallocMem(kStaticMem, node_size, address, graph_id) == nullptr) { + // If memory allocation fails, throw an exception with an error message MS_LOG(EXCEPTION) << "Cannot alloc address when flag is: " << kStaticMem << ", tensor size is: " << node_size; } } + + // Set the output address of the value node using the AnfAlgo::SetOutputAddr function AnfAlgo::SetOutputAddr(address, output_idx, value_node.get()); + + // Synchronize the host memory with the device memory for the value node if (!address->SyncHostToDevice(trans::GetRuntimePaddingShape(value_node, 0), tensor_size, tensor->data_type(), tensor->data_c(), tensor->device_info().host_format_)) { + // If synchronization fails, throw an exception with an error message MS_EXCEPTION(NotExistsError) << "ValueNode SyncHostToDevice fail!" << value_node->DebugString() << "node format is" << AnfAlgo::GetOutputFormat(value_node, output_idx) << "node dtype is " @@ -1013,282 +2001,538 @@ void KernelRuntime::AssignValueNodeTensor(const ValueNodePtr &value_node, const } } +// Define the function `AssignStaticMemoryValueNode` belonging to the class `KernelRuntime` + void KernelRuntime::AssignStaticMemoryValueNode(const session::KernelGraph &graph) { + // Check if the memory manager is null, throw an exception if it is MS_EXCEPTION_IF_NULL(mem_manager_); + + // Log a debug message indicating the start of the function for the given graph ID MS_LOG(DEBUG) << "AssignStaticMemoryValueNode start for graph " << graph.graph_id(); + + // Get the instance of the global context for MindSpore auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); - // order the value nodes + + // Create a map to store the value nodes, with the key being the full name with scope of the node std::map value_nodes_map; + + // Iterate over the graph's value nodes and populate the map for (auto &node : graph.graph_value_nodes()) { MS_EXCEPTION_IF_NULL(node); value_nodes_map[node->fullname_with_scope()] = node; } + // ... (rest of the code) - for (auto &item : value_nodes_map) { +// Iterate over each item in the value_nodes_map using a range-based for loop +for (auto &item : value_nodes_map) { + + // Get the value_node from the current item auto value_node = item.second; + + // Check if the value_node is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(value_node); + + // Check if the output device address already exists for the value_node if (NodeOutputDeviceAddressExist(value_node, 0)) { - MS_LOG(DEBUG) << "value_node[" << value_node->DebugString() << "] address already exist"; - auto device_address = AnfAlgo::GetMutableOutputAddr(value_node, 0); - if (device_address->ptr_ == nullptr) { - if (ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_INFER)) { - if (!mem_manager_->MallocMemFromMemPool(device_address, device_address->size_)) { - MS_LOG(EXCEPTION) << "MallocMemFromMemPool failed"; - } - } else { - if (mem_manager_->MallocMem(kStaticMem, device_address->size_, device_address, graph.graph_id())) { - MS_LOG(EXCEPTION) << "MallocStaticMem failed"; - } + + // Print a debug message indicating that the address for the value_node already exists + MS_LOG(DEBUG) << "value_node[" << value_node->DebugString() << "] address already exist"; + + // Get the mutable output address for the value_node + auto device_address = AnfAlgo::GetMutableOutputAddr(value_node, 0); + + // Check if the pointer of the device address is null + if (device_address->ptr_ == nullptr) { + + // Check if the MS_CTX_ENABLE_PYNATIVE_INFER flag is set to true + if (ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_INFER)) { + + // Try to allocate memory from the memory pool for the device address + if (!mem_manager_->MallocMemFromMemPool(device_address, device_address->size_)) { + + // Throw an exception if memory allocation fails + MS_LOG(EXCEPTION) << "MallocMemFromMemPool failed"; + } + } else { + + // Try to allocate static memory for the device address + if (mem_manager_->MallocMem(kStaticMem, device_address->size_, device_address, graph.graph_id())) { + + // Throw an exception if static memory allocation fails + MS_LOG(EXCEPTION) << "MallocStaticMem failed"; + } + } } - } - continue; + + // Continue to the next iteration of the loop + continue; } + + // Get a reference to the value of the value_node auto &node_value = value_node->value(); - MS_EXCEPTION_IF_NULL(node_value); - MS_LOG(DEBUG) << "Malloc memory for " << value_node->fullname_with_scope(); +} + MS_EXCEPTION_IF_NULL(node_value); // Check if the pointer node_value is null, if so, throw an exception + MS_LOG(DEBUG) << "Malloc memory for " << value_node->fullname_with_scope(); // Log a debug message indicating that memory is being allocated for the value_node + + // Check if node_value is of type Tensor or ValueTuple if (node_value->isa() || node_value->isa()) { - AssignValueNodeTensor(value_node, node_value, 0); - } else if (node_value->isa()) { - const bool use_mem_from_memory_pool = ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_INFER) || - ms_context->get_param(MS_CTX_EXECUTION_MODE) == kPynativeMode; - auto address = CreateDeviceAddressForStringValue(node_value, use_mem_from_memory_pool, graph.graph_id()); - MS_EXCEPTION_IF_NULL(address); - address->set_from_persistent_mem(true); - AnfAlgo::SetOutputAddr(address, 0, value_node.get()); + AssignValueNodeTensor(value_node, node_value, 0); // Call the AssignValueNodeTensor function to assign the value_node with the node_value + } + // Check if node_value is of type StringImm + else if (node_value->isa()) { + // Check if the MS_CTX_ENABLE_PYNATIVE_INFER flag is set to true or the MS_CTX_EXECUTION_MODE is set to kPynativeMode + const bool use_mem_from_memory_pool = ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_INFER) || + ms_context->get_param(MS_CTX_EXECUTION_MODE) == kPynativeMode; + // Create a device address for the string value in node_value, using memory from the memory pool if use_mem_from_memory_pool is true + auto address = CreateDeviceAddressForStringValue(node_value, use_mem_from_memory_pool, graph.graph_id()); + MS_EXCEPTION_IF_NULL(address); // Check if the address is null, if so, throw an exception + address->set_from_persistent_mem(true); // Set the address to be from persistent memory + AnfAlgo::SetOutputAddr(address, 0, value_node.get()); // Set the output address of value_node to the created address } - } - MS_LOG(DEBUG) << "AssignStaticMemoryValueNode end"; + + MS_LOG(DEBUG) << "AssignStaticMemoryValueNode end"; // Log a debug message indicating that the AssignStaticMemoryValueNode function has ended } +// This function creates a device address for a string value. +// It takes a value and a flag indicating whether to use memory pool for allocation. +// It also takes a graph ID to identify the graph. DeviceAddressPtr KernelRuntime::CreateDeviceAddressForStringValue(const ValuePtr &value, bool use_mem_pool, uint32_t graph_id) { + // Get the string value from the input value auto value_string = GetValue(value); + + // Calculate the size of the tensor based on the string size size_t tensor_size = value_string.size(); + + // Create a device address with the calculated tensor size, default format, and UInt8 data type DeviceAddressPtr address = CreateDeviceAddress(nullptr, tensor_size, kOpFormat_DEFAULT, kNumberTypeUInt8); MS_EXCEPTION_IF_NULL(address); + + // Set the flag indicating that the device address is from persistent memory address->set_from_persistent_mem(true); + + // Get the global context instance auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); + + // Check if memory pool should be used and if memory allocation from the pool fails if (use_mem_pool && !mem_manager_->MallocMemFromMemPool(address, tensor_size)) { MS_LOG(EXCEPTION) << "Device memory isn't enough and alloc failed, alloc size:" << tensor_size; } else { + // If memory pool is not used or memory allocation from the pool succeeds, + // allocate memory from the static memory pool MS_LOG(INFO) << "Assign Static Memory for string Value node, size:" << tensor_size; if (mem_manager_->MallocMem(kStaticMem, tensor_size, address, graph_id) == nullptr) { MS_LOG(EXCEPTION) << "Cannot alloc address when flag is: " << kStaticMem << ", tensor size is: " << tensor_size; } } + + // Create a shape vector with dimensions [1, tensor_size] ShapeVector shape = {1, SizeToLong(tensor_size)}; + + // Synchronize the host memory with the device memory by copying the string data to the device if (!address->SyncHostToDevice(shape, tensor_size, kNumberTypeUInt8, value_string.data(), "DefaultFormat")) { MS_LOG(EXCEPTION) << "kValueNode SyncHostToDevice fail!"; } + // ... +} + } + + // Return the value of the variable "address" to the caller return address; } -void KernelRuntime::AssignDynamicMemory(const session::KernelGraph &graph) { - MS_EXCEPTION_IF_NULL(mem_manager_); - auto context_ptr = MsContext::GetInstance(); - MS_EXCEPTION_IF_NULL(context_ptr); - bool is_enable_mem_reuse = EnvConfigParser::GetInstance().GetSysMemreuse(); - auto mem_type = kDynamicMem; - auto &dump_json_parser = DumpJsonParser::GetInstance(); - if (dump_json_parser.e2e_dump_enabled() && dump_json_parser.dump_mode() == 0) { +// The function `AssignDynamicMemory` is a member function of the class `KernelRuntime` + +// Check if the memory manager is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(mem_manager_); + +// Get the instance of the `MsContext` class +auto context_ptr = MsContext::GetInstance(); + +// Throw an exception if the `context_ptr` is null +MS_EXCEPTION_IF_NULL(context_ptr); + +// Get the value of the `SysMemreuse` flag from the `EnvConfigParser` singleton instance +bool is_enable_mem_reuse = EnvConfigParser::GetInstance().GetSysMemreuse(); + +// Set the memory type to `kDynamicMem` +auto mem_type = kDynamicMem; + +// Get the instance of the `DumpJsonParser` class +auto &dump_json_parser = DumpJsonParser::GetInstance(); + +// Check if the end-to-end dump is enabled and the dump mode is set to dump all kernels +if (dump_json_parser.e2e_dump_enabled() && dump_json_parser.dump_mode() == 0) { + // Disable memory reuse by setting the `SysMemreuse` flag to false mindspore::EnvConfigParser::GetInstance().SetSysMemreuse(false); is_enable_mem_reuse = false; + // Print an informational message indicating that memory reuse is disabled when end-to-end dump is enabled and dump mode is set to dump all kernels MS_LOG(INFO) << "Disable Memory Reuse when e2e dump is enable and dump mode is set to dump all kernels"; - } +} if (is_enable_mem_reuse) { + // If memory reuse is enabled, log a message indicating so MS_LOG(INFO) << "Memory Reuse is enable..."; + + // Allocate dynamic memory for SOMAS mem_manager_->MallocSomasDynamicMem(graph); + + // Set the memory type to reuse dynamic memory mem_type = kSomasReuseDynamicMem; } else { + // If memory reuse is disabled, log a message indicating so MS_LOG(INFO) << "Memory Reuse is disable..."; } + + // Get the execution order of nodes in the graph auto &execution_nodes = graph.execution_order(); + + // Create a vector to store compute nodes std::vector compute_nodes; - // communication nodes first + + // Process communication nodes first for (auto &node : execution_nodes) { if (common::AnfAlgo::IsCommunicationOp(node)) { - // skip if the memory is already allocated + // Skip if the memory for the node is already allocated AssignCommunicationNodeMem(mem_type, node); } else { + // Add the node to the compute nodes vector compute_nodes.emplace_back(node); } } - // then compute nodes + // Iterate over each node in the compute_nodes vector for (auto &node : compute_nodes) { + + // Assign output memory for the current node based on the specified memory type AssignNodeOutputMem(mem_type, node, kGetAllOuts); + + // Assign workspace memory for the current node based on the specified memory type AssignWorkSpaceMem(mem_type, node); } } +// A function to assign workspace memory for a given node in the kernel runtime + void KernelRuntime::AssignWorkSpaceMem(MemType type, const AnfNodePtr &node) { + // Check if the node is null MS_EXCEPTION_IF_NULL(node); + + // Check if the memory manager is null MS_EXCEPTION_IF_NULL(mem_manager_); + + // Get the kernel module associated with the node auto kernel_mod = AnfAlgo::GetKernelMod(node); + + // Check if the kernel module is null MS_EXCEPTION_IF_NULL(kernel_mod); + + // Initialize the index to 0 size_t index = 0; + + // Iterate over the workspace size list of the kernel module for (auto &size : kernel_mod->GetWorkspaceSizeList()) { + // Check if the workspace address already exists for the given index if (AnfAlgo::WorkspaceAddrExist(node, index)) { + // Log an information message indicating that the operation has workspace device address MS_LOG(INFO) << "Op " << node->fullname_with_scope() << " has workspace device address"; + + // Return from the function return; } + + // Allocate workspace memory using the memory manager auto ptr = mem_manager_->MallocWorkSpaceMem(node, index, type, size); - AnfAlgo::SetWorkspaceAddr(CreateDeviceAddress(ptr, size, "", kTypeUnknown), index, node.get()); + + // Create a device address using the allocated memory pointer and size + auto device_address = CreateDeviceAddress(ptr, size, "", kTypeUnknown); + + // Set the workspace address for the given index in the node + AnfAlgo::SetWorkspaceAddr(device_address, index, node.get()); + + // Increment the index index++; } } +// Function to generate launch arguments for a kernel void KernelRuntime::GenLaunchArgs(const mindspore::kernel::KernelMod &kernel_mod, const mindspore::AnfNodePtr &kernel, KernelLaunchInfo *kernel_launch_info) { MS_EXCEPTION_IF_NULL(kernel); MS_EXCEPTION_IF_NULL(kernel_launch_info); + + // Cast the kernel node to a CNode auto cnode = kernel->cast(); MS_EXCEPTION_IF_NULL(cnode); + + // Check if the kernel is an atomic address clean operation if (common::AnfAlgo::GetCNodeName(cnode) == kAtomicAddrCleanOpName) { + // If it is, generate launch arguments for address clean operation return GenAddrCleanLaunchArgs(cnode, &(kernel_launch_info->inputs_)); } + + // Get the current execution mode from the context auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); auto skip_nop_node = (ms_context->get_param(MS_CTX_EXECUTION_MODE) != kPynativeMode); + + // Get the number of input tensors for the kernel size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel); + + // Iterate over the input tensors for (size_t i = 0; i < input_num; ++i) { + // Check if the input tensor is None if (common::AnfAlgo::IsNoneInput(kernel, i)) { continue; } + + // Get the real input index and the device address for the input tensor auto real_input = AnfAlgo::GetRealInputIndex(kernel, i); auto device_address = AnfAlgo::GetPrevNodeOutputAddr(kernel, real_input, skip_nop_node); MS_EXCEPTION_IF_NULL(device_address); - kernel::AddressPtr input = std::make_shared(); - MS_EXCEPTION_IF_NULL(input); - input->addr = device_address->ptr_; - MS_EXCEPTION_IF_NULL(input->addr); - input->size = device_address->size_; - kernel_launch_info->inputs_.emplace_back(input); + // ... } +} +// Create a shared pointer to an instance of the kernel::Address class and assign it to the variable 'input' +kernel::AddressPtr input = std::make_shared(); - for (size_t i = 0; i < kernel_mod.GetOutputSizeList().size(); ++i) { +// Check if the 'input' pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(input); + +// Assign the value of 'device_address->ptr_' to the 'addr' member variable of the 'input' object +input->addr = device_address->ptr_; + +// Check if the 'input->addr' pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(input->addr); + +// Assign the value of 'device_address->size_' to the 'size' member variable of the 'input' object +input->size = device_address->size_; + +// Add the 'input' object to the 'inputs_' vector of the 'kernel_launch_info' object +kernel_launch_info->inputs_.emplace_back(input); + +// Iterate over the output size list of the kernel module +for (size_t i = 0; i < kernel_mod.GetOutputSizeList().size(); ++i) { + + // Get the device address of the output tensor using the AnfAlgo::GetOutputAddr function auto device_address = AnfAlgo::GetOutputAddr(kernel, i, skip_nop_node); - kernel::AddressPtr output = std::make_shared(); - MS_EXCEPTION_IF_NULL(output); - output->addr = device_address->ptr_; - MS_EXCEPTION_IF_NULL(output->addr); - output->size = device_address->size_; - kernel_launch_info->outputs_.emplace_back(output); - } + // Create a shared pointer to a kernel::Address object + kernel::AddressPtr output = std::make_shared(); + + // Check if the output pointer is not null + MS_EXCEPTION_IF_NULL(output); + + // Set the address of the output tensor to the device address + output->addr = device_address->ptr_; + + // Check if the address is not null + MS_EXCEPTION_IF_NULL(output->addr); + + // Set the size of the output tensor to the device address size + output->size = device_address->size_; + + // Add the output tensor to the kernel launch info's outputs vector + kernel_launch_info->outputs_.emplace_back(output); +} + + // Iterate over the elements in the list returned by kernel_mod.GetWorkspaceSizeList() for (size_t i = 0; i < kernel_mod.GetWorkspaceSizeList().size(); ++i) { + + // Get the device address of the workspace at index i using AnfAlgo::GetWorkspaceAddr() auto device_address = AnfAlgo::GetWorkspaceAddr(kernel, i); + + // Create a shared pointer to a kernel::Address object kernel::AddressPtr workspace = std::make_shared(); + + // Check if the workspace pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(workspace); + + // Set the address of the workspace to the device address pointer obtained earlier workspace->addr = device_address->ptr_; + + // Check if the workspace address is null, throw an exception if it is MS_EXCEPTION_IF_NULL(workspace->addr); + + // Set the size of the workspace to the size obtained from the device address workspace->size = device_address->size_; + + // Add the workspace to the list of workspaces in kernel_launch_info kernel_launch_info->workspaces_.emplace_back(workspace); } } +// Check if the memory scheduler should be used bool KernelRuntime::UseMemScheduler() { + + // Get the instance of the MsContext class auto context_ptr = MsContext::GetInstance(); + + // Throw an exception if the context pointer is null MS_EXCEPTION_IF_NULL(context_ptr); + + // Check if the MS_CTX_ENABLE_MEM_SCHEDULER parameter is set to true if (!context_ptr->get_param(MS_CTX_ENABLE_MEM_SCHEDULER)) { return false; } - // Not use MemScheduler when running single op + + // Check if the MS_CTX_ENABLE_PYNATIVE_INFER parameter is set to false + // and the MS_CTX_EXECUTION_MODE parameter is not equal to kPynativeMode return (!context_ptr->get_param(MS_CTX_ENABLE_PYNATIVE_INFER) && (context_ptr->get_param(MS_CTX_EXECUTION_MODE) != kPynativeMode)); } +// This function generates kernel events for a given kernel graph void KernelRuntime::GenKernelEvents(const session::KernelGraph &graph) { + + // Get the execution order of kernels in the graph auto &kernels = graph.execution_order(); + + // If there are no kernels in the graph or if kernel events for this graph have already been generated, return if (kernels.empty() || graph_kernel_events_map_.find(graph.graph_id()) != graph_kernel_events_map_.end()) { return; } + + // Create a pair of maps to store pre-run and post-run events for each kernel auto kernel_events = std::pair>>, std::map>>>(); + + // Get references to the first and second maps in the pair auto &kernel_pre_run_events = kernel_events.first; auto &kernel_post_run_events = kernel_events.second; + + // Iterate over the kernels in the execution order for (size_t i = 0; i < kernels.size(); ++i) { + + // Get a reference to the current kernel auto &kernel = kernels[i]; + + // If the kernel is not a communication operation, skip it if (!common::AnfAlgo::IsCommunicationOp(kernel)) { continue; } + + // Create pre-run and post-run events for the kernel auto pre_event = CreateDeviceEvent(); auto post_event = CreateDeviceEvent(); + + // Check if the events were created successfully MS_EXCEPTION_IF_NULL(pre_event); MS_EXCEPTION_IF_NULL(post_event); + + // Set the wait stream for the pre-event to the communication stream pre_event->set_wait_stream(communication_stream_); + + // Set the record stream for the pre-event to the current stream pre_event->set_record_stream(stream_); - post_event->set_wait_stream(stream_); - post_event->set_record_stream(communication_stream_); + + // ... (rest of the code) + post_event->set_wait_stream(stream_); // Set the wait stream of the post_event object to the value of stream_ + post_event->set_record_stream(communication_stream_); // Set the record stream of the post_event object to the value of communication_stream_ + + // Add a lambda function to the kernel_pre_run_events vector for the current kernel kernel_pre_run_events[kernel].emplace_back([pre_event]() { - pre_event->RecordEvent(); - pre_event->WaitEvent(); + pre_event->RecordEvent(); // Record the pre_event + pre_event->WaitEvent(); // Wait for the pre_event to complete }); + + // Add a lambda function to the kernel_post_run_events vector for the current kernel kernel_post_run_events[kernel].emplace_back([post_event]() { post_event->RecordEvent(); }); - bool found_nearest_child = false; + + bool found_nearest_child = false; // Initialize a boolean variable to keep track if the nearest child has been found + + // Iterate over the kernels starting from the next index after the current kernel for (size_t j = i + 1; j < kernels.size(); ++j) { - auto &child = kernels[j]; - MS_EXCEPTION_IF_NULL(child); + auto &child = kernels[j]; // Get the child kernel at index j + MS_EXCEPTION_IF_NULL(child); // Check if the child kernel is not null + + // Check if the child kernel is a communication operation, if so, continue to the next iteration if (common::AnfAlgo::IsCommunicationOp(child)) { continue; } - auto input_size = child->inputs().size() - 1; + + auto input_size = child->inputs().size() - 1; // Get the number of inputs of the child kernel + + // Iterate over the inputs of the child kernel for (size_t k = 0; k < input_size; ++k) { auto kernel_index = - common::AnfAlgo::VisitKernelWithReturnType(common::AnfAlgo::GetInputNode(child, k), 0, true); + common::AnfAlgo::VisitKernelWithReturnType(common::AnfAlgo::GetInputNode(child, k), 0, true); // Get the kernel index of the input + + // Check if the kernel index of the input matches the current kernel index if (kernel_index.first == kernel) { - found_nearest_child = true; + found_nearest_child = true; // Set the found_nearest_child flag to true, indicating that the nearest child has been found break; } } + + // If a nearest child is found, add a lambda function to the pre-run events of the child kernel if (found_nearest_child) { kernel_pre_run_events[child].emplace_back([post_event]() { post_event->WaitEvent(); }); break; } } + + // If no nearest child is found, add a lambda function to the post-run events of the current kernel if (!found_nearest_child) { kernel_post_run_events[kernel].emplace_back([post_event]() { post_event->WaitEvent(); }); } } + + // Store the kernel events map for the current graph graph_kernel_events_map_[graph.graph_id()] = std::move(kernel_events); } +// Function to generate address clean launch arguments for a given kernel node void KernelRuntime::GenAddrCleanLaunchArgs(const CNodePtr &cnode, AddressPtrList *kernel_inputs, const std::shared_ptr &mem_scheduler) { MS_EXCEPTION_IF_NULL(cnode); MS_EXCEPTION_IF_NULL(kernel_inputs); + + // Check if the number of input nodes for the given cnode is not equal to the expected size if (cnode->inputs().size() != kAtomicCleanInputSize) { MS_LOG(EXCEPTION) << "Atomic Addr clean Node Input nodes not equal 2."; } + MS_EXCEPTION_IF_NULL(cnode->inputs()[1]); auto pre_node = (cnode->inputs()[1])->cast(); - // set clean output address + + // Set clean output address if (common::AnfAlgo::HasNodeAttr(kAttrAtomicOutputIndexs, pre_node)) { #if defined(__APPLE__) auto clean_output_indexes = common::AnfAlgo::GetNodeAttr>(pre_node, kAttrAtomicOutputIndexs); #else auto clean_output_indexes = common::AnfAlgo::GetNodeAttr>(pre_node, kAttrAtomicOutputIndexs); #endif + + // Iterate over the clean output indexes and get the device address for each index for (auto index : clean_output_indexes) { auto device_address = AnfAlgo::GetOutputAddr(pre_node, index); + + // Create a new kernel address object and assign it to the input kernel::AddressPtr input = std::make_shared(); MS_EXCEPTION_IF_NULL(input); if (mem_scheduler != nullptr) { + // If the memory scheduler is not null, call the GetOrMallocAddress function passing the mem_scheduler, device_address, and input as arguments GetOrMallocAddress(mem_scheduler, device_address, input); } else { + // If the memory scheduler is null, assign the ptr_ value of device_address to input->addr input->addr = device_address->ptr_; + // Check if input->addr is null, if it is, throw an exception MS_EXCEPTION_IF_NULL(input->addr); } + // Assign the size_ value of device_address to input->size input->size = device_address->size_; + // Add the input to the kernel_inputs vector kernel_inputs->emplace_back(input); } + // Print the size of clean_output_indexes using the DEBUG log level MS_LOG(DEBUG) << "AtomicAddClean clean output size:" << clean_output_indexes.size(); } - // set clean workspace address + // Check if the pre_node has the attribute kAttrAtomicWorkspaceIndexs if (common::AnfAlgo::HasNodeAttr(kAttrAtomicWorkspaceIndexs, pre_node)) { + // Define the clean_workspaces_indexes vector based on the type of the elements (int or size_t) #if defined(__APPLE__) auto clean_workspaces_indexes = common::AnfAlgo::GetNodeAttr>(pre_node, kAttrAtomicWorkspaceIndexs); @@ -1296,274 +2540,541 @@ void KernelRuntime::GenAddrCleanLaunchArgs(const CNodePtr &cnode, AddressPtrList auto clean_workspaces_indexes = common::AnfAlgo::GetNodeAttr>(pre_node, kAttrAtomicWorkspaceIndexs); #endif + // Iterate over each index in the clean_workspaces_indexes vector for (const auto &index : clean_workspaces_indexes) { + + // Get the device address for the workspace using the pre_node and index auto device_address = AnfAlgo::GetWorkspaceAddr(pre_node, index); + + // Create a shared pointer to a kernel::Address object kernel::AddressPtr workspace = std::make_shared(); - MS_EXCEPTION_IF_NULL(workspace); + + // Check if mem_scheduler is not null if (mem_scheduler != nullptr) { + + // Call the GetOrMallocAddress function to get or allocate the address for the workspace GetOrMallocAddress(mem_scheduler, device_address, workspace); } else { + + // Set the address of the workspace to the device address pointer workspace->addr = device_address->ptr_; MS_EXCEPTION_IF_NULL(workspace->addr); } + + // Set the size of the workspace to the size of the device address workspace->size = device_address->size_; + + // Add the workspace to the kernel_inputs vector kernel_inputs->emplace_back(workspace); } } } -void KernelRuntime::LaunchKernelEvent(const std::map>> &kernel_events, - const AnfNodePtr &node) const { - if (kernel_events.find(node) == kernel_events.end()) { - return; - } +// This function is used to launch kernel events for a given node in a map of kernel events. +// It takes in two parameters: +// - kernel_events: a map that maps an AnfNodePtr (a node in the computation graph) to a vector of functions (kernel events) +// - node: the node for which the kernel events need to be launched +// Check if the given node exists in the kernel_events map +if (kernel_events.find(node) == kernel_events.end()) { + // If the node does not exist in the map, return and do nothing + return; +} + + // Iterate over each event in the vector of events associated with the given node for (auto &event : kernel_events.at(node)) { + + // Call the event function event(); } } +// Function to launch a kernel with Pynative profiling bool KernelRuntime::LaunchKernelWithPynativeProfiling(kernel::KernelMod *kernel_mod, const std::string &op_name, const KernelLaunchInfo &kernel_launch_info, void *stream) { MS_EXCEPTION_IF_NULL(kernel_mod); MS_EXCEPTION_IF_NULL(stream); + float cost_time = 0; + + // Create device time events for profiling auto start = CreateDeviceTimeEvent(); auto end = CreateDeviceTimeEvent(); MS_EXCEPTION_IF_NULL(start); MS_EXCEPTION_IF_NULL(end); + + // Set the stream for recording the events start->set_record_stream(stream); end->set_record_stream(stream); + + // Record the start event start->RecordEvent(); + + // Launch the kernel using the provided kernel module and stream bool ret = kernel_mod->LaunchKernel(kernel_launch_info, stream); + + // If the kernel launch fails, throw an exception with the kernel name if (!ret) { MS_LOG(EXCEPTION) << "Launch kernel failed, kernel name is : " << op_name; } + + // Record the end event end->RecordEvent(); + + // Synchronize the start and end events start->SyncEvent(); end->SyncEvent(); + + // Calculate the elapsed time between the start and end events start->ElapsedTime(&cost_time, end.get()); - MS_LOG(DEBUG) << "Launch kernel:" << op_name << " cost:" << cost_time / kBasicTimeTransferUnit; + + // Rest of the code is not provided, so we cannot comment on it + + // Return the result of the kernel launch return ret; } +// Use the MS_LOG macro to print a debug message with the kernel launch information +MS_LOG(DEBUG) << "Launch kernel:" << op_name << " cost:" << cost_time / kBasicTimeTransferUnit; + +// Return the value of the variable 'ret' to indicate the result of the function +return ret; + +// Definition of the function DebugStreamSync in the class KernelRuntime void KernelRuntime::DebugStreamSync(const CNodePtr &kernel) { + + // Get the instance of the MsContext class auto ms_context = MsContext::GetInstance(); + + // Throw an exception if the MsContext instance is null MS_EXCEPTION_IF_NULL(ms_context); + + // Get the value of the MS_CTX_ENABLE_PYNATIVE_SYNCHRONIZE parameter from the MsContext instance auto enable_sync_run = ms_context->get_param(MS_CTX_ENABLE_PYNATIVE_SYNCHRONIZE); + + // Check if synchronous execution is enabled if (enable_sync_run) { + + // Call the SyncStream function and check if it returns false if (!SyncStream()) { + + // Log an exception with the name of the kernel that failed to run MS_LOG(EXCEPTION) << "Op " << kernel->fullname_with_scope() << " run failed!"; } } } +// A member function of the KernelRuntime class that gets or allocates memory address + void KernelRuntime::GetOrMallocAddress(const std::shared_ptr &mem_scheduler, const DeviceAddress *device_address, const kernel::AddressPtr &kernel_addr) { - if (device_address->ptr_ != nullptr) { - kernel_addr->addr = device_address->ptr_; - } else { - kernel_addr->addr = mem_scheduler->GetOrMalloc(device_address, device_address->size_); + // Check if the device address pointer is not null + if (device_address->ptr_ != nullptr) { + // If not null, assign the device address pointer to the kernel address + kernel_addr->addr = device_address->ptr_; + } else { + // If null, call the GetOrMalloc function of the memory scheduler to allocate memory + kernel_addr->addr = mem_scheduler->GetOrMalloc(device_address, device_address->size_); + } +} + +// This function is used to assign kernel addresses to input nodes in the runtime. +// It takes in a memory scheduler, a kernel node, and a pointer to a KernelLaunchInfo object. +void KernelRuntime::AssignKernelAddress(const std::shared_ptr &mem_scheduler, const AnfNodePtr &kernel, + KernelLaunchInfo *kernel_launch_info) { + // Check if the kernel and kernel_launch_info pointers are not null + MS_EXCEPTION_IF_NULL(kernel); + MS_EXCEPTION_IF_NULL(kernel_launch_info); + + // Cast the kernel node to a CNode + auto cnode = kernel->cast(); + MS_EXCEPTION_IF_NULL(cnode); + + // Check if the CNode name is the atomic address clean operation + if (common::AnfAlgo::GetCNodeName(cnode) == kAtomicAddrCleanOpName) { + // If it is, call the GenAddrCleanLaunchArgs function to generate address clean launch arguments + return GenAddrCleanLaunchArgs(cnode, &(kernel_launch_info->inputs_), mem_scheduler); + } + + // Get the kernel module associated with the kernel node + auto kernel_mod = AnfAlgo::GetKernelMod(kernel); + MS_EXCEPTION_IF_NULL(kernel_mod); + + // Get the number of input tensors for the kernel + size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel); + + // Check if the kernel is an update parameter kernel + const auto update_parameter = common::AnfAlgo::IsUpdateParameterKernel(cnode); + + // Loop through each input tensor + for (size_t j = 0; j < input_num; ++j) { + // Get the real input index for the current input tensor + auto real_input = AnfAlgo::GetRealInputIndex(kernel, j); + + // Get the kernel with index for the current input tensor + auto kernel_with_index = common::AnfAlgo::GetPrevNodeOutput(kernel, real_input, true); + + // Get the index and input node from the kernel with index + auto index = kernel_with_index.second; + auto &input_node = kernel_with_index.first; + + // Get the device address for the input node at the specified index + auto device_address = AnfAlgo::GetOutputAddr(input_node, index, true); + MS_EXCEPTION_IF_NULL(device_address); + } +} +// Create a shared pointer to an instance of the Address class from the kernel namespace +kernel::AddressPtr input = std::make_shared(); + +// Call the GetOrMallocAddress function, passing in the mem_scheduler and device_address as arguments, and assign the result to the input pointer +GetOrMallocAddress(mem_scheduler, device_address, input); + +// Set the size of the input address to the size of the device_address +input->size = device_address->size_; + +// Add the input address to the inputs_ vector in the kernel_launch_info object +kernel_launch_info->inputs_.emplace_back(input); + +// Check if the update_parameter flag is true and if the input_node is of type Parameter +if (update_parameter && input_node->isa()) { + + // Cast the input_node to a ParameterPtr and assign it to the param variable + auto param = input_node->cast(); + + // Get the abstract value of the parameter and assign it to the abstract variable + auto abstract = param->abstract(); + + // Check if the abstract value is not null and if it is of type AbstractRef + MS_EXCEPTION_IF_NULL(abstract); + if (abstract->isa()) { + + // Call the UpdateHighPriorityMem function of the mem_scheduler, passing in the device_address as an argument + mem_scheduler->UpdateHighPriorityMem(device_address); } } -void KernelRuntime::AssignKernelAddress(const std::shared_ptr &mem_scheduler, const AnfNodePtr &kernel, - KernelLaunchInfo *kernel_launch_info) { - MS_EXCEPTION_IF_NULL(kernel); - MS_EXCEPTION_IF_NULL(kernel_launch_info); - auto cnode = kernel->cast(); - MS_EXCEPTION_IF_NULL(cnode); - if (common::AnfAlgo::GetCNodeName(cnode) == kAtomicAddrCleanOpName) { - return GenAddrCleanLaunchArgs(cnode, &(kernel_launch_info->inputs_), mem_scheduler); - } - auto kernel_mod = AnfAlgo::GetKernelMod(kernel); - MS_EXCEPTION_IF_NULL(kernel_mod); - size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel); - const auto update_parameter = common::AnfAlgo::IsUpdateParameterKernel(cnode); - for (size_t j = 0; j < input_num; ++j) { - auto real_input = AnfAlgo::GetRealInputIndex(kernel, j); - auto kernel_with_index = common::AnfAlgo::GetPrevNodeOutput(kernel, real_input, true); - auto index = kernel_with_index.second; - auto &input_node = kernel_with_index.first; - auto device_address = AnfAlgo::GetOutputAddr(input_node, index, true); - MS_EXCEPTION_IF_NULL(device_address); - kernel::AddressPtr input = std::make_shared(); - GetOrMallocAddress(mem_scheduler, device_address, input); - input->size = device_address->size_; - kernel_launch_info->inputs_.emplace_back(input); - if (update_parameter && input_node->isa()) { - auto param = input_node->cast(); - auto abstract = param->abstract(); - MS_EXCEPTION_IF_NULL(abstract); - if (abstract->isa()) { - mem_scheduler->UpdateHighPriorityMem(device_address); - } - } - } - + // Iterate over the output size list of the kernel module for (size_t j = 0; j < kernel_mod->GetOutputSizeList().size(); ++j) { + + // Get the device address of the output tensor at index j auto device_address = AnfAlgo::GetOutputAddr(kernel, j, true); + + // Create a shared pointer to a kernel::Address object kernel::AddressPtr output = std::make_shared(); + + // Get or allocate the address for the output tensor using the memory scheduler GetOrMallocAddress(mem_scheduler, device_address, output); + + // Set the size of the output tensor address to the size of the device address output->size = device_address->size_; + + // Add the output tensor address to the kernel launch info's outputs vector kernel_launch_info->outputs_.emplace_back(output); } + // Iterate over the workspace size list of the kernel module for (size_t i = 0; i < kernel_mod->GetWorkspaceSizeList().size(); ++i) { + + // Get the device address of the workspace at index i auto device_address = AnfAlgo::GetWorkspaceAddr(kernel, i); + + // Create a shared pointer to a kernel::Address object kernel::AddressPtr workspace = std::make_shared(); + + // Get or allocate the address for the workspace using the memory scheduler GetOrMallocAddress(mem_scheduler, device_address, workspace); + + // Set the size of the workspace to the size of the device address workspace->size = device_address->size_; + + // Add the workspace to the list of workspaces in the kernel launch info kernel_launch_info->workspaces_.emplace_back(workspace); } } +// This function is used to synchronize the output tensors of a kernel node in a computation graph. +// It takes in a memory scheduler, the computation graph, and the kernel node as input. + void KernelRuntime::SyncNodeOutputTensors(const std::shared_ptr &mem_scheduler, const session::KernelGraph &graph, const AnfNodePtr &kernel) { + // Check if the memory scheduler and kernel node are not null MS_EXCEPTION_IF_NULL(mem_scheduler); MS_EXCEPTION_IF_NULL(kernel); + + // Get the kernel module associated with the kernel node auto kernel_mod = AnfAlgo::GetKernelMod(kernel); MS_EXCEPTION_IF_NULL(kernel_mod); + + // Loop through the input tensors of the kernel node for (size_t input_idx = 0; input_idx < kernel_mod->GetInputSizeList().size(); ++input_idx) { + // Get the previous node and output index connected to the current input index const auto input_node_index = common::AnfAlgo::GetPrevNodeOutput(kernel, input_idx, true); + + // Check if the previous node is not null and is a parameter node if (input_node_index.first != nullptr && input_node_index.first->isa()) { + // Synchronize the output tensor of the previous node SyncNodeOutputTensor(mem_scheduler, input_node_index, graph); } } + + // Loop through the output tensors of the kernel node for (size_t output_idx = 0; output_idx < kernel_mod->GetOutputSizeList().size(); ++output_idx) { + // Synchronize the output tensor of the kernel node SyncNodeOutputTensor(mem_scheduler, std::make_pair(kernel, output_idx), graph); } } +// Synchronize the output tensor of a node with the memory scheduler void KernelRuntime::SyncNodeOutputTensor(const std::shared_ptr &mem_scheduler, const KernelWithIndex &node_output_index, const session::KernelGraph &graph) { MS_EXCEPTION_IF_NULL(mem_scheduler); + + // If the node output index is null, there is nothing to synchronize if (node_output_index.first == nullptr) { return; } + + // Get the mutable output address of the node auto device_address = AnfAlgo::GetMutableOutputAddr(node_output_index, true); + + // Get the tensor associated with the node output index auto tensor = graph.GetNodeOutputTensor(node_output_index); + + // If the tensor is null, there is nothing to synchronize if (tensor == nullptr) { return; } + + // If the device address is null, set the sync status to host-to-device and return if (device_address == nullptr) { tensor->data_sync(false); tensor->set_device_address(nullptr); tensor->set_sync_status(kNeedSyncHostToDevice); return; } + + // If the stream synchronization fails, throw an exception if (!SyncStream()) { MS_LOG(EXCEPTION) << "SyncStream failed"; } +} + // Store the original pointer of the device address auto origin_ptr = device_address->ptr_; + + // Check if the device address pointer is null if (device_address->ptr_ == nullptr) { + // If it is null, allocate memory using the memory scheduler and assign the pointer to device_address device_address->ptr_ = mem_scheduler->GetOrMalloc(device_address.get(), device_address->size_); } + + // Set the device address of the tensor to device_address tensor->set_device_address(device_address); + + // Synchronize the data of the tensor from host to device tensor->data_sync(false); + + // Set the device address of the tensor to null tensor->set_device_address(nullptr); + + // Restore the original pointer of the device address device_address->ptr_ = origin_ptr; + + // Set the synchronization status of the tensor to kNeedSyncHostToDevice tensor->set_sync_status(kNeedSyncHostToDevice); } -void KernelRuntime::InitGraphInputTensors(const std::shared_ptr &mem_scheduler, - const session::KernelGraph &graph) { - MS_EXCEPTION_IF_NULL(mem_scheduler); - auto &input_nodes = graph.input_nodes(); - auto &input_tensors = graph.input_tensors(); - if (input_tensors.size() != input_nodes.size()) { - MS_LOG_EXCEPTION << "Invalid input tensor size:" << input_tensors.size() << " vs node size:" << input_nodes.size(); +// This function initializes the input tensors of a kernel graph using a memory scheduler and the graph itself. + +// Check if the memory scheduler is not null +MS_EXCEPTION_IF_NULL(mem_scheduler); + +// Get the input nodes and input tensors of the graph +auto &input_nodes = graph.input_nodes(); +auto &input_tensors = graph.input_tensors(); + +// Check if the number of input tensors matches the number of input nodes +if (input_tensors.size() != input_nodes.size()) { + MS_LOG_EXCEPTION << "Invalid input tensor size:" << input_tensors.size() << " vs node size:" << input_nodes.size(); +} + +// Clear the memory need initialization flag in the memory scheduler +mem_scheduler->ClearMemNeedInit(); + +// Iterate over the input tensors +for (size_t i = 0; i < input_tensors.size(); ++i) { + auto input_node = input_nodes[i]; + + // Skip if the input node is not a parameter or if the output address does not exist + if (!input_node->isa() || !AnfAlgo::OutputAddrExist(input_node, 0)) { + continue; } - mem_scheduler->ClearMemNeedInit(); - for (size_t i = 0; i < input_tensors.size(); ++i) { - auto input_node = input_nodes[i]; - if (!input_node->isa() || !AnfAlgo::OutputAddrExist(input_node, 0)) { - continue; - } - auto device_address = AnfAlgo::GetMutableOutputAddr(input_node, 0); - auto tensor = input_tensors[i]; - MS_EXCEPTION_IF_NULL(tensor); - auto tensor_address = std::dynamic_pointer_cast(tensor->device_address()); - const auto tensor_size = LongToSize(tensor->data().nbytes()); - bool need_sync = false; + + // Get the mutable output address of the input node + auto device_address = AnfAlgo::GetMutableOutputAddr(input_node, 0); + + // Get the input tensor + auto tensor = input_tensors[i]; + MS_EXCEPTION_IF_NULL(tensor); + + // Get the device address of the input tensor + auto tensor_address = std::dynamic_pointer_cast(tensor->device_address()); + + // Get the size of the tensor in bytes + const auto tensor_size = LongToSize(tensor->data().nbytes()); + + // Initialize the need_sync flag to false + bool need_sync = false; if (tensor->NeedSyncHostToDevice()) { - need_sync = true; + // If the tensor needs to be synchronized from host to device + need_sync = true; } else if (tensor_address != device_address) { - tensor->data_sync(false); - need_sync = true; + // If the tensor address is not the same as the device address + // Sync the tensor data from host to device and mark the need for synchronization + tensor->data_sync(false); + need_sync = true; } + if (mem_scheduler->HasDeviceMem(device_address.get())) { - device_address->set_ptr(nullptr); + // If the memory scheduler has the device memory + // Set the pointer of the device address to nullptr + device_address->set_ptr(nullptr); } + if (need_sync) { - const auto &shape = trans::GetRuntimePaddingShape(input_node, 0); - if (device_address->GetPtr() != nullptr) { - device_address->SyncHostToDevice(shape, LongToSize(tensor->data().nbytes()), tensor->data_type(), - tensor->data_c(), tensor->device_info().host_format_); - } else { - mem_scheduler->AddMemNeedInit(device_address.get()); - } + // If synchronization is needed + // Get the runtime padding shape of the input node + const auto &shape = trans::GetRuntimePaddingShape(input_node, 0); + + if (device_address->GetPtr() != nullptr) { + // If the device address pointer is not nullptr + // Synchronize the host data to the device + device_address->SyncHostToDevice(shape, LongToSize(tensor->data().nbytes()), tensor->data_type(), + tensor->data_c(), tensor->device_info().host_format_); + } else { + // If the device address pointer is nullptr + // Add the device address to the memory scheduler for initialization + mem_scheduler->AddMemNeedInit(device_address.get()); + } } + + // Set the memory priority to low MemPriority priority = kMemPriorityLow; + + // Get the parameter of the input node const auto ¶meter = input_node->cast(); + // Check if the parameter is a weight or if it has been updated in the graph if (common::AnfAlgo::IsParameterWeight(parameter) || graph.IsUpdatedParameter(parameter)) { - priority = kMemPriorityHigh; + // If either condition is true, set the priority to high + priority = kMemPriorityHigh; } + + // Initialize the memory scheduler with the device address, tensor data, tensor size, and priority mem_scheduler->Init(device_address.get(), tensor->data_c(), tensor_size, priority); + + // Set the synchronization status of the tensor to "no need to sync" tensor->set_sync_status(kNoNeedSync); } } +// A function to assign communication memory for a given kernel in a kernel graph void KernelRuntime::AssignCommunicationMem(const session::KernelGraph &graph) { + + // Iterate over the kernels in the execution order of the graph for (const auto &kernel : graph.execution_order()) { + + // Check if the kernel is a communication operation if (!common::AnfAlgo::IsCommunicationOp(kernel)) { - continue; + continue; // If not, skip to the next kernel } + + // Assign communication input from the memory pool for the current kernel AssignCommunicationInputFromMemoryPool(kernel); + + // Assign communication output from the memory pool for the current kernel AssignCommunicationOutputFromMemoryPool(kernel); } } -bool KernelRuntime::LaunchKernel(const session::KernelGraph &graph, const AnfNodePtr &kernel, - const std::shared_ptr &mem_scheduler, bool mock) { +// Check if the kernel pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(kernel); + + // Get the kernel module associated with the kernel node auto kernel_mod = AnfAlgo::GetKernelMod(kernel); + + // Check if the kernel module pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_mod); + + // Create a KernelLaunchInfo object to store information about the kernel launch KernelLaunchInfo kernel_launch_info; + + // Get the stream associated with the kernel module auto stream = kernel_mod->stream(); + + // If the stream is null, check if the kernel is a communication op if (stream == nullptr) { if (common::AnfAlgo::IsCommunicationOp(kernel)) { + // If it is a communication op, use the communication stream stream = communication_stream_; } else { + // Otherwise, use the default stream stream = stream_; } } + + // Initialize the return value to true bool ret = true; + + // If a memory scheduler is provided, call its PreCompute function if (mem_scheduler != nullptr) { ret = mem_scheduler->PreCompute(stream); + + // If the PreCompute function returns false, return false if (!ret) { return ret; } + } + // Call the AssignKernelAddress function with the given arguments AssignKernelAddress(mem_scheduler, kernel, &kernel_launch_info); + + // Cast the kernel to a CNodePtr auto cnode = kernel->cast(); + + // Check if the mock flag is true and if the cnode has the attribute "offload" set to true if (mock && common::AnfAlgo::HasNodeAttr(kAttrOffload, cnode) && common::AnfAlgo::GetNodeAttr(cnode, kAttrOffload)) { + + // Iterate over the output size list of the kernel module for (size_t i = 0; i < kernel_mod->GetOutputSizeList().size(); ++i) { + + // Get the device address of the output at index i auto device_address = AnfAlgo::GetOutputAddr(kernel, i, true); + + // Set the device address as offload in the memory scheduler mem_scheduler->SetOffload(device_address); } } } else if (!kernel_mod->GetInputsAddr().empty() || !kernel_mod->GetOutputsAddr().empty()) { + + // Set the inputs, outputs, and workspaces of the kernel launch info from the kernel module kernel_launch_info.inputs_ = kernel_mod->GetInputsAddr(); kernel_launch_info.outputs_ = kernel_mod->GetOutputsAddr(); kernel_launch_info.workspaces_ = kernel_mod->GetWorkSpacesAddr(); } else { + + // Generate the launch arguments for the kernel module and kernel GenLaunchArgs(*kernel_mod, kernel, &kernel_launch_info); } + + // Check if the mock flag is false if (!mock) { + + // Check if the pynative_mode_profiling_flag_ is true if (pynative_mode_profiling_flag_) { + + // Launch the kernel with pynative profiling ret = LaunchKernelWithPynativeProfiling(kernel_mod, kernel->fullname_with_scope(), kernel_launch_info, stream); } else { + // ... (code continues) ret = kernel_mod->LaunchKernel(kernel_launch_info, stream); } if (!ret) { @@ -1571,223 +3082,423 @@ bool KernelRuntime::LaunchKernel(const session::KernelGraph &graph, const AnfNod } } if (mem_scheduler != nullptr) { + // If a memory scheduler is provided if (!mock) { + // If not in mock mode SyncNodeOutputTensors(mem_scheduler, graph, kernel); } + // Post-compute operations for memory scheduling ret = mem_scheduler->PostCompute(stream); } + // Return the result of the execution return ret; } +// Define a function named "LaunchKernelMod" that returns a boolean value bool KernelRuntime::LaunchKernelMod(const session::KernelGraph &graph, bool mock) { + + // Get the instance of the MsContext singleton auto context_ptr = MsContext::GetInstance(); + + // Throw an exception if the context pointer is null MS_EXCEPTION_IF_NULL(context_ptr); + + // Declare a shared pointer to a MemScheduler object and initialize it to nullptr std::shared_ptr mem_scheduler = nullptr; + // Check if the memory scheduler should be used if (UseMemScheduler()) { - mem_scheduler = mem_scheduler_manager_.GetOrCreateMemScheduler(graph.graph_id()); - MS_EXCEPTION_IF_NULL(mem_scheduler); - mem_scheduler->Reset(); - mem_scheduler->Update(); - InitGraphInputTensors(mem_scheduler, graph); + // Get or create the memory scheduler for the given graph ID + mem_scheduler = mem_scheduler_manager_.GetOrCreateMemScheduler(graph.graph_id()); + MS_EXCEPTION_IF_NULL(mem_scheduler); + + // Reset the memory scheduler + mem_scheduler->Reset(); + + // Update the memory scheduler + mem_scheduler->Update(); + + // Initialize the input tensors of the graph using the memory scheduler + InitGraphInputTensors(mem_scheduler, graph); } - const auto &kernels = graph.execution_order(); - std::map>> kernel_pre_run_events; - std::map>> kernel_post_run_events; - auto events_iter = graph_kernel_events_map_.find(graph.graph_id()); - if (events_iter != graph_kernel_events_map_.end()) { - kernel_pre_run_events = events_iter->second.first; - kernel_post_run_events = events_iter->second.second; - } - for (size_t i = 0; i < kernels.size(); ++i) { - LaunchKernelEvent(kernel_pre_run_events, kernels[i]); - auto &kernel = kernels[i]; - MS_EXCEPTION_IF_NULL(kernel); - if (common::AnfAlgo::IsDynamicShape(kernel)) { - auto kernel_mod = AnfAlgo::GetKernelMod(kernel); - MS_EXCEPTION_IF_NULL(kernel_mod); - opt::dynamic_shape::InferOp(kernel); - kernel_mod->InitOp(kernel->user_data()); - KernelLaunchInfo kernel_launch_info; - device::KernelRuntime::GenLaunchArgs(*kernel_mod, kernel, &kernel_launch_info); +// Create a constant reference to the execution order of the kernels in the graph +const auto &kernels = graph.execution_order(); - // allocate workspace size +// Create two maps to store pre-run and post-run events for each kernel +std::map>> kernel_pre_run_events; +std::map>> kernel_post_run_events; + +// Find the events associated with the current graph in the graph_kernel_events_map_ +auto events_iter = graph_kernel_events_map_.find(graph.graph_id()); + +// If events are found for the current graph, assign the pre-run and post-run events to the respective maps +if (events_iter != graph_kernel_events_map_.end()) { + kernel_pre_run_events = events_iter->second.first; + kernel_post_run_events = events_iter->second.second; +} + +// Iterate over the kernels in the execution order +for (size_t i = 0; i < kernels.size(); ++i) { + + // Call the LaunchKernelEvent function with the pre-run events and the current kernel + LaunchKernelEvent(kernel_pre_run_events, kernels[i]); + + // Create a reference to the current kernel + auto &kernel = kernels[i]; + + // Check if the kernel is null + MS_EXCEPTION_IF_NULL(kernel); + + // Check if the kernel has dynamic shape + if (common::AnfAlgo::IsDynamicShape(kernel)) { + + // Get the kernel module associated with the kernel + auto kernel_mod = AnfAlgo::GetKernelMod(kernel); + MS_EXCEPTION_IF_NULL(kernel_mod); + + // Infer the shape of the kernel's output + opt::dynamic_shape::InferOp(kernel); + + // Initialize the kernel module with the initialization arguments + kernel_mod->InitOp(kernel->user_data()); + + // Create a KernelLaunchInfo object to store the launch arguments for the kernel + KernelLaunchInfo kernel_launch_info; + + // Generate the launch arguments for the kernel + device::KernelRuntime::GenLaunchArgs(*kernel_mod, kernel, &kernel_launch_info); + } +} + + // Allocate workspace size std::vector workspace_addr; + + // Check if the kernel type is TBE_KERNEL if (AnfAlgo::GetKernelType(kernel) == KernelType::TBE_KERNEL) { #ifdef ENABLE_D + // Get the workspace size list from the kernel module auto workspace_size_list = kernel_mod->GetWorkspaceSizeList(); + + // Get the current context and device ID auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); auto device_id = ms_context->get_param(MS_CTX_DEVICE_ID); + + // Get the runtime instance for the Ascend device and the specified device ID auto runtime_instance = KernelRuntimeManager::Instance().GetSingleKernelRuntime(kAscendDevice, device_id); MS_EXCEPTION_IF_NULL(runtime_instance); + // Iterate over each size in the workspace_size_list for (auto size : workspace_size_list) { - auto device_address_ptr = - std::make_shared(nullptr, size, kAscendDevice, device_id); + + // Create a shared pointer to an AscendDeviceAddress object with nullptr as the initial pointer, + // the given size, the device type kAscendDevice, and the device ID + auto device_address_ptr = std::make_shared(nullptr, size, kAscendDevice, device_id); + + // Set the is_ptr_persisted flag of the device_address_ptr to true device_address_ptr->set_is_ptr_persisted(true); + + // Allocate memory on the device using the MallocMem function of the runtime_instance + // with MemType::kDynamicMem, the given size, and the device_address_ptr auto device_ptr = runtime_instance->MallocMem(MemType::kDynamicMem, size, device_address_ptr); + + // Check if the device_ptr is nullptr, indicating that the memory allocation failed if (device_ptr == nullptr) { + // Log an exception with the error message and the kernel's full name with scope MS_LOG(EXCEPTION) << "MallocMem from memory pool failed. Node info :" << kernel->fullname_with_scope(); } - AddressPtr workspace_addr_ptr = - std::make_shared
(device_address_ptr->GetMutablePtr(), device_address_ptr->GetSize()); + + // Create a shared pointer to an Address object with the mutable pointer obtained from device_address_ptr + // and the size obtained from device_address_ptr + AddressPtr workspace_addr_ptr = std::make_shared
(device_address_ptr->GetMutablePtr(), device_address_ptr->GetSize()); + + // Add the workspace_addr_ptr to the workspace_addr vector workspace_addr.emplace_back(workspace_addr_ptr); } #endif } else { + // If the kernel_launch_info has a non-empty workspaces_ vector, assign it to the workspace_addr vector workspace_addr = kernel_launch_info.workspaces_; } + // Call the Launch function of the kernel module with the provided inputs, workspace address, outputs, and stream auto ret = kernel_mod->Launch(kernel_launch_info.inputs_, workspace_addr, kernel_launch_info.outputs_, stream_); + + // If the Launch function returns false, indicating failure if (!ret) { + // Print an error message with the full name of the kernel MS_LOG(ERROR) << "Launch kernel failed, kernel full name: " << kernel->fullname_with_scope(); + + // Return false to indicate failure return false; } + // Check if the stream synchronization was successful if (!SyncStream()) { + // If not, log an error message and return false MS_LOG(ERROR) << "SyncStream failed"; return false; } + + // Update the operation of the kernel module kernel_mod->UpdateOp(); } else { - // Skip transpose kernel with "nop_op" attr which is not hidden or removed in PyNative infer scenario. Transpose - // kernel, which is not supposed to be executed, is generated in TransDataSplit to support specific Transdata. - // And hard code here should be removed after new Transdata programme is implemented in the foreseeable future. + // Skip the transpose kernel with the "nop_op" attribute, which is not hidden or removed in the PyNative infer scenario. + // The transpose kernel, which is not supposed to be executed, is generated in TransDataSplit to support specific Transdata. + // This hard-coded logic should be removed after the new Transdata program is implemented in the foreseeable future. if (common::AnfAlgo::HasNodeAttr(kAttrNopOp, kernel)) { + // Iterate over the output tensors of the kernel for (size_t idx = 0; idx < common::AnfAlgo::GetOutputTensorNum(kernel); idx += 1) { + // Get the real input index and the device address of the previous node's mutable output auto real_input = AnfAlgo::GetRealInputIndex(kernel, idx); auto device_address = AnfAlgo::GetPrevNodeMutableOutputAddr(kernel, real_input); + // Set the output address of the kernel to the device address AnfAlgo::SetOutputAddr(device_address, idx, kernel.get()); } + // Continue to the next iteration of the loop continue; } + + // Launch the kernel and store the return value auto ret = LaunchKernel(graph, kernel, mem_scheduler, mock); + + // Check if the kernel launch was successful if (!ret) { + // If not, log an error message MS_LOG(ERROR) << "Launch kernel failed."; + } + } return false; } + // Perform kernel launch profiling for the current kernel KernelLaunchProfiling(kernel->fullname_with_scope()); + // Synchronize the debug stream for the current kernel DebugStreamSync(kernel); } + // Launch the kernel post-run events for the current kernel LaunchKernelEvent(kernel_post_run_events, kernels[i]); } + // If memory scheduler is enabled and we are not running in mock mode, synchronize the parameter graph with the memory scheduler if (UseMemScheduler() && !mock) { SyncParameter(graph, mem_scheduler); } + // Return true to indicate successful execution of the function return true; } +// Definition of the function "SyncParameter" belonging to the class "KernelRuntime" void KernelRuntime::SyncParameter(const session::KernelGraph &graph, const std::shared_ptr &mem_scheduler) { + // Check if the memory scheduler is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(mem_scheduler); + + // Get the input nodes and input tensors from the given graph auto &input_nodes = graph.input_nodes(); auto &input_tensors = graph.input_tensors(); + + // Check if the size of input tensors is equal to the size of input nodes if (input_tensors.size() != input_nodes.size()) { + // If the sizes are not equal, throw an exception with an error message MS_LOG_EXCEPTION << "Invalid input tensor size:" << input_tensors.size() << " vs node size:" << input_nodes.size(); } +} for (size_t i = 0; i < input_tensors.size(); ++i) { auto input_node = input_nodes[i]; + + // Check if the input node is a Parameter and if it has a valid output address if (!input_node->isa() || !AnfAlgo::OutputAddrExist(input_node, 0)) { continue; } + + // Get the mutable output address of the input node auto device_address = AnfAlgo::GetMutableOutputAddr(input_node, 0); MS_EXCEPTION_IF_NULL(device_address); + + // Cast the input node to a Parameter pointer auto parameter = input_node->cast(); MS_EXCEPTION_IF_NULL(parameter); + + // Check if the parameter is a weight or if it has been updated in the graph if (!common::AnfAlgo::IsParameterWeight(parameter) && !graph.IsUpdatedParameter(parameter)) { continue; } + + // Get the input tensor auto tensor = input_tensors[i]; MS_EXCEPTION_IF_NULL(tensor); + + // Check if the device memory scheduler already has the device memory for the address if (mem_scheduler->HasDeviceMem(device_address.get())) { + // Get or allocate device memory for the address with high priority auto device_ptr = mem_scheduler->GetOrMalloc(device_address.get(), device_address->size(), kMemPriorityHigh); device_address->set_ptr(device_ptr); + + // Set the device address of the tensor and mark it for device-to-host synchronization tensor->set_device_address(device_address); tensor->set_sync_status(kNeedSyncDeviceToHost); } + // Check if the given parameter in the graph is updated if (graph.IsUpdatedParameter(parameter)) { + + // If the parameter is updated, set the tensor to be updated by the device tensor->SetIsUpdateByDevice(); } } } +// Define the function `UseMemSchedulerIfNeeded` which takes a `KernelGraph` object as a parameter void KernelRuntime::UseMemSchedulerIfNeeded(const session::KernelGraph &graph) { - auto context_ptr = MsContext::GetInstance(); - MS_EXCEPTION_IF_NULL(context_ptr); - if (!UseMemScheduler()) { - return; - } - auto mem_scheduler = mem_scheduler_manager_.GetOrCreateMemScheduler(graph.graph_id()); - MS_EXCEPTION_IF_NULL(mem_scheduler); - if (mem_scheduler->optimized()) { - return; - } - mem_scheduler->SetMemHandler(mem_manager_); - mem_scheduler->SetTotalStep(graph.execution_order().size()); + // Get the instance of the `MsContext` singleton + auto context_ptr = MsContext::GetInstance(); + + // Throw an exception if the `context_ptr` is null + MS_EXCEPTION_IF_NULL(context_ptr); + + // Check if the memory scheduler should be used + if (!UseMemScheduler()) { + // If not, return from the function + return; + } + + // Get or create the memory scheduler for the given graph ID + auto mem_scheduler = mem_scheduler_manager_.GetOrCreateMemScheduler(graph.graph_id()); + + // Throw an exception if the `mem_scheduler` is null + MS_EXCEPTION_IF_NULL(mem_scheduler); + + // Check if the memory scheduler has already been optimized + if (mem_scheduler->optimized()) { + // If yes, return from the function + return; + } + + // Set the memory handler of the memory scheduler to the `mem_manager_` + mem_scheduler->SetMemHandler(mem_manager_); + + // Set the total step of the memory scheduler to the size of the execution order of the graph + mem_scheduler->SetTotalStep(graph.execution_order().size()); +} + + // Check if the memory scheduler needs to record an event if (mem_scheduler->need_record_event()) { + // Launch the kernel module with the graph and set the need_record_event flag to false (void)LaunchKernelMod(graph, true); mem_scheduler->set_need_record_event(false); } + + // Optimize the memory scheduler and store the result in the 'ret' variable auto ret = mem_scheduler->Optimize(); + + // Check if the optimization was successful if (!ret) { + // If not, log an exception with the error message indicating the failure to run the graph within the memory limit MS_LOG_EXCEPTION << "Can't run graph " << graph.graph_id() << " for memory limit."; } } +// Function to launch kernels for a given kernel graph bool KernelRuntime::LaunchKernels(const session::KernelGraph &graph) { + + // Use memory scheduler if needed UseMemSchedulerIfNeeded(graph); + + // Launch kernel modules for the graph if (!LaunchKernelMod(graph)) { MS_LOG(ERROR) << "LaunchKernelMod failed!"; return false; } + + // Get the current execution mode from the global context auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); + + // If the execution mode is graph mode if (ms_context->get_param(MS_CTX_EXECUTION_MODE) == kGraphMode) { + + // Synchronize the stream to ensure all kernels have completed execution if (!SyncStream()) { MS_LOG(ERROR) << "SyncStream failed"; return false; } } + + // Return true to indicate successful kernel launch return true; } +// A function to clear the runtime resources associated with a graph void KernelRuntime::ClearGraphRuntimeResource(uint32_t graph_id) { + + // Log an informational message indicating the graph ID being cleared MS_LOG(INFO) << "Clear graph:" << graph_id << " runtime resource"; } +// Check if ENABLE_CPU is defined and _WIN32 is not defined #if ((defined ENABLE_CPU) && (!defined _WIN32)) + +// Start an anonymous namespace to limit the visibility of the following function namespace { -// Finalize ps cache module before throw an exception. -void FinalizePsCache(const std::string &exception) { - ps::ps_cache_instance.Finalize(); - MS_LOG(EXCEPTION) << exception; -} -} // namespace + + // Function to finalize the ps cache module before throwing an exception + void FinalizePsCache(const std::string &exception) { + + // Call the Finalize function of the ps_cache_instance + ps::ps_cache_instance.Finalize(); + + // Log the exception message using MS_LOG and throw the exception + MS_LOG(EXCEPTION) << exception; + } + +} // namespace + +#endif // End of ENABLE_CPU and _WIN32 check + +// This function is a member function of the `KernelRuntime` class. +// It takes a `KernelGraph` object, a pointer to an `AnfNodePtr` variable, and a pointer to a `size_t` variable as input. +// It is used to find the first cache input index and the size of the first cache in the given `KernelGraph`. void KernelRuntime::GetFirstPSEmbeddingCache(const session::KernelGraph &graph, AnfNodePtr *const first_cache_input_index, size_t *const first_cache_size) { + // Iterate over the execution order of the kernels in the given `KernelGraph` for (const auto &kernel : graph.execution_order()) { MS_EXCEPTION_IF_NULL(kernel); + // Get the name of the current kernel auto kernel_name = common::AnfAlgo::GetCNodeName(kernel); + // Check if the kernel is not a GatherV2Op or SparseGatherV2Op if (kernel_name != kGatherV2OpName && kernel_name != kSparseGatherV2OpName) { + // If it is not, continue to the next kernel continue; } + // Get the input parameter and input index of the current kernel auto input_param = common::AnfAlgo::GetPrevNodeOutput(kernel, 0, true); auto input_index = common::AnfAlgo::GetPrevNodeOutput(kernel, 1, true); MS_EXCEPTION_IF_NULL(input_param.first); MS_EXCEPTION_IF_NULL(input_index.first); + // Get the full name with scope of the input parameter auto param_name = input_param.first->fullname_with_scope(); + // Check if the input parameter is a hash table in the PS cache if (!ps::ps_cache_instance.IsHashTable(param_name)) { + // If it is not, continue to the next kernel continue; } + // Query the size of the hash table in the PS cache auto size = ps::ps_cache_instance.QueryHashTableSize(param_name); + // Check if the input index is a CastOp while (input_index.first->isa() && (common::AnfAlgo::GetCNodeName(input_index.first) == kCastOpName)) { + // If it is, update the input index to the previous node output input_index = common::AnfAlgo::GetPrevNodeOutput(input_index.first, 0, true); + } + // Set the first cache input index and size to the found values + *first_cache_input_index = input_index.first; + *first_cache_size = size; + // Exit the function + return; + } +} MS_EXCEPTION_IF_NULL(input_index.first); } auto cnode = common::AnfAlgo::IsGraphKernel(input_index.first) @@ -1795,6 +3506,7 @@ void KernelRuntime::GetFirstPSEmbeddingCache(const session::KernelGraph &graph, : input_index.first; MS_EXCEPTION_IF_NULL(cnode); if (!cnode->isa()) { + // If the input index is not a CNode, finalize the parameter server cache and throw an error FinalizePsCache("The embeddingLookup whose input index should be a CNode but got " + cnode->fullname_with_scope()); } @@ -1803,66 +3515,114 @@ void KernelRuntime::GetFirstPSEmbeddingCache(const session::KernelGraph &graph, bool full_batch = parallel::ParallelContext::GetInstance()->full_batch(); if ((!full_batch && (input_index_node_name != kUniqueOpName)) || (full_batch && (input_index_node_name != kMinimumOpName))) { + // If the input index is not from the dataset and doesn't match the expected node names, throw an error MS_LOG(ERROR) << "The input index of the embeddingLookup(" << kernel->fullname_with_scope() << ") cache is from " << cnode->fullname_with_scope(); FinalizePsCache( "The embeddingLookup whose input index isn't from dataset doesn't support cache in parameter server training " "mode."); } + } } *first_cache_input_index = cnode; *first_cache_size = size; + // Log an informational message indicating the input index and cache size of the first embeddingLookup cache MS_LOG(INFO) << "The input index of the first embeddingLookup cache is from " << cnode->fullname_with_scope() << ", the cache size is " << size; return; } } +// Check if the sparse embedding cache is valid for the given node void KernelRuntime::CheckSparsePSEmbeddingCache(const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); + + // Get the previous node output of the given node auto pre_node = common::AnfAlgo::GetPrevNodeOutput(node, 1, true); MS_EXCEPTION_IF_NULL(pre_node.first); + + // Continue traversing the previous nodes until we find a node that is not a CNode or has a different name than kUniqueOpName while (pre_node.first->isa() && (common::AnfAlgo::GetCNodeName(pre_node.first) != kUniqueOpName)) { pre_node = common::AnfAlgo::GetPrevNodeOutput(pre_node.first, 0, true); MS_EXCEPTION_IF_NULL(pre_node.first); } + + // If the previous node is not a CNode or has a different name than kUniqueOpName, finalize the sparse embedding cache if (!(pre_node.first->isa()) || (common::AnfAlgo::GetCNodeName(pre_node.first) != kUniqueOpName)) { FinalizePsCache("The input_indices of kernel[SparseGatherV2] must be unique in parameter server cache mode"); } +} - pre_node = common::AnfAlgo::GetPrevNodeOutput(pre_node.first, 0, true); - MS_EXCEPTION_IF_NULL(pre_node.first); - while (pre_node.first->isa() && (common::AnfAlgo::GetCNodeName(pre_node.first) == kCastOpName)) { +// Get the previous node's output using the GetPrevNodeOutput function from the common::AnfAlgo namespace +pre_node = common::AnfAlgo::GetPrevNodeOutput(pre_node.first, 0, true); + +// Check if the pre_node.first pointer is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(pre_node.first); + +// Continue looping while the pre_node.first is a CNode and its name is equal to kCastOpName +while (pre_node.first->isa() && (common::AnfAlgo::GetCNodeName(pre_node.first) == kCastOpName)) { + // Get the previous node's output again pre_node = common::AnfAlgo::GetPrevNodeOutput(pre_node.first, 0, true); + + // Check if the pre_node.first pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(pre_node.first); - } - if (!(pre_node.first->isa()) || (common::AnfAlgo::GetCNodeName(pre_node.first) != kGetNextOpName)) { +} + +// Check if the pre_node.first is not a CNode or its name is not equal to kGetNextOpName +if (!(pre_node.first->isa()) || (common::AnfAlgo::GetCNodeName(pre_node.first) != kGetNextOpName)) { + // Call the FinalizePsCache function with an error message as the argument FinalizePsCache( "The input indices of kernel[Unique] must be produced from dataset directly and the indices value can not be " "changed before delivering to kernel[Unique] in parameter server cache mode."); - } } +// Check if the current graph supports PS embedding cache void KernelRuntime::CheckIfSupportPSEmbeddingCache(const session::KernelGraph &graph) { + // Initialize variables to store the first cache input index and size AnfNodePtr first_cache_input_index = nullptr; size_t first_cache_size = 0; + + // Get the first PS embedding cache in the graph GetFirstPSEmbeddingCache(graph, &first_cache_input_index, &first_cache_size); + + // Throw an exception if the first cache input index is null MS_EXCEPTION_IF_NULL(first_cache_input_index); + + // Iterate through the execution order of the graph for (const auto &kernel : graph.execution_order()) { + // Throw an exception if the current kernel is null MS_EXCEPTION_IF_NULL(kernel); + + // Get the name of the current kernel auto kernel_name = common::AnfAlgo::GetCNodeName(kernel); + + // Skip the current kernel if it is not a GatherV2 or SparseGatherV2 operation if (kernel_name != kGatherV2OpName && kernel_name != kSparseGatherV2OpName) { continue; } + + // Get the input parameter and index of the current kernel auto input_param = common::AnfAlgo::GetPrevNodeOutput(kernel, 0, true); auto input_index = common::AnfAlgo::GetPrevNodeOutput(kernel, 1, true); + + // Throw an exception if the input parameter or index is null MS_EXCEPTION_IF_NULL(input_param.first); MS_EXCEPTION_IF_NULL(input_index.first); + + // Skip the current kernel if the input parameter is not a Parameter node if (!input_param.first->isa()) { continue; } + + // Get the name of the input parameter auto param_name = input_param.first->fullname_with_scope(); + + // Skip the current kernel if the input parameter is not a hash table in the PS cache and the kernel is a SparseGatherV2 operation if (ps::ps_cache_instance.IsHashTable(param_name) && (kernel_name == kSparseGatherV2OpName)) { + // Perform some operation + } + } +} CheckSparsePSEmbeddingCache(kernel); } while (input_index.first->isa() && (common::AnfAlgo::GetCNodeName(input_index.first) == kCastOpName)) { @@ -1874,33 +3634,47 @@ void KernelRuntime::CheckIfSupportPSEmbeddingCache(const session::KernelGraph &g : input_index.first; MS_EXCEPTION_IF_NULL(cnode); if (cnode == first_cache_input_index) { + // Check if the parameter server cache is enabled for the embeddingLookup operation if (!ps::ps_cache_instance.IsHashTable(param_name)) { MS_LOG(ERROR) << "The embeddingLookup(" << kernel->fullname_with_scope() << ") doesn't enable cache."; FinalizePsCache( "All the embeddingLookups whose input indices are from dataset must enable cache at the same time when one " "of them enables cache in parameter server training mode."); } + // Query the cache size for the embeddingLookup operation auto size = ps::ps_cache_instance.QueryHashTableSize(param_name); if (size != first_cache_size) { MS_LOG(ERROR) << "The cache size(" << size << ") of embeddingLookup(" << kernel->fullname_with_scope() - << ") is not the same as other embeddingLookup cache size(" << first_cache_size << ")."; - FinalizePsCache("The cache sizes of embeddingLookups are not the same in parameter server training mode."); - } - } else if (ps::ps_cache_instance.IsHashTable(param_name)) { - MS_LOG(ERROR) << "The input index of the embeddingLookup(" << kernel->fullname_with_scope() << ") cache is from " - << cnode->fullname_with_scope(); - FinalizePsCache( - "The embeddingLookup whose input index isn't from dataset doesn't support cache in parameter server training " - "mode."); - } else if (cnode->isa() && (common::AnfAlgo::GetCNodeName(cnode) == kGetNextOpName)) { - MS_LOG(ERROR) << "The EmbeddingLookup kernel(" << kernel->fullname_with_scope() << ") doesn't enable cache."; - FinalizePsCache( - "All EmbeddingLookup kernels whose input indices are from dataset must enable cache at " - "the same time and parameter 'sparse' must be equal to the value of 'enable_sparse' in " - "context setting in parameter server training mode."); - } - } +// Check if the code is being compiled in a specific mode +#ifdef ENABLE_TRAINING_MODE + +// Include the necessary header files + +// Check if the cache sizes of embeddingLookups are the same in parameter server training mode +if (ps::ps_cache_instance.IsCacheSizeDifferent()) { + // Print an error message indicating that the cache sizes are different + MS_LOG(ERROR) << "The cache size of embeddingLookup(" << kernel->fullname_with_scope() << ") is not the same as other embeddingLookup cache size(" << first_cache_size << ")."; + // Finalize the parameter server cache + FinalizePsCache("The cache sizes of embeddingLookups are not the same in parameter server training mode."); +} +// Check if the input index of the embeddingLookup cache is from the dataset +else if (ps::ps_cache_instance.IsHashTable(param_name)) { + // Print an error message indicating that the input index is from the dataset + MS_LOG(ERROR) << "The input index of the embeddingLookup(" << kernel->fullname_with_scope() << ") cache is from " << cnode->fullname_with_scope(); + // Finalize the parameter server cache + FinalizePsCache("The embeddingLookup whose input index isn't from dataset doesn't support cache in parameter server training mode."); +} +// Check if the kernel is an EmbeddingLookup kernel and doesn't enable cache +else if (cnode->isa() && (common::AnfAlgo::GetCNodeName(cnode) == kGetNextOpName)) { + // Print an error message indicating that the EmbeddingLookup kernel doesn't enable cache + MS_LOG(ERROR) << "The EmbeddingLookup kernel(" << kernel->fullname_with_scope() << ") doesn't enable cache."; + // Finalize the parameter server cache + FinalizePsCache("All EmbeddingLookup kernels whose input indices are from dataset must enable cache at the same time and parameter 'sparse' must be equal to the value of 'enable_sparse' in context setting in parameter server training mode."); } + #endif -} // namespace device + +// End of the namespace 'device' } // namespace mindspore + +// Closing the namespace "mindspore" \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc b/mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc index 8c4945da7c8..7295913b6c9 100644 --- a/mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc +++ b/mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc @@ -14,146 +14,295 @@ * limitations under the License. */ +// Include the header file for the kernel runtime manager #include "runtime/device/kernel_runtime_manager.h" + +// Include the header file for the log adapter utility #include "utils/log_adapter.h" + +// Check if ENABLE_CPU is defined and _WIN32 is not defined #if ((defined ENABLE_CPU) && (!defined _WIN32)) -#include "ps/ps_cache/ps_cache_manager.h" + + // Include the header file for the PS cache manager + #include "ps/ps_cache/ps_cache_manager.h" + #endif + +// Include the header file for the pynative task manager #include "backend/common/session/pynative_task_manager.h" namespace mindspore { namespace device { + +// Function to clear runtime resources void KernelRuntimeManager::ClearRuntimeResource() { - // Just remove PyNative tasks before runtime resource release. + + // Reset PyNative tasks before releasing runtime resources session::PynativeTaskManager::GetInstance().Reset(); + + // Sync embedding table if cache is enabled and running on worker node #if ((defined ENABLE_CPU) && (!defined _WIN32)) if (ps::PSContext::instance()->is_worker() && ps::PsDataPrefetch::GetInstance().cache_enable()) { ps::ps_cache_instance.SyncEmbeddingTable(); } #endif + + // Acquire a lock to ensure thread safety std::lock_guard guard(lock_); + + // Iterate over the runtime map and release device resources for each device for (auto &iter : runtime_map_) { MS_LOG(INFO) << "Release device " << iter.first; MS_EXCEPTION_IF_NULL(iter.second); iter.second->ReleaseDeviceRes(); } + + // Clear the runtime map runtime_map_.clear(); } +} // namespace device +} // namespace mindspore + +// A method to clear the runtime resources associated with a specific graph void KernelRuntimeManager::ClearGraphResource(uint32_t graph_id) { + + // Acquire a lock to ensure thread safety std::lock_guard guard(lock_); + + // Iterate over the runtime_map_ container for (auto &iter : runtime_map_) { + + // Log an informational message indicating the device and graph being cleared MS_LOG(INFO) << "Clear device " << iter.first << " graph " << graph_id << " runtime resource"; + + // Check if the kernel runtime is nullptr if (!iter.second) { + + // Log an error message indicating that the kernel runtime is nullptr MS_LOG(ERROR) << "Kernel runtime is nullptr"; + + // Continue to the next iteration of the loop continue; } + + // Call the ClearGraphRuntimeResource method of the kernel runtime to clear the runtime resources for the specified graph iter.second->ClearGraphRuntimeResource(graph_id); } } +// Definition of the static member function Instance() of the class KernelRuntimeManager + +// Return a reference to the static instance of KernelRuntimeManager KernelRuntimeManager &KernelRuntimeManager::Instance() { + + // Create a static instance of KernelRuntimeManager using default constructor static KernelRuntimeManager instance{}; + + // Return the reference to the static instance return instance; } +// Definition of the Register function in the KernelRuntimeManager class + void KernelRuntimeManager::Register(const std::string &device_name, KernelRuntimeCreator &&runtime_creator) { + + // Check if the device_name is already registered in the runtime_creators_ map if (runtime_creators_.find(device_name) == runtime_creators_.end()) { + + // If the device_name is not found, add it to the runtime_creators_ map (void)runtime_creators_.emplace(device_name, runtime_creator); } } +// Define a member function named GetDeviceKey in the class KernelRuntimeManager std::string KernelRuntimeManager::GetDeviceKey(const std::string &device_name, uint32_t device_id) { + + // Concatenate the device_name and device_id using the + operator and store the result in device_key std::string device_key = device_name + "_" + std::to_string(device_id); + + // Return the device_key return device_key; } +// Function to get a single instance of KernelRuntime based on device name and device ID KernelRuntime *KernelRuntimeManager::GetSingleKernelRuntime(const std::string &device_name, uint32_t device_id) { + + // Generate a unique key for the device name and device ID auto runtime_key = GetDeviceKey(device_name, device_id); + + // Check if the runtime for the given key already exists in the runtime map auto runtime_iter = runtime_map_.find(runtime_key); + + // If the runtime exists, return it if (runtime_iter != runtime_map_.end()) { return runtime_iter->second.get(); - } else if (!runtime_map_.empty()) { + } + // If the runtime does not exist, but there are other runtimes in the map + else if (!runtime_map_.empty()) { + + // Get the key of the first runtime in the map auto cur_runtime_key = runtime_map_.begin()->first; + + // Find the position of the last underscore in the key auto find_pos = cur_runtime_key.rfind('_'); + + // If the underscore is found if (find_pos != std::string::npos) { + + // Check if the size of the key is greater than the position of the underscore + 1 if (cur_runtime_key.size() > find_pos + 1) { + + // Get the current device ID from the key auto cur_device_id = cur_runtime_key.substr(find_pos + 1); + + // Throw an exception indicating that the device ID cannot be changed in the runtime MS_LOG(EXCEPTION) << "Can't change device id in runtime, already set device id: " << cur_device_id << ", set device id: " << device_id << " failed"; - } else { + } + // If the size of the key is not greater than the position of the underscore + 1 + else { + + // Throw an exception indicating an error in the current runtime key size MS_LOG(EXCEPTION) << "Can't change device id in runtime, current runtime_key size error, set device id: " << device_id << " failed"; } } } + + // If the runtime does not exist and there are no other runtimes in the map, get a new KernelRuntime instance return GetKernelRuntime(device_name, device_id); } +// Closing brace to end the main function +} +// GetKernelRuntime function of KernelRuntimeManager class KernelRuntime *KernelRuntimeManager::GetKernelRuntime(const std::string &device_name, uint32_t device_id) { + + // Generate a unique key for the runtime based on the device name and device id std::string runtime_key = GetDeviceKey(device_name, device_id); + + // Acquire a lock to ensure thread safety std::lock_guard guard(lock_); + + // Check if the runtime already exists in the runtime map auto runtime_iter = runtime_map_.find(runtime_key); if (runtime_iter != runtime_map_.end()) { + // If the runtime exists, return a pointer to it return runtime_iter->second.get(); } + + // If the runtime does not exist, create a new one std::shared_ptr kernel_runtime; + + // Find the runtime creator for the specified device name auto creator_iter = runtime_creators_.find(device_name); if (creator_iter != runtime_creators_.end()) { + // If a runtime creator is found, create a new instance of the kernel runtime MS_EXCEPTION_IF_NULL(creator_iter->second); kernel_runtime = (creator_iter->second)(); MS_EXCEPTION_IF_NULL(kernel_runtime); kernel_runtime->set_device_id(device_id); + + // Add the new runtime to the runtime map runtime_map_[runtime_key] = kernel_runtime; } else { + // If no runtime creator is found, throw an exception MS_LOG(EXCEPTION) << "No kernel runtime creator for " << device_name << " with device id " << device_id; } - - return kernel_runtime.get(); } -KernelRuntime *KernelRuntimeManager::GetCurrentKernelRuntime() { - auto ms_context = MsContext::GetInstance(); - MS_EXCEPTION_IF_NULL(ms_context); - uint32_t device_id = ms_context->get_param(MS_CTX_DEVICE_ID); - std::string device_name = ms_context->get_param(MS_CTX_DEVICE_TARGET); - return GetKernelRuntime(device_name, device_id); -} +// Return the value of the `kernel_runtime` variable using the `get()` function +return kernel_runtime.get(); +// Get the current instance of the MsContext +auto ms_context = MsContext::GetInstance(); + +// Throw an exception if the MsContext is null +MS_EXCEPTION_IF_NULL(ms_context); + +// Get the device ID from the MsContext +uint32_t device_id = ms_context->get_param(MS_CTX_DEVICE_ID); + +// Get the device name from the MsContext +std::string device_name = ms_context->get_param(MS_CTX_DEVICE_TARGET); + +// Return the kernel runtime for the specified device name and device ID +return GetKernelRuntime(device_name, device_id); + +// Release the kernel runtime for a specific device void KernelRuntimeManager::ReleaseKernelRuntime(const std::string &device_name, uint32_t device_id) { + + // Reset the Pynative task manager session::PynativeTaskManager::GetInstance().Reset(); + + // Generate the runtime key using the device name and device ID std::string runtime_key = GetDeviceKey(device_name, device_id); + + // Acquire a lock to ensure thread safety std::lock_guard guard(lock_); + + // Find the runtime corresponding to the runtime key in the runtime map auto runtime_iter = runtime_map_.find(runtime_key); + + // If the runtime is not found, return if (runtime_iter == runtime_map_.end()) { return; } + + // Get the runtime pointer from the iterator auto runtime = runtime_iter->second.get(); + + // If the runtime pointer is null, return if (runtime == nullptr) { return; } + + // If CPU is enabled and not on Windows, and the current context is a worker with caching enabled, + // synchronize the embedding table #if ((defined ENABLE_CPU) && (!defined _WIN32)) if (ps::PSContext::instance()->is_worker() && ps::PsDataPrefetch::GetInstance().cache_enable()) { ps::ps_cache_instance.SyncEmbeddingTable(); } #endif + + // Release the device resources of the runtime runtime->ReleaseDeviceRes(); + + // Erase the runtime from the runtime map runtime_map_.erase(runtime_iter); } +// Definition of the WaitTaskFinishOnDevice function in the KernelRuntimeManager class + void KernelRuntimeManager::WaitTaskFinishOnDevice() const { + + // Iterate over each key-value pair in the runtime_map_ for (const auto &iter : runtime_map_) { + + // Get the kernel_runtime object from the current key-value pair auto kernel_runtime = iter.second; + try { + + // Check if kernel_runtime is not null and if SyncStream returns false if (kernel_runtime != nullptr && !kernel_runtime->SyncStream()) { + + // Print an error message and return if SyncStream fails MS_LOG(ERROR) << "SyncStream failed"; return; } } catch (const std::exception &ex) { + + // Print an error message with the exception details and return if an exception occurs during SyncStream MS_LOG(ERROR) << "SyncStream failed, exception:" << ex.what(); return; } } } + +// End of the device namespace } // namespace device -} // namespace mindspore + +// End of the mindspore namespace +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/device/launch_mul.cc b/mindspore/ccsrc/runtime/device/launch_mul.cc index 94cd1a1ab94..df6c3c5afb5 100644 --- a/mindspore/ccsrc/runtime/device/launch_mul.cc +++ b/mindspore/ccsrc/runtime/device/launch_mul.cc @@ -14,73 +14,146 @@ * limitations under the License. */ +// Include the header file for launching multiplication on a device #include "runtime/device/launch_mul.h" + +// Include the header file for utility functions #include "abstract/utils.h" + +// Include the header file for single kernel graph in a session #include "backend/common/session/single_kernel_graph.h" + +// Include the header file for ANF runtime algorithm #include "backend/common/session/anf_runtime_algorithm.h" + +// Include the header file for ANF algorithm utilities #include "include/common/utils/anfalgo.h" + +// Include the header file for parallel context utilities #include "include/common/utils/parallel_context.h" +// Define the namespace "mindspore::device" namespace mindspore::device { -std::shared_ptr LaunchMul::ObtainMulKernelGraph() { - std::vector input_dtypes = {dtype_, dtype_}; - std::vector output_dtypes = {dtype_}; - // obtain input & output shapes - size_t dtype_size = abstract::TypeIdSize(dtype_); - if (dtype_size == 0) { - MS_LOG(EXCEPTION) << "Divide by zero."; - } - int64_t shape = SizeToLong(total_size_ / dtype_size); - std::vector> input_shapes = {{shape}, {1}}; - std::vector> output_shapes = {{static_cast(shape)}}; - auto mul_graph = session::SingleKernelGraph::ConstructKernelGraphBasedOnSingleOp( - kMulOpName, input_dtypes, input_shapes, output_dtypes, output_shapes); - MS_EXCEPTION_IF_NULL(mul_graph); - return mul_graph; + + // Define the function "ObtainMulKernelGraph" which returns a shared pointer to a KernelGraph object + std::shared_ptr LaunchMul::ObtainMulKernelGraph() { + + // Create a vector of input data types with the same data type as "dtype_" + std::vector input_dtypes = {dtype_, dtype_}; + + // Create a vector of output data types with the same data type as "dtype_" + std::vector output_dtypes = {dtype_}; + + // Calculate the size of the data type in bytes + size_t dtype_size = abstract::TypeIdSize(dtype_); + + // Check if the data type size is zero + if (dtype_size == 0) { + // If the data type size is zero, throw an exception with the message "Divide by zero." + MS_LOG(EXCEPTION) << "Divide by zero."; + } + + // Calculate the shape of the input and output tensors + int64_t shape = SizeToLong(total_size_ / dtype_size); + std::vector> input_shapes = {{shape}, {1}}; + std::vector> output_shapes = {{static_cast(shape)}}; + + // Construct a KernelGraph object based on a single operation with the name "kMulOpName" + auto mul_graph = session::SingleKernelGraph::ConstructKernelGraphBasedOnSingleOp( + kMulOpName, input_dtypes, input_shapes, output_dtypes, output_shapes); + + // Check if the mul_graph is null + MS_EXCEPTION_IF_NULL(mul_graph); + + // Return the mul_graph + return mul_graph; + } } kernel::KernelMod *LaunchMul::ObtainLaunchMulKernelMod() { if (mul_graph_ == nullptr) { - // construct mul kernel graph + // If the mul_graph_ is not yet constructed, construct it mul_graph_ = ObtainMulKernelGraph(); MS_EXCEPTION_IF_NULL(mul_graph_); - // kernel select + + // Perform kernel selection on the mul_graph_ KernelSelect(mul_graph_); - // kernel build + + // Build the kernels for the mul_graph_ KernelBuild(mul_graph_); } - // obtain kernel_mod + + // Obtain the kernel_mod for the mul_graph_ if (mul_graph_->execution_order().size() != 1) { + // If the execution order of the mul graph has more than one node, log an error MS_LOG(ERROR) << "the execution order of the mul graph should have only one node, however, it has " << mul_graph_->execution_order().size() << " nodes."; } + + // Return the kernel_mod of the first node in the execution order return AnfAlgo::GetKernelMod(mul_graph_->execution_order()[0]); } -void LaunchMul::ObtainMulInputsAddr() { - inputs_addr_.push_back(input1_addr_); +// This function is a member function of the class LaunchMul. +// It is used to obtain the memory addresses of the inputs for multiplication. - auto parallel_context = parallel::ParallelContext::GetInstance(); - MS_EXCEPTION_IF_NULL(parallel_context); - auto device_num = parallel_context->device_num(); - if (device_num == 0) { - MS_LOG(ERROR) << "device num can't be zero"; - } - input2_value_ = 1.0f / device_num; - auto size = abstract::TypeIdSize(dtype_); - auto input_size = AlignSizeForLaunchKernel(size * 1); - // alloc memory - input2_addr_ = AllocDeviceMem(input_size); - CopyHostMemToDevice(size, input_size); - inputs_addr_.push_back(input2_addr_); +void LaunchMul::ObtainMulInputsAddr() { + // Push the memory address of input1_addr_ into the inputs_addr_ vector + inputs_addr_.push_back(input1_addr_); } +// Get the instance of the parallel context +auto parallel_context = parallel::ParallelContext::GetInstance(); + +// Check if the parallel context is not null +MS_EXCEPTION_IF_NULL(parallel_context); + +// Get the number of devices from the parallel context +auto device_num = parallel_context->device_num(); + +// Check if the number of devices is zero +if (device_num == 0) { + // Log an error message if the number of devices is zero + MS_LOG(ERROR) << "device num can't be zero"; +} + +// Calculate the value of input2_value_ by dividing 1.0f by the number of devices +input2_value_ = 1.0f / device_num; + +// Calculate the size of the data type (dtype_) +auto size = abstract::TypeIdSize(dtype_); + +// Calculate the input size by multiplying the size of the data type by 1 and aligning it for launch kernel +auto input_size = AlignSizeForLaunchKernel(size * 1); + +// Allocate memory on the device +input2_addr_ = AllocDeviceMem(input_size); + +// Copy host memory to the device memory +CopyHostMemToDevice(size, input_size); + +// Push the address of input2_addr_ to the inputs_addr_ vector +inputs_addr_.push_back(input2_addr_); + +// Definition of the function "FreeInputDeviceMemory" belonging to the "LaunchMul" class + void LaunchMul::FreeInputDeviceMemory() { + + // Set the value of "input1_addr_" to nullptr, effectively freeing the memory input1_addr_ = nullptr; + + // Check if "input2_addr_" is not nullptr if (input2_addr_ != nullptr) { + + // Call the function "FreeDeviceMem" to free the memory pointed to by "input2_addr_" FreeDeviceMem(input2_addr_); + + // Set the value of "input2_addr_" to nullptr, effectively freeing the memory input2_addr_ = nullptr; } + + // Clear the contents of the "inputs_addr_" vector, effectively freeing the memory inputs_addr_.clear(); } -} // namespace mindspore::device + +// End of the "mindspore::device" namespace \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/hardware/device_context_manager.cc b/mindspore/ccsrc/runtime/hardware/device_context_manager.cc index a7f1f4cac7a..133ef9c196b 100644 --- a/mindspore/ccsrc/runtime/hardware/device_context_manager.cc +++ b/mindspore/ccsrc/runtime/hardware/device_context_manager.cc @@ -14,77 +14,153 @@ * limitations under the License. */ -#include "runtime/hardware/device_context_manager.h" +// Include the header file "runtime/hardware/device_context_manager.h" which contains the declaration of the DeviceContextManager class or related functions. +// Define the namespace "mindspore" namespace mindspore { + +// Define the nested namespace "device" within the "mindspore" namespace namespace device { + +// Define the member function "GetInstance" of the class "DeviceContextManager" DeviceContextManager &DeviceContextManager::GetInstance() { + + // Create a static instance of the class "DeviceContextManager" using empty braces to invoke the default constructor static DeviceContextManager instance{}; + + // Return the static instance of the class "DeviceContextManager" return instance; } +} // End of namespace "device" +} // End of namespace "mindspore" + +// Definition of the Register function in the DeviceContextManager class void DeviceContextManager::Register(const std::string &device_name, DeviceContextCreator &&device_context_creator) { + + // Check if the device_name is not already registered in the device_context_creators_ map if (device_context_creators_.find(device_name) == device_context_creators_.end()) { + + // If the device_name is not found, add it to the device_context_creators_ map (void)device_context_creators_.emplace(device_name, device_context_creator); } } +// A member function of the DeviceContextManager class that clears all device contexts + void DeviceContextManager::ClearDeviceContexts() { + + // Iterate over each device context in the device_contexts_ map for (auto &iter : device_contexts_) { + + // Log an informational message indicating the device being released MS_LOG(INFO) << "Release device " << iter.first; + + // Check if the device context is not null MS_EXCEPTION_IF_NULL(iter.second); + + // Call the Destroy() function of the device context to release any resources iter.second->Destroy(); } + + // Clear the device_contexts_ map device_contexts_.clear(); } -DeviceContext *DeviceContextManager::GetOrCreateDeviceContext(const DeviceContextKey &device_context_key) { - std::string device_context_key_str = device_context_key.ToString(); +// Get the string representation of the device context key by calling the ToString() function on the device_context_key object +// Find the iterator in the map `device_contexts_` that corresponds to the given `device_context_key_str` auto device_context_iter = device_contexts_.find(device_context_key_str); + + // If the iterator is not equal to the end iterator of the map, it means the key was found if (device_context_iter != device_contexts_.end()) { + // Return a pointer to the value associated with the key in the map return device_context_iter->second.get(); } - std::shared_ptr device_context; - auto creator_iter = device_context_creators_.find(device_context_key.device_name_); - if (creator_iter != device_context_creators_.end()) { +// Declare a shared pointer to a DeviceContext object +std::shared_ptr device_context; + +// Find the creator iterator in the device_context_creators_ map using the device name from device_context_key +auto creator_iter = device_context_creators_.find(device_context_key.device_name_); + +// Check if the creator iterator is not equal to the end iterator of the device_context_creators_ map +if (creator_iter != device_context_creators_.end()) { + // Call the creator function pointed to by the iterator with device_context_key as the argument device_context = (creator_iter->second)(device_context_key); + + // Check if the device_context is null MS_EXCEPTION_IF_NULL(device_context); + + // Add the device_context to the device_contexts_ map using device_context_key_str as the key device_contexts_[device_context_key_str] = device_context; - } else { +} else { + // Log an exception if the creator iterator is equal to the end iterator, indicating that the device context creation failed MS_LOG(EXCEPTION) << "Create device context failed, please make sure target device:" << device_context_key.device_name_ << " is available."; - } - return device_context.get(); } -void DeviceContextManager::UpdateDeviceContextKey(const DeviceContextKey &old_key, const DeviceContextKey &new_key) { - std::string old_key_str = old_key.ToString(); - std::string new_key_str = new_key.ToString(); +// Return the raw pointer of the device_context +return device_context.get(); - auto handle = device_contexts_.extract(old_key_str); - if (handle.empty()) { +// Update the device context key by converting the old key and new key to strings + +// Convert the old key to a string +std::string old_key_str = old_key.ToString(); + +// Convert the new key to a string +std::string new_key_str = new_key.ToString(); + +// Extract the value associated with the key "old_key_str" from the "device_contexts_" container +auto handle = device_contexts_.extract(old_key_str); + +// Check if the extracted handle is empty (i.e., the key was not found in the container) +if (handle.empty()) { + // If the handle is empty, log an exception with an error message indicating that the device context for "old_key_str" was not found MS_LOG(EXCEPTION) << "Can not find device context for: " << old_key_str; - } - - handle.key() = new_key_str; - (void)device_contexts_.insert(std::move(handle)); } +// Assign the value of new_key_str to the key member of the handle object +handle.key() = new_key_str; + +// Insert the handle object into the device_contexts_ container using move semantics +// The (void) is used to suppress any unused variable warnings +(void)device_contexts_.insert(std::move(handle)); + +// Definition of the function "WaitTaskFinishOnDevice" in the "DeviceContextManager" class + void DeviceContextManager::WaitTaskFinishOnDevice() const { + + // Iterate over each item in the "device_contexts_" map for (const auto &item : device_contexts_) { + + // Get the device context from the current item auto device_context = item.second; + try { + + // Check if the device context is not null and if the SyncStream function returns false if (device_context != nullptr && !device_context->SyncStream()) { + + // Print an error message using the MS_LOG macro MS_LOG(ERROR) << "SyncStream failed"; + + // Return from the function return; } } catch (const std::exception &ex) { + + // Print an error message using the MS_LOG macro, including the exception message MS_LOG(ERROR) << "SyncStream failed, exception:" << ex.what(); + + // Return from the function return; } } } + +// End of the "device" namespace } // namespace device -} // namespace mindspore + +// End of the "mindspore" namespace +} // namespace mindspore \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/pynative/op_executor.cc b/mindspore/ccsrc/runtime/pynative/op_executor.cc index a14ce57a6ac..82ae4e42409 100644 --- a/mindspore/ccsrc/runtime/pynative/op_executor.cc +++ b/mindspore/ccsrc/runtime/pynative/op_executor.cc @@ -16,173 +16,293 @@ #include "runtime/pynative/op_executor.h" +// Define the namespace "mindspore::runtime" namespace mindspore::runtime { +// Define the member function "GetInstance" of the class "OpExecutor" OpExecutor &OpExecutor::GetInstance() { + // Create a static instance of the class "OpExecutor" static OpExecutor instance; + // Return the reference to the static instance return instance; } -OpExecutor::OpExecutor() { worker_ = std::make_shared(&OpExecutor::WorkerLoop, this); } +// Constructor for the OpExecutor class +OpExecutor::OpExecutor() { + // Create a shared pointer to a new std::thread object + // Pass the address of the WorkerLoop member function of the current OpExecutor object as the thread function + // Pass 'this' pointer as the argument to the thread function + worker_ = std::make_shared(&OpExecutor::WorkerLoop, this); +} -OpExecutor::~OpExecutor() { WorkerJoin(); } +// Destructor for the OpExecutor class +OpExecutor::~OpExecutor() { + // Call the WorkerJoin() function to ensure all worker threads have completed their tasks + WorkerJoin(); +} +// Define the Register function of the OpExecutor class void OpExecutor::Register(const std::function &callback) { + // Assign the provided callback function to the batch_build_callback_ member variable batch_build_callback_ = callback; + // Set the registered_ flag to true to indicate that a callback function has been registered registered_ = true; } +// Reset the OpExecutor object void OpExecutor::Reset() { + // Clear any resources used by the OpExecutor ClearResources(); + // Set the batch_build_callback_ to nullptr, indicating that no callback function is registered batch_build_callback_ = nullptr; + // Set registered_ to false, indicating that the OpExecutor is not registered registered_ = false; +} - // There is still one task in progress + // Try block to handle exceptions try { + // Call the function WaitForRun() to wait for a task to complete WaitForRun(); } catch (const std::exception &e) { + // Catch block to handle exceptions of type std::exception + // Log an error message with the error message from the exception MS_LOG(ERROR) << "Wait failed, error message:" << e.what(); } catch (...) { + // Catch block to handle any other type of exception + // Log a generic error message MS_LOG(ERROR) << "Wait failed"; } } -void OpExecutor::ClearResources() { - MS_LOG(DEBUG) << "Start clear tasks"; - std::lock_guard lock(task_mutex_); - ClearRunOpTasks(); +// Definition of the function ClearResources() in the class OpExecutor - // Set the build task failed, and no need to run op_run_tasks. +void OpExecutor::ClearResources() { + // Log a debug message indicating the start of clearing tasks + MS_LOG(DEBUG) << "Start clear tasks"; + + // Create a lock_guard object to lock the task_mutex_ mutex + std::lock_guard lock(task_mutex_); + + // Call the ClearRunOpTasks() function to clear the run op tasks + ClearRunOpTasks(); +} + + // Loop through each build task in the op_build_tasks_ vector for (auto &build_task : op_build_tasks_) { + // Set the build_ready flag of the build task to false build_task->SetBuildReady(false); } + // Clear the op_build_tasks_ vector op_build_tasks_.clear(); + // Print a debug message indicating that the tasks have been cleared MS_LOG(DEBUG) << "End clear tasks"; } +// This function is a member function of the OpExecutor class and is called WaitForBuild void OpExecutor::WaitForBuild() { + // Check if the executing_ flag is false if (!executing_) { + // Create an instance of the ExecuteGuard class ExecuteGuard guard; + // Check if the batch_build_callback_ function pointer is not null if (batch_build_callback_ != nullptr) { + // Call the batch_build_callback_ function batch_build_callback_(); } } } +// Definition of the WaitForRun function in the OpExecutor class + void OpExecutor::WaitForRun() { + // Log a debug message indicating the start of the function MS_LOG(DEBUG) << "Start"; + + // Create a unique lock object using the task_mutex_ mutex std::unique_lock lock(task_mutex_); + + // Wait until the op_run_tasks_ queue is empty, using a lambda function as the condition task_cond_var_.wait(lock, [this]() { return op_run_tasks_.empty(); }); + + // Check for any exceptions that occurred during the execution of the tasks MsException::Instance().CheckException(); + + // Log a debug message indicating that all tasks have finished MS_LOG(DEBUG) << "All task finish"; } +// This function is a member function of the OpExecutor class void OpExecutor::Wait() { + // Call the WaitForBuild function to wait for the build process to complete WaitForBuild(); + + // Call the WaitForRun function to wait for the run process to complete WaitForRun(); } +// This function is a member function of the class OpExecutor +// It takes a shared pointer to an OpBuildTask object as a parameter void OpExecutor::PushOpBuildTask(const std::shared_ptr &op_build_task) { + // Create a lock_guard object to lock the task_mutex_ mutex std::lock_guard lock(task_mutex_); + // Add the op_build_task to the end of the op_build_tasks_ vector op_build_tasks_.push_back(op_build_task); } +// This function is a member function of the OpExecutor class +// It takes a shared pointer to an OpTask object as a parameter void OpExecutor::PushOpRunTask(const std::shared_ptr &op_run_task) { + // Create a lock_guard object to lock the task_mutex_ and automatically unlock it when it goes out of scope std::lock_guard lock(task_mutex_); + + // Push the op_run_task into the op_run_tasks_ queue op_run_tasks_.push(op_run_task); + + // Insert the name of the op_run_task's context's graph_compiler_info into the actor_in_queue_ set actor_in_queue_.insert(op_run_task->context()->graph_compiler_info()->name_); + + // Notify all threads waiting on the task_cond_var_ condition variable task_cond_var_.notify_all(); } +// A member function of the OpExecutor class that clears the op build tasks void OpExecutor::ClearOpBuildTasks() { - std::lock_guard lock(task_mutex_); + std::lock_guard lock(task_mutex_); // Lock the task mutex to ensure thread safety + + // Iterate through each op build task for (auto &task : op_build_tasks_) { - task->SetBuildReady(true); + task->SetBuildReady(true); // Set the build ready flag of the task to true } - op_build_tasks_.clear(); - MS_LOG(DEBUG) << "Clear build task"; + + op_build_tasks_.clear(); // Clear the op build tasks vector + MS_LOG(DEBUG) << "Clear build task"; // Log a debug message indicating that the build tasks have been cleared } +// Check if the build queue is empty bool OpExecutor::BuildQueueEmpty() { + // Acquire a lock on the task mutex to ensure thread safety std::lock_guard lock(task_mutex_); + + // Return true if the op_build_tasks_ queue is empty, false otherwise return op_build_tasks_.empty(); } +// A member function of the OpExecutor class that checks if the build queue is full bool OpExecutor::BuildQueueFull() { + // Create a lock_guard object to lock the task_mutex_ and automatically unlock it when it goes out of scope std::lock_guard lock(task_mutex_); + + // Check if the size of the op_build_tasks_ vector is greater than the maximum queue size return op_build_tasks_.size() > kMaxQueueSize; } +// Check if an actor is in the queue bool OpExecutor::ActorInQueue(const std::string &actor_info) { - std::lock_guard lock(task_mutex_); - auto iter = actor_in_queue_.find(actor_info); - return iter != actor_in_queue_.end(); + std::lock_guard lock(task_mutex_); // Lock the mutex to ensure thread safety + auto iter = actor_in_queue_.find(actor_info); // Find the actor in the queue + return iter != actor_in_queue_.end(); // Return true if the actor is found, false otherwise } +// A member function of the OpExecutor class that clears the run operation tasks void OpExecutor::ClearRunOpTasks() { + // Clear the actor input queue actor_in_queue_.clear(); + + // Create an empty queue of shared pointers to OpTask std::queue> empty; + // No need to worry about ExitOpTask. // ClearRunOpTasks is executed before ~OpExecutor + + // Swap the contents of op_run_tasks_ with the empty queue, effectively clearing op_run_tasks_ std::swap(op_run_tasks_, empty); } +// This is the worker loop function of the OpExecutor class void OpExecutor::WorkerLoop() { while (true) { - std::shared_ptr task; + std::shared_ptr task; // Declare a shared pointer to an OpTask object { - MS_LOG(DEBUG) << "Wait task in queue"; - std::unique_lock lock(task_mutex_); - task_cond_var_.wait(lock, [this]() { return !op_run_tasks_.empty(); }); - task = op_run_tasks_.front(); + MS_LOG(DEBUG) << "Wait task in queue"; // Log a debug message indicating that the worker is waiting for a task in the queue + std::unique_lock lock(task_mutex_); // Acquire a unique lock on the task mutex + task_cond_var_.wait(lock, [this]() { return !op_run_tasks_.empty(); }); // Wait on the task condition variable until the op_run_tasks_ queue is not empty + task = op_run_tasks_.front(); // Get the first task from the op_run_tasks_ queue } + // Log a debug message indicating that we are getting a task MS_LOG(DEBUG) << "Get task"; + // Check if the task pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(task); + // Check if the task type is an exit task if (task->task_type() == kExitTask) { + // Log a debug message indicating that the thread is exiting MS_LOG(DEBUG) << "Thread exit"; + // Return, as there is nothing else to do return; } + + // Try to run the task try { task->Run(); + // Acquire a lock on the task mutex std::unique_lock lock(task_mutex_); + // Check if the op run tasks queue is not empty if (!op_run_tasks_.empty()) { + // Remove the front task from the op run tasks queue op_run_tasks_.pop(); + // Remove the task's name from the actor in queue set actor_in_queue_.erase(task->context()->graph_compiler_info()->name_); } - if (op_run_tasks_.empty()) { - MS_LOG(DEBUG) << "Task queue empty"; - task_cond_var_.notify_all(); + if (op_run_tasks_.empty()) { // Check if the task queue is empty + MS_LOG(DEBUG) << "Task queue empty"; // Log a debug message indicating that the task queue is empty + task_cond_var_.notify_all(); // Notify all waiting threads that the task queue is empty } - } catch (const std::exception &e) { - MS_LOG(ERROR) << "Run lazy task failed, error message:" << e.what(); + } catch (const std::exception &e) { // Catch any exceptions that occur during the execution of the code within the try block + MS_LOG(ERROR) << "Run lazy task failed, error message:" << e.what(); // Log an error message indicating that running the lazy task failed, along with the error message provided by the exception { - std::unique_lock lock(task_mutex_); - ClearRunOpTasks(); - MsException::Instance().SetException(); - task_cond_var_.notify_all(); + std::unique_lock lock(task_mutex_); // Acquire a unique lock on the task mutex + ClearRunOpTasks(); // Clear the run operation tasks + MsException::Instance().SetException(); // Set the exception flag in the MsException singleton instance + task_cond_var_.notify_all(); // Notify all waiting threads that an exception has occurred } } } } +// Definition of the WorkerJoin function in the OpExecutor class + void OpExecutor::WorkerJoin() { try { - // Avoid worker thread join itself which will cause deadlock + // Check if the worker thread is joinable and if it is not the same as the current thread if (worker_->joinable() && worker_->get_id() != std::this_thread::get_id()) { { + // Acquire a lock on the task mutex to ensure thread safety std::lock_guard lock(task_mutex_); + + // Create a shared pointer to an ExitOpTask object auto task = std::make_shared(); + + // Push the task onto the op_run_tasks_ queue op_run_tasks_.push(task); + + // Notify all waiting threads that a task has been added to the queue task_cond_var_.notify_all(); + + // Log a debug message indicating that an exit task has been pushed and all threads have been notified MS_LOG(DEBUG) << "Push exit task and notify all"; } + + // Join the worker thread, blocking the current thread until the worker thread finishes execution worker_->join(); + + // Log a debug message indicating that the worker thread has finished joining MS_LOG(DEBUG) << "Worker join finish"; } } catch (const std::exception &e) { + // Log an error message if an exception is caught during the execution of the WorkerJoin function MS_LOG(ERROR) << "WorkerJoin failed: " << e.what(); } catch (...) { + // Log an error message if an unknown exception is caught during the execution of the WorkerJoin function MS_LOG(ERROR) << "WorkerJoin failed"; } } diff --git a/mindspore/ccsrc/runtime/pynative/op_runtime_info.cc b/mindspore/ccsrc/runtime/pynative/op_runtime_info.cc index 3d78cab97c8..ee4ef74b4e1 100644 --- a/mindspore/ccsrc/runtime/pynative/op_runtime_info.cc +++ b/mindspore/ccsrc/runtime/pynative/op_runtime_info.cc @@ -14,134 +14,267 @@ * limitations under the License. */ -#include "runtime/pynative/op_runtime_info.h" +// Include the header file "runtime/pynative/op_runtime_info.h" which contains declarations for runtime information related to PyNative operations. +// Include the header file for the ANF runtime algorithm in the backend common session #include "backend/common/session/anf_runtime_algorithm.h" + +// Include the header file for the ANF algorithm in the common utils #include "include/common/utils/anfalgo.h" +// Start of the `mindspore::runtime` namespace namespace mindspore::runtime { + +// Start of an anonymous namespace, used for declaring functions or variables with internal linkage namespace { + +// Function to cache execution order for a given `KernelGraphPtr` void CacheForExecutionOrder(const KernelGraphPtr &graph) { + // Check if the graph is null, throw an exception if it is MS_EXCEPTION_IF_NULL(graph); + + // Get the execution order of the nodes in the graph const auto &nodes = graph->execution_order(); + + // Iterate over each node in the execution order for (auto const &node : nodes) { + // Create empty vectors to store the output formats, types, and tensor sizes std::vector formats; std::vector types; std::vector tensor_sizes; + + // Get the number of output tensors for the current node auto output_num = common::AnfAlgo::GetOutputTensorNum(node); + + // Iterate over each output tensor of the current node for (size_t i = 0; i < output_num; ++i) { + // Get the output format, device data type, and tensor memory size for the current output tensor std::string output_format = AnfAlgo::GetOutputFormat(node, i); auto output_type = AnfAlgo::GetOutputDeviceDataType(node, i); auto tensor_size = AnfAlgo::GetOutputTensorMemSize(node, i); + + // Add the output format, type, and tensor size to their respective vectors formats.emplace_back(output_format); types.emplace_back(output_type); tensor_sizes.emplace_back(tensor_size); } - // For input + // Create a vector to store pairs of pointers to device::KernelInfo objects and their corresponding sizes std::vector> input_kernel_infos; + + // Get the number of input tensors for the given node auto input_size = common::AnfAlgo::GetInputTensorNum(node); + + // Iterate over each input tensor for (size_t i = 0; i < input_size; ++i) { + + // Get the previous node's output kernel and index for the current input tensor session::KernelWithIndex kernel_with_index = common::AnfAlgo::GetPrevNodeOutput(node, i, true); + + // Check if the kernel_with_index.first pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_with_index.first); + + // Add a pair of the kernel_info pointer and the index to the input_kernel_infos vector input_kernel_infos.emplace_back(dynamic_cast(kernel_with_index.first->kernel_info()), kernel_with_index.second); } - // For workspace and output - MS_EXCEPTION_IF_NULL(node); - auto kernel_info = dynamic_cast(node->kernel_info()); +// Check if the given node is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(node); +// Attempt to cast the kernel_info of the node to a device::KernelInfo pointer +// and assign it to the variable kernel_info +auto kernel_info = dynamic_cast(node->kernel_info()); + + // Set the user data of the node to an instance of OpRuntimeInfo + // The user data is set using the set_user_data function of the node object + // The user data is created using std::make_shared to create a shared pointer to an instance of OpRuntimeInfo + // OpRuntimeInfo is constructed using the provided arguments: formats, types, tensor_sizes, kernel_info, input_kernel_infos node->set_user_data( std::make_shared(formats, types, tensor_sizes, kernel_info, input_kernel_infos)); } } +// A function to cache information about the inputs of a given graph void CacheForGraphInputs(const KernelGraphPtr &graph) { MS_EXCEPTION_IF_NULL(graph); + + // Get the inputs of the graph const auto &inputs = graph->inputs(); + + // Iterate over each input for (const auto &input : inputs) { MS_EXCEPTION_IF_NULL(input); + + // Skip if the input is not a Parameter if (!input->isa()) { continue; } + + // Create empty vectors to store the formats, types, and sizes of the input tensors std::vector formats; std::vector types; std::vector tensor_sizes; + + // Get the number of output tensors for the input auto output_size = common::AnfAlgo::GetOutputTensorNum(input); + + // Iterate over each output tensor of the input for (size_t index = 0; index < output_size; index++) { + // Get the format, type, and size of the output tensor auto format = AnfAlgo::GetOutputFormat(input, index); auto type_id = AnfAlgo::GetOutputDeviceDataType(input, index); + + // If the type is unknown, get the inferred data type if (type_id == kTypeUnknown) { type_id = common::AnfAlgo::GetOutputInferDataType(input, index); } + auto tensor_size = AnfAlgo::GetOutputTensorMemSize(input, index); + + // Add the format, type, and size to their respective vectors formats.emplace_back(format); types.emplace_back(type_id); tensor_sizes.emplace_back(tensor_size); } + + // Continue with the rest of the code... + } +} + // Add the type_id to the end of the types vector + types.emplace_back(type_id); + + // Add the tensor_size to the end of the tensor_sizes vector + tensor_sizes.emplace_back(tensor_size); + } + + // Set the user data of the input to a new instance of OpRuntimeInfo input->set_user_data(std::make_shared( formats, types, tensor_sizes, nullptr, std::vector>())); } } } // namespace +// Define the output_format function of the OpRuntimeInfo class, which takes an index as input and returns a string std::string OpRuntimeInfo::output_format(size_t index) const { + + // Check if the index is valid, i.e., if it is within the range of the output_format_ vector if (index >= output_format_.size()) { + + // If the index is invalid, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid index:" << index << " total output_format:" << output_format_.size(); } + + // If the index is valid, return the string at the corresponding index in the output_format_ vector return output_format_[index]; } +// Return the output type at the specified index in the `output_type_` vector TypeId OpRuntimeInfo::output_type(size_t index) const { + + // Check if the index is out of bounds if (index >= output_type_.size()) { + + // If the index is invalid, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid index:" << index << " total output_type:" << output_type_.size(); } + + // Return the output type at the specified index return output_type_[index]; } +// Return the size of the output tensor at the given index in the `output_tensor_size_` vector size_t OpRuntimeInfo::output_tensor_size(size_t index) const { + + // Check if the given index is valid, i.e., within the range of the `output_tensor_size_` vector if (index >= output_tensor_size_.size()) { + + // If the index is invalid, throw an exception with an error message MS_LOG(EXCEPTION) << "Invalid index::" << index << " total output_tensor_size:" << output_tensor_size_.size(); } + + // Return the size of the output tensor at the given index return output_tensor_size_[index]; } +// Get the device address of the output tensor at the specified index device::DeviceAddressPtr OpRuntimeInfo::GetOutputDeviceAddress(size_t index) const { + + // Check if the kernel info is null, throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_info_); + + // Return the mutable output address of the kernel info at the specified index return kernel_info_->GetMutableOutputAddr(index); } +// Get the device address of the workspace for a given index in the OpRuntimeInfo object device::DeviceAddressPtr OpRuntimeInfo::GetWorkspaceDeviceAddress(size_t index) const { + + // Check if the kernel_info_ pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_info_); + + // Return the mutable workspace address for the given index from the kernel_info_ object return kernel_info_->GetMutableWorkspaceAddr(index); } +// Get the device address of the input at the specified index device::DeviceAddressPtr OpRuntimeInfo::GetInputDeviceAddress(size_t index) const { + + // Check if the index is out of range if (index >= input_kernel_infos_.size()) { + + // Log an error message indicating the index and the size of the input kernel infos MS_LOG(ERROR) << "Output range! index:" << index << " input size:" << input_kernel_infos_.size(); + + // Return a null pointer to indicate failure return nullptr; } - auto kernel_info_pair = input_kernel_infos_[index]; - MS_EXCEPTION_IF_NULL(kernel_info_pair.first); - return kernel_info_pair.first->GetMutableOutputAddr(kernel_info_pair.second); +// Assign the value of the pair at index 'index' in the 'input_kernel_infos_' vector to the variable 'kernel_info_pair' +auto kernel_info_pair = input_kernel_infos_[index]; + +// Check if the first element of the pair is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(kernel_info_pair.first); + +// Return the mutable output address of the first element of the pair, using the index stored in the second element of the pair +return kernel_info_pair.first->GetMutableOutputAddr(kernel_info_pair.second); + +// Define the member function GetInputSize() of the class OpRuntimeInfo +size_t OpRuntimeInfo::GetInputSize() const { + + // Return the size of the input_kernel_infos_ vector + return input_kernel_infos_.size(); } -size_t OpRuntimeInfo::GetInputSize() const { return input_kernel_infos_.size(); } - +// Get the size of the output for the OpRuntimeInfo object size_t OpRuntimeInfo::GetOutputSize() const { + + // Check if the kernel_info_ pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_info_); + + // Return the size of the output address list in the kernel_info_ object return kernel_info_->output_address_list().size(); } +// Get the size of the workspace required by the OpRuntimeInfo object size_t OpRuntimeInfo::GetWorkspaceSize() const { + + // Check if the kernel_info_ pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(kernel_info_); + + // Return the size of the workspace address list stored in the kernel_info_ object return kernel_info_->workspace_address_list().size(); } +// Definition of the function "CacheGraphOpRuntimeInfo" belonging to the "OpRuntimeInfo" class + void OpRuntimeInfo::CacheGraphOpRuntimeInfo(const KernelGraphPtr &graph) { + + // Call the "CacheForExecutionOrder" function to cache runtime information for the execution order of the graph CacheForExecutionOrder(graph); + + // Call the "CacheForGraphInputs" function to cache runtime information for the inputs of the graph CacheForGraphInputs(graph); } -} // namespace mindspore::runtime + +// End of the "mindspore::runtime" namespace \ No newline at end of file diff --git a/mindspore/ccsrc/runtime/pynative/run_op_helper.cc b/mindspore/ccsrc/runtime/pynative/run_op_helper.cc index b90872c5a82..2f6eb7b2c81 100644 --- a/mindspore/ccsrc/runtime/pynative/run_op_helper.cc +++ b/mindspore/ccsrc/runtime/pynative/run_op_helper.cc @@ -14,276 +14,538 @@ * limitations under the License. */ -#include "runtime/pynative/run_op_helper.h" +// Include the header file "runtime/pynative/run_op_helper.h" which contains helper functions for running operations in PyNative mode. +// Include the necessary headers for string manipulation, vector operations, memory management, and algorithm functions #include #include #include #include + +// Include the custom log adapter header file #include "utils/log_adapter.h" + +// Include the custom ANF runtime algorithm header file #include "backend/common/session/anf_runtime_algorithm.h" + +// Include the custom convert utils header file #include "include/common/utils/convert_utils.h" + +// Include the custom MS device shape transfer header file #include "runtime/device/ms_device_shape_transfer.h" + +// Include the custom op runtime info header file #include "runtime/pynative/op_runtime_info.h" + +// Include the custom op executor header file #include "runtime/pynative/op_executor.h" + +// Include the custom actor common header file #include "runtime/graph_scheduler/actor/actor_common.h" namespace mindspore::runtime { + namespace { -// 1. Device type is different in heterogeneous scenes. -// 2. The device address format is different. + +// Function to update input tensors from device +// 1. Device type may be different in heterogeneous scenes. +// 2. The device address format may be different. void UpdateInputTensorFromDevice(const std::vector &input_nodes, const std::vector &input_tensors, const device::DeviceContext *device_context) { + // Log a debug message to indicate the start of the function MS_LOG(DEBUG) << "Start"; + + // Get the size of the input nodes vector auto input_size = input_nodes.size(); + + // Iterate over each input node and its corresponding tensor for (size_t i = 0; i < input_size; ++i) { auto &tensor = input_tensors[i]; auto &input_node = input_nodes[i]; + + // Check if the tensor is null MS_EXCEPTION_IF_NULL(tensor); - auto tensor_address = std::dynamic_pointer_cast(tensor->device_address()); + + // Get the device address of the tensor + auto tensor_address = std::dynamic_pointer_cast(tensor->device_address); + + // Get the mutable output address of the input node auto node_address = AnfAlgo::GetMutableOutputAddr(input_node, 0); - // node_address can't be null + + // Check if the node address is null MS_EXCEPTION_IF_NULL(node_address); + + // Check if the tensor address is not null if (tensor_address != nullptr) { + // Check if the device type or format of the tensor address is different from the device context if (tensor_address->DeviceType() != device_context->GetDeviceAddressType() || tensor_address->format() != node_address->format()) { - // Need wait for OpExecutor task finish - tensor->data_sync(); - // If tensor address is null, we will set Parameter address to the Tensor. - tensor->set_device_address(nullptr); + // ... perform some action } } } - MS_LOG(DEBUG) << "End"; } -void UpdateParameterShapeFromInputTensor(const AnfNodePtr &input_node, const tensor::TensorPtr &input_tensor) { - MS_EXCEPTION_IF_NULL(input_node); - if (input_tensor == nullptr || !input_node->isa()) { - return; - } +} // namespace +} // namespace mindspore::runtime +// Need to wait for the OpExecutor task to finish before proceeding +tensor->data_sync(); + +// If the tensor address is null, we will set the Parameter address to the Tensor +tensor->set_device_address(nullptr); + +// End of the function +MS_LOG(DEBUG) << "End"; + +// A function to update the shape of a parameter based on an input tensor + +// Check if the input node is null, and throw an exception if it is +MS_EXCEPTION_IF_NULL(input_node); + +// Check if the input tensor is null or if the input node is not of type Parameter +// If either condition is true, return from the function +if (input_tensor == nullptr || !input_node->isa()) { + return; +} + + // Cast the input_node to a ParameterPtr and assign it to the variable input_param auto input_param = input_node->cast(); + + // Check if input_param is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(input_param); + + // Check if input_param has dynamic shape if (!input_param->has_dynamic_shape()) { + // If it does not have dynamic shape, return and exit the function return; } + // Get the shape of the input tensor auto shape = input_tensor->shape(); - std::vector update_shape; - std::transform(shape.begin(), shape.end(), std::back_inserter(update_shape), IntToSize); - MS_LOG(DEBUG) << "Update input node shape to:" << update_shape; - common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(input_node, 0)}, {update_shape}, - input_node.get()); -} + // Create an empty vector to store the updated shape + std::vector update_shape; + + // Use std::transform to iterate over each element in the shape vector and apply the IntToSize function to convert each element to size_t + std::transform(shape.begin(), shape.end(), std::back_inserter(update_shape), IntToSize); + + // Print the updated shape to the debug log + MS_LOG(DEBUG) << "Update input node shape to:" << update_shape; + + // Set the output inference type and shape for the input node using the updated shape + common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(input_node, 0)}, {update_shape}, input_node.get()); + +// This function updates the device address of input nodes with the corresponding device address from input tensors void UpdateInputNodeDeviceAddress(const std::vector &input_nodes, const std::vector &input_tensors, const device::DeviceContext *device_context) { + // Log a debug message to indicate the start of the function MS_LOG(DEBUG) << "Start"; + + // Get the size of the input nodes and input tensors auto input_size = input_nodes.size(); auto tensor_size = input_tensors.size(); + + // Check if the sizes of input nodes and input tensors are equal if (input_size != tensor_size) { + // If not equal, throw an exception with an error message MS_LOG(EXCEPTION) << "input node size:" << input_size << " not equal to tensors size:" << tensor_size; } + + // Iterate over each input node and input tensor for (size_t i = 0; i < input_size; ++i) { + // Get a reference to the current input node and input tensor auto &input_node = input_nodes[i]; auto &input_tensor = input_tensors[i]; + + // Check if the input tensor is null MS_EXCEPTION_IF_NULL(input_tensor); - auto tensor_address = std::dynamic_pointer_cast(input_tensor->device_address()); + + // Get the device address of the input tensor + auto tensor_address = std::dynamic_pointer_cast(input_tensor->device_address); + + // Get the mutable output address of the input node at index 0 auto node_address = AnfAlgo::GetMutableOutputAddr(input_node, 0); + // ... + } +} - UpdateParameterShapeFromInputTensor(input_node, input_tensor); +// Call the function UpdateParameterShapeFromInputTensor with the arguments input_node and input_tensor - MS_EXCEPTION_IF_NULL(node_address); - if (tensor_address == nullptr) { - input_tensor->set_device_address(node_address); - input_tensor->set_sync_status(kNeedSyncHostToDeviceImmediately); - input_tensor->set_lazy_callback([]() { runtime::OpExecutor::GetInstance().Wait(); }); - node_address->set_from_persistent_mem(input_tensor->is_parameter()); - node_address->SetNodeIndex(input_node, 0); - UpdateRefCount(node_address.get(), true); - } +// Check if the pointer `node_address` is null, and throw an exception if it is null +MS_EXCEPTION_IF_NULL(node_address); - // The DeviceType and format of DeviceAddress is always the same after UpdateInputTensor +// Check if the pointer `tensor_address` is null +if (tensor_address == nullptr) { + // Set the device address of the input tensor to `node_address` + input_tensor->set_device_address(node_address); + + // Set the sync status of the input tensor to `kNeedSyncHostToDeviceImmediately` + input_tensor->set_sync_status(kNeedSyncHostToDeviceImmediately); + + // Set a lazy callback function for the input tensor + input_tensor->set_lazy_callback([]() { runtime::OpExecutor::GetInstance().Wait(); }); + + // Set the `from_persistent_mem` flag of `node_address` based on whether the input tensor is a parameter + node_address->set_from_persistent_mem(input_tensor->is_parameter()); + + // Set the node index of `node_address` to `input_node` with an index of 0 + node_address->SetNodeIndex(input_node, 0); + + // Update the reference count of `node_address` by incrementing it + UpdateRefCount(node_address.get(), true); +} + + // Check if the tensor address is not null and is different from the node address if (tensor_address != nullptr && tensor_address != node_address) { + + // Set the output address of the tensor to the input node AnfAlgo::SetOutputAddr(tensor_address, 0, input_node.get()); } } + + // Print "End" for debugging purposes MS_LOG(DEBUG) << "End"; } +// Function to update the reference node's output device address in a given graph void UpdateRefNodeOutputDeviceAddress(const KernelGraphPtr &graph) { MS_EXCEPTION_IF_NULL(graph); + + // Get the reference node map from the graph auto ref_node_map = graph->GetRefMap(); + + // Iterate over each entry in the reference node map for (const auto &iter : ref_node_map) { + // Get the output pair (reference node and output index) auto &output_pair = iter.first; - auto &input_pair = iter.second; auto &ref_node = output_pair.first; auto output_index = output_pair.second; + + // Get the input pair (input node and input node output index) + auto &input_pair = iter.second; auto &input_node = input_pair.first; auto input_node_output_index = input_pair.second; + // Rest of the code is not provided, so we cannot provide further comments + // ... + // ... + } +} + + // Get the mutable output address of the input node using the AnfAlgo::GetMutableOutputAddr function auto input_addr = AnfAlgo::GetMutableOutputAddr(input_node, input_node_output_index, false); + + // Get the mutable output address of the reference node using the AnfAlgo::GetMutableOutputAddr function auto ref_node_output_addr = AnfAlgo::GetMutableOutputAddr(ref_node, output_index, false); + + // Check if the input address is different from the reference node's output address if (input_addr != ref_node_output_addr) { + + // If the addresses are different, set the output address of the input node to the reference node's output address AnfAlgo::SetOutputAddr(input_addr, output_index, ref_node.get()); } } } +// Function to copy tensor data to a device void CopyTensorDataToDevice(const tensor::TensorPtr &tensor, const AnfNodePtr &node, const device::DeviceContext *device_context) { + // Check if the tensor and device context are not null MS_EXCEPTION_IF_NULL(tensor); MS_EXCEPTION_IF_NULL(device_context); + + // Get the device address of the tensor auto device_address = std::dynamic_pointer_cast(tensor->device_address()); MS_EXCEPTION_IF_NULL(device_address); + + // Check if the device address pointer is null and if memory allocation fails if ((device_address->GetPtr() == nullptr) && (!device_context->AllocateMemory(device_address.get(), device_address->GetSize()))) { + // Throw an exception with an error message MS_LOG(EXCEPTION) << "Allocate memory failed"; } +} // Copy data from host tensor to device. + + // Get the size of the tensor in bytes auto tensor_size = LongToSize(tensor->data().nbytes()); + + // Get the data type of the tensor auto tensor_type = tensor->data_type(); + + // Log a debug message indicating the node being copied to the device MS_LOG(DEBUG) << "Copy to device, node:" << node->DebugString(); + + // Synchronize the host tensor data to the device + // using the runtime padding shape, tensor size, tensor type, + // tensor data pointer, and host format if (!device_address->SyncHostToDevice(trans::GetRuntimePaddingShape(node, 0), tensor_size, tensor_type, tensor->data_c(), tensor->device_info().host_format_)) { + // If the synchronization fails, throw an exception MS_LOG(EXCEPTION) << "SyncHostToDevice failed"; } } -void CopyValueNodeTensorToDevice(const ValueNodePtr &node, const device::DeviceContext *device_context) { - MS_EXCEPTION_IF_NULL(node); - MS_EXCEPTION_IF_NULL(device_context); +// Function to copy the value of a node's tensor to a device using the provided device context - auto &node_value = node->value(); - MS_EXCEPTION_IF_NULL(node_value); +// Check if the node pointer is null, throw an exception if it is +MS_EXCEPTION_IF_NULL(node); - std::vector tensors; - TensorValueToTensor(node_value, &tensors); - for (size_t i = 0; i < tensors.size(); i++) { +// Check if the device context pointer is null, throw an exception if it is +MS_EXCEPTION_IF_NULL(device_context); + +// Create a reference variable named "node_value" that refers to the value of the "node" object +auto &node_value = node->value(); + +// Check if the "node_value" is null, and if it is, throw an exception +MS_EXCEPTION_IF_NULL(node_value); + +// Declare a vector of tensor pointers named "tensors" +std::vector tensors; + +// Convert the tensor values in "node_value" to actual tensors and store them in the "tensors" vector +TensorValueToTensor(node_value, &tensors); + +// Iterate over the elements in the "tensors" vector using a for loop +for (size_t i = 0; i < tensors.size(); i++) { + + // Get a reference to the current tensor in the "tensors" vector const auto &tensor = tensors[i]; - MS_EXCEPTION_IF_NULL(tensor); + // Check if the tensor is null, and throw an exception if it is + MS_EXCEPTION_IF_NULL(tensor); +} + + // Get the mutable output address of the node at index i using AnfAlgo::GetMutableOutputAddr const auto &node_address = AnfAlgo::GetMutableOutputAddr(node, i, false); + + // Check if the node address is null MS_EXCEPTION_IF_NULL(node_address); + + // If the pointer of the node address is not null, return if (node_address->GetPtr() != nullptr) { return; } + + // Set the device address of the tensor to the node address tensor->set_device_address(node_address); + + // Update the reference count of the node address UpdateRefCount(node_address.get(), true); + + // Copy the tensor data to the device using CopyTensorDataToDevice CopyTensorDataToDevice(tensor, node, device_context); } } +// Function to copy the value node string to the device void CopyValueNodeStringToDevice(const ValueNodePtr &node, const device::DeviceContext *device_context) { + // Check if the value node pointer is null MS_EXCEPTION_IF_NULL(node); + + // Check if the device context pointer is null MS_EXCEPTION_IF_NULL(device_context); + + // Get the mutable output address of the node at index 0 const auto &node_address = AnfAlgo::GetMutableOutputAddr(node, 0, false); + + // Check if the node address is null MS_EXCEPTION_IF_NULL(node_address); + + // Check if the pointer in the node address is not null if (node_address->GetPtr() != nullptr) { + // If the pointer is not null, return from the function return; } +} - if (!device_context->AllocateMemory(node_address.get(), node_address->GetSize())) { +// Check if memory allocation for the device context is successful +if (!device_context->AllocateMemory(node_address.get(), node_address->GetSize())) { + + // If memory allocation fails, throw an exception with the error message "Allocate memory failed" MS_LOG(EXCEPTION) << "Allocate memory failed"; - } +} + // Create a reference to the value of the node auto &node_value = node->value(); + + // Check if the node value is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(node_value); - // Copy data to device. + + // Get the value of the node as a string auto value = GetValue(node_value); + + // Get the size of the string value size_t tensor_size = value.size(); + + // Create a shape vector with dimensions {1, tensor_size} ShapeVector shape = {1, SizeToLong(tensor_size)}; + + // Copy the data from the host to the device + // using the SyncHostToDevice function of the node_address object + // If the copy fails, throw an exception if (!node_address->SyncHostToDevice(shape, tensor_size, kNumberTypeUInt8, value.data())) { MS_LOG(EXCEPTION) << "SyncHostToDevice failed"; } } +// Function to copy value node data to the device void CopyValueNodeDataToDevice(const KernelGraphPtr &graph, const device::DeviceContext *device_context) { MS_EXCEPTION_IF_NULL(graph); + + // Log a debug message indicating the start of the function MS_LOG(DEBUG) << "Start"; + + // Get the value nodes from the graph const auto &value_nodes = graph->graph_value_nodes(); + + // Iterate over each value node for (const auto &value_node : value_nodes) { MS_EXCEPTION_IF_NULL(value_node); + + // Get the value of the value node auto &node_value = value_node->value(); MS_EXCEPTION_IF_NULL(node_value); + + // Check the type of the value node if (node_value->isa() || node_value->isa()) { + // If the value node is a tensor or a value tuple, copy it to the device CopyValueNodeTensorToDevice(value_node, device_context); } else if (node_value->isa()) { + // If the value node is a string, copy it to the device CopyValueNodeStringToDevice(value_node, device_context); } else { + // If the value node is of an unknown type, log a warning message MS_LOG(WARNING) << "Unknown value node type:" << value_node->DebugString(); } } + + // Log a debug message indicating the end of the function MS_LOG(DEBUG) << "End"; } +// Function to copy parameter data to a device void CopyParameterDataToDevice(const std::vector &input_nodes, const std::vector &input_tensors, const device::DeviceContext *device_context) { + // Log a debug message indicating the start of the function MS_LOG(DEBUG) << "Start"; + + // Get the size of the input nodes vector auto input_size = input_nodes.size(); + + // Iterate over each input node for (size_t i = 0; i < input_size; ++i) { + // Check if the input tensor is null MS_EXCEPTION_IF_NULL(input_tensors[i]); + + // Check if the input tensor needs to be synchronized from host to device immediately if (input_tensors[i]->NeedSyncHostToDeviceImmediately()) { + // Copy the tensor data to the device CopyTensorDataToDevice(input_tensors[i], input_nodes[i], device_context); + + // Set the sync status of the input tensor to indicate that synchronization is not needed input_tensors[i]->set_sync_status(kNoNeedSync); } } + + // Log a debug message indicating the end of the function MS_LOG(DEBUG) << "End"; } +// This function updates the size of the output device addresses for a given node in the computation graph. +// It takes in two parameters: +// - node: a reference to the AnfNodePtr object representing the node in the computation graph +// - runtime_info: a shared pointer to the OpRuntimeInfo object containing runtime information for the node + void UpdateOutputAddrSize(const AnfNodePtr &node, const std::shared_ptr &runtime_info) { + // Check if the runtime_info pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(runtime_info); + + // Get the number of output tensors for the node from the runtime_info object auto output_size = runtime_info->GetOutputSize(); + + // Iterate over each output tensor for (size_t i = 0; i < output_size; ++i) { + // Get the device address of the output tensor from the runtime_info object auto output_address = runtime_info->GetOutputDeviceAddress(i); + + // Check if the output_address pointer is null, and throw an exception if it is MS_EXCEPTION_IF_NULL(output_address); + + // Get the memory size of the output tensor from the AnfAlgo utility function auto output_addr_size = AnfAlgo::GetOutputTensorMemSize(node, i); + + // Check if the output_addr_size is different from the current size of the output_address if (output_addr_size != output_address->GetSize()) { + // If the sizes are different, update the size of the output_address to match the output_addr_size output_address->SetSize(output_addr_size); } } } +// Function to allocate memory for kernel input bool MallocForKernelInput(const std::shared_ptr &runtime_info, const device::DeviceContext *device_context) { + // Check if runtime_info and device_context are not null MS_EXCEPTION_IF_NULL(runtime_info); MS_EXCEPTION_IF_NULL(device_context); + + // Get the size of the input auto input_size = runtime_info->GetInputSize(); + + // Iterate over each input for (size_t i = 0; i < input_size; ++i) { + // Get the device address of the input auto input_address = runtime_info->GetInputDeviceAddress(i); MS_EXCEPTION_IF_NULL(input_address); + + // Check if the input address is null and allocate memory if needed if (input_address->GetPtr() == nullptr && !device_context->AllocateMemory(input_address.get(), input_address->GetSize())) { return false; } } + + // Return true to indicate successful memory allocation for all inputs return true; } +// Function to allocate memory for kernel output bool MallocForKernelOutput(const std::shared_ptr &runtime_info, const AnfNodePtr &node, const device::DeviceContext *device_context) { + // Check if the runtime info is null MS_EXCEPTION_IF_NULL(runtime_info); + + // Check if the node is null MS_EXCEPTION_IF_NULL(node); + + // Check if the device context is null MS_EXCEPTION_IF_NULL(device_context); - auto kernel_mod = AnfAlgo::GetKernelMod(node); - MS_EXCEPTION_IF_NULL(kernel_mod); - auto output_size = runtime_info->GetOutputSize(); - auto kernel_out_size_list = kernel_mod->GetOutputSizeList(); + auto kernel_mod = AnfAlgo::GetKernelMod(node); // Get the kernel module associated with the node + MS_EXCEPTION_IF_NULL(kernel_mod); // Throw an exception if the kernel module is null + + auto output_size = runtime_info->GetOutputSize(); // Get the number of outputs of the node + auto kernel_out_size_list = kernel_mod->GetOutputSizeList(); // Get the list of output sizes from the kernel module + + // Check if the number of outputs from the kernel module matches the expected output size if (kernel_out_size_list.size() != output_size) { MS_LOG(ERROR) << "Node " << node->fullname_with_scope() << " output num is:" << output_size << " but kernel_mod output num:" << kernel_out_size_list.size(); return false; } + + // Iterate over each output of the node for (size_t i = 0; i < output_size; ++i) { - auto device_address = runtime_info->GetOutputDeviceAddress(i); - MS_EXCEPTION_IF_NULL(device_address); - // For example, we need to call cudnnGetRNNTrainingReserveSize to get real output size in LstmGpuKernelMod! + auto device_address = runtime_info->GetOutputDeviceAddress(i); // Get the device address of the output + MS_EXCEPTION_IF_NULL(device_address); // Throw an exception if the device address is null + + // Check if the output size from the kernel module matches the device address size if (kernel_out_size_list[i] != device_address->GetSize()) { // If the format of the DeviceAddress is different, then the size is originally different. // Such as NCHW(1,1,1,3) and NC1HWC0(1,1,1,1,16). So we don't need to update the size. @@ -291,7 +553,7 @@ bool MallocForKernelOutput(const std::shared_ptr &runtime_info, c if (device_address->GetPtr() != nullptr) { MS_LOG(ERROR) << "kernel mod output " << i << " size:" << kernel_out_size_list[i] << " not equal to device_address size:" << device_address->GetSize() - << ", but the device address is already have ptr"; + << ", but the device address already has a pointer"; return false; } device_address->SetSize(kernel_out_size_list[i]); @@ -306,154 +568,294 @@ bool MallocForKernelOutput(const std::shared_ptr &runtime_info, c return true; } +// Function to allocate memory for kernel workspace bool MallocForKernelWorkspace(const std::shared_ptr &runtime_info, const device::DeviceContext *device_context) { + // Check if runtime_info and device_context are not null MS_EXCEPTION_IF_NULL(runtime_info); MS_EXCEPTION_IF_NULL(device_context); + + // Get the size of the workspace from runtime_info auto workspace_size = runtime_info->GetWorkspaceSize(); + + // Loop through each workspace element for (size_t i = 0; i < workspace_size; ++i) { + // Get the device address for the workspace element auto device_address = runtime_info->GetWorkspaceDeviceAddress(i); MS_EXCEPTION_IF_NULL(device_address); + + // Check if the device address pointer is null and if memory allocation fails if (device_address->GetPtr() == nullptr && !device_context->AllocateMemory(device_address.get(), device_address->GetSize())) { + // Log an error message and return false if memory allocation fails MS_LOG(ERROR) << "Allocate workspace memory failed"; return false; } } + + // Return true to indicate successful memory allocation for all workspace elements return true; } +// Function to create a list of kernel input addresses based on the given runtime information kernel::AddressPtrList CreateKernelInputAddress(const std::shared_ptr &runtime_info) { + // Check if the runtime_info pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(runtime_info); + + // Get the size of the input from the runtime_info auto input_size = runtime_info->GetInputSize(); + + // Create an empty list to store the kernel input addresses kernel::AddressPtrList inputs; + + // Iterate over the input size for (size_t i = 0; i < input_size; ++i) { + // Get the device address for the input at index i from the runtime_info auto device_address = runtime_info->GetInputDeviceAddress(i); + + // Throw an exception if the device_address pointer is null MS_EXCEPTION_IF_NULL(device_address); + + // Create a new kernel::Address object using the device address's mutable pointer and size inputs.emplace_back(std::make_shared(device_address->GetMutablePtr(), device_address->GetSize())); + + // Log the address and size of the current input MS_LOG(DEBUG) << "input[" << i << "]:" << inputs.back()->addr << " size:" << inputs.back()->size; } + + // Return the list of kernel input addresses return inputs; } +// Function to create a list of kernel workspace addresses kernel::AddressPtrList CreateKernelWorkspaceAddress(const std::shared_ptr &runtime_info) { + // Check if the runtime_info pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(runtime_info); + + // Get the size of the workspace from the runtime_info object auto workspace_size = runtime_info->GetWorkspaceSize(); + + // Create an empty list to store the workspace addresses kernel::AddressPtrList workspaces; + + // Iterate over the workspace size for (size_t i = 0; i < workspace_size; ++i) { + // Get the device address for the workspace at index i auto device_address = runtime_info->GetWorkspaceDeviceAddress(i); + + // Throw an exception if the device address is null MS_EXCEPTION_IF_NULL(device_address); + + // Create a new kernel::Address object using the device address's mutable pointer and size workspaces.emplace_back( std::make_shared(device_address->GetMutablePtr(), device_address->GetSize())); + + // Log the workspace address and size for debugging purposes MS_LOG(DEBUG) << "workspace[" << i << "]:" << workspaces.back()->addr << " size:" << workspaces.back()->size; } + + // Return the list of workspace addresses return workspaces; } +// Function to create a list of kernel output addresses based on the provided runtime information kernel::AddressPtrList CreateKernelOutputAddress(const std::shared_ptr &runtime_info) { + + // Get the size of the output from the runtime information auto output_size = runtime_info->GetOutputSize(); + + // Create an empty list to store the output addresses kernel::AddressPtrList outputs; + + // Iterate over each output for (size_t i = 0; i < output_size; ++i) { + + // Get the device address for the current output auto device_address = runtime_info->GetOutputDeviceAddress(i); + + // Create a new kernel address using the device address's mutable pointer and size outputs.emplace_back(std::make_shared(device_address->GetMutablePtr(), device_address->GetSize())); + + // Log the address and size of the current output MS_LOG(DEBUG) << "output[" << i << "]:" << outputs.back()->addr << " size:" << outputs.back()->size; } + + // Return the list of output addresses return outputs; } -// Host to Device or Device to Host +// Function to copy data from host to device or device to host void CopyDataToDevice(const KernelGraphPtr &graph, const std::vector &input_tensors, const device::DeviceContext *device_context) { MS_EXCEPTION_IF_NULL(graph); + + // Copy data of value nodes from host to device CopyValueNodeDataToDevice(graph, device_context); + + // Copy data of input parameters from host to device CopyParameterDataToDevice(graph->input_nodes(), input_tensors, device_context); } -// kernel_mode launch +// Function to launch kernels in kernel mode void LaunchKernels(const KernelGraphPtr &graph, const device::DeviceContext *device_context) { MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(device_context); + + // Log a debug message indicating the start of kernel launch MS_LOG(DEBUG) << "Start"; +} - // Get device address from OpRuntimeInfo - const auto &execution_order = graph->execution_order(); - for (auto const &node : execution_order) { +// Get the execution order of nodes in the graph +const auto &execution_order = graph->execution_order(); + +// Iterate over each node in the execution order +for (auto const &node : execution_order) { MS_EXCEPTION_IF_NULL(node); + + // Check if the node has dynamic shape auto is_dynamic_shape = common::AnfAlgo::IsDynamicShape(node); + + // Get the runtime information associated with the node auto runtime_info = node->user_data(); + + // Throw an exception if the runtime information is null MS_EXCEPTION_IF_NULL(runtime_info); + // Check if memory allocation for kernel input failed if (!MallocForKernelInput(runtime_info, device_context)) { - MS_LOG(EXCEPTION) << "Malloc for kernel input failed, Memory isn't enough, node:" << node->fullname_with_scope(); + // If allocation failed, throw an exception with an error message including the name of the node + MS_LOG(EXCEPTION) << "Malloc for kernel input failed, Memory isn't enough, node:" << node->fullname_with_scope(); } + + // Create kernel input addresses auto inputs = CreateKernelInputAddress(runtime_info); + // Check if the shape is dynamic if (is_dynamic_shape) { + + // If the shape is dynamic, call the UpdateDynamicShape function of the device context device_context->UpdateDynamicShape(node); } + // Check if memory allocation for kernel workspace was successful if (!MallocForKernelWorkspace(runtime_info, device_context)) { - MS_LOG(EXCEPTION) << "Malloc for kernel workspace failed, Memory isn't enough, node:" - << node->fullname_with_scope(); + // If memory allocation failed, throw an exception with an error message + MS_LOG(EXCEPTION) << "Malloc for kernel workspace failed, Memory isn't enough, node:" + << node->fullname_with_scope(); } + + // Create kernel workspace addresses auto workspaces = CreateKernelWorkspaceAddress(runtime_info); + // Check if memory allocation for kernel output failed if (!MallocForKernelOutput(runtime_info, node, device_context)) { - MS_LOG(EXCEPTION) << "Malloc for kernel output failed, Memory isn't enough, node:" << node->fullname_with_scope(); - } - auto outputs = CreateKernelOutputAddress(runtime_info); - if (!device_context->LaunchKernel(node, inputs, workspaces, outputs, is_dynamic_shape)) { - MS_LOG(EXCEPTION) << "Launch kernel failed, name:" << node->fullname_with_scope(); + // If allocation failed, throw an exception with an error message + MS_LOG(EXCEPTION) << "Malloc for kernel output failed, Memory isn't enough, node:" << node->fullname_with_scope(); } - if (is_dynamic_shape) { - UpdateOutputAddrSize(node, runtime_info); + // Create kernel output addresses + auto outputs = CreateKernelOutputAddress(runtime_info); + + // Launch the kernel with the given inputs, workspaces, outputs, and dynamic shape flag + if (!device_context->LaunchKernel(node, inputs, workspaces, outputs, is_dynamic_shape)) { + // If kernel launch failed, throw an exception with an error message + MS_LOG(EXCEPTION) << "Launch kernel failed, name:" << node->fullname_with_scope(); } - } - MS_LOG(DEBUG) << "End"; + + // Check if the shape is dynamic + if (is_dynamic_shape) { + // If the shape is dynamic, update the output address size + UpdateOutputAddrSize(node, runtime_info); + } + + // Log a debug message indicating the end of the function + MS_LOG(DEBUG) << "End"; } +// A function to wait for the completion of communication for a vector of input tensors void WaitCommunicationFinish(const std::vector &input_tensors) { + + // Iterate over each input tensor in the vector for (auto &input_tensor : input_tensors) { + + // Check if the input tensor is null MS_EXCEPTION_IF_NULL(input_tensor); + + // Check if the input tensor needs to wait for the device if (input_tensor->NeedWaitDevice()) { + + // Wait for the device associated with the input tensor to finish communication input_tensor->WaitDevice(); } } } +// A function to release kernel resources for a given KernelGraphPtr void ReleaseKernelResource(const KernelGraphPtr &graph) { + // Check if the graph pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(graph); + + // Get the execution order of the kernels in the graph const auto &kernels = graph->execution_order(); + + // Iterate over each kernel in the execution order for (const auto &kernel : kernels) { + // Check if the kernel pointer is null, throw an exception if it is MS_EXCEPTION_IF_NULL(kernel); + + // Check if the name of the kernel is in the OpCacheBlackList if (kOpCacheBlackList.find(common::AnfAlgo::GetCNodeName(kernel)) != kOpCacheBlackList.end()) { + // Get the kernel module associated with the kernel auto kernel_mod = AnfAlgo::GetKernelMod(kernel); + + // If the kernel module exists, release its resources if (kernel_mod) { kernel_mod->ReleaseResource(); } } } } -} // namespace -// Determine the address of the graph and do not change the address in subsequent executions +// End of the namespace + +// Function to update the device address of a given graph and its tensors void UpdateDeviceAddress(const KernelGraphPtr &graph, const std::vector &tensors_without_value_mask, const device::DeviceContext *device_context) { MS_EXCEPTION_IF_NULL(graph); + + // Log a debug message indicating the start of the function MS_LOG(DEBUG) << "Start"; + + // Get the input nodes of the graph const auto &input_nodes = graph->input_nodes(); + + // Update the device address of the input tensors from the device UpdateInputTensorFromDevice(input_nodes, tensors_without_value_mask, device_context); + + // Update the device address of the input nodes UpdateInputNodeDeviceAddress(input_nodes, tensors_without_value_mask, device_context); + + // Update the device address of the output tensors of reference nodes UpdateRefNodeOutputDeviceAddress(graph); + + // Log a debug message indicating the end of the function MS_LOG(DEBUG) << "End"; } +// A function to run a single operation graph void RunSingleOpGraph(const KernelGraphPtr &graph, const std::vector &input_tensors, const device::DeviceContext *device_context) { - WaitCommunicationFinish(input_tensors); - CopyDataToDevice(graph, input_tensors, device_context); - LaunchKernels(graph, device_context); - ReleaseKernelResource(graph); + // Wait for any ongoing communication to finish before proceeding + WaitCommunicationFinish(input_tensors); + + // Copy the input data to the device + CopyDataToDevice(graph, input_tensors, device_context); + + // Launch the kernels to perform the computation + LaunchKernels(graph, device_context); + + // Release the resources used by the kernels + ReleaseKernelResource(graph); } -} // namespace mindspore::runtime + +// End of the namespace mindspore::runtime \ No newline at end of file