mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-24 19:43:35 +00:00
Fix a misaligned error in CUDA GEMM (#16130)
### Description Fix an issue that FusedMatMulOpTest.FloatTypeTransposeBatch fails to run on GPUs with TF32 support. Authored-by: Tianlei Wu <tlwu@microsoft.com>
This commit is contained in:
parent
f67f7c0f0b
commit
d19e5c0abb
3 changed files with 26 additions and 9 deletions
|
|
@ -371,9 +371,24 @@ class MatMulComputeHelper {
|
|||
return right_zp_offsets_;
|
||||
}
|
||||
|
||||
static bool IsAligned(const std::vector<size_t>& offsets) {
|
||||
constexpr size_t alignment = 16;
|
||||
const auto len = offsets.size();
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
if ((offsets[i] % alignment) != 0) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool IsBatchedGemmAligned() const {
|
||||
return IsAligned(left_offsets_) && IsAligned(right_offsets_) && IsAligned(output_offsets_);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void OffsetToArrays(T* p, const std::vector<size_t>& offsets, gsl::span<T*> arrays) {
|
||||
auto len = offsets.size();
|
||||
const auto len = offsets.size();
|
||||
ORT_ENFORCE(arrays.size() == len);
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
arrays[i] = p + offsets[i];
|
||||
|
|
@ -382,7 +397,7 @@ class MatMulComputeHelper {
|
|||
|
||||
template <typename T>
|
||||
static void OffsetToArrays(const T* p, const std::vector<size_t>& offsets, gsl::span<const T*> arrays) {
|
||||
auto len = offsets.size();
|
||||
const auto len = offsets.size();
|
||||
ORT_ENFORCE(arrays.size() == len);
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
arrays[i] = p + offsets[i];
|
||||
|
|
|
|||
|
|
@ -180,6 +180,14 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
|
|||
ORT_RETURN_IF_ERROR(right_arrays.CopyToGpu(ctx->GetComputeStream()));
|
||||
ORT_RETURN_IF_ERROR(output_arrays.CopyToGpu(ctx->GetComputeStream()));
|
||||
|
||||
// TF32 provides a huge performance gain for training and inference while preserving FP32 levels of accuracy.
|
||||
// It requires Ampere or newer GPU, and pointers of matrics shall be aligned (ideal alignment is 16-byte).
|
||||
// Assume that start memory of input/output tensor is aligned, we only check offsets of sub-matrix per batch here.
|
||||
cublasMath_t mode = (std::is_same<T, float>::value && device_prop.major >= 8 && helper.IsBatchedGemmAligned())
|
||||
? CUBLAS_TF32_TENSOR_OP_MATH
|
||||
: CUBLAS_DEFAULT_MATH;
|
||||
CublasMathModeSetter math_mode_setter(device_prop, GetCublasHandle(ctx), mode);
|
||||
|
||||
// note that onnxruntime OrtValue is row major, while cublas is column major,
|
||||
// so swap left/right operands
|
||||
CUBLAS_RETURN_IF_ERROR(cublasGemmBatchedHelper(
|
||||
|
|
|
|||
|
|
@ -185,13 +185,7 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle,
|
|||
const float* beta,
|
||||
float* Carray[], int ldc,
|
||||
int batch_count,
|
||||
const cudaDeviceProp& prop) {
|
||||
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
|
||||
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TF32_TENSOR_OP_MATH);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(prop);
|
||||
#endif
|
||||
|
||||
const cudaDeviceProp&) {
|
||||
return cublasSgemmBatched(handle,
|
||||
transa,
|
||||
transb,
|
||||
|
|
|
|||
Loading…
Reference in a new issue