| 1 | #pragma once |
| 2 | |
| 3 | #include <cute/tensor.hpp> |
| 4 | namespace cute { |
| 5 | |
| 6 | //////////////////////////////////////////////////////////////////// |
| 7 | // layout utils |
| 8 | //////////////////////////////////////////////////////////////////// |
| 9 | |
| 10 | // Permute layout based on indices, example: |
| 11 | // permute_layout<1, 0>(layout) will swap the two dimensions |
| 12 | // permute_layout<0, 2, 1>(layout) will swap the last two dimensions |
| 13 | template <size_t... I, typename Layout> |
| 14 | CUTE_HOST_DEVICE static constexpr auto permute_layout(Layout l) { |
| 15 | static_assert(rank(l) == sizeof...(I), "Invalid permutation, rank mismatch"); |
| 16 | return cute::make_layout(cute::get<I>(l)...); |
| 17 | } |
| 18 | |
| 19 | // is the layout f(x) = x |
| 20 | template <typename Layout> |
| 21 | CUTE_HOST_DEVICE static constexpr bool is_identity_layout() { |
| 22 | if constexpr (std::is_same_v<Layout, void>) { |
| 23 | return true; |
| 24 | } else { |
| 25 | constexpr auto coalesced_layout = coalesce(Layout{}); |
| 26 | if constexpr (rank(coalesced_layout) == 1 && |
| 27 | stride<0>(coalesced_layout) == 1) { |
| 28 | return true; |
| 29 | } |
| 30 | return false; |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | //////////////////////////////////////////////////////////////////// |
| 35 | // Pointer utils |
| 36 | //////////////////////////////////////////////////////////////////// |
| 37 | |
| 38 | template <class PointerType> |
| 39 | static constexpr auto get_logical_ptr(PointerType* ptr) { |
| 40 | if constexpr (cute::sizeof_bits_v<PointerType> < 8) { |
| 41 | return cute::subbyte_iterator<PointerType>(ptr); |
| 42 | } else { |
| 43 | return ptr; |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | //////////////////////////////////////////////////////////////////// |
| 48 | // Misc utils |
| 49 | //////////////////////////////////////////////////////////////////// |
| 50 | |
| 51 | template <typename T, typename Elements> |
| 52 | CUTE_HOST_DEVICE static constexpr auto create_auto_vectorizing_copy() { |
| 53 | constexpr auto bits = sizeof_bits_v<T> * Elements{}; |
| 54 | if constexpr (bits % 128 == 0) { |
| 55 | return AutoVectorizingCopyWithAssumedAlignment<128>{}; |
| 56 | } else if constexpr (bits % 64 == 0) { |
| 57 | return AutoVectorizingCopyWithAssumedAlignment<64>{}; |
| 58 | } else if constexpr (bits % 32 == 0) { |
| 59 | return AutoVectorizingCopyWithAssumedAlignment<32>{}; |
| 60 | } else if constexpr (bits % 16 == 0) { |
| 61 | return AutoVectorizingCopyWithAssumedAlignment<16>{}; |
| 62 | } else { |
| 63 | return AutoVectorizingCopyWithAssumedAlignment<8>{}; |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | }; // namespace cute |