This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch refactor
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/refactor by this push:
new 78ced6c7bb bring back topi gpu sort and etc
78ced6c7bb is described below
commit 78ced6c7bb602f2ca4450270c48037e60a96c506
Author: Siyuan Feng <[email protected]>
AuthorDate: Sat Feb 15 21:06:13 2025 +0800
bring back topi gpu sort and etc
---
python/tvm/relax/backend/dispatch_sort_scan.py | 16 +-
python/tvm/topi/__init__.py | 1 +
python/tvm/topi/gpu/__init__.py | 2 +
python/tvm/topi/gpu/scan.py | 729 ++++++++++++++++
python/tvm/topi/gpu/sort.py | 939 +++++++++++++++++++++
.../relax/test_backend_dispatch_sort_scan.py | 22 +-
6 files changed, 1690 insertions(+), 19 deletions(-)
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
index e37869c40c..b5a94619c2 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -79,10 +79,10 @@ class SortScanDispatcher(BackendDispatcher):
kwargs = {}
with tgt:
if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
- te_func = topi.cuda.sort_thrust
+ te_func = topi.gpu.sort_thrust
kwargs["workspace"] = self.allocate_workspace(call)
elif self.is_gpu_target(tgt):
- te_func = topi.cuda.sort
+ te_func = topi.gpu.sort
return self.builder_.call_te(
te_func, call.args[0], call.attrs.axis, not
call.attrs.descending, **kwargs
)
@@ -92,10 +92,10 @@ class SortScanDispatcher(BackendDispatcher):
kwargs = {}
with tgt:
if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
- te_func = topi.cuda.argsort_thrust
+ te_func = topi.gpu.argsort_thrust
kwargs["workspace"] = self.allocate_workspace(call)
elif self.is_gpu_target(tgt):
- te_func = topi.cuda.argsort
+ te_func = topi.gpu.argsort
return self.builder_.call_te(
te_func,
call.args[0],
@@ -109,10 +109,10 @@ class SortScanDispatcher(BackendDispatcher):
te_func = topi.topk
kwargs = {}
if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
- te_func = topi.cuda.topk_thrust
+ te_func = topi.gpu.topk_thrust
kwargs["workspace"] = self.allocate_workspace(call)
elif self.is_gpu_target(tgt):
- te_func = topi.cuda.topk
+ te_func = topi.gpu.topk
tir_call = self.builder_.call_te(
te_func,
call.args[0],
@@ -176,11 +176,11 @@ class SortScanDispatcher(BackendDispatcher):
with tgt:
if call.op.name == "relax.cumsum":
- te_func = topi.cuda.cumsum if self.is_gpu_target(tgt) else
topi.cumsum
+ te_func = topi.gpu.cumsum if self.is_gpu_target(tgt) else
topi.cumsum
if can_use_thrust(tgt, "tvm.contrib.thrust.sum_scan"):
kwargs["workspace"] = self.allocate_workspace(call)
elif call.op.name == "relax.cumprod":
- te_func = topi.cuda.cumprod if self.is_gpu_target(tgt)
else topi.cumprod
+ te_func = topi.gpu.cumprod if self.is_gpu_target(tgt) else
topi.cumprod
else:
raise ValueError(f"Unsupported op: {call.op.name}")
tir_call = self.builder_.call_te(
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index 2bd5964fef..3588c04d8f 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -52,6 +52,7 @@ from . import utils
from . import vision
from . import image
from . import random
+from . import gpu
# error reporting
from .utils import InvalidShapeError
diff --git a/python/tvm/topi/gpu/__init__.py b/python/tvm/topi/gpu/__init__.py
new file mode 100644
index 0000000000..e5a65dc054
--- /dev/null
+++ b/python/tvm/topi/gpu/__init__.py
@@ -0,0 +1,2 @@
+from .scan import cumsum, cumprod
+from .sort import *
\ No newline at end of file
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
new file mode 100644
index 0000000000..be7f7dbb17
--- /dev/null
+++ b/python/tvm/topi/gpu/scan.py
@@ -0,0 +1,729 @@
+# 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, too-many-locals, too-many-statements
+"Scan related operators"
+from typing import Callable, Optional, Union
+
+import tvm
+from tvm import te
+from tvm.contrib.thrust import can_use_rocthrust, can_use_thrust
+
+from .. import tag
+from ..math import cast, ceil_log2
+from ..transform import expand_dims, reshape, squeeze, transpose
+from ..utils import ceil_div, get_const_int, prod, swap
+
+
+def _get_thrust_func_name(tvmop):
+ tvmop_to_thrust_func_name = {tvm.tir.generic.add:
"tvm.contrib.thrust.sum_scan"}
+ assert tvmop in tvmop_to_thrust_func_name, f"{tvmop} not supported by
thrust"
+ return tvmop_to_thrust_func_name[tvmop]
+
+
+def _can_use_scan_thrust(binop):
+ """
+ Check if scan_thrust can be utilized based on the current target and
binary op.
+ """
+ target = tvm.target.Target.current()
+ if target is None:
+ return False
+ return binop == tvm.tir.generic.add and any(
+ [
+ can_use_thrust(target, "tvm.contrib.thrust.sum_scan"),
+ can_use_rocthrust(target, "tvm.contrib.thrust.sum_scan"),
+ ]
+ )
+
+
+def exclusive_scan_ir(data, output, reduction=None, binop=tvm.tir.generic.add,
identity_value=0):
+ """Low level IR to do exclusive sum scan along rows of 2D input.
+
+ Parameters
+ ----------
+ data : Buffer
+ Input N-D Buffer. Scan is done over the innermost axis.
+
+ output: Buffer
+ A buffer to store the output scan, of the same shape as data
+
+ reduction: Buffer, optional
+ (N-1)-D Buffer, to store the sum of each scan axis.
+
+ binop: function, optional
+ A binary associative op to use for scan. The function takes two TIR
expressions
+ and produce a new TIR expression. By default it uses
tvm.tir.generic.add to compute
+ prefix sum.
+
+ identity_value: int or float
+ A value for the binary operation which provides the identity property.
E.g. if * is
+ your operator and i is the identity_value then a * i = a for all a in
the domain of
+ your operation.
+ """
+
+ batch_size = cast(prod(data.shape[:-1]), "int32")
+ scan_axis_size = cast(data.shape[-1], "int32")
+
+ ib = tvm.tir.ir_builder.create()
+
+ data = ib.buffer_ptr(data)
+ output = ib.buffer_ptr(output)
+
+ out_dtype = output.dtype
+
+ if reduction is not None:
+ reduction = ib.buffer_ptr(reduction)
+
+ max_threads =
int(tvm.target.Target.current(allow_none=False).max_num_threads)
+
+ with ib.if_scope(scan_axis_size == 0):
+ with ib.new_scope():
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(bx, "thread_extent", batch_size)
+ with ib.if_scope(bx < batch_size):
+ if reduction is not None:
+ reduction[bx] = cast(identity_value, out_dtype)
+ with ib.else_scope():
+ with ib.new_scope():
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(scan_axis_size, max_threads)
+ nthread_by = batch_size
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.x")
+ by = te.thread_axis("blockIdx.y")
+ ib.scope_attr(tx, "thread_extent", nthread_tx)
+ ib.scope_attr(bx, "thread_extent", nthread_bx)
+ ib.scope_attr(by, "thread_extent", nthread_by)
+ tid = bx * nthread_tx + tx
+ with ib.if_scope(tid < scan_axis_size):
+ output[by * scan_axis_size + tid] = cast(data[by *
scan_axis_size + tid], out_dtype)
+
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(scan_axis_size, max_threads)
+ nthread_by = batch_size
+
+ # The following algorithm performs parallel exclusive scan
+ # Up Sweep of exclusive scan
+ lim = ceil_log2(scan_axis_size)
+
+ with ib.for_range(0, cast(lim, "int32"), dtype="int32") as l2_width:
+ width = 2 << l2_width
+
+ with ib.new_scope():
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(tx, "thread_extent", nthread_tx)
+ ib.scope_attr(
+ bx,
+ "thread_extent",
+ tvm.tir.generic.cast(ceil_div(scan_axis_size, max_threads
* width), "int32"),
+ )
+ tid = bx * nthread_tx + tx
+
+ by = te.thread_axis("blockIdx.y")
+ ib.scope_attr(by, "thread_extent", nthread_by)
+ start = ib.allocate("int32", (1,), name="start", scope="local")
+ middle = ib.allocate("int32", (1,), name="middle",
scope="local")
+ end = ib.allocate("int32", (1,), name="end", scope="local")
+ start[0] = width * tid
+ with ib.if_scope(start[0] < scan_axis_size):
+ middle[0] = start[0] + tvm.tir.indexdiv(width, 2)
+ end[0] = tvm.te.min(start[0] + width, scan_axis_size)
+ with ib.if_scope(middle[0] < scan_axis_size):
+ output[by * scan_axis_size + end[0] - 1] = binop(
+ output[by * scan_axis_size + end[0] - 1],
+ output[by * scan_axis_size + middle[0] - 1],
+ )
+
+ # Down Sweep of exclusive scan
+ with ib.new_scope():
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(bx, "thread_extent", batch_size)
+ with ib.if_scope(bx < batch_size):
+ if reduction is not None:
+ reduction[bx] = output[(bx + 1) * scan_axis_size - 1]
+ output[(bx + 1) * scan_axis_size - 1] = cast(identity_value,
out_dtype)
+
+ with ib.for_range(0, cast(lim, "int32"), dtype="int32") as l2_width:
+ width = 2 << (lim - l2_width - 1)
+
+ with ib.new_scope():
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(tx, "thread_extent", nthread_tx)
+ ib.scope_attr(
+ bx,
+ "thread_extent",
+ tvm.tir.generic.cast(ceil_div(scan_axis_size, max_threads
* width), "int32"),
+ )
+ tid = bx * nthread_tx + tx
+
+ by = te.thread_axis("blockIdx.y")
+ ib.scope_attr(by, "thread_extent", nthread_by)
+ start = ib.allocate("int32", (1,), name="start", scope="local")
+ middle = ib.allocate("int32", (1,), name="middle",
scope="local")
+ end = ib.allocate("int32", (1,), name="end", scope="local")
+ tmp = ib.allocate(out_dtype, (1,), name="end", scope="local")
+ start[0] = width * tid
+ with ib.if_scope(tvm.tir.all(start[0] < scan_axis_size)):
+ middle[0] = start[0] + tvm.tir.indexdiv(width, 2)
+ end[0] = tvm.tir.min(start[0] + width, scan_axis_size)
+ with ib.if_scope(middle[0] < scan_axis_size):
+ tmp[0] = output[by * scan_axis_size + middle[0] - 1]
+ output[by * scan_axis_size + middle[0] - 1] = output[
+ by * scan_axis_size + end[0] - 1
+ ]
+ output[by * scan_axis_size + end[0] - 1] = binop(
+ output[by * scan_axis_size + end[0] - 1], tmp[0]
+ )
+ return ib.get()
+
+
+def get_reduction_from_exclusive_scan(data, ex_scan_output,
binop=tvm.tir.generic.add):
+ """Return the sum of the last element of data and the exclusive scan
output.
+ The is the reduction of data along each row (for 2-D case).
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ Input data of any shape
+
+ ex_scan_output : tvm.te.Tensor
+ The output of exclusive scan on data
+
+ binop: function, optional
+ A binary associative op to use for scan. The function takes two TIR
expressions
+ and produce a new TIR expression. By default it uses
tvm.tir.generic.add to compute
+ prefix sum.
+
+ Returns
+ -------
+ reduction : tvm.te.Tensor
+ (N-1)-D tensor storing the reduction of each scan axis.
+ """
+ ndim = len(data.shape)
+ if ndim == 1:
+ data = expand_dims(data, axis=0)
+ ex_scan_output = expand_dims(ex_scan_output, axis=0)
+
+ def ir(data, data_ex_scan, reduction):
+ batch_size = cast(prod(data.shape[:-1]), "int32")
+ scan_axis_size = cast(data.shape[-1], "int32")
+
+ ib = tvm.tir.ir_builder.create()
+
+ data = ib.buffer_ptr(data)
+ data_ex_scan = ib.buffer_ptr(data_ex_scan)
+ reduction = ib.buffer_ptr(reduction)
+
+ max_threads =
int(tvm.target.Target.current(allow_none=False).max_num_threads)
+ with ib.new_scope():
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(batch_size, max_threads)
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(tx, "thread_extent", nthread_tx)
+ ib.scope_attr(bx, "thread_extent", nthread_bx)
+ tid = bx * max_threads + tx
+ with ib.if_scope(tid < batch_size):
+ with ib.if_scope(scan_axis_size > 0):
+ reduction[tid] = binop(
+ data_ex_scan[tid * scan_axis_size + scan_axis_size -
1],
+ data[tid * scan_axis_size + scan_axis_size - 1],
+ )
+ with ib.else_scope():
+ reduction[tid] = cast(0, reduction.dtype)
+
+ return ib.get()
+
+ data_buf = tvm.tir.decl_buffer(data.shape, data.dtype,
"valid_indices_buf", data_alignment=8)
+ ex_scan_output_buf = tvm.tir.decl_buffer(
+ ex_scan_output.shape, ex_scan_output.dtype, "ex_scan_output_buf",
data_alignment=8
+ )
+
+ reduction = te.extern(
+ [data.shape[:-1]],
+ [data, ex_scan_output],
+ lambda ins, outs: ir(ins[0], ins[1], outs[0]),
+ dtype=[ex_scan_output.dtype],
+ in_buffers=[data_buf, ex_scan_output_buf],
+ name="ex_scan_reduction",
+ tag="ex_scan_reduction_gpu",
+ )
+
+ if ndim == 1:
+ return squeeze(reduction, 0)
+
+ return reduction
+
+
+def scan_thrust(
+ data,
+ output_dtype,
+ exclusive=True,
+ return_reduction=False,
+ binop=tvm.tir.generic.add,
+ workspace=None,
+):
+ """Do exclusive or inclusive scan on 1D or multidimensional input, using
thrust.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ Input data of any shape. The scan is done over the innermost axis.
+
+ output_dtype: string
+ The dtype of the output scan tensor.
+
+ exclusive: bool, optional
+ Whether or not do exclusive or inclusive scan.
+
+ return_reduction: bool, optional
+ Whether or not return a (N-1)-D tensor storing the reduction of each
scan axis.
+ Reductions are computed as part of the upsweep pass, so there is no
extra cost.
+ If False, reductions are ignored. It must be False when exclusive is
False.
+
+ binop: function, optional
+ A binary associative op to use for scan. Since we need to lookup the
corresponding
+ thrust function, arbitrariy callables are not supported. Currently only
+ tvm.tir.generic.add can be passed in.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results. The size of the workspace
should be sufficiently
+ large, this can be obtained by overestimation or memory usage
profiling. If None, it will
+ fallback to use thrust internal memory allocation.
+
+ Returns
+ -------
+ output : tvm.te.Tensor
+ A N-D tensor of the same rank N and shape as the input data.
+
+ reduction : tvm.te.Tensor, optional
+ (N-1)-D tensor storing the reduction of each scan axis.
+ Returned if return_reduction is True.
+ """
+ data_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "data_buf",
data_alignment=8)
+ output_buf = tvm.tir.decl_buffer(data.shape, output_dtype, "output_buf",
data_alignment=8)
+
+ workspace_buf = (
+ tvm.tir.decl_buffer(workspace.shape, workspace.dtype, "workspace_buf",
data_alignment=8)
+ if workspace is not None
+ else None
+ )
+
+ def f_compute(ins, outs):
+ args = [_get_thrust_func_name(binop), ins[0], outs[0], exclusive]
+ if workspace is not None:
+ args.append(ins[1])
+ return tvm.tir.call_packed(*args)
+
+ output = te.extern(
+ [data.shape],
+ [data] if workspace is None else [data, workspace],
+ f_compute,
+ dtype=[output_dtype],
+ in_buffers=[data_buf] if workspace is None else [data_buf,
workspace_buf],
+ out_buffers=[output_buf],
+ name="exclusive_scan_thrust",
+ tag="exclusive_scan_thrust_gpu",
+ )
+
+ if return_reduction:
+ assert exclusive, "return_reduction should be False for inclusive scan"
+ reduction = get_reduction_from_exclusive_scan(data, output, binop)
+ return output, reduction
+
+ return output
+
+
+def exclusive_scan(
+ data,
+ axis=-1,
+ return_reduction=False,
+ output_dtype=None,
+ binop=tvm.tir.generic.add,
+ identity_value=0,
+ workspace=None,
+):
+ """Do exclusive scan on 1D or multidimensional input.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ Input data of any shape.
+
+ axis: int, optional
+ The axis to do scan on. By default, scan is done on the innermost axis.
+
+ return_reduction: bool, optional
+ Whether or not return a tensor storing the reduction over each scan
axis.
+ If the input rank is N, this tensor is of rank N - 1.
+ Reductions are computed as part of the upsweep pass, so there is no
extra cost.
+ If False, reductions are ignored.
+
+ output_dtype: string, optional
+ The dtype of the output scan tensor. If not provided, the dtype of the
input is used.
+
+ binop: function, optional
+ A binary associative op to use for scan. The function takes two TIR
expressions
+ and produce a new TIR expression. By default it uses
tvm.tir.generic.add to compute
+ prefix sum.
+
+ identity_value: int or float
+ A value for the binary operation which provides the identity property.
E.g. if * is
+ your operator and i is the identity_value then a * i = a for all a in
the domain of
+ your operation.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results if thrust is enabled. The size
of the workspace
+ should be sufficiently large, this can be obtained by overestimation
or memory usage
+ profiling. If None, it will fallback to use thrust internal memory
allocation.
+
+ Returns
+ -------
+ output : tvm.te.Tensor
+ A N-D tensor of the same rank N and shape as the input data.
+
+ reduction : tvm.te.Tensor, optional
+ (N-1)-D tensor storing the reduction of each scan axis.
+ Returned if return_reduction is True.
+ """
+
+ def do_scan(data, output_dtype):
+ # TODO: add support for a prod_scan
+ if _can_use_scan_thrust(binop):
+ return scan_thrust(
+ data,
+ output_dtype,
+ exclusive=True,
+ return_reduction=return_reduction,
+ binop=binop,
+ workspace=workspace,
+ )
+
+ if ndim == 1:
+ # TIR exclusive scan accepts only 2D or higher-rank inputs.
+ data = expand_dims(data, axis=0)
+
+ data_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "data_buf",
data_alignment=8)
+ output_buf = tvm.tir.decl_buffer(data.shape, output_dtype,
"output_buf", data_alignment=8)
+
+ if return_reduction:
+ output, reduction = te.extern(
+ [data.shape, data.shape[:-1]],
+ [data],
+ lambda ins, outs: exclusive_scan_ir(
+ ins[0], outs[0], outs[1], binop=binop,
identity_value=identity_value
+ ),
+ dtype=[output_dtype, output_dtype],
+ in_buffers=[data_buf],
+ name="exclusive_scan",
+ tag="exclusive_scan_gpu",
+ )
+ else:
+ output = te.extern(
+ [data.shape],
+ [data],
+ lambda ins, outs: exclusive_scan_ir(
+ ins[0], outs[0], binop=binop, identity_value=identity_value
+ ),
+ dtype=[output_dtype],
+ in_buffers=[data_buf],
+ out_buffers=[output_buf],
+ name="exclusive_scan",
+ tag="exclusive_scan_gpu",
+ )
+ reduction = None
+
+ if ndim == 1:
+ output = squeeze(output, 0)
+ if return_reduction:
+ reduction = squeeze(reduction, 0)
+
+ if return_reduction:
+ return output, reduction
+
+ return output
+
+ if output_dtype is None or output_dtype == "":
+ output_dtype = data.dtype
+
+ ndim = len(data.shape)
+ if axis < 0:
+ axis += ndim
+
+ # If scan axis is not the innermost one, swap the scan and the innermost
axes
+ # Scan is always done on the innermost axis, for performance reason.
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ if return_reduction:
+ output, reduction = do_scan(data, output_dtype)
+ else:
+ output = do_scan(data, output_dtype)
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ output = transpose(output, axes)
+
+ if return_reduction:
+ return output, reduction
+
+ return output
+
+
+def inclusive_scan(
+ data, axis=-1, output_dtype=None, binop=tvm.tir.generic.add,
identity_value=0, workspace=None
+):
+ """Do inclusive scan on 1D or multidimensional input.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ Input data of any shape.
+
+ axis: int, optional
+ The axis to do scan on. By default, scan is done on the innermost axis.
+
+ output_dtype: string, optional
+ The dtype of the output scan tensor. If not provided, the dtype of the
input is used.
+
+ binop: function, optional
+ A binary associative op to use for scan. The function takes two TIR
expressions
+ and produce a new TIR expression. By default it uses
tvm.tir.generic.add to compute
+ prefix sum.
+
+ identity_value: int or float
+ A value for the binary operation which provides the identity property.
E.g. if * is
+ your operator and i is the identity_value then a * i = a for all a in
the domain of
+ your operation.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results if thrust is enabled. The size
of the workspace
+ should be sufficiently large, this can be obtained by overestimation
or memory usage
+ profiling. If None, it will fallback to use thrust internal memory
allocation.
+
+ Returns
+ -------
+ output : tvm.te.Tensor
+ A N-D tensor of the same rank N as the input data.
+ """
+
+ if _can_use_scan_thrust(binop):
+ if output_dtype is None or output_dtype == "":
+ output_dtype = data.dtype
+ ndim = len(data.shape)
+ if axis < 0:
+ axis += ndim
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+ output = scan_thrust(data, output_dtype, exclusive=False, binop=binop,
workspace=workspace)
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ output = transpose(output, axes)
+ return output
+
+ ex_scan = exclusive_scan(
+ data,
+ axis,
+ output_dtype=output_dtype,
+ binop=binop,
+ identity_value=identity_value,
+ workspace=workspace,
+ )
+
+ if output_dtype is not None and data.dtype != output_dtype and
output_dtype != "":
+ data = cast(data, output_dtype)
+
+ return binop(data, ex_scan)
+
+
+def scanop(
+ data: tvm.te.Tensor,
+ binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
+ identity_value: Union[float, int],
+ axis: Optional[int] = None,
+ dtype: Optional[str] = None,
+ exclusive: Optional[bool] = None,
+ workspace: Optional[tvm.te.Tensor] = None,
+) -> tvm.te.Tensor:
+ """Cumulative binary operator (scan) with similar axis behavior as
np.cumsum and np.cumprod.
+
+ See cumprod and cumsum for an example of use.
+
+ E.g. if * is your binary operator and the input tensor is [1, 2, 3, 4] the
output may be
+ [1, 1 * 2, 1 * 2 * 3, 1 * 2 * 3 * 4]
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ The input data to the operator.
+
+ binop: Callable (tvm.Expr, tvm.Expr) -> tvm.Expr
+ A binary operator which should be associative and commutative. E.g. if
* is your
+ operator then a * (b * c) = (a * b) * c and a * b = b * a
+
+ identity_value: int or float
+ A value for the binary operation which provides the identity property.
E.g. if * is
+ your operator and i is the identity_value then a * i = a for all a in
the domain of
+ your operation.
+
+ axis : int, optional
+ Axis along which the operation is computed. The default (None) is to
compute
+ the cumulative operation over the flattened array.
+
+ dtype : string, optional
+ Type of the returned array and of the accumulator in which the
elements are computed.
+ If dtype is not specified, it defaults to the dtype of data.
+
+ exclusive : bool, optional
+ If true will return exclusive cumulative operation in which the first
element is not
+ included. In other terms, if true, the j-th output element would be
+ the cumulative operation of the first (j-1) elements. Otherwise, it
would be the
+ cumulative operation of the first j elements.
+
+ workspace: Optional[tvm.te.Tensor]
+
+ Returns
+ -------
+ result : tvm.te.Tensor
+ The result has the same size as data, and the same shape as data if
axis is not None.
+ If axis is None, the result is a 1-d array.
+ """
+ if axis is None:
+ axis = 0
+ data = reshape(data, (prod(data.shape),))
+ axis = get_const_int(axis)
+ if exclusive is not None and exclusive:
+ return exclusive_scan(
+ data,
+ axis,
+ output_dtype=dtype,
+ binop=binop,
+ identity_value=identity_value,
+ workspace=workspace,
+ )
+ return inclusive_scan(
+ data,
+ axis,
+ output_dtype=dtype,
+ binop=binop,
+ identity_value=identity_value,
+ workspace=workspace,
+ )
+
+
+def cumsum(
+ data: tvm.te.Tensor,
+ axis: Optional[int] = None,
+ dtype: Optional[int] = None,
+ exclusive: Optional[bool] = None,
+ workspace: Optional[tvm.te.Tensor] = None,
+) -> tvm.te.Tensor:
+ """Numpy style cumsum op. Return the cumulative sum of the elements along
a given axis.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ The input data to the operator.
+
+ axis : int, optional
+ Axis along which the cumulative sum is computed. The default (None) is
to compute
+ the cumsum over the flattened array.
+
+ dtype : string, optional
+ Type of the returned array and of the accumulator in which the
elements are summed.
+ If dtype is not specified, it defaults to the dtype of data.
+
+ exclusive : bool, optional
+ If true will return exclusive sum in which the first element is not
+ included. In other terms, if true, the j-th output element would be
+ the sum of the first (j-1) elements. Otherwise, it would be the sum of
+ the first j elements.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results if thrust is enabled. The size
of the workspace
+ should be sufficiently large, this can be obtained by overestimation
or memory usage
+ profiling. If None, it will fallback to use thrust internal memory
allocation.
+
+ Returns
+ -------
+ result : tvm.te.Tensor
+ The result has the same size as data, and the same shape as data if
axis is not None.
+ If axis is None, the result is a 1-d array.
+ """
+ return scanop(
+ data=data,
+ binop=tvm.tir.generic.add,
+ identity_value=0,
+ axis=axis,
+ dtype=dtype,
+ exclusive=exclusive,
+ workspace=workspace,
+ )
+
+
+def cumprod(
+ data: tvm.te.Tensor,
+ axis: Optional[int] = None,
+ dtype: Optional[int] = None,
+ exclusive: Optional[bool] = None,
+ workspace: Optional[tvm.te.Tensor] = None,
+):
+ """Numpy style cumprod op. Return the cumulative product of the elements
along a given axis.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ The input data to the operator.
+
+ axis : int, optional
+ Axis along which the cumulative product is computed. The default
(None) is to compute
+ the cumproduct over the flattened array.
+
+ dtype : string, optional
+ Type of the returned array and of the accumulator in which the
elements are multiplied.
+ If dtype is not specified, it defaults to the dtype of data.
+
+ exclusive : bool, optional
+ If True, will return exclusive product in which the first element is
not
+ included. In other terms, if True, the j-th output element would be
+ the product of the first (j-1) elements. Otherwise, it would be the
product of
+ the first j elements.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results if thrust is enabled. The size
of the workspace
+ should be sufficiently large, this can be obtained by overestimation
or memory usage
+ profiling. If None, it will fallback to use thrust internal memory
allocation.
+
+ Returns
+ -------
+ result : tvm.te.Tensor
+ The result has the same size as data, and the same shape as data if
axis is not None.
+ If axis is None, the result is a 1-d array.
+ """
+ return scanop(
+ data=data,
+ binop=tvm.tir.generic.multiply,
+ identity_value=1,
+ axis=axis,
+ dtype=dtype,
+ exclusive=exclusive,
+ workspace=workspace,
+ )
\ No newline at end of file
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
new file mode 100644
index 0000000000..71854e4399
--- /dev/null
+++ b/python/tvm/topi/gpu/sort.py
@@ -0,0 +1,939 @@
+# 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, no-member, too-many-locals,
too-many-arguments, too-many-statements, singleton-comparison, unused-argument,
no-else-return
+"""Sort related operators """
+import tvm
+from tvm import te
+
+from ..transform import strided_slice, transpose
+from ..utils import ceil_div, swap
+from ..math import cast, ceil_log2
+
+
+def _get_threads(ib, nthread_tx, nthread_bx, nthread_by):
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.x")
+ ib.scope_attr(tx, "thread_extent", nthread_tx)
+ ib.scope_attr(bx, "thread_extent", nthread_bx)
+
+ by = te.thread_axis("blockIdx.y")
+ ib.scope_attr(by, "thread_extent", nthread_by)
+
+ return tx, bx, by
+
+
+def _sort_init(ib, shape, axis, keys_in, keys_out, values_out=None,
value_init_func=None):
+ """Initialize the output buffers by copying from inputs"""
+ axis_mul_before = 1
+ axis_mul_after = 1
+ if axis < 0:
+ axis = len(shape) + axis
+ for i, value in enumerate(shape, 0):
+ if i < axis:
+ axis_mul_before *= value
+ elif i > axis:
+ axis_mul_after *= value
+
+ # Set up threading
+ max_threads =
int(tvm.target.Target.current(allow_none=False).max_num_threads)
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(shape[axis], max_threads)
+ nthread_by = axis_mul_before * axis_mul_after
+
+ # Copy the keys_in to initial output
+ with ib.new_scope():
+ tx, bx, by = _get_threads(ib, nthread_tx, nthread_bx, nthread_by)
+ tid = bx * nthread_tx + tx
+ by, bz = by % axis_mul_before, by // axis_mul_before
+ idx = (by * shape[axis] + tid) * axis_mul_after + bz
+ with ib.if_scope(tid < shape[axis]):
+ keys_out[idx] = keys_in[idx]
+ if values_out is not None:
+ values_out[idx] = value_init_func(idx, tid)
+
+ return axis_mul_before, axis_mul_after
+
+
+## TODO(mbrookhart): These are effective optimziation hyperparametrs
+## Perhaps we can autotune?
+block_size = 128
+thread_work = 4
+
+
+def _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before,
+ axis_mul_after,
+ is_ascend,
+ keys,
+ keys_swap,
+ values=None,
+ values_swap=None,
+):
+ nthread_tx = block_size // 2
+ nthread_bx = ceil_div(size, block_size)
+ nthread_by = axis_mul_before * axis_mul_after
+ with ib.new_scope():
+ ib.scope_attr(tvm.tir.const(0), "hand_threaded", 0)
+ tx, bx, by = _get_threads(ib, nthread_tx, nthread_bx, nthread_by)
+ by, bz = by % axis_mul_before, by // axis_mul_before
+ tid = 2 * tx
+ start = bx * block_size
+
+ ## Create shared memory as syncable thread scratch space
+ tmp_keys_swap = ib.allocate(
+ keys_swap.dtype,
+ (block_size,),
+ name="temp_keys_swap",
+ scope="shared",
+ )
+ if values_swap is not None:
+ tmp_values_swap = ib.allocate(
+ values_swap.dtype,
+ (block_size,),
+ name="temp_values_swap",
+ scope="shared",
+ )
+
+ ## Create thread local data for swapping
+ temp_keys = ib.allocate(keys_swap.dtype, (1,), name="temp_keys",
scope="local")
+ if values_swap is not None:
+ temp_values = ib.allocate(values_swap.dtype, (1,),
name="temp_values", scope="local")
+
+ temp_cond1 = ib.allocate(keys_swap.dtype, (1,), name="temp_cond1",
scope="local")
+ temp_cond2 = ib.allocate(keys_swap.dtype, (1,), name="temp_cond2",
scope="local")
+ # Copy data to scratch space
+ base_idx = by * size * axis_mul_after + bz
+ with ib.for_range(0, 2) as n:
+ with ib.if_scope((tid + n + start) < size):
+ tmp_keys_swap[tid + n] = keys[base_idx + (tid + n + start) *
axis_mul_after]
+ if values_swap is not None:
+ tmp_values_swap[tid + n] = values[base_idx + (tid + n +
start) * axis_mul_after]
+
+ ib.emit(tvm.tir.Call(None, "tir.tvm_storage_sync",
tvm.runtime.convert(["shared"])))
+
+ idxm = tvm.tir.indexmod
+ # OddEvenTransposeSort
+ current_sort_num = tvm.tir.min(block_size, size - start)
+ with ib.for_range(0, current_sort_num) as k:
+ n = idxm(tid + k, 2)
+ with ib.if_scope(tid + n < current_sort_num - 1):
+ temp_cond1[0] = tmp_keys_swap[tid + n]
+ temp_cond2[0] = tmp_keys_swap[tid + n + 1]
+ if is_ascend:
+ cond = temp_cond1[0] > temp_cond2[0]
+ else:
+ cond = temp_cond1[0] < temp_cond2[0]
+ with ib.if_scope(cond):
+ temp_keys[0] = tmp_keys_swap[tid + n]
+ tmp_keys_swap[tid + n] = tmp_keys_swap[tid + n + 1]
+ tmp_keys_swap[tid + n + 1] = temp_keys[0]
+ if values_swap is not None:
+ temp_values[0] = tmp_values_swap[tid + n]
+ tmp_values_swap[tid + n] = tmp_values_swap[tid + n + 1]
+ tmp_values_swap[tid + n + 1] = temp_values[0]
+ ib.emit(tvm.tir.Call(None, "tir.tvm_storage_sync",
tvm.runtime.convert(["shared"])))
+
+ ## Copy sorted data to output
+ with ib.for_range(0, 2) as n:
+ with ib.if_scope(tid + n + start < size):
+ keys[base_idx + (tid + n + start) * axis_mul_after] =
tmp_keys_swap[tid + n]
+ keys_swap[base_idx + (tid + n + start) * axis_mul_after] =
tmp_keys_swap[tid + n]
+ if values_swap is not None:
+ values[base_idx + (tid + n + start) * axis_mul_after] =
tmp_values_swap[tid + n]
+ values_swap[base_idx + (tid + n + start) * axis_mul_after]
= tmp_values_swap[
+ tid + n
+ ]
+
+
+def _sort_common(
+ ib,
+ size,
+ axis_mul_before,
+ axis_mul_after,
+ is_ascend,
+ keys,
+ keys_swap,
+ values=None,
+ values_swap=None,
+):
+ """Either sort only values or sort values by keys."""
+
+ ## This function performs a multi-level mergesort
+ ## For blocks of length <= block_size, it does odd-even transpose sort
+ ## in GPU shared memory
+ ## For intermediate block sizes (>block_size, < max_threads * thread_work)
+ ## it uses the mergpath algorthim https://arxiv.org/abs/1406.2628
+ ## to merge blocks in parallel
+ ## At some point, the size of the blocks to be merged is too big for
max_threads
+ ## and we switch to using a dual-level mergepath where the outer
mergepath
+ ## finds the start/end locations of the inner mergepath so that we can
split
+ ## the merge into more blocks
+
+ max_threads =
int(tvm.target.Target.current(allow_none=False).max_num_threads)
+ nthread_by = axis_mul_before * axis_mul_after
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(size, nthread_tx)
+
+ def compare(a, b):
+ """
+ Compare a and b in proper ascending or descending order
+ """
+ if is_ascend:
+ out = a <= b
+ else:
+ out = b <= a
+ return out
+
+ # Sort the lower levels of the merge using odd-even sort, it's fast for
small inputs
+ lower_lim = ceil_log2(block_size)
+
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = ceil_log2(size)
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart,
diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid)]
+ b = source[base_idx + (bStart + diag - 1 - mid)]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid + 1
+ with ib.else_scope():
+ last[0] = mid
+ return first[0], last[0]
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first
+ j[0] = bStart + diag - last
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count))
as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart
+ bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to
the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, cast(upper_lim - lower_lim, "int64"), dtype="int64")
as l2_width:
+ width = 2 << (l2_width + lower_lim)
+ # Define and launch the cuda kernel
+ with ib.new_scope():
+ target = tvm.target.Target.current()
+ if "vulkan" in str(target):
+ # Vulkan can't handle dynamic nthread, so we thread slightly
differently
+ # for vulkan. We don't do this generally because it causes a
15% perf
+ # regression on other platforms
+ ntx = max_threads
+ nbx = tvm.tir.generic.cast(ceil_div(width, max_threads *
thread_work), "int32")
+ nbz = tvm.tir.generic.cast(ceil_div(size, width), "int32")
+ tx, bx, by = _get_threads(ib, ntx, nbx, nthread_by * nbz)
+ else:
+ ntx = tvm.tir.generic.cast(tvm.te.min(max_threads, width),
"int32")
+ nbx = tvm.tir.generic.cast(ceil_div(width, max_threads *
thread_work), "int32")
+ nbz = tvm.tir.generic.cast(ceil_div(size, width), "int32")
+ tx, bx, by = _get_threads(ib, ntx, nbx, nthread_by * nbz)
+ by, bz = by % nthread_by, by // nthread_by
+
+ def mergepath(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ step_count,
+ even,
+ ):
+ # pylint: disable=arguments-out-of-order
+ def merge(source, dest, source_idx, dest_idx):
+ diag = tx * step_count
+ first, last = get_merge_begin(
+ source,
+ by * size,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ diag,
+ step_count,
+ )
+ # iterate over the output loop
+ serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ by * size,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ )
+
+ with ib.if_scope(even):
+ merge(source, dest, source_idx, dest_idx)
+ with ib.else_scope():
+ merge(dest, source, dest_idx, source_idx)
+
+ def mergesort(source, dest, source_idx, dest_idx, size, width,
even):
+ # calculate the start, mid, and end points of this section
+ start = width * bz
+ middle = cast(tvm.te.min(start + tvm.tir.indexdiv(width, 2),
size), "int64")
+ end = cast(tvm.te.min(start + width, size), "int64")
+ with ib.if_scope(start < size):
+ with ib.if_scope(nbx == 1):
+ ## merge the start->middle and middle->end arrays
+ aCount = middle - start
+ bCount = end - middle
+ mergepath(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ aCount,
+ bCount,
+ start,
+ middle,
+ start,
+ ceil_div(width, ntx),
+ even,
+ )
+ with ib.else_scope():
+ step_count = max_threads * thread_work
+ diag = bx * step_count
+
+ def do_merge(first, last):
+ aStart = start + first
+ bStart = middle + diag - last
+ aCount = tvm.te.min(middle - aStart, step_count)
+ bCount = tvm.te.min(end - bStart, step_count)
+ mergepath(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ start + diag,
+ thread_work,
+ even,
+ )
+
+ with ib.if_scope(even):
+ first, last = get_merge_begin(
+ source,
+ by * size,
+ middle - start,
+ end - middle,
+ start,
+ middle,
+ diag,
+ step_count,
+ )
+ do_merge(first, last)
+ with ib.else_scope():
+ first, last = get_merge_begin(
+ dest,
+ by * size,
+ middle - start,
+ end - middle,
+ start,
+ middle,
+ diag,
+ step_count,
+ )
+ do_merge(first, last)
+
+ # Call the kernel
+ mergesort(
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ size,
+ width,
+ tvm.tir.indexmod(l2_width, 2) == 0,
+ )
+ nthread_by = axis_mul_before * axis_mul_after
+ nthread_tx = max_threads
+ nthread_bx = ceil_div(size, nthread_tx)
+ ## if the final sorted data ended up in the swap, copy it to the real
output
+ with ib.if_scope(
+ tvm.tir.all(upper_lim > lower_lim, tvm.tir.indexmod(upper_lim -
lower_lim, 2) == 1)
+ ):
+ with ib.new_scope():
+ tx, bx, by = _get_threads(ib, nthread_tx, nthread_bx, nthread_by)
+ tid = bx * nthread_tx + tx
+ idx = by * size + tid
+ with ib.if_scope(tid < size):
+ keys[idx] = keys_swap[idx]
+ if values is not None:
+ values[idx] = values_swap[idx]
+
+
+def sort_ir(
+ data, values_out, values_out_swap, axis, is_ascend, indices_out=None,
indices_out_swap=None
+):
+ """Low level IR to do sorting on the GPU, same usage as
tvm.contrib.sort.argsort on the CPU.
+
+ Parameters
+ ----------
+ data: Buffer
+ Buffer of input data. Data will be sorted in place.
+
+ values_out : Buffer
+ Output buffer of values of sorted tensor with same shape as data.
+
+ values_out_swap : Buffer
+ Output buffer of values with same shape as data to use as swap.
+
+ axis : Int
+ Axis long which to sort the input tensor.
+
+ is_ascend : Boolean
+ Whether to sort in ascending or descending order.
+
+ indicess_out : Buffer
+ Output buffer of indices of sorted tensor with same shape as data.
+
+ indices_out_swap : Buffer
+ Output buffer of indices with same shape as data to use as swap.
+
+ Returns
+ -------
+ stmt : Stmt
+ The result IR statement.
+ """
+ ib = tvm.tir.ir_builder.create()
+ shape = data.shape
+
+ data = ib.buffer_ptr(data)
+ values_out = ib.buffer_ptr(values_out)
+ values_out_swap = ib.buffer_ptr(values_out_swap)
+ if indices_out is not None:
+ indices_out = ib.buffer_ptr(indices_out)
+ assert indices_out_swap is not None
+ indices_out_swap = ib.buffer_ptr(indices_out_swap)
+
+ with ib.if_scope(shape[axis] > 0):
+ axis_mul_before, axis_mul_after = _sort_init(
+ ib,
+ shape,
+ axis,
+ data,
+ values_out,
+ indices_out,
+ value_init_func=lambda _, tid: tvm.tir.generic.cast(tid,
indices_out.dtype),
+ )
+
+ _sort_common(
+ ib,
+ shape[axis],
+ axis_mul_before,
+ axis_mul_after,
+ is_ascend,
+ values_out,
+ values_out_swap,
+ values=indices_out,
+ values_swap=indices_out_swap,
+ )
+
+ return ib.get()
+
+
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with the same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf",
data_alignment=8)
+ value_buf_swap = tvm.tir.decl_buffer(data.shape, data.dtype,
"value_buf_swap", data_alignment=8)
+
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], outs[1], -1, is_ascend),
+ out_buffers=[value_buf, value_buf_swap],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ out = transpose(out, axes)
+
+ return out
+
+
+def sort_thrust(data, axis=-1, is_ascend=1, workspace=None):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with the same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ workspace: Optional[tvm.te.Tensor]
+ A buffer to store intermediate results. The size of the workspace
should be sufficiently
+ large, this can be obtained by overestimation or memory usage
profiling. If None, it will
+ fallback to use thrust internal memory allocation.
+
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf",
data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf",
data_alignment=8)
+
+ def f_compute(ins, outs):
+ args = ["tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend]
+ if workspace is not None:
+ args.append(ins[1])
+ return tvm.tir.call_packed(*args)
+
+ out = te.extern(
+ [data.shape, data.shape],
+ [data] if workspace is None else [data, workspace],
+ ## TODO(mbrookhart): This thrust function is actually doing argsort,
not sort
+ ## For performance, we should probably rename the contrib function and
add
+ ## a pure sort
+ f_compute,
+ out_buffers=[value_buf, indices_buf],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ out = transpose(out, axes)
+ return out
+
+
+def argsort(data, axis=-1, is_ascend=1, dtype="float32", ret_type="indices"):
+ """Performs sorting along the given axis and returns an array of indices
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ ret_type : string, optional
+ The return type [both, indices].
+ "both": return both sorted data and indices.
+ "indices": return sorted indices only.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf",
data_alignment=8)
+ value_swap_buf = tvm.tir.decl_buffer(data.shape, data.dtype,
"value_swap_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf",
data_alignment=8)
+ indices_swap_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_swap_buf",
data_alignment=8)
+
+ outs = te.extern(
+ [data.shape, data.shape, data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(
+ ins[0],
+ outs[0],
+ outs[2],
+ -1,
+ is_ascend,
+ indices_out=outs[1],
+ indices_out_swap=outs[3],
+ ),
+ out_buffers=[value_buf, indices_buf, value_swap_buf, indices_swap_buf],
+ name="argsort_gpu",
+ tag="argsort_gpu",
+ )
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ outs = [transpose(out, axes) for out in outs]
+
+ if ret_type == "indices":
+ return outs[1]
+
+ return outs[0], outs[1]
+
+
+def argsort_thrust(data, axis=-1, is_ascend=1, dtype="float32",
ret_type="indices", workspace=None):
+ """Performs sorting along the given axis and returns an array of indices
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ ret_type : string, optional
+ The return type [both, indices].
+ "both": return both sorted data and indices.
+ "indices": return sorted indices only.
+
+ workspace : Optional[tvm.te.Tensor]
+ A buffer to store intermediate results. The size of the workspace
should be sufficiently
+ large, this can be obtained by overestimation or memory usage
profiling. If None, it will
+ fallback to use thrust internal memory allocation.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ return topk_thrust(data, 0, axis, ret_type, is_ascend, dtype, workspace)
+
+
+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, 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"]
+ ndim = len(data.shape)
+ axis = axis + ndim if axis < 0 else axis
+ assert 0 <= axis < ndim
+ dshape = data.shape
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ values_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "values_buf",
data_alignment=8)
+ values_swap_buf = tvm.tir.decl_buffer(
+ data.shape, data.dtype, "values_swap_buf", data_alignment=8
+ )
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "indices_buf",
data_alignment=8)
+ indices_swap_buf = tvm.tir.decl_buffer(data.shape, dtype,
"indies_swap_buf", data_alignment=8)
+
+ if ret_type == "values":
+ output = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], outs[1], -1, is_ascend),
+ out_buffers=[values_buf, values_swap_buf],
+ name="topk_gpu",
+ tag="topk_gpu",
+ )[0]
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ output = transpose(output, axes)
+ else:
+ output = te.extern(
+ [data.shape, data.shape, data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(
+ ins[0],
+ outs[0],
+ outs[2],
+ -1,
+ is_ascend,
+ indices_out=outs[1],
+ indices_out_swap=outs[3],
+ ),
+ out_buffers=[values_buf, indices_buf, values_swap_buf,
indices_swap_buf],
+ name="topk_gpu",
+ tag="topk_gpu",
+ )[0:2]
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ output[0] = transpose(output[0], axes)
+ output[1] = transpose(output[1], axes)
+
+ if isinstance(k, int) and k < 1:
+ if ret_type == "indices":
+ return output[1]
+ return output
+ beg = [0] * ndim
+ end = []
+ strides = [1] * ndim
+ for i in range(ndim):
+ if i == axis:
+ end.append(k if isinstance(k, int) else tvm.te.size_var("dim"))
+ else:
+ end.append(dshape[i])
+ if ret_type == "both":
+ values_out, indices_out = output
+ values_out = strided_slice(values_out, beg, end, strides)
+ indices_out = strided_slice(indices_out, beg, end, strides)
+ output = [values_out, indices_out]
+ elif ret_type == "values":
+ output = [strided_slice(output, beg, end, strides)]
+ else: # ret_type == "indices"
+ indices_out = output[1]
+ output = [strided_slice(indices_out, beg, end, strides)]
+ return output
+
+
+def topk_thrust(
+ data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64",
workspace=None
+):
+ """Get the top k elements in an input tensor along the given axis.
+
+ Parameters
+ ----------
+ data : tvm.te.Tensor
+ The input tensor.
+
+ k : int, 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.
+
+ workspace : Optional[tvm.te.Tensor]
+ A buffer to store intermediate results. The size of the workspace
should be sufficiently
+ large, this can be obtained by overestimation or memory usage
profiling. If None, it will
+ fallback to use thrust internal memory allocation.
+
+ Returns
+ -------
+ out : tvm.te.Tensor or List[tvm.te.Tensor]
+ The computed result.
+ """
+ assert ret_type in ["both", "values", "indices"]
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ data_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "data_buf",
data_alignment=8)
+ if workspace is not None:
+ workspace_buf = tvm.tir.decl_buffer(
+ workspace.shape, workspace.dtype, "workspace_buf", data_alignment=8
+ )
+ else:
+ workspace_buf = None
+ out_bufs = [
+ tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf",
data_alignment=8),
+ tvm.tir.decl_buffer(data.shape, dtype, "indices_buf",
data_alignment=8),
+ ]
+
+ def f_compute(ins, outs):
+ args = ["tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend]
+ if workspace is not None:
+ args.append(ins[1])
+ return tvm.tir.call_packed(*args)
+
+ is_ascend = 1 if is_ascend else 0
+
+ out = te.extern(
+ [data.shape, data.shape],
+ [data] if workspace is None else [data, workspace],
+ f_compute,
+ in_buffers=[data_buf] if workspace is None else [data_buf,
workspace_buf],
+ out_buffers=out_bufs,
+ name="topk_gpu",
+ tag="topk_gpu",
+ )
+
+ if isinstance(k, tvm.tir.IntImm):
+ k = k.value
+
+ if not isinstance(k, int) or k > 0:
+ beg = [0] * ndim
+ end = data.shape[:-1] + [k if isinstance(k, int) else
tvm.te.size_var("dim")]
+ strides = [1] * ndim
+ out = [strided_slice(o, beg, end, strides) for o in out]
+
+ if axis != ndim - 1:
+ axes = swap(list(range(ndim)), axis)
+ out = [transpose(o, axes) for o in out]
+
+ if ret_type == "values":
+ out = out[0]
+ elif ret_type == "indices":
+ out = out[1]
+
+ return out
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index 1efbd690f0..4fe6de9e09 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -93,13 +93,13 @@ def test_dispatch_scanop_cuda():
with bb.function("main", (x,), {"global_symbol": "main"}):
with bb.dataflow():
lv = bb.emit_te(
- topi.cuda.cumsum,
+ topi.gpu.cumsum,
x,
axis=1,
exclusive=True,
)
out = bb.emit_te(
- topi.cuda.cumprod,
+ topi.gpu.cumprod,
lv,
axis=1,
)
@@ -178,7 +178,7 @@ def test_dispatch_sort_cuda():
with bb.function("foo", (x,), {"global_symbol": "foo"}):
with bb.dataflow():
out = bb.emit_te(
- topi.cuda.sort,
+ topi.gpu.sort,
x,
axis=1,
)
@@ -193,14 +193,14 @@ def test_dispatch_sort_cuda():
)
)
out = bb.emit_te(
- topi.cuda.sort_thrust,
+ topi.gpu.sort_thrust,
y,
axis=0,
is_ascend=False,
workspace=workspace,
)
else:
- out = bb.emit_te(topi.cuda.sort, y, axis=0,
is_ascend=False)
+ out = bb.emit_te(topi.gpu.sort, y, axis=0, is_ascend=False)
out = bb.emit_output(out)
bb.emit_func_output(out)
expected_mod = bb.finalize()
@@ -273,7 +273,7 @@ def test_dispatch_argsort_cuda():
with target:
with bb.function("foo", (x,), {"global_symbol": "foo"}):
with bb.dataflow():
- out = bb.emit_te(topi.cuda.argsort, x, axis=1, is_ascend=True,
dtype="int32")
+ out = bb.emit_te(topi.gpu.argsort, x, axis=1, is_ascend=True,
dtype="int32")
out = bb.emit_output(out)
bb.emit_func_output(out)
with bb.function("foo2", (y,), {"global_symbol": "foo2"}):
@@ -285,7 +285,7 @@ def test_dispatch_argsort_cuda():
)
)
out = bb.emit_te(
- topi.cuda.argsort_thrust,
+ topi.gpu.argsort_thrust,
y,
axis=0,
is_ascend=False,
@@ -293,7 +293,7 @@ def test_dispatch_argsort_cuda():
workspace=workspace,
)
else:
- out = bb.emit_te(topi.cuda.argsort, y, axis=0,
is_ascend=False, dtype="int64")
+ out = bb.emit_te(topi.gpu.argsort, y, axis=0,
is_ascend=False, dtype="int64")
out = bb.emit_output(out)
bb.emit_func_output(out)
expected_mod = bb.finalize()
@@ -357,7 +357,7 @@ def test_dispatch_topk_cuda():
with target:
with bb.function("foo", (x,), {"global_symbol": "foo"}):
with bb.dataflow():
- out = bb.emit_te(topi.cuda.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
+ out = bb.emit_te(topi.gpu.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
out = bb.emit_output(out)
bb.emit_func_output(out)
expected_mod = bb.finalize()
@@ -393,8 +393,8 @@ def test_dispatch_topk_gpu():
with target:
with bb.function("foo", (x,), {"global_symbol": "foo"}):
with bb.dataflow():
- lv0 = bb.emit_te(topi.cuda.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
- lv1 = bb.emit_te(topi.cuda.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
+ lv0 = bb.emit_te(topi.gpu.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
+ lv1 = bb.emit_te(topi.gpu.topk, x, k=2, axis=1,
is_ascend=False, dtype="int32")
out = (lv0, lv1)
out = bb.emit_output(out)
bb.emit_func_output(out)
