/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.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/batched_contraction/pipeline/batched_contraction_problem.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>
 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; 
 };
  
  
 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; 
  
     using KernelArgs =
         BatchedContractionKernelArgs<NumDimG, NumDimM, NumDimN, NumDimK, NumDTensor>; 
  
     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_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];
         }
  
         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);
         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;
         });
  
         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);
  
         const ADataType* a_ptr_final = a_ptr + splitk_batch_offset.as_k_split_offset[0];
         const BDataType* b_ptr_final = b_ptr + splitk_batch_offset.bs_k_split_offset[0];
         __shared__ char smem_ptr[GetSmemSize()];
  
         UniversalGemmKernel::RunGemm({a_ptr_final},
                                      {b_ptr_final},
                                      ds_batch_ptr,
                                      e_ptr,
                                      smem_ptr,
                                      gemm_kargs,
                                      splitk_batch_offset,
                                      i_m,
                                      i_n);
     }
 };
  
 } // namespace ck_tile