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

Split topk implementation into per-type translation units to speed up compilation (#12861)

Browse source
This commit is contained in:
cloudhan 2022-09-14 19:36:54 +08:00 committed by GitHub
parent da07c83948
commit b8e34fbd91
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 80 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
onnxruntime/core/providers/cuda/math/topk_impl.cu → onnxruntime/core/providers/cuda/math/topk_impl.cuh
View file

@ -192,49 +192,49 @@ __global__ void BitonicTopK(const T* X, T* V, int64_t* I, const TArray<int64_t>
template <typename T>
__device__ __inline__ bool Equal(const T& t0, const T& t1) {
__device__ __forceinline__ bool Equal(const T& t0, const T& t1) {
return t0 == t1;
}
__device__ __inline__ bool Equal(const float& t0, const float& t1) {
__device__ __forceinline__ bool Equal(const float& t0, const float& t1) {
return !(t0 > t1 || t1 > t0);
}
__device__ __inline__ bool Equal(const double& t0, const double& t1) {
__device__ __forceinline__ bool Equal(const double& t0, const double& t1) {
return !(t0 > t1 || t1 > t0);
}
template<typename T>
__device__ bool SamePrefix(const T* t0, const T* t1, int64_t skip) {
__device__ __forceinline__ bool SamePrefix(const T* t0, const T* t1, int64_t skip) {
return ((*t0)^(*t1))>>skip == 0;
}
__device__ bool SamePrefix(const half* f0, const half* f1, int64_t skip) {
__device__ __forceinline__ bool SamePrefix(const half* f0, const half* f1, int64_t skip) {
return SamePrefix((const int16_t*)f0, (const int16_t*)f1, skip);
}
__device__ bool SamePrefix(const float* f0, const float* f1, int64_t skip) {
__device__ __forceinline__ bool SamePrefix(const float* f0, const float* f1, int64_t skip) {
return SamePrefix((const int32_t*)f0, (const int32_t*)f1, skip);
}
__device__ bool SamePrefix(const double* d0, const double* d1, int64_t skip) {
__device__ __forceinline__ bool SamePrefix(const double* d0, const double* d1, int64_t skip) {
return SamePrefix((const int64_t*)d0, (const int64_t*)d1, skip);
}
template<typename T>
__device__ int32_t Radix(const T* t, int64_t skip) {
__device__ __forceinline__ int32_t Radix(const T* t, int64_t skip) {
return ((*t)>>skip)&255;
}
__device__ int32_t Radix(const half* f, int64_t skip) {
__device__ __forceinline__ int32_t Radix(const half* f, int64_t skip) {
return Radix((const int16_t*)f, skip);
}
__device__ int32_t Radix(const float* f, int64_t skip) {
__device__ __forceinline__ int32_t Radix(const float* f, int64_t skip) {
return Radix((const int32_t*)f, skip);
}
__device__ int32_t Radix(const double* d, int64_t skip) {
__device__ __forceinline__ int32_t Radix(const double* d, int64_t skip) {
return Radix((const int64_t*)d, skip);
}
@ -243,15 +243,15 @@ __device__ void SetByte(T* t, int64_t byte) {
(*t) |= byte;
}
__device__ void SetByte(half* f, int64_t byte) {
__device__ __forceinline__ void SetByte(half* f, int64_t byte) {
SetByte((int16_t*)f, byte);
}
__device__ void SetByte(float* f, int64_t byte) {
__device__ __forceinline__ void SetByte(float* f, int64_t byte) {
SetByte((int32_t*)f, byte);
}
__device__ void SetByte(double* d, int64_t byte) {
__device__ __forceinline__ void SetByte(double* d, int64_t byte) {
SetByte((int64_t*)d, byte);
}
@ -300,13 +300,13 @@ __global__ void RadixTopK(const T* X, T* V, int64_t* I, const TArray<int64_t> el
if (KK > positive) {
KK = dimension - KK + 1;
sign = (T)-1;
}
}
} else {
if (KK > negative) {
KK = dimension - KK + 1;
} else {
sign = (T)-1;
}
}
}
__syncthreads();
#pragma unroll
@ -427,7 +427,9 @@ __global__ void FillOutput(const T* input_v, const int64_t* input_i, T* output_v
output_i[output_offset] = input_i[id];
}
__global__ void ExcludeOutput(int64_t* output_i, int64_t K, int64_t dimension) {
// template is used to avoid linking issue, since __global__ function cannot be inline-ed
template <typename T>
__global__ void ExcludeOutput(T* output_i, T K, T dimension) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, dimension);
if (id >= K) {
output_i[id] = dimension;
@ -472,16 +474,16 @@ Status TopKImpl(const CudaKernel* kernel, cudaStream_t stream, const T* input_x,
auto blocks_per_grid_K = (int)(ceil(static_cast<float>(K) / BT));
for (int64_t i = 0; i < N; i++) {
FillInput<CudaT><<<blocks_per_grid_D, BT, 0, stream>>>(input_x_ptr, input_key, input_value, elem_nums, size, axis, K, i, dimension);
CUDA_RETURN_IF_ERROR(1 == largest ? cub::DeviceRadixSort::SortPairsDescending(temp_storage, temp_bytes, input_key, output_key, input_value, output_value, dimension, 0, sizeof(T)*8, stream)
CUDA_RETURN_IF_ERROR(1 == largest ? cub::DeviceRadixSort::SortPairsDescending(temp_storage, temp_bytes, input_key, output_key, input_value, output_value, dimension, 0, sizeof(T)*8, stream)
: cub::DeviceRadixSort::SortPairs(temp_storage, temp_bytes, input_key, output_key, input_value, output_value, dimension, 0, sizeof(T)*8, stream));
if (1 == sorted) {
FillOutput<CudaT><<<blocks_per_grid_K, BT, 0, stream>>>(output_key, output_value, output_v_ptr, output_i, elem_nums, size, axis, K, i, dimension);
} else { //reorder by ascending index
ExcludeOutput<<<blocks_per_grid_D, BT, 0, stream>>>(output_value, K, dimension);
ExcludeOutput<int64_t><<<blocks_per_grid_D, BT, 0, stream>>>(output_value, K, dimension);
CUDA_RETURN_IF_ERROR(cub::DeviceRadixSort::SortPairs(temp_storage, temp_bytes, output_value, input_value, output_key, input_key, dimension, 0, sizeof(T)*8, stream));
FillOutput<CudaT><<<blocks_per_grid_K, BT, 0, stream>>>(input_key, input_value, output_v_ptr, output_i, elem_nums, size, axis, K, i, dimension);
}
}
}
}
return Status::OK();
}
@ -500,17 +502,9 @@ Status TopKImpl(const CudaKernel* kernel, cudaStream_t stream, const T* input_x,
int64_t N, \
int64_t dimension)
TOPKIMPLE(uint8_t);
TOPKIMPLE(uint16_t);
TOPKIMPLE(uint32_t);
TOPKIMPLE(uint64_t);
TOPKIMPLE(int8_t);
TOPKIMPLE(int16_t);
TOPKIMPLE(int32_t);
TOPKIMPLE(int64_t);
TOPKIMPLE(float);
TOPKIMPLE(MLFloat16);
TOPKIMPLE(double);
// This file is causing excessive long compilation time in ROCm EP. Split all those compilation into multiple
// translation units to speed it up.
TOPKIMPLE(TOPK_IMPL_TYPE);
} // namespace cuda
} // namespace onnxruntime

5
onnxruntime/core/providers/cuda/math/topk_impl_f16.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE MLFloat16
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_f32.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE float
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_f64.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE double
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_i16.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE int16_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_i32.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE int32_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_i64.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE int64_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_i8.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE int8_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_u16.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE uint16_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_u32.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE uint32_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_u64.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE uint64_t
#include "topk_impl.cuh"

5
onnxruntime/core/providers/cuda/math/topk_impl_u8.cu Normal file
View file

@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#define TOPK_IMPL_TYPE uint8_t
#include "topk_impl.cuh"