From eb827bd3e5cf4110023ccebaf1789ce5545561ec Mon Sep 17 00:00:00 2001 From: Xinya Zhang Date: Thu, 11 Aug 2022 21:32:08 -0500 Subject: [PATCH] [ROCm] NGramRepeatBlock, LongformerAttention and DecoderAttention Ops (#11971) * [ROCm] enable NGramRepeatBlock Op * [ROCm] Enable testing ROCm in NGramRepeatBlockTest.NGramSize_3 Also link onnxruntime_test_all with amdhip64 when USE_ROCM=1 * [ROCm] add LongformerAttention Op * [ROCm] Enable LongformerAttentionTest * [ROCm] Add DecoderAttention Op * Enable DecoderAttention Test for ROCm. * [ROCM] Updates according to reviews --- cmake/onnxruntime_unittests.cmake | 5 + .../contrib_ops/cuda/bert/attention.cc | 4 +- .../contrib_ops/cuda/bert/attention_concat.cu | 2 +- .../contrib_ops/cuda/bert/attention_impl.cu | 6 +- .../contrib_ops/cuda/bert/attention_impl.h | 1 + .../cuda/bert/attention_transpose.cu | 2 +- .../cuda/bert/decoder_attention.cc | 6 +- .../cuda/bert/longformer_attention.cc | 10 +- .../cuda/bert/longformer_attention_impl.cu | 6 +- .../cuda/bert/longformer_attention_softmax.cu | 4 +- .../cuda/bert/longformer_global_impl.cu | 1 + .../cuda/bert/transformer_cuda_common.h | 4 +- .../contrib_ops/rocm/bert/attention_impl.cu | 244 ++++++++++++++++++ .../contrib_ops/rocm/bert/attention_impl.h | 200 ++++++++++++++ .../contrib_ops/rocm/rocm_contrib_kernels.cc | 15 +- .../providers/rocm/shared_inc/fpgeneric.h | 143 ++++++---- .../contrib_ops/decoder_attention_op_test.cc | 10 +- .../longformer_attention_op_test.cc | 16 +- .../contrib_ops/ngram_repeat_block_op_test.cc | 7 + tools/ci_build/amd_hipify.py | 32 ++- 20 files changed, 622 insertions(+), 96 deletions(-) create mode 100644 onnxruntime/contrib_ops/rocm/bert/attention_impl.h diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 32c2642f6d..3ae12479de 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -586,6 +586,11 @@ if(onnxruntime_USE_XNNPACK) list(APPEND onnxruntime_test_providers_libs onnxruntime_providers_xnnpack) endif() +if(onnxruntime_USE_ROCM) + find_library(HIP_LIB amdhip64 REQUIRED) + list(APPEND onnxruntime_test_providers_libs ${HIP_LIB}) +endif() + if(WIN32) if (onnxruntime_USE_NUPHAR_TVM) diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc index e13ede66d4..580e6b22e3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc @@ -1,8 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "attention.h" -#include "attention_impl.h" +#include "contrib_ops/cuda/bert/attention.h" +#include "contrib_ops/cuda/bert/attention_impl.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_concat.cu b/onnxruntime/contrib_ops/cuda/bert/attention_concat.cu index 88a75f36d8..3f9ed50e89 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_concat.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_concat.cu @@ -2,7 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cuda/cuda_common.h" -#include "attention_impl.h" +#include "contrib_ops/cuda/bert/attention_impl.h" using namespace onnxruntime::cuda; diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 368f520515..835213e9b7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -25,9 +25,9 @@ limitations under the License. #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" -#include "attention_impl.h" -#include "attention_softmax.h" -#include "transformer_common.h" +#include "contrib_ops/cuda/bert/attention_impl.h" +#include "contrib_ops/cuda/bert/attention_softmax.h" +#include "contrib_ops/cuda/bert/transformer_common.h" using namespace onnxruntime::cuda; using namespace cub; diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h index 9affcf2f49..9c1e6da801 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h @@ -3,6 +3,7 @@ #pragma once #include "core/providers/cuda/shared_inc/cuda_utils.h" +#include #include namespace onnxruntime { diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_transpose.cu b/onnxruntime/contrib_ops/cuda/bert/attention_transpose.cu index 1c9c7a8532..ce203ac8d8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_transpose.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_transpose.cu @@ -18,7 +18,7 @@ limitations under the License. */ #include "core/providers/cuda/cuda_common.h" -#include "attention_impl.h" +#include "contrib_ops/cuda/bert/attention_impl.h" using namespace onnxruntime::cuda; diff --git a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc index 53392c7979..411f3ff6f2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "attention_impl.h" -#include "decoder_attention.h" -#include "transformer_cuda_common.h" +#include "contrib_ops/cuda/bert/attention_impl.h" +#include "contrib_ops/cuda/bert/decoder_attention.h" +#include "contrib_ops/cuda/bert/transformer_cuda_common.h" #include "core/framework/op_kernel.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc index 5e783b0efa..004fd98e08 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc @@ -3,11 +3,11 @@ #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "core/platform/env_var_utils.h" -#include "longformer_attention.h" -#include "longformer_global_impl.h" -#include "longformer_attention_impl.h" -#include "transformer_cuda_common.h" -#include "transformer_common.h" +#include "contrib_ops/cuda/bert/longformer_attention.h" +#include "contrib_ops/cuda/bert/longformer_global_impl.h" +#include "contrib_ops/cuda/bert/longformer_attention_impl.h" +#include "contrib_ops/cuda/bert/transformer_cuda_common.h" +#include "contrib_ops/cuda/bert/transformer_common.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu index f301c80d5a..bf4c23328b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu @@ -26,9 +26,9 @@ limitations under the License. #include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" -#include "longformer_attention_impl.h" -#include "attention_impl.h" -#include "longformer_attention_softmax.h" +#include "contrib_ops/cuda/bert/longformer_attention_impl.h" +#include "contrib_ops/cuda/bert/attention_impl.h" +#include "contrib_ops/cuda/bert/longformer_attention_softmax.h" #include "core/common/safeint.h" using namespace onnxruntime::cuda; diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu index 0389c64a16..831d4867b1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu @@ -25,8 +25,8 @@ limitations under the License. #include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" -#include "longformer_attention_softmax.h" -#include "attention_impl.h" +#include "contrib_ops/cuda/bert/longformer_attention_softmax.h" +#include "contrib_ops/cuda/bert/attention_impl.h" using namespace onnxruntime::cuda; using namespace cub; diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu index 395c19f39b..492c24330f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu @@ -15,6 +15,7 @@ limitations under the License. */ #include +#include #include "core/providers/cuda/cuda_common.h" #include "longformer_global_impl.h" diff --git a/onnxruntime/contrib_ops/cuda/bert/transformer_cuda_common.h b/onnxruntime/contrib_ops/cuda/bert/transformer_cuda_common.h index f5828c9970..faf9310c4c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/transformer_cuda_common.h +++ b/onnxruntime/contrib_ops/cuda/bert/transformer_cuda_common.h @@ -17,7 +17,7 @@ class AutoDestoryCudaEvent { ~AutoDestoryCudaEvent() { if (cuda_event_ != nullptr) - cudaEventDestroy(cuda_event_); + (void)cudaEventDestroy(cuda_event_); } cudaEvent_t& Get() { @@ -30,4 +30,4 @@ class AutoDestoryCudaEvent { } // namespace cuda } // namespace contrib -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu b/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu index 83514aff9b..c87c12a1f9 100644 --- a/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/rocm/bert/attention_impl.cu @@ -208,6 +208,250 @@ bool LaunchAttentionKernel( reinterpret_cast(present), use_persistent_softmax); } } + + +template +bool DecoderQkvToContext( + const hipDeviceProp_t& prop, + hipStream_t stream, + rocblas_handle& rocblas, + const size_t element_size, + const int batch_size, + const int sequence_length, + const int kv_sequence_length, + const int num_heads, + const int head_size, + const bool static_kv, + const bool use_past, + const bool has_layer_state, + const bool has_key_padding_mask, + const T* gemm_query_buffer, + const T* gemm_kv_buffer, + const bool* key_padding_mask, + const T* key_cache, + const T* value_cache, + T* qkv_buffer, + T* workspace_buffer, + T* output, + T* new_key_cache, + T* new_value_cache) { + const int max_threads_per_block = prop.maxThreadsPerBlock; + const int BN = batch_size * num_heads; + const int BHN = BN * head_size; + const int BNS = BN * sequence_length; + const int k_buffer_offset = sequence_length * BHN; + const int v_buffer_offset = (sequence_length + kv_sequence_length) * BHN; + + T* temp_qkv_buffer = workspace_buffer; + + const T* q = qkv_buffer; + // transpose q and copy them to qkv_buffer + if (!LaunchTransQkv(stream, 1, sequence_length, batch_size, head_size, + num_heads, max_threads_per_block, true, gemm_query_buffer, qkv_buffer)) { + return false; + } + + const T* k = qkv_buffer + k_buffer_offset; + const T* v = qkv_buffer + v_buffer_offset; + if (!has_layer_state || !use_past) { + if (!static_kv) { + // transpose kv and copy them to qkv_buffer + if (!LaunchTransQkv(stream, 2, sequence_length, batch_size, head_size, num_heads, + max_threads_per_block, true, gemm_kv_buffer, qkv_buffer + k_buffer_offset)) { + return false; + } + } else { + // transpose kv and copy them to qkv_buffer + if (!LaunchTransQkv(stream, 2, kv_sequence_length, batch_size, head_size, num_heads, + max_threads_per_block, true, gemm_kv_buffer, qkv_buffer + k_buffer_offset)) { + return false; + } + } + } else { + if (!static_kv) { + // transpose kv and copy them to temp_buffer + if (!LaunchTransQkv(stream, 2, sequence_length, batch_size, head_size, num_heads, + max_threads_per_block, true, gemm_kv_buffer, temp_qkv_buffer)) { + return false; + } + // concat cache-k with k and copy to qkv_buffer + if (nullptr != key_cache && !LaunchConcatTensorToTensor(stream, kv_sequence_length, sequence_length, + batch_size, head_size, num_heads, + max_threads_per_block, 1, key_cache, + temp_qkv_buffer, qkv_buffer + k_buffer_offset)) { + return false; + } + // concat cache-v with v and copy to qkv_buffer + if (nullptr != value_cache && !LaunchConcatTensorToTensor(stream, kv_sequence_length, sequence_length, + batch_size, head_size, num_heads, + max_threads_per_block, 1, value_cache, + temp_qkv_buffer + k_buffer_offset, + qkv_buffer + v_buffer_offset)) { + return false; + } + } + } + + if (has_layer_state) { + if (use_past && static_kv) { + CHECK_ROCM(hipMemcpyAsync(new_key_cache, key_cache, + kv_sequence_length * BHN * sizeof(T), hipMemcpyDeviceToDevice, stream)); + CHECK_ROCM(hipMemcpyAsync(new_value_cache, value_cache, + kv_sequence_length * BHN * sizeof(T), hipMemcpyDeviceToDevice, stream)); + } else { + CHECK_ROCM(hipMemcpyAsync(new_key_cache, k, + kv_sequence_length * BHN * sizeof(T), hipMemcpyDeviceToDevice, stream)); + CHECK_ROCM(hipMemcpyAsync(new_value_cache, v, + kv_sequence_length * BHN * sizeof(T), hipMemcpyDeviceToDevice, stream)); + } + } + + // scratch1: BxNxSxS* buffer + // scratch2: BxNxSxS* buffer + // scratch3: BxNxSxH buffer + T* scratch1 = temp_qkv_buffer + 3 * BHN * sequence_length; + T* scratch2 = scratch1 + BNS * kv_sequence_length; + T* scratch3 = scratch2 + BNS * kv_sequence_length; + + // compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS* + // Q: BxNxSxH, K (present_k): BxNxS*xH, Q*K': BxNxSxS* + const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); + const int temp_matrix_size = sequence_length * kv_sequence_length; + float one = 1.0f; + float zero = 0.f; + + float alpha = rsqrt_head_size; + const int strideA = kv_sequence_length * head_size; + const int strideB = sequence_length * head_size; + if (use_past && static_kv) { + if (!ROCBLAS_CALL(rocblasGemmStridedBatchedHelper( + rocblas, rocblas_operation_transpose, rocblas_operation_none, + kv_sequence_length, sequence_length, head_size, &alpha, key_cache, head_size, strideA, + q, head_size, strideB, &zero, scratch1, kv_sequence_length, temp_matrix_size, BN))) { + return false; + } + } else { + if (!ROCBLAS_CALL(rocblasGemmStridedBatchedHelper( + rocblas, rocblas_operation_transpose, rocblas_operation_none, + kv_sequence_length, sequence_length, head_size, &alpha, k, head_size, strideA, + q, head_size, strideB, &zero, scratch1, kv_sequence_length, temp_matrix_size, BN))) { + return false; + } + } + + if (has_key_padding_mask) { + if (!ComputeSoftmaxWithRawMask(stream, kv_sequence_length, sequence_length, batch_size, + num_heads, nullptr, key_padding_mask, nullptr, scratch1, scratch2, + false, 1, 2, static_cast(0), false, nullptr)) { + return false; + } + } else { + if (!ComputeSoftmax(stream, kv_sequence_length, sequence_length, batch_size, + num_heads, nullptr, scratch1, scratch2, false)) { + return false; + } + } + + // compute P*V (as V*P), and store in scratch3: BxNxSxH + if (use_past && static_kv) { + if (!ROCBLAS_CALL(rocblasGemmStridedBatchedHelper( + rocblas, rocblas_operation_none, rocblas_operation_none, + head_size, sequence_length, kv_sequence_length, &one, value_cache, head_size, strideA, + scratch2, kv_sequence_length, temp_matrix_size, &zero, scratch3, head_size, strideB, BN))) { + return false; + } + } else { + if (!ROCBLAS_CALL(rocblasGemmStridedBatchedHelper( + rocblas, rocblas_operation_none, rocblas_operation_none, + head_size, sequence_length, kv_sequence_length, &one, v, head_size, strideA, + scratch2, kv_sequence_length, temp_matrix_size, &zero, scratch3, head_size, strideB, BN))) { + return false; + } + } + + // scratch3 is BxNxSxH, transpose to output SxBxNxH + return LaunchTransCtx(stream, sequence_length, batch_size, head_size, + num_heads, max_threads_per_block, true, scratch3, output); +} + + +bool LaunchDecoderAttentionKernel( + const hipDeviceProp_t& prop, + hipStream_t stream, + rocblas_handle& rocblas, + const size_t element_size, + const int batch_size, + const int sequence_length, + const int kv_sequence_length, + const int num_heads, + const int head_size, + const bool static_kv, + const bool use_past, + const bool has_layer_state, + const bool has_key_padding_mask, + const void* gemm_query_buffer, + const void* gemm_kv_buffer, + const bool* key_padding_mask, + const void* key_cache, + const void* value_cache, + void* qkv_buffer, + void* workspace_buffer, + void* output, + void* new_key_cache, + void* new_value_cache) { + if (element_size == 2) { + return DecoderQkvToContext( + prop, + stream, + rocblas, + element_size, + batch_size, + sequence_length, + kv_sequence_length, + num_heads, + head_size, + static_kv, + use_past, + has_layer_state, + has_key_padding_mask, + reinterpret_cast(gemm_query_buffer), + reinterpret_cast(gemm_kv_buffer), + key_padding_mask, + reinterpret_cast(key_cache), + reinterpret_cast(value_cache), + reinterpret_cast(qkv_buffer), + reinterpret_cast(workspace_buffer), + reinterpret_cast(output), + reinterpret_cast(new_key_cache), + reinterpret_cast(new_value_cache)); + } else { + return DecoderQkvToContext( + prop, + stream, + rocblas, + element_size, + batch_size, + sequence_length, + kv_sequence_length, + num_heads, + head_size, + static_kv, + use_past, + has_layer_state, + has_key_padding_mask, + reinterpret_cast(gemm_query_buffer), + reinterpret_cast(gemm_kv_buffer), + key_padding_mask, + reinterpret_cast(key_cache), + reinterpret_cast(value_cache), + reinterpret_cast(qkv_buffer), + reinterpret_cast(workspace_buffer), + reinterpret_cast(output), + reinterpret_cast(new_key_cache), + reinterpret_cast(new_value_cache)); + } +} + } // namespace rocm } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_impl.h b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h new file mode 100644 index 0000000000..7a17c2c38b --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h @@ -0,0 +1,200 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include "core/providers/rocm/shared_inc/rocm_utils.h" + +namespace onnxruntime { +namespace contrib { +namespace rocm { +size_t GetAttentionScratchSize(size_t element_size, int batch_size, + int num_heads, int sequence_length, int all_sequence_length); + +size_t GetAttentionWorkspaceSize( + size_t element_size, + int batchsize, + int num_heads, + int head_size, + int sequence_length, + int past_sequence_length); + +bool LaunchAttentionKernel( + const hipDeviceProp_t& prop, // Device Properties + hipStream_t stream, // rocm stream + const void* input, // Input tensor + const int* mask_index, // Attention mask raw data or index + // (end position of each sequence, + // or end positions and start positions). + // NULL means no mask. + gsl::span mask_index_dims, // Mask index shape + void* output, // Output tensor + int batch_size, // Batch size (B) + int sequence_length, // Sequence length (S) + int num_heads, // Number of attention heads (N) + int head_size, // Hidden layer size per head (H) + void* workspace, // Temporary buffer + rocblas_handle& rocblas, // Rocblas handle + const size_t element_size, // Element size of input tensor + bool is_unidirectional, // Whether there is unidirecitonal mask. + int past_sequence_length, // Sequence length in past state + const void* past, // Past state input + const void* extra_add_qk, // Additional Add + void* present // Present state output +); + +bool LaunchDecoderAttentionKernel( + const hipDeviceProp_t& prop, // Device Properties + hipStream_t stream, // Cuda stream + rocblas_handle& rocblas, // Rocblas handle + const size_t element_size, // Element size of input tensor + const int batch_size, // Batch size (B) + const int sequence_length, // Sequence length (S) + const int kv_sequence_length, // Key/Value/Cache sequence length + const int num_heads, // Number of attention heads (N) + const int head_size, // Hidden layer size per head (H) + const bool static_kv, // Whether cross attention or not + const bool use_past, // Whether use cache or not + const bool has_layer_state, // Whether output cache or not + const bool has_key_padding_mask, // Whether use key_padding_mask or not + const void* gemm_query_buffer, // Query buffer + const void* gemm_kv_buffer, // Key and value buffer + const bool* key_padding_mask, // Key padding mask + const void* key_cache, // Input key cache + const void* value_cache, // Input value cache + void* qkv_buffer, // Temporary buffer + void* workspace_buffer, // Temporary buffer + void* output, // Output tensor + void* new_key_cache, // New_key_cache tensor + void* new_value_cache // New_value_cache tensor +); + +bool LaunchTransCtx(hipStream_t stream, + const int sequence_length, const int batch_size, const int head_size, const int num_heads, + const int max_threads_per_block, const bool reversed_bs, const float* input, float* output); + +bool LaunchTransCtx(hipStream_t stream, + const int sequence_length, const int batch_size, const int head_size, const int num_heads, + const int max_threads_per_block, const bool reversed_bs, const half* input, half* output); + +bool LaunchTransQkv(hipStream_t stream, const int matrix_num, + const int sequence_length, const int batch_size, const int head_size, const int num_heads, + const int max_threads_per_block, const bool reversed_bs, const float* input, float* output); + +bool LaunchTransQkv(hipStream_t stream, const int matrix_num, + const int sequence_length, const int batch_size, const int head_size, const int num_heads, + const int max_threads_per_block, const bool reversed_bs, const half* input, half* output); + +bool LaunchConcatTensorToTensor(hipStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const int max_threads_per_block, + const int matrix_num, + const float* tensor_in, + const float* tensor_add, + float* tensor_out); + +bool LaunchConcatTensorToTensor(hipStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const int max_threads_per_block, + const int matrix_num, + const half* tensor_in, + const half* tensor_add, + half* tensor_out); + +bool LaunchConcatPastToPresent(hipStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const int max_threads_per_block, + const float* past, + const float* k_v, + float* present); + +bool LaunchConcatPastToPresent(hipStream_t stream, + const int all_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const int max_threads_per_block, + const half* past, + const half* k_v, + half* present); + +inline rocblas_status _compat_rocblas_gemm_strided_batched_ex(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, + int n, + int k, + const void* alpha, + const void* A, + rocblas_datatype a_type, + rocblas_int lda, + rocblas_stride stride_A, + const void* b, + rocblas_datatype b_type, + rocblas_int ldb, + rocblas_stride stride_b, + const void* beta, + void* c, + rocblas_datatype c_type, + rocblas_int ldc, + rocblas_stride stride_c, + rocblas_int batch_count, + rocblas_datatype compute_type, + rocblas_gemm_algo algo) { + return rocblas_gemm_strided_batched_ex(handle, + transa, + transb, + m, // m + n, // n + k, // k + alpha, // alpha + A, // A + a_type, // A type + lda, // lda + stride_A, // strideA + b, // B + b_type, // B type + ldb, // ldb + stride_b, // strideB + beta, // beta + c, // C + c_type, // C type + ldc, // ldc + stride_c, // strideC + c, // D = C + c_type, // D type = C type + ldc, // ldd = ldc + stride_c, // strideD = strideC + batch_count, // batch count + compute_type, + algo, + 0, 0); +} + +// Compatible for CublasMathModeSetter +class CompatRocblasMathModeSetter { + public: + CompatRocblasMathModeSetter(const hipDeviceProp_t&, + rocblas_handle, + int) { + } +}; + +} // namespace rocm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc index eecfb57477..2966d378bf 100644 --- a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc @@ -138,7 +138,8 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility @@ -151,8 +152,10 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, @@ -160,8 +163,10 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h b/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h index cf6971a697..b930f7a9bd 100644 --- a/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h @@ -68,33 +68,56 @@ inline rocblas_status rocblasGemmHelper(rocblas_handle handle, rocblas_gemm_algo_standard, 0, 0); } -inline rocblas_status rocblasGemmHelper(rocblas_handle handle, - rocblas_operation transa, - rocblas_operation transb, - int m, int n, int k, - const BFloat16* alpha, +inline rocblas_status rocblasGemmHelper(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, int n, int k, + const BFloat16* alpha, const BFloat16* A, int lda, - const BFloat16* B, int ldb, - const BFloat16* beta, + const BFloat16* B, int ldb, + const BFloat16* beta, BFloat16* C, int ldc) { float h_a = alpha->ToFloat(); float h_b = beta->ToFloat(); // accumulating in FP32 - return rocblas_gemm_ex(handle, - transa, + return rocblas_gemm_ex(handle, + transa, transb, - m, n, k, - &h_a, - A, rocblas_datatype_bf16_r, lda, - B, rocblas_datatype_bf16_r, ldb, - &h_b, + m, n, k, + &h_a, + A, rocblas_datatype_bf16_r, lda, + B, rocblas_datatype_bf16_r, ldb, + &h_b, + C, rocblas_datatype_bf16_r, ldc, C, rocblas_datatype_bf16_r, ldc, - C, rocblas_datatype_bf16_r, ldc, rocblas_datatype_f32_r, rocblas_gemm_algo_standard, 0, 0); } +// Compatible for function call with the extra hipDeviceProp_t argument +template +rocblas_status rocblasGemmHelper(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, int n, int k, + const Scalar* alpha, + const Scalar* A, int lda, + const Scalar* B, int ldb, + const Scalar* beta, + Scalar* C, int ldc, + const hipDeviceProp_t&) { + return rocblasGemmHelper(handle, + transa, + transb, + m, n, k, + alpha, + A, lda, + B, ldb, + beta, + C, ldc); +} + // batched gemm inline rocblas_status rocblasGemmBatchedHelper(rocblas_handle handle, rocblas_operation transa, @@ -159,31 +182,31 @@ inline rocblas_status rocblasGemmBatchedHelper(rocblas_handle handle, rocblas_gemm_algo_standard, 0, 0); } -inline rocblas_status rocblasGemmBatchedHelper(rocblas_handle handle, - rocblas_operation transa, +inline rocblas_status rocblasGemmBatchedHelper(rocblas_handle handle, + rocblas_operation transa, rocblas_operation transb, - int m, int n, int k, - const BFloat16* alpha, - const BFloat16* Aarray[], int lda, - const BFloat16* Barray[], int ldb, + int m, int n, int k, + const BFloat16* alpha, + const BFloat16* Aarray[], int lda, + const BFloat16* Barray[], int ldb, const BFloat16* beta, - BFloat16* Carray[], int ldc, + BFloat16* Carray[], int ldc, int batch_count) { float h_a = alpha->ToFloat(); float h_b = beta->ToFloat(); // accumulating in FP32 - return rocblas_gemm_batched_ex(handle, - transa, - transb, - m, n, k, - &h_a, + return rocblas_gemm_batched_ex(handle, + transa, + transb, + m, n, k, + &h_a, (const void**)Aarray, rocblas_datatype_bf16_r, lda, - (const void**)Barray, rocblas_datatype_bf16_r, ldb, - &h_b, + (const void**)Barray, rocblas_datatype_bf16_r, ldb, + &h_b, (void**)Carray, rocblas_datatype_bf16_r, ldc, (void**)Carray, rocblas_datatype_bf16_r, ldc, - batch_count, + batch_count, rocblas_datatype_f32_r, rocblas_gemm_algo_standard, 0, 0); } @@ -263,33 +286,61 @@ inline rocblas_status rocblasGemmStridedBatchedHelper(rocblas_handle handle, rocblas_gemm_algo_standard, 0, 0); } -inline rocblas_status rocblasGemmStridedBatchedHelper(rocblas_handle handle, +inline rocblas_status rocblasGemmStridedBatchedHelper(rocblas_handle handle, + rocblas_operation transa, + rocblas_operation transb, + int m, int n, int k, + const float* alpha, + const __half* A, int lda, + intmax_t strideA, + const __half* B, int ldb, + intmax_t strideB, + const float* beta, + __half* C, int ldc, + intmax_t strideC, + int batchCount) { + return rocblas_gemm_strided_batched_ex(handle, + transa, + transb, + m, n, k, + &alpha, + A, rocblas_datatype_f16_r, lda, strideA, + B, rocblas_datatype_f16_r, ldb, strideB, + &beta, + C, rocblas_datatype_f16_r, ldc, strideC, + C, rocblas_datatype_f16_r, ldc, strideC, + batchCount, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, 0, 0); +} + +inline rocblas_status rocblasGemmStridedBatchedHelper(rocblas_handle handle, rocblas_operation transa, - rocblas_operation transb, + rocblas_operation transb, int m, int n, int k, - const BFloat16* alpha, + const BFloat16* alpha, const BFloat16* A, int lda, - long long int strideA, + intmax_t strideA, const BFloat16* B, int ldb, - long long int strideB, - const BFloat16* beta, + intmax_t strideB, + const BFloat16* beta, BFloat16* C, int ldc, - long long int strideC, + intmax_t strideC, int batch_count) { float h_a = alpha->ToFloat(); float h_b = beta->ToFloat(); // accumulating in FP32 - return rocblas_gemm_strided_batched_ex(handle, - transa, - transb, - m, n, k, - &h_a, - A, rocblas_datatype_bf16_r, lda, strideA, - B, rocblas_datatype_bf16_r, ldb, strideB, - &h_b, + return rocblas_gemm_strided_batched_ex(handle, + transa, + transb, + m, n, k, + &h_a, + A, rocblas_datatype_bf16_r, lda, strideA, + B, rocblas_datatype_bf16_r, ldb, strideB, + &h_b, C, rocblas_datatype_bf16_r, ldc, strideC, C, rocblas_datatype_bf16_r, ldc, strideC, - batch_count, + batch_count, rocblas_datatype_f32_r, rocblas_gemm_algo_standard, 0, 0); } diff --git a/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc b/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc index 9b7710a1b5..f8d559216e 100644 --- a/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/decoder_attention_op_test.cc @@ -36,9 +36,10 @@ static void RunAttentionTest( ) { int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()); bool enable_cpu = false; - if (enable_cpu || enable_cuda) { + if (enable_cpu || enable_cuda || enable_rocm) { OpTester tester("DecoderAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); @@ -99,11 +100,14 @@ static void RunAttentionTest( tester.AddOutput("new_value_cache", output_cache_dims, *new_value_cache); } + std::vector> execution_providers; if (enable_cuda) { - std::vector> execution_providers; execution_providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } + if (enable_rocm) { + execution_providers.push_back(DefaultRocmExecutionProvider()); + } + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } } diff --git a/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc index 82f97eede0..704b4898c6 100644 --- a/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc @@ -29,6 +29,7 @@ static void RunAttentionTest( int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()); bool enable_cpu = false; if (enable_cpu || enable_cuda) { OpTester tester("LongformerAttention", 1, onnxruntime::kMSDomain); @@ -64,17 +65,18 @@ static void RunAttentionTest( tester.AddInput("global", global_dims, global_data); + std::vector> execution_providers; if (enable_cuda) { - std::vector> execution_providers; execution_providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + if (enable_rocm) { + execution_providers.push_back(DefaultRocmExecutionProvider()); + } + if (enable_cpu) { + execution_providers.push_back(DefaultCpuExecutionProvider()); } - if (enable_cpu) { - std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); - } + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } } diff --git a/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc b/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc index 8b062a09f6..5b7706df55 100644 --- a/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc +++ b/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc @@ -31,6 +31,13 @@ TEST(NGramRepeatBlockTest, NGramSize_3) { execution_providers.push_back(DefaultCudaExecutionProvider()); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +#ifdef USE_ROCM + if (nullptr != DefaultRocmExecutionProvider().get()) { + std::vector> execution_providers; + execution_providers.push_back(DefaultRocmExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +#endif std::vector> execution_providers; execution_providers.push_back(DefaultCpuExecutionProvider()); diff --git a/tools/ci_build/amd_hipify.py b/tools/ci_build/amd_hipify.py index db57509fbe..0ff9e221fe 100644 --- a/tools/ci_build/amd_hipify.py +++ b/tools/ci_build/amd_hipify.py @@ -19,8 +19,6 @@ contrib_ops_excluded_files = [ "bert/attention.cc", "bert/attention_impl.cu", "bert/attention_softmax.h", - "bert/decoder_attention.h", - "bert/decoder_attention.cc", "bert/embed_layer_norm.cc", "bert/embed_layer_norm.h", "bert/embed_layer_norm_impl.cu", @@ -29,20 +27,10 @@ contrib_ops_excluded_files = [ "bert/fast_gelu_impl.h", "bert/fast_gelu.cc", "bert/fast_gelu.h", - "bert/layer_norm.cuh", "bert/skip_layer_norm.cc", "bert/skip_layer_norm.h", "bert/skip_layer_norm_impl.cu", "bert/skip_layer_norm_impl.h", - "bert/longformer_attention.cc", - "bert/longformer_attention.h", - "bert/longformer_attention_softmax.cu", - "bert/longformer_attention_softmax.h", - "bert/longformer_attention_impl.cu", - "bert/longformer_attention_impl.h", - "bert/longformer_global_impl.cu", - "bert/longformer_global_impl.h", - "bert/transformer_cuda_common.h", "bert/transformer_common.h", "bert/transformer_common.cc", "math/complex_mul.cc", @@ -199,7 +187,9 @@ def hipify(src_file_path, dst_file_path): if not os.path.exists(dir_name): os.makedirs(dir_name, exist_ok=True) # Run hipify-perl first, capture output - s = subprocess.run([get_hipify_path(), src_file_path], stdout=subprocess.PIPE, universal_newlines=True).stdout + s = subprocess.run( + [get_hipify_path(), "-roc", src_file_path], stdout=subprocess.PIPE, universal_newlines=True, check=False + ).stdout # Additional exact-match replacements. # Order matters for all of the following replacements, reglardless of appearing in logical sections. @@ -272,6 +262,9 @@ def hipify(src_file_path, dst_file_path): s = s.replace('#include "cub/util_allocator.cuh"', "#include ") s = s.replace("#include ", "#include ") s = s.replace('#include "cub/util_type.cuh"', "#include ") + s = s.replace("#include ", "#include ") + s = s.replace("#include ", "#include ") + s = s.replace("#include ", "") # Doesn't exist s = s.replace("typedef half MappedType", "typedef __half MappedType") # CUBLAS -> HIPBLAS @@ -285,6 +278,19 @@ def hipify(src_file_path, dst_file_path): s = s.replace("Cublas", "Rocblas") s = s.replace("cublas", "rocblas") + # Undefined ROCMRT constants -> std::numeric_limits + s = s.replace("ROCMRT_INF_F", "std::numeric_limits::infinity()") + + # HIPBLAS -> rocblas + s = s.replace("HIPBLAS_R_16F", "rocblas_datatype_f16_r") + s = s.replace("HIPBLAS_R_32F", "rocblas_datatype_f32_r") + s = s.replace("ROCBLAS_GEMM_DEFAULT_TENSOR_OP", "rocblas_gemm_algo_standard") + s = s.replace("ROCBLAS_TENSOR_OP_MATH", "0 /* CUBLAS_TENSOR_OP_MATH is deprecated */") + + # compatible layer + s = s.replace("rocblas_gemm_strided_batched_ex", "_compat_rocblas_gemm_strided_batched_ex") + s = s.replace("RocblasMathModeSetter", "CompatRocblasMathModeSetter") + # CURAND -> HIPRAND s = s.replace("CURAND", "HIPRAND") s = s.replace("Curand", "Hiprand")