| #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 |
