//
// This confidential and proprietary software may be used only as
// authorised by a licensing agreement from ARM Limited
// (C) COPYRIGHT 2020 ARM Limited
// ALL RIGHTS RESERVED
// The entire notice above must be reproduced on all authorised
// copies and copies may only be made to the extent permitted
// by a licensing agreement from ARM Limited.

=== Scatter/Gather Operators

==== GATHER

Generate a tensor for which each element in the output is a subtensor of the values tensor based on the indices.
N is the number of batches, W the number of indices in each batch, K the range of each index and C the number data channels for each index.

*Arguments:*

|===
|Argument|Type|Name|Shape|Description

|Input|value_t*|values|[N,K,C]|3D value tensor
|Input|index_t*|indices|[N,W]|2D index tensor
|Output|value_t*|output|[N,W,C]|3D output tensor
|===

*Quantization Parameters:*

None

*Operation Function:*

[source,c]
----
for_each(0<=n<N, 0<=w<W, 0<=c<C) {
    index_t k = tensor_read<index_t>(indices, [N,W], [n,w])
    assert(0<=k && k<K)
    value_t value = tensor_read<value_t>(values, [N,K,C], [n, k, c])
    tensor_write<value_t>(output, [N,W,C], [n,w,c], value)
}
----

*Supported Data Types:*

|===
|Profile|Mode|index_t|value_t

|Any|signed 8|int32|aint8
|Any|signed 16|int32|int16
|Any|signed 32|int32|int32
|MI,MT|float|int32|float
|===

==== SCATTER

The values_out tensor is set to the values_in tensor with data modified as follows: data from the input tensor is inserted at the positions specified by the indices tensor.
N is the number of batches, W the number of indices in each batch, K the range of each index and C the number data channels for each index.
It is not permitted to repeat the same output index within a single SCATTER operation and so each output index occurs at most once.
In use cases that require multiple updates to the same output position, these must be decomposed into multiple SCATTER operations.

*Arguments:*

|===
|Argument|Type|Name|Shape|Description

|Input|value_t*|values_in|[N,K,C]|3D values in tensor
|Input|index_t*|indices|[N,W]|2D index tensor
|Input|value_t*|input|[N,W,C]|3D input tensor
|Output|value_t*|values_out|[N,K,C]|3D values out tensor
|===

*Quantization Parameters:*

None

*Operation Function:*

[source,c]
----

// The following array is used to check compliance that an output position
// is modified at most once.
bool output_modified[N,K,C];

// Copy the values_in tensor to the values_out tensor.
// Values not written by the scatter operation are unchanged in the output.
for_each(0<=n<N, 0<=k<K, 0<=c<C) {
    value_t value = tensor_read<value_t>(values_in, [N,K,C], [n,k,c])
    tensor_write<value_t>(values_out, [N,K,C], [n, k, c], value)
    output_modified[n,k,c]=false;
}

// Now perform the SCATTER operation, modifying the positions from the indices tensor
for_each(0<=n<N, 0<=w<W, 0<=c<C) {
    index_t k = tensor_read<index_t>(indices, [N,W], [n,w])
    assert(0<=k && k<K)
    assert(output_modified[n,k,c]==false);
    value_t value = tensor_read<value_t>(input, [N,W,C], [n,w,c])
    tensor_write<value_t>(values_out, [N,K,C], [n, k, c], value)
    output_modified[n,k,c]=true;
}
----

*Supported Data Types:*

|===
|Profile|Mode|index_t|value_t

|Any|signed 8|int32|aint8
|Any|signed 16|int32|int16
|Any|signed 32|int32|int32
|MI,MT|float|int32|float
|===