mindspore device类图

@startuml

class KernelRuntimeManager {
  Register(const std::string &device_name, KernelRuntimeCreator &&runtime_creator);
  KernelRuntime *GetKernelRuntime(const std::string &device_name, uint32_t device_id);
  KernelRuntime *GetSingleKernelRuntime(const std::string &device_name, uint32_t device_id);
  void ClearRuntimeResource();
  void ClearGraphResource(uint32_t graph_id);
  std::map<std::string, std::shared_ptr<KernelRuntime> > runtime_map_;
  std::map<std::string, KernelRuntimeCreator> runtime_creators_;
}
class KernelRuntimeRegistrar


interface KernelRuntime {
  Init()
  AssignMemory(session::KernelGraph *graph)
  RunOpAssignMemory(const std::vector<tensor::TensorPtr> &input_tensors, session::KernelGraph *graph)
  RunOpClearMemory(const session::KernelGraph *graph)
  Run(session::KernelGraph *graph)
  DumpData(session::KernelGraph *graph)
  LoadData(session::KernelGraph *graph, Debugger *debugger)
  RunTask(const session::KernelGraph *graph)
  GenTask(const session::KernelGraph *graph)
  LaunchKernel(const session::KernelGraph *graph)
  AssignStaticMemoryInput(const session::KernelGraph *graph)
  AssignStaticMemoryValueNode(session::KernelGraph *graph)
  ClearGraphRuntimeResource(uint32_t graph_id)
  SyncStream()
  virtual bool LoadTask(const session::KernelGraph *graph);
  virtual void ReleaseDeviceRes() {}
  void set_device_id(uint32_t device_id) { device_id_ = device_id; }
}

class GPUKernelRuntime {
  Init
  ReleaseDeviceRes
  AssignMemory
  Run
}

class CPUKernelRuntime {
  Init
  Run
  AssignKernelAddress
  BindInputOutput
  IncreaseSummaryRefCount
  DecreaseSummaryRefCount
}

class AscendKernelRuntime {
  Init
  DumpData
  LoadData
  GenTask
  RunTask
  LoadTask
  ClearGraphRuntimeResource
  SyncStream
}

class DeviceAddress {
   SyncDeviceToHost
   SyncHostToDevice
   DeviceType
}
class GPUDeviceAddress
class AscendDeviceAddress
class CPUDeviceAddress
class GPUDeviceManager {
  GetInstance
  InitDevice
  ReleaseDevice
  device_count
  set_cur_device_id
  cur_device_id
  CreateStream
  SyncStream
  default_stream
  GetCudnnHandle
  GetCublasHandle
  CopyDeviceMemToHost
  CopyHostMemToDevice
  CopyDeviceMemToHostAsync
  CopyHostMemToDeviceAsync
}

class CudaDriver {
  AllocDeviceMem
  FreeDeviceMem
  AllocHostPinnedMem
  FreeHostPinnedMem
  CopyHostMemToDevice
  CopyDeviceMemToHost
  CopyHostMemToDeviceAsync
  CopyDeviceMemToHostAsync
  CreateStream
  DestroyStream
  SyncStream
  CreateEvent
  DestroyEvent
  RecordEvent
  SyncEvent
  QueryEvent
  device_count
  set_current_device
}

interface MemoryManager {
    MallocDeviceMemory
    FreeDeviceMemory
    ResetDynamicMemory
    MallocReusedDynamicMem
    MallocOutputMem
    MallocWorkSpaceMem
    MallocMem
    MallocMemFromMemPool
    FreeMemFromMemPool
    MallocContinuousMemFromMemPool
    GetCommonAlignSize
    GetCommunicationAlignSize
}

class GPUMemoryManager {
  MallocDeviceMemory
  FreeDeviceMemory
  MallocMemFromMemPool
  FreeMemFromMemPool
  MallocContinuousMemFromMemPool
}

class AscendMemoryManager {
    MallocDeviceMemory
    FreeDeviceMemory
    ResetDynamicMemory
    MallocMemFromMemPool
}


class DynamicMemPoolBestFit {
    AllocTensorMem
    AllocContinuousTensorMem
    FreeTensorMem
    ReleaseDeviceRes
    DumpDynamicMemPoolInfo
    total_mem_statistics
    used_mem_statistics
    used_mem_peak_statistics
    AllocDeviceMem
    FreeDeviceMem
    free_mem_size
    total_mem_size
}

class GPUMemoryAllocator {
    GetInstance
    Init
    CheckMaxDeviceMemory
    Finalize
    AllocBufferQueueMem
    AllocDeviceMem
    FreeDeviceMem
    free_mem_size
    total_mem_size
}

class AscendMemoryPool {
    GetInstance
    AllocDeviceMem
    FreeDeviceMem
    free_mem_size
    total_mem_size
    set_device_mem_pool_base
    set_device_mem_pool_size
}

KernelRuntime <|.. GPUKernelRuntime
KernelRuntime <|.. CPUKernelRuntime
KernelRuntime <|.. AscendKernelRuntime

KernelRuntime *-- MemoryManager
MemoryManager <|.. GPUMemoryManager
MemoryManager <|.. AscendMemoryManager

DynamicMemPoolBestFit <|-- GPUMemoryAllocator
DynamicMemPoolBestFit <|-- AscendMemoryPool

AscendMemoryPool <.. AscendMemoryManager
GPUMemoryAllocator <.. GPUMemoryManager


CudaDriver <.. GPUDeviceManager
GPUDeviceManager <.. GPUKernelRuntime
GPUDeviceManager <.. GPUDeviceAddress

DeviceAddress <|.. GPUDeviceAddress
DeviceAddress <|.. AscendDeviceAddress
DeviceAddress <|.. CPUDeviceAddress
@enduml