| # 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 .injective import schedule_injective_from_existing |
| from ..transform import strided_slice, transpose |
| from .. import tag |
| from ..utils import ceil_div, swap |
| from ..math import cast, ceil_log2 |
| |
| |
| def _schedule_sort(outs): |
| """Schedule for argsort operator. |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of argsort |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for the op. |
| """ |
| outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs |
| s = te.create_schedule([x.op for x in outs]) |
| scheduled_ops = [] |
| |
| def traverse(op): |
| if tag.is_injective(op.tag): |
| schedule_injective_from_existing(s, op.output(0)) |
| for tensor in op.input_tensors: |
| if tensor.op.input_tensors and tensor.op not in scheduled_ops: |
| traverse(tensor.op) |
| scheduled_ops.append(op) |
| |
| for out in outs: |
| traverse(out.op) |
| return s |
| |
| |
| def _get_threads(ib, nthread_tx, nthread_bx, nthread_by, nthread_bz): |
| 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") |
| bz = te.thread_axis("blockIdx.z") |
| ib.scope_attr(by, "thread_extent", nthread_by) |
| ib.scope_attr(bz, "thread_extent", nthread_bz) |
| |
| return tx, bx, by, bz |
| |
| |
| 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 |
| nthread_bz = axis_mul_after |
| |
| # Copy the keys_in to initial output |
| with ib.new_scope(): |
| tx, bx, by, bz = _get_threads(ib, nthread_tx, nthread_bx, nthread_by, nthread_bz) |
| tid = bx * nthread_tx + tx |
| 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 |
| nthread_bz = axis_mul_after |
| with ib.new_scope(): |
| ib.scope_attr(tvm.tir.const(0), "hand_threaded", 0) |
| tx, bx, by, bz = _get_threads(ib, nthread_tx, nthread_bx, nthread_by, nthread_bz) |
| 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_bz = 1 |
| 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, bz = _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, bz = _get_threads(ib, ntx, nbx, nthread_by, nbz) |
| |
| 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 |
| nthread_bz = 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, bz = _get_threads(ib, nthread_tx, nthread_bx, nthread_by, nthread_bz) |
| tid = bx * nthread_tx + tx |
| idx = (by * axis_mul_after + bz) * 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_by_key_ir( |
| keys_in, values_in, keys_out, values_out, keys_out_swap, values_out_swap, axis, is_ascend |
| ): |
| """Low level IR to do sort by key on the GPU. |
| |
| Parameters |
| ---------- |
| keys_in: Buffer |
| Buffer of input keys. |
| |
| values_in: Buffer |
| Buffer of input keys. |
| |
| keys_out : Buffer |
| Buffer of output sorted keys. |
| |
| values_out : Buffer |
| Buffer of output sorted values. |
| |
| keys_out_swap : Buffer |
| Output buffer of values with same shape as keys_in to use as swap. |
| |
| values_out_swap : Buffer |
| Output buffer of values with same shape as values_in 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 keys_in. |
| |
| values_out_swap : Buffer |
| Output buffer of indices with same shape as keys_in to use as swap. |
| |
| Returns |
| ------- |
| stmt : Stmt |
| The result IR statement. |
| """ |
| ib = tvm.tir.ir_builder.create() |
| shape = keys_in.shape |
| |
| keys_in = ib.buffer_ptr(keys_in) |
| values_in = ib.buffer_ptr(values_in) |
| keys_out = ib.buffer_ptr(keys_out) |
| keys_out_swap = ib.buffer_ptr(keys_out_swap) |
| values_out = ib.buffer_ptr(values_out) |
| values_out_swap = ib.buffer_ptr(values_out_swap) |
| |
| with ib.if_scope(shape[axis] > 0): |
| axis_mul_before, axis_mul_after = _sort_init( |
| ib, |
| shape, |
| axis, |
| keys_in, |
| keys_out, |
| values_out, |
| value_init_func=lambda idx, _: values_in[idx], |
| ) |
| |
| _sort_common( |
| ib, |
| shape[axis], |
| axis_mul_before, |
| axis_mul_after, |
| is_ascend, |
| keys_out, |
| keys_out_swap, |
| values=values_out, |
| values_swap=values_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): |
| """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. |
| """ |
| 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) |
| out = te.extern( |
| [data.shape, data.shape], |
| [data], |
| ## 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 |
| lambda ins, outs: tvm.tir.call_packed( |
| "tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend |
| ), |
| 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"): |
| """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. |
| """ |
| return topk_thrust(data, 0, axis, ret_type, is_ascend, dtype) |
| |
| |
| def schedule_sort(outs): |
| """Schedule for sort operator. |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of argsort |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for the op. |
| """ |
| return _schedule_sort(outs) |
| |
| |
| def schedule_argsort(outs): |
| """Schedule for argsort operator. |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of argsort |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for the op. |
| """ |
| return _schedule_sort(outs) |
| |
| |
| 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"): |
| """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 = 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) |
| 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), |
| ] |
| |
| is_ascend = 1 if is_ascend else 0 |
| |
| out = te.extern( |
| [data.shape, data.shape], |
| [data], |
| lambda ins, outs: tvm.tir.call_packed( |
| "tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend |
| ), |
| in_buffers=[data_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 |
| |
| |
| def schedule_topk(outs): |
| """Schedule for argsort operator. |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of argsort |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for the op. |
| """ |
| return _schedule_sort(outs) |
| |
| |
| def sort_by_key(keys, values, axis=-1, is_ascend=1): |
| """Sort values with respect to keys. Both keys and values will |
| be sorted and returned. |
| |
| Parameters |
| ---------- |
| keys: tvm.te.Tensor |
| The input keys. |
| |
| values : tvm.te.Tensor, |
| The input values. |
| |
| axis : int, optional |
| Axis long which to sort the input tensor. |
| |
| is_ascend : boolean, optional |
| Whether to sort in ascending or descending order. |
| |
| Returns |
| ------- |
| keys_sorted : tvm.te.Tensor |
| The sorted keys |
| |
| values_sorted : tvm.te.Tensor |
| The values sorted with respect to the keys |
| """ |
| keys_buf = tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8) |
| values_buf = tvm.tir.decl_buffer(values.shape, values.dtype, "values_buf", data_alignment=8) |
| |
| out_bufs = [ |
| tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8), |
| tvm.tir.decl_buffer(values.shape, values.dtype, "values_buf", data_alignment=8), |
| tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_swap_buf", data_alignment=8), |
| tvm.tir.decl_buffer(values.shape, values.dtype, "values_swap_buf", data_alignment=8), |
| ] |
| out = te.extern( |
| [keys.shape, values.shape, keys.shape, values.shape], |
| [keys, values], |
| lambda ins, outs: sort_by_key_ir( |
| ins[0], ins[1], outs[0], outs[1], outs[2], outs[3], axis, is_ascend |
| ), |
| in_buffers=[keys_buf, values_buf], |
| out_buffers=out_bufs, |
| dtype=[keys.dtype, values.dtype], |
| name="sort_by_key", |
| tag="sort_by_key", |
| ) |
| return out[0], out[1] |
| |
| |
| def stable_sort_by_key_thrust(keys, values, for_scatter=False): |
| """Sort values with respect to keys using thrust. |
| Both keys and values will be sorted and returned. |
| Sorting is done via stable sort, so relative ordering among |
| ties are preserved. |
| |
| Parameters |
| ---------- |
| keys: tvm.te.Tensor |
| The 1D input keys. |
| |
| values : tvm.te.Tensor, |
| The 1D input values. |
| |
| for_scatter: bool, optional |
| If True, negative keys are interpreted as negative indices. |
| Before sorting, negative indices are converted to corresponding positive indices. |
| The output keys (indices) are all positive. |
| This option is introduced to optimize the scatter implementation. |
| |
| Returns |
| ------- |
| keys_sorted : tvm.te.Tensor |
| The sorted keys |
| |
| values_sorted : tvm.te.Tensor |
| The values sorted with respect to the keys |
| """ |
| keys_buf = tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8) |
| values_buf = tvm.tir.decl_buffer(values.shape, values.dtype, "values_buf", data_alignment=8) |
| out_bufs = [ |
| tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8), |
| tvm.tir.decl_buffer(keys.shape, values.dtype, "values_buf", data_alignment=8), |
| ] |
| out = te.extern( |
| [keys.shape, values.shape], |
| [keys, values], |
| lambda ins, outs: tvm.tir.call_packed( |
| "tvm.contrib.thrust.stable_sort_by_key", ins[0], ins[1], outs[0], outs[1], for_scatter |
| ), |
| in_buffers=[keys_buf, values_buf], |
| out_buffers=out_bufs, |
| dtype=[keys.dtype, values.dtype], |
| name="stable_sort_by_key", |
| tag="stable_sort_by_key", |
| ) |
| return out[0], out[1] |
