/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_grouped_gemm_multi_abd.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_grouped_gemm_multi_abd.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_grouped_gemm_multi_abd.hpp Source File

Go to the documentation of this file.

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include <iostream>
 #include <vector>
  
 #include "device_base.hpp"
  
 namespace ck {
 namespace tensor_operation {
 namespace device {
  
 struct GemmMultiABDDesc
 {
     ck::index_t M_, N_, K_;
  
     std::vector<ck::index_t> stride_As_;
     std::vector<ck::index_t> stride_Bs_;
     std::vector<ck::index_t> stride_Ds_;
  
     ck::index_t stride_C_;
 };
  
 /*
  * \brief Grouped Gemm Multi ABD
  *
  * C = a_op(A, A1...) * b_op(B, B1...)
  * E = cde_op(C, D0, D1, ...)
  *
  * \tparam AsLayout A layouts (tuple).
  * \tparam BsLayout B layouts (tuple).
  * \tparam DsLayout Ds layouts (tuple).
  * \tparam ELayout Output layout.
  * \tparam AsDataType A data types (tuple).
  * \tparam BsDataType B data types (tuple).
  * \tparam DsDataType D data types (tuple).
  * \tparam EDataType Output data type.
  * \tparam AElementwiseOperation A elementwise operation.
  * \tparam BElementwiseOperation B elementwise operation.
  * \tparam CDEElementwiseOperation C elementwise operation.
  */
 template <typename AsLayout,
           typename BsLayout,
           typename DsLayout,
           typename ELayout,
           typename AsDataType,
           typename BsDataType,
           typename DsDataType,
           typename EDataType,
           typename AElementwiseOperation,
           typename BElementwiseOperation,
           typename CDEElementwiseOperation>
 struct DeviceGroupedGemmMultiABD : public BaseOperator
 {
     static constexpr index_t NumATensor = AsDataType::Size();
     static constexpr index_t NumBTensor = BsDataType::Size();
     static constexpr index_t NumDTensor = DsDataType::Size();
  
     static_assert(AsLayout::Size() == AsDataType::Size(), "wrong! inconsistent NumATensor");
     static_assert(BsLayout::Size() == BsDataType::Size(), "wrong! inconsistent NumBTensor");
     static_assert(DsLayout::Size() == DsDataType::Size(), "wrong! inconsistent NumDTensor");
  
     /*
      * \brief Make argument pointer for grouped gemm multi abd.
      *
      * \param p_as A pointers to the A.
      * \param p_bs A pointers to the B.
      * \param p_ds A pointers to the Ds.
      * \param p_e A pointers to the E.
      * \param gemm_desc Gemm descriptors for each group.
      * \param a_element_op A elementwise operation object.
      * \param b_element_op B elementwise operation object.
      * \param cde_element_op CDE elementwise operation object.
      * \return Pointer to the argument.
      */
     virtual std::unique_ptr<BaseArgument>
     MakeArgumentPointer(std::vector<std::array<const void*, NumATensor>>& p_as,
                         std::vector<std::array<const void*, NumBTensor>>& p_bs,
                         std::vector<std::array<const void*, NumDTensor>>& p_ds,
                         std::vector<void*>& p_e,
                         std::vector<GemmMultiABDDesc>& gemm_desc,
                         AElementwiseOperation a_element_op   = AElementwiseOperation{},
                         BElementwiseOperation b_element_op   = BElementwiseOperation{},
                         CDEElementwiseOperation c_element_op = CDEElementwiseOperation{}) = 0;
  
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
  
     virtual void SetElementwiseOps(BaseArgument* p_arg,
                                    AElementwiseOperation a_element_op,
                                    BElementwiseOperation b_element_op,
                                    CDEElementwiseOperation cde_element_op) const = 0;
 };
  
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck