1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
|
#include "caffe2/mpi/mpi_ops.h"
namespace caffe2 {
OPERATOR_SCHEMA(MPICreateCommonWorld)
.NumInputs(0)
.NumOutputs(1);
OPERATOR_SCHEMA(MPIBroadcast)
.NumInputs(2)
.NumOutputs(1)
.EnforceInplace({{1, 0}});
OPERATOR_SCHEMA(MPIReduce)
.NumInputs(2)
.NumOutputs(1);
OPERATOR_SCHEMA(MPIAllgather)
.NumInputs(2)
.NumOutputs(1);
OPERATOR_SCHEMA(MPIAllreduce)
.NumInputs(2)
.NumOutputs(1)
.AllowInplace({{1, 0}});
OPERATOR_SCHEMA(MPISendTensor);
OPERATOR_SCHEMA(MPIReceiveTensor);
REGISTER_CPU_OPERATOR(MPICreateCommonWorld, MPICreateCommonWorldOp<CPUContext>);
REGISTER_CPU_OPERATOR(MPIBroadcast, MPIBroadcastOp<CPUContext>);
REGISTER_CPU_OPERATOR(MPIReduce, MPIReduceOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(MPIAllgather, MPIAllgatherOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(MPIAllreduce, MPIAllreduceOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(MPISendTensor, MPISendTensorOp<CPUContext>);
REGISTER_CPU_OPERATOR(MPIReceiveTensor, MPIReceiveTensorOp<CPUContext>);
} // namespace caffe2
|
