mindspore 2.6.0rc1__cp310-none-any.whl → 2.7.0rc1__cp310-none-any.whl

mindspore/include/mindspore/ccsrc/include/backend/anf_runtime_algorithm.h

@@ -71,7 +71,7 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   // get input format select of anf node
   static std::string GetInputFormat(const AnfNodePtr &node, size_t input_idx);
   // Judge whether the format is equivalent by converting between default format and real format.
-  static bool IsEquivalentFormat(const std::string &src_format, const std::string &dst_format);
+  static bool IsEquivalentFormat(const Format &src_format, const Format &dst_format);
   // get output format from prev node,input_index is the input index of current node related to prev node
   static std::string GetPrevNodeOutputFormat(const AnfNodePtr &anf_node, size_t input_idx);
   // get reshape_type of from the output of input node.
@@ -124,16 +124,17 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   // Get shape, devie type and value information.
   static std::tuple<abstract::BaseShapePtr, TypePtr, ValuePtr> GetAbstractInfo(const AnfNodePtr &node,
                                                                                size_t output_idx);
-
   static bool ExistOutputKernelTensor(const AnfNodePtr &node, size_t output_idx);
 
   // Get output kernel tensor if exists, otherwise throw a exception.
-  static const KernelTensorPtr &GetOutputKernelTensor(const AnfNodePtr &node, size_t output_idx);
+  static const KernelTensorPtr &GetOutputKernelTensor(const AnfNodePtr &node, size_t output_idx,
+                                                      bool skip_nop_node = true);
   // Get output kernel tensor if exists, otherwise create a new one and set into node.
   static const KernelTensorPtr &GetOrCreateOutputKernelTensor(const AnfNodePtr &node, size_t output_idx);
 
   // Get input kernel tensor if exists, otherwise throw a exception.
-  static const KernelTensorPtr &GetPrevNodeOutputKernelTensor(const AnfNodePtr &node, size_t input_idx);
+  static const KernelTensorPtr &GetPrevNodeOutputKernelTensor(const AnfNodePtr &node, size_t input_idx,
+                                                              bool skip_nop_node = true);
   // Get input kernel tensor if exists, otherwise create a new one and set into node.
   static const KernelTensorPtr &GetOrCreatePrevNodeOutputKernelTensor(const AnfNodePtr &node, size_t input_idx);
 
@@ -157,23 +158,26 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
                                        std::vector<kernel::KernelAttr> *selected_kernel_attrs);
   // Create output kernel tensor for node using node's shape, type and value,
   // and set device information to kernel tensor.
-  static KernelTensorPtr CreateOutputKernelTensorWithDeviceInfo(const AnfWithOutIndex &node_with_index,
-                                                                void *const device_ptr, size_t size,
-                                                                const string &format, TypeId dtype_id,
-                                                                const ShapeVector &host_shape,
-                                                                const std::string &device_name, uint32_t device_id,
-                                                                const UserDataPtr &user_data = nullptr);
+  static KernelTensorPtr CreateOutputKernelTensorWithDeviceInfo(
+    const AnfWithOutIndex &node_with_index, void *const device_ptr, size_t size, const string &format, TypeId dtype_id,
+    const ShapeVector &host_shape, const std::string &device_name, uint32_t device_id,
+    const UserDataPtr &user_data = nullptr, uint32_t stream_id = 0);
 
   // Get all input memory size list for node.
   static std::vector<size_t> GetNodeInputSizeList(const AnfNodePtr &node);
 
   static size_t GetOutputAddressNum(const AnfNodePtr &node);
   // set output device addr of anf_node
-  static void SetOutputAddr(const DeviceAddressPtr &addr, size_t output_idx, AnfNode *node);
+  static void SetOutputAddr(const DeviceAddressPtr &addr, size_t output_idx, const AnfNodePtr &node);
+  // set output kernel tensor of anf node
+  static void SetOutputKernelTensor(const KernelTensorPtr &kernel_tensor, size_t output_idx, AnfNode *node);
   // set workspace device addr of anf_node
-  static void SetWorkspaceAddr(const DeviceAddressPtr &addr, size_t output_idx, AnfNode *node);
+  static void SetWorkspaceAddr(const DeviceAddressPtr &addr, size_t output_idx, const AnfNodePtr &node);
+  // set workspace kernel tensor of anf_node
+  static void SetWorkspaceKernelTensor(const KernelTensorPtr &kernel_tensor, size_t output_idx, AnfNode *node);
   // get workspace device addr of anf_node
   static DeviceAddress *GetWorkspaceAddr(const AnfNodePtr &node, size_t output_idx);
