mindspore 2.6.0rc1__cp311-cp311-manylinux1_x86_64.whl → 2.7.0__cp311-cp311-manylinux1_x86_64.whl

mindspore/include/mindspore/ccsrc/pynative/op_function/comm_handle_py.h

@@ -19,6 +19,7 @@
 
 #include <string>
 #include <memory>
+#include <set>
 #include "mindspore/ccsrc/pyboost/comm_handle.h"
 
 namespace mindspore {
@@ -42,6 +43,7 @@ class CommHandlePy {
  private:
   const device::DeviceContext *device_ctx_;
   kernel::pyboost::CommHandlePtr comm_handle_;
+  std::set<size_t> wait_streams_;
 };
 
 using CommHandlePyPtr = std::shared_ptr<CommHandlePy>;

mindspore/include/mindspore/ccsrc/pynative/op_function/converter.h

@@ -141,6 +141,9 @@ struct PYNATIVE_EXPORT ParserArgs {
   void ClearArgs();
   const int &GetOvertLoadIndex() { return signature_->index_; }
   void PrintConvertError(size_t index);
+  // convert to basic type
+  std::vector<int64_t> ToBasicIntVector(size_t index);
+  int64_t ToBasicInt(size_t index);
 
   template <typename T>
   std::shared_ptr<T> Convert(size_t index) {
@@ -241,7 +244,15 @@ class PYNATIVE_EXPORT Converter {
   std::optional<Int64ImmPtr> ToDtypeOptional(const py::list &python_args, size_t i);
   ValuePtr ConvertByCastDtype(const py::object &input, const ops::OpInputArg &op_arg, size_t i);
   ValueTuplePtr ConvertValueTupleByCastDtype(const py::list &python_args, const ops::OpInputArg &op_arg, size_t index);
+  std::vector<int64_t> ConvertIntVectorByCastDtype(const py::list &python_args, const ops::OpInputArg &op_arg,
+                                                   size_t index);
+  int64_t ConvertIntByCastDtype(const py::list &python_args, const ops::OpInputArg &op_arg, size_t index);
   const std::vector<ops::OP_DTYPE> &source_type() const { return source_type_; }
+  // basic type
+  int64_t ToBasicInt(const py::list &python_args, size_t i);
+  std::optional<int64_t> ToBasicIntOptional(const py::list &python_args, size_t i);
+  std::vector<int64_t> ToBasicIntVector(const py::list &python_args, size_t i);
+  std::optional<std::vector<int64_t>> ToBasicIntVectorOptional(const py::list &python_args, size_t i);
 
  private:
   ops::OpDefPtr op_def_;

mindspore/include/mindspore/ccsrc/pynative/op_function/customize/direct_ops.h

@@ -26,15 +26,5 @@
 #include "mindspore/core/include/ir/value.h"
 
 namespace py = pybind11;
-namespace mindspore::pynative {
-
-PYNATIVE_EXPORT py::object Empty(const py::list &args);
-PYNATIVE_EXPORT py::object EmptyLike(const py::list &args);
-PYNATIVE_EXPORT py::object NewEmpty(const py::list &args);
-PYNATIVE_EXPORT py::object Pyboost_Empty_Base(const PrimitivePtr &prim, const py::list &args);
-PYNATIVE_EXPORT py::object Pyboost_Empty_OP(const PrimitivePtr &prim,
-                                            const std::vector<mindspore::ops::OP_DTYPE> &source_type,
-                                            const ValueTuplePtr &shape, const std::optional<Int64ImmPtr> &dtype,
-                                            const std::optional<StringImmPtr> &device);
-}  // namespace mindspore::pynative
-#endif  // MINDSPORE_CCSRC_PIPELINE_PYNATIVE_OP_FUNCTION_CUSTOMIZE_DIRECT_OPS_H
+namespace mindspore::pynative {}  // namespace mindspore::pynative
+#endif                            // MINDSPORE_CCSRC_PIPELINE_PYNATIVE_OP_FUNCTION_CUSTOMIZE_DIRECT_OPS_H

mindspore/include/mindspore/ccsrc/pynative/predict_out_type_map.h

@@ -56,6 +56,7 @@ inline static PredictOutTypeMap out_type_prediction = {{"ActsULQ", kTupleTensor4
                                                        {"AdaptiveMaxPool2D", kTupleTensor2},
                                                        {"AdaptiveMaxPool3D", kTupleTensor2},
                                                        {"TopKRouter", kTypeAny},
+                                                       {"TopPRouter", kTypeAny},
                                                        {"ApplyAdaMax", kTupleTensor3},
                                                        {"ApplyAdadelta", kTupleTensor3},
                                                        {"ApplyAdagrad", kTupleTensor2},
@@ -305,6 +306,8 @@ inline static PredictOutTypeMap out_type_prediction = {{"ActsULQ", kTupleTensor4
 
 TypePtr PredictOutTypeByName(const std::string &op_name);
 
+TypePtr GetPredictOutTypeByName(const std::string &op_run_info_op_name);
+TypePtr GetPredictOutTypeByOutputNum(const PrimitivePtr &op_run_info_op_prim, const TypePtr &type);
 TypePtr PredictOutType(const FrontendOpRunInfoPtr &op_run_info);
 }  // namespace pynative
 }  // namespace mindspore

mindspore/include/mindspore/ccsrc/pynative/pynative_execute.h

@@ -28,6 +28,7 @@
 #include "ir/anf.h"
 #include "include/fork_utils.h"
 #include "include/common/visible.h"
+#include "include/common/pynative/variable.h"
 
 namespace mindspore::pynative {
 namespace py = pybind11;
@@ -38,6 +39,7 @@ class PYNATIVE_EXPORT PyNativeExecutor : public std::enable_shared_from_this<PyN
     std::lock_guard<std::mutex> i_lock(instance_lock_);
     if (executor_ == nullptr) {
       executor_ = std::shared_ptr<PyNativeExecutor>(new (std::nothrow) PyNativeExecutor());
+      MS_EXCEPTION_IF_NULL(executor_);
       Init();
     }
     return executor_;
@@ -54,6 +56,7 @@ class PYNATIVE_EXPORT PyNativeExecutor : public std::enable_shared_from_this<PyN
     MS_EXCEPTION_IF_NULL(forward_executor_);
     return forward_executor_;
   }
+  void StoreAsyncStatus(const PyboostOpRunInfoPtr &op_run_info) const;
   void StoreAsyncStatus(const FrontendOpRunInfoPtr &op_run_info) const;
   // Generate stub tensor and dispatch async task.
   py::object RunOpStub(const py::args &args) const;
@@ -73,13 +76,13 @@ class PYNATIVE_EXPORT PyNativeExecutor : public std::enable_shared_from_this<PyN
   py::object RunGrad(const prim::GradOperationPtr &grad, const py::object &cell, const py::object &weights,
                      const py::object &grad_position, const py::object &has_aux, const py::args &args) const;
   py::object GradJit(const py::args &args) const;
-  void CallCustomBprop(const py::object &cell_obj, const py::object &out, const py::args &args) const;
+  py::object CallCustomBprop(const py::object &cell_obj, const py::object &out, const py::args &args) const;
   void set_forward_use_dynamic_shape_process(bool flag) const;
   void SetDynamicInput(const py::object &obj, const py::args &args) const;
   py::object GetDynamicInput(const py::object &actual_input) const;
   bool IsHighOrder() const;
   py::object CheckAlreadyRun(const prim::GradOperationPtr &grad, const py::object &obj, const py::object &weights,
-                             const py::object &grad_hash_id, const py::args &args) const;
+                             const py::object &grad_hash_id, const py::args &args, const py::kwargs &kwargs) const;
   void ClearRes() const;
   // Sync stream
   void Sync() const;
@@ -91,6 +94,7 @@ class PYNATIVE_EXPORT PyNativeExecutor : public std::enable_shared_from_this<PyN
   void ChildAfterFork();
   py::object RunSliceOpStub(const std::vector<ValuePtr> &input_v,
                             const std::vector<SliceOpInfoPtr> &slice_op_infos) const;
+  void SetCreationType(const py::object &obj, autograd::CreationType creation_type);
 
  private:
   PyNativeExecutor() = default;

mindspore/include/mindspore/ccsrc/pynative/pynative_utils.h

@@ -47,7 +47,7 @@ struct Common {
   static bool IsTensor(const ValuePtr &v, bool include_sequence = false);
   static bool IsControlFlowGraph(const FuncGraphPtr &func_graph);
   static ValuePtr FilterSensValues(const ValuePtr &value, bool dict_convert_to_tuple);
-  static tensor::BaseTensorPtr GetTensorFromParam(const AnfNodePtr &param_node);
+  static tensor::TensorPtr GetTensorFromParam(const AnfNodePtr &param_node);
   static TypeId GetTypeFromAbstract(const abstract::AbstractBasePtr &abs);
   static ShapeVector GetShapeFromAbstract(const abstract::AbstractBasePtr &abs);
   static std::pair<TypePtr, TypeId> GetTypeFromValue(const ValuePtr &v);
@@ -57,20 +57,17 @@ struct Common {
   static const std::shared_ptr<PyNativeExecutor> &GetPyNativeExecutor();
   static ValuePtr StubNodeToValue(const ValuePtr &val);
   static void StubNodeToValue(const FrontendOpRunInfoPtr &op_run_info);
-  static tensor::BaseTensorPtr StubNodeToTensor(const ValuePtr &value);
-  PYNATIVE_EXPORT static tensor::BaseTensorPtr ConvertStubNodeToTensor(const ValuePtr &v, bool need_contiguous,
-                                                                       bool requires_grad);
-  PYNATIVE_EXPORT static std::optional<tensor::BaseTensorPtr> ConvertStubNodeToTensor(const std::optional<ValuePtr> &v,
-                                                                                      bool need_contiguous,
-                                                                                      bool requires_grad);
+  static tensor::TensorPtr StubNodeToTensor(const ValuePtr &value);
+  PYNATIVE_EXPORT static tensor::TensorPtr ConvertStubNodeToTensor(const ValuePtr &v, bool need_contiguous,
+                                                                   bool requires_grad);
+  PYNATIVE_EXPORT static std::optional<tensor::TensorPtr> ConvertStubNodeToTensor(const std::optional<ValuePtr> &v,
+                                                                                  bool need_contiguous,
+                                                                                  bool requires_grad);
   static ValueTuplePtr ConvertStubNodeToValueTuple(const ValueListPtr &v, bool need_contiguous, bool requires_grad);
   static ValueTuplePtr ConvertStubNodeToValueTuple(const ValueTuplePtr &v, bool need_contiguous, bool requires_grad);
   static std::optional<ValueTuplePtr> ConvertStubNodeToValueTuple(const std::optional<ValueTuplePtr> &v,
                                                                   bool need_contiguous, bool requires_grad);
   static ValuePtr StubNodeToValueInner(const ValuePtr &v);
-  static void GetConstInputToAttr(const PrimitivePtr &op_prim, const std::string &op_name,
-                                  const std::string &device_target, bool is_dynamic_shape,
-                                  mindspore::HashSet<size_t> *input_to_attr_index);
   static ValueNodePtr CreateValueNodeByValue(const ValuePtr &v, const abstract::AbstractBasePtr &abs = nullptr);
   static void SetOutputUsedInBpropGraph(const ValuePtr &value);
   static ValuePtr CreateFakeValueWithoutDeviceAddress(const ValuePtr &value, bool is_force_create_fake = false);
@@ -79,10 +76,10 @@ struct Common {
   static inline bool IsConstant(InputType grad_type) { return grad_type == InputType::kConstant; }
   static void SetGraphInputAndWeightsInfo(const FrontendOpRunInfoPtr &op_run_info, const FuncGraphPtr &func_graph);
   static void FreeFuncGraphForwardNodes(const FuncGraphPtr &func_graph);
-  static tensor::BaseTensorPtr ConvertToContiguousTensor(const tensor::BaseTensorPtr &tensor, bool requires_grad);
+  static tensor::TensorPtr ConvertToContiguousTensor(const tensor::TensorPtr &tensor, bool requires_grad);
   static ValuePtr ConvertToContiguousValue(const ValuePtr &v, bool requires_grad);
   static ValuePtr CreateTensorByConstantValue(const ValuePtr &value);
-
+  static tensor::TensorPtr CaculateGradNorm(const tensor::TensorPtr &grad);
   template <typename T>
   static std::string PrintDebugInfo(std::vector<T> items, const std::string &info_header = "",
                                     bool is_print_tensor_data = false) {
@@ -94,12 +91,13 @@ struct Common {
         MS_LOG(DEBUG) << "The " << i << "'th item is nullptr!";
         continue;
       }
-      if (items[i]->template isa<tensor::BaseTensor>() && is_print_tensor_data) {
-        auto tensor = items[i]->template cast<tensor::BaseTensorPtr>();
+      if (items[i]->template isa<tensor::Tensor>() && is_print_tensor_data) {
+        auto tensor = items[i]->template cast<tensor::TensorPtr>();
         auto grad = std::make_shared<tensor::Tensor>(*tensor);
-        grad->data_sync();
+        auto norm_val = CaculateGradNorm(grad);
+        norm_val->data_sync();
         buf << i << "th: "
-            << "ptr " << items[i].get() << ", " << grad->ToStringRepr() << ", ";
+            << "ptr " << items[i].get() << ", " << norm_val->ToStringRepr() << ", ";
       } else {
         buf << i << "th: "
             << "ptr " << items[i].get() << ", " << items[i]->ToString() << ", ";
@@ -111,11 +109,10 @@ struct Common {
   static bool IsVmOp(const std::string &op_name);
   static std::vector<int64_t> BuildShape(const abstract::AbstractBasePtr &abs);
   static void ClearRes();
-  static TopCellInfo *FindPreTopcell(const GradExecutor *grad_executor, const OpGradInfoPtr &op_grad_info,
-                                     const std::string &op_info, const ValuePtr &value);
-  static void UpdateGradOpInfo(const GradExecutor *grad_executor, const OpGradInfoPtr &op_grad_info,
-                               TopCellInfo *pre_top_cell, bool is_jit_graph);
   static OperatorType GetOpTypeFromOpdef(const ops::OpDef &op_def);
+  static void DoGradInner(runtime::OpRunnerInfo *op_runner_info, VectorRef *op_outputs);
+  static tensor::TensorPtr GetTensorFromSparseTensor(const ValuePtr &val);
+  static void WaitBprop();
 };
 
 // Parser python
@@ -140,37 +137,34 @@ struct PyParser {
 
 // Data convert
 struct DataConvert {
-  static ValuePtrList FlattenOnlyTensor(const ValuePtr &v);
   static void FlattenArgs(const std::vector<ValuePtr> &v_vec, std::vector<ValuePtr> *flatten_v, bool has_sens);
-  static ValuePtrList FlattenTensorSeqInValue(const ValuePtr &v);
-  static void GetInputTensor(const FrontendOpRunInfoPtr &op_run_info, const TopCellInfoPtr &top_cell);
+  static void GetInputTensor(const FrontendOpRunInfoPtr &op_run_info);
   static void ConvertCSRTensorToTensorList(const FrontendOpRunInfoPtr &op_run_info,
-                                           const tensor::CSRTensorPtr &csr_tensor, const TopCellInfoPtr &top_cell,
-                                           size_t index);
+                                           const tensor::CSRTensorPtr &csr_tensor, size_t index);
   static void ConvertMapTensor(const FrontendOpRunInfoPtr &op_run_info, const tensor::MapTensorPtr &map_tensor,
-                               const TopCellInfoPtr &top_cell, size_t index);
+                               size_t index);
   static ValuePtr ConvertValueDictToValueTuple(const ValuePtr &v);
   static void PlantTensorTupleToVector(const FrontendOpRunInfoPtr &op_run_info, const ValueSequencePtr &value_seq,
-                                       size_t index, const TopCellInfoPtr &top_cell);
-  static void GetTensorIdFromOutputValue(const ValuePtr &value, std::vector<std::string> *converted_tensor_id);
+                                       size_t index);
   static void ConvertTupleValueToTensor(const FrontendOpRunInfoPtr &op_run_info, const ValueSequencePtr &value_seq,
-                                        size_t index, const TopCellInfoPtr &top_cell);
-  static void MarkInputs(const FrontendOpRunInfoPtr &op_run_info, const ValuePtr &v, size_t index,
-                         const TopCellInfoPtr &top_cell);
+                                        size_t index);
+  static void MarkInputs(const FrontendOpRunInfoPtr &op_run_info, const ValuePtr &v, size_t index);
   static bool RunOpConvertConstInputToAttr(const FrontendOpRunInfoPtr &op_run_info, const ValuePtr &v,
                                            size_t input_index);
-  static void TransformValueNodeBaseTensorToTensor(const ValueNodePtr &value_node);
+  static void TransformValueNodeTensorToTensor(const ValueNodePtr &value_node);
   static ValuePtr ValueListToValue(const ValuePtrList &values, const abstract::AbstractBasePtr &abs);
-  static ValuePtrList TensorListToValueList(const tensor::BaseTensorPtrList &tensor_list);
+  static ValuePtrList TensorListToValueList(const tensor::TensorPtrList &tensor_list);
 };
 
 struct PyBoost {
+  static PyboostOpRunInfoPtr Init_Pyboost(const PrimitivePtr &prim);
   static FrontendOpRunInfoPtr Init(const PrimitivePtr &prim);
   static void DoGrad(const kernel::pyboost::OpPtr &op, const OpGradInfoPtr &grad_info, const AsyncStatus &async_status);
   static void SetAnyValueForAbstract(const kernel::pyboost::OpPtr &op);
   static void UpdateStubOutput(const kernel::pyboost::OpPtr &op, const stub::StubNodePtr &stub_output,
                                const AbstractBasePtr &abstract, const ValuePtr &real_out);
   static PrimitivePtr ConvertPrimitive(const py::object &obj);
+
   static py::object RunPyFunction(const PrimitivePtr &prim, const py::list &args);
   template <typename T>
   static ValuePtr OptionalToValue(const std::optional<T> &val) {
@@ -193,31 +187,33 @@ struct PyBoost {
   }
 
   template <size_t N, typename... T>
-  static auto SetPyBoostCastForInputs(const FrontendOpRunInfoPtr &op_run_info,
+  static auto SetPyBoostCastForInputs(const PyboostOpRunInfoPtr &op_run_info, const std::string &op_name,
                                       const std::vector<std::vector<size_t>> &same_type_table, T... t) {
     MS_EXCEPTION_IF_NULL(op_run_info);
-    op_run_info->input_size = sizeof...(t);
-    if (op_run_info->op_grad_info->op_prim->name() == kCast) {
+    const size_t &input_size = sizeof...(t);
+    if (op_name == kCast) {
       return std::make_tuple(t...);
     }
     const auto &pyboost_cast_operation = Common::GetPyNativeExecutor()->forward_executor()->pyboost_cast_operation();
-    const auto &ret = pyboost_cast_operation->DoMixPrecisionCast(op_run_info, t...);
+    const auto &ret = pyboost_cast_operation->DoMixPrecisionCast(op_run_info, input_size, t...);
     if constexpr (N != 0) {
-      return pyboost_cast_operation->DoImplicitCast<N>(op_run_info, same_type_table, ret);
+      return pyboost_cast_operation->DoImplicitCast<N>(op_run_info, input_size, same_type_table, ret);
     }
     return ret;
   }
+
   static void DataSyncForGraph(const kernel::pyboost::OpPtr &op);
-  static void MarkPyBoostInputs(const OpGradInfoPtr &op_grad_info, const TopCellInfoPtr &top_cell);
-  static void BumpVersionAsync(const tensor::BaseTensorPtr &tensor);
-  static ValuePtr OutputToValue(const BaseTensorPtr &output) { return output; }
-  static ValuePtr MultiOutputToValue(const std::vector<BaseTensorPtr> &outputs) {
+  static void MarkPyBoostInputs(const OpGradInfoPtr &op_grad_info);
+  static void BumpVersionAsync(const tensor::TensorPtr &tensor);
+  static ValuePtr OutputToValue(const tensor::TensorPtr &output) { return output; }
+  static ValuePtr MultiOutputToValue(const std::vector<TensorPtr> &outputs) {
     std::vector<ValuePtr> output_values;
     output_values.reserve(outputs.size());
     (void)std::transform(outputs.begin(), outputs.end(), std::back_inserter(output_values),
-                         [](const BaseTensorPtr &value) -> ValuePtr { return value; });
+                         [](const TensorPtr &value) -> ValuePtr { return value; });
     return std::make_shared<ValueTuple>(output_values);
   }
+
   template <typename... Args>
   static ValuePtr MultiOutputToValue(const std::tuple<Args...> &outputs) {
     std::vector<ValuePtr> output_values = TupleToVector(outputs);

mindspore/include/mindspore/ccsrc/runtime/collective/collective_communication_lib.h

@@ -22,12 +22,22 @@
 #include <memory>
 #include <vector>
 #include <string>
+#include <variant>
+#include <unordered_map>
 #include "ir/dtype/type_id.h"
 #include "runtime/collective/communication_group.h"
 #include "include/backend/visible.h"
 
 namespace mindspore {
 namespace device {
+// This the config passed to 'CreateCommunicationGroup' method. It controls initialization mode for communication group.
+struct BACKEND_COMMON_EXPORT GroupOptions {
+  // Whether creating communication group asynchonizely.
+  bool async = false;
+  // Used to specify some hccl settings.
+  std::unordered_map<std::string, std::variant<uint32_t, std::string>> hccl_config = {};
+};
+
 // The reduce type of collective operations.
 enum CollectiveOpReduceType : int64_t {
   Reduce_Mean = 0,
@@ -63,7 +73,8 @@ class BACKEND_COMMON_EXPORT CollectiveCommunicationLib {
 
   // Create communication group. This is the precondition for all collective operations on both host and device side.
   virtual bool CreateCommunicationGroup(const std::string &group_name, const std::vector<uint32_t> &group_ranks,
-                                        uint32_t local_group_rank, uint32_t local_group_size) {
+                                        uint32_t local_group_rank, uint32_t local_group_size,
+                                        const GroupOptions &config = {}) {
     return true;
   }
 
@@ -81,6 +92,11 @@ class BACKEND_COMMON_EXPORT CollectiveCommunicationLib {
   // Destroy the communication group.
   virtual bool DestroyCommunicationGroup(const std::string &group_name);
 
+  // Switch network interface card due to borrowing or switchback
+  virtual bool CommSwitchNic(const std::vector<uint32_t> &global_ranks, const std::vector<bool> &use_backup) {
+    return true;
+  }
+
   // Get the inner communicator name of the specified group.
   virtual std::string CommName(const std::string &group_name) { return ""; }
 

mindspore/include/mindspore/ccsrc/runtime/collective/communication_group.h

@@ -50,6 +50,11 @@ class BACKEND_COMMON_EXPORT CommunicationGroup {
   // For NCCL group, it returns a pointer to 'ncclUniqueId'. For HCCL group, it returns a pointer to 'HcclRootInfo'.
   virtual void *GenerateRootInfo(size_t *root_info_size) { return nullptr; }
 
+  virtual bool SetGlobalCommInfo(uint32_t master_ip, uint32_t master_port, uint32_t total_rank_size, uint32_t node_rank,
+                                 uint32_t local_rank_size) {
+    return true;
+  }
+
   // Get group or global rank for the given rank.
   uint32_t GetGroupRank(uint32_t global_rank);
   uint32_t GetGlobalRank(uint32_t group_rank);

mindspore/include/mindspore/ccsrc/runtime/collective/dummy_collective_communication_lib.h

@@ -66,7 +66,8 @@ class BACKEND_COMMON_EXPORT DummyCollectiveCommunicationLib : public CollectiveC
   bool CreateDeviceCommunicationGroup(const std::string &group_name, const std::vector<uint32_t> &group_ranks) override;
 
   bool CreateCommunicationGroup(const std::string &group_name, const std::vector<uint32_t> &group_ranks,
-                                uint32_t local_group_rank, uint32_t local_group_size) override;
+                                uint32_t local_group_rank, uint32_t local_group_size,
+                                const GroupOptions &config = {}) override;
 
   uint32_t GetRankId(const std::string &group_name) override;
 

mindspore/include/mindspore/ccsrc/runtime/device/device_address_utils.h

@@ -35,16 +35,13 @@ struct is_optional<std::optional<T>> : public std::true_type {};
 
 namespace mindspore {
 using device::DeviceContext;
+using KernelTensorPtr = kernel::KernelTensorPtr;
 namespace runtime {
 // Extract the methods related to DeviceAddress in GraphCompiler to the DeviceAddressUtils class.
 class BACKEND_EXPORT DeviceAddressUtils {
  public:
-  static void CreateKernelTensor(const device::DeviceAddressPtr &device_address, const tensor::BaseTensor *tensor);
-  static void CreateKernelTensor(const device::DeviceAddressPtr &device_address, const AbstractBasePtr &abs);
-  static void CreateKernelTensor(const ValuePtr &input_value);
-  static void CreateKernelTensor(const tensor::TensorPtr &input_tensor);
   static void CopyNoneTensorDataToDevice(const device::DeviceContext *device_context,
-                                         const device::DeviceAddressPtr &device_address, const ShapeVector &shape = {});
+                                         const KernelTensorPtr &kernel_tensor, const ShapeVector &shape = {});
   static void CreateParameterDeviceAddress(const DeviceContext *device_context, const KernelGraphPtr &graph);
   static device::DeviceAddressPtrList CreateDeviceAddressForTensorValue(const DeviceContext *device_context,
                                                                         const ValuePtr &node_value, size_t output_idx,
@@ -53,8 +50,9 @@ class BACKEND_EXPORT DeviceAddressUtils {
   static void CreateKernelOutputDeviceAddress(const DeviceContext *device_context, const KernelGraphPtr &graph,
                                               bool is_gradient_out);
 
-  static std::vector<device::DeviceAddressPtr> CreateGraphOutputDeviceAddress(
-    const OpCompilerInfoPtr &op_compiler_info, const abstract::AbstractBasePtr &out_abstract, size_t stream_id);
+  static std::vector<KernelTensorPtr> CreateGraphOutputKernelTensor(const OpCompilerInfoPtr &op_compiler_info,
+                                                                    const abstract::AbstractBasePtr &out_abstract,
+                                                                    size_t stream_id);
 
   static void CreateKernelWorkspaceDeviceAddress(const DeviceContext *device_context, const KernelGraphPtr &graph);
   static void CreateDeviceAddressByMapTensorNode(const DeviceContext *device_context, const AnfNodePtr &node,
@@ -66,22 +64,22 @@ class BACKEND_EXPORT DeviceAddressUtils {
   static void UpdateDeviceAddressForRefNodeForSingleOp(const KernelGraphPtr &graph);
   static void UpdateDeviceAddressMonadFlag(const session::AnfWithOutIndex &cur_pair,
                                            const session::AnfWithOutIndex &origin_pair);
-  static device::DeviceAddressPtr CloneEmptyDeviceAddress(const device::DeviceAddressPtr &old_device_address,
-                                                          const DeviceContext *device_context);
+  static KernelTensorPtr CloneEmptyKernelTensor(const KernelTensorPtr &old_kernel_tensor,
+                                                const DeviceContext *device_context);
   static void CreateGraphOutputDeviceAddress(const DeviceContext *device_context, const KernelGraphPtr &graph);
   static size_t GetTensorDeviceSize(const DeviceContext *device_context, const AnfNodePtr &node,
                                     const ShapeVector &shape, const string &format, TypeId dtype, size_t output_index);
 
   // Overloading
   static void CreateInputTensorAddress(const DeviceContext *device_context, size_t stream_id, size_t index,
-                                       const tensor::BaseTensorPtr &tensor);
-  static void MallocForInput(const DeviceContext *device_context, const tensor::BaseTensorPtr &tensor, bool is_view);
-  static void MallocForInput(const DeviceContext *device_context, const std::optional<tensor::BaseTensorPtr> &val,
+                                       const tensor::TensorPtr &tensor);
+  static void MallocForInput(const DeviceContext *device_context, const tensor::TensorPtr &tensor, bool is_view);
+  static void MallocForInput(const DeviceContext *device_context, const std::optional<tensor::TensorPtr> &val,
                              bool is_view);
-  static void MallocForInput(const DeviceContext *device_context, const std::vector<tensor::BaseTensorPtr> &tensors,
+  static void MallocForInput(const DeviceContext *device_context, const std::vector<tensor::TensorPtr> &tensors,
                              bool is_view);
   static void CreateInputTensorAddress(const DeviceContext *device_context, size_t stream_id, size_t index,
-                                       const std::optional<tensor::BaseTensorPtr> &val);
+                                       const std::optional<tensor::TensorPtr> &val);
   template <typename T>
   static void CreateInputTensorAddress(const DeviceContext *device_context, size_t stream_id, size_t index,
                                        const std::vector<T> &inputs) {
@@ -90,24 +88,24 @@ class BACKEND_EXPORT DeviceAddressUtils {
     }
   }
 
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index,
-                                                     const tensor::BaseTensorPtr &tensor);
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index,
-                                                     const std::optional<tensor::BaseTensorPtr> &val);
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index,
-                                                     const ScalarPtr &scalar_value);
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index,
-                                                     const StringImmPtr &string_imm);
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index,
-                                                     const TypePtr &type_ptr);
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index,
+                                                 const tensor::TensorPtr &tensor);
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index,
+                                                 const std::optional<tensor::TensorPtr> &val);
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index,
+                                                 const ScalarPtr &scalar_value);
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index,
+                                                 const StringImmPtr &string_imm);
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index,
+                                                 const TypePtr &type_ptr);
   template <typename T>
-  static device::DeviceAddressPtr CreateInputAddress(const DeviceContext *device_context, size_t stream_id,
-                                                     const abstract::AbstractBasePtr &abs, size_t index, const T &t) {
+  static KernelTensorPtr CreateInputKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                 const abstract::AbstractBasePtr &abs, size_t index, const T &t) {
     MS_EXCEPTION_IF_NULL(device_context);
     auto tmp_abs = abs;
     if (abs == nullptr) {
@@ -116,43 +114,46 @@ class BACKEND_EXPORT DeviceAddressUtils {
     auto shape = tmp_abs->GetShape();
     auto type = tmp_abs->GetType();
     auto value = tmp_abs->GetValue();
-    auto kernel_tensor = std::make_shared<kernel::KernelTensor>(shape, type, value);
-    auto device_address = device_context->device_res_manager_->CreateDeviceAddress(kernel_tensor);
+    auto device_address = device_context->device_res_manager_->CreateDeviceAddress();
+    auto kernel_tensor = std::make_shared<kernel::KernelTensor>(device_address, shape, type, value, ShapeVector{});
     device_address->set_from_persistent_mem(true);
     device_address->set_new_ref_count(SIZE_MAX);
 
     if (device_address->GetPtr() == nullptr) {
-      CopyNoneTensorDataToDevice(device_context, device_address);
+      CopyNoneTensorDataToDevice(device_context, kernel_tensor);
     }
     MS_LOG(DEBUG) << "Create input " << tmp_abs->ToString() << " device address for " << index
                   << "th input, Shape: " << shape->ToString() << ", Type: " << type->ToString()
-                  << ", Value: " << (value ? value->ToString() : "nullptr") << " device address:" << device_address;
-    return device_address;
+                  << ", Value: " << (value ? value->ToString() : "nullptr") << " device address:" << device_address
+                  << ", kernel tensor: " << kernel_tensor.get();
+    return kernel_tensor;
   }
 
   static void CreateOutputTensorAddress(const DeviceContext *device_context, size_t stream_id,
-                                        const std::vector<tensor::BaseTensorPtr> &outputs);
+                                        const std::vector<tensor::TensorPtr> &outputs);
   static void CreateOutputTensorAddress(const DeviceContext *device_context, size_t stream_id,
-                                        const tensor::BaseTensorPtr &output_tensor, size_t size);
+                                        const tensor::TensorPtr &output_tensor, size_t size);
 
-  static void MallocForOutputs(const DeviceContext *device_context, const std::vector<tensor::BaseTensorPtr> &outputs);
+  static void MallocForOutputs(const DeviceContext *device_context, const std::vector<tensor::TensorPtr> &outputs);
 
   static device::DeviceAddressPtr CreateWorkspaceAddressWithoutKernelTensor(const DeviceContext *device_context,
                                                                             size_t stream_id,
                                                                             const size_t &workspace_size,
                                                                             bool no_exception = false);
 
-  static device::DeviceAddressPtr CreateWorkspaceAddress(const DeviceContext *device_context, size_t stream_id,
-                                                         const size_t &workspace_size);
+  static KernelTensorPtr CreateWorkspaceKernelTensor(const DeviceContext *device_context, size_t stream_id,
+                                                     const size_t &workspace_size);
 
   static void UpdateDeviceAddressHostInfoByNode(const device::DeviceAddressPtr &addr, const AnfNodePtr &node,
                                                 size_t output_idx);
   static device::DeviceAddressPtr CreateDeviceAddress(const DeviceContext *device_context,
-                                                      const tensor::BaseTensorPtr &tensor,
-                                                      const ShapeVector &real_shape, const size_t &stream_id);
+                                                      const tensor::TensorPtr &tensor, const ShapeVector &real_shape,
+                                                      const size_t &stream_id);
+  static KernelTensorPtr CreateKernelTensor(const DeviceContext *device_context, const tensor::TensorPtr &tensor,
+                                            const ShapeVector &real_shape, const size_t &stream_id);
 
   // Convert tensor to contiguous tensor.
-  static void ConvertContiguousTensorSync(const tensor::BaseTensorPtr &tensor);
+  static void ConvertContiguousTensorSync(const tensor::TensorPtr &tensor);
 
   // Convert old_device_address to contiguous device address.
   static device::DeviceAddressPtr ConvertContiguousDeviceAddress(const DeviceContext *device_context,
@@ -160,7 +161,7 @@ class BACKEND_EXPORT DeviceAddressUtils {
                                                                  bool is_sync);
 
   // Convert view tensor to contiguous tensor.
-  static tensor::BaseTensorPtr TensorContiguous(const tensor::BaseTensorPtr &tensor);
+  static tensor::TensorPtr TensorContiguous(const tensor::TensorPtr &tensor);
 
   template <typename... T>
   static void ProcessCrossStreamAddress(const std::string &op_name, const DeviceContext *device_context,
@@ -172,7 +173,8 @@ class BACKEND_EXPORT DeviceAddressUtils {
       return;
     }
 
-    auto &controller = device::HalResManager::GetInstance().GetMultiStreamController(device_context->DeviceName());
+    auto &controller =
+      device::HalResManager::GetInstance().GetMultiStreamController(device_context->device_context_key().device_name_);
     controller->Refresh();
     auto task_id_on_stream = controller->LaunchTaskIdOnStream(op_stream_id);
     MS_LOG(DEBUG) << "Launch stream_id:" << op_stream_id << ", task id:" << task_id_on_stream << ", op_name:" << op_name
@@ -190,7 +192,8 @@ class BACKEND_EXPORT DeviceAddressUtils {
       return;
     }
 
-    auto &controller = device::HalResManager::GetInstance().GetMultiStreamController(device_context->DeviceName());
+    auto &controller =
+      device::HalResManager::GetInstance().GetMultiStreamController(device_context->device_context_key().device_name_);
     controller->Refresh();
     auto task_id_on_stream = controller->LaunchTaskIdOnStream(op_stream_id);
     MS_LOG(DEBUG) << "Launch stream_id:" << op_stream_id << ", task id:" << task_id_on_stream << ", op_name:" << op_name
@@ -206,7 +209,7 @@ class BACKEND_EXPORT DeviceAddressUtils {
 
   static void GetCrossStreamAddressInfoFromInput(size_t op_stream_id,
                                                  std::vector<std::pair<uint32_t, void *>> *cross_stream_addresses,
-                                                 const tensor::BaseTensorPtr &tensor);
+                                                 const tensor::TensorPtr &tensor);
 
   static void GetCrossStreamAddressInfoFromInput(size_t op_stream_id,
                                                  std::vector<std::pair<uint32_t, void *>> *cross_stream_addresses,
@@ -229,7 +232,7 @@ class BACKEND_EXPORT DeviceAddressUtils {
   static void GetCrossStreamAddressInfo(size_t op_stream_id,
                                         std::vector<std::pair<uint32_t, void *>> *cross_stream_addresses,
                                         const std::vector<T> &inputs) {
-    if constexpr (!std::is_same_v<T, tensor::BaseTensorPtr> && !std::is_same_v<T, tensor::TensorPtr> &&
+    if constexpr (!std::is_same_v<T, tensor::TensorPtr> && !std::is_same_v<T, tensor::TensorPtr> &&
                   !std::is_same_v<T, mindspore::kernel::KernelTensor *> &&
                   !std::is_same_v<T, device::DeviceAddressPtr>) {
       return;
@@ -243,12 +246,14 @@ class BACKEND_EXPORT DeviceAddressUtils {
   static void GetCrossStreamAddressInfo(size_t op_stream_id,
                                         std::vector<std::pair<uint32_t, void *>> *cross_stream_addresses,
                                         const T &input) {
-    if constexpr (std::is_same_v<T, tensor::BaseTensorPtr> || std::is_same_v<T, tensor::TensorPtr> ||
+    if constexpr (std::is_same_v<T, tensor::TensorPtr> || std::is_same_v<T, tensor::TensorPtr> ||
                   std::is_same_v<T, mindspore::kernel::KernelTensor *> || std::is_same_v<T, device::DeviceAddressPtr>) {
       GetCrossStreamAddressInfoFromInput(op_stream_id, cross_stream_addresses, input);
     }
   }
 };
+
+void CheckAutoH2D(const DeviceContext *device_context, const tensor::TensorPtr &tensor);
 }  // namespace runtime
 }  // namespace mindspore
 #endif  // MINDSPORE_MINDSPORE_CCSRC_RUNTIME_GRAPH_SCHEDULER_COMMON_UTILS_H_

mindspore/include/mindspore/ccsrc/runtime/device/memory_scheduler.h

@@ -24,10 +24,11 @@
 #include <utility>
 #include "runtime/device/memory_offload_strategy.h"
 #include "runtime/device/res_manager/auto_mem_offload.h"
+#include "include/backend/visible.h"
 
 namespace mindspore {
 namespace device {
-class MemScheduler {
+class BACKEND_EXPORT MemScheduler {
  public:
   MemScheduler() = default;
   ~MemScheduler() = default;

mindspore/include/mindspore/ccsrc/runtime/device/move_to.h

@@ -16,9 +16,12 @@
 
 #ifndef MINDSPORE_MINDSPORE_CCSRC_RUNTIME_DEVICE_MOVE_TO_H
 #define MINDSPORE_MINDSPORE_CCSRC_RUNTIME_DEVICE_MOVE_TO_H
+
 #include <string>
+
 #include "ir/tensor.h"
 #include "include/backend/visible.h"
+
 namespace mindspore {
 namespace device {
 void BACKEND_EXPORT MoveTo(const tensor::TensorPtr &src_tensor, const tensor::TensorPtr &dst_tensor,