saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-04 04:07:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] UTs and code clean up (#10511)

Browse source 
* Fix UT

* UT

* UTs

* enable ROCm UT

* fix build attempt

* minor

* fix UT

* fix UT

* fix UTs

Co-authored-by: Ethan Tao <ettao@microsoft.com@orttrainingdev7.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
Co-authored-by: root <root@GCRAMDRR1-MI100-087.redmond.corp.microsoft.com>
This commit is contained in:
ytaous 2022-02-11 08:23:25 -08:00 committed by GitHub
parent 2002a96594
commit 4e2a974090
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 103 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu
View file

@ -92,45 +92,17 @@ bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int i
return CUDA_CALL(cudaPeekAtLastError());
}
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template <unsigned TPB>
__global__ void FastGeluKernel2(const nv_bfloat162 a, const nv_bfloat162 b, const nv_bfloat162 c, int input_length,
int bias_length, const nv_bfloat162* input, const nv_bfloat162* bias,
nv_bfloat162* output) {
const int idx = blockIdx.x * TPB + threadIdx.x;
if (idx < input_length) {
const nv_bfloat162 x = input[idx];
const nv_bfloat162 in = (bias == nullptr) ? x : (x + bias[idx % bias_length]);
const nv_bfloat162 cdf = a + a * _Tanh(in * (c * in * in + b));
output[idx] = in * cdf;
}
}
#endif
template <>
bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int input_length, int bias_length,
const BFloat16* input, const BFloat16* bias, BFloat16* output, bool /*use_half2*/) {
constexpr int blockSize = 256;
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
if (0 == (bias_length & 1) && prop.major >= 7) {
const int n = input_length / 2;
const int gridSize = (n + blockSize - 1) / blockSize;
const nv_bfloat162 A2 = __floats2bfloat162_rn(A, A);
const nv_bfloat162 B2 = __floats2bfloat162_rn(B, B);
const nv_bfloat162 C2 = __floats2bfloat162_rn(C, C);
const nv_bfloat162* input2 = reinterpret_cast<const nv_bfloat162*>(input);
const nv_bfloat162* bias2 = reinterpret_cast<const nv_bfloat162*>(bias);
nv_bfloat162* output2 = reinterpret_cast<nv_bfloat162*>(output);
FastGeluKernel2<blockSize>
<<<gridSize, blockSize, 0, stream>>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2);
} else {
#endif
const int gridSize = (input_length + blockSize - 1) / blockSize;
FastGeluKernel<BFloat16, blockSize>
<<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
}
#endif
// remove nv_bfloat162 implementation for now to fix build issue
// we can decide whether to add it back if there's perf concern
const int gridSize = (input_length + blockSize - 1) / blockSize;
FastGeluKernel<BFloat16, blockSize>
<<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
return CUDA_CALL(cudaPeekAtLastError());
}

10
onnxruntime/core/providers/cuda/cu_inc/common.cuh
View file