+  static KernelTensorPtr GetWorkspaceKernelTensor(const AnfNodePtr &node, size_t output_idx);
   // get workspace device mutable addr of anf_node
   static DeviceAddressPtr GetMutableWorkspaceAddr(const AnfNodePtr &node, size_t index);
   // get op pattern of the node
@@ -266,6 +270,7 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   // Whether the kernel is not supported by other device and need be backed off on the CPU device.
   static bool IsNodeSupportKernelSelectBackoff(const AnfNodePtr &node, const KernelGraphPtr &graph);
   static bool IsKernelSelectBackoffOp(const AnfNodePtr &node);
+  static bool IsNeedContinuesMemoryOp(const AnfNodePtr &node);
   static void SetKernelSelectBackoffInfo(const CNodePtr &node,
                                          const std::pair<std::string, ExceptionType> &failure_info);
   static std::pair<std::string, ExceptionType> GetKernelSelectBackoffInfo(const AnfNodePtr &node);
@@ -305,7 +310,7 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   static bool IsSequenceOutputOfScalar(const AnfNodePtr &node);
 
   // The tensor related interfaces.
-  static tensor::TensorPtr CreateMapTensor(const DeviceAddressPtr &output_device_address);
+  static tensor::TensorPtr CreateMapTensor(const KernelTensorPtr &output_kernel_tensor);
   static tensor::TensorPtr CreateMapTensor(const AnfNodePtr &output_node, size_t output_index);
   static tensor::TensorPtr SequenceToTensor(const ValuePtr &value);
   static void FlattenDynamicInputArg(const BaseRef &arg, const AnfNodePtr &node,
@@ -320,6 +325,7 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   static bool NeedEraseCache(const PrimitivePtr &prim);
 
   static abstract::AbstractBasePtr GetNodeAbstractByIndex(const AnfNodePtr &node, size_t index);
+  static abstract::AbstractBasePtr GetNodeAbstractByIndex(AnfNode *node, size_t index);
 
   static inline ValueNodePtr ConvertValueToNode(const KernelGraphPtr &kernel_graph, const ValuePtr &value) {
     MS_EXCEPTION_IF_NULL(kernel_graph);
@@ -341,6 +347,16 @@ class BACKEND_COMMON_EXPORT AnfRuntimeAlgorithm {
   // Only used for ascend ops.
   static bool IsLaunchIgnoredInputAddressIdx(const AnfNodePtr &node, size_t input_idx);
   static std::string GetValueByDeviceAddress(DeviceAddress *const device_address, size_t element_num);
+  static void PrintKernelTensor(const std::vector<KernelTensor *> &kernel_tensors, const std::string &info,
+                                size_t element_num);
+  static KernelTensorPtr CreateKernelTensor(const abstract::BaseShapePtr &shape, const TypePtr &type,
+                                            const ValuePtr &value, void *device_ptr, size_t size,
+                                            const std::string &format, TypeId dtype_id, const ShapeVector &host_shape,
+                                            const string &device_name, uint32_t device_id,
+                                            const UserDataPtr &user_data = nullptr);
+  static KernelTensorPtr CreateKernelTensor(void *device_ptr, size_t size, Format format, TypeId dtype_id,
+                                            const ShapeVector &host_shape, const string &device_name,
+                                            uint32_t device_id, const UserDataPtr &user_data = nullptr);
 };
 }  // namespace session
 

mindspore/include/mindspore/ccsrc/include/backend/data_queue/data_queue_mgr.h

@@ -167,7 +167,7 @@ BACKEND_EXPORT void RetryPeakItemFromDataQueue(const AnfNodePtr &data_kernel,
     device_name##DataQueueCreatorClass() {                                        \
       DataQueueMgr::GetInstance().RegisterDataQueueCreator(device_name, creator); \
     }                                                                             \
-  } g_##device_name##_data_queue_creator;
+  } g_##device_name##_data_queue_creator
 }  // namespace device
 }  // namespace mindspore
 

mindspore/include/mindspore/ccsrc/include/backend/debug/data_dump/dump_control.h

