/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/inner_product_dpp8.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/inner_product_dpp8.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/utility/inner_product_dpp8.hpp Source File

Go to the documentation of this file.

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "amd_gemm_dpp.hpp"
 #include "data_type.hpp"
 #include "type_convert.hpp"
  
 namespace ck {
  
 namespace dpp8 {
  
 constexpr index_t lane_group_size = 8;
  
 template <int SrcLaneIdx>
 __device__ void inline_v_dot2c_dpp8_instr(const half2_t& a, const half2_t& b, float& c);
  
 // clang-format off
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<0>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[0, 0, 0, 0, 0, 0, 0, 0]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<1>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[1, 1, 1, 1, 1, 1, 1, 1]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<2>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[2, 2, 2, 2, 2, 2, 2, 2]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<3>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[3, 3, 3, 3, 3, 3, 3, 3]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<4>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[4, 4, 4, 4, 4, 4, 4, 4]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<5>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[5, 5, 5, 5, 5, 5, 5, 5]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<6>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[6, 6, 6, 6, 6, 6, 6, 6]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 template <>
 __device__ void inline_v_dot2c_dpp8_instr<7>(const half2_t& a, const half2_t& b, float& c){
     asm volatile("\n v_dot2c_f32_f16_dpp %0, %1, %2 dpp8:[7, 7, 7, 7, 7, 7, 7, 7]" : "=v"(c) : "v"(a), "v"(b), "0"(c));
 }
 // clang-format on
  
 template <int SrcLaneIdx, bool ShareA>
 __device__ void inline_v_dot2c_dpp8(const half2_t& a, const half2_t& b, float& c)
 {
     static_assert(SrcLaneIdx >= 0 && SrcLaneIdx < dpp8::lane_group_size,
                   "DPP8 src broadcast lane out of range <0, 7>.");
     if constexpr(ShareA)
     {
         inline_v_dot2c_dpp8_instr<SrcLaneIdx>(a, b, c);
     }
     else
     {
         inline_v_dot2c_dpp8_instr<SrcLaneIdx>(b, a, c);
     }
 }
  
 constexpr std::array<int, dpp8::lane_group_size> IntrinsicMaskDpp8 = {
     0,        // 0, 0, 0, 0, 0, 0, 0, 0
     2396745,  // 1, 1, 1, 1, 1, 1, 1, 1
     4793490,  // 2, 2, 2, 2, 2, 2, 2, 2
     7190235,  // 3, 3, 3, 3, 3, 3, 3, 3
     9586980,  // 4, 4, 4, 4, 4, 4, 4, 4
     11983725, // 5, 5, 5, 5, 5, 5, 5, 5
     14380470, // 6, 6, 6, 6, 6, 6, 6, 6
     16777215, // 7, 7, 7, 7, 7, 7, 7, 7
 };
  
 template <int SrcLaneIdx>
 constexpr int get_dpp_sel_mask_broadcast()
 {
     static_assert(SrcLaneIdx >= 0 && SrcLaneIdx < dpp8::lane_group_size,
                   "DPP8 src broadcast lane out of range <0, 7>.");
     return IntrinsicMaskDpp8[SrcLaneIdx];
 }
  
 template <int SrcLaneIdx>
 __device__ void intrinsic_fdot2_impl(const half2_t& a, const half2_t& b, float& c)
 {
     constexpr int sel_mask = get_dpp_sel_mask_broadcast<SrcLaneIdx>();
     const half2_t val_from_other_lane =
         bit_cast<half2_t>(__builtin_amdgcn_mov_dpp8(bit_cast<int>(a), sel_mask));
     c = __builtin_amdgcn_fdot2(val_from_other_lane, b, c, false);
 }
  
 template <int SrcLaneIdx, bool ShareA>
 __device__ void intrinsic_fdot2(const half2_t& a, const half2_t& b, float& c)
 {
     if constexpr(ShareA)
     {
         intrinsic_fdot2_impl<SrcLaneIdx>(a, b, c);
     }
     else
     {
         intrinsic_fdot2_impl<SrcLaneIdx>(b, a, c);
     }
 }
  
 template <typename TA, typename TB, typename TC, int SrcLaneIdx, bool ShareA>
 __device__ void inner_product_dpp(const TA& a, const TB& b, TC& c)
 {
 #if CK_USE_AMD_V_DOT_DPP8_INLINE_ASM
     inline_v_dot2c_dpp8<SrcLaneIdx, ShareA>(a, b, c);
 #else
     intrinsic_fdot2<SrcLaneIdx, ShareA>(a, b, c);
 #endif
 }
  
 } // namespace dpp8
  
 } // namespace ck