| 2612 | } |
| 2613 | |
| 2614 | StatusOr<XlaComputation> XlaBuilder::BuildConstantSubGraph( |
| 2615 | XlaOp root_op, bool dynamic_dimension_is_minus_one) { |
| 2616 | TF_ASSIGN_OR_RETURN(bool is_constant, IsConstant(root_op)); |
| 2617 | if (!is_constant) { |
| 2618 | auto op_status = LookUpInstruction(root_op); |
| 2619 | string op_string = |
| 2620 | op_status.ok() ? op_status.ValueOrDie()->name() : "<unknown operation>"; |
| 2621 | return InvalidArgument( |
| 2622 | "Operand to BuildConstantSubGraph depends on a parameter.\n\n" |
| 2623 | " op requested for constant subgraph: %s\n\n" |
| 2624 | "This is an internal error that typically happens when the XLA user " |
| 2625 | "(e.g. TensorFlow) is attempting to determine a value that must be a " |
| 2626 | "compile-time constant (e.g. an array dimension) but it is not capable " |
| 2627 | "of being evaluated at XLA compile time.\n\n" |
| 2628 | "Please file a usability bug with the framework being used (e.g. " |
| 2629 | "TensorFlow).", |
| 2630 | op_string); |
| 2631 | } |
| 2632 | |
| 2633 | TF_ASSIGN_OR_RETURN(const HloInstructionProto* root, |
| 2634 | LookUpInstruction(root_op)); |
| 2635 | |
| 2636 | HloComputationProto entry; |
| 2637 | SetProtoIdAndName(&entry, StrCat(name_, "_compute_constant"), kNameSeparator, |
| 2638 | GetNextId()); |
| 2639 | entry.set_root_id(root->id()); |
| 2640 | ProgramShapeProto* program_shape = entry.mutable_program_shape(); |
| 2641 | *program_shape->mutable_result() = root->shape(); |
| 2642 | |
| 2643 | // We use std::set to keep the instruction ids in ascending order (which is |
| 2644 | // also a valid dependency order). The related ops will be added to the |
| 2645 | // subgraph in the same order. |
| 2646 | std::set<int64> related_ops; |
| 2647 | absl::flat_hash_set<int64> related_calls; // Related computations. |
| 2648 | std::queue<int64> worklist; |
| 2649 | worklist.push(root->id()); |
| 2650 | related_ops.insert(root->id()); |
| 2651 | while (!worklist.empty()) { |
| 2652 | int64 handle = worklist.front(); |
| 2653 | worklist.pop(); |
| 2654 | TF_ASSIGN_OR_RETURN(const HloInstructionProto* instr_proto, |
| 2655 | LookUpInstructionByHandle(handle)); |
| 2656 | |
| 2657 | if (instr_proto->opcode() == |
| 2658 | HloOpcodeString(HloOpcode::kGetDimensionSize)) { |
| 2659 | // At this point, BuildConstantSubGraph should never encounter a |
| 2660 | // GetDimensionSize with a dynamic dimension. IsConstant check would have |
| 2661 | // failed at the beginning of this function. |
| 2662 | // |
| 2663 | // Replace GetDimensionSize with a Constant representing the static bound |
| 2664 | // of the shape. |
| 2665 | int64 dimension = instr_proto->dimensions(0); |
| 2666 | int64 operand_handle = instr_proto->operand_ids(0); |
| 2667 | TF_ASSIGN_OR_RETURN(const HloInstructionProto* operand_proto, |
| 2668 | LookUpInstructionByHandle(operand_handle)); |
| 2669 | |
| 2670 | int32 constant_dimension_size = -1; |
| 2671 | if (!(operand_proto->shape().is_dynamic_dimension(dimension) && |
nothing calls this directly
no test coverage detected