| 1530 | } |
| 1531 | |
| 1532 | XlaOp XlaBuilder::CustomCall( |
| 1533 | const string& call_target_name, absl::Span<const XlaOp> operands, |
| 1534 | const Shape& shape, const string& opaque, |
| 1535 | absl::optional<absl::Span<const Shape>> operand_shapes_with_layout) { |
| 1536 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1537 | HloInstructionProto instr; |
| 1538 | if (absl::StartsWith(call_target_name, "$")) { |
| 1539 | return InvalidArgument( |
| 1540 | "Invalid custom_call_target \"%s\": Call targets that start with '$' " |
| 1541 | "are reserved for internal use.", |
| 1542 | call_target_name); |
| 1543 | } |
| 1544 | *instr.mutable_shape() = shape.ToProto(); |
| 1545 | instr.set_custom_call_target(call_target_name); |
| 1546 | instr.set_backend_config(opaque); |
| 1547 | if (operand_shapes_with_layout.has_value()) { |
| 1548 | if (!LayoutUtil::HasLayout(shape)) { |
| 1549 | return InvalidArgument( |
| 1550 | "Result shape must have layout for custom call with constrained " |
| 1551 | "layout."); |
| 1552 | } |
| 1553 | if (operands.size() != operand_shapes_with_layout->size()) { |
| 1554 | return InvalidArgument( |
| 1555 | "Must specify a shape with layout for each operand for custom call " |
| 1556 | "with constrained layout; given %d shapes, expected %d", |
| 1557 | operand_shapes_with_layout->size(), operands.size()); |
| 1558 | } |
| 1559 | instr.set_constrain_layout(true); |
| 1560 | int64 operand_num = 0; |
| 1561 | for (const Shape& operand_shape : *operand_shapes_with_layout) { |
| 1562 | if (!LayoutUtil::HasLayout(operand_shape)) { |
| 1563 | return InvalidArgument( |
| 1564 | "No layout specified for operand %d for custom call with " |
| 1565 | "constrained layout.", |
| 1566 | operand_num); |
| 1567 | } |
| 1568 | *instr.add_operand_shapes_with_layout() = operand_shape.ToProto(); |
| 1569 | ++operand_num; |
| 1570 | } |
| 1571 | } |
| 1572 | return AddInstruction(std::move(instr), HloOpcode::kCustomCall, operands); |
| 1573 | }); |
| 1574 | } |
| 1575 | |
| 1576 | XlaOp XlaBuilder::Transpose(XlaOp operand, |
| 1577 | absl::Span<const int64> permutation) { |