/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp Source File
Go to the documentation of this file.
 // Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 // SPDX-License-Identifier: MIT
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/batched_contraction/pipeline/batched_contraction_problem.hpp"
 #include "ck_tile/ops/batched_contraction/utils/tensor_descriptor_utils.hpp"
 #include "ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp"
  
 namespace ck_tile {
  
 template <ck_tile::index_t NumDTensor = 0>
 struct BatchedContractionHostArgs
 {
     CK_TILE_HOST
     BatchedContractionHostArgs(
         const void* a_ptr_,
         const void* b_ptr_,
         const std::array<const void*, NumDTensor>& ds_ptr_,
         void* e_ptr_,
         ck_tile::index_t k_batch_,
         const std::vector<ck_tile::index_t>& A_dims_, // [G0, G1, ..., M0, M1, ... , K0, K1, ...]
         const std::vector<ck_tile::index_t>& B_dims_, // [G0, G1, ..., N0, N1, ... , K0, K1, ...]
         const std::array<std::vector<ck_tile::index_t>, NumDTensor>&
             Ds_dims_, // [G0, G1, ..., M0, M1, ... , N0, N1, ...][NumDTensor]
         const std::vector<ck_tile::index_t>& E_dims_, // [G0, G1, ..., M0, M1, ... , N0, N1, ...]
  
         const std::vector<ck_tile::index_t>& A_strides_, // [G0, G1, ..., M0, M1, ...,K0, K1, ...]
         const std::vector<ck_tile::index_t>& B_strides_, // [G0, G1, ..., N0, N1, ...,K0, K1, ...]
         const std::array<std::vector<ck_tile::index_t>, NumDTensor>&
             Ds_strides_, // [G0, G1, ..., M0, M1, ...,N0, N1, ...]
         const std::vector<ck_tile::index_t>&
             E_strides_) // [G0, G1, ..., M0, M1, ...,N0, N1, ...][NumDTensor]
  
         : a_ptr(a_ptr_),
           b_ptr(b_ptr_),
           ds_ptr(ds_ptr_),
           e_ptr(e_ptr_),
           k_batch(k_batch_),
           A_dims(A_dims_),
           B_dims(B_dims_),
           Ds_dims(Ds_dims_),
           E_dims(E_dims_),
           A_strides(A_strides_),
           B_strides(B_strides_),
           Ds_strides(Ds_strides_),
           E_strides(E_strides_)
     {
     }
  
     const void* a_ptr;                          
     const void* b_ptr;                          
     std::array<const void*, NumDTensor> ds_ptr; 
     void* e_ptr;                                
     ck_tile::index_t k_batch;                   
     const std::vector<ck_tile::index_t>
         A_dims; 
     const std::vector<ck_tile::index_t>
         B_dims; 
     const std::array<std::vector<ck_tile::index_t>, NumDTensor>
         Ds_dims; 
     const std::vector<ck_tile::index_t>
         E_dims; 
     const std::vector<ck_tile::index_t>
         A_strides; 
     const std::vector<ck_tile::index_t>
         B_strides; 
     const std::array<std::vector<ck_tile::index_t>, NumDTensor>
         Ds_strides; 
     const std::vector<ck_tile::index_t>
         E_strides; 
 };
  
  
 template <ck_tile::index_t NumDimG,
           ck_tile::index_t NumDimM,
           ck_tile::index_t NumDimN,
           ck_tile::index_t NumDimK,
           ck_tile::index_t NumDTensor  = 0,
           ck_tile::index_t VectorSizeA = 1,
           ck_tile::index_t VectorSizeB = 1,
           ck_tile::index_t VectorSizeE = 1>
 struct BatchedContractionKernelArgs
 {
     const void* a_ptr;                          
     const void* b_ptr;                          
     std::array<const void*, NumDTensor> ds_ptr; 
     void* e_ptr;                                
     ck_tile::index_t k_batch;                   
  
     ck_tile::index_t M_dims[NumDimM]; 
     ck_tile::index_t N_dims[NumDimN]; 
     ck_tile::index_t K_dims[NumDimK]; 
     ck_tile::index_t
         G_dims[NumDimG]; 
  
     // Batch strides for efficient offset calculation
     ck_tile::index_t batch_stride_A;                          
     ck_tile::index_t batch_stride_B;                          
     ck_tile::index_t batch_stride_E;                          
     std::array<ck_tile::index_t, NumDTensor> batch_stride_Ds; 
  
     ck_tile::index_t G_total; 
     ck_tile::index_t M_total; 
     ck_tile::index_t N_total; 
     ck_tile::index_t K_total; 
  
     ck_tile::index_t
         stride_A; 
     ck_tile::index_t
         stride_B; 
     std::array<ck_tile::index_t, NumDTensor>
         stride_Ds; 
     ck_tile::index_t
         stride_E; 
  
     // Tensor descriptors (encode full multi-dimensional stride information with vectorization)
     using AGridDesc_M_K_ =
         decltype(TensorDescriptorUtils<NumDimG,
                                        NumDimM,
                                        NumDimN,
                                        NumDimK,
                                        VectorSizeA,
                                        VectorSizeB,
                                        VectorSizeE>::Make_A_GridDescriptor_M_K({}, {}));
     using BGridDesc_N_K_ =
         decltype(TensorDescriptorUtils<NumDimG,
                                        NumDimM,
                                        NumDimN,
                                        NumDimK,
                                        VectorSizeA,
                                        VectorSizeB,
                                        VectorSizeE>::Make_B_GridDescriptor_N_K({}, {}));
     using EGridDesc_M_N_ =
         decltype(TensorDescriptorUtils<NumDimG,
                                        NumDimM,
                                        NumDimN,
                                        NumDimK,
                                        VectorSizeA,
                                        VectorSizeB,
                                        VectorSizeE>::Make_E_GridDescriptor_M_N({}, {}));
  
     AGridDesc_M_K_ a_grid_desc_m_k; 
     BGridDesc_N_K_ b_grid_desc_n_k; 
     EGridDesc_M_N_ e_grid_desc_m_n; 
     std::array<EGridDesc_M_N_, NumDTensor>
         ds_grid_desc_m_n; 
 };
  
  
 template <typename Problem_,
           typename TilePartitioner_,
           typename GemmPipeline_,
           typename EpiloguePipeline_>
 struct BatchedContractionKernel
 {
     // Type aliases for cleaner code and better readability
     using Problem = ck_tile::remove_cvref_t<Problem_>; 
     using ADataType =
         ck_tile::remove_cvref_t<typename Problem::ADataType>; 
     using BDataType =
         ck_tile::remove_cvref_t<typename Problem::BDataType>; 
     using DsDataType =
         ck_tile::remove_cvref_t<typename Problem::DsDataType>; 
     using EDataType =
         ck_tile::remove_cvref_t<typename Problem::EDataType>; 
  
     // Compile-time dimension constants extracted from problem specification
     static constexpr ck_tile::index_t NumDimG = Problem::NumDimG; 
     static constexpr ck_tile::index_t NumDimM =
         Problem::NumDimM; 
     static constexpr ck_tile::index_t NumDimN =
         Problem::NumDimN; 
     static constexpr ck_tile::index_t NumDimK =
         Problem::NumDimK; 
     static constexpr ck_tile::index_t NumDTensor =
         Problem::NumDTensor; 
  
     // Pipeline and partitioning strategy types
     using TilePartitioner =
         ck_tile::remove_cvref_t<TilePartitioner_>; 
     using GemmPipeline = ck_tile::remove_cvref_t<GemmPipeline_>; 
     using EpiloguePipeline =
         ck_tile::remove_cvref_t<EpiloguePipeline_>; 
  
     // Underlying GEMM kernel that performs the actual computation
     using UniversalGemmKernel =
         ck_tile::UniversalGemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>;
  
     static constexpr ck_tile::index_t kBlockSize =
         UniversalGemmKernel::kBlockSize; 
  
     // Tensor descriptor utilities with vectorization support
     using DescriptorUtils = TensorDescriptorUtils<NumDimG,
                                                   NumDimM,
                                                   NumDimN,
                                                   NumDimK,
                                                   GemmPipeline::GetVectorSizeA(),
                                                   GemmPipeline::GetVectorSizeB(),
                                                   EpiloguePipeline::GetVectorSizeC()>;
  
     // Kernel arguments with vectorization support
     using KernelArgs = BatchedContractionKernelArgs<NumDimG,
                                                     NumDimM,
                                                     NumDimN,
                                                     NumDimK,
                                                     NumDTensor,
                                                     GemmPipeline::GetVectorSizeA(),
                                                     GemmPipeline::GetVectorSizeB(),
                                                     EpiloguePipeline::GetVectorSizeC()>;
  
     CK_TILE_HOST static constexpr auto GetKernelName() { return "batched_contraction_kernel"; }
  
     CK_TILE_HOST static constexpr bool IsSupportedArguments(const KernelArgs& kargs)
     {
         typename UniversalGemmKernel::KernelArgs gemm_kargs{{kargs.a_ptr},
                                                             {kargs.b_ptr},
                                                             kargs.ds_ptr,
                                                             kargs.e_ptr,
                                                             kargs.M_total,
                                                             kargs.N_total,
                                                             kargs.K_total,
                                                             {kargs.stride_A},
                                                             {kargs.stride_B},
                                                             kargs.stride_Ds,
                                                             kargs.stride_E,
                                                             kargs.k_batch};
  
         return UniversalGemmKernel::IsSupportedArgument(gemm_kargs) && kargs.G_total > 0;
     }
  
     CK_TILE_HOST static constexpr ck_tile::index_t GetSmemSize()
     {
         return UniversalGemmKernel::GetSmemSize();
     }
  
     CK_TILE_HOST static constexpr auto GetBlockSize()
     {
         return dim3(UniversalGemmKernel::kBlockSize);
     }
  
     CK_TILE_HOST static constexpr auto GridSize(const KernelArgs& kargs)
     {
         return dim3(
             TilePartitioner::GridSize(kargs.M_total, kargs.N_total), kargs.G_total, kargs.k_batch);
     }
  
     CK_TILE_DEVICE static void RunGemm(const ADataType* a_ptr,
                                        const BDataType* b_ptr,
                                        const std::array<const void*, NumDTensor>& ds_ptr,
                                        EDataType* e_ptr,
                                        void* smem_ptr,
                                        const KernelArgs& kargs,
                                        const index_t k_size,
                                        const index_t i_m,
                                        const index_t i_n)
     {
         // Create tensor views from descriptors (supports arbitrary stride patterns)
         auto a_tensor_view =
             make_tensor_view<address_space_enum::global>(a_ptr, kargs.a_grid_desc_m_k);
         auto b_tensor_view =
             make_tensor_view<address_space_enum::global>(b_ptr, kargs.b_grid_desc_n_k);
         auto e_tensor_view =
             make_tensor_view<address_space_enum::global>(e_ptr, kargs.e_grid_desc_m_n);
  
         // Pad views for boundary handling and optimization (like UniversalGemmKernel)
         auto a_pad_view = pad_tensor_view(
             a_tensor_view,
             make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::KPerBlock>{}),
             sequence<false, GemmPipeline::kPadK>{});
  
         auto b_pad_view = pad_tensor_view(
             b_tensor_view,
             make_tuple(number<TilePartitioner::NPerBlock>{}, number<TilePartitioner::KPerBlock>{}),
             sequence<false, GemmPipeline::kPadK>{});
  
         auto e_pad_view = pad_tensor_view(
             e_tensor_view,
             make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::NPerBlock>{}),
             sequence<false, GemmPipeline::kPadN>{});
  
         // Create tile windows from PADDED views
         auto a_block_window = make_tile_window(
             a_pad_view,
             make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::KPerBlock>{}),
             {i_m, 0});
  
         auto b_block_window = make_tile_window(
             b_pad_view,
             make_tuple(number<TilePartitioner::NPerBlock>{}, number<TilePartitioner::KPerBlock>{}),
             {i_n, 0});
  
         auto e_block_window = make_tile_window(
             e_pad_view,
             make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::NPerBlock>{}),
             {i_m, i_n});
  
         // Calculate number of K loops
         const index_t num_loop =
             __builtin_amdgcn_readfirstlane(TilePartitioner::GetLoopNum(k_size));
  
         // Run GEMM Pipeline (same as UniversalGemmKernel, but with descriptor-based windows)
         using AElementWise = remove_cvref_t<typename GemmPipeline::AElementWise>;
         using BElementWise = remove_cvref_t<typename GemmPipeline::BElementWise>;
  
         const auto& c_block_tile = GemmPipeline{}(
             a_block_window, AElementWise{}, b_block_window, BElementWise{}, num_loop, smem_ptr);
  
         // Create D windows from descriptors (for each D tensor)
         auto ds_block_windows = generate_tuple(
             [&](auto i) {
                 using DDataType        = remove_cvref_t<std::tuple_element_t<i.value, DsDataType>>;
                 const DDataType* d_ptr = static_cast<const DDataType*>(ds_ptr[i]);
  
                 auto d_tensor_view =
                     make_tensor_view<address_space_enum::global>(d_ptr, kargs.ds_grid_desc_m_n[i]);
  
                 return make_tile_window(d_tensor_view,
                                         make_tuple(number<TilePartitioner::MPerBlock>{},
                                                    number<TilePartitioner::NPerBlock>{}),
                                         {i_m, i_n});
             },
             number<NumDTensor>{});
  
         // Run Epilogue Pipeline with descriptor-based D windows
         EpiloguePipeline{}(e_block_window, c_block_tile, ds_block_windows, smem_ptr);
     }
  
     CK_TILE_HOST static constexpr KernelArgs
     MakeKernelArgs(const BatchedContractionHostArgs<NumDTensor>& host_args)
     {
         const auto expected_A_dims = NumDimG + NumDimM + NumDimK;
         const auto expected_B_dims = NumDimG + NumDimN + NumDimK;
         const auto expected_E_dims = NumDimG + NumDimM + NumDimN;
  
         if(host_args.A_dims.size() != expected_A_dims ||
            host_args.A_strides.size() != expected_A_dims)
         {
             throw std::invalid_argument("A dimension size mismatch");
         }
         if(host_args.B_dims.size() != expected_B_dims ||
            host_args.B_strides.size() != expected_B_dims)
         {
             throw std::invalid_argument("B dimension size mismatch");
         }
         if(host_args.E_dims.size() != expected_E_dims ||
            host_args.E_strides.size() != expected_E_dims)
         {
             throw std::invalid_argument("E dimension size mismatch");
         }
  
         for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
         {
             if(host_args.Ds_dims[d].size() != expected_E_dims ||
                host_args.Ds_strides[d].size() != expected_E_dims)
             {
                 throw std::invalid_argument("D dimension size mismatch");
             }
         }
  
         KernelArgs kargs;
         kargs.a_ptr   = host_args.a_ptr;
         kargs.b_ptr   = host_args.b_ptr;
         kargs.ds_ptr  = host_args.ds_ptr;
         kargs.e_ptr   = host_args.e_ptr;
         kargs.k_batch = host_args.k_batch;
  
         // Validate and set G dimensions (must be identical across all tensors)
         for(ck_tile::index_t i = 0; i < NumDimG; ++i)
         {
             // All tensors must have same G dimensions for valid contraction
             if(host_args.A_dims[i] != host_args.B_dims[i] ||
                host_args.A_dims[i] != host_args.E_dims[i])
             {
                 throw std::invalid_argument(
                     "All tensors must have identical G dimensions for valid contraction");
             }
  
             // Store G dimensions (same for all tensors)
             kargs.G_dims[i] = host_args.A_dims[i];
         }
  
         // Set batch strides from the stride of last G dimension
         kargs.batch_stride_A = host_args.A_strides[NumDimG - 1];
         kargs.batch_stride_B = host_args.B_strides[NumDimG - 1];
         kargs.batch_stride_E = host_args.E_strides[NumDimG - 1];
  
         for(ck_tile::index_t i = 0; i < NumDimM; ++i)
         {
             kargs.M_dims[i] = host_args.A_dims[NumDimG + i];
             if(kargs.M_dims[i] != host_args.E_dims[NumDimG + i])
             {
                 throw std::invalid_argument("M dimension mismatch between A and E tensors");
             }
         }
         for(ck_tile::index_t i = 0; i < NumDimN; ++i)
         {
             kargs.N_dims[i] = host_args.B_dims[NumDimG + i];
             if(kargs.N_dims[i] != host_args.E_dims[NumDimG + NumDimM + i])
             {
                 throw std::invalid_argument("N dimension mismatch between B and E tensors");
             }
         }
         for(ck_tile::index_t i = 0; i < NumDimK; ++i)
         {
             kargs.K_dims[i] = host_args.A_dims[NumDimG + NumDimM + i];
             if(kargs.K_dims[i] != host_args.B_dims[NumDimG + NumDimN + i])
             {
                 throw std::invalid_argument("K dimension mismatch between A and B tensors");
             }
         }
  
         // Calculate total dimensions from individual dimension arrays
         kargs.G_total = 1;
         for(ck_tile::index_t i = 0; i < NumDimG; ++i)
         {
             kargs.G_total *= kargs.G_dims[i];
         }
  
         kargs.M_total = 1;
         for(ck_tile::index_t i = 0; i < NumDimM; ++i)
         {
             kargs.M_total *= kargs.M_dims[i];
         }
  
         kargs.N_total = 1;
         for(ck_tile::index_t i = 0; i < NumDimN; ++i)
         {
             kargs.N_total *= kargs.N_dims[i];
         }
  
         kargs.K_total = 1;
         for(ck_tile::index_t i = 0; i < NumDimK; ++i)
         {
             kargs.K_total *= kargs.K_dims[i];
         }
  
         // Create tensor descriptors on host using actual dims and strides
         kargs.a_grid_desc_m_k =
             DescriptorUtils::Make_A_GridDescriptor_M_K(host_args.A_dims, host_args.A_strides);
         kargs.b_grid_desc_n_k =
             DescriptorUtils::Make_B_GridDescriptor_N_K(host_args.B_dims, host_args.B_strides);
         kargs.e_grid_desc_m_n =
             DescriptorUtils::Make_E_GridDescriptor_M_N(host_args.E_dims, host_args.E_strides);
  
         // Create D descriptors with their own strides (same shape as E, independent strides)
         for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
         {
             kargs.ds_grid_desc_m_n[d] = DescriptorUtils::Make_E_GridDescriptor_M_N(
                 host_args.Ds_dims[d], host_args.Ds_strides[d]);
         }
  
         // Keep simple strides for backward compatibility
         kargs.stride_A = kargs.K_total;
         kargs.stride_B = kargs.K_total;
         kargs.stride_E = kargs.N_total;
  
         // Validate D tensors have same G dimensions and set their batch strides
         for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
         {
             for(ck_tile::index_t i = 0; i < NumDimG; ++i)
             {
                 if(host_args.Ds_dims[d][i] != host_args.A_dims[i])
                 {
                     throw std::invalid_argument(
                         "D tensor G dimensions must match A/B/E tensor G dimensions");
                 }
             }
             // Set batch stride for D tensor
             kargs.batch_stride_Ds[d] = host_args.Ds_strides[d][NumDimG - 1];
             kargs.stride_Ds[d]       = kargs.N_total; // D tensors same shape as E
         }
  
         return kargs;
     }
  
     CK_TILE_DEVICE void operator()(const KernelArgs& kargs) const
     {
  
         const auto [iM, iN] =
             TilePartitioner{kargs.M_total, kargs.N_total}.GetOutputTileIndex(blockIdx.x);
         const ck_tile::index_t i_m =
             __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock);
         const ck_tile::index_t i_n =
             __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock);
  
         const auto i_batch_flat              = __builtin_amdgcn_readfirstlane(blockIdx.y);
         [[maybe_unused]] const auto i_splitk = __builtin_amdgcn_readfirstlane(blockIdx.z);
  
         // Calculate batch offsets for each tensor
         const auto batch_offset_A = i_batch_flat * kargs.batch_stride_A;
         const auto batch_offset_B = i_batch_flat * kargs.batch_stride_B;
         const auto batch_offset_E = i_batch_flat * kargs.batch_stride_E;
  
         const ADataType* a_ptr = static_cast<const ADataType*>(kargs.a_ptr) + batch_offset_A;
         const BDataType* b_ptr = static_cast<const BDataType*>(kargs.b_ptr) + batch_offset_B;
         EDataType* e_ptr       = static_cast<EDataType*>(kargs.e_ptr) + batch_offset_E;
  
         std::array<const void*, NumDTensor> ds_batch_ptr;
         static_for<0, NumDTensor, 1>{}([&](auto i) {
             using DDataType           = typename std::tuple_element<i.value, DsDataType>::type;
             const auto batch_offset_D = i_batch_flat * kargs.batch_stride_Ds[i];
             ds_batch_ptr[i] = static_cast<const DDataType*>(kargs.ds_ptr[i]) + batch_offset_D;
         });
  
         // Allocate shared memory
         __shared__ char smem_ptr[GetSmemSize()];
  
         // Use UniversalGemmKernel's SplitKBatchOffset for split-K calculation
         typename UniversalGemmKernel::KernelArgs gemm_kargs{{a_ptr},
                                                             {b_ptr},
                                                             ds_batch_ptr,
                                                             e_ptr,
                                                             kargs.M_total,
                                                             kargs.N_total,
                                                             kargs.K_total,
                                                             {kargs.stride_A},
                                                             {kargs.stride_B},
                                                             kargs.stride_Ds,
                                                             kargs.stride_E,
                                                             kargs.k_batch};
  
         const typename UniversalGemmKernel::SplitKBatchOffset splitk_batch_offset(gemm_kargs,
                                                                                   i_splitk);
  
         // Apply K-split offsets and run descriptor-based RunGemm
         const ADataType* a_ptr_split = a_ptr + splitk_batch_offset.as_k_split_offset[0];
         const BDataType* b_ptr_split = b_ptr + splitk_batch_offset.bs_k_split_offset[0];
  
         RunGemm(a_ptr_split,
                 b_ptr_split,
                 ds_batch_ptr,
                 e_ptr,
                 smem_ptr,
                 kargs,
                 splitk_batch_offset.splitted_k,
                 i_m,
                 i_n);
     }
 };
  
 } // namespace ck_tile