diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 85c39dbdf3..53c79598ab 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -135,8 +135,8 @@ __global__ void MaskedSoftmaxKernel(const int sequence_length, const int* mask_i } template -bool ComputeMaskedSoftmax(cudaStream_t stream, const int sequence_length, const int batch_size, const int num_heads, - const int* mask_index, const T* input, T* output) { +bool ComputeMaskedSoftmax(cudaStream_t stream, const int sequence_length, const int batch_size, const int num_heads, + const int* mask_index, const T* input, T* output) { // Mask is of length batch_size and assumes the valid region is contiguous starting // from the beginning of the sequence @@ -310,23 +310,6 @@ cublasStatus_t inline CublasGemmStridedBatched( handle, transa, transb, m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, batchCount); } -struct CublasConfigHelper { - cublasPointerMode_t pointer_mode_; - cublasMath_t math_mode_; - cublasHandle_t cublas_; - CublasConfigHelper(cublasHandle_t cublas) - : cublas_(cublas) { - cublasGetPointerMode(cublas_, &pointer_mode_); - cublasGetMathMode(cublas_, &math_mode_); - cublasSetPointerMode(cublas_, CUBLAS_POINTER_MODE_HOST); - cublasSetMathMode(cublas_, CUBLAS_TENSOR_OP_MATH); - } - ~CublasConfigHelper() { - cublasSetMathMode(cublas_, math_mode_); - cublasSetPointerMode(cublas_, pointer_mode_); - } -}; - template bool QkvToContext( cublasHandle_t& cublas, cudaStream_t stream, @@ -337,7 +320,7 @@ bool QkvToContext( T* scratch1 = workspace; T* scratch2 = scratch1 + (bytes / element_size); T* scratch3 = scratch2 + (bytes / element_size); - + // input should be BxSx3xNxH => scratch3: 3xBxNxSxH if (!LaunchTransQkv(stream, sequence_length, batch_size, head_size, num_heads, input, scratch3)) { return false; @@ -354,17 +337,16 @@ bool QkvToContext( const T* v = k + total_size; cublasSetStream(cublas, stream); - CublasConfigHelper helper(cublas); + CublasMathModeSetter helper(cublas, CUBLAS_TENSOR_OP_MATH); // compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); if (!CUBLAS_CALL(CublasGemmStridedBatched( - cublas, CUBLAS_OP_T, CUBLAS_OP_N, sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, size_per_batch, - q, head_size, size_per_batch, 0.f, scratch1, sequence_length, temp_matrix_size, batches))) { + cublas, CUBLAS_OP_T, CUBLAS_OP_N, sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, size_per_batch, + q, head_size, size_per_batch, 0.f, scratch1, sequence_length, temp_matrix_size, batches))) { return false; } - // apply softmax and store result P to scratch2: BxNxSxS if (!ComputeMaskedSoftmax(stream, sequence_length, batch_size, num_heads, mask_index, scratch1, scratch2)) { return false; @@ -372,8 +354,8 @@ bool QkvToContext( // compute P*V (as V*P), and store in scratch3: BxNxSxH if (!CUBLAS_CALL(CublasGemmStridedBatched( - cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, sequence_length, 1.f, v, head_size, size_per_batch, - scratch2, sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) { + cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, sequence_length, 1.f, v, head_size, size_per_batch, + scratch2, sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) { return false; } diff --git a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu index dfa6dfbcbd..50fb6b3aa9 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu @@ -383,7 +383,7 @@ void HostApplyLayerNorm( const T* gamma, const T* beta) { const dim3 threads(32, 4, 1); - const cudaDeviceProp& prop = GridDim::GetDeviceProps(); + const cudaDeviceProp& prop = DeviceProp::GetDeviceProps(); const uint64_t maxGridY = prop.maxGridSize[1]; const int warp_size = prop.warpSize; // const uint64_t maxGridY = 32; diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index c58113c67d..1062be123f 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -7,6 +7,7 @@ #include #include #include +#include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/cuda_call.h" namespace onnxruntime { @@ -224,7 +225,7 @@ struct GridDim { N = 1; // get device information - const auto& props = GetDeviceProps(); + const auto& props = DeviceProp::GetDeviceProps(); CUDA_LONG numProcs = props.multiProcessorCount; CUDA_LONG warpSize = props.warpSize; @@ -246,40 +247,10 @@ struct GridDim { assert(blocks_per_grid_ * threads_per_block_ >= N); } - static const std::vector& GetCachedDeviceProps() { - std::call_once(s_cachedDevicePropsInitFlag, [=] { - int numDevices; - // must wait GPU idle, otherwise cudaGetDeviceProperties might fail - CUDA_CALL_THROW(cudaDeviceSynchronize()); - CUDA_CALL_THROW(cudaGetDeviceCount(&numDevices)); - s_cachedDeviceProps.resize(numDevices); - for (int i = 0; i < numDevices; i++) - CUDA_CALL_THROW(cudaGetDeviceProperties(&s_cachedDeviceProps[i], i)); - }); - - return s_cachedDeviceProps; - } - - static size_t GetCurrentDeviceId() { - int deviceId; - cudaGetDevice(&deviceId); - return (size_t)deviceId; - } - - // get device properties of current device - static const cudaDeviceProp& GetDeviceProps() { - const auto& cachedDevicesProps = GetCachedDeviceProps(); - return cachedDevicesProps[GetCurrentDeviceId()]; - } - // compute our location on the grid static __device__ CUDA_LONG GetLinearThreadId() { return blockDim.x * blockIdx.x + threadIdx.x; } - - private: - static std::vector s_cachedDeviceProps; - static std::once_flag s_cachedDevicePropsInitFlag; }; #define CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N) \ diff --git a/onnxruntime/core/providers/cuda/cuda_common.h b/onnxruntime/core/providers/cuda/cuda_common.h index 8c467980a3..8807ede3c4 100644 --- a/onnxruntime/core/providers/cuda/cuda_common.h +++ b/onnxruntime/core/providers/cuda/cuda_common.h @@ -90,7 +90,7 @@ class CudaKernel : public OpKernel { memcpy(CpuPtr(), vec.data(), vec.size() * sizeof(T)); } - void AllocCpuPtr( size_t count) { + void AllocCpuPtr(size_t count) { cpu_pinned_copy_ = op_kernel_->AllocateBufferOnCPUPinned(count); if (cpu_pinned_copy_ == nullptr) throw std::runtime_error("alloc failed"); @@ -187,5 +187,52 @@ inline bool CalculateFdmStrides(gsl::span p, const std::vector& GetCachedDeviceProps() { + std::call_once(s_cachedDevicePropsInitFlag, [=] { + int numDevices; + // must wait GPU idle, otherwise cudaGetDeviceProperties might fail + CUDA_CALL_THROW(cudaDeviceSynchronize()); + CUDA_CALL_THROW(cudaGetDeviceCount(&numDevices)); + s_cachedDeviceProps.resize(numDevices); + for (int i = 0; i < numDevices; i++) + CUDA_CALL_THROW(cudaGetDeviceProperties(&s_cachedDeviceProps[i], i)); + }); + + return s_cachedDeviceProps; + } + + static size_t GetCurrentDeviceId() { + int deviceId; + cudaGetDevice(&deviceId); + return (size_t)deviceId; + } + + // get device properties of current device + static const cudaDeviceProp& GetDeviceProps() { + const auto& cachedDevicesProps = GetCachedDeviceProps(); + return cachedDevicesProps[GetCurrentDeviceId()]; + } + + private: + static std::vector s_cachedDeviceProps; + static std::once_flag s_cachedDevicePropsInitFlag; +}; + +class CublasMathModeSetter { + public: + CublasMathModeSetter(cublasHandle_t handle, cublasMath_t mode) : handle_(handle) { + cublasGetMathMode(handle, &mode_); + cublasSetMathMode(handle, mode); + } + ~CublasMathModeSetter() { + cublasSetMathMode(handle_, mode_); + } + + private: + cublasHandle_t handle_; + cublasMath_t mode_; +}; + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_utils.cu b/onnxruntime/core/providers/cuda/cuda_utils.cu index 5dc61cfa18..2a6f17b609 100644 --- a/onnxruntime/core/providers/cuda/cuda_utils.cu +++ b/onnxruntime/core/providers/cuda/cuda_utils.cu @@ -10,9 +10,6 @@ namespace onnxruntime { namespace cuda { -std::once_flag GridDim::s_cachedDevicePropsInitFlag; -std::vector GridDim::s_cachedDeviceProps; - template __global__ void _Fill( T* output_data, @@ -68,8 +65,8 @@ template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); -#define SPECIALIZED_FILL(T) \ -template void Fill(T* output, T value, int64_t count); +#define SPECIALIZED_FILL(T) \ + template void Fill(T * output, T value, int64_t count); SPECIALIZED_FILL(int8_t) SPECIALIZED_FILL(int16_t) diff --git a/onnxruntime/core/providers/cuda/cudnn_common.cc b/onnxruntime/core/providers/cuda/cudnn_common.cc index 135b61999c..089313a383 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.cc +++ b/onnxruntime/core/providers/cuda/cudnn_common.cc @@ -144,5 +144,8 @@ const float Consts::Zero = 0; const float Consts::One = 1; +std::vector DeviceProp::s_cachedDeviceProps; +std::once_flag DeviceProp::s_cachedDevicePropsInitFlag; + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index bb0d7ba56f..d6ede7cf76 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -26,11 +26,13 @@ inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t } inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const half* alpha, const half* A, int lda, const half* B, int ldb, const half* beta, half* C, int ldc) { // This does true FP16 computation which is slow for non-Volta GPUs - //return cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc); + if (onnxruntime::cuda::DeviceProp().GetDeviceProps().major >= 7) { + onnxruntime::cuda::CublasMathModeSetter math_mode_setter( handle, CUBLAS_TENSOR_OP_MATH ); + return cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc); + } // This does pseudo FP16 computation (input/output in fp16, computation in fp32) float h_a = onnxruntime::math::halfToFloat(*reinterpret_cast(alpha)); float h_b = onnxruntime::math::halfToFloat(*reinterpret_cast(beta)); - cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH); return cublasGemmEx(handle, transa, transb, m, n, k, &h_a, A, CUDA_R_16F, lda, B, CUDA_R_16F, ldb, &h_b, C, CUDA_R_16F, ldc, CUDA_R_32F, CUBLAS_GEMM_DFALT); }