Correct wrapping code generation for DynamicUpdateSlice.
Currently if slice_index falls outside of the input, those values are not updated.
With wrapping behavior:
for each rank:
// Index for update:
index
// Indexes for input:
slice_start_index = slice_start_index % input_size //in-bounds
slice_limit_index = slice_start_index + update_size //may be out of bounds
// If in bounds:
if (slice_limit_index <= input_size) {
slice intersects = (index > slice_start_index) && (index <= slice_limit_index)
}
// If out of bounds:
else {
slice intersects = (index > slice_start_index) || (index <= slice_limit_index%input_size)
}
if (slice intersects)
for all ranks, compute index for intersection:
if slice_index was in bounds
intersection index [rank] = index - slice_start_index
else (out of bounds) && (index <= slice_limit_index%input_size)
intersection index [rank] = index + update_size - slice_start_index
return update [intersection index]
else
return input [index]
The above algorithms does not account for negative indexes given as inputs.
PiperOrigin-RevId: 174915327
Loading
Please sign in to comment