mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
34 lines
1020 B
C++
34 lines
1020 B
C++
#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
|