saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

exclude split3inner kernel on rocm ep (#21238)

Browse source 
### Description
There is an issue when using split3inner kernel on rocm-6.0.3, exclude
these code from rocm EP.
This commit is contained in:
kailums 2024-07-04 14:32:28 +08:00 committed by GitHub
parent 7d9b12a2e3
commit 40d4b2ec75
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 4 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/cuda/tensor/split.cc
View file

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

2
onnxruntime/core/providers/cuda/tensor/split_impl.cu
View file

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