@startuml

class GpuQueue {
    RegisterRelease
    IsEmpty
    IsFull
    Push
    Front
    Pop
    Destroy
}
class BlockingQueue {
    Create
    RegisterRelease
    Push
    Front
    Pop
    Destroy
    std::shared_ptr<GpuQueue> queue_
}

class GpuBufferMgr {
    GetInstance()
    Create(device_id, channel_name, *addr, shape,capacity)
    Open(device_id, channel_name, shape,
    std::function<void(void *)> func)
    Push(handle, data, unsigned int timeout_in_sec);
    Front(unsigned int handle, void **addr, size_t *len);
    Pop(unsigned int handle);
    set_device_id(int device_id);
    Close(unsigned int handle) noexcept;
    IsInit()
    IsClosed()
    Destroy()
    CloseNotify()
    CloseConfirm()
    std::map<unsigned int, std::shared_ptr<BlockingQueue>> handle_queue_map_
    std::map<std::string, std::shared_ptr<BlockingQueue>> name_queue_map_

}

interface KernelMod {
  GetInputSizeList()
  GetOutputSizeList()
  GetWorkspaceSizeList()
  bool Launch(inputs, workspace,
              outputs, stream_ptr)
  GenParameters()
}
class AscendKernelMod {
    GenTask
    block_dim
    stream_id
}
class AicpuOpKernelMod{
    SetInputSizeList
    SetOutputSizeList
    SetWorkspaceSizeList
    GetInputSizeList
    GetOutputSizeList
    GetWorkspaceSizeList
    Launch
    GenTask
    SetInputList
    SetOutputList
    SetAnfNode
    SetNodeDef
    SetNodeName
    CreateCpuKernelInfo(inputs, outputs)    
}

class TbeKernelMod{
    SetInputSizeList
    SetOutputSizeList
    SetWorkspaceSizeList
    GetInputSizeList
    GetOutputSizeList
    GetWorkspaceSizeList
    Launch
    GenTask
    GenParameters
}
class AkgKernelMod
class HcclKernel

class RtKernel
class ProfilingKernelMod
class ProfilingKernelMod
class SendKernel
class RecvKernel
class StreamActiveKernel
class StreamSwitchKernel
class MemCpyAsyncKernel
class AssignKernel
class LabelGotoKernel
class LabelSetKernel
class LabelSwitchKernel
  
class CPUKernel
class GpuKernel {
  Init(const CNodePtr &kernel_node)
}
class AddNGpuFwdKernel{
    Init()
    Launch()
}
class xxxKernel {
    Init()
    Launch()
}
class Conv2dGpuFwdKernel {
    Init()
    Launch()
}


KernelMod <|.. GpuKernel
KernelMod <|.. AscendKernelMod
KernelMod <|.. CPUKernel
GpuKernel <|-- AddNGpuFwdKernel
GpuKernel <|-- xxxKernel
GpuKernel <|-- Conv2dGpuFwdKernel

AscendKernelMod <|-- AicpuOpKernelMod
AscendKernelMod <|-- TbeKernelMod
AscendKernelMod <|-- AkgKernelMod
AscendKernelMod <|-- HcclKernel
AscendKernelMod <|-- RtKernel
RtKernel <|-- ProfilingKernelMod
RtKernel <|-- SendKernel
RtKernel <|-- RecvKernel
RtKernel <|-- StreamActiveKernel
RtKernel <|-- StreamSwitchKernel
RtKernel <|-- MemCpyAsyncKernel
RtKernel <|-- AssignKernel
RtKernel <|-- LabelGotoKernel
RtKernel <|-- LabelSetKernel
RtKernel <|-- LabelSwitchKernel


class GPUDeviceManager
class GPUKernelRuntime
class DatasetInitKernel
class DatasetIteratorKernel

GpuBufferMgr *-right- BlockingQueue
BlockingQueue o-right- GpuQueue
GpuBufferMgr <.down. GPUDeviceManager
GpuBufferMgr <.down. GPUKernelRuntime
GpuBufferMgr <.down. DatasetInitKernel
GpuBufferMgr <.. DatasetIteratorKernel

@enduml