Body of the while loop that performs the scatter operation using other HLOs.
| 217 | |
| 218 | // Body of the while loop that performs the scatter operation using other HLOs. |
| 219 | static StatusOr<std::vector<HloInstruction*>> ScatterLoopBody( |
| 220 | HloInstruction* scatter, HloInstruction* induction_var, |
| 221 | const std::vector<HloInstruction*>& loop_state) { |
| 222 | const ScatterDimensionNumbers& dim_numbers = |
| 223 | scatter->scatter_dimension_numbers(); |
| 224 | CHECK_EQ(loop_state.size(), 3); |
| 225 | HloInstruction* operand = loop_state[0]; |
| 226 | HloInstruction* scatter_indices = loop_state[1]; |
| 227 | HloInstruction* updates = loop_state[2]; |
| 228 | |
| 229 | bool has_scalar_indices = scatter_indices->shape().dimensions_size() == 1; |
| 230 | |
| 231 | // Build a vector form of the induction variable of the while loop. |
| 232 | HloInstruction* induction_var_as_vector = |
| 233 | MakeBroadcastHlo(induction_var, /*broadcast_dimensions=*/{}, |
| 234 | /*result_shape_bounds=*/{1}); |
| 235 | |
| 236 | // Pick the index to scatter from scatter_indices based on the induction_var |
| 237 | // and transform that to an index into the `operand` space. |
| 238 | HloInstruction* index_vector; |
| 239 | if (has_scalar_indices) { |
| 240 | TF_ASSIGN_OR_RETURN( |
| 241 | index_vector, |
| 242 | MakeDynamicSliceHlo(scatter_indices, induction_var_as_vector, {1})); |
| 243 | } else { |
| 244 | TF_ASSIGN_OR_RETURN( |
| 245 | HloInstruction * index_into_scatter_indices, |
| 246 | PadVectorWithZeros(induction_var_as_vector, |
| 247 | /*zeros_to_prepend=*/0, /*zeros_to_append=*/1)); |
| 248 | int index_vector_size = scatter_indices->shape().dimensions(1); |
| 249 | TF_ASSIGN_OR_RETURN( |
| 250 | HloInstruction * index_vector_2d, |
| 251 | MakeDynamicSliceHlo(scatter_indices, index_into_scatter_indices, |
| 252 | {1, index_vector_size})); |
| 253 | TF_ASSIGN_OR_RETURN(index_vector, |
| 254 | ElideDegenerateDims(index_vector_2d, {0})); |
| 255 | } |
| 256 | TF_ASSIGN_OR_RETURN( |
| 257 | HloInstruction * scatter_slice_start, |
| 258 | ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, |
| 259 | operand->shape().dimensions_size())); |
| 260 | |
| 261 | // Extract the slice to be used to update from `updates` tensor for the |
| 262 | // induction_var corresponding to this iteration of the while loop. |
| 263 | TF_ASSIGN_OR_RETURN( |
| 264 | HloInstruction * index_into_updates, |
| 265 | PadVectorWithZeros( |
| 266 | induction_var_as_vector, /*zeros_to_prepend=*/0, |
| 267 | /*zeros_to_append=*/updates->shape().dimensions_size() - 1)); |
| 268 | std::vector<int64> update_slice_bounds(updates->shape().dimensions().begin(), |
| 269 | updates->shape().dimensions().end()); |
| 270 | update_slice_bounds[0] = 1; |
| 271 | TF_ASSIGN_OR_RETURN( |
| 272 | HloInstruction * update_slice, |
| 273 | MakeDynamicSliceHlo(updates, index_into_updates, update_slice_bounds)); |
| 274 | TF_ASSIGN_OR_RETURN(HloInstruction * update_slice_for_scatter, |
| 275 | ElideDegenerateDims(update_slice, {0})); |
| 276 | TF_ASSIGN_OR_RETURN( |
no test coverage detected