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