wehub-resource-sync
166 lines
6.1 KiB
Python
166 lines
6.1 KiB
Python
# 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.
|
|
# pylint: disable=invalid-name
|
|
"""IndexPut operator"""
|
|
|
|
from tvm import te, tirx
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import tirx as T
|
|
|
|
from . import utils
|
|
|
|
|
|
def index_put(data, indices, values, accumulate=False):
|
|
"""Put values into an array according to indices.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
The source array to be modified.
|
|
|
|
indices : Tuple[tvm.te.Tensor]
|
|
Tuple of index tensors (can be multi-dimensional) specifying positions.
|
|
Index tensors are broadcast together following NumPy broadcasting rules.
|
|
|
|
values : tvm.te.Tensor
|
|
The values to place at the specified indices.
|
|
|
|
accumulate : bool, optional
|
|
Whether to accumulate (add) values rather than replace.
|
|
If True, performs tensor[indices] += values
|
|
If False, performs tensor[indices] = values
|
|
Default is False.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.te.Tensor
|
|
"""
|
|
if not isinstance(indices, list | tuple):
|
|
indices = [indices]
|
|
|
|
# Check indices match data dimensions
|
|
if len(indices) != len(data.shape):
|
|
raise ValueError(
|
|
f"Number of index tensors ({len(indices)}) must match "
|
|
f"data dimensions ({len(data.shape)})"
|
|
)
|
|
|
|
# Prepare ranges and strides
|
|
shape = data.shape
|
|
full_range = 1
|
|
for dim in shape:
|
|
full_range *= dim
|
|
|
|
index_shapes = [idx.shape for idx in indices]
|
|
broadcast_ndim = max(len(s) for s in index_shapes)
|
|
broadcast_shape = []
|
|
|
|
for i in range(broadcast_ndim):
|
|
max_dim = 1
|
|
for idx_shape in index_shapes:
|
|
# Right-align shapes
|
|
dim_idx = len(idx_shape) - broadcast_ndim + i
|
|
if dim_idx >= 0:
|
|
dim_size = idx_shape[dim_idx]
|
|
if not utils.equal_const_int(dim_size, 1):
|
|
if utils.equal_const_int(max_dim, 1):
|
|
max_dim = dim_size
|
|
elif not utils.equal_const_int(dim_size, max_dim):
|
|
raise ValueError(f"Cannot broadcast index shapes: {index_shapes}")
|
|
broadcast_shape.append(max_dim)
|
|
|
|
# Compute total number of elements after broadcasting
|
|
index_len = 1
|
|
for dim in broadcast_shape:
|
|
index_len *= dim
|
|
|
|
def gen_ir(data_ptr, index_ptrs, values_ptr, out_ptr, reduce_func):
|
|
data = T.buffer_proxy(data_ptr)
|
|
indices = [T.buffer_proxy(idx) for idx in index_ptrs]
|
|
values = T.buffer_proxy(values_ptr)
|
|
out = T.buffer_proxy(out_ptr)
|
|
|
|
with IRBuilder() as ib:
|
|
with T.seq_scope():
|
|
with T.parallel(0, full_range) as i:
|
|
out[i] = data[i]
|
|
|
|
with T.parallel(0, index_len) as k:
|
|
# Decompose k into multi-dimensional broadcast index
|
|
k_temp = k
|
|
broadcast_indices = []
|
|
for i in range(broadcast_ndim - 1, -1, -1):
|
|
broadcast_indices.insert(0, k_temp % broadcast_shape[i])
|
|
k_temp = k_temp // broadcast_shape[i]
|
|
|
|
flat_index = 0
|
|
stride = 1
|
|
for dim in range(len(shape) - 1, -1, -1):
|
|
# Get the index for this dimension using broadcasting
|
|
idx_shape = index_shapes[dim]
|
|
idx_ndim = len(idx_shape)
|
|
|
|
# Compute the linear index into this index tensor
|
|
idx_offset = 0
|
|
idx_stride = 1
|
|
for i in range(broadcast_ndim - 1, -1, -1):
|
|
# Right-align the index shape with broadcast shape
|
|
dim_idx = idx_ndim - broadcast_ndim + i
|
|
if dim_idx >= 0:
|
|
dim_size = idx_shape[dim_idx]
|
|
# Use broadcasting: if size is 1, use index 0
|
|
# otherwise use broadcast_indices[i]
|
|
if utils.equal_const_int(dim_size, 1):
|
|
idx_in_dim = 0
|
|
else:
|
|
idx_in_dim = broadcast_indices[i]
|
|
idx_offset += idx_in_dim * idx_stride
|
|
idx_stride *= dim_size
|
|
|
|
idx_val = indices[dim][idx_offset]
|
|
shifted_idx = idx_val + (idx_val < 0) * shape[dim]
|
|
flat_index += shifted_idx * stride
|
|
stride *= shape[dim]
|
|
|
|
reduce_func(out, flat_index, values[k])
|
|
|
|
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
|
|
|
|
reduce_func = add_func if accumulate else update_func
|
|
|
|
# Prepare input buffers
|
|
in_buffers = [data]
|
|
in_buffers.extend(indices)
|
|
in_buffers.append(values)
|
|
|
|
out_buf = tirx.decl_buffer(data.shape, data.dtype, "out_buf", layout=None)
|
|
return te.extern(
|
|
[data.shape],
|
|
in_buffers,
|
|
lambda ins, outs: gen_ir(ins[0], ins[1:-1], ins[-1], outs[0], reduce_func),
|
|
dtype=data.dtype,
|
|
out_buffers=[out_buf],
|
|
name="index_put.generic",
|
|
tag="index_put.generic",
|
|
)
|