This generates the body of the while that implements the main data movement behavior of gather using dynamic-slice and dynamic-update-slice.
| 148 | // This generates the body of the while that implements the main data movement |
| 149 | // behavior of gather using dynamic-slice and dynamic-update-slice. |
| 150 | static StatusOr<std::vector<HloInstruction*>> GatherLoopBody( |
| 151 | const HloInstruction& gather, HloInstruction* induction_var, |
| 152 | const std::vector<HloInstruction*>& incoming_loop_state) { |
| 153 | const GatherDimensionNumbers& dim_numbers = gather.gather_dimension_numbers(); |
| 154 | CHECK_EQ(incoming_loop_state.size(), 3); |
| 155 | HloInstruction* const operand = incoming_loop_state[0]; |
| 156 | HloInstruction* const start_indices = incoming_loop_state[1]; |
| 157 | HloInstruction* const output_accumulator = incoming_loop_state[2]; |
| 158 | |
| 159 | bool has_scalar_indices = start_indices->shape().dimensions_size() == 1; |
| 160 | CHECK_EQ(has_scalar_indices, |
| 161 | dim_numbers.index_vector_dim() == |
| 162 | gather.operand(1)->shape().dimensions_size()); |
| 163 | |
| 164 | HloInstruction* induction_var_as_vector = |
| 165 | MakeBroadcastHlo(induction_var, /*broadcast_dimensions=*/{}, |
| 166 | /*result_shape_bounds=*/{1}); |
| 167 | |
| 168 | HloInstruction* index_vector; |
| 169 | |
| 170 | if (has_scalar_indices) { |
| 171 | // In this case start_indices has rank 1 and induction_var_as_vector (of |
| 172 | // shape {1}) is an index into this rank 1 tensor. |
| 173 | TF_ASSIGN_OR_RETURN( |
| 174 | index_vector, |
| 175 | MakeDynamicSliceHlo(start_indices, induction_var_as_vector, {1})); |
| 176 | } else { |
| 177 | // In this case start_indices has rank 2 and induction_var_as_vector (of |
| 178 | // shape {1}) is an index into just the first dimension of this rank 2 |
| 179 | // tensor. |
| 180 | TF_ASSIGN_OR_RETURN( |
| 181 | HloInstruction * index_into_start_indices, |
| 182 | PadVectorWithZeros(induction_var_as_vector, |
| 183 | /*zeros_to_prepend=*/0, /*zeros_to_append=*/1)); |
| 184 | |
| 185 | int64 index_vector_size = start_indices->shape().dimensions(1); |
| 186 | TF_ASSIGN_OR_RETURN( |
| 187 | HloInstruction * index_vector_2d, |
| 188 | MakeDynamicSliceHlo(start_indices, index_into_start_indices, |
| 189 | {1, index_vector_size})); |
| 190 | |
| 191 | TF_ASSIGN_OR_RETURN(index_vector, |
| 192 | ElideDegenerateDims(index_vector_2d, {0})); |
| 193 | } |
| 194 | |
| 195 | TF_ASSIGN_OR_RETURN( |
| 196 | HloInstruction * gathered_slice_start, |
| 197 | ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, |
| 198 | operand->shape().dimensions_size())); |
| 199 | |
| 200 | TF_ASSIGN_OR_RETURN(HloInstruction * gathered_slice, |
| 201 | MakeDynamicSliceHlo(operand, gathered_slice_start, |
| 202 | gather.gather_slice_sizes())); |
| 203 | |
| 204 | TF_ASSIGN_OR_RETURN( |
| 205 | HloInstruction* const gathered_slice_with_dims_collapsed, |
| 206 | ElideDegenerateDims(gathered_slice, |
| 207 | AsInt64Slice(dim_numbers.collapsed_slice_dims()))); |
no test coverage detected