[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

@@ -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)

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"
 

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;
 

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;

onnxruntime/contrib_ops/cuda/bert/attention_impl.h

@@ -3,6 +3,7 @@
 
 #pragma once
 #include "core/providers/cuda/shared_inc/cuda_utils.h"
+#include <cuda_fp16.h>
 #include <cublas_v2.h>
 
 namespace onnxruntime {

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;
 

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"
 

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;

onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu

@@ -26,9 +26,9 @@ limitations under the License.
 #include <library_types.h>
 #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;

onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu

@@ -25,8 +25,8 @@ limitations under the License.
 #include <math_constants.h>
 #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;

onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 #include <cub/cub.cuh>
+#include <cub/device/device_partition.cuh>
 #include "core/providers/cuda/cuda_common.h"
 #include "longformer_global_impl.h"
 

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
+}  // namespace onnxruntime

onnxruntime/contrib_ops/rocm/bert/attention_impl.cu

@@ -208,6 +208,250 @@ bool LaunchAttentionKernel(
                         reinterpret_cast<float*>(present), use_persistent_softmax);
   }
 }
+
+
+template <typename T>
+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<float>(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<T>(stream, kv_sequence_length, sequence_length, batch_size,
+                                      num_heads, nullptr, key_padding_mask, nullptr, scratch1, scratch2,
+        false, 1, 2, static_cast<int>(0), false, nullptr)) {
+      return false;
+    }
+  } else {
+    if (!ComputeSoftmax<T>(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<const half*>(gemm_query_buffer),
+      reinterpret_cast<const half*>(gemm_kv_buffer),
+      key_padding_mask,
+      reinterpret_cast<const half*>(key_cache),
+      reinterpret_cast<const half*>(value_cache),
+      reinterpret_cast<half*>(qkv_buffer),
+      reinterpret_cast<half*>(workspace_buffer),
+      reinterpret_cast<half*>(output),
+      reinterpret_cast<half*>(new_key_cache),
+      reinterpret_cast<half*>(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<const float*>(gemm_query_buffer),
+      reinterpret_cast<const float*>(gemm_kv_buffer),
+      key_padding_mask,
+      reinterpret_cast<const float*>(key_cache),
+      reinterpret_cast<const float*>(value_cache),
+      reinterpret_cast<float*>(qkv_buffer),
+      reinterpret_cast<float*>(workspace_buffer),
+      reinterpret_cast<float*>(output),
+      reinterpret_cast<float*>(new_key_cache),
+      reinterpret_cast<float*>(new_value_cache));
+  }
+}
+
 }  // namespace rocm
 }  // namespace contrib
 }  // namespace onnxruntime

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 <hip/hip_fp16.h>
+#include <rocblas.h>
+#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<const int64_t> 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

onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc

@@ -138,7 +138,8 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) {
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, ComplexMulConj)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, NGramRepeatBlock)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
+                                                            1, NGramRepeatBlock)>,
 
       // 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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, Crop)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, Crop)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, Crop)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, DecoderAttention)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, DecoderAttention)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
+                                                                  1, float, DecoderAttention)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
+                                                                  1, MLFloat16, DecoderAttention)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, int64_t, DynamicSlice)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization)>,
@@ -160,8 +163,10 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) {
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, ImageScaler)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, ImageScaler)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, ImageScaler)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, LongformerAttention)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, LongformerAttention)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
+                                                                  1, float, LongformerAttention)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
+                                                                  1, MLFloat16, LongformerAttention)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, ParametricSoftplus)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, ParametricSoftplus)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, ParametricSoftplus)>,

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<typename Scalar>
+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);
 }

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<int64_t>("num_heads", static_cast<int64_t>(num_heads));
 
@@ -99,11 +100,14 @@ static void RunAttentionTest(
       tester.AddOutput<float>("new_value_cache", output_cache_dims, *new_value_cache);
     }
 
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
     if (enable_cuda) {
-      std::vector<std::unique_ptr<IExecutionProvider>> 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);
   }
 }
 

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<int32_t>("global", global_dims, global_data);
 
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
     if (enable_cuda) {
-      std::vector<std::unique_ptr<IExecutionProvider>> 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<std::unique_ptr<IExecutionProvider>> execution_providers;
-      execution_providers.push_back(DefaultCpuExecutionProvider());
-      tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
-    }
+    tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
   }
 }
 

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<std::unique_ptr<IExecutionProvider>> execution_providers;
+    execution_providers.push_back(DefaultRocmExecutionProvider());
+    tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+  }
+#endif
 
   std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
   execution_providers.push_back(DefaultCpuExecutionProvider());

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 <hipcub/util_allocator.hpp>")
     s = s.replace("#include <cub/util_type.cuh>", "#include <hipcub/backend/rocprim/util_type.hpp>")
     s = s.replace('#include "cub/util_type.cuh"', "#include <hipcub/backend/rocprim/util_type.hpp>")
+    s = s.replace("#include <cub/device/device_partition.cuh>", "#include <hipcub/device/device_partition.hpp>")
+    s = s.replace("#include <math_constants.h>", "#include <limits>")
+    s = s.replace("#include <library_types.h>", "")  # 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<float>::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")