@ -351,16 +351,6 @@ __device__ __inline__ BFloat16 _Log(BFloat16 a) { return logf(static_cast<float>
template <>
__device__ __inline__ BFloat16 _Tanh(BFloat16 a) { return tanhf(static_cast<float>(a)); }
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template <>
__device__ __inline__ nv_bfloat162 _Tanh(nv_bfloat162 a) {
float2 tmp = (__bfloat1622float2(a));
tmp.x = tanhf(tmp.x);
tmp.y = tanhf(tmp.y);
return __float22bfloat162_rn(tmp);
}
#endif
template <>
__device__ __inline__ BFloat16 _Normcdf(BFloat16 a) { return normcdff(static_cast<float>(a)); }

25
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
View file

@ -513,7 +513,9 @@ Status ReduceComputeCore(CUDAExecutionProvider& cuda_ep, const Tensor& input, Pr
IAllocatorUniquePtr<float> temp_X;
cudnnDataType_t cudnn_type_X = CudnnTensor::GetDataType<CudaT>();
if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same<T, MLFloat16>::value) {
// Reducesum with BFP16 is not supported by cudnn, so convert input to fp32 then call cudnn
if ((ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same<T, MLFloat16>::value) ||
(ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES && std::is_same<T, BFloat16>::value)) {
// ArgMax/ArgMin with FP16 are not supported by cudnn, so convert input to fp32 then call cudnn
temp_X = cuda_ep.GetScratchBuffer<float>(input_count);
cudnn_type_X = CUDNN_DATA_FLOAT;
@ -652,11 +654,22 @@ Status ReduceComputeCore(CUDAExecutionProvider& cuda_ep, const Tensor& input, Pr
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output.template MutableData<T>(), input.template Data<T>(), input_count * sizeof(T), cudaMemcpyDeviceToDevice, stream));
}
} else {
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
cuda_ep.PerThreadCudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes,
workspace_cuda.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const CudaT*>(input.template Data<T>()),
&zero, output_tensor, reinterpret_cast<CudaT*>(output.template MutableData<T>())));
if (temp_X) {
auto temp_output = cuda_ep.GetScratchBuffer<float>(output_count);
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
cuda_ep.PerThreadCudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes,
workspace_cuda.get(), workspace_bytes,
&one, input_tensor, temp_X.get(),
&zero, output_tensor, temp_output.get()));
Impl_Cast<float, CudaT>(stream, temp_output.get(), reinterpret_cast<CudaT*>(output.template MutableData<T>()), output_count);
} else {
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
cuda_ep.PerThreadCudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes,
workspace_cuda.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const CudaT*>(input.template Data<T>()),
&zero, output_tensor, reinterpret_cast<CudaT*>(output.template MutableData<T>())));
}
}
}
} else {

26
onnxruntime/core/providers/rocm/reduction/reduction_ops.cc
View file

@ -511,7 +511,10 @@ Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, Pr
IAllocatorUniquePtr<float> temp_X;
miopenDataType_t miopen_type_X = MiopenTensor::GetDataType<HipT>();
if (ReduceTensorIndices == MIOPEN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same<T, MLFloat16>::value) {
// unlike bfp16 not supported in cudnn, miopen call for bfp16 succeeded below, however, UT shows data error
// so for now, follow the same logic in cudnn and convert input to fp32 then call miopen
if ((ReduceTensorIndices == MIOPEN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same<T, MLFloat16>::value) ||
(ReduceTensorIndices == MIOPEN_REDUCE_TENSOR_NO_INDICES && std::is_same<T, BFloat16>::value)) {
// ArgMax/ArgMin with FP16 are not supported by miopen, so convert input to fp32 then call miopen
temp_X = rocm_ep.GetScratchBuffer<float>(input_count);
miopen_type_X = miopenFloat;
@ -651,11 +654,22 @@ Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, Pr
HIP_RETURN_IF_ERROR(hipMemcpyAsync(output.template MutableData<T>(), input.template Data<T>(), input_count * sizeof(T), hipMemcpyDeviceToDevice, stream));
}
} else {
MIOPEN_RETURN_IF_ERROR(miopenReduceTensor(
rocm_ep.PerThreadMiopenHandle(), reduce_desc, indices_rocm.get(), indices_bytes,
workspace_rocm.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const HipT*>(input.template Data<T>()),
&zero, output_tensor, reinterpret_cast<HipT*>(output.template MutableData<T>())));
if (temp_X) {
auto temp_output = rocm_ep.GetScratchBuffer<float>(output_count);
MIOPEN_RETURN_IF_ERROR(miopenReduceTensor(
rocm_ep.PerThreadMiopenHandle(), reduce_desc, indices_rocm.get(), indices_bytes,
workspace_rocm.get(), workspace_bytes,
&one, input_tensor, temp_X.get(),
&zero, output_tensor, temp_output.get()));
Impl_Cast<float, HipT>(stream, temp_output.get(), reinterpret_cast<HipT*>(output.template MutableData<T>()), output_count);
} else {
MIOPEN_RETURN_IF_ERROR(miopenReduceTensor(
rocm_ep.PerThreadMiopenHandle(), reduce_desc, indices_rocm.get(), indices_bytes,
workspace_rocm.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const HipT*>(input.template Data<T>()),
&zero, output_tensor, reinterpret_cast<HipT*>(output.template MutableData<T>())));
}
}
}
} else {

3
onnxruntime/test/contrib_ops/element_wise_ops_test.cc
View file

@ -146,9 +146,8 @@ TEST(BiasGeluTest, Two_One_Dim_fp16) {
}
#endif
// failed test for CUDA (therefore ROCM as well) to be investigated
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(BiasGeluTest, DISABLED_Two_One_Dim_bfloat16) {
TEST(BiasGeluTest, Two_One_Dim_bfloat16) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {

4
onnxruntime/test/contrib_ops/fastgelu_op_test.cc
View file

@ -197,11 +197,9 @@ TEST(FastGeluTest, FastGeluWithoutBiasFloat16) {
RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, false, true);
}
// failed with device error, disabled for now
// CUDA only, ROCM has not been supported yet
#ifdef USE_CUDA
TEST(FastGeluTest, DISABLED_FastGeluWithBias_BFloat16) {
TEST(FastGeluTest, FastGeluWithBias_BFloat16) {
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {
LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16";

46
onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
View file

@ -1450,6 +1450,31 @@ TEST(ReductionOpTest, ReduceSumHalfHalf) {
test.Run();
}
TEST(ReductionOpTest, ReduceSumHalfHalf_2) {
OpTester test("ReduceSum");
test.AddAttribute("keepdims", (int64_t)0);
test.AddAttribute("axes", std::vector<int64_t>{0, 2});
std::vector<float> data = {1.0f, 2.0f,
3.0f, 4.0f,
5.0f, 6.0f,
7.0f, 8.0f,
9.0f, 10.0f,
11.0f, 12.0f};
std::vector<MLFloat16> data_half(12);
ConvertFloatToMLFloat16(data.data(), data_half.data(), 12);
std::vector<float> result = {33.0f, 45.0f};
std::vector<MLFloat16> result_half(2);
ConvertFloatToMLFloat16(result.data(), result_half.data(), 2);
test.AddInput<MLFloat16>("data", {3, 2, 2}, data_half);
test.AddOutput<MLFloat16>("reduced", {2}, result_half);
test.Run();
}
void test_half_reduce_sum(
int64_t m, int64_t n) {
OpTester test("ReduceSum");
@ -1509,6 +1534,27 @@ TEST(ReductionOpTest, ReduceSumBFloat16) {
}
#endif
// on CUDA - this UT, with axes {0,2}, will go thru cudnn lib only if ATenOp is not initialized
// on ROCM - miopen call succeeded, but results in data error, thus follow the same logic done in cudnn for now
// TODO - try ROCm 4.5.2 and/or double check the source code on BFloat16 support
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(ReductionOpTest, ReduceSumBFloat16_2) {
OpTester test("ReduceSum", 14);
test.AddAttribute("keepdims", (int64_t)0);
test.AddInput<BFloat16>("data", {3, 2, 2},
MakeBFloat16({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f}));
test.AddInput<int64_t>("axes", {2}, std::vector<int64_t>{0, 2});
test.AddOutput<BFloat16>("reduced", {2}, MakeBFloat16({33.0f, 45.0f}));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
#ifdef USE_CUDA
execution_providers.push_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.push_back(DefaultRocmExecutionProvider());
#endif
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
#endif
TEST(ReductionOpTest, ReduceSum_apex_reduction) {
OpTester test("ReduceSum");
test.AddAttribute("keepdims", (int64_t)0);

2
onnxruntime/test/providers/provider_test_utils.cc
View file

@ -330,7 +330,7 @@ struct TensorCheck<BFloat16> {
/// XXX: May need to adjust threshold as BFloat is coarse
float threshold = 0.001f;
#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING) || defined(USE_CUDA) || defined(USE_ROCM)
threshold = 0.008f;
threshold = 0.05f; // expect at least 95% close
#endif
for (int i = 0; i < size; ++i) {
if (std::isnan(f_expected[i])) {

17
orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc
View file

@ -22,13 +22,13 @@ struct MixedPrecisionScaleInputOutput {
input1_bf16.resize(input1.size());
output1_bf16.resize(output1.size());
std::vector<BFloat16> input1_bf16 = FloatsToBFloat16s(input1);
std::vector<BFloat16> output1_bf16 = FloatsToBFloat16s(output1);
input1_bf16 = FloatsToBFloat16s(input1);
output1_bf16 = FloatsToBFloat16s(output1);
input2_bf16.resize(input2.size());
output2_bf16.resize(output2.size());
std::vector<BFloat16> input2_bf16 = FloatsToBFloat16s(input2);
std::vector<BFloat16> output2_bf16 = FloatsToBFloat16s(output2);
input2_bf16 = FloatsToBFloat16s(input2);
output2_bf16 = FloatsToBFloat16s(output2);
}
// Fp32 Inputs/Output
@ -172,8 +172,7 @@ TEST(CudaKernelTest, MixedPrecisionScale_bfloat16_bfloat16) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
// failed with data error, disabled for now
TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_float_bfloat16) {
TEST(CudaKernelTest, MixedPrecisionScale_float_bfloat16) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {
@ -197,7 +196,7 @@ TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_float_bfloat16) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_bfloat16_float) {
TEST(CudaKernelTest, MixedPrecisionScale_bfloat16_float) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {
@ -221,7 +220,7 @@ TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_bfloat16_float) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_half_bfloat16) {
TEST(CudaKernelTest, MixedPrecisionScale_half_bfloat16) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {
@ -245,7 +244,7 @@ TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_half_bfloat16) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_bfloat16_half) {
TEST(CudaKernelTest, MixedPrecisionScale_bfloat16_half) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!HasCudaEnvironment(min_cuda_architecture)) {

21
orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc
View file

@ -62,27 +62,6 @@ Status SoftMaxGradComputeHelper(
return Status::OK();
}
#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \
template <> \
Status SoftMaxGradComputeHelper<BFloat16, is_log_softmax>(hipStream_t stream, const BFloat16* dY, \
const TensorShape& input_shape, const BFloat16* Y, \
BFloat16* dX, miopenHandle_t, int64_t axis) { \
typedef typename ToHipType<BFloat16>::MappedType HipT; \
const int64_t normalized_axis = HandleNegativeAxis(axis, input_shape.NumDimensions()); \
int64_t N = input_shape.SizeToDimension(normalized_axis); \
int64_t D = input_shape.SizeFromDimension(normalized_axis); \
auto dY_data = reinterpret_cast<const HipT*>(dY); \
auto Y_data = reinterpret_cast<const HipT*>(Y); \
auto dX_data = reinterpret_cast<HipT*>(dX); \
dispatch_softmax_backward<HipT, HipT, AccumulationType_t<HipT>, is_log_softmax>( \
stream, dX_data, dY_data, Y_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), \
gsl::narrow_cast<int>(N)); \
return Status::OK(); \
}
SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true)
SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false)
#define REGISTER_GRADIENT_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
SoftmaxGrad, \