CompositeRandomAccessor.h 929 B

1234567891011121314151617181920212223242526272829303132333435
  1. #pragma once
  2. #include <ATen/native/CompositeRandomAccessorCommon.h>
  3. #include <thrust/tuple.h>
  4. namespace at { namespace native {
  5. struct TupleInfoCPU {
  6. template <typename ...Types>
  7. using tuple = thrust::tuple<Types...>;
  8. template <typename ...Types>
  9. static constexpr auto tie(Types&... args) noexcept {
  10. return thrust::tie(args...);
  11. }
  12. };
  13. template <typename KeyAccessor, typename ValueAccessor>
  14. using CompositeRandomAccessorCPU =
  15. CompositeRandomAccessor<KeyAccessor, ValueAccessor, TupleInfoCPU>;
  16. template <typename Values, typename References>
  17. void swap(
  18. references_holder<Values, References> rh1,
  19. references_holder<Values, References> rh2
  20. ) {
  21. return thrust::swap(rh1.data(), rh2.data());
  22. }
  23. template <int N, typename Values, typename References>
  24. auto get(references_holder<Values, References> rh) -> decltype(thrust::get<N>(rh.data())) {
  25. return thrust::get<N>(rh.data());
  26. }
  27. }} // namespace at::native