ops.org

Better current

 class  CommunicationOp << (I, #9922DD) abstract >>

 CommunicationOp <|-- BaseSendOp
 CommunicationOp <|-- BaseRecvOp

 CommunicationOp : __init__([CommunicationOp] paired)
 /' CommunicationOp : is_communication_op() '/
 CommunicationOp : Dict metadata
 /' 
  ' CommunicationOp : String metadata['device']
  ' CommunicationOp : String metadata['device_id']
  ' CommunicationOp : String metadata['transformer']
  ' CommunicationOp : String metadata['host_transformer']
  ' CommunicationOp : String metadata['marker']
  ' CommunicationOp : String metadata['parallel']
  '/
 CommunicationOp : [CommunicationOp] paired
 CommunicationOp : [CommunicationOp] PartionFn(Axes axes)
 CommunicationOp : [Value] ReductionFn(Values)

class AbstractCommunicationOpFactory << abstract >>
AbstractCommunicationOpFactory : CommunicationOpFactory getFactory()
AbstractCommunicationOpFactory <|-- CommunicationOpFactory
AbstractCommunicationOpFactory <|-- CommunicationOpFactoryDecorator

class CommunicationOpFactory
CommunicationOpFactory : ([CommunicationOp], [CommunicationOp]) getSendRecvPairs()
CommunicationOpFactory .. "instantiates" SendRecvPairs

class CommunicationOpFactoryDecorator
CommunicationOpFactoryDecorator : ([AbstractCommunicationOp], [AbstractCommunicationOp]) getSendRecvPairs()
/' 
 ' CommunicationOpFactoryDecorator <|-- GPUFactory
 ' CommunicationOpFactoryDecorator <|-- CPUFactory
 '/
CommunicationOpFactoryDecorator <|-- GatherOpFactory
CommunicationOpFactoryDecorator <|-- ScatterOpFactory
CommunicationOpFactoryDecorator <|-- AllReduceOpFactory
CommunicationOpFactoryDecorator <|-- BroadcastOpFactory
CommunicationOpFactoryDecorator --* AbstractCommunicationOpFactory

class GatherOpFactory
GatherOpFactory : ([CommunicationOp], [CommunicationOp]) getSendRecvPairs()
/' GatherNodeFactory : CommNodeFactory getFactory '/

class ScatterOpFactory
ScatterOpFactory : ([CommunicationOp], [CommunicationOp]) getSendRecvPairs()
/' ScatterNodeFactory : CommNodeFactory getFactory '/

class SendRecvPairs
SendRecvPairs *-- CommunicationOp
 /' 
  ' class  SendOp << (I, #9922DD) abstract >>
  ' SendOp <|-- BaseSendOp
  ' SendOp : RecvOp recv_op
  ' 
  ' class  RecvOp << (I, #9922DD) abstract >>
  ' RecvOp <|-- BaseRecvOp
  ' RecvOp : SendOp send_op
  '/

 class BaseSendOp
 BaseSendOp : init([CommunicationOp] paired)
 BaseSendOp : send()

 class BaseRecvOp
 BaseRecvOp : init([CommunicationOp] paired)
 BaseRecvOp : recv()

 class CommPattern << (I, #9922DD) abstract >>
 CommunicationOp *-- CommPattern
 CommunicationOp <|-- CommPattern
 CommPattern <|-- Reduction
 CommPattern <|-- Partition
 /' 
  ' CommPattern <|-- ScatterOp
  ' CommPattern <|-- GatherOp
  '/

 /' 
  ' class ScatterOp
  ' ScatterOp : init()
  ' /\' ScatterOp : Axes slices '\/
  ' 
  ' class GatherOp
  ' GatherOp : init(from_id, from_axes)
  ' GatherOp : from_slices
  ' GatherOp : Int from_id
  ' GatherOp : Node send_nodes
  '/

 class Reduction
 Reduction : init(reduce_fn)

 class Partition
 Partition : init(partion_fn)
 /' Partition : [CommunicationOp] partition_fn(Axes) '/