batch_gemm.hpp

Go to the documentation of this file.

/*******************************************************************************
* Copyright (c) 2022-2023 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
 
 
#pragma once
 
#include "xetla.hpp"
 
namespace gpu::xetla::kernel {
 
template <typename gemm_t_, typename epilogue_t_, gpu_arch arch_tag_>
class batch_gemm_t {
    using gemm_t = gemm_t_;
    using epilogue_t = epilogue_t_;
    using gemm_args_t = typename gemm_t::arguments_t;
    using epilogue_args_t = typename epilogue_t::arguments_t;
 
    using tile_shape = typename gemm_t::tile_shape;
    static constexpr uint32_t wg_tile_m = tile_shape::wg_tile_size_y;
    static constexpr uint32_t wg_tile_n = tile_shape::wg_tile_size_x;
    static constexpr uint32_t sg_tile_m = tile_shape::sg_tile_size_y;
    static constexpr uint32_t sg_tile_n = tile_shape::sg_tile_size_x;
    static constexpr uint32_t wg_size_y = tile_shape::wg_size_y;
    static constexpr uint32_t wg_size_x = tile_shape::wg_size_x;
    static constexpr uint32_t real_wg_tile_m = sg_tile_m * wg_size_y;
    static constexpr uint32_t real_wg_tile_n = sg_tile_n * wg_size_x;
 
    static constexpr uint32_t k_stride = gemm_t::k_stride;
    using work_group_t = typename gemm_t::work_group_t;
 
    static constexpr gpu_arch arch_tag = arch_tag_;
    static_assert(arch_tag == gemm_t::arch_tag, "arch_tag should be the same");
    static_assert(
            arch_tag == epilogue_t::arch_tag, "arch_tag should be the same");
    static_assert(std::is_same<typename gemm_t::tile_shape,
                          typename epilogue_t::tile_shape>::value,
            "tile_shape should be the same");
 
    using mem_desc_a_t = typename gemm_t::mem_desc_a_t;
    using mem_desc_b_t = typename gemm_t::mem_desc_b_t;
    using mem_desc_c_t = typename epilogue_t::mem_desc_c_t;
    using matA_base_t = typename mem_desc_a_t::base_t;
    using matB_base_t = typename mem_desc_b_t::base_t;
    using matC_base_t = typename mem_desc_c_t::base_t;
    using dtype_a = typename mem_desc_a_t::dtype;
    using dtype_b = typename mem_desc_b_t::dtype;
    using dtype_c = typename mem_desc_c_t::dtype;
    using matAcc_t = typename gemm_t::matAcc_t;
 
    static_assert((!gemm_t::is_col_major_a) && (!gemm_t::is_col_major_b),
            "only support row-major");
 
public:
    struct arguments_t {
        uint32_t batch_size;
        uint32_t matrix_m;
        uint32_t matrix_k;
        uint32_t matrix_n;
        uint32_t matA_ld;
        uint32_t matB_ld;
        uint32_t matC_ld;
        matA_base_t matA_base;
        matB_base_t matB_base;
        matC_base_t matC_base;
        epilogue_args_t epilogue_args;
 
        inline arguments_t() = default;
 
        static constexpr bool host_callable = true;
 
        inline arguments_t(uint32_t batch_size_, uint32_t matrix_m_,
                uint32_t matrix_k_, uint32_t matrix_n_, matA_base_t matA_base_,
                uint32_t matA_ld_, matB_base_t matB_base_, uint32_t matB_ld_,
                matC_base_t matC_base_, uint32_t matC_ld_,
                epilogue_args_t epilogue_args_ = {})
            : batch_size(batch_size_)
            , matrix_m(matrix_m_)
            , matrix_k(matrix_k_)
            , matrix_n(matrix_n_)
            , matA_ld(matA_ld_)
            , matB_ld(matB_ld_)
            , matC_ld(matC_ld_)
            , matA_base(matA_base_)
            , matB_base(matB_base_)
            , matC_base(matC_base_)
            , epilogue_args(epilogue_args_) {}
        // Be aware of the risks: Rule of three (copy constructor, copy assignment, destructor)
        // Please check if you need to add self-define destructor
        // inline ~arguments_t(){}
        inline arguments_t(const arguments_t &args)
            : batch_size(args.batch_size)
            , matrix_m(args.matrix_m)
            , matrix_k(args.matrix_k)
            , matrix_n(args.matrix_n)
            , matA_ld(args.matA_ld)
            , matB_ld(args.matB_ld)
            , matC_ld(args.matC_ld)
            , matA_base(args.matA_base)
            , matB_base(args.matB_base)
            , matC_base(args.matC_base)
            , epilogue_args(args.epilogue_args) {}
        inline arguments_t &operator=(const arguments_t &args) {
            this->batch_size = args.batch_size;
            this->matrix_m = args.matrix_m;
            this->matrix_k = args.matrix_k;
            this->matrix_n = args.matrix_n;
            this->matA_base = args.matA_base;
            this->matA_ld = args.matA_ld;
            this->matB_base = args.matB_base;
            this->matB_ld = args.matB_ld;
            this->matC_base = args.matC_base;
            this->matC_ld = args.matC_ld;
            this->epilogue_args = args.epilogue_args;
            return *this;
        }
    };
 
    __XETLA_API static constexpr uint32_t get_barrier_count() {
        constexpr uint32_t count
                = gemm_t::barrier_count + epilogue_t::barrier_count;
        static_assert(
                count <= 32, "The named_barrier count should be less than 32!");
        return count;
    }
 
    __XETLA_API static constexpr uint32_t get_slm_size() {
        constexpr uint32_t size = gemm_t::slm_size + epilogue_t::slm_size;
        static_assert(size <= (128 * 1024),
                "The local memory size should be less than 128KB!");
        return size;
    };
 
    static cl::sycl::range<3> get_local_range() {
        uint32_t local_range_m = (wg_tile_m + sg_tile_m - 1) / sg_tile_m;
        uint32_t local_range_n = (wg_tile_n + sg_tile_n - 1) / sg_tile_n;
        std::cout << "Local range: {" << 1 << ", " << local_range_m << ", "
                  << local_range_n << "} \n";
        assert(local_range_m * local_range_n <= 32);
        return cl::sycl::range<3> {1, local_range_m, local_range_n};
    };
 
    static cl::sycl::range<3> get_group_range(
            uint32_t batch_size, uint32_t matrix_m, uint32_t matrix_n) {
        uint32_t group_range_m = (matrix_m + wg_tile_m - 1) / wg_tile_m;
        uint32_t group_range_n = (matrix_n + wg_tile_n - 1) / wg_tile_n;
        std::cout << "Group range: {" << batch_size << ", " << group_range_m
                  << ", " << group_range_n << "} \n";
        return cl::sycl::range<3> {batch_size, group_range_m, group_range_n};
    };
 
    static cl::sycl::nd_range<3> get_nd_range(arguments_t &args) {
        cl::sycl::range<3> local_range = get_local_range();
        cl::sycl::range<3> group_range = get_group_range(
                args.batch_size, args.matrix_m, args.matrix_n);
        return cl::sycl::nd_range<3> {group_range * local_range, local_range};
    };
 
    static bool can_implement(arguments_t &args) {
        bool implementable = true;
        if (gemm_t::msg_type_a != msg_type::unaligned_2d) {
            if (gemm_t::msg_type_a == msg_type::block_2d) {
                implementable &= kernel::block_2d<gpu_arch::Xe,
                        dtype_a>::check_tensor((uint64_t)(args.matA_base.base),
                        args.matrix_k, args.matrix_m * args.batch_size,
                        args.matA_ld);
            } else {
                implementable &= kernel::general_1d<gpu_arch::Xe,
                        dtype_a>::check_alignment(args.matA_base.base,
                        args.matA_ld);
            }
        }
        if (gemm_t::msg_type_b != msg_type::unaligned_2d) {
            if (gemm_t::msg_type_b == msg_type::block_2d) {
                implementable &= kernel::block_2d<gpu_arch::Xe,
                        dtype_b>::check_tensor((uint64_t)(args.matB_base.base),
                        args.matrix_n, args.matrix_k * args.batch_size,
                        args.matB_ld);
            } else {
                implementable &= kernel::general_1d<gpu_arch::Xe,
                        dtype_b>::check_alignment(args.matB_base.base,
                        args.matB_ld);
            }
        }
        if (epilogue_t::msg_type_c != msg_type::unaligned_2d) {
            if (epilogue_t::msg_type_c == msg_type::block_2d) {
                implementable &= kernel::block_2d<gpu_arch::Xe,
                        dtype_c>::check_tensor((uint64_t)(args.matC_base.base),
                        args.matrix_n, args.matrix_m * args.batch_size,
                        args.matC_ld);
            } else {
                implementable &= kernel::general_1d<gpu_arch::Xe,
                        dtype_c>::check_alignment(args.matC_base.base,
                        args.matC_ld);
            }
        }
 
        return implementable;
    }
 
    __XETLA_API KERNEL_FUNC void operator()(sycl::nd_item<3> &item,
            const arguments_t &args, uint32_t slm_base = 0,
            uint32_t nbarrier_base = 0) {
        // set up workgroup level coordinates and boundaries
        int batch_id = item.get_group(0);
        int start_n = item.get_group(2) * wg_tile_n;
        int start_m = item.get_group(1) * wg_tile_m + batch_id * args.matrix_m;
        int start_k_a = 0;
        int start_k_b = batch_id * args.matrix_k;
        uint32_t wg_tile_k = args.matrix_k;
        uint32_t boundary_n = (start_n + wg_tile_n) > args.matrix_n
                ? args.matrix_n
                : (start_n + wg_tile_n);
        uint32_t boundary_m = (start_m + wg_tile_m)
                        > (args.matrix_m + batch_id * args.matrix_m)
                ? (args.matrix_m + batch_id * args.matrix_m)
                : (start_m + wg_tile_m);
        uint32_t boundary_k_a = wg_tile_k;
        uint32_t boundary_k_b = wg_tile_k + batch_id * args.matrix_k;
 
        uint32_t gemm_slm_base = slm_base;
        uint32_t gemm_nbarr_base = nbarrier_base;
        uint32_t epilogue_slm_base = gemm_slm_base + gemm_t::slm_size;
        uint32_t epilogue_nbarr_base = gemm_nbarr_base + gemm_t::barrier_count;
 
        // set up arguments
        work_group_t g;
        g.init(item.get_local_linear_id());
        mem_desc_a_t mem_desc_a;
        mem_desc_b_t mem_desc_b;
        mem_desc_c_t mem_desc_c;
        //setup for matA
        mem_desc_a.init(args.matA_base,
                {boundary_k_a, boundary_m, args.matA_ld}, {start_k_a, start_m});
 
        //setup for matB
        mem_desc_b.init(args.matB_base,
                {boundary_n, boundary_k_b, args.matB_ld}, {start_n, start_k_b});
 
        //setup for matC
        mem_desc_c.init(args.matC_base, {boundary_n, boundary_m, args.matC_ld},
                {start_n, start_m});
 
        uint32_t inner_loop_count = (wg_tile_k + k_stride - 1) / k_stride;
        gemm_args_t gemm_args(mem_desc_a, mem_desc_b, inner_loop_count);
        gemm_t gemm;
        epilogue_t epilogue;
 
        matAcc_t matAcc(0);
        gemm(g, matAcc, gemm_args, gemm_slm_base, gemm_nbarr_base);
        epilogue(g, matAcc, mem_desc_c, args.epilogue_args, epilogue_slm_base,
                epilogue_nbarr_base);
    }
};
 
 
} // namespace gpu::xetla::kernel