1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifdef TENSORFLOW_USE_MPI 17 18 #include "tensorflow/core/framework/op.h" 19 #include "tensorflow/core/framework/shape_inference.h" 20 21 namespace tensorflow { 22 namespace contrib { 23 namespace mpi_collectives { 24 25 REGISTER_OP("MPIInit").Doc(R"doc( 26 Initialize MPI for the current process. 27 28 If this is run on a GPU, then that GPU must be used for all future MPI 29 operations. If it is run on CPU, then all future MPI operations must also 30 run on CPU. 31 )doc"); 32 33 REGISTER_OP("MPISize") 34 .Output("size: int32") __anon107d6f1b0102(shape_inference::InferenceContext* c) 35 .SetShapeFn([](shape_inference::InferenceContext* c) { 36 c->set_output(0, c->Scalar()); 37 return Status::OK(); 38 }) 39 .Doc(R"doc( 40 Returns the number of running MPI processes. 41 42 More precisely, returns the number of MPI processes in the group associated 43 with the MPI_COMM_WORLD communicator. 44 45 size: Size of the MPI group. 46 )doc"); 47 48 REGISTER_OP("MPIRank") 49 .Output("rank: int32") __anon107d6f1b0202(shape_inference::InferenceContext* c) 50 .SetShapeFn([](shape_inference::InferenceContext* c) { 51 c->set_output(0, c->Scalar()); 52 return Status::OK(); 53 }) 54 .Doc(R"doc( 55 Returns the index of the current process in the MPI group. 56 57 More precisely, returns the rank of the calling process in the MPI_COMM_WORLD 58 communicator. 59 60 rank: Rank of the calling process. 61 )doc"); 62 63 REGISTER_OP("MPILocalRank") 64 .Output("rank: int32") __anon107d6f1b0302(shape_inference::InferenceContext* c) 65 .SetShapeFn([](shape_inference::InferenceContext* c) { 66 c->set_output(0, c->Scalar()); 67 return Status::OK(); 68 }) 69 .Doc(R"doc( 70 Returns the index of the current process in the node it is on. 71 72 More precisely, returns the rank of the calling process in communicator that 73 only spans the MPI processes running on that node. 74 75 rank: Rank of the calling process on the node it is on. 76 )doc"); 77 78 REGISTER_OP("MPIAllreduce") 79 .Attr("T: {int32, int64, float32}") 80 .Input("tensor: T") 81 .Output("sum: T") __anon107d6f1b0402(shape_inference::InferenceContext* c) 82 .SetShapeFn([](shape_inference::InferenceContext* c) { 83 c->set_output(0, c->input(0)); 84 return Status::OK(); 85 }) 86 .Doc(R"doc( 87 Perform an MPI Allreduce on a tensor. All other processes that do a reduction 88 on a tensor with the same name must have the same dimension for that tensor. 89 Tensors are reduced with other tensors that have the same node name for the 90 allreduce. 91 92 Arguments 93 tensor: A tensor to reduce. 94 95 Output 96 sum: A tensor with the same shape as `tensor`, summed across all 97 MPI processes. 98 )doc"); 99 100 REGISTER_OP("MPIAllgather") 101 .Attr("T: {int32, int64, float32}") 102 .Attr("S: {int64}") 103 .Input("tensor: T") 104 .Input("sizes: S") 105 .Output("gathered: T") __anon107d6f1b0502(shape_inference::InferenceContext* c) 106 .SetShapeFn([](shape_inference::InferenceContext* c) { 107 shape_inference::ShapeHandle output; 108 TF_RETURN_IF_ERROR( 109 c->ReplaceDim(c->input(0), 0, c->UnknownDim(), &output)); 110 c->set_output(0, output); 111 return Status::OK(); 112 }) 113 .Doc(R"doc( 114 Perform an MPI Allgather on a tensor. All other processes that do a gather on a 115 tensor with the same name must have the same rank for that tensor, and have the 116 same dimension on all but the first dimension. 117 118 Arguments 119 tensor: A tensor to gather. 120 sizes: A tensor containing the first-dimension sizes of tensors to be 121 gathered from other ranks 122 123 Output 124 gathered: A tensor with the same shape as `tensor` except for the first 125 dimension, which is the sum of dimensions in `sizes`. 126 )doc"); 127 128 } // namespace mpi_collectives 129 } // namespace contrib 130 } // namespace tensorflow 131 132 #endif // TENSORFLOW_USE_MPI 133