| 1703 | } |
| 1704 | |
| 1705 | Status HloEvaluator::HandleGather(HloInstruction* gather) { |
| 1706 | Literal result = Literal::CreateFromShape(gather->shape()); |
| 1707 | const Shape& shape = gather->shape(); |
| 1708 | const GatherDimensionNumbers& dim_numbers = |
| 1709 | gather->gather_dimension_numbers(); |
| 1710 | const Literal& operand = GetEvaluatedLiteralFor(gather->operand(0)); |
| 1711 | Literal reshaped_start_indices; |
| 1712 | TF_ASSIGN_OR_RETURN( |
| 1713 | const Literal& start_indices, |
| 1714 | ReshapedGatherIndices(dim_numbers.index_vector_dim(), |
| 1715 | GetEvaluatedLiteralFor(gather->operand(1)), |
| 1716 | &reshaped_start_indices)); |
| 1717 | |
| 1718 | // We iterate over the gather dimensions in the output shape in an outer loop |
| 1719 | // nest, and iterate over the window dimensions in the output shape in an |
| 1720 | // inner loop nest. |
| 1721 | |
| 1722 | ShapeUtil::IndexIterationSpace start_indices_iteration_space = |
| 1723 | IterationSpaceForOutputBatchIndices(shape, dim_numbers); |
| 1724 | ShapeUtil::IndexIterationSpace offset_indices_iteration_space = |
| 1725 | IterationSpaceForOutputOffsetIndices( |
| 1726 | shape.dimensions_size(), gather->gather_slice_sizes(), dim_numbers); |
| 1727 | |
| 1728 | // Scratch buffers that hold an index in the output shape and the |
| 1729 | // corresponding index in the input shape. |
| 1730 | std::vector<int64> input_index(operand.shape().dimensions_size()); |
| 1731 | std::vector<int64> output_index(gather->shape().dimensions_size()); |
| 1732 | std::vector<int64> input_index_clamped(operand.shape().dimensions_size()); |
| 1733 | |
| 1734 | OutputBatchIndexToInputIndex output_batch_index_to_input_index( |
| 1735 | &gather->gather_dimension_numbers(), /*input_shape=*/operand.shape(), |
| 1736 | /*output_shape=*/shape, &start_indices); |
| 1737 | OutputOffsetIndexToInputIndex output_offset_index_to_input_index( |
| 1738 | gather->gather_dimension_numbers(), /*input_shape=*/operand.shape(), |
| 1739 | /*output_shape=*/shape); |
| 1740 | |
| 1741 | const Shape& operand_shape = operand.shape(); |
| 1742 | if (ShapeUtil::IsZeroElementArray(operand_shape)) { |
| 1743 | evaluated_[gather] = std::move(result); |
| 1744 | return Status::OK(); |
| 1745 | } |
| 1746 | |
| 1747 | auto gather_inner_loop_body = |
| 1748 | [&](absl::Span<const int64> output_window_index, |
| 1749 | absl::Span<const int64> input_gather_index, |
| 1750 | absl::Span<const int64> output_gather_index) -> StatusOr<bool> { |
| 1751 | TF_ASSIGN_OR_RETURN( |
| 1752 | absl::Span<const int64> input_window_index, |
| 1753 | output_offset_index_to_input_index(output_window_index)); |
| 1754 | for (int i = 0, e = output_index.size(); i < e; i++) { |
| 1755 | output_index[i] = output_gather_index[i] + output_window_index[i]; |
| 1756 | DCHECK_LT(output_index[i], shape.dimensions(i)); |
| 1757 | } |
| 1758 | for (int i = 0, e = input_gather_index.size(); i < e; i++) { |
| 1759 | int64 output_dim = |
| 1760 | output_offset_index_to_input_index.input_dim_value_to_output_index(i); |
| 1761 | // If 'output_dim' is -1, it means 'i' is an elided window dim. This means |
| 1762 | // we set the iteration index to 0, so for the purpose of the following |
nothing calls this directly
no test coverage detected