From 40d4b2ec75a17d48e16d2794d3d6ced39333a265 Mon Sep 17 00:00:00 2001
From: kailums <109063327+kailums@users.noreply.github.com>
Date: Thu, 4 Jul 2024 14:32:28 +0800
Subject: [PATCH] exclude split3inner kernel on rocm ep (#21238)

### Description
There is an issue when using split3inner kernel on rocm-6.0.3, exclude
these code from rocm EP.
---
 onnxruntime/core/providers/cuda/tensor/split.cc      | 2 ++
 onnxruntime/core/providers/cuda/tensor/split_impl.cu | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc
index ca82387600..52775b2e8b 100644
--- a/onnxruntime/core/providers/cuda/tensor/split.cc
+++ b/onnxruntime/core/providers/cuda/tensor/split.cc
@@ -76,6 +76,7 @@ Status SplitKernel::ComputeInternal(OpKernelContext* ctx) const {
   auto input_dims = input_shape.GetDims();
   auto output_dimensions{input_shape.AsShapeVector()};
 
+#ifndef USE_ROCM
   if (split_sizes.size() == 3 && ((axis + 1) == gsl::narrow_cast<int64_t>(input_shape.NumDimensions()))) {
     // we use (axis + 1) == num_dimensions to check if we are splitting on inner most axis.
     // only when split on inner axis and output size is 3, we can use Split3Inner.
@@ -100,6 +101,7 @@ Status SplitKernel::ComputeInternal(OpKernelContext* ctx) const {
                        output2->MutableDataRaw(),
                        input_dims);
   }
+#endif
 
   CudaAsyncBuffer<void*> output_ptr(this, num_outputs);
   gsl::span<void*> output_ptr_span = output_ptr.CpuSpan();
diff --git a/onnxruntime/core/providers/cuda/tensor/split_impl.cu b/onnxruntime/core/providers/cuda/tensor/split_impl.cu
index 00f94694f8..e2f42e4d58 100644
--- a/onnxruntime/core/providers/cuda/tensor/split_impl.cu
+++ b/onnxruntime/core/providers/cuda/tensor/split_impl.cu
@@ -157,6 +157,7 @@ Status SplitImpl(cudaStream_t stream, const size_t element_size, const int block
   return Status::OK();
 }
 
+#ifndef USE_ROCM
 template <typename T>
 __global__ void _Split3InnerKernel(const int64_t size0_in_byte,
                                    const int64_t size1_in_byte,
@@ -263,6 +264,7 @@ Status Split3Inner(cudaStream_t stream, const size_t element_size, const int64_t
 
   return Status::OK();
 }
+#endif
 
 }  // namespace cuda
 }  // namespace onnxruntime