chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,219 @@
|
||||
# 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=too-many-arguments
|
||||
"""Argsort operator"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .utils import get_const_tuple
|
||||
|
||||
|
||||
def sort(data, axis=-1, is_ascend=1):
|
||||
"""Performs sorting along the given axis and returns an array
|
||||
in sorted order.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
The input tensor.
|
||||
|
||||
axis : int, optional
|
||||
Axis along which to sort the input tensor.
|
||||
By default the flattened array is used.
|
||||
|
||||
is_ascend : boolean, optional
|
||||
Whether to sort in ascending or descending order.
|
||||
|
||||
dtype : string, optional
|
||||
DType of the output indices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor
|
||||
Sorted index tensor.
|
||||
|
||||
"""
|
||||
data_buf = tvm.tirx.decl_buffer(
|
||||
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
|
||||
)
|
||||
out_buf = tvm.tirx.decl_buffer(data.shape, data.dtype, "out_buf", data_alignment=8, layout=None)
|
||||
out = te.extern(
|
||||
data.shape,
|
||||
[data],
|
||||
lambda ins, outs: tvm.tirx.call_packed(
|
||||
"tvm.contrib.sort.sort", ins[0], outs[0], axis, is_ascend
|
||||
),
|
||||
dtype=data.dtype,
|
||||
in_buffers=[data_buf],
|
||||
out_buffers=out_buf,
|
||||
name="sort_cpu",
|
||||
tag="sort_cpu",
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def argsort(data, valid_count=None, axis=-1, is_ascend=1, dtype="float32"):
|
||||
"""Performs sorting along the given axis and returns an array
|
||||
of indices having the same shape as an input array that index
|
||||
data in sorted order.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
The input tensor.
|
||||
|
||||
valid_count : tvm.te.Tensor, optional
|
||||
1-D tensor for valid number of boxes.
|
||||
|
||||
axis : int, optional
|
||||
Axis along which to sort the input tensor.
|
||||
By default the flattened array is used.
|
||||
|
||||
is_ascend : boolean, optional
|
||||
Whether to sort in ascending or descending order.
|
||||
|
||||
dtype : string, optional
|
||||
DType of the output indices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor
|
||||
Sorted index tensor.
|
||||
|
||||
Example
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
# An example to use argsort
|
||||
dshape = (1, 5, 6)
|
||||
data = te.placeholder(dshape, name="data")
|
||||
axis = 0
|
||||
is_ascend = False
|
||||
out = argsort(data, axis=axis, is_ascend=is_ascend)
|
||||
np_data = np.random.uniform(dshape)
|
||||
s = topi.generic.schedule_argsort(out)
|
||||
f = tvm.compile(s, [data, out], "llvm")
|
||||
dev = tvm.cpu()
|
||||
tvm_data = tvm.runtime.tensor(np_data, dev)
|
||||
tvm_out = tvm.runtime.tensor(np.zeros(dshape, dtype=data.dtype.dtype), dev)
|
||||
f(tvm_data, tvm_out)
|
||||
"""
|
||||
data_buf = tvm.tirx.decl_buffer(
|
||||
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
|
||||
)
|
||||
if valid_count is not None:
|
||||
valid_count_buf = tvm.tirx.decl_buffer(
|
||||
valid_count.shape, valid_count.dtype, "valid_count_buf", data_alignment=4, layout=None
|
||||
)
|
||||
out_buf = tvm.tirx.decl_buffer(
|
||||
data.shape, "int32", "out_buf", data_alignment=8, layout=None
|
||||
)
|
||||
out = te.extern(
|
||||
data.shape,
|
||||
[data, valid_count],
|
||||
lambda ins, outs: tvm.tirx.call_packed(
|
||||
"tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
|
||||
),
|
||||
dtype="int32",
|
||||
in_buffers=[data_buf, valid_count_buf],
|
||||
out_buffers=out_buf,
|
||||
name="argsort_nms_cpu",
|
||||
tag="argsort_nms_cpu",
|
||||
)
|
||||
else:
|
||||
out_buf = tvm.tirx.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8, layout=None)
|
||||
out = te.extern(
|
||||
data.shape,
|
||||
[data],
|
||||
lambda ins, outs: tvm.tirx.call_packed(
|
||||
"tvm.contrib.sort.argsort", ins[0], outs[0], axis, is_ascend
|
||||
),
|
||||
dtype=dtype,
|
||||
in_buffers=[data_buf],
|
||||
out_buffers=out_buf,
|
||||
name="argsort_cpu",
|
||||
tag="argsort_cpu",
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"):
|
||||
"""Get the top k elements in an input tensor along the given axis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
The input tensor.
|
||||
|
||||
k : int or tvm.te.Tensor, optional
|
||||
Number of top elements to select. Return all elements if k < 1.
|
||||
|
||||
axis : int, optional
|
||||
Axis long which to sort the input tensor.
|
||||
|
||||
ret_type: str, optional
|
||||
The return type [both, values, indices].
|
||||
"both": return both top k data and indices.
|
||||
"values": return top k data only.
|
||||
"indices": return top k indices only.
|
||||
|
||||
is_ascend : boolean, optional
|
||||
Whether to sort in ascending or descending order.
|
||||
|
||||
dtype : string, optional
|
||||
The data type of the indices output.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor or List[tvm.te.Tensor]
|
||||
The computed result.
|
||||
"""
|
||||
assert ret_type in ["both", "values", "indices"]
|
||||
data_buf = tvm.tirx.decl_buffer(
|
||||
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
|
||||
)
|
||||
out_shape = list(get_const_tuple(data.shape))
|
||||
kvar = tvm.te.var("k")
|
||||
if not isinstance(k, int):
|
||||
out_shape[axis] = kvar
|
||||
elif k >= 1:
|
||||
out_shape[axis] = k
|
||||
out_bufs = []
|
||||
if ret_type in ["both", "values"]:
|
||||
out_bufs.append(
|
||||
tvm.tirx.decl_buffer(out_shape, data.dtype, "value_buf", data_alignment=8, layout=None)
|
||||
)
|
||||
if ret_type in ["both", "indices"]:
|
||||
out_bufs.append(
|
||||
tvm.tirx.decl_buffer(out_shape, dtype, "indices_buf", data_alignment=8, layout=None)
|
||||
)
|
||||
out_shapes = [out_shape] * len(out_bufs)
|
||||
|
||||
kv = kvar if not isinstance(k, int) else k
|
||||
out = te.extern(
|
||||
out_shapes,
|
||||
[data],
|
||||
lambda ins, outs: tvm.tirx.call_packed(
|
||||
"tvm.contrib.sort.topk", ins[0], *outs, kv, axis, ret_type, is_ascend
|
||||
),
|
||||
in_buffers=[data_buf],
|
||||
out_buffers=out_bufs,
|
||||
name="topk_cpu",
|
||||
tag="topk_cpu",
|
||||
)
|
||||
return out
|
||||
Reference in New Issue
Block a user