Repository: numpy/numpy
File: numpy/_core/src/npysort/binsearch.cpp (C++)

Code before:
        |keys| reads.
        3. In the first iteration, all elements are compared against the
        median. So we can store it in a variable and use it for all keys.

    This initial block replaces the initialization loop that is used for the
    arr_len==0 case. Note that when arr_len = 1, then half is 0 so the
    following block initializes the array as with 0s.
    */
    npy_intp interval_length = arr_len;
    npy_intp half = interval_length >> 1;
    interval_length -= half; // length -> ceil(length / 2)

    npy_intp base = 0;
    const T mid_val = *(const T *)(arr + (base + half) * arr_str);

    for (npy_intp i = 0; i < key_len; ++i) {
        const T key_val = *(const T *)(key + i * key_str);
        *(npy_intp *)(ret + i * ret_str) = cmp(mid_val, key_val) * half;
    }

    while (interval_length > 1) {
        npy_intp half = interval_length >> 1;
        interval_length -= half; // length -> ceil(length / 2)

        for (npy_intp i = 0; i < key_len; ++i) {
            npy_intp &base = *(npy_intp *)(ret + i * ret_str);
            const T mid_val = *(const T *)(arr + (base + half) * arr_str);
            const T key_val = *(const T *)(key + i * key_str);
            base += cmp(mid_val, key_val) * half;
        }
    }

    /*
    At this point interval_length == 1, so the candidates are in interval
    [base, base + 1].

    We have two options:
        If cmp(arr[base], key_val) == true, insertion index is base + 1
        Otherwise the insertion order is just base
    */
    for (npy_intp i = 0; i < key_len; ++i) {
        npy_intp &base = *(npy_intp *)(ret + i * ret_str);
        const T key_val = *(const T *)(key + i * key_str);
        base += cmp(*(const T *)(arr + base * arr_str), key_val);
    }
}

template <class Tag, side_t side>
static int
argbinsearch(const char *arr, const char *key, const char *sort, char *ret,
             npy_intp arr_len, npy_intp key_len, npy_intp arr_str,

Review comment (performance):
This is perhaps the reason why the benchmarks for a single key are a bit slower. We on each iteration we are updating the memory, which is not needed on current main.

If performance is a concern, we could special case the code for the case `key_len == 1`. But the difference is quite small, so to keep the code simple I am fine with leaving as is.

Code after:
        |keys| reads.
        3. In the first iteration, all elements are compared against the
        median. So we can store it in a variable and use it for all keys.

    This initial block replaces the initialization loop that is used for the
    arr_len==0 case. Note that when arr_len = 1, then half is 0 so the
    following block initializes the array as with 0s.
    */
    npy_intp interval_length = arr_len;
    npy_intp half = interval_length >> 1;
    interval_length -= half; // length -> ceil(length / 2)

    npy_intp base = 0;
    const T mid_val = *(const T *)(arr + (base + half) * arr_str);

    for (npy_intp i = 0; i < key_len; ++i) {
        const T key_val = *(const T *)(key + i * key_str);
        *(npy_intp *)(ret + i * ret_str) = cmp(mid_val, key_val) * half;
    }

    while (interval_length > 1) {
        npy_intp half = interval_length >> 1;
        interval_length -= half; // length -> ceil(length / 2)

        for (npy_intp i = 0; i < key_len; ++i) {
            npy_intp &base = *(npy_intp *)(ret + i * ret_str);
            const T mid_val = *(const T *)(arr + (base + half) * arr_str);
            const T key_val = *(const T *)(key + i * key_str);
            base += cmp(mid_val, key_val) * half;
        }
    }

    /*
    At this point interval_length == 1, so the candidates are in the 
    interval [base, base + 1].

    We have two options:
        If cmp(arr[base], key_val) == true, insertion index is base + 1
        Otherwise the insertion order is just base
    */
    for (npy_intp i = 0; i < key_len; ++i) {
        npy_intp &base = *(npy_intp *)(ret + i * ret_str);
        const T key_val = *(const T *)(key + i * key_str);
        base += cmp(*(const T *)(arr + base * arr_str), key_val);
    }
}

template <class Tag, side_t side>
static int
argbinsearch(const char *arr, const char *key, const char *sort, char *ret,
             npy_intp arr_len, npy_intp key_len, npy_intp arr_str,