Function ir

python/tvm/topi/gpu/sort.py:1122–1163

(sorted_sequence_buf, values_buf, indices_buf)

            )

    def ir(sorted_sequence_buf, values_buf, indices_buf):
        with IRBuilder() as ib:
            sorted_sequence_shape = sorted_sequence_buf.shape
            values_shape = values_buf.shape
            num_search = prod(values_shape)
            search_range = sorted_sequence_shape[-1]

            sorted_sequence_ptr = T.buffer_proxy(sorted_sequence_buf)
            values_ptr = T.buffer_proxy(values_buf)
            indices_ptr = T.buffer_proxy(indices_buf)

            max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"])
            nthread_tx = max_threads
            nthread_bx = ceil_div(num_search, nthread_tx)
            tx = te.thread_axis("threadIdx.x")
            bx = te.thread_axis("blockIdx.x")
            with T.frame_scope(
                [
                    T.attr(tx, "thread_extent", nthread_tx),
                    T.attr(bx, "thread_extent", nthread_bx),
                ]
            ):
                tid = bx * nthread_tx + tx

                with T.If(tid < num_search):
                    with T.Then():
                        if len(sorted_sequence_shape) == 1:
                            sequence_offset = 0
                        else:
                            sequence_id = tid // values_shape[-1]
                            sequence_offset = sequence_id * search_range

                        indices_ptr[tid] = binary_search(
                            sequence_offset,
                            search_range,
                            sorted_sequence_ptr,
                            values_ptr[tid],
                            right,
                            out_dtype,
                        )

            return ib.get()

    return te.extern(
        values.shape,