use cublasHgemm for Volta GPU (#2074)

* use cublasHgemm for Volta GPU
This commit is contained in:
Yufeng Li 2019-10-14 17:29:13 -07:00 committed by GitHub
parent 8c5db7f973
commit 2536553136
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 68 additions and 66 deletions

View file

@ -135,8 +135,8 @@ __global__ void MaskedSoftmaxKernel(const int sequence_length, const int* mask_i
}
template <typename T>
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 <typename T>
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<float>(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<T>(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;
}

View file

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

View file

@ -7,6 +7,7 @@
#include <mutex>
#include <assert.h>
#include <cuda_runtime.h>
#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<cudaDeviceProp>& 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<cudaDeviceProp> s_cachedDeviceProps;
static std::once_flag s_cachedDevicePropsInitFlag;
};
#define CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N) \

View file

@ -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<T>(count);
if (cpu_pinned_copy_ == nullptr)
throw std::runtime_error("alloc failed");
@ -187,5 +187,52 @@ inline bool CalculateFdmStrides(gsl::span<fast_divmod> p, const std::vector<int6
return true;
}
struct DeviceProp {
static const std::vector<cudaDeviceProp>& 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<cudaDeviceProp> 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

View file

@ -10,9 +10,6 @@
namespace onnxruntime {
namespace cuda {
std::once_flag GridDim::s_cachedDevicePropsInitFlag;
std::vector<cudaDeviceProp> GridDim::s_cachedDeviceProps;
template <typename T>
__global__ void _Fill(
T* output_data,
@ -68,8 +65,8 @@ template std::unique_ptr<IConstantBuffer<float>> CreateConstantOnes<float>();
template std::unique_ptr<IConstantBuffer<double>> CreateConstantOnes<double>();
template std::unique_ptr<IConstantBuffer<half>> CreateConstantOnes<half>();
#define SPECIALIZED_FILL(T) \
template void Fill<T>(T* output, T value, int64_t count);
#define SPECIALIZED_FILL(T) \
template void Fill<T>(T * output, T value, int64_t count);
SPECIALIZED_FILL(int8_t)
SPECIALIZED_FILL(int16_t)

View file

@ -144,5 +144,8 @@ const float Consts<half>::Zero = 0;
const float Consts<half>::One = 1;
std::vector<cudaDeviceProp> DeviceProp::s_cachedDeviceProps;
std::once_flag DeviceProp::s_cachedDevicePropsInitFlag;
} // namespace cuda
} // namespace onnxruntime

View file

@ -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<const uint16_t*>(alpha));
float h_b = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(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);
}