batched_contraction_kernel.hpp Source File

batched_contraction_kernel.hpp Source File#

Composable Kernel: batched_contraction_kernel.hpp Source 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