You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
127 lines
3.6 KiB
C++
127 lines
3.6 KiB
C++
#include "searchsorted_cpu_wrapper.h"
|
|
#include <stdio.h>
|
|
|
|
template<typename scalar_t>
|
|
int eval(scalar_t val, scalar_t *a, int64_t row, int64_t col, int64_t ncol, bool side_left)
|
|
{
|
|
/* Evaluates whether a[row,col] < val <= a[row, col+1]*/
|
|
|
|
if (col == ncol - 1)
|
|
{
|
|
// special case: we are on the right border
|
|
if (a[row * ncol + col] <= val){
|
|
return 1;}
|
|
else {
|
|
return -1;}
|
|
}
|
|
bool is_lower;
|
|
bool is_next_higher;
|
|
|
|
if (side_left) {
|
|
// a[row, col] < v <= a[row, col+1]
|
|
is_lower = (a[row * ncol + col] < val);
|
|
is_next_higher = (a[row*ncol + col + 1] >= val);
|
|
} else {
|
|
// a[row, col] <= v < a[row, col+1]
|
|
is_lower = (a[row * ncol + col] <= val);
|
|
is_next_higher = (a[row * ncol + col + 1] > val);
|
|
}
|
|
if (is_lower && is_next_higher) {
|
|
// we found the right spot
|
|
return 0;
|
|
} else if (is_lower) {
|
|
// answer is on the right side
|
|
return 1;
|
|
} else {
|
|
// answer is on the left side
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
template<typename scalar_t>
|
|
int64_t binary_search(scalar_t*a, int64_t row, scalar_t val, int64_t ncol, bool side_left)
|
|
{
|
|
/* Look for the value `val` within row `row` of matrix `a`, which
|
|
has `ncol` columns.
|
|
|
|
the `a` matrix is assumed sorted in increasing order, row-wise
|
|
|
|
returns:
|
|
* -1 if `val` is smaller than the smallest value found within that row of `a`
|
|
* `ncol` - 1 if `val` is larger than the largest element of that row of `a`
|
|
* Otherwise, return the column index `res` such that:
|
|
- a[row, col] < val <= a[row, col+1]. (if side_left), or
|
|
- a[row, col] < val <= a[row, col+1] (if not side_left).
|
|
*/
|
|
|
|
//start with left at 0 and right at number of columns of a
|
|
int64_t right = ncol;
|
|
int64_t left = 0;
|
|
|
|
while (right >= left) {
|
|
// take the midpoint of current left and right cursors
|
|
int64_t mid = left + (right-left)/2;
|
|
|
|
// check the relative position of val: are we good here ?
|
|
int rel_pos = eval(val, a, row, mid, ncol, side_left);
|
|
// we found the point
|
|
if(rel_pos == 0) {
|
|
return mid;
|
|
} else if (rel_pos > 0) {
|
|
if (mid==ncol-1){return ncol-1;}
|
|
// the answer is on the right side
|
|
left = mid;
|
|
} else {
|
|
if (mid==0){return -1;}
|
|
right = mid;
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
void searchsorted_cpu_wrapper(
|
|
at::Tensor a,
|
|
at::Tensor v,
|
|
at::Tensor res,
|
|
bool side_left)
|
|
{
|
|
|
|
// Get the dimensions
|
|
auto nrow_a = a.size(/*dim=*/0);
|
|
auto ncol_a = a.size(/*dim=*/1);
|
|
auto nrow_v = v.size(/*dim=*/0);
|
|
auto ncol_v = v.size(/*dim=*/1);
|
|
|
|
auto nrow_res = fmax(nrow_a, nrow_v);
|
|
|
|
//auto acc_v = v.accessor<float, 2>();
|
|
//auto acc_res = res.accessor<float, 2>();
|
|
|
|
AT_DISPATCH_ALL_TYPES(a.type(), "searchsorted cpu", [&] {
|
|
|
|
scalar_t* a_data = a.data_ptr<scalar_t>();
|
|
scalar_t* v_data = v.data_ptr<scalar_t>();
|
|
int64_t* res_data = res.data<int64_t>();
|
|
|
|
for (int64_t row = 0; row < nrow_res; row++)
|
|
{
|
|
for (int64_t col = 0; col < ncol_v; col++)
|
|
{
|
|
// get the value to look for
|
|
int64_t row_in_v = (nrow_v == 1) ? 0 : row;
|
|
int64_t row_in_a = (nrow_a == 1) ? 0 : row;
|
|
|
|
int64_t idx_in_v = row_in_v * ncol_v + col;
|
|
int64_t idx_in_res = row * ncol_v + col;
|
|
|
|
// apply binary search
|
|
res_data[idx_in_res] = (binary_search(a_data, row_in_a, v_data[idx_in_v], ncol_a, side_left) + 1);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("searchsorted_cpu_wrapper", &searchsorted_cpu_wrapper, "searchsorted (CPU)");
|
|
}
|