searchsorted_cpu_contiguous Class — pytorch Architecture

Source Code

aten/src/ATen/native/Bucketization.cpp lines 87–115

template<typename input_t, typename output_t>
void searchsorted_cpu_contiguous(Tensor& result, const Tensor& input, const Tensor& boundaries, const bool& right, const Tensor& sorter) {
  int64_t numel_in = input.numel();
  bool is_scalar_input = input.dim() == 0 && numel_in == 1;
  // inner most dim size of input and boundaries
  int64_t idim_in = is_scalar_input ? 1 : input.sizes().back();
  int64_t idim_bd = boundaries.sizes().back();

  const input_t *data_in = input.const_data_ptr<input_t>();
  const input_t *data_bd = boundaries.const_data_ptr<input_t>();
  const int64_t *data_st = sorter.defined() ? sorter.const_data_ptr<int64_t>() : nullptr;
  output_t *data_out = result.data_ptr<output_t>();

  bool is_1d_boundaries = boundaries.dim() == 1;
  at::parallel_for(0, numel_in, SEARCHSORTED_GRAIN_SIZE, [&](int64_t start, int64_t end) {
    for (const auto i : c10::irange(start, end)) {
      // If boundaries tensor is 1d, we always search the entire boundary tensor
      int64_t start_bd = is_1d_boundaries ? 0 : i / idim_in * idim_bd;
      int64_t end_bd = start_bd + idim_bd;

      int64_t pos = !right ?
        cus_lower_bound(start_bd, end_bd, data_in[i], data_bd, data_st) - start_bd :
        cus_upper_bound(start_bd, end_bd, data_in[i], data_bd, data_st) - start_bd;

      // type conversion might happen here
      data_out[i] = pos;
    }
  });
}