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:
Changming Sun 2023-05-31 18:10:17 -07:00 committed by GitHub
parent f67f7c0f0b
commit d19e5c0abb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 26 additions and 9 deletions

View file

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

View file

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

View file

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