| 1 | #pragma once |
| 2 | |
| 3 | #include "cutlass/cutlass.h" |
| 4 | #include "cute/layout.hpp" |
| 5 | |
| 6 | namespace machete { |
| 7 | |
| 8 | using namespace cute; |
| 9 | |
| 10 | // get an interleaved block layout where each element consecutive element has a |
| 11 | // stride of bit_stride and the block width is blk_bit_width, |
| 12 | // examples: |
| 13 | // size_bits<T> = 8, bit_stride = 8, blk_bit_width = 32 -> 4:1 |
| 14 | // size_bits<T> = 8, bit_stride = 16, blk_bit_width = 32 -> (2, 2):(2, 1) |
| 15 | // size_bits<T> = 4, bit_stride = 8, blk_bit_width = 32 -> (4, 2):(2, 1) |
| 16 | // size_bits<T> = 4, bit_stride = 16, blk_bit_width = 32 -> (2, 4):(4, 1) |
| 17 | template <typename T, int bit_stride, int blk_bit_width> |
| 18 | CUTE_HOST_DEVICE static constexpr auto get_interleaved_blk_layout() { |
| 19 | static_assert(blk_bit_width % bit_stride == 0); |
| 20 | static_assert(bit_stride % cute::sizeof_bits_v<T> == 0); |
| 21 | |
| 22 | constexpr auto elems_per_blk = blk_bit_width / cute::sizeof_bits_v<T>; |
| 23 | |
| 24 | if constexpr (cute::sizeof_bits_v<T> == bit_stride) { |
| 25 | // identity layout |
| 26 | return Layout<Shape<Int<elems_per_blk>>>{}; |
| 27 | } else { |
| 28 | constexpr auto elems_per_stride = bit_stride / cute::sizeof_bits_v<T>; |
| 29 | constexpr auto num_strides = elems_per_blk / elems_per_stride; |
| 30 | return Layout<Shape<Int<num_strides>, Int<elems_per_stride>>, |
| 31 | Stride<Int<elems_per_stride>, Int<1>>>{}; |
| 32 | } |
| 33 | } |
| 34 | |
| 35 | }; // namespace machete |