# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""ScatterElements operator"""
from tvm import te
from tvm import tir
from . import utils
from .math import cast
[文档]
def scatter_elements(data, indices, updates, axis=0, reduction="update"):
"""Scatter elements from updates to corresponding indices of copied data.
Data, indices, updates and output have the same shape.
Indices can not have duplicates (if idx1 != idx2, then indices[idx1] != indices[idx2])
if reduction == "update".
.. code-block::
output[indices[i][j]][j] = f(output[indices[i][j]][j], updates[i][j]) if axis = 0
output[i][indices[i][j]] = f(output[i][indices[i][j]], updates[i][j]) if axis = 1
where the update function f is determined by the reduction.
Five types of the function are supported: "update", "add", "mul", "min" and "max" (see below)
Parameters
----------
data : tvm.te.Tensor
The source array.
indices : tvm.te.Tensor
The indices of the values to extract.
updates : tvm.te.Tensor
The updates to apply at the Indices
axis : optional, int
The axis to scatter on. It is zero by default.
reduction : optional, string
The update mode for the algorithm, either "update", "add", "mul", "min" or "max"
If update, the update values will replace the input data
If add, the update values will be added to the input data
If mul, the input data will be multiplied on the update values
If mean, the input data will be mean between the update values and the input data
If min, there is choice of minimal between the update values and the input data
If max, there is choice of maximal between the update values and the input data
It is "update" by default
Returns
-------
ret : tvm.te.Tensor
"""
if not isinstance(axis, int):
axis = utils.get_const_int(axis)
# Prepare ranges and strides
shape = data.shape
if axis < 0:
axis = len(shape) + axis
axis_range = cast(shape[axis], indices.dtype)
full_range = 1
after_axis_range = 1
for i, value in enumerate(shape, 0):
full_range *= value
if i > axis:
after_axis_range *= value
before_axis_stride = axis_range * after_axis_range
ind_shape = indices.shape
ind_axis_range = ind_shape[axis]
ind_before_axis_range = 1
ind_after_axis_range = 1
for i, value in enumerate(ind_shape, 0):
if i < axis:
ind_before_axis_range *= value
elif i > axis:
ind_after_axis_range *= value
ind_before_axis_stride = ind_axis_range * ind_after_axis_range
def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr, reduce_func):
# pylint: disable=invalid-name
ib = tir.ir_builder.create()
data = ib.buffer_ptr(data_ptr)
indices = ib.buffer_ptr(indices_ptr)
updates = ib.buffer_ptr(updates_ptr)
out = ib.buffer_ptr(out_ptr)
# Copy initial input data to output
with ib.for_range(0, full_range, "i", kind="parallel") as i:
out[i] = data[i]
with ib.for_range(
0, ind_before_axis_range * ind_after_axis_range, "fused", kind="parallel"
) as fused:
i = fused // ind_after_axis_range
j = fused % ind_after_axis_range
pre_index1 = i * ind_before_axis_stride + j
pre_index2 = i * before_axis_stride + j
with ib.for_range(0, ind_axis_range, "k") as k:
# Offset along indices or updates
index1 = pre_index1 + k * ind_after_axis_range
# Get index and shift to positive side if need
k_new = indices[index1]
shifted_index = k_new + (k_new < 0) * axis_range
# Offset along data
index2 = pre_index2 + shifted_index * after_axis_range
reduce_func(out, index2, updates[index1])
return ib.get()
def update_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = update
def add_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] += update
def mul_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] *= update
def mean_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = (dst_ptr[dst_index] + update) / 2
def min_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = tir.min(dst_ptr[dst_index], update)
def max_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = tir.max(dst_ptr[dst_index], update)
reduce_func = None
if reduction == "update":
reduce_func = update_func
elif reduction == "add":
reduce_func = add_func
elif reduction == "mul":
reduce_func = mul_func
elif reduction == "mean":
reduce_func = mean_func
elif reduction == "min":
reduce_func = min_func
elif reduction == "max":
reduce_func = max_func
else:
raise NotImplementedError(
"scatter_elements reduction not in [update, add, mul, mean, min, max]:", reduction
)
out_buf = tir.decl_buffer(data.shape, data.dtype, "out_buf")
return te.extern(
[data.shape],
[data, indices, updates],
lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0], reduce_func),
dtype=data.dtype,
out_buffers=[out_buf],
name="scatter_elements.generic",
tag="scatter_elements.generic",
)
