From 633dec0b17ae99b2817271d5b1d13ae8c6147ec3 Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Fri, 21 Apr 2023 12:57:08 -0700 Subject: [PATCH] refactor some code (#15566) ### Description 1. moved onnxruntime/contrib_ops/cuda/decoder to onnxruntime/contrib_ops/cuda/bert 2. create utils.cuh under /bert for shared implementations in decoder_masked_multihead_attention_impl_utils.h and rotary_embedding_util.h 3. refactored relative_attn_bias_impl.cu by reusing the template specializations in utils.cuh ### Motivation and Context --------- Co-authored-by: Ubuntu --- cmake/onnxruntime_rocm_hipify.cmake | 15 +- .../cuda/bert/add_bias_transpose.cu | 6 +- .../decoder_masked_multihead_attention.cc | 4 +- .../decoder_masked_multihead_attention.h | 0 .../decoder_masked_self_attention.cc | 4 +- .../decoder_masked_self_attention.h | 0 .../decoder_masked_multihead_attention_128.cu | 0 .../decoder_masked_multihead_attention_32.cu | 0 .../decoder_masked_multihead_attention_64.cu | 0 ...decoder_masked_multihead_attention_impl.cu | 0 .../decoder_masked_multihead_attention_impl.h | 0 ...er_masked_multihead_attention_impl_utils.h | 173 ++++ .../cuda/bert/relative_attn_bias_impl.cu | 30 +- .../cuda/bert/rotary_embedding_util.h | 817 +----------------- .../utils.cuh} | 442 ++++++---- 15 files changed, 485 insertions(+), 1006 deletions(-) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/decoder_masked_multihead_attention.cc (98%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/decoder_masked_multihead_attention.h (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/decoder_masked_self_attention.cc (98%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/decoder_masked_self_attention.h (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu (100%) rename onnxruntime/contrib_ops/cuda/{decoder => bert}/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h (100%) create mode 100644 onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h rename onnxruntime/contrib_ops/cuda/{decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h => bert/utils.cuh} (70%) diff --git a/cmake/onnxruntime_rocm_hipify.cmake b/cmake/onnxruntime_rocm_hipify.cmake index df2d68a6a4..77c79ed551 100644 --- a/cmake/onnxruntime_rocm_hipify.cmake +++ b/cmake/onnxruntime_rocm_hipify.cmake @@ -9,6 +9,11 @@ set(contrib_ops_excluded_files "bert/attention.h" "bert/attention_impl.cu" "bert/attention_softmax.h" + "bert/decoder_masked_multihead_attention.h" + "bert/decoder_masked_multihead_attention.cc" + "bert/decoder_masked_self_attention.h" + "bert/decoder_masked_self_attention.cc" + "bert/fastertransformer_decoder_attention/*" "bert/multihead_attention.cc" "bert/multihead_attention.h" "bert/fast_gelu_impl.cu" @@ -92,16 +97,6 @@ set(contrib_ops_excluded_files "cuda_contrib_kernels.h" "inverse.cc" "fused_conv.cc" - "decoder/decoder_masked_multihead_attention.h" - "decoder/decoder_masked_multihead_attention.cc" - "decoder/decoder_masked_self_attention.h" - "decoder/decoder_masked_self_attention.cc" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention.h" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu" - "decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu" ) if (NOT onnxruntime_ENABLE_ATEN) diff --git a/onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu b/onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu index cd3c33f6c7..8782201a64 100644 --- a/onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu +++ b/onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu @@ -276,9 +276,9 @@ __global__ void AddBiasTransposeQKV(int M, const T* input, const T* biases, T* o v_bias = *reinterpret_cast(&biases[2 * NH + n * H + head_idx]); } - q = add(q, q_bias); - k = add(k, k_bias); - v = add(v, v_bias); + q = add_vec(q, q_bias); + k = add_vec(k, k_bias); + v = add_vec(v, v_bias); const bool do_rotary = !is_masked && vec_size * tidx < rotary_embedding_dim; diff --git a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.cc similarity index 98% rename from onnxruntime/contrib_ops/cuda/decoder/decoder_masked_multihead_attention.cc rename to onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.cc index a990cb6557..5a1fbe2ad8 100644 --- a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.cc @@ -5,8 +5,8 @@ #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "core/platform/env_var_utils.h" #include "contrib_ops/cpu/bert/multihead_attention_helper.h" -#include "contrib_ops/cuda/decoder/decoder_masked_multihead_attention.h" -#include "contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h" +#include "contrib_ops/cuda/bert/decoder_masked_multihead_attention.h" +#include "contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; diff --git a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_multihead_attention.h b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.h similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/decoder_masked_multihead_attention.h rename to onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.h diff --git a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_self_attention.cc b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_self_attention.cc similarity index 98% rename from onnxruntime/contrib_ops/cuda/decoder/decoder_masked_self_attention.cc rename to onnxruntime/contrib_ops/cuda/bert/decoder_masked_self_attention.cc index 37febf5f43..5552bdebfc 100644 --- a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_self_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_self_attention.cc @@ -4,8 +4,8 @@ #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "core/platform/env_var_utils.h" -#include "contrib_ops/cuda/decoder/decoder_masked_self_attention.h" -#include "contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h" +#include "contrib_ops/cuda/bert/decoder_masked_self_attention.h" +#include "contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; diff --git a/onnxruntime/contrib_ops/cuda/decoder/decoder_masked_self_attention.h b/onnxruntime/contrib_ops/cuda/bert/decoder_masked_self_attention.h similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/decoder_masked_self_attention.h rename to onnxruntime/contrib_ops/cuda/bert/decoder_masked_self_attention.h diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu rename to onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu rename to onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu rename to onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu rename to onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h similarity index 100% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h rename to onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h diff --git a/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h new file mode 100644 index 0000000000..56258073e7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h @@ -0,0 +1,173 @@ +/* + * The implementation of this file is based on code provided by https://github.com/NVIDIA/FasterTransformer + * + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Modifications Copyright (c) Microsoft. +// Licensed under the MIT License. + +// Modifications: +// (1) Minor routine name changes for integration into the ORT code-base + +#pragma once + +#include "contrib_ops/cuda/bert/utils.cuh" + +using namespace onnxruntime::cuda; + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace decoder_masked_self_attention_details { + +//------------------------------------------------------------ +// Block_sum +//------------------------------------------------------------ + +template +inline __device__ float block_sum(float* red_smem, float sum) { + // Decompose the thread index into warp / lane. + int warp = threadIdx.x / WARP_SIZE; + int lane = threadIdx.x % WARP_SIZE; + + // Compute the sum per warp. +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Warp leaders store the data to shared memory. + if (lane == 0) { + red_smem[warp] = sum; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The warps compute the final sums. + if (lane < WARPS_PER_BLOCK) { + sum = red_smem[lane]; + } + + // Parallel reduction inside the warp. +#pragma unroll + for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Broadcast to other threads. + return __shfl_sync(uint32_t(-1), sum, 0); +} + +//------------------------------------------------------------ +// Shfl_Mask +//------------------------------------------------------------ + +inline __device__ constexpr uint32_t shfl_mask(int threads) { + return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u; +} + +//------------------------------------------------------------ +// Dot +//------------------------------------------------------------ + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +//------------------------------------------------------------ +// Qk_Dot +//------------------------------------------------------------ + +template +inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) { + using K_vec_acum = K_vec; + + // Compute the parallel products for Q*K^T (treat vector lanes separately). + K_vec_acum qk_vec = mul(q[0], k[0]); +#pragma unroll + for (int ii = 1; ii < N; ++ii) { + qk_vec = onnxruntime::cuda::fma(q[ii], k[ii], qk_vec); + } + + // Finalize the reduction across lanes. + float qk = sum(qk_vec); +#pragma unroll + for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) { + qk += __shfl_xor_sync(uint32_t(-1), qk, mask); + } + return qk; +} + +template +struct Qk_dot { + template + static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) { + return qk_dot_(q, k); + } +}; + +//------------------------------------------------------------ +// ThreadsPerValue +//------------------------------------------------------------ + +template +struct ThreadsPerValue { + static const int value = head_size * sizeof(T) / 16; +}; + +//------------------------------------------------------------ +// CalcDynamicBlockMemory +//------------------------------------------------------------ + +template +inline size_t CalcDynamicBlockMemory(const DecoderMaskedMultiHeadAttentionParams& params, + int threads_per_value, int threads_per_block) { + // The amount of shared memory needed to store the Q*K^T values in float. + + const int total_sequence_length = params.total_sequence_length; + size_t qk_sz = (((total_sequence_length + 3) / 4) * 16); + + // The extra memory needed if we are not using floats for the final logits. + size_t logits_sz = 0; + + if (sizeof(T) != 4) { + logits_sz = (((total_sequence_length + 3) / 4) * 4 * sizeof(T)); + } + + // The total size needed during softmax. + size_t softmax_sz = qk_sz + logits_sz; + + // The number of partial rows to reduce in the final reduction. + int rows_per_red = threads_per_block / threads_per_value; + + // The amount of storage needed to finalize the outputs. + size_t red_sz = rows_per_red * params.head_size * sizeof(T) / 2; + + // The max. + return std::max(softmax_sz, red_sz); +} + +} // namespace decoder_masked_self_attention_details +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu index ab21da6c31..e72828b8fc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/relative_attn_bias_impl.cu @@ -21,7 +21,7 @@ Licensed under the MIT License. #include #include "core/providers/cuda/cu_inc/common.cuh" #include "contrib_ops/cuda/bert/relative_attn_bias_impl.h" -#include "contrib_ops/cuda/bert/rotary_embedding_util.h" +#include "contrib_ops/cuda/bert/utils.cuh" namespace onnxruntime { namespace contrib { @@ -149,29 +149,11 @@ template Status LaunchRelPosAttnBiasKernel(cudaStream_t stream, const int max_threads_per_block); namespace { + template -struct TypeMapper; - -template <> -struct TypeMapper { - using type = float2; -}; - -template <> -struct TypeMapper { - using type = float4; -}; - -template <> -struct TypeMapper { - using type = half2; -}; - -template <> -struct TypeMapper { - using type = Half4; -}; +struct TypeMapper : public V_vec_m_ {}; +// The following operator overriding is not common so we put it in anonymous namespace #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ > 530 inline __device__ half2 operator*(const float a, const half2 b) { return __hmul2_rn(__float2half2_rn(a), b); @@ -292,13 +274,13 @@ Status LaunchGatedRelativePositionBiasKernel( dim3 grid(seq_len, num_heads, batch_size); if (seq_len % 4 == 0) { - using vec_type = typename TypeMapper::type; + using vec_type = typename TypeMapper::Type; GatedRelativePositionBiasKernelSmallD<<>>( reinterpret_cast(output), reinterpret_cast(rel_pos), qw, bias, eco_a, D, ldqw, equiv_seq_len); } else if (seq_len & 1 == 0) { - using vec_type = typename TypeMapper::type; + using vec_type = typename TypeMapper::Type; GatedRelativePositionBiasKernelSmallD<<>>( reinterpret_cast(output), reinterpret_cast(rel_pos), diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_util.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_util.h index 29ea8f331b..a35f626cad 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_util.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_util.h @@ -20,821 +20,14 @@ #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cu_inc/common.cuh" +#include "contrib_ops/cuda/bert/utils.cuh" + +using namespace onnxruntime::cuda; namespace onnxruntime { namespace cuda { -struct __align__(8) Half4 { - half2 x; - half2 y; -}; - -__device__ __forceinline__ Half4 operator+(const Half4& a, const Half4& b) { - Half4 r; - r.x = a.x + b.x; - r.y = a.y + b.y; - return r; -} - -__device__ __forceinline__ float2 operator+(const float2& a, const float2& b) { - return make_float2(a.x + b.x, a.y + b.y); -} - -__device__ __forceinline__ float4 operator+(const float4& a, const float4& b) { - return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); -} - #ifndef USE_ROCM -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct Float8_ { - float2 x; - float2 y; - float2 z; - float2 w; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct Float4_ { - float2 x; - float2 y; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct num_elems; -template <> -struct num_elems { - static constexpr int value = 1; -}; -template <> -struct num_elems { - static constexpr int value = 2; -}; -template <> -struct num_elems { - static constexpr int value = 4; -}; -template <> -struct num_elems { - static constexpr int value = 4; -}; -template <> -struct num_elems { - static constexpr int value = 8; -}; - -template <> -struct num_elems { - static constexpr int value = 2; -}; -template <> -struct num_elems { - static constexpr int value = 4; -}; -template <> -struct num_elems { - static constexpr int value = 8; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct Vec_t { - static constexpr int size = 0; -}; - -template <> -struct Vec_t { - using Type = float2; - static constexpr int size = 2; -}; - -template <> -struct Vec_t { - using Type = float4; - static constexpr int size = 4; -}; - -template <> -struct Vec_t { - using Type = Float8_; - static constexpr int size = 8; -}; - -template <> -struct Vec_t { - using Type = uint32_t; - static constexpr int size = 2; -}; - -template <> -struct Vec_t { - using Type = uint2; - static constexpr int size = 4; -}; - -template <> -struct Vec_t { - using Type = uint4; - static constexpr int size = 8; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float add(float a, float b) { - return a + b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 add(float2 a, float2 b) { - float2 c; - c.x = add(a.x, b.x); - c.y = add(a.y, b.y); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float4 add(float4 a, float4 b) { - float4 c; - c.x = add(a.x, b.x); - c.y = add(a.y, b.y); - c.z = add(a.z, b.z); - c.w = add(a.w, b.w); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float8_ add(Float8_ a, Float8_ b) { - Float8_ c; - c.x = add(a.x, b.x); - c.y = add(a.y, b.y); - c.z = add(a.z, b.z); - c.w = add(a.w, b.w); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint16_t add(uint16_t a, uint16_t b) { - uint16_t c; - asm volatile("add.f16 %0, %1, %2;\n" - : "=h"(c) - : "h"(a), "h"(b)); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint32_t add(uint32_t a, uint32_t b) { - uint32_t c; - asm volatile("add.f16x2 %0, %1, %2;\n" - : "=r"(c) - : "r"(a), "r"(b)); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint2 add(uint2 a, uint2 b) { - uint2 c; - c.x = add(a.x, b.x); - c.y = add(a.y, b.y); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint4 add(uint4 a, uint4 b) { - uint4 c; - c.x = add(a.x, b.x); - c.y = add(a.y, b.y); - c.z = add(a.z, b.z); - c.w = add(a.w, b.w); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint16_t float_to_half(float f) { - union { - uint32_t u32; - uint16_t u16[2]; - } tmp; -#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 // Is it better? - float zero = 0.f; - asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f)); -#else - asm volatile("cvt.rn.f16.f32 %0, %1;\n" - : "=h"(tmp.u16[0]) - : "f"(f)); -#endif - return tmp.u16[0]; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint32_t float2_to_half2(float2 f) { - union { - uint32_t u32; - uint16_t u16[2]; - } tmp; -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" - : "=r"(tmp.u32) - : "f"(f.y), "f"(f.x)); -#else - asm volatile("cvt.rn.f16.f32 %0, %1;\n" - : "=h"(tmp.u16[0]) - : "f"(f.x)); - asm volatile("cvt.rn.f16.f32 %0, %1;\n" - : "=h"(tmp.u16[1]) - : "f"(f.y)); -#endif - return tmp.u32; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float half_to_float(uint16_t h) { - float f; - asm volatile("cvt.f32.f16 %0, %1;\n" - : "=f"(f) - : "h"(h)); - return f; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 half2_to_float2(uint32_t v) { - uint16_t lo, hi; - asm volatile("mov.b32 {%0, %1}, %2;\n" - : "=h"(lo), "=h"(hi) - : "r"(v)); - return make_float2(half_to_float(lo), half_to_float(hi)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float add(float a, uint16_t b) { - return a + half_to_float(b); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 add(uint32_t a, float2 fb) { - float2 fa = half2_to_float2(a); - return add(fa, fb); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float4_ add(uint2 a, Float4_ fb) { - Float4_ fc; - fc.x = add(a.x, fb.x); - fc.y = add(a.y, fb.y); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float8_ add(uint4 a, Float8_ fb) { - Float8_ fc; - fc.x = add(a.x, fb.x); - fc.y = add(a.y, fb.y); - fc.z = add(a.z, fb.z); - fc.w = add(a.w, fb.w); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint32_t h0_h0(uint16_t a) { - uint32_t b; - asm volatile("mov.b32 %0, {%1, %1};" - : "=r"(b) - : "h"(a)); - return b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float fma(float a, float b, float c) { - return a * b + c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 fma(float2 a, float2 b, float2 c) { - float2 d; - d.x = fma(a.x, b.x, c.x); - d.y = fma(a.y, b.y, c.y); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 fma(float a, float2 b, float2 c) { - float2 d; - d.x = fma(a, b.x, c.x); - d.y = fma(a, b.y, c.y); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float4 fma(float4 a, float4 b, float4 c) { - float4 d; - d.x = fma(a.x, b.x, c.x); - d.y = fma(a.y, b.y, c.y); - d.z = fma(a.z, b.z, c.z); - d.w = fma(a.w, b.w, c.w); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float4 fma(float a, float4 b, float4 c) { - float4 d; - d.x = fma(a, b.x, c.x); - d.y = fma(a, b.y, c.y); - d.z = fma(a, b.z, c.z); - d.w = fma(a, b.w, c.w); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) { - Float4_ d; - d.x = fma(a, b.x, c.x); - d.y = fma(a, b.y, c.y); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) { - Float8_ d; - d.x = fma(a, b.x, c.x); - d.y = fma(a, b.y, c.y); - d.z = fma(a, b.z, c.z); - d.w = fma(a, b.w, c.w); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) { - uint32_t d; - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(d) - : "r"(a), "r"(b), "r"(c)); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) { - return fma(h0_h0(a), b, c); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) { - uint2 d; - d.x = fma(a.x, b.x, c.x); - d.y = fma(a.y, b.y, c.y); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) { - uint32_t s = h0_h0(a); - uint2 d; - d.x = fma(s, b.x, c.x); - d.y = fma(s, b.y, c.y); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) { - uint4 d; - d.x = fma(a.x, b.x, c.x); - d.y = fma(a.y, b.y, c.y); - d.z = fma(a.z, b.z, c.z); - d.w = fma(a.w, b.w, c.w); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) { - uint32_t s = h0_h0(a); - uint4 d; - d.x = fma(s, b.x, c.x); - d.y = fma(s, b.y, c.y); - d.z = fma(s, b.z, c.z); - d.w = fma(s, b.w, c.w); - return d; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float fma(uint16_t a, uint16_t b, float fc) { - float fa = half_to_float(a); - float fb = half_to_float(b); - return fa * fb + fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) { - float2 fa = half2_to_float2(a); - float2 fb = half2_to_float2(b); - return fma(fa, fb, fc); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) { - return fma(h0_h0(a), b, fc); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) { - Float4_ fd; - fd.x = fma(a.x, b.x, fc.x); - fd.y = fma(a.y, b.y, fc.y); - return fd; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) { - uint32_t s = h0_h0(a); - Float4_ fd; - fd.x = fma(s, b.x, fc.x); - fd.y = fma(s, b.y, fc.y); - return fd; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) { - Float8_ fd; - fd.x = fma(a.x, b.x, fc.x); - fd.y = fma(a.y, b.y, fc.y); - fd.z = fma(a.z, b.z, fc.z); - fd.w = fma(a.w, b.w, fc.w); - return fd; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) { - uint32_t s = h0_h0(a); - Float8_ fd; - fd.x = fma(s, b.x, fc.x); - fd.y = fma(s, b.y, fc.y); - fd.z = fma(s, b.z, fc.z); - fd.w = fma(s, b.w, fc.w); - return fd; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ Acc mul(A a, B b) { - return a * b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float mul(float a, float b) { - return a * b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float2 mul(float2 a, float2 b) { - float2 c; - c.x = a.x * b.x; - c.y = a.y * b.y; - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float2 mul(float a, float2 b) { - float2 c; - c.x = a * b.x; - c.y = a * b.y; - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float4 mul(float4 a, float4 b) { - float4 c; - c.x = a.x * b.x; - c.y = a.y * b.y; - c.z = a.z * b.z; - c.w = a.w * b.w; - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float4 mul(float a, float4 b) { - float4 c; - c.x = a * b.x; - c.y = a * b.y; - c.z = a * b.z; - c.w = a * b.w; - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ Float8_ mul(float a, Float8_ b) { - Float8_ c; - c.x = make_float2(a * b.x.x, a * b.x.y); - c.y = make_float2(a * b.y.x, a * b.y.y); - c.z = make_float2(a * b.z.x, a * b.z.y); - c.w = make_float2(a * b.w.x, a * b.w.y); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint16_t mul(uint16_t a, uint16_t b) { - uint16_t c; - asm volatile("mul.f16 %0, %1, %2;\n" - : "=h"(c) - : "h"(a), "h"(b)); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint32_t mul(uint32_t a, uint32_t b) { - uint32_t c; - asm volatile("mul.f16x2 %0, %1, %2;\n" - : "=r"(c) - : "r"(a), "r"(b)); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint32_t mul(uint16_t a, uint32_t b) { - return mul(h0_h0(a), b); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint2 mul(uint2 a, uint2 b) { - uint2 c; - c.x = mul(a.x, b.x); - c.y = mul(a.y, b.y); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint2 mul(uint16_t a, uint2 b) { - uint32_t s = h0_h0(a); - uint2 c; - c.x = mul(s, b.x); - c.y = mul(s, b.y); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint4 mul(uint4 a, uint4 b) { - uint4 c; - c.x = mul(a.x, b.x); - c.y = mul(a.y, b.y); - c.z = mul(a.z, b.z); - c.w = mul(a.w, b.w); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ uint4 mul(uint16_t a, uint4 b) { - uint32_t s = h0_h0(a); - uint4 c; - c.x = mul(s, b.x); - c.y = mul(s, b.y); - c.z = mul(s, b.z); - c.w = mul(s, b.w); - return c; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float mul(uint16_t a, uint16_t b) { - float fa = half_to_float(a); - float fb = half_to_float(b); - return fa * fb; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float mul(uint16_t a, float b) { - return half_to_float(a) * b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float2 mul(uint32_t a, uint32_t b) { - float2 fa = half2_to_float2(a); - float2 fb = half2_to_float2(b); - return mul(fa, fb); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ float2 mul(uint16_t a, uint32_t b) { - return mul(h0_h0(a), b); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ Float4_ mul(uint2 a, uint2 b) { - Float4_ fc; - fc.x = mul(a.x, b.x); - fc.y = mul(a.y, b.y); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ Float4_ mul(uint16_t a, uint2 b) { - uint32_t s = h0_h0(a); - Float4_ fc; - fc.x = mul(s, b.x); - fc.y = mul(s, b.y); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ Float8_ mul(uint4 a, uint4 b) { - Float8_ fc; - fc.x = mul(a.x, b.x); - fc.y = mul(a.y, b.y); - fc.z = mul(a.z, b.z); - fc.w = mul(a.w, b.w); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline __device__ Float8_ mul(uint16_t a, uint4 b) { - uint32_t s = h0_h0(a); - Float8_ fc; - fc.x = mul(s, b.x); - fc.y = mul(s, b.y); - fc.z = mul(s, b.z); - fc.w = mul(s, b.w); - return fc; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(float v) { - return v; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(float2 v) { - return v.x + v.y; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(float4 v) { - return v.x + v.y + v.z + v.w; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(uint16_t v) { - return half_to_float(v); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(uint32_t v) { - float2 tmp = half2_to_float2(v); - return tmp.x + tmp.y; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(uint2 v) { - uint32_t c = add(v.x, v.y); - return sum(c); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(uint4 v) { -#if 1 - uint32_t c = add(v.x, v.y); - c = add(c, v.z); - c = add(c, v.w); -#else - uint32_t c = add(v.x, v.y); - uint32_t d = add(v.z, v.w); - c = add(c, d); -#endif - return sum(c); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(Float4_ v) { - return v.x.x + v.x.y + v.y.x + v.y.y; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ float sum(Float8_ v) { - return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ float dot(T a, T b) { - return sum(mul(a, b)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ float dot(T a, T b) { - return sum(mul(a, b)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ void zero(uint16_t& dst) { - dst = uint16_t(0); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ void zero(T& dst) { - constexpr int WORDS = sizeof(T) / 4; - union { - T raw; - uint32_t words[WORDS]; - } tmp; -#pragma unroll - for (int ii = 0; ii < WORDS; ++ii) { - tmp.words[ii] = 0u; - } - dst = tmp.raw; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// inline __device__ float2 rotary_embedding_coefficient(const int zid, const int rot_embed_dim, const float t_step) { const float inv_freq = t_step / pow(10000.0f, zid / (float)rot_embed_dim); @@ -849,9 +42,9 @@ inline __device__ float2 rotary_embedding_transform(const float2 v, const float2 } inline __device__ uint32_t rotary_embedding_transform(const uint32_t v, const float2 coef) { - float2 fv = half2_to_float2(v); + float2 fv = Half2ToFloat2(v); float2 rot_fv = rotary_embedding_transform(fv, coef); - return float2_to_half2(rot_fv); + return Float2ToHalf2(rot_fv); } inline __device__ void apply_rotary_embedding(float& q, int zid, int rot_embed_dim, int t_step) { diff --git a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h b/onnxruntime/contrib_ops/cuda/bert/utils.cuh similarity index 70% rename from onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h rename to onnxruntime/contrib_ops/cuda/bert/utils.cuh index 42d54e38d4..dd61288a3c 100644 --- a/onnxruntime/contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl_utils.h +++ b/onnxruntime/contrib_ops/cuda/bert/utils.cuh @@ -25,12 +25,123 @@ #pragma once #include "core/providers/cuda/cuda_common.h" -#include "decoder_masked_multihead_attention_impl_utils.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +using namespace onnxruntime::cuda; namespace onnxruntime { -namespace contrib { namespace cuda { -namespace decoder_masked_self_attention_details { + +struct __align__(8) Half4 { + half2 x; + half2 y; +}; + +__device__ __forceinline__ Half4 operator+(const Half4& a, const Half4& b) { + Half4 r; + r.x = a.x + b.x; + r.y = a.y + b.y; + return r; +} + +__device__ __forceinline__ float2 operator+(const float2& a, const float2& b) { + return make_float2(a.x + b.x, a.y + b.y); +} + +__device__ __forceinline__ float4 operator+(const float4& a, const float4& b) { + return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); +} + +struct Float8_ { + float2 x; + float2 y; + float2 z; + float2 w; +}; + +struct Float4_ { + float2 x; + float2 y; +}; + +#ifndef USE_ROCM + +template +struct num_elems; +template <> +struct num_elems { + static constexpr int value = 1; +}; +template <> +struct num_elems { + static constexpr int value = 2; +}; +template <> +struct num_elems { + static constexpr int value = 4; +}; +template <> +struct num_elems { + static constexpr int value = 4; +}; +template <> +struct num_elems { + static constexpr int value = 8; +}; + +template <> +struct num_elems { + static constexpr int value = 2; +}; +template <> +struct num_elems { + static constexpr int value = 4; +}; +template <> +struct num_elems { + static constexpr int value = 8; +}; + +template +struct Vec_t { + static constexpr int size = 0; +}; + +template <> +struct Vec_t { + using Type = float2; + static constexpr int size = 2; +}; + +template <> +struct Vec_t { + using Type = float4; + static constexpr int size = 4; +}; + +template <> +struct Vec_t { + using Type = Float8_; + static constexpr int size = 8; +}; + +template <> +struct Vec_t { + using Type = uint32_t; + static constexpr int size = 2; +}; + +template <> +struct Vec_t { + using Type = uint2; + static constexpr int size = 4; +}; + +template <> +struct Vec_t { + using Type = uint4; + static constexpr int size = 8; +}; //------------------------------------------------------------ // Qk_vec @@ -163,6 +274,16 @@ struct V_vec_m_ { using Type = uint4; }; +template <> +struct V_vec_m_ { + using Type = half2; +}; + +template <> +struct V_vec_m_ { + using Type = Half4; +}; + template struct V_vec_k_ { using Type = typename V_vec_m_::Type; @@ -190,18 +311,6 @@ struct V_vec_acum_fp32_ { using Type = float4; }; -struct Float8_ { - float2 x; - float2 y; - float2 z; - float2 w; -}; - -struct Float4_ { - float2 x; - float2 y; -}; - template <> struct V_vec_acum_fp32_ { using Type = float2; @@ -219,7 +328,9 @@ struct V_vec_acum_fp32_ { // Zero //------------------------------------------------------------ -// TODO: fp16 ? +inline __device__ void zero(uint16_t& dst) { + dst = uint16_t(0); +} template inline __device__ void zero(T& dst) { @@ -235,6 +346,14 @@ inline __device__ void zero(T& dst) { dst = tmp.raw; } +inline __device__ uint32_t h0_h0(uint16_t a) { + uint32_t b; + asm volatile("mov.b32 %0, {%1, %1};" + : "=r"(b) + : "h"(a)); + return b; +} + //------------------------------------------------------------ // vec_conversion //------------------------------------------------------------ @@ -268,6 +387,15 @@ inline __device__ float4 add_vec(float4 a, float4 b) { return c; } +inline __device__ Float8_ add_vec(Float8_ a, Float8_ b) { + Float8_ c; + c.x = add_vec(a.x, b.x); + c.y = add_vec(a.y, b.y); + c.z = add_vec(a.z, b.z); + c.w = add_vec(a.w, b.w); + return c; +} + inline __device__ float HalfToFloat(uint16_t h) { float f; asm volatile("cvt.f32.f16 %0, %1;\n" @@ -378,6 +506,14 @@ inline __device__ float sum(uint4 v) { return sum(c); } +inline __device__ float sum(Float4_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y; +} + +inline __device__ float sum(Float8_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y; +} + //------------------------------------------------------------ // Mul //------------------------------------------------------------ @@ -400,6 +536,14 @@ inline __device__ float2 mul(float2 a, float2 b) { return c; } +template <> +inline __device__ float2 mul(float a, float2 b) { + float2 c; + c.x = a * b.x; + c.y = a * b.y; + return c; +} + template <> inline __device__ float4 mul(float4 a, float4 b) { float4 c; @@ -410,6 +554,26 @@ inline __device__ float4 mul(float4 a, float4 b) { return c; } +template <> +inline __device__ float4 mul(float a, float4 b) { + float4 c; + c.x = a * b.x; + c.y = a * b.y; + c.z = a * b.z; + c.w = a * b.w; + return c; +} + +template <> +inline __device__ Float8_ mul(float a, Float8_ b) { + Float8_ c; + c.x = make_float2(a * b.x.x, a * b.x.y); + c.y = make_float2(a * b.y.x, a * b.y.y); + c.z = make_float2(a * b.z.x, a * b.z.y); + c.w = make_float2(a * b.w.x, a * b.w.y); + return c; +} + template <> inline __device__ uint16_t mul(uint16_t a, uint16_t b) { uint16_t c; @@ -428,6 +592,11 @@ inline __device__ uint32_t mul(uint32_t a, uint32_t b) { return c; } +template <> +inline __device__ uint32_t mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + template <> inline __device__ uint2 mul(uint2 a, uint2 b) { uint2 c; @@ -436,6 +605,15 @@ inline __device__ uint2 mul(uint2 a, uint2 b) { return c; } +template <> +inline __device__ uint2 mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + uint2 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + return c; +} + template <> inline __device__ uint4 mul(uint4 a, uint4 b) { uint4 c; @@ -446,6 +624,79 @@ inline __device__ uint4 mul(uint4 a, uint4 b) { return c; } +template <> +inline __device__ uint4 mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + uint4 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + c.z = mul(s, b.z); + c.w = mul(s, b.w); + return c; +} + +template <> +inline __device__ float mul(uint16_t a, float b) { + return HalfToFloat(a) * b; +} + +template <> +inline __device__ float mul(uint16_t a, uint16_t b) { + float fa = HalfToFloat(a); + float fb = HalfToFloat(b); + return fa * fb; +} + +template <> +inline __device__ float2 mul(uint32_t a, uint32_t b) { + float2 fa = Half2ToFloat2(a); + float2 fb = Half2ToFloat2(b); + return mul(fa, fb); +} + +template <> +inline __device__ float2 mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + +template <> +inline __device__ Float4_ mul(uint2 a, uint2 b) { + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +template <> +inline __device__ Float4_ mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint4 a, uint4 b) { + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + //------------------------------------------------------------ // Fma //------------------------------------------------------------ @@ -485,14 +736,6 @@ inline __device__ float4 fma(float a, float4 b, float4 c) { return d; } -inline __device__ uint32_t h0_h0(uint16_t a) { - uint32_t b; - asm volatile("mov.b32 %0, {%1, %1};" - : "=r"(b) - : "h"(a)); - return b; -} - inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) { uint32_t d; asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" @@ -555,6 +798,13 @@ inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) { return fma(h0_h0(a), b, fc); } +inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) { + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) { uint32_t s = h0_h0(a); Float4_ fd; @@ -563,6 +813,15 @@ inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) { return fd; } +inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) { + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) { uint32_t s = h0_h0(a); Float8_ fd; @@ -573,103 +832,6 @@ inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) { return fd; } -//------------------------------------------------------------ -// Block_sum -//------------------------------------------------------------ - -template -inline __device__ float block_sum(float* red_smem, float sum) { - // Decompose the thread index into warp / lane. - int warp = threadIdx.x / WARP_SIZE; - int lane = threadIdx.x % WARP_SIZE; - - // Compute the sum per warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - sum += __shfl_xor_sync(uint32_t(-1), sum, mask); - } - - // Warp leaders store the data to shared memory. - if (lane == 0) { - red_smem[warp] = sum; - } - - // Make sure the data is in shared memory. - __syncthreads(); - - // The warps compute the final sums. - if (lane < WARPS_PER_BLOCK) { - sum = red_smem[lane]; - } - - // Parallel reduction inside the warp. -#pragma unroll - for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { - sum += __shfl_xor_sync(uint32_t(-1), sum, mask); - } - - // Broadcast to other threads. - return __shfl_sync(uint32_t(-1), sum, 0); -} - -//------------------------------------------------------------ -// Shfl_Mask -//------------------------------------------------------------ - -inline __device__ constexpr uint32_t shfl_mask(int threads) { - return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u; -} - -//------------------------------------------------------------ -// Dot -//------------------------------------------------------------ - -template -inline __device__ float dot(T a, T b) { - return sum(mul(a, b)); -} - -//------------------------------------------------------------ -// Qk_Dot -//------------------------------------------------------------ - -template -inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) { - using K_vec_acum = K_vec; - - // Compute the parallel products for Q*K^T (treat vector lanes separately). - K_vec_acum qk_vec = mul(q[0], k[0]); -#pragma unroll - for (int ii = 1; ii < N; ++ii) { - qk_vec = fma(q[ii], k[ii], qk_vec); - } - - // Finalize the reduction across lanes. - float qk = sum(qk_vec); -#pragma unroll - for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) { - qk += __shfl_xor_sync(uint32_t(-1), qk, mask); - } - return qk; -} - -template -struct Qk_dot { - template - static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) { - return qk_dot_(q, k); - } -}; - -//------------------------------------------------------------ -// ThreadsPerValue -//------------------------------------------------------------ - -template -struct ThreadsPerValue { - static const int value = head_size * sizeof(T) / 16; -}; - //------------------------------------------------------------ // ConvertFromFloat //------------------------------------------------------------ @@ -691,9 +853,14 @@ inline __device__ uint16_t FloatToHalf(float f) { uint32_t u32; uint16_t u16[2]; } tmp; +#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 // Is it better? + float zero = 0.f; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f)); +#else asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f)); +#endif return tmp.u16[0]; } @@ -736,39 +903,8 @@ inline __device__ void ConvertFromFloat(uint4& dst, Float8_ src) { dst.z = Float2ToHalf2(src.z); dst.w = Float2ToHalf2(src.w); } -//------------------------------------------------------------ -// CalcDynamicBlockMemory -//------------------------------------------------------------ -template -inline size_t CalcDynamicBlockMemory(const DecoderMaskedMultiHeadAttentionParams& params, - int threads_per_value, int threads_per_block) { - // The amount of shared memory needed to store the Q*K^T values in float. +#endif - const int total_sequence_length = params.total_sequence_length; - size_t qk_sz = (((total_sequence_length + 3) / 4) * 16); - - // The extra memory needed if we are not using floats for the final logits. - size_t logits_sz = 0; - - if (sizeof(T) != 4) { - logits_sz = (((total_sequence_length + 3) / 4) * 4 * sizeof(T)); - } - - // The total size needed during softmax. - size_t softmax_sz = qk_sz + logits_sz; - - // The number of partial rows to reduce in the final reduction. - int rows_per_red = threads_per_block / threads_per_value; - - // The amount of storage needed to finalize the outputs. - size_t red_sz = rows_per_red * params.head_size * sizeof(T) / 2; - - // The max. - return std::max(softmax_sz, red_sz); -} - -} // namespace decoder_masked_self_attention_details } // namespace cuda -} // namespace contrib } // namespace onnxruntime