| 2163 | } |
| 2164 | |
| 2165 | XlaOp XlaBuilder::AllReduce(XlaOp operand, const XlaComputation& computation, |
| 2166 | absl::Span<const ReplicaGroup> replica_groups, |
| 2167 | const absl::optional<ChannelHandle>& channel_id, |
| 2168 | const absl::optional<Shape>& shape_with_layout) { |
| 2169 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2170 | HloInstructionProto instr; |
| 2171 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 2172 | std::vector<const Shape*> operand_shapes; |
| 2173 | std::vector<XlaOp> operands; |
| 2174 | if (operand_shape->IsTuple()) { |
| 2175 | if (operand_shape->tuple_shapes_size() == 0) { |
| 2176 | return Unimplemented("0 element tuple AllReduce is not supported"); |
| 2177 | } |
| 2178 | for (int64 i = 0; i < operand_shape->tuple_shapes_size(); ++i) { |
| 2179 | if (operand_shape->tuple_shapes(i).element_type() != |
| 2180 | operand_shape->tuple_shapes(0).element_type()) { |
| 2181 | return Unimplemented( |
| 2182 | "All the shapes of a tuple input of AllReduce must have the same " |
| 2183 | "element type"); |
| 2184 | } |
| 2185 | operand_shapes.push_back(&operand_shape->tuple_shapes(i)); |
| 2186 | operands.push_back(GetTupleElement(operand, i)); |
| 2187 | } |
| 2188 | } else { |
| 2189 | operand_shapes.push_back(operand_shape); |
| 2190 | operands.push_back(operand); |
| 2191 | } |
| 2192 | |
| 2193 | TF_ASSIGN_OR_RETURN(Shape inferred_shape, |
| 2194 | ShapeInference::InferAllReduceShape(operand_shapes)); |
| 2195 | if (shape_with_layout) { |
| 2196 | if (!LayoutUtil::HasLayout(*shape_with_layout)) { |
| 2197 | return InvalidArgument("shape_with_layout must have the layout set: %s", |
| 2198 | shape_with_layout->ToString()); |
| 2199 | } |
| 2200 | if (!ShapeUtil::Compatible(*shape_with_layout, *operand_shape)) { |
| 2201 | return InvalidArgument( |
| 2202 | "Provided shape_with_layout must be compatible with the " |
| 2203 | "operand shape: %s vs %s", |
| 2204 | shape_with_layout->ToString(), operand_shape->ToString()); |
| 2205 | } |
| 2206 | instr.set_constrain_layout(true); |
| 2207 | if (operand_shape->IsTuple() && !inferred_shape.IsTuple()) { |
| 2208 | // For a single-element tuple, take the tuple element shape. |
| 2209 | TF_RET_CHECK(shape_with_layout->tuple_shapes_size() == 1); |
| 2210 | *instr.mutable_shape() = shape_with_layout->tuple_shapes(0).ToProto(); |
| 2211 | } else { |
| 2212 | *instr.mutable_shape() = shape_with_layout->ToProto(); |
| 2213 | } |
| 2214 | } else { |
| 2215 | *instr.mutable_shape() = inferred_shape.ToProto(); |
| 2216 | } |
| 2217 | |
| 2218 | for (const ReplicaGroup& group : replica_groups) { |
| 2219 | *instr.add_replica_groups() = group; |
| 2220 | } |
| 2221 | |
| 2222 | if (channel_id.has_value()) { |
no test coverage detected