| 1 | #pragma once |
| 2 | |
| 3 | #include "machete_mm_kernel.cuh" |
| 4 | #include "cutlass_extensions/cute_utils.cuh" |
| 5 | #include "libtorch_stable/cutlass_extensions/torch_utils.hpp" |
| 6 | #include <torch/headeronly/util/Exception.h> |
| 7 | |
| 8 | namespace machete { |
| 9 | |
| 10 | template <int threads, typename PrepackedLayoutB, typename BInTensor, |
| 11 | typename ElementB> |
| 12 | static __global__ void prepack_B_kernel(BInTensor B_in, ElementB* B_out_ptr) { |
| 13 | auto constexpr block_size = |
| 14 | Int<size(typename PrepackedLayoutB::PPBlockShape_NK{})>{}; |
| 15 | auto constexpr eles_per_thread = Int<block_size / threads>{}; |
| 16 | static_assert(block_size % threads == 0, |
| 17 | "block_size must be divisible by the number of threads"); |
| 18 | |
| 19 | // Which pre-packed are we responsible for |
| 20 | auto blk_coord = make_coord(blockIdx.x, blockIdx.y, blockIdx.z); |
| 21 | auto tB_in = local_tile( |
| 22 | B_in, append(typename PrepackedLayoutB::PPBlockShape_NK{}, _1{}), |
| 23 | blk_coord); |
| 24 | |
| 25 | // Find the start offset in the output for this pre-packed block |
| 26 | auto bNbKL_to_offset = PrepackedLayoutB::bNbKL_to_offset(shape(B_in)); |
| 27 | |
| 28 | // Tensor representing a 1:1 mapping to the output space in 1D |
| 29 | auto tB_out_linear = |
| 30 | make_tensor(get_logical_ptr(B_out_ptr) + bNbKL_to_offset(blk_coord), |
| 31 | make_layout(make_shape(block_size))); |
| 32 | // Mapping from output space (1D) to input space |
| 33 | auto tB_in_linear = make_tensor( |
| 34 | tB_in.data(), |
| 35 | tB_in.layout() |
| 36 | .compose(right_inverse(PrepackedLayoutB::ppblock_ilvd_NK_to_offset())) |
| 37 | .with_shape(make_shape(block_size))); |
| 38 | |
| 39 | // Tile for this specific thread (could have used a TiledCopy but these work |
| 40 | // best with 2d layouts, this is a simple 1d layout so local_tile is enough, |
| 41 | // we are also not that concerned with performance for this kernel) |
| 42 | auto thr_tB_in_linear = |
| 43 | local_tile(tB_in_linear, make_shape(eles_per_thread), threadIdx.x); |
| 44 | auto thr_tB_out_linear = |
| 45 | local_tile(tB_out_linear, make_shape(eles_per_thread), threadIdx.x); |
| 46 | |
| 47 | // Construct a register-backed Tensor with the same shape as each thread's |
| 48 | // partition |
| 49 | auto fragment = make_tensor<ElementB>(shape(thr_tB_in_linear)); |
| 50 | |
| 51 | copy(thr_tB_in_linear, fragment); |
| 52 | copy(Copy_Atom<DefaultCopy, uint8_t>{}, fragment, thr_tB_out_linear); |
| 53 | } |
| 54 | |
| 55 | template <typename PrepackedLayoutB, typename InLayout> |
| 56 | static void prepack_B_template( |
| 57 | cudaStream_t stream, typename PrepackedLayoutB::ElementB const* B_in_ptr, |
| 58 | InLayout B_layout, typename PrepackedLayoutB::ElementB* B_out_ptr) { |
| 59 | using TileShapeNKL = |
| 60 | decltype(append(typename PrepackedLayoutB::PPBlockShape_NK{}, _1{})); |
| 61 | auto ilvd_NKbNbKL_to_offset = |
| 62 | PrepackedLayoutB::ilvd_NKbNbKL_to_offset(shape(B_layout)); |
| 63 | |
| 64 | STD_TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0); |
| 65 | STD_TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0); |
| 66 | |
| 67 | auto N_tiles = size<0>(B_layout) / size<0>(TileShapeNKL{}); |
| 68 | auto K_tiles = size<1>(B_layout) / size<1>(TileShapeNKL{}); |
| 69 | auto L_tiles = size<2>(B_layout); |
| 70 | |
| 71 | auto B_in = make_tensor(get_logical_ptr(B_in_ptr), B_layout); |
| 72 | |
| 73 | prepack_B_kernel<128, PrepackedLayoutB> |
| 74 | <<<dim3(N_tiles, K_tiles, L_tiles), 128, 0, stream>>>(B_in, B_out_ptr); |
| 75 | } |
| 76 | |
| 77 | }; // namespace machete |