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