iterator_factory.cuh

/*
 * Copyright (c) 2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <cuspatial/detail/functors.cuh>
#include <cuspatial/error.hpp>
#include <cuspatial/geometry/box.hpp>
#include <cuspatial/geometry/vec_2d.hpp>
#include <cuspatial/traits.hpp>
 
#include <thrust/binary_search.h>
#include <thrust/detail/raw_reference_cast.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/transform_output_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#include <thrust/type_traits/is_contiguous_iterator.h>
 
#include <type_traits>
 
namespace cuspatial {
namespace detail {
 
template <typename IndexType, typename UnaryFunction>
inline CUSPATIAL_HOST_DEVICE auto make_counting_transform_iterator(IndexType start, UnaryFunction f)
{
  return thrust::make_transform_iterator(thrust::make_counting_iterator(start), f);
}
 
template <typename T, typename VectorType = vec_2d<T>>
struct tuple_to_vec_2d {
  __device__ VectorType operator()(thrust::tuple<T, T> const& pos)
  {
    return VectorType{thrust::get<0>(pos), thrust::get<1>(pos)};
  }
};
 
template <typename T, typename VectorType = vec_2d<T>>
struct vec_2d_to_tuple {
  __device__ thrust::tuple<T, T> operator()(VectorType const& xy)
  {
    return thrust::make_tuple(xy.x, xy.y);
  }
};
 
template <typename T, typename Box = box<T>>
struct box_to_tuple {
  __device__ thrust::tuple<T, T, T, T> operator()(Box const& box)
  {
    return thrust::make_tuple(box.v1.x, box.v1.y, box.v2.x, box.v2.y);
  }
};
 
template <typename T, typename Vertex = vec_2d<T>>
struct vec_2d_tuple_to_box {
  __device__ box<T, Vertex> operator()(thrust::tuple<Vertex, Vertex> pair)
  {
    auto v1 = thrust::get<0>(pair);
    auto v2 = thrust::get<1>(pair);
    return {v1, v2};
  }
};
 
template <typename Iter, typename Enable = void>
struct interleaved_to_vec_2d {
  using element_t  = typename std::iterator_traits<Iter>::value_type;
  using value_type = vec_2d<element_t>;
  Iter it;
  constexpr interleaved_to_vec_2d(Iter it) : it{it} {}
 
  CUSPATIAL_HOST_DEVICE value_type operator()(std::size_t i)
  {
    return vec_2d<element_t>{it[2 * i], it[2 * i + 1]};
  }
};
 
template <typename Iter>
struct interleaved_to_vec_2d<Iter,
                             typename std::enable_if_t<thrust::is_contiguous_iterator_v<Iter>>> {
  using element_t  = typename std::iterator_traits<Iter>::value_type;
  using value_type = vec_2d<element_t>;
 
  element_t const* ptr;
 
  constexpr interleaved_to_vec_2d(Iter it) : ptr{&thrust::raw_reference_cast(*it)}
  {
    CUSPATIAL_EXPECTS(!((intptr_t)ptr % alignof(vec_2d<element_t>)),
                      "Misaligned interleaved data.");
  }
 
  CUSPATIAL_HOST_DEVICE value_type operator()(std::size_t i)
  {
    auto const aligned =
      static_cast<element_t const*>(__builtin_assume_aligned(ptr + 2 * i, 2 * sizeof(element_t)));
    return vec_2d<element_t>{aligned[0], aligned[1]};
  }
};
 
template <int stride>
struct strided_functor {
  auto __device__ operator()(std::size_t i) { return i * stride; }
};
 
template <class IndexT, class GeometryIter>
struct index_to_geometry_id {
  GeometryIter geometry_begin;
  GeometryIter geometry_end;
 
  index_to_geometry_id(GeometryIter begin, GeometryIter end)
    : geometry_begin(begin), geometry_end(end)
  {
  }
 
  CUSPATIAL_HOST_DEVICE auto operator()(IndexT idx)
  {
    return thrust::distance(geometry_begin,
                            thrust::upper_bound(thrust::seq, geometry_begin, geometry_end, idx));
  }
};
 
}  // namespace detail
 
template <typename FirstIter, typename SecondIter>
auto make_vec_2d_iterator(FirstIter first, SecondIter second)
{
  using T = typename std::iterator_traits<FirstIter>::value_type;
  static_assert(is_same<T, iterator_value_type<SecondIter>>(), "Iterator value_type mismatch");
 
  auto zipped = thrust::make_zip_iterator(first, second);
  return thrust::make_transform_iterator(zipped, detail::tuple_to_vec_2d<T>());
}
 
template <typename Iter>
auto make_vec_2d_iterator(Iter xy_begin)
{
  return detail::make_counting_transform_iterator(0, detail::interleaved_to_vec_2d<Iter>{xy_begin});
}
 
template <typename FirstIter, typename SecondIter>
auto make_vec_2d_output_iterator(FirstIter first, SecondIter second)
{
  using T         = typename std::iterator_traits<FirstIter>::value_type;
  auto zipped_out = thrust::make_zip_iterator(thrust::make_tuple(first, second));
  return thrust::transform_output_iterator(zipped_out, detail::vec_2d_to_tuple<T>());
}
 
template <typename Iter>
auto make_vec_2d_output_iterator(Iter d_points_begin)
{
  using T                     = typename std::iterator_traits<Iter>::value_type;
  auto fixed_stride_2_functor = detail::strided_functor<2>();
  auto even_positions         = thrust::make_permutation_iterator(
    d_points_begin, detail::make_counting_transform_iterator(0, fixed_stride_2_functor));
  auto odd_positions = thrust::make_permutation_iterator(
    thrust::next(d_points_begin),
    detail::make_counting_transform_iterator(0, fixed_stride_2_functor));
  auto zipped_outputs =
    thrust::make_zip_iterator(thrust::make_tuple(even_positions, odd_positions));
  return thrust::make_transform_output_iterator(zipped_outputs, detail::vec_2d_to_tuple<T>());
}
 
template <typename FirstIter, typename SecondIter>
auto make_box_iterator(FirstIter first, SecondIter second)
{
  using Vertex = typename cuspatial::iterator_value_type<FirstIter>;
  using T      = typename Vertex::value_type;
 
  static_assert(is_same<Vertex, cuspatial::iterator_value_type<SecondIter>>(),
                "Iterator value_type mismatch");
 
  auto zipped = thrust::make_zip_iterator(first, second);
  return thrust::make_transform_iterator(zipped, detail::vec_2d_tuple_to_box<T>());
}
 
template <typename MinXIter, typename MinYIter, typename MaxXIter, typename MaxYIter>
auto make_box_output_iterator(MinXIter min_x, MinYIter min_y, MaxXIter max_x, MaxYIter max_y)
{
  using T         = typename std::iterator_traits<MinXIter>::value_type;
  auto zipped_out = thrust::make_zip_iterator(min_x, min_y, max_x, max_y);
  return thrust::transform_output_iterator(zipped_out, detail::box_to_tuple<T>());
}
 
template <typename IndexT, typename GeometryIter>
auto make_geometry_id_iterator(GeometryIter offsets_begin, GeometryIter offsets_end)
{
  return detail::make_counting_transform_iterator(
    IndexT{0}, detail::index_to_geometry_id<IndexT, GeometryIter>{offsets_begin, offsets_end});
}
 
template <typename IndexT, typename GeometryIter, typename PartIter>
auto make_geometry_id_iterator(GeometryIter geometry_offsets_begin,
                               GeometryIter geometry_offsets_end,
                               PartIter part_offsets_begin)
{
  auto first_part_offsets_begin =
    thrust::make_permutation_iterator(part_offsets_begin, geometry_offsets_begin);
 
  return make_geometry_id_iterator<IndexT>(
    first_part_offsets_begin,
    thrust::next(first_part_offsets_begin,
                 std::distance(geometry_offsets_begin, geometry_offsets_end)));
}
 
template <typename OffsetIterator>
auto make_count_iterator_from_offset_iterator(OffsetIterator it)
{
  auto zipped_offsets_it = thrust::make_zip_iterator(it, thrust::next(it));
  return thrust::make_transform_iterator(zipped_offsets_it, detail::offset_pair_to_count_functor{});
}
 
}  // namespace cuspatial