@@ -37,6 +37,7 @@ class BACKEND_COMMON_EXPORT DumpControl {
   void DynamicDumpStart();
   void DynamicDumpStop();
   void SetInitialIteration(std::uint32_t initial_iteration);
+  void UpdateUserDumpStep(const std::uint32_t step);
 
  private:
   DumpControl() = default;

mindspore/include/mindspore/ccsrc/include/backend/debug/data_dump/dump_json_parser.h

@@ -51,7 +51,6 @@ class BACKEND_COMMON_EXPORT DumpJsonParser {
   static bool DumpToFile(const std::string &filename, const void *data, size_t len, const ShapeVector &shape,
                          TypeId type);
   void CopyDumpJsonToDir(uint32_t rank_id);
-  void CopyHcclJsonToDir(uint32_t rank_id);
   void CopyMSCfgJsonToDir(uint32_t rank_id);
   bool NeedDump(const std::string &op_full_name);
   void MatchKernel(const std::string &kernel_name);
@@ -77,10 +76,11 @@ class BACKEND_COMMON_EXPORT DumpJsonParser {
   bool save_args_flag() const { return save_args_flag_; }
   uint32_t sample_mode() const { return sample_mode_; }
   uint32_t sample_num() const { return sample_num_; }
-  uint32_t cur_dump_iter() const { return cur_dump_iter_; }
+  uint32_t cur_dump_iter() const;
   uint32_t input_output() const { return input_output_; }
   void UpdateDumpIter() { ++cur_dump_iter_; }
   void UpdateDumpIter(int cur_step_count);
+  void UpdateUserDumpStep(const uint32_t step);
   bool GetDatasetSink() { return is_dataset_sink_; }
   void SetDatasetSink(bool is_dataset_sink) { is_dataset_sink_ = is_dataset_sink; }
   bool FileFormatIsNpy() const { return file_format_ == JsonFileFormat::FORMAT_NPY; }
@@ -88,14 +88,11 @@ class BACKEND_COMMON_EXPORT DumpJsonParser {
   bool DumpEnabledForIter() const;
   bool InputNeedDump() const;
   bool OutputNeedDump() const;
-  std::string GetOpOverflowBinPath(uint32_t graph_id) const;
   void GetCellDumpFlag(const session::KernelGraph &kernel_graph);
   void UpdateNeedDumpKernels(const session::KernelGraph &kernel_graph);
   bool IsDumpEnabled();
   bool IsDeviceCalcStats() const;
-  void PyNativeModeCheck();
   void CheckE2eSetting();
-  bool IsHCCLKernelInput(const std::string &kernel_name) const;
   bool IsCallbackRegistered() { return dumpdatacallback_registered_; }
   void SetCallbackRegistered() { dumpdatacallback_registered_ = true; }
   void SetInitialIteration(uint32_t initial_iteration);
@@ -153,8 +150,10 @@ class BACKEND_COMMON_EXPORT DumpJsonParser {
   uint32_t sample_mode_{0};
   uint32_t sample_num_{100};
   uint32_t cur_dump_iter_{0};
+  uint32_t cur_user_dump_iter_{0};
   uint32_t initial_dump_iter_{0};
   bool already_parsed_{false};
+  bool dump_user_step_flag_{false};
   std::string dump_layer_{""};
   std::string stat_calc_mode_{"host"};
   std::string device_stat_precision_mode_{"high"};

mindspore/include/mindspore/ccsrc/include/backend/debug/data_dump/dump_utils.h

@@ -24,6 +24,7 @@
 #include <numeric>
 
 #include "include/backend/kernel_graph.h"
+#include "include/common/utils/contract.h"
 #include "common/device_address.h"
 
 using DeviceTensor = mindspore::device::DeviceAddress;

mindspore/include/mindspore/ccsrc/include/backend/debug/debugger/debugger.h

@@ -81,8 +81,6 @@ class BACKEND_COMMON_EXPORT Debugger : public std::enable_shared_from_this<Debug
   // don't need a graph_ptr because it is saved during pre_execute
   void PostExecute();
 
-  static uint32_t GetRankID();
-
   void DumpConstantDataAscend(const KernelGraphPtr &graph);
 
   void DumpSingleNode(const CNodePtr &node, uint32_t graph_id, const DeviceContext *device_context = nullptr) const;

mindspore/include/mindspore/ccsrc/include/backend/debug/execute_order_tracker/execute_order_tracker.h

@@ -22,7 +22,9 @@
 #include <mutex>
 #include <tuple>
 #include <unordered_map>
+#include <utility>
 #include "include/backend/visible.h"
+#include "mindspore/ccsrc/pyboost/op_runner.h"
 #include "ir/anf.h"
 
 namespace mindspore {
@@ -70,6 +72,10 @@ class BACKEND_COMMON_EXPORT ExecuteOrderTracker {
 
   void ProcessNode(const CNodePtr &cnode);
 
+  void ProcessPyboostCommOp(const std::shared_ptr<kernel::pyboost::OpRunner> &op, const std::string &group,
+                            size_t comm_stream_id, const tensor::TensorPtr &input_tensor,
+                            const tensor::TensorPtr &output_tensor, int64_t rank);
+
   // Data is written to disk and cleaned up
   void Clear();
 
@@ -87,11 +93,15 @@ class BACKEND_COMMON_EXPORT ExecuteOrderTracker {
 
   bool IsCommunicationOp(const CNodePtr &cnode) const;
 
-  std::tuple<std::string, std::string, std::string, std::string, size_t, size_t> GetInputOutputShapeAndType(
-    const CNodePtr &cnode) const;
+  CommOrderInfoPtr CreateCommOrderInfo(const std::string &index, const std::string &group,
+                                       const std::string &primitive_str, const CNodePtr &cnode,
+                                       const tensor::TensorPtr &input_tensor = nullptr,
+                                       const tensor::TensorPtr &output_tensor = nullptr, int64_t direct_rank = -1);
+
+  void GetInputOutputShapeAndType(const CNodePtr &cnode, const CommOrderInfoPtr &comm_info) const;
 
-  std::tuple<std::string, std::string, std::string> GetCommunicationRanks(
-    const CNodePtr &cnode, const std::vector<uint32_t> &comm_ranks) const;
+  std::string GetCommunicationRanks(const std::variant<int64_t, std::pair<const CNodePtr &, const char *>> &input,
+                                    const std::vector<uint32_t> &comm_ranks) const;
 
   std::mutex mutex_;
   std::string order_path_;

mindspore/include/mindspore/ccsrc/include/backend/debug/tensor_data.h

@@ -25,9 +25,7 @@
 #include "base/bfloat16.h"
 #include "utils/log_adapter.h"
 #include "mindapi/base/type_id.h"
-#ifndef OFFLINE_DBG_MODE
 #include "ir/tensor.h"
-#endif
 
 constexpr size_t kFloat32Size = 4;
 constexpr size_t kFloat64Size = 8;
@@ -128,10 +126,8 @@ class TensorData {
     this->root_graph_id_ = obj.root_graph_id_;
     this->is_output_ = obj.is_output_;
     this->time_stamp_ = obj.time_stamp_;
-#ifndef OFFLINE_DBG_MODE
     this->format_ = obj.format_;
     this->tensor_ptr_ = obj.tensor_ptr_;
-#endif
   }
 
   TensorData &operator=(const TensorData &other) {
@@ -151,10 +147,8 @@ class TensorData {
       this->root_graph_id_ = other.root_graph_id_;
       this->is_output_ = other.is_output_;
       this->time_stamp_ = other.time_stamp_;
-#ifndef OFFLINE_DBG_MODE
       this->format_ = other.format_;
       this->tensor_ptr_ = other.tensor_ptr_;
-#endif
     }
     return *this;
   }
@@ -162,16 +156,8 @@ class TensorData {
   ~TensorData() { DeleteDataPtr(); }
 
   void DeleteDataPtr() noexcept {
-#ifndef OFFLINE_DBG_MODE
     this->tensor_ptr_ = nullptr;
     this->data_ptr_ = nullptr;
-#else
-    if (this->data_ptr_ != nullptr) {
-      delete[] this->data_ptr_;
-      this->data_ptr_ = nullptr;
-      this->size_ = 0;
-    }
-#endif
   }
 
   std::string GetName() const { return this->name_; }
@@ -188,13 +174,11 @@ class TensorData {
 
   void SetTimeStamp(const std::string &time_stamp) { this->time_stamp_ = time_stamp; }
 
-#ifndef OFFLINE_DBG_MODE
   void SetTensor(const mindspore::tensor::TensorPtr &out_tensor) { this->tensor_ptr_ = out_tensor; }
 
   void SetFormat(const std::string &format) { this->format_ = format; }
 
   std::string GetFormat() { return this->format_; }
-#endif
 
   void SetSlot(size_t slot) { this->slot_ = slot; }
 
@@ -462,11 +446,8 @@ class TensorData {
   int execution_order_{-1};
   std::string time_stamp_;
   const float INT4_SIZE = 0.5;
-
-#ifndef OFFLINE_DBG_MODE
   std::string format_{""};
   mindspore::tensor::TensorPtr tensor_ptr_{nullptr};
-#endif
 };
 }  // namespace mindspore
 #endif  // MINDSPORE_CCSRC_DEBUG_TENSOR_DATA_H_

mindspore/include/mindspore/ccsrc/include/backend/distributed/cluster/tcp_store.h

@@ -0,0 +1,53 @@
+/**
+ * Copyright 2025 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_CCSRC_DISTRIBUTED_TCP_STORE_H_
+#define MINDSPORE_CCSRC_DISTRIBUTED_TCP_STORE_H_
+
+#include <string>
+#include <memory>
+#include <vector>
+#if defined(__linux__) && defined(WITH_BACKEND)
+#include "include/backend/distributed/cluster/cluster_context.h"
+#include "include/backend/distributed/cluster/topology/compute_graph_node.h"
+#else
+#include "include/backend/distributed/cluster/dummy_cluster_context.h"
+#endif
+#include "pybind11/pybind11.h"
+namespace py = pybind11;
+
+namespace mindspore {
+namespace distributed {
+namespace cluster {
+class BACKEND_EXPORT TCPStoreClient {
+ public:
+  TCPStoreClient();
+  ~TCPStoreClient();
+  static std::shared_ptr<TCPStoreClient> instance();
+
+  // Get the rank id of this process in the specified group.
+  py::bytes GetKey(const std::string &key);
+
+  // Get the size of the specified group.
+  void SetKey(const std::string &key, const std::string &value);
+
+  bool DeleteKey(const std::string &key);
+};
+
+}  // namespace cluster
+}  // namespace distributed
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_DISTRIBUTED_TCP_STORE_H_

mindspore/include/mindspore/ccsrc/include/backend/distributed/collective/collective_manager.h

@@ -48,6 +48,7 @@ using DeviceContextKey = device::DeviceContextKey;
 using DeviceContextManager = device::DeviceContextManager;
 using CollectiveCommunicationLib = device::CollectiveCommunicationLib;
 using CommunicationGroupPtr = device::CommunicationGroupPtr;
+using GroupOptions = device::GroupOptions;
 
 // Data types for creating communication groups function.
 
@@ -65,13 +66,6 @@ using GroupToResultMap = std::unordered_map<std::string, std::pair<bool, std::st
 // Interval of initializing each communicator in queue is 300 milliseconds.
 const uint32_t kInitCommInterval = 300;
 
-// This the config passed to 'CreateCommunicationGroup' method. It controls initialization mode for communication group.
-struct CreateGroupConfig {
-  bool async = false;      // Whether creating communication group asynchonizely.
-  bool submit_now = true;  // For sync manner, this key means whether submit init task immediately for this group. If
-                           // set to false, caller has to call 'SubmitCreateDeviceCommTask' itself.
-};
-
 // The collective communication API.
 // MindSpore uses OpenMPI on CPU, NCCL on GPU, HCCL on Ascend, to achieve distributed training.
 // Besides, MindSpore also has its own communication library which is implemented on the CPU side.
@@ -89,14 +83,19 @@ class BACKEND_COMMON_EXPORT CollectiveManager {
 
   // Create communication group.
   bool CreateCommunicationGroup(const std::string &group_name, const std::vector<uint32_t> &group_ranks,
-                                const CreateGroupConfig &config = {});
+                                const GroupOptions &config = {});
 
   // Destroy the communication group.
   bool DestroyCommunicationGroup(const std::string &group_name);
 
+  void RemoveGroupInfoForARF(const std::string &group_name);
+
   // Destroy device communication group for resume training.
   bool DestroyDeviceSideCommunicationGroup(const std::string &group_name);
 
+  // Call HcclSetGlobalCommInfo to pass necessary info in world group.
+  void SetGlobalCommInfo(CommunicationGroupPtr group, const std::string &group_name);
+
   // Get the inner comm name of the specified group.
   std::string GetCommName(const std::string &group_name);
 
@@ -142,14 +141,22 @@ class BACKEND_COMMON_EXPORT CollectiveManager {
 
   CollectiveCommunicationLib *device_comm_lib() { return device_comm_lib_instance_; }
 
+  void CacheInitedGroups(const std::string &name);
+  void ClearCacheInitedGroups();
+  size_t InitedGroupSize();
+
   // Initialize and finalize Dummy communication lib.
   bool InitializeDummyCommLib();
   bool FinalizeDummyCommLib();
 
   bool ResumeHcclComm();
+
   // clean unique if after group build success when using ARF
   void ClearUniqueID(const std::string &group_name);
 
+  // switch network interface card between the primary and the secondary NIC.
+  bool CommSwitchNic(const std::vector<uint32_t> &global_ranks, const std::vector<bool> &use_backup);
+
   // Return whether initializing global comm asynchronizely.
   bool IsAsyncInitGlobalComm();
 
@@ -277,6 +284,8 @@ class BACKEND_COMMON_EXPORT CollectiveManager {
   std::condition_variable result_blocker_;
   GroupToResultMap group_name_to_result_;
   std::vector<std::pair<std::string, std::vector<uint32_t>>> group_infos_;
+  mutable std::mutex cache_mutes_;
+  std::vector<std::string> inited_groups_;
 };
 
 // For scheduler node, CollectiveManager is not initialized. Return 0 as rank id.

mindspore/include/mindspore/ccsrc/include/backend/distributed/constants.h

@@ -37,8 +37,7 @@ namespace distributed {
 constexpr char kDetailedFailureReason[] =
   "Maybe you are trying to call 'mindspore.communication.init()' without using 'mpirun', which will make MindSpore "
   "load several environment variables and check their validation. Please use 'mpirun' to launch this process to fix "
-  "this issue, or refer to this link if you want to run distributed training without using 'mpirun': "
-  "https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/dynamic_cluster.html";
+  "this issue, or refer to this link if you want to run distributed training using 'msrun'";
 
 constexpr char kWorkerProcessNotEnoughError[] = "Spawned worker process number is not as expected.";
 constexpr char kSchedPortOccupiedError[] = "Configured scheduler port MS_SCHED_PORT is occupied by other processes.";
@@ -125,14 +124,6 @@ constexpr char kControlDstOpName[] = "ControlDst";
 static const char URL_PROTOCOL_IP_SEPARATOR[] = "://";
 static const char URL_IP_PORT_SEPARATOR[] = ":";
 
-constexpr char kEnableRDMA[] = "enable_rdma";
-constexpr char kRDMADevName[] = "rdma_dev";
-constexpr char kRDMAIP[] = "rdma_ip";
-
-constexpr char kDefaultIP[] = "1.1.8.203";
-constexpr char kDefaultIfName[] = "hrn0_2";
-constexpr uint16_t kDefaultPort = 10969;
-
 // The interval of retrying connecting for rpc clients.
 constexpr uint32_t kRetryConnectInterval = 2;
 

mindspore/include/mindspore/ccsrc/include/backend/distributed/embedding_cache/embedding_hash_map.h

@@ -23,8 +23,6 @@
 #include <vector>
 #include <list>
 #include <mutex>
-#include "utils/hash_map.h"
-#include "utils/convert_utils_base.h"
 #include "include/backend/visible.h"
 #include "distributed/embedding_cache/cache_strategy/cache.h"
 

mindspore/include/mindspore/ccsrc/include/backend/distributed/ps/ps_context.h

@@ -97,9 +97,6 @@ class BACKEND_COMMON_EXPORT PSContext {
 
   core::ClusterConfig &cluster_config();
 
-  void set_scheduler_manage_port(uint16_t sched_port);
-  uint16_t scheduler_manage_port() const;
-
   void set_config_file_path(const std::string &path);
   std::string config_file_path() const;
 
@@ -154,9 +151,6 @@ class BACKEND_COMMON_EXPORT PSContext {
   // The cluster config read through environment variables, the value does not change.
   std::unique_ptr<core::ClusterConfig> cluster_config_;
 
-  // The port used by scheduler to receive http requests for scale out or scale in.
-  uint16_t scheduler_manage_port_;
-
   // The path of the configuration file, used to configure the certification path and persistent storage type, etc.
   std::string config_file_path_;
 

mindspore/include/mindspore/ccsrc/include/backend/distributed/rpc/tcp/constants.h

@@ -42,7 +42,8 @@ constexpr int SENDMSG_DROPED = -1;
 constexpr size_t MAX_KMSG_FROM_LEN = 1024;
 constexpr size_t MAX_KMSG_TO_LEN = 1024;
 constexpr size_t MAX_KMSG_NAME_LEN = 1024;
-constexpr size_t MAX_KMSG_BODY_LEN = 1073741824;
+// MAX_KMSG_BODY_LEN=2GB-4k, Protobuf limits 2G size and reserves 4k for message headers.
+constexpr size_t MAX_KMSG_BODY_LEN = 2147479552;
 
 enum ParseType { kTcpMsg = 1, kHttpReq, kHttpRsp, kUnknown };
 enum State { kMsgHeader, kBody };

mindspore/include/mindspore/ccsrc/include/backend/kernel_graph.h

@@ -29,7 +29,6 @@
 #include "utils/hash_set.h"
 #include "ir/func_graph.h"
 #include "ir/anf.h"
-#include "ir/graph_utils.h"
 #include "include/common/utils/contract.h"
 #include "common/device_type.h"
 #include "include/backend/kernel_info.h"

mindspore/include/mindspore/ccsrc/include/backend/kernel_info.h

@@ -37,8 +37,8 @@ class BACKEND_COMMON_EXPORT KernelInfo : public KernelInfoDevice {
     kernel_mod_ = nullptr;
     is_feature_map_ = false;
     select_kernel_build_info_ = nullptr;
-    output_address_list_ = {};
-    workspace_address_list_ = {};
+    output_kernel_tensor_list_ = {};
+    workspace_kernel_tensor_list_ = {};
     stream_id_ = 0;
     stream_distinction_label_ = kInvalidDistincLabel;
     graph_id_ = kInvalidGraphId;
@@ -61,12 +61,11 @@ class BACKEND_COMMON_EXPORT KernelInfo : public KernelInfoDevice {
   bool SetOutputAddr(const DeviceAddressPtr &output_address, size_t index);
   DeviceAddress *GetWorkspaceAddr(size_t index) const;
   DeviceAddressPtr GetMutableWorkspaceAddr(size_t index) const;
+  KernelTensorPtr GetWorkspaceKernelTensor(size_t index) const;
   bool WorkspaceAddrExist(size_t index) const;
   bool SetWorkspaceAddr(const DeviceAddressPtr &output_address, size_t index);
-  // The number of workspace may change after kernel Resize.
-  void set_workspace_address_list(const std::vector<DeviceAddressPtr> &workspace_address_list) {
-    workspace_address_list_ = workspace_address_list;
-  }
+  bool SetWorkspaceKernelTensor(const KernelTensorPtr &kernel_tensor, size_t index);
+  bool WorkspaceKernelTensorExist(size_t index) const;
   void set_kernel_mod(const kernel::KernelModPtr &kernel_mod);
   kernel::KernelMod *MutableKernelMod() const;
   kernel::KernelModPtr GetKernelMod() const;
@@ -82,8 +81,8 @@ class BACKEND_COMMON_EXPORT KernelInfo : public KernelInfoDevice {
   bool operator==(const KernelInfo &other) const;
   bool is_feature_map() const { return is_feature_map_; }
 
-  const std::vector<std::shared_ptr<DeviceAddress>> &output_address_list() const { return output_address_list_; }
-  const std::vector<std::shared_ptr<DeviceAddress>> &workspace_address_list() const { return workspace_address_list_; }
+  const std::vector<KernelTensorPtr> &output_kernel_tensor_list() const { return output_kernel_tensor_list_; }
+  const std::vector<KernelTensorPtr> &workspace_kernel_tensor_list() const { return workspace_kernel_tensor_list_; }
 
   // Set output and input reference. If all_ref is set true, each output is a reference to the input with the same
   // index.
@@ -105,8 +104,7 @@ class BACKEND_COMMON_EXPORT KernelInfo : public KernelInfoDevice {
   bool is_feature_map_;
   kernel::KernelBuildInfoPtr select_kernel_build_info_;
   std::vector<KernelTensorPtr> output_kernel_tensor_list_;
-  std::vector<std::shared_ptr<DeviceAddress>> output_address_list_;
-  std::vector<std::shared_ptr<DeviceAddress>> workspace_address_list_;
+  std::vector<KernelTensorPtr> workspace_kernel_tensor_list_;
   // pair<size_t, size_t> : (offset, aligned_size)
   // aligned_size of 0 means no memory allocation
   std::vector<std::pair<size_t, size_t>> somas_output_result_;

mindspore/include/mindspore/ccsrc/include/backend/mbuf_device_address.h

@@ -18,7 +18,7 @@
 
 #include <string>
 #include <memory>
-#include "common/device_address.h"
+#include "common/kernel.h"
 
 namespace mindspore {
 namespace device {
@@ -31,13 +31,13 @@ class MbufDeviceAddress : public device::DeviceAddress {
     auto tensor_shape = std::make_shared<abstract::TensorShape>();
     tensor_shape->SetShapeVector(shape);
     auto tensor_type = std::make_shared<TensorType>(TypeIdToType(type));
-    kernel_tensor_ = std::make_shared<KernelTensor>(tensor_shape, tensor_type, nullptr, ptr, size, "DefaultFormat",
-                                                    type, shape, device_name, device_id, nullptr);
-    address_common_ = kernel_tensor_->address_common();
+    address_common_ = std::make_shared<AddressCommon>(ptr, size, shape, kernel::GetFormatFromStrToEnum("DefaultFormat"),
+                                                      type, device_name, device_id, 0);
   }
   void SetData(void *data) { set_ptr(data); }
 
-  bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override {
+  bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr,
+                        bool sync_on_demand = false) const override {
     MS_LOG(ERROR) << "Mbuf address does not support sync data from device to host, please use graph mode";
     return false;
   }

mindspore/include/mindspore/ccsrc/include/backend/mem_reuse/abstract_dynamic_mem_pool.h

@@ -31,7 +31,6 @@
 #include "include/backend/mem_reuse/dynamic_mem_pool.h"
 #include "include/backend/visible.h"
 #include "include/common/utils/stream_util.h"
-#include "utils/ms_utils.h"
 
 namespace mindspore {
 namespace device {
@@ -256,6 +255,8 @@ class BACKEND_EXPORT MemBufAllocator {
   MemBuf *MallocExpandBlock(size_t size);
   const std::pair<size_t, size_t> FreeIdleMemsByEagerFree();
 
+  size_t ReleaseFreeBlocks();
+
   std::string DumpStateInfo() const;
   std::string DumpDebugInfo() const;
 
@@ -290,15 +291,15 @@ class BACKEND_EXPORT MemBufAllocator {
   bool enable_eager_free_;
 
   std::list<MemBlock *> mem_blocks_;
-  std::set<MemBuf *, MemBufComparator> free_mem_bufs_;
-  std::set<MemBuf *, MemBufComparator> eager_free_mem_bufs_;
+  using MemAllocator = memory::mem_pool::PooledAllocator<MemBuf *>;
+  std::set<MemBuf *, MemBufComparator, MemAllocator> free_mem_bufs_;
+  std::set<MemBuf *, MemBufComparator, MemAllocator> eager_free_mem_bufs_;
 
  private:
   MemBuf *search_key_;
 
   uint32_t stream_id_;
   bool is_persistent_;
-  bool has_do_eager_free_{false};
 
   friend AbstractDynamicMemPool;
 };
@@ -394,6 +395,8 @@ class BACKEND_EXPORT AbstractDynamicMemPool : virtual public DynamicMemPool {
 
   const std::pair<size_t, size_t> FreeIdleMemsByEagerFree() override;
 
+  size_t ReleaseFreeBlocks() override;
+
   MemStat &mem_stat() { return mem_stat_; }
 
   Lock &lock() { return lock_; }
@@ -441,7 +444,9 @@ class BACKEND_EXPORT AbstractEnhancedDynamicMemPool : public AbstractDynamicMemP
 
   // Report memory pool stat info for enhanced processing.
   virtual void ReportMemoryPoolInfo();
-
+  // Report memory pool stat info for mstx
+  virtual void ReportMemoryPoolMallocInfoToMstx(void *ptr, size_t size);
+  virtual void ReportMemoryPoolFreeInfoToMstx(void *ptr);
   bool IsEnableTimeEvent() override { return enable_time_event_; }
 
   void SetEnableTimeEvent(bool enable_time_event) override { enable_time_event_ = enable_time_event; }