Add documentation and tests

This commit is contained in:
Frederic Bastien 2019-08-23 06:36:32 -07:00
parent 5fea6eecde
commit c77c165dcd
3 changed files with 28 additions and 1 deletions

View File

@ -1382,6 +1382,10 @@ For a more intuitive description, see the "Informal Description" section below.
| `indices_are_sorted` | `bool` | Whether the indices are |
: : : guaranteed to be sorted by :
: : : the caller. :
| `use_atomic` | `bool` | Whether to use atomic |
: : : operation for the update. To :
: : : use only when the the caller :
: : : guarante no duplicate indices :
For convenience, we label dimensions in the output array not in `offset_dims`
as `batch_dims`.
@ -1450,6 +1454,9 @@ If `indices_are_sorted` is set to true then XLA can assume that `start_indices`
are sorted (in ascending `start_index_map` order) by the user. If they are not
then the semantics is implementation defined.
If `use_atomic` is set to false then XLA will not use atomic operation. This
is only safe when there is no duplicate indices.
### Informal Description and Examples
Informally, every index `Out` in the output array corresponds to an element `E`

View File

@ -1529,7 +1529,8 @@ TEST_F(HloInstructionTest, StringifyScatter) {
/*inserted_window_dims=*/{},
/*scatter_dims_to_operand_dims=*/{0, 1, 2, 3, 4},
/*index_vector_dim=*/2),
/*indices_are_sorted=*/false));
/*indices_are_sorted=*/false,
/*use_atomic=*/true));
module->AddEntryComputation(builder.Build());
EXPECT_EQ(

View File

@ -934,6 +934,25 @@ ENTRY %Scatter (input_tensor: f32[50,49,48,47,46], scatter_indices: s64[10,9,8,7
ROOT %scatter = f32[50,49,48,47,46]{4,3,2,1,0} scatter(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, s64[10,9,8,7,5]{4,3,2,1,0} %scatter_indices, f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} %updates), update_window_dims={4,5,6,7,8}, inserted_window_dims={}, scatter_dims_to_operand_dims={0,1,2,3,4}, index_vector_dim=4, indices_are_sorted=true, to_apply=%add_F32.v3
}
)"
},
{
"AtomicScatter",
R"(HloModule StringifyAtomicScatter
%add_F32.v3 (lhs: f32[], rhs: f32[]) -> f32[] {
%lhs = f32[] parameter(0)
%rhs = f32[] parameter(1)
ROOT %add = f32[] add(f32[] %lhs, f32[] %rhs)
}
ENTRY %Scatter (input_tensor: f32[50,49,48,47,46], scatter_indices: s64[10,9,8,7,5], updates: f32[10,9,8,7,30,29,28,27,26]) -> f32[50,49,48,47,46] {
%input_tensor = f32[50,49,48,47,46]{4,3,2,1,0} parameter(0)
%scatter_indices = s64[10,9,8,7,5]{4,3,2,1,0} parameter(1)
%updates = f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} parameter(2)
ROOT %scatter = f32[50,49,48,47,46]{4,3,2,1,0} scatter(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, s64[10,9,8,7,5]{4,3,2,1,0} %scatter_indices, f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} %updates), update_window_dims={4,5,6,7,8}, inserted_window_dims={}, scatter_dims_to_operand_dims={0,1,2,3,4}, index_vector_dim=4, use_atomic=false, to_apply=%add_F32.v3
}
)"
},
{