device2

@startuml

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 *-up- MemoryManager
MemoryManager <|.. GPUMemoryManager
MemoryManager <|.. AscendMemoryManager