mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
parent
8c5db7f973
commit
2536553136
7 changed files with 68 additions and 66 deletions
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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) \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue