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,
(data, axis=-1, is_ascend=1, workspace=None)
| 705 | |
| 706 | |
| 707 | def sort_thrust(data, axis=-1, is_ascend=1, workspace=None): |
| 708 | """Performs sorting along the given axis and returns an array of |
| 709 | sorted values with the same shape as the input data. |
| 710 | |
| 711 | Parameters |
| 712 | ---------- |
| 713 | data: tvm.te.Tensor |
| 714 | The input array. |
| 715 | |
| 716 | axis : int, optional |
| 717 | Axis long which to sort the input tensor. |
| 718 | |
| 719 | is_ascend : boolean, optional |
| 720 | Whether to sort in ascending or descending order. |
| 721 | |
| 722 | workspace: Optional[tvm.te.Tensor] |
| 723 | A buffer to store intermediate results. The size of the workspace should be sufficiently |
| 724 | large, this can be obtained by overestimation or memory usage profiling. If None, it will |
| 725 | fallback to use thrust internal memory allocation. |
| 726 | |
| 727 | |
| 728 | Returns |
| 729 | ------- |
| 730 | out : tvm.te.Tensor |
| 731 | The output of this function. |
| 732 | """ |
| 733 | dtype = "float32" |
| 734 | ndim = len(data.shape) |
| 735 | axis = ndim + axis if axis < 0 else axis |
| 736 | |
| 737 | if axis != ndim - 1: |
| 738 | # Prepare for sorting along axis -1. |
| 739 | axes = swap(list(range(ndim)), axis) |
| 740 | data = transpose(data, axes) |
| 741 | |
| 742 | value_buf = tvm.tirx.decl_buffer( |
| 743 | data.shape, data.dtype, "value_buf", data_alignment=8, layout=None |
| 744 | ) |
| 745 | indices_buf = tvm.tirx.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8, layout=None) |
| 746 | |
| 747 | def f_compute(ins, outs): |
| 748 | args = ["tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend] |
| 749 | if workspace is not None: |
| 750 | args.append(ins[1]) |
| 751 | return tvm.tirx.call_packed(*args) |
| 752 | |
| 753 | out = te.extern( |
| 754 | [data.shape, data.shape], |
| 755 | [data] if workspace is None else [data, workspace], |
| 756 | ## TODO(mbrookhart): This thrust function is actually doing argsort, not sort |
| 757 | ## For performance, we should probably rename the contrib function and add |
| 758 | ## a pure sort |
| 759 | f_compute, |
| 760 | out_buffers=[value_buf, indices_buf], |
| 761 | name="sort_gpu", |
| 762 | tag="sort_gpu", |
| 763 | )[0] |
| 764 |