Function serial_merge

python/tvm/topi/gpu/sort.py:257–313

(
        source,
        dest,
        source_idx,
        dest_idx,
        base_idx,
        aCount,
        bCount,
        aStart,
        bStart,
        kStart,
        diag,
        step_count,
        first,
        last,
        i_buf,
        j_buf,
    )

                    last[0] = mid

    def serial_merge(
        source,
        dest,
        source_idx,
        dest_idx,
        base_idx,
        aCount,
        bCount,
        aStart,
        bStart,
        kStart,
        diag,
        step_count,
        first,
        last,
        i_buf,
        j_buf,
    ):
        i_val = aStart + first[0]
        j_val = bStart + diag - last[0]
        if is_webgpu:
            i_buf[0] = cast(i_val, target_dtype)
            j_buf[0] = cast(j_val, target_dtype)
        else:
            i_buf[0] = i_val
            j_buf[0] = j_val

        with T.serial(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
            i_idx = base_idx + i_buf[0]
            j_idx = base_idx + j_buf[0]
            k_idx = base_idx + (kStart + diag + count)

            with T.If(tvm.tirx.all(i_buf[0] < aStart + aCount, j_buf[0] < bStart + bCount)):
                with T.Then():
                    with T.If(compare(source[i_idx], source[j_idx])):
                        with T.Then():
                            dest[k_idx] = source[i_idx]
                            if values is not None:
                                dest_idx[k_idx] = source_idx[i_idx]
                            i_buf[0] = i_buf[0] + 1
                        with T.Else():
                            dest[k_idx] = source[j_idx]
                            if values is not None:
                                dest_idx[k_idx] = source_idx[j_idx]
                            j_buf[0] = j_buf[0] + 1
                with T.Else():
                    with T.If(i_buf[0] < aStart + aCount):
                        with T.Then():
                            dest[k_idx] = source[i_idx]
                            if values is not None:
                                dest_idx[k_idx] = source_idx[i_idx]
                            i_buf[0] = i_buf[0] + 1
                        with T.Else():
                            dest[k_idx] = source[j_idx]
                            if values is not None:
                                dest_idx[k_idx] = source_idx[j_idx]
                            j_buf[0] = j_buf[0] + 1