| 1422 | } |
| 1423 | |
| 1424 | void XlaBuilder::Outfeed(XlaOp operand, const Shape& shape_with_layout, |
| 1425 | const string& outfeed_config) { |
| 1426 | ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1427 | HloInstructionProto instr; |
| 1428 | |
| 1429 | *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); |
| 1430 | |
| 1431 | // Check and set outfeed shape. |
| 1432 | if (!LayoutUtil::HasLayout(shape_with_layout)) { |
| 1433 | return InvalidArgument("Given shape to Outfeed must have a layout"); |
| 1434 | } |
| 1435 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 1436 | if (!ShapeUtil::Compatible(*operand_shape, shape_with_layout)) { |
| 1437 | return InvalidArgument( |
| 1438 | "Outfeed shape %s must be compatible with operand shape %s", |
| 1439 | ShapeUtil::HumanStringWithLayout(shape_with_layout), |
| 1440 | ShapeUtil::HumanStringWithLayout(*operand_shape)); |
| 1441 | } |
| 1442 | *instr.mutable_outfeed_shape() = shape_with_layout.ToProto(); |
| 1443 | |
| 1444 | instr.set_outfeed_config(outfeed_config); |
| 1445 | |
| 1446 | // Outfeed takes a token as its second operand. Generate the token to pass |
| 1447 | // to the outfeed. |
| 1448 | HloInstructionProto token_instr; |
| 1449 | *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); |
| 1450 | TF_ASSIGN_OR_RETURN(XlaOp token, AddInstruction(std::move(token_instr), |
| 1451 | HloOpcode::kAfterAll, {})); |
| 1452 | |
| 1453 | TF_RETURN_IF_ERROR( |
| 1454 | AddInstruction(std::move(instr), HloOpcode::kOutfeed, {operand, token}) |
| 1455 | .status()); |
| 1456 | |
| 1457 | // The outfeed instruction produces a token. However, existing users expect |
| 1458 | // a nil shape (empty tuple). This should only be relevant if the outfeed is |
| 1459 | // the root of a computation. |
| 1460 | // TODO(b/80000000): Remove this when clients have been updated to handle |
| 1461 | // tokens. |
| 1462 | HloInstructionProto tuple_instr; |
| 1463 | *tuple_instr.mutable_shape() = ShapeUtil::MakeNil().ToProto(); |
| 1464 | |
| 1465 | // The dummy tuple should have no sharding. |
| 1466 | { |
| 1467 | XlaScopedShardingAssignment scoped_sharding(this, OpSharding()); |
| 1468 | TF_ASSIGN_OR_RETURN( |
| 1469 | XlaOp empty_tuple, |
| 1470 | AddInstruction(std::move(tuple_instr), HloOpcode::kTuple, {})); |
| 1471 | return empty_tuple; |
| 1472 | } |
| 1473 | }); |
| 1474 | } |
| 1475 | |
| 1476 | XlaOp XlaBuilder::OutfeedWithToken(XlaOp operand, XlaOp token, |
| 1477 | const Shape& shape_with_layout, |
no test coverage detected