From 44e2db939711cf1c6108bee3088f69aef38553e4 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Fri, 14 Jan 2022 19:38:56 +0800 Subject: [PATCH] CUDA BFloat16 Refactor (#10085) --- docs/ContribOperators.md | 2 +- docs/OperatorKernels.md | 2 +- include/onnxruntime/core/framework/float16.h | 87 ++++---- .../contrib_ops/cuda/bert/fast_gelu.cc | 2 - .../contrib_ops/cuda/bert/fast_gelu_impl.cu | 24 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 8 +- onnxruntime/contrib_ops/cuda/layer_norm.cc | 2 - .../contrib_ops/cuda/layer_norm_impl.cu | 6 +- .../contrib_ops/cuda/math/bias_dropout.cc | 8 +- .../cuda/math/bias_dropout_impl.cu | 4 +- .../cuda/math/binary_elementwise_ops.cc | 3 +- .../cuda/math/binary_elementwise_ops_impl.cu | 3 +- .../contrib_ops/cuda/math/fused_matmul.cc | 5 +- .../contrib_ops/cuda/math/isfinite.cuh | 18 +- .../core/framework/tensorprotoutils.cc | 2 +- .../core/graph/contrib_ops/contrib_defs.cc | 2 +- .../providers/cuda/activation/activations.cc | 14 +- .../cuda/activation/activations_impl.cu | 15 +- .../core/providers/cuda/atomic/common.cuh | 21 +- .../core/providers/cuda/cu_inc/common.cuh | 112 ++++++++-- onnxruntime/core/providers/cuda/cuda_common.h | 12 - .../providers/cuda/cuda_execution_provider.cc | 12 +- .../providers/cuda/cuda_execution_provider.h | 10 +- onnxruntime/core/providers/cuda/cuda_utils.cu | 8 +- .../core/providers/cuda/cudnn_common.cc | 6 +- .../core/providers/cuda/cudnn_common.h | 4 +- onnxruntime/core/providers/cuda/fpgeneric.cu | 12 +- .../cuda/math/binary_elementwise_ops.cc | 40 +--- .../cuda/math/binary_elementwise_ops_impl.cu | 20 +- onnxruntime/core/providers/cuda/math/gemm.cc | 2 - .../core/providers/cuda/math/matmul.cc | 2 - .../core/providers/cuda/math/softmax.cc | 34 ++- .../core/providers/cuda/math/softmax_impl.cu | 20 +- .../cuda/math/unary_elementwise_ops_impl.cu | 26 +-- .../cuda/math/variadic_elementwise_ops.cc | 26 +-- .../math/variadic_elementwise_ops_impl.cu | 10 +- onnxruntime/core/providers/cuda/nn/dropout.cc | 8 +- onnxruntime/core/providers/cuda/nn/dropout.h | 12 - .../core/providers/cuda/nn/dropout_impl.cu | 4 +- .../cuda/reduction/reduction_functions.cu | 18 +- .../providers/cuda/reduction/reduction_ops.cc | 14 +- .../cuda/shared_inc/accumulation_type.h | 5 +- .../providers/cuda/shared_inc/fpgeneric.h | 54 ++--- .../core/providers/cuda/tensor/cast_op.cc | 6 - .../core/providers/cuda/tensor/gather_nd.cc | 41 +--- .../providers/cuda/tensor/gather_nd_impl.cu | 4 +- .../core/providers/rocm/atomic/common.cuh | 18 +- .../core/providers/rocm/cu_inc/common.cuh | 91 ++++++++ onnxruntime/test/framework/data_types_test.cc | 2 +- .../cpu/math/element_wise_ops_test.cc | 209 ++++++++++-------- .../test/providers/cpu/math/gemm_test.cc | 23 ++ .../cpu/reduction/reduction_ops_test.cc | 15 ++ .../test/providers/provider_test_utils.cc | 3 + .../test/providers/provider_test_utils.h | 6 + .../core/graph/training_op_defs.cc | 2 +- .../cuda/activation/bias_gelu_grad.cc | 12 +- .../cuda/activation/bias_gelu_grad_impl.cu | 6 +- .../cuda/cuda_training_kernels.cc | 12 +- .../loss/softmax_cross_entropy_loss_impl.cc | 4 + .../loss/softmax_cross_entropy_loss_impl.cu | 3 + .../cuda/math/mixed_precision_scale.cc | 37 +--- .../cuda/math/mixed_precision_scale.cu | 13 +- .../training_ops/cuda/math/softmax_grad.cc | 42 ++-- .../cuda/math/softmax_grad_impl.cu | 5 +- .../training_ops/cuda/nn/dropout_grad.cc | 25 +-- .../training_ops/cuda/nn/dropout_grad_impl.cu | 4 +- .../training_ops/cuda/nn/layer_norm.cc | 2 - .../training_ops/cuda/nn/layer_norm_impl.cu | 6 +- .../training_ops/cuda/optimizer/adam.cc | 3 - .../training_ops/cuda/optimizer/adam.cu | 21 +- .../cuda/optimizer/gradient_control.cc | 2 - .../cuda/optimizer/gradient_control.cu | 18 +- .../training_ops/cuda/optimizer/lamb.cc | 3 - .../training_ops/cuda/optimizer/lamb.cu | 53 ++--- .../cuda/reduction/reduction_all.cc | 2 - .../cuda/reduction/reduction_all.cu | 14 +- .../training_ops/cuda/tensor/gather_grad.cc | 33 +-- .../cuda/tensor/gather_grad_impl.cu | 4 +- .../cuda/tensor/gather_nd_grad.cc | 30 +-- .../cuda/tensor/gather_nd_grad_impl.cu | 4 +- 80 files changed, 710 insertions(+), 767 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index d8630076c1..ac2130f36c 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -486,7 +486,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
T : tensor(float16), tensor(float), tensor(double)
+
T : tensor(float16), tensor(float), tensor(double), tensor(bfloat16)
Constrain input and output types to float tensors.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index c3e2ff94d3..7e311e15c0 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -729,7 +729,7 @@ Do not modify directly.* |**Operator Domain:** *com.microsoft*|||| |Attention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* extra_add:**T**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(float), tensor(float16)| |BiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T2**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)| -|BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| +|BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |BiasSoftmax|*in* data:**T**
*in* bias:**T**
*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| diff --git a/include/onnxruntime/core/framework/float16.h b/include/onnxruntime/core/framework/float16.h index 5613b66c48..8de851d9ae 100644 --- a/include/onnxruntime/core/framework/float16.h +++ b/include/onnxruntime/core/framework/float16.h @@ -3,9 +3,19 @@ #pragma once #include "endian.h" +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#include "cuda_bf16.h" +#endif -namespace onnxruntime -{ +#include "core/common/common.h" + +namespace onnxruntime { + +#if defined(__CUDACC__) || defined(__HIPCC__) +#define ORT_HOST_DEVICE __host__ __device__ +#else +#define ORT_HOST_DEVICE +#endif // MLFloat16 struct MLFloat16 { @@ -17,53 +27,64 @@ struct MLFloat16 { float ToFloat() const; - operator float() const { - return ToFloat(); - } + operator float() const { return ToFloat(); } }; -inline bool operator==(const MLFloat16& left, const MLFloat16& right) { - return left.val == right.val; -} +inline bool operator==(const MLFloat16& left, const MLFloat16& right) { return left.val == right.val; } +inline bool operator!=(const MLFloat16& left, const MLFloat16& right) { return left.val != right.val; } +inline bool operator<(const MLFloat16& left, const MLFloat16& right) { return left.val < right.val; } -inline bool operator!=(const MLFloat16& left, const MLFloat16& right) { - return left.val != right.val; -} - -inline bool operator<(const MLFloat16& left, const MLFloat16& right) { - return left.val < right.val; -} - -//BFloat16 +// BFloat16 struct BFloat16 { uint16_t val{0}; - explicit BFloat16() = default; - explicit BFloat16(uint16_t v) : val(v) {} - explicit BFloat16(float v) { +#if defined(USE_ROCM) + ORT_HOST_DEVICE BFloat16() = default; +#else + BFloat16() = default; +#endif + + struct FromBitsT {}; + static constexpr ORT_HOST_DEVICE FromBitsT FromBits() { return FromBitsT(); } + constexpr ORT_HOST_DEVICE BFloat16(unsigned short bits, FromBitsT) : val(bits){}; + + inline ORT_HOST_DEVICE BFloat16(float v) { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + val = __bfloat16_as_ushort(__float2bfloat16(v)); +#else ORT_IF_CONSTEXPR(endian::native == endian::little) { std::memcpy(&val, reinterpret_cast(&v) + sizeof(uint16_t), sizeof(uint16_t)); - } else { + } + else { std::memcpy(&val, &v, sizeof(uint16_t)); } +#endif } - float ToFloat() const { + inline ORT_HOST_DEVICE float ToFloat() const { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + return __bfloat162float(*reinterpret_cast(&val)); +#else float result; char* const first = reinterpret_cast(&result); char* const second = first + sizeof(uint16_t); ORT_IF_CONSTEXPR(endian::native == endian::little) { std::memset(first, 0, sizeof(uint16_t)); std::memcpy(second, &val, sizeof(uint16_t)); - } else { + } + else { std::memcpy(first, &val, sizeof(uint16_t)); std::memset(second, 0, sizeof(uint16_t)); } return result; +#endif } - operator float() const { - return ToFloat(); - } + inline ORT_HOST_DEVICE operator float() const { return ToFloat(); } + +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + ORT_HOST_DEVICE BFloat16(const __nv_bfloat16& value) { val = *reinterpret_cast(&value); } + explicit ORT_HOST_DEVICE operator __nv_bfloat16() const { return *reinterpret_cast(&val); } +#endif }; inline void BFloat16ToFloat(const BFloat16* blf, float* flt, size_t size) { @@ -82,16 +103,4 @@ inline void FloatToBFloat16(const float* flt, BFloat16* blf, size_t size) { } } -inline bool operator==(const BFloat16& left, const BFloat16& right) { - return left.val == right.val; -} - -inline bool operator!=(const BFloat16& left, const BFloat16& right) { - return left.val != right.val; -} - -inline bool operator<(const BFloat16& left, const BFloat16& right) { - return left.val < right.val; -} - -} \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc index 864c614b99..29f823d994 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc @@ -25,9 +25,7 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_KERNEL_TYPED(BFloat16) -#endif using namespace ONNX_NAMESPACE; diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu index 4e9863edd2..61879024c2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu @@ -94,11 +94,10 @@ bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int i #if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) template -__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) { +__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]); @@ -106,11 +105,13 @@ __global__ void FastGeluKernel2(const nv_bfloat162 a, const nv_bfloat162 b, cons output[idx] = in * cdf; } } +#endif template <> -bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int input_length, int bias_length, const nv_bfloat16* input, const nv_bfloat16* bias, nv_bfloat16* output, bool /*use_half2*/) { +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; @@ -120,15 +121,18 @@ bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int i const nv_bfloat162* input2 = reinterpret_cast(input); const nv_bfloat162* bias2 = reinterpret_cast(bias); nv_bfloat162* output2 = reinterpret_cast(output); - FastGeluKernel2<<>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2); + FastGeluKernel2 + <<>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2); } else { +#endif const int gridSize = (input_length + blockSize - 1) / blockSize; - FastGeluKernel<<>>(A, B, C, input_length, bias_length, input, bias, output); + FastGeluKernel + <<>>(A, B, C, input_length, bias_length, input, bias, output); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) } - +#endif return CUDA_CALL(cudaPeekAtLastError()); } -#endif } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 19b226858a..442138d3ad 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -18,6 +18,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, BiasGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, BiasGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, BiasGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, BiasGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, TransposeMatMul); // backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, TransposeMatMul); // backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, TransposeMatMul); // backward compatibility @@ -86,13 +87,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int8_t, QAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FusedConv); - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FastGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul); // backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float, LayerNormalization); -#endif template <> KernelCreateInfo BuildKernelCreateInfo() { @@ -112,6 +110,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, // backward compatibility BuildKernelCreateInfo, // backward compatibility BuildKernelCreateInfo, // backward compatibility @@ -180,14 +179,11 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 BuildKernelCreateInfo, // TransposedMatMul is still here for backward compatibility BuildKernelCreateInfo, // backward compatibility BuildKernelCreateInfo, BuildKernelCreateInfo, -#endif BuildKernelCreateInfo, }; diff --git a/onnxruntime/contrib_ops/cuda/layer_norm.cc b/onnxruntime/contrib_ops/cuda/layer_norm.cc index 3095ebf437..3f7360bd77 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/layer_norm.cc @@ -35,9 +35,7 @@ namespace cuda { REGISTER_KERNEL_TYPED(float, float) REGISTER_KERNEL_TYPED(double, double) REGISTER_KERNEL_TYPED(MLFloat16, float) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_KERNEL_TYPED(BFloat16, float) -#endif template LayerNorm::LayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { diff --git a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu index 740aae4d45..d476ca3e25 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu @@ -394,10 +394,8 @@ LAYERNORM_LINEAR_IMPL(half, float, false) LAYERNORM_LINEAR_IMPL(double, double, false) //LAYERNORM_LINEAR_IMPL(half, half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -LAYERNORM_LINEAR_IMPL(nv_bfloat16, float, true) -LAYERNORM_LINEAR_IMPL(nv_bfloat16, float, false) -#endif +LAYERNORM_LINEAR_IMPL(BFloat16, float, true) +LAYERNORM_LINEAR_IMPL(BFloat16, float, false) } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc b/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc index 192c28c86d..a16a1b88ef 100644 --- a/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc +++ b/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc @@ -17,8 +17,8 @@ ONNX_OPERATOR_KERNEL_EX( 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) - .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) + .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T1", BuildKernelDefConstraints()) .TypeConstraint("T2", DataTypeImpl::GetTensorType()) .InputMemoryType(OrtMemTypeCPUInput, 3) .InputMemoryType(OrtMemTypeCPUInput, 4), @@ -96,7 +96,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const { float ratio_data = default_ratio_; auto ratio = context->Input(3); if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); t_disp.Invoke(ratio, ratio_data); } @@ -117,7 +117,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const { const fast_divmod fdm_dim(gsl::narrow_cast(dim)); PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default(); - utils::MLTypeCallDispatcher t_disp(X->GetElementType()); + utils::MLTypeCallDispatcher t_disp(X->GetElementType()); return t_disp.InvokeRet( GetDeviceProp(), Stream(), N, fdm_dim, ratio_data, generator, *X, *bias, residual, *Y, mask_data, has_same_shape_bias); } diff --git a/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu b/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu index ad2ff05a58..727d4f4878 100644 --- a/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu +++ b/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu @@ -238,9 +238,7 @@ void BiasDropoutKernelImpl( SPECIALIZED_BIAS_DROPOUT_IMPL(float) SPECIALIZED_BIAS_DROPOUT_IMPL(double) SPECIALIZED_BIAS_DROPOUT_IMPL(half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_BIAS_DROPOUT_IMPL(nv_bfloat16) -#endif +SPECIALIZED_BIAS_DROPOUT_IMPL(BFloat16) } // namespace cuda } // namespace contrib { diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc index 0df46fb763..6ce2446e44 100644 --- a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc @@ -61,7 +61,8 @@ namespace cuda { #define CONTRIB_BINARY_OP_HFD(name, ver) \ CONTRIB_BINARY_OP_TYPED(name, ver, MLFloat16) \ CONTRIB_BINARY_OP_TYPED(name, ver, float) \ - CONTRIB_BINARY_OP_TYPED(name, ver, double) + CONTRIB_BINARY_OP_TYPED(name, ver, double) \ + CONTRIB_BINARY_OP_TYPED(name, ver, BFloat16) CONTRIB_BINARY_OP_HFD(BiasGelu, 1) diff --git a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu index 01791ed94c..bb522fe9e2 100644 --- a/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu @@ -63,7 +63,8 @@ namespace cuda { #define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(x) \ CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ - CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) \ + CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, BFloat16) // create declarations for op and impl #define CONTRIB_BINARY_OP_NAME_EXPR(name, expr) \ diff --git a/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc b/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc index 86ce8d3efa..a2fad0ac4d 100644 --- a/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc +++ b/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc @@ -22,15 +22,12 @@ namespace cuda { REGISTER_KERNEL_TYPED(TransposeMatMul, float) REGISTER_KERNEL_TYPED(TransposeMatMul, double) REGISTER_KERNEL_TYPED(TransposeMatMul, MLFloat16) +REGISTER_KERNEL_TYPED(TransposeMatMul, BFloat16) REGISTER_KERNEL_TYPED(FusedMatMul, float) REGISTER_KERNEL_TYPED(FusedMatMul, double) REGISTER_KERNEL_TYPED(FusedMatMul, MLFloat16) - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -REGISTER_KERNEL_TYPED(TransposeMatMul, BFloat16) REGISTER_KERNEL_TYPED(FusedMatMul, BFloat16) -#endif } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/math/isfinite.cuh b/onnxruntime/contrib_ops/cuda/math/isfinite.cuh index 1116228a9d..d27a6631c1 100644 --- a/onnxruntime/contrib_ops/cuda/math/isfinite.cuh +++ b/onnxruntime/contrib_ops/cuda/math/isfinite.cuh @@ -5,10 +5,6 @@ #include "core/providers/cuda/cu_inc/common.cuh" #include "contrib_ops/cuda/math/isfinite.h" -#if CUDA_VERSION >= 11000 -#include "cuda_bf16.h" -#endif - namespace onnxruntime { namespace cuda { @@ -54,22 +50,20 @@ __device__ __forceinline__ bool IsNaNScalar(const half value) { #endif } -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) template <> -__device__ __forceinline__ bool IsFiniteScalar(const nv_bfloat16 value) { - return !__hisinf(value) && !__hisnan(value); +__device__ __forceinline__ bool IsFiniteScalar(const BFloat16 value) { + return isfinite(static_cast(value)); } template <> -__device__ __forceinline__ bool IsInfScalar(const nv_bfloat16 value) { - return __hisinf(value); +__device__ __forceinline__ bool IsInfScalar(const BFloat16 value) { + return isinf(static_cast(value)); } template <> -__device__ __forceinline__ bool IsNaNScalar(const nv_bfloat16 value) { - return __hisnan(value); +__device__ __forceinline__ bool IsNaNScalar(const BFloat16 value) { + return isnan(static_cast(value)); } -#endif } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 2b6d2bbed1..08e86d815f 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -389,7 +389,7 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d if (v < 0 || v > max_value) { return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "data overflow"); } - p_data[i] = BFloat16(static_cast(v)); + p_data[i] = BFloat16(static_cast(v), BFloat16::FromBits()); } return Status::OK(); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 502244fa4b..fcc219799b 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -2874,7 +2874,7 @@ It's an extension of Gelu. It takes the sum of input A and bias input B as the i .Output(0, "C", "The output.", "T") .TypeConstraint( "T", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); diff --git a/onnxruntime/core/providers/cuda/activation/activations.cc b/onnxruntime/core/providers/cuda/activation/activations.cc index 08e0f86106..df02e333b3 100644 --- a/onnxruntime/core/providers/cuda/activation/activations.cc +++ b/onnxruntime/core/providers/cuda/activation/activations.cc @@ -58,25 +58,17 @@ namespace cuda { REGISTER_ACTIVATION_KERNEL(name, ver, T) \ UNARY_ACTIVATION_COMPUTE(name, T) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) UNARY_ACTIVATION_OP_TYPED(name, ver, BFloat16) -#define UNARY_ACTIVATION_OP_VERSIONED_TYPED_BF16(name, startver, endver) UNARY_ACTIVATION_OP_VERSIONED_TYPED(name, startver, endver, BFloat16) -#else -#define UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) -#define UNARY_ACTIVATION_OP_VERSIONED_TYPED_BF16(name, startver, endver) -#endif - #define UNARY_ACTIVATION_OP_VERSIONED_HFD_WITH_BF16(name, startver, endver) \ UNARY_ACTIVATION_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ UNARY_ACTIVATION_OP_VERSIONED_TYPED(name, startver, endver, float) \ UNARY_ACTIVATION_OP_VERSIONED_TYPED(name, startver, endver, double) \ - UNARY_ACTIVATION_OP_VERSIONED_TYPED_BF16(name, startver, endver) + UNARY_ACTIVATION_OP_VERSIONED_TYPED(name, startver, endver, BFloat16) #define UNARY_ACTIVATION_OP_HFD(name, ver) \ UNARY_ACTIVATION_OP_TYPED(name, ver, MLFloat16) \ - UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) \ UNARY_ACTIVATION_OP_TYPED(name, ver, float) \ - UNARY_ACTIVATION_OP_TYPED(name, ver, double) + UNARY_ACTIVATION_OP_TYPED(name, ver, double) \ + UNARY_ACTIVATION_OP_TYPED(name, ver, BFloat16) UNARY_ACTIVATION_OP_HFD(Elu, 6); UNARY_ACTIVATION_OP_HFD(HardSigmoid, 6); diff --git a/onnxruntime/core/providers/cuda/activation/activations_impl.cu b/onnxruntime/core/providers/cuda/activation/activations_impl.cu index 2ff5a4748f..1a4dc3cf79 100644 --- a/onnxruntime/core/providers/cuda/activation/activations_impl.cu +++ b/onnxruntime/core/providers/cuda/activation/activations_impl.cu @@ -91,20 +91,15 @@ struct OP_ThresholdedRelu : public CtxThresholdedRelu { count); \ } -#define SPECIALIZED_UNARY_ACTIVATION_IMPL(name, T) \ - template void Impl_##name(cudaStream_t stream, const T* input_data, T* output_data, const Ctx##name* func_ctx, size_t count); - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -#define SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) SPECIALIZED_UNARY_ACTIVATION_IMPL(name, nv_bfloat16) -#else -#define SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) -#endif +#define SPECIALIZED_UNARY_ACTIVATION_IMPL(name, T) \ + template void Impl_##name(cudaStream_t stream, const T* input_data, T* output_data, const Ctx##name* func_ctx, \ + size_t count); #define SPECIALIZED_UNARY_ACTIVATIONL_HFD(name) \ SPECIALIZED_UNARY_ACTIVATION_IMPL(name, half) \ - SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) \ SPECIALIZED_UNARY_ACTIVATION_IMPL(name, float) \ - SPECIALIZED_UNARY_ACTIVATION_IMPL(name, double) + SPECIALIZED_UNARY_ACTIVATION_IMPL(name, double) \ + SPECIALIZED_UNARY_ACTIVATION_IMPL(name, BFloat16) #define UNARY_ACTIVATION_OP_NAME(name) \ UNARY_ACTIVATION_IMPL(name); \ diff --git a/onnxruntime/core/providers/cuda/atomic/common.cuh b/onnxruntime/core/providers/cuda/atomic/common.cuh index 098f650986..d6751f9bd0 100644 --- a/onnxruntime/core/providers/cuda/atomic/common.cuh +++ b/onnxruntime/core/providers/cuda/atomic/common.cuh @@ -20,10 +20,7 @@ #include "cuda.h" #include "cuda_fp16.h" #include "cuda_runtime.h" - -#if CUDA_VERSION >= 11000 -#include "cuda_bf16.h" -#endif +#include "core/framework/float16.h" namespace onnxruntime { namespace cuda { @@ -72,22 +69,20 @@ __device__ __forceinline__ void atomic_add(half *address, half value) { #endif } -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -__device__ __forceinline__ void atomic_add(nv_bfloat16 *address, nv_bfloat16 value) { - unsigned int * base_address = reinterpret_cast(reinterpret_cast(address) - (reinterpret_cast(address) & 2)); +__device__ __forceinline__ void atomic_add(BFloat16* address, BFloat16 value) { + unsigned int* base_address = + reinterpret_cast(reinterpret_cast(address) - (reinterpret_cast(address) & 2)); unsigned int old = *base_address; unsigned int assumed; - unsigned short x; - + BFloat16 bsum; do { assumed = old; - x = reinterpret_cast(address) & 2 ? (old >> 16) : (old & 0xffff); - x = __bfloat16_as_short(__float2bfloat16(__bfloat162float(*reinterpret_cast(&x)) + __bfloat162float(value))); - old = reinterpret_cast(address) & 2 ? (old & 0xffff) | (x << 16) : (old & 0xffff0000) | x; + bsum.val = reinterpret_cast(address) & 2 ? (old >> 16) : (old & 0xffff); + bsum = bsum + value; + old = reinterpret_cast(address) & 2 ? (old & 0xffff) | (bsum.val << 16) : (old & 0xffff0000) | bsum.val; old = atomicCAS(base_address, assumed, old); } while (assumed != old); } -#endif } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index 7b283731e9..ea9ddf0450 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -81,6 +81,82 @@ __device__ __forceinline__ half2 operator*(const half2& lh, const half2& rh) { h __device__ __forceinline__ half2 operator/(const half2& lh, const half2& rh) { half2 r; r.x = lh.x / rh.x; r.y = lh.y / rh.y; return r; } #endif +/// Arithmetic for BFloat16 + +__device__ __forceinline__ BFloat16 operator+(const BFloat16& a, const BFloat16& b) { + return static_cast(a) + static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator-(const BFloat16& a, const BFloat16& b) { + return static_cast(a) - static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator*(const BFloat16& a, const BFloat16& b) { + return static_cast(a) * static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator/(const BFloat16& a, const BFloat16& b) { + return static_cast(a) / static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator-(const BFloat16& a) { return -static_cast(a); } + +__device__ __forceinline__ BFloat16& operator+=(BFloat16& a, const BFloat16& b) { + a = a + b; + return a; +} + +__device__ __forceinline__ BFloat16& operator-=(BFloat16& a, const BFloat16& b) { + a = a - b; + return a; +} + +__device__ __forceinline__ BFloat16& operator*=(BFloat16& a, const BFloat16& b) { + a = a * b; + return a; +} + +__device__ __forceinline__ BFloat16& operator/=(BFloat16& a, const BFloat16& b) { + a = a / b; + return a; +} + +/// Arithmetic with floats + +__device__ __forceinline__ float operator+(BFloat16 a, float b) { return static_cast(a) + b; } +__device__ __forceinline__ float operator-(BFloat16 a, float b) { return static_cast(a) - b; } +__device__ __forceinline__ float operator*(BFloat16 a, float b) { return static_cast(a) * b; } +__device__ __forceinline__ float operator/(BFloat16 a, float b) { return static_cast(a) / b; } + +__device__ __forceinline__ float operator+(float a, BFloat16 b) { return a + static_cast(b); } +__device__ __forceinline__ float operator-(float a, BFloat16 b) { return a - static_cast(b); } +__device__ __forceinline__ float operator*(float a, BFloat16 b) { return a * static_cast(b); } +__device__ __forceinline__ float operator/(float a, BFloat16 b) { return a / static_cast(b); } + +__device__ __forceinline__ float& operator+=(float& a, const BFloat16& b) { return a += static_cast(b); } +__device__ __forceinline__ float& operator-=(float& a, const BFloat16& b) { return a -= static_cast(b); } +__device__ __forceinline__ float& operator*=(float& a, const BFloat16& b) { return a *= static_cast(b); } +__device__ __forceinline__ float& operator/=(float& a, const BFloat16& b) { return a /= static_cast(b); } + +/// Arithmetic with doubles + +__device__ __forceinline__ double operator+(BFloat16 a, double b) { return static_cast(a) + b; } +__device__ __forceinline__ double operator-(BFloat16 a, double b) { return static_cast(a) - b; } +__device__ __forceinline__ double operator*(BFloat16 a, double b) { return static_cast(a) * b; } +__device__ __forceinline__ double operator/(BFloat16 a, double b) { return static_cast(a) / b; } + +__device__ __forceinline__ double operator+(double a, BFloat16 b) { return a + static_cast(b); } +__device__ __forceinline__ double operator-(double a, BFloat16 b) { return a - static_cast(b); } +__device__ __forceinline__ double operator*(double a, BFloat16 b) { return a * static_cast(b); } +__device__ __forceinline__ double operator/(double a, BFloat16 b) { return a / static_cast(b); } + +// Overloading < and > operators + +__device__ __forceinline__ bool operator==(BFloat16& lhs, BFloat16& rhs) { return float(lhs) == float(rhs); } +__device__ __forceinline__ bool operator!=(BFloat16& lhs, BFloat16& rhs) { return float(lhs) != float(rhs); } +__device__ __forceinline__ bool operator>(BFloat16& lhs, BFloat16& rhs) { return float(lhs) > float(rhs); } +__device__ __forceinline__ bool operator<(BFloat16& lhs, BFloat16& rhs) { return float(lhs) < float(rhs); } + template __device__ __inline__ T _Ceil(T a); @@ -263,24 +339,19 @@ __device__ __inline__ double _Normcdf(double a) { return normcdf(a); } template <> __device__ __inline__ half _Normcdf(half a) { return half(normcdff((float)a)); } -template -__device__ __inline__ T _Gelu(T a) { - return a * _Normcdf(a); -} +template <> +__device__ __inline__ BFloat16 _Sqrt(BFloat16 a) { return sqrtf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Exp(BFloat16 a) { return expf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Log(BFloat16 a) { return logf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Tanh(BFloat16 a) { return tanhf(static_cast(a)); } #if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -template <> -__device__ __inline__ nv_bfloat16 _Sqrt(nv_bfloat16 a) { return nv_bfloat16(sqrtf(static_cast(a))); } - -template <> -__device__ __inline__ nv_bfloat16 _Exp(nv_bfloat16 a) { return nv_bfloat16(expf(static_cast(a))); } - -template <> -__device__ __inline__ nv_bfloat16 _Log(nv_bfloat16 a) { return nv_bfloat16(logf(static_cast(a))); } - -template <> -__device__ __inline__ nv_bfloat16 _Tanh(nv_bfloat16 a) { return nv_bfloat16(tanhf(static_cast(a))); } - template <> __device__ __inline__ nv_bfloat162 _Tanh(nv_bfloat162 a) { float2 tmp = (__bfloat1622float2(a)); @@ -288,10 +359,15 @@ __device__ __inline__ nv_bfloat162 _Tanh(nv_bfloat162 a) { tmp.y = tanhf(tmp.y); return __float22bfloat162_rn(tmp); } +#endif template <> -__device__ __inline__ nv_bfloat16 _Normcdf(nv_bfloat16 a) { return nv_bfloat16(normcdff(static_cast(a))); } -#endif +__device__ __inline__ BFloat16 _Normcdf(BFloat16 a) { return normcdff(static_cast(a)); } + +template +__device__ __inline__ T _Gelu(T a) { + return a * _Normcdf(a); +} // We would like to use 64-bit integer to support large matrices. However, CUDA seems to support only 32-bit integer // For now, use int32_t to ensure that both Linux and Windows see this as 32 bit integer type. diff --git a/onnxruntime/core/providers/cuda/cuda_common.h b/onnxruntime/core/providers/cuda/cuda_common.h index 94f56a346c..1d41b47536 100644 --- a/onnxruntime/core/providers/cuda/cuda_common.h +++ b/onnxruntime/core/providers/cuda/cuda_common.h @@ -69,18 +69,6 @@ class ToCudaType { } }; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -template <> -class ToCudaType { - public: - typedef nv_bfloat16 MappedType; - static MappedType FromFloat(float f) { - uint16_t h = BFloat16(f).val; - return *reinterpret_cast(&h); - } -}; -#endif - inline bool CalculateFdmStrides(gsl::span p, const std::vector& dims) { int stride = 1; if (dims.empty() || p.size() < dims.size()) diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index cd5e7ad6b9..5ebdb8dd1f 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -1116,8 +1116,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, bool, Pad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, SpaceToDepth); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, DepthToSpace); - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Add); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Sub); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Mul); @@ -1130,7 +1128,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum); -#endif // OpSet 14 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum); @@ -1187,14 +1184,11 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, uint8_t, ReduceMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, int64_t, ReduceMin); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, Trilu); - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Add); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Sub); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Mul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Div); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Relu); -#endif // OpSet 15 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 15, Pow); @@ -1966,8 +1960,6 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -1980,7 +1972,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, -#endif + // OpSet 14 BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2035,13 +2027,11 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, -#endif BuildKernelCreateInfo, // OpSet 15 diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index 67120de11c..7e5326ddc7 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -143,13 +143,11 @@ class CUDAExecutionProvider : public IExecutionProvider { constant_ones_half_ = cuda::CreateConstantOnes(); } return reinterpret_cast(constant_ones_half_->GetBuffer(stream_, count)); -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 - } else if (std::is_same::value) { + } else if (std::is_same::value) { if (!constant_ones_bfloat16_) { - constant_ones_bfloat16_ = cuda::CreateConstantOnes(); + constant_ones_bfloat16_ = cuda::CreateConstantOnes(); } return reinterpret_cast(constant_ones_bfloat16_->GetBuffer(stream_, count)); -#endif } else { return nullptr; } @@ -172,9 +170,7 @@ class CUDAExecutionProvider : public IExecutionProvider { std::unique_ptr> constant_ones_float_; std::unique_ptr> constant_ones_double_; std::unique_ptr> constant_ones_half_; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 - std::unique_ptr> constant_ones_bfloat16_; -#endif + std::unique_ptr> constant_ones_bfloat16_; AllocatorPtr allocator_; }; diff --git a/onnxruntime/core/providers/cuda/cuda_utils.cu b/onnxruntime/core/providers/cuda/cuda_utils.cu index 2296b127ae..cc9eaf9ae3 100644 --- a/onnxruntime/core/providers/cuda/cuda_utils.cu +++ b/onnxruntime/core/providers/cuda/cuda_utils.cu @@ -71,9 +71,7 @@ std::unique_ptr> CreateConstantOnes() { template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -template std::unique_ptr> CreateConstantOnes(); -#endif +template std::unique_ptr> CreateConstantOnes(); #define SPECIALIZED_FILL(T) \ template void Fill(cudaStream_t stream, T * output, T value, int64_t count); @@ -85,9 +83,7 @@ SPECIALIZED_FILL(int64_t) SPECIALIZED_FILL(float) SPECIALIZED_FILL(double) SPECIALIZED_FILL(__half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_FILL(nv_bfloat16) -#endif +SPECIALIZED_FILL(BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cudnn_common.cc b/onnxruntime/core/providers/cuda/cudnn_common.cc index 52889ff8c9..d85ca2d7cf 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.cc +++ b/onnxruntime/core/providers/cuda/cudnn_common.cc @@ -161,10 +161,8 @@ const float Consts::Zero = 0; const float Consts::One = 1; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -const float Consts::Zero = 0; -const float Consts::One = 1; -#endif +const float Consts::Zero = 0; +const float Consts::One = 1; template <> const int8_t Consts::Zero = 0; diff --git a/onnxruntime/core/providers/cuda/cudnn_common.h b/onnxruntime/core/providers/cuda/cudnn_common.h index 6526fc0aae..f8cee8fa89 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.h +++ b/onnxruntime/core/providers/cuda/cudnn_common.h @@ -124,13 +124,11 @@ struct Consts { static const float One; }; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 template <> -struct Consts { +struct Consts { static const float Zero; static const float One; }; -#endif inline double ClampCudnnBatchNormEpsilon(double epsilon) { if (epsilon < CUDNN_BN_MIN_EPSILON) { diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 695c6038c0..3ae0b9c8eb 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -55,13 +55,12 @@ __global__ void CopyVectorHalf(const half* x, int incx, half* y, int incy, int n y[id * incy] = x[id * incx]; } -#if CUDA_VERSION >= 11000 -__global__ void CopyVectorBFloat16(const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy, int n) { +__global__ void CopyVectorBFloat16(const onnxruntime::BFloat16* x, int incx, onnxruntime::BFloat16* y, int incy, + int n) { int id = blockIdx.x * blockDim.x + threadIdx.x; if (id >= n) return; y[id * incy] = x[id * incx]; } -#endif } // namespace @@ -84,13 +83,10 @@ cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t, int n, cons return CUBLAS_STATUS_SUCCESS; } -#if CUDA_VERSION >= 11000 -cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t, int n, const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy) { +cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t, int n, const onnxruntime::BFloat16* x, int incx, + onnxruntime::BFloat16* y, int incy) { dim3 dimGrid((unsigned int)(n + COPY_BLOCK_DIM - 1) / COPY_BLOCK_DIM, 1, 1); dim3 dimBlock(COPY_BLOCK_DIM, 1, 1); CopyVectorBFloat16<<>>(x, incx, y, incy, n); return CUBLAS_STATUS_SUCCESS; } - - -#endif diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc index fd669a72a5..e98e639d67 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc @@ -189,22 +189,6 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin // D: double // O: bool -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define BINARY_OP_TYPED_BF16(name, ver) BINARY_OP_TYPED(name, ver, BFloat16) -#define BINARY_OP_VERSIONED_TYPED_BF16(name, startver, endver) BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, BFloat16) -#define BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, BFloat16) -#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, BFloat16) -#define BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, BFloat16) -#define BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, BFloat16) -#else -#define BINARY_OP_TYPED_BF16(name, ver) -#define BINARY_OP_VERSIONED_TYPED_BF16(name, startver, endver) -#define BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) -#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) -#define BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) -#define BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) -#endif - #define BINARY_OP_VERSIONED_HFD(name, startver, endver) \ BINARY_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ BINARY_OP_VERSIONED_TYPED(name, startver, endver, float) \ @@ -222,14 +206,14 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin BINARY_OP_VERSIONED_TYPED(name, startver, endver, uint64_t) \ BINARY_OP_VERSIONED_TYPED(name, startver, endver, int32_t) \ BINARY_OP_VERSIONED_TYPED(name, startver, endver, int64_t) \ - BINARY_OP_VERSIONED_TYPED_BF16(name, startver, endver) \ - BINARY_OP_VERSIONED_HFD(name, startver, endver) + BINARY_OP_VERSIONED_HFD(name, startver, endver) \ + BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, BFloat16) #define BINARY_OP_HFD(name, ver) \ BINARY_OP_TYPED(name, ver, MLFloat16) \ - BINARY_OP_TYPED_BF16(name, ver) \ BINARY_OP_TYPED(name, ver, float) \ - BINARY_OP_TYPED(name, ver, double) + BINARY_OP_TYPED(name, ver, double) \ + BINARY_OP_TYPED(name, ver, BFloat16) #define BINARY_OP_UZILHFD(name, ver) \ BINARY_OP_TYPED(name, ver, uint32_t) \ @@ -250,9 +234,9 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin #define BINARY_OP_REGISTER_HFD(name, ver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, MLFloat16) \ - BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, float) \ - BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, double) + BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, double) \ + BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, BFloat16) #define BINARY_OP_REGISTER_UZILHFD(name, ver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, uint32_t) \ @@ -267,21 +251,21 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, int32_t) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, int64_t) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, MLFloat16) \ - BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, float) \ - BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, double) + BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, double) \ + BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, BFloat16) #define BINARY_OP_REGISTER_VERSIONED_HFD(name, startver, endver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ - BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, float) \ - BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, double) + BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, double) \ + BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, BFloat16) #define BINARY_OP_REGISTER_VERSIONED_CLASS_HFD(name, class_name, startver, endver) \ BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, MLFloat16) \ - BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) \ BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, float) \ - BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, double) + BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, double) \ + BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, BFloat16) #define BINARY_OP_REGISTER_VERSIONED_UZILHFD(name, startver, endver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, uint32_t) \ diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu index 39ae0acfd8..e408cbbc04 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu @@ -79,23 +79,15 @@ namespace cuda { const TArray* rhs_padded_strides, const T2* rhs_data, \ const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, nv_bfloat16) -#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, nv_bfloat16, nv_bfloat16) -#else -#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) -#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) -#endif - #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) + SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, BFloat16) #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1_ILHFD(x, T) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, int32_t) \ @@ -111,9 +103,9 @@ namespace cuda { #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) + SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, BFloat16) // create declarations for impl #define BINARY_OP_NAME_EXPR(name, expr) \ @@ -172,9 +164,9 @@ BINARY_OPS2() SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int32_t, int32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int64_t, int64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, half, half) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, float, float) \ - SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, double, double) + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, double, double) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, BFloat16, BFloat16) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD2(Greater) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD2(Equal) diff --git a/onnxruntime/core/providers/cuda/math/gemm.cc b/onnxruntime/core/providers/cuda/math/gemm.cc index def63e4db2..5e61bf4e08 100644 --- a/onnxruntime/core/providers/cuda/math/gemm.cc +++ b/onnxruntime/core/providers/cuda/math/gemm.cc @@ -53,9 +53,7 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) REGISTER_KERNEL_TYPED(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_KERNEL_TYPED(BFloat16) -#endif template Status Gemm::ComputeInternal(OpKernelContext* ctx) const { diff --git a/onnxruntime/core/providers/cuda/math/matmul.cc b/onnxruntime/core/providers/cuda/math/matmul.cc index c7632dcdc5..abdef98751 100644 --- a/onnxruntime/core/providers/cuda/math/matmul.cc +++ b/onnxruntime/core/providers/cuda/math/matmul.cc @@ -41,9 +41,7 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) REGISTER_KERNEL_TYPED(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_KERNEL_TYPED(BFloat16) -#endif // StridedBatchedGemm can be used for the following GEMM computation // C[pnm] = A[pnk]*B[km] or C[pnm] = A[pnk]*B[pkm] diff --git a/onnxruntime/core/providers/cuda/math/softmax.cc b/onnxruntime/core/providers/cuda/math/softmax.cc index bddf60947d..9116178dda 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.cc +++ b/onnxruntime/core/providers/cuda/math/softmax.cc @@ -44,29 +44,23 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL(float) SPECIALIZED_SOFTMAX_HELPER_IMPL(double) SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 // cudnnSoftmaxForward/Backward doesn't support BFloat16. -#define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \ - template <> \ - Status SoftMaxComputeHelper( \ - cudaStream_t stream, \ - const BFloat16* X, \ - const TensorShape& input_shape, \ - BFloat16* Y, \ - int64_t axis) { \ - typedef typename ToCudaType::MappedType CudaT; \ - int64_t N = input_shape.SizeToDimension(axis); \ - int64_t D = input_shape.SizeFromDimension(axis); \ - auto Y_data = reinterpret_cast(Y); \ - auto X_data = reinterpret_cast(X); \ - dispatch_warpwise_softmax_forward, is_log_softmax>( \ - stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ - return Status::OK(); \ +#define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \ + template <> \ + Status SoftMaxComputeHelper(cudaStream_t stream, const BFloat16* X, \ + const TensorShape& input_shape, BFloat16* Y, int64_t axis) { \ + typedef typename ToCudaType::MappedType CudaT; \ + int64_t N = input_shape.SizeToDimension(axis); \ + int64_t D = input_shape.SizeFromDimension(axis); \ + auto Y_data = reinterpret_cast(Y); \ + auto X_data = reinterpret_cast(X); \ + dispatch_warpwise_softmax_forward, is_log_softmax>( \ + stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ + return Status::OK(); \ } SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true) - SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) -#endif +SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ @@ -225,9 +219,7 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true) SPECIALIZED_COMPUTE(float) SPECIALIZED_COMPUTE(double) SPECIALIZED_COMPUTE(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 SPECIALIZED_COMPUTE(BFloat16) -#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/softmax_impl.cu b/onnxruntime/core/providers/cuda/math/softmax_impl.cu index 8a3396a448..4ef7f6674b 100644 --- a/onnxruntime/core/providers/cuda/math/softmax_impl.cu +++ b/onnxruntime/core/providers/cuda/math/softmax_impl.cu @@ -109,9 +109,7 @@ template void dispatch_warpwise_softmax_forward( SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(float, float, float) SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(half, half, float) SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(double, double, double) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(nv_bfloat16, nv_bfloat16, float) -#endif +SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(BFloat16, BFloat16, float) template void dispatch_blockwise_softmax_forward(cudaStream_t stream, output_t* output, const input_t* input, int softmax_elements, int softmax_elements_stride, int batch_count) { @@ -127,17 +125,17 @@ void dispatch_blockwise_softmax_forward(cudaStream_t stream, output_t* output, c } } -#define SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(input_t, output_t, acc_t) \ -template void dispatch_blockwise_softmax_forward(cudaStream_t stream, output_t* output, const input_t* src, int softmax_elements, int softmax_elements_stride, int batch_count); \ -template void dispatch_blockwise_softmax_forward(cudaStream_t stream, output_t* output, const input_t* src, int softmax_elements, int softmax_elements_stride, int batch_count); +#define SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(input_t, output_t, acc_t) \ + template void dispatch_blockwise_softmax_forward( \ + cudaStream_t stream, output_t * output, const input_t* src, int softmax_elements, int softmax_elements_stride, \ + int batch_count); \ + template void dispatch_blockwise_softmax_forward( \ + cudaStream_t stream, output_t * output, const input_t* src, int softmax_elements, int softmax_elements_stride, \ + int batch_count); SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(float, float, float) SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, half, float) SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(double, double, double) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(nv_bfloat16, nv_bfloat16, float) -#endif - - +SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(BFloat16, BFloat16, float) } } diff --git a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu index 66d50dafa5..9bb5a2d855 100644 --- a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu @@ -94,24 +94,18 @@ struct ViaTypeMap { typedef float ViaT; }; -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) template <> -struct ViaTypeMap { +struct ViaTypeMap { typedef float ViaT; }; -#endif template struct OP_Cast { __device__ __inline__ OutT operator()(const InT& a) const { const bool any_float16 = std::is_same::value || std::is_same::value; -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) - const bool any_bf16 = std::is_same::value || std::is_same::value; - typedef typename std::conditional::type T1; + const bool any_bf16 = std::is_same::value || std::is_same::value; + typedef typename std::conditional::type T1; typedef typename std::conditional::type T; -#else - typedef typename std::conditional::type T; -#endif typedef typename ViaTypeMap::ViaT ViaT; return (OutT)((ViaT)a); } @@ -133,15 +127,8 @@ void Impl_Cast( #define SPECIALIZED_CAST_IMPL2(InT, OutT) \ template void Impl_Cast(cudaStream_t stream, const InT* input_data, OutT* output_data, size_t count); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -#define SPECIALIZED_CAST_IMPL2_BF16(T) SPECIALIZED_CAST_IMPL2(T, nv_bfloat16) -#else -#define SPECIALIZED_CAST_IMPL2_BF16(T) -#endif - #define SPECIALIZED_CAST_FROM(T) \ SPECIALIZED_CAST_IMPL2(T, half) \ - SPECIALIZED_CAST_IMPL2_BF16(T) \ SPECIALIZED_CAST_IMPL2(T, float) \ SPECIALIZED_CAST_IMPL2(T, double) \ SPECIALIZED_CAST_IMPL2(T, int8_t) \ @@ -152,7 +139,8 @@ void Impl_Cast( SPECIALIZED_CAST_IMPL2(T, uint16_t) \ SPECIALIZED_CAST_IMPL2(T, uint32_t) \ SPECIALIZED_CAST_IMPL2(T, uint64_t) \ - SPECIALIZED_CAST_IMPL2(T, bool) + SPECIALIZED_CAST_IMPL2(T, bool) \ + SPECIALIZED_CAST_IMPL2(T, BFloat16) SPECIALIZED_CAST_FROM(half) SPECIALIZED_CAST_FROM(float) @@ -166,9 +154,7 @@ SPECIALIZED_CAST_FROM(uint16_t) SPECIALIZED_CAST_FROM(uint32_t) SPECIALIZED_CAST_FROM(uint64_t) SPECIALIZED_CAST_FROM(bool) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_CAST_FROM(nv_bfloat16) -#endif +SPECIALIZED_CAST_FROM(BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc index dc4db5a830..81d9e15ce3 100644 --- a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc @@ -228,26 +228,18 @@ Status VariadicElementwiseOp namespace { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double, BFloat16 -#else -#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double -#endif +using SumOp = VariadicElementwiseOp; -using SumOp = VariadicElementwiseOp< - variadic_elementwise_ops::Sum, - ALL_IEEE_FLOAT_DATA_TYPES>; +using MinOp = VariadicElementwiseOp; -using MinOp = VariadicElementwiseOp< - variadic_elementwise_ops::Min, - uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>; +using MaxOp = VariadicElementwiseOp; -using MaxOp = VariadicElementwiseOp< - variadic_elementwise_ops::Max, - uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>; - -const DeleteOnUnloadPtr> k_uzilhfd_datatypes = new std::vector(BuildKernelDefConstraints()); -const DeleteOnUnloadPtr> k_hfd_datatypes = new std::vector(BuildKernelDefConstraints()); +const DeleteOnUnloadPtr> k_uzilhfd_datatypes = new std::vector( + BuildKernelDefConstraints()); +const DeleteOnUnloadPtr> k_hfd_datatypes = + new std::vector(BuildKernelDefConstraints()); } // namespace diff --git a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu index da1f228e5b..a17dc0df13 100644 --- a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu @@ -129,17 +129,11 @@ void Impl_NoBroadcastInputBatch( // D: double // O: bool -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -#define SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) SPECIALIZE_IMPL(nv_bfloat16, VariadicElementwiseOpTag) -#else -#define SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) -#endif - #define SPECIALIZE_IMPL_HFD(VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(half, VariadicElementwiseOpTag) \ - SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(float, VariadicElementwiseOpTag) \ - SPECIALIZE_IMPL(double, VariadicElementwiseOpTag) + SPECIALIZE_IMPL(double, VariadicElementwiseOpTag) \ + SPECIALIZE_IMPL(BFloat16, VariadicElementwiseOpTag) #define SPECIALIZE_IMPL_UZILHFD(VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(uint32_t, VariadicElementwiseOpTag) \ diff --git a/onnxruntime/core/providers/cuda/nn/dropout.cc b/onnxruntime/core/providers/cuda/nn/dropout.cc index 91d9560582..40c4826b56 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.cc +++ b/onnxruntime/core/providers/cuda/nn/dropout.cc @@ -25,8 +25,8 @@ ONNX_OPERATOR_KERNEL_EX( 13, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) - .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) + .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T1", BuildKernelDefConstraints()) .TypeConstraint("T2", DataTypeImpl::GetTensorType()) .InputMemoryType(OrtMemTypeCPUInput, 1) .InputMemoryType(OrtMemTypeCPUInput, 2), @@ -50,7 +50,7 @@ Status Dropout::ComputeInternal(OpKernelContext* context) const { float ratio_data = default_ratio_; auto ratio = context->Input(1); if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); t_disp.Invoke(ratio, ratio_data); } @@ -80,7 +80,7 @@ Status Dropout::ComputeInternal(OpKernelContext* context) const { PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default(); - utils::MLTypeCallDispatcher t_disp(X->GetElementType()); + utils::MLTypeCallDispatcher t_disp(X->GetElementType()); t_disp.Invoke(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data); return Status::OK(); diff --git a/onnxruntime/core/providers/cuda/nn/dropout.h b/onnxruntime/core/providers/cuda/nn/dropout.h index af04a36b40..78ef4e3237 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.h +++ b/onnxruntime/core/providers/cuda/nn/dropout.h @@ -10,18 +10,6 @@ namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double, BFloat16 -#else -#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() -#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double -#endif - template struct GetRatioDataImpl { void operator()(const Tensor* ratio, float& ratio_data) const { diff --git a/onnxruntime/core/providers/cuda/nn/dropout_impl.cu b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu index 27e1145431..a7418a17d0 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout_impl.cu +++ b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu @@ -158,9 +158,7 @@ void DropoutKernelImpl( SPECIALIZED_DROPOUT_IMPL(float) SPECIALIZED_DROPOUT_IMPL(double) SPECIALIZED_DROPOUT_IMPL(half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_DROPOUT_IMPL(nv_bfloat16) -#endif +SPECIALIZED_DROPOUT_IMPL(BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index e48ad4f4f5..9488454e5f 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -348,6 +348,8 @@ INSTANTIATE_REDUCE_SUM(half, half); INSTANTIATE_REDUCE_SUM(half, float); INSTANTIATE_REDUCE_SUM(float, float); INSTANTIATE_REDUCE_SUM(double, double); +INSTANTIATE_REDUCE_SUM(BFloat16, BFloat16); +INSTANTIATE_REDUCE_SUM(BFloat16, float); #undef INSTANTIATE_REDUCE_SUM #define INSTANTIATE_REDUCE_SQUARE_SUM(TIn, TOut) \ @@ -355,9 +357,7 @@ INSTANTIATE_REDUCE_SUM(double, double); INSTANTIATE_REDUCE_SQUARE_SUM(half, float); INSTANTIATE_REDUCE_SQUARE_SUM(float, float); INSTANTIATE_REDUCE_SQUARE_SUM(double, double); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_REDUCE_SQUARE_SUM(nv_bfloat16, float); -#endif +INSTANTIATE_REDUCE_SQUARE_SUM(BFloat16, float); #undef INSTANTIATE_REDUCE_SQUARE_SUM #define INSTANTIATE_REDUCE_L2_NORM(TIn, TOut) \ @@ -480,9 +480,7 @@ void UnaryDiv(cudaStream_t stream, const T* input, T* output, T denominator, siz INSTANTIATE_UNARY_DIV(half); INSTANTIATE_UNARY_DIV(float); INSTANTIATE_UNARY_DIV(double); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_UNARY_DIV(nv_bfloat16); -#endif +INSTANTIATE_UNARY_DIV(BFloat16); #undef INSTANTIATE_UNARY_DIV template @@ -496,9 +494,7 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i INSTANTIATE_REDUCE_MATRIX_ROWS(half); INSTANTIATE_REDUCE_MATRIX_ROWS(float); INSTANTIATE_REDUCE_MATRIX_ROWS(double); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_REDUCE_MATRIX_ROWS(nv_bfloat16); -#endif +INSTANTIATE_REDUCE_MATRIX_ROWS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_ROWS template @@ -512,9 +508,7 @@ Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output INSTANTIATE_REDUCE_MATRIX_COLUMNS(half); INSTANTIATE_REDUCE_MATRIX_COLUMNS(float); INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_REDUCE_MATRIX_COLUMNS(nv_bfloat16); -#endif +INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 88e3af5ff6..8ce2c268ae 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -835,7 +835,6 @@ SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int64_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int8_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(uint8_t) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 template <> template <> Status ReduceKernel::ComputeImpl( @@ -940,7 +939,6 @@ Status ReduceKernel::ComputeImpl return Status::OK(); } -#endif namespace ReductionOps { @@ -997,17 +995,11 @@ template std::unique_ptr ReduceCompute= 11000 -#define REGISTER_KERNEL_TYPED_BF16(name) REGISTER_KERNEL_TYPED(name, BFloat16) -#else -#define REGISTER_KERNEL_TYPED_BF16(name) -#endif - #define REGISTER_KERNEL_HFD(name) \ REGISTER_KERNEL_TYPED(name, MLFloat16) \ - REGISTER_KERNEL_TYPED_BF16(name) \ REGISTER_KERNEL_TYPED(name, float) \ - REGISTER_KERNEL_TYPED(name, double) + REGISTER_KERNEL_TYPED(name, double) \ + REGISTER_KERNEL_TYPED(name, BFloat16) #define REGISTER_KERNEL_HFD_11(name) \ REGISTER_KERNEL_TYPED_11(name, MLFloat16) \ @@ -1052,9 +1044,7 @@ REGISTER_KERNEL_TYPED_13(ReduceSum, float) REGISTER_KERNEL_TYPED_13(ReduceSum, double) REGISTER_KERNEL_TYPED_13(ReduceSum, int32_t) REGISTER_KERNEL_TYPED_13(ReduceSum, int64_t) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_KERNEL_TYPED_13(ReduceSum, BFloat16) -#endif REGISTER_KERNEL_HFD(ReduceLogSum) REGISTER_KERNEL_HFD(ReduceSumSquare) diff --git a/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h b/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h index d0eafc9d23..3aa7e477e4 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h +++ b/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h @@ -4,6 +4,7 @@ #pragma once #include +#include "core/framework/float16.h" namespace onnxruntime { namespace cuda { @@ -17,10 +18,8 @@ template <> struct AccumulationType { using type = float; }; template <> struct AccumulationType { using type = double; }; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 template <> -struct AccumulationType { using type = float; }; -#endif +struct AccumulationType { using type = float; }; template using AccumulationType_t = typename AccumulationType::type; diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index 2e6f96ded4..9e1e0fe52d 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -15,6 +15,7 @@ #include "core/providers/cuda/cuda_common.h" +using namespace onnxruntime; using namespace onnxruntime::cuda; // Generalize library calls to be use in template functions @@ -111,19 +112,18 @@ inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, } } -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, - const nv_bfloat16* alpha, - const nv_bfloat16* A, int lda, - const nv_bfloat16* B, int ldb, - const nv_bfloat16* beta, - nv_bfloat16* C, int ldc, + const BFloat16* alpha, + const BFloat16* A, int lda, + const BFloat16* B, int ldb, + const BFloat16* beta, + BFloat16* C, int ldc, const cudaDeviceProp& /*prop*/) { - float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); - float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + float h_a = alpha->ToFloat(); + float h_b = beta->ToFloat(); // accumulating in FP32 return cublasGemmEx(handle, @@ -138,7 +138,6 @@ inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); } -#endif // batched gemm inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, @@ -237,20 +236,19 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, } } -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, - const nv_bfloat16* alpha, - const nv_bfloat16* Aarray[], int lda, - const nv_bfloat16* Barray[], int ldb, - const nv_bfloat16* beta, - nv_bfloat16* Carray[], int ldc, + const BFloat16* alpha, + const BFloat16* Aarray[], int lda, + const BFloat16* Barray[], int ldb, + const BFloat16* beta, + BFloat16* Carray[], int ldc, int batch_count, const cudaDeviceProp& /*prop*/) { - float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); - float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + float h_a = alpha->ToFloat(); + float h_b = beta->ToFloat(); // accumulating in FP32 return cublasGemmBatchedEx(handle, @@ -266,7 +264,6 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, CUDA_R_32F, CUBLAS_GEMM_DEFAULT); } -#endif // strided batched gemm inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, @@ -428,23 +425,22 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, } } -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, - const nv_bfloat16* alpha, - const nv_bfloat16* A, int lda, + const BFloat16* alpha, + const BFloat16* A, int lda, long long int strideA, - const nv_bfloat16* B, int ldb, + const BFloat16* B, int ldb, long long int strideB, - const nv_bfloat16* beta, - nv_bfloat16* C, int ldc, + const BFloat16* beta, + BFloat16* C, int ldc, long long int strideC, int batch_count, const cudaDeviceProp& /*prop*/) { - float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); - float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + float h_a = alpha->ToFloat(); + float h_b = beta->ToFloat(); // accumulating in FP32 return cublasGemmStridedBatchedEx(handle, transa, @@ -459,7 +455,6 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, CUDA_R_32F, CUBLAS_GEMM_DEFAULT); } -#endif // transpose using geam inline cublasStatus_t cublasTransposeHelper(cudaStream_t, cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, const float* alpha, const float* A, int lda, const float* beta, const float* B, int ldb, float* C, int ldc) { @@ -482,7 +477,4 @@ inline cublasStatus_t cublasCopyHelper(cudaStream_t, cublasHandle_t handle, int } cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t handle, int n, const half* x, int incx, half* y, int incy); - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t handle, int n, const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy); -#endif +cublasStatus_t cublasCopyHelper(cudaStream_t stream, cublasHandle_t handle, int n, const BFloat16* x, int incx, BFloat16* y, int incy); diff --git a/onnxruntime/core/providers/cuda/tensor/cast_op.cc b/onnxruntime/core/providers/cuda/tensor/cast_op.cc index 27f1363f77..7e30a5f2c6 100644 --- a/onnxruntime/core/providers/cuda/tensor/cast_op.cc +++ b/onnxruntime/core/providers/cuda/tensor/cast_op.cc @@ -11,9 +11,7 @@ namespace cuda { const DeleteOnUnloadPtr> castOpTypeConstraints = new std::vector { DataTypeImpl::GetTensorType(), -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 DataTypeImpl::GetTensorType(), -#endif DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), @@ -81,9 +79,7 @@ Status Cast::ComputeInternal(OpKernelContext* context) const { switch (to_) { CASE(TensorProto_DataType_FLOAT16, MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 CASE(TensorProto_DataType_BFLOAT16, BFloat16) -#endif CASE(TensorProto_DataType_FLOAT, float) CASE(TensorProto_DataType_DOUBLE, double) CASE(TensorProto_DataType_INT8, int8_t) @@ -121,9 +117,7 @@ SPECIALIZE_IMPL(uint16_t) SPECIALIZE_IMPL(uint32_t) SPECIALIZE_IMPL(uint64_t) SPECIALIZE_IMPL(bool) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 SPECIALIZE_IMPL(BFloat16) -#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc index 214a2f527b..f1866ae85e 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc +++ b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc @@ -112,35 +112,12 @@ Status GatherNDBase::PrepareCompute( .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherND); -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define GATHER_ND_T_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16, bool -#else -#define GATHER_ND_T_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, bool -#endif - -#define REGISTER_KERNEL_TYPED_GATHER_ND(TIndex, ver) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatherND, \ - kOnnxDomain, \ - ver, \ - TIndex, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", GATHER_ND_T_TENSOR_TYPES) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ +#define REGISTER_KERNEL_TYPED_GATHER_ND(TIndex, ver) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatherND, kOnnxDomain, ver, TIndex, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", BuildKernelDefConstraints()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherND); // TODO: decprecate GatherND-1 after updating training models to opset-12 @@ -211,9 +188,9 @@ Status GatherND::ComputeInternal(OpKernelContext* context) const { const void* const kernel_input_data = input_tensor->DataRaw(); void* const kernel_output_data = output_tensor->MutableDataRaw(); - utils::MLTypeCallDispatcher t_disp(input_tensor->GetElementType()); - t_disp.Invoke( - Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get()); + utils::MLTypeCallDispatcher t_disp(input_tensor->GetElementType()); + t_disp.Invoke(Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data, + input_slice_offsets_buffer.get()); return Status::OK(); } diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu index 272d3a3886..ce3dfc68d1 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu @@ -113,9 +113,7 @@ SPECIALIZED_IMPL(int64_t) #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600 SPECIALIZED_IMPL(half) SPECIALIZED_IMPL(double) -#endif -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_IMPL(nv_bfloat16) +SPECIALIZED_IMPL(BFloat16) #endif } // namespace cuda diff --git a/onnxruntime/core/providers/rocm/atomic/common.cuh b/onnxruntime/core/providers/rocm/atomic/common.cuh index 303b1aaa5a..7cf4e31e9b 100644 --- a/onnxruntime/core/providers/rocm/atomic/common.cuh +++ b/onnxruntime/core/providers/rocm/atomic/common.cuh @@ -4,6 +4,7 @@ #pragma once #include #include +#include "core/framework/float16.h" typedef __half half; @@ -36,5 +37,20 @@ __device__ __forceinline__ void atomic_add(half *address, half value) { } while (assumed != old); } +__device__ __forceinline__ void atomic_add(BFloat16* address, BFloat16 value) { + unsigned int* base_address = + reinterpret_cast(reinterpret_cast(address) - (reinterpret_cast(address) & 2)); + unsigned int old = *base_address; + unsigned int assumed; + BFloat16 bsum; + do { + assumed = old; + bsum.val = reinterpret_cast(address) & 2 ? (old >> 16) : (old & 0xffff); + bsum = bsum + value; + old = reinterpret_cast(address) & 2 ? (old & 0xffff) | (bsum.val << 16) : (old & 0xffff0000) | bsum.val; + old = atomicCAS(base_address, assumed, old); + } while (assumed != old); +} + } // namespace rocm -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/cu_inc/common.cuh b/onnxruntime/core/providers/rocm/cu_inc/common.cuh index 637ba2e9bb..3ad6d9db46 100644 --- a/onnxruntime/core/providers/rocm/cu_inc/common.cuh +++ b/onnxruntime/core/providers/rocm/cu_inc/common.cuh @@ -14,6 +14,82 @@ namespace onnxruntime { namespace rocm { +/// Arithmetic for BFloat16 + +__device__ __forceinline__ BFloat16 operator+(const BFloat16& a, const BFloat16& b) { + return static_cast(a) + static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator-(const BFloat16& a, const BFloat16& b) { + return static_cast(a) - static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator*(const BFloat16& a, const BFloat16& b) { + return static_cast(a) * static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator/(const BFloat16& a, const BFloat16& b) { + return static_cast(a) / static_cast(b); +} + +__device__ __forceinline__ BFloat16 operator-(const BFloat16& a) { return -static_cast(a); } + +__device__ __forceinline__ BFloat16& operator+=(BFloat16& a, const BFloat16& b) { + a = a + b; + return a; +} + +__device__ __forceinline__ BFloat16& operator-=(BFloat16& a, const BFloat16& b) { + a = a - b; + return a; +} + +__device__ __forceinline__ BFloat16& operator*=(BFloat16& a, const BFloat16& b) { + a = a * b; + return a; +} + +__device__ __forceinline__ BFloat16& operator/=(BFloat16& a, const BFloat16& b) { + a = a / b; + return a; +} + +/// Arithmetic with floats + +__device__ __forceinline__ float operator+(BFloat16 a, float b) { return static_cast(a) + b; } +__device__ __forceinline__ float operator-(BFloat16 a, float b) { return static_cast(a) - b; } +__device__ __forceinline__ float operator*(BFloat16 a, float b) { return static_cast(a) * b; } +__device__ __forceinline__ float operator/(BFloat16 a, float b) { return static_cast(a) / b; } + +__device__ __forceinline__ float operator+(float a, BFloat16 b) { return a + static_cast(b); } +__device__ __forceinline__ float operator-(float a, BFloat16 b) { return a - static_cast(b); } +__device__ __forceinline__ float operator*(float a, BFloat16 b) { return a * static_cast(b); } +__device__ __forceinline__ float operator/(float a, BFloat16 b) { return a / static_cast(b); } + +__device__ __forceinline__ float& operator+=(float& a, const BFloat16& b) { return a += static_cast(b); } +__device__ __forceinline__ float& operator-=(float& a, const BFloat16& b) { return a -= static_cast(b); } +__device__ __forceinline__ float& operator*=(float& a, const BFloat16& b) { return a *= static_cast(b); } +__device__ __forceinline__ float& operator/=(float& a, const BFloat16& b) { return a /= static_cast(b); } + +/// Arithmetic with doubles + +__device__ __forceinline__ double operator+(BFloat16 a, double b) { return static_cast(a) + b; } +__device__ __forceinline__ double operator-(BFloat16 a, double b) { return static_cast(a) - b; } +__device__ __forceinline__ double operator*(BFloat16 a, double b) { return static_cast(a) * b; } +__device__ __forceinline__ double operator/(BFloat16 a, double b) { return static_cast(a) / b; } + +__device__ __forceinline__ double operator+(double a, BFloat16 b) { return a + static_cast(b); } +__device__ __forceinline__ double operator-(double a, BFloat16 b) { return a - static_cast(b); } +__device__ __forceinline__ double operator*(double a, BFloat16 b) { return a * static_cast(b); } +__device__ __forceinline__ double operator/(double a, BFloat16 b) { return a / static_cast(b); } + +// Overloading < and > operators + +__device__ __forceinline__ bool operator==(BFloat16& lhs, BFloat16& rhs) { return float(lhs) == float(rhs); } +__device__ __forceinline__ bool operator!=(BFloat16& lhs, BFloat16& rhs) { return float(lhs) != float(rhs); } +__device__ __forceinline__ bool operator>(BFloat16& lhs, BFloat16& rhs) { return float(lhs) > float(rhs); } +__device__ __forceinline__ bool operator<(BFloat16& lhs, BFloat16& rhs) { return float(lhs) < float(rhs); } + template __device__ __inline__ T _Ceil(T a); @@ -184,6 +260,21 @@ __device__ __inline__ double _Normcdf(double a) { return normcdf(a); } template <> __device__ __inline__ half _Normcdf(half a) { return half(normcdff((float)a)); } +template <> +__device__ __inline__ BFloat16 _Sqrt(BFloat16 a) { return sqrtf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Exp(BFloat16 a) { return expf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Log(BFloat16 a) { return logf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Tanh(BFloat16 a) { return tanhf(static_cast(a)); } + +template <> +__device__ __inline__ BFloat16 _Normcdf(BFloat16 a) { return normcdff(static_cast(a)); } + template __device__ __inline__ T _Gelu(T a) { return a * _Normcdf(a); diff --git a/onnxruntime/test/framework/data_types_test.cc b/onnxruntime/test/framework/data_types_test.cc index 15362cbabf..9d6a564404 100644 --- a/onnxruntime/test/framework/data_types_test.cc +++ b/onnxruntime/test/framework/data_types_test.cc @@ -422,7 +422,7 @@ TEST_F(DataTypeTest, BFloat16Test) { constexpr float sample = 1.0f; BFloat16 flt16(sample); auto int_rep = flt16.val; - BFloat16 flt_from_int(int_rep); + BFloat16 flt_from_int(int_rep, BFloat16::FromBits()); const double diff = std::fabs(sample - flt_from_int.ToFloat()); if (diff > FLT_EPSILON || (std::isnan(diff) && !std::isnan(sample))) { EXPECT_TRUE(false); diff --git a/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc b/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc index 841e075855..4e7ae2faa6 100644 --- a/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc +++ b/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc @@ -14,12 +14,37 @@ namespace test { std::vector MakeMLFloat16(const std::initializer_list& input) { std::vector output; std::transform(input.begin(), input.end(), std::back_inserter(output), - [](float fl) { - return MLFloat16(math::floatToHalf(fl)); - }); + [](float fl) { return MLFloat16(math::floatToHalf(fl)); }); return output; } +#ifdef USE_CUDA +void TestFloat16(const char* op_name, const std::vector& lhs_dim, + const std::initializer_list& lhs_values, const std::vector& rhs_dim, + const std::initializer_list& rhs_values, const std::vector& out_dim, + const std::initializer_list& out_values) { + { + OpTester tester(op_name, 14); + tester.AddInput("A", lhs_dim, MakeMLFloat16(lhs_values)); + tester.AddInput("B", rhs_dim, MakeMLFloat16(rhs_values)); + tester.AddOutput("C", out_dim, MakeMLFloat16(out_values)); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + + { + OpTester tester(op_name, 14); + tester.AddInput("A", lhs_dim, MakeBFloat16(lhs_values)); + tester.AddInput("B", rhs_dim, MakeBFloat16(rhs_values)); + tester.AddOutput("C", out_dim, MakeBFloat16(out_values)); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} +#endif + TEST(MathOpTest, DimWithZeroHandling) { auto run = [](OpTester& tester) { // exclude TensorRT and NNAPI as this isn't handled by those EPs @@ -89,24 +114,23 @@ TEST(MathOpTest, Add_int64) { TEST(MathOpTest, Add_float) { OpTester test("Add"); std::vector dims{3, 3}; - test.AddInput("A", dims, - {1.0f, 2.0f, -1.0f, - 0.0f, 1.5f, -100.0f, - -5.4f, 9.3f, -10000.0f}); - test.AddInput("B", dims, - {-1.0f, 4.4f, 432.3f, - 0.0f, 3.5f, 64.0f, - -5.4f, 9.3f, 10000.0f}); - test.AddOutput("C", dims, - {0.0f, 6.4f, 431.3f, - 0.0f, 5.0f, -36.0f, - -10.8f, 18.6f, 0.0f}); + std::initializer_list lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.4f, 9.3f, -10000.0f}; + std::initializer_list rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.3f, 10000.0f}; + std::initializer_list out_values{0.0f, 6.4f, 431.3f, 0.0f, 5.0f, -36.0f, -10.8f, 18.6f, 0.0f}; + test.AddInput("A", dims, lhs_values); + test.AddInput("B", dims, rhs_values); + test.AddOutput("C", dims, out_values); #if defined(OPENVINO_CONFIG_MYRIAD) || defined(OPENVINO_CONFIG_GPU_GP16) || defined(OPENVINO_CONFIG_VAD_M) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); // OpenVINO: Disabled due to accuracy mismatch for FP16 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kOpenVINOExecutionProvider}); // OpenVINO: Disabled due to accuracy mismatch for FP16 #else test.Run(); #endif + +#ifdef USE_CUDA + TestFloat16("Add", dims, lhs_values, dims, rhs_values, dims, out_values); +#endif } TEST(MathOpTest, Add_double) { @@ -131,58 +155,61 @@ TEST(MathOpTest, Add_Broadcast_Axis) { OpTester test("Add"); std::vector dims{3, 3}; - test.AddInput("A", dims, - {1.0f, 2.0f, 3.0f, - 4.0f, 5.0f, 6.0f, - 7.0f, 8.0f, 9.0f}); - test.AddInput("B", {3, 1}, - {3.0f, - 2.0f, - 1.0f}); - test.AddOutput("C", dims, - {4.0f, 5.0f, 6.0f, - 6.0f, 7.0f, 8.0f, - 8.0f, 9.0f, 10.0f}); + std::initializer_list lhs_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::initializer_list rhs_values{3.0f, 2.0f, 1.0f}; + std::initializer_list out_values{4.0f, 5.0f, 6.0f, 6.0f, 7.0f, 8.0f, 8.0f, 9.0f, 10.0f}; + test.AddInput("A", dims, lhs_values); + test.AddInput("B", {3, 1}, rhs_values); + test.AddOutput("C", dims, out_values); test.Run(OpTester::ExpectResult::kExpectSuccess, ""); + +#ifdef USE_CUDA + TestFloat16("Add", dims, lhs_values, {3, 1}, rhs_values, dims, out_values); +#endif } TEST(MathOpTest, Add_Broadcast_MultidirectionalAB) { OpTester test("Add"); - - test.AddInput("A", {3, 1}, - {3.0f, - 2.0f, - 1.0f}); - test.AddInput("B", {3}, - {1.0f, 2.0f, 3.0f}); - test.AddOutput("C", {3, 3}, - {4.0f, 5.0f, 6.0f, - 3.0f, 4.0f, 5.0f, - 2.0f, 3.0f, 4.0f}); -#if defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_MYRIAD) || defined(OPENVINO_CONFIG_VAD_M) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: disabled temporarily due to accurarcy issues + std::initializer_list lhs_values{3.0f, 2.0f, 1.0f}; + std::initializer_list rhs_values{1.0f, 2.0f, 3.0f}; + std::initializer_list out_values{4.0f, 5.0f, 6.0f, 3.0f, 4.0f, 5.0f, 2.0f, 3.0f, 4.0f}; + test.AddInput("A", {3, 1}, lhs_values); + test.AddInput("B", {3}, rhs_values); + test.AddOutput("C", {3, 3}, out_values); +#if defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_MYRIAD) || \ + defined(OPENVINO_CONFIG_VAD_M) + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, + kOpenVINOExecutionProvider}); // OpenVINO: disabled temporarily due to accurarcy issues #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: got C with shape [3, 1] + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider}); // TensorRT: got C with shape [3, 1] +#endif + +#ifdef USE_CUDA + TestFloat16("Add", {3, 1}, lhs_values, {3}, rhs_values, {3, 3}, out_values); #endif } TEST(MathOpTest, Add_Broadcast_MultidirectionalBA) { OpTester test("Add"); - - test.AddInput("A", {3}, - {1.0f, 2.0f, 3.0f}); - test.AddInput("B", {3, 1}, - {3.0f, - 2.0f, - 1.0f}); - test.AddOutput("C", {3, 3}, - {4.0f, 5.0f, 6.0f, - 3.0f, 4.0f, 5.0f, - 2.0f, 3.0f, 4.0f}); + std::initializer_list lhs_values{1.0f, 2.0f, 3.0f}; + std::initializer_list rhs_values{3.0f, 2.0f, 1.0f}; + std::initializer_list out_values{4.0f, 5.0f, 6.0f, 3.0f, 4.0f, 5.0f, 2.0f, 3.0f, 4.0f}; + test.AddInput("A", {3}, lhs_values); + test.AddInput("B", {3, 1}, rhs_values); + test.AddOutput("C", {3, 3}, out_values); #if defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_GPU_FP16) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: disabled temporarily due to accurarcy issues + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, + kOpenVINOExecutionProvider}); // OpenVINO: disabled temporarily due to accurarcy issues #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: got C with shape [3, 1] + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider}); // TensorRT: got C with shape [3, 1] +#endif + +#ifdef USE_CUDA + TestFloat16("Add", {3}, lhs_values, {3, 1}, rhs_values, {3, 3}, out_values); #endif } @@ -364,23 +391,22 @@ TEST(MathOpTest, Sub_int64) { TEST(MathOpTest, Sub) { OpTester test("Sub"); std::vector dims{3, 3}; - test.AddInput("A", dims, - {1.0f, 2.0f, -1.0f, - 0.0f, 1.5f, -100.0f, - -5.4f, 9.3f, -10000.0f}); - test.AddInput("B", dims, - {-1.0f, 4.4f, 432.3f, - 0.0f, 3.5f, 64.0f, - -5.4f, 9.3f, 10000.0f}); - test.AddOutput("C", dims, - {2.0f, -2.4f, -433.3f, - 0.0f, -2.0f, -164.0f, - 0.0f, 0.0f, -20000.0f}); + std::initializer_list lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.4f, 9.3f, -10000.0f}; + std::initializer_list rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.3f, 10000.0f}; + std::initializer_list out_values{2.0f, -2.4f, -433.3f, 0.0f, -2.0f, -164.0f, 0.0f, 0.0f, -20000.0f}; + test.AddInput("A", dims, lhs_values); + test.AddInput("B", dims, rhs_values); + test.AddOutput("C", dims, out_values); #if defined(OPENVINO_CONFIG_MYRIAD) || defined(OPENVINO_CONFIG_VAD_M) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); // OpenVINO EP: Disabled due to accuracy mismatch for FP16 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kOpenVINOExecutionProvider}); // OpenVINO EP: Disabled due to accuracy mismatch for FP16 #else test.Run(); #endif + +#ifdef USE_CUDA + TestFloat16("Sub", dims, lhs_values, dims, rhs_values, dims, out_values); +#endif } TEST(MathOpTest, Sub_Broadcast_Scalar) { @@ -422,24 +448,23 @@ TEST(MathOpTest, Mul_int64) { TEST(MathOpTest, Mul) { OpTester test("Mul"); std::vector dims{3, 3}; - test.AddInput("A", dims, - {1.0f, 2.0f, -1.0f, - 0.0f, 1.5f, -100.0f, -5.4f, - 9.30f, -10000.0f}); - test.AddInput("B", dims, - {-1.0f, 4.4f, 432.3f, - 0.0f, 3.5f, 64.0f, -5.4f, - 9.30f, 10000.0f}); - test.AddOutput("C", dims, - {-1.0f, 8.8f, -432.3f, - 0.0f, 5.25f, -6400.0f, - 29.16f, 86.49f, -100000000.0f}); + std::initializer_list lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.0f, 9.30f, -10000.0f}; + std::initializer_list rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.0f, 10000.0f}; + std::initializer_list out_values{-1.0f, 8.8f, -432.3f, 0.0f, 5.25f, -6400.0f, 27.0f, 83.7f, -100000000.0f}; + test.AddInput("A", dims, lhs_values); + test.AddInput("B", dims, rhs_values); + test.AddOutput("C", dims, out_values); #if defined(OPENVINO_CONFIG_MYRIAD) || defined(OPENVINO_CONFIG_VAD_M) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); // OpenVINO: Disabled due to accuracy issues for MYRIAD FP16 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kOpenVINOExecutionProvider}); // OpenVINO: Disabled due to accuracy issues for MYRIAD FP16 #else test.Run(); #endif + +#ifdef USE_CUDA + TestFloat16("Mul", dims, lhs_values, dims, rhs_values, dims, out_values); +#endif } TEST(MathOpTest, Div_int32) { @@ -463,20 +488,22 @@ TEST(MathOpTest, Div_int64) { TEST(MathOpTest, Div) { OpTester test("Div"); std::vector dims{2, 3}; - test.AddInput("A", dims, - {1000.0f, 1.0f, 6.0f, - 0.0f, -10.0f, -1.0f}); - test.AddInput("B", dims, - {1000.0f, 2.0f, 3.0f, - 1.0f, -1.0f, 4.0f}); - test.AddOutput("C", dims, - {1.0f, 0.5f, 2.0f, - 0.0f, 10.0f, -0.25f}); + std::initializer_list lhs_values{1000.0f, 1.0f, 6.0f, 0.0f, -10.0f, -1.0f}; + std::initializer_list rhs_values{1000.0f, 2.0f, 3.0f, 1.0f, -1.0f, 4.0f}; + std::initializer_list out_values{1.0f, 0.5f, 2.0f, 0.0f, 10.0f, -0.25f}; + test.AddInput("A", dims, lhs_values); + test.AddInput("B", dims, rhs_values); + test.AddOutput("C", dims, out_values); #if defined(OPENVINO_CONFIG_MYRIAD) || defined(OPENVINO_CONFIG_VAD_M) - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); // OpenVINO EP: Hardware limitation + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kOpenVINOExecutionProvider}); // OpenVINO EP: Hardware limitation #else test.Run(); #endif + +#ifdef USE_CUDA + TestFloat16("Div", dims, lhs_values, dims, rhs_values, dims, out_values); +#endif } TEST(MathOpTest, Abs) { diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index 6c1573b2a0..5bdd5c3dea 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -4,6 +4,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include "test/common/cuda_op_test_utils.h" +#include "test/util/include/default_providers.h" namespace onnxruntime { namespace test { @@ -85,6 +86,28 @@ TEST(GemmOpTest, GemmNoTrans_f16) { } #endif +#ifdef USE_CUDA +TEST(GemmOpTest, GemmNoTrans_bfloat16) { + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } + OpTester test("Gemm", 14); + test.AddAttribute("transA", (int64_t)0); + test.AddAttribute("transB", (int64_t)0); + test.AddAttribute("alpha", 1.0f); + test.AddAttribute("beta", 1.0f); + test.AddInput("A", {2, 4}, MakeBFloat16({1.0f, 2.0f, 3.0f, 4.0f, -1.0f, -2.0f, -3.0f, -4.0f})); + test.AddInput("B", {4, 3}, MakeBFloat16({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f})); + test.AddInput("C", {2, 3}, MakeBFloat16({1.f, 1.f, 1.f, 1.f, 1.f, 1.f})); + test.AddOutput("Y", {2, 3}, MakeBFloat16({11.0f, 11.0f, 11.0f, -9.0f, -9.0f, -9.0f})); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + template void TestGemmBroadcast() { auto run_test = [](bool b_is_initializer, bool c_is_initializer) { diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index cc5a2b706f..8660f65060 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -9,6 +9,7 @@ #include "test/providers/provider_test_utils.h" #include "test/providers/cpu/reduction/reduction_test_cases.h" #include "core/providers/cpu/reduction/reduction_ops.h" +#include "test/util/include/default_providers.h" namespace onnxruntime { namespace test { @@ -1490,6 +1491,20 @@ TEST(ReductionOpTest, ReduceSum_half_bert) { // Add more UTs for half as needed #endif +#ifdef USE_CUDA +TEST(ReductionOpTest, ReduceSumBFloat16) { + OpTester test("ReduceSum", 14); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("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("axes", {2}, std::vector{0, 1}); + test.AddOutput("reduced", {2}, MakeBFloat16({36.0f, 42.0f})); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(ReductionOpTest, ReduceSum_apex_reduction) { OpTester test("ReduceSum"); test.AddAttribute("keepdims", (int64_t)0); diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index 4b666f621c..f74768aee7 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -329,6 +329,9 @@ struct TensorCheck { /// 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; +#endif for (int i = 0; i < size; ++i) { if (std::isnan(f_expected[i])) { EXPECT_TRUE(std::isnan(f_expected[i])) << "Expected NaN. i:" << i << ", provider_type: " << provider_type; diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index c31d37157d..65eb7c97e7 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -1088,6 +1088,12 @@ inline std::vector FloatsToMLFloat16s(const std::vector& f) { return m; } +inline std::vector MakeBFloat16(const std::initializer_list& input) { + std::vector output; + std::transform(input.begin(), input.end(), std::back_inserter(output), [](float f) { return BFloat16(f); }); + return output; +} + inline CheckParams MakeCheckParams(const OpTester::Data& d) { return CheckParams{d.sort_output_, d.absolute_error_, d.relative_error_}; } diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index d732f0b4ae..013a35d0af 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -2814,7 +2814,7 @@ Return true if all elements are true and false otherwise. .Output(0, "dX", "Gradient of the input.", "T") .TypeConstraint( "T", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); diff --git a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc index 30c174dd77..f76ff9aa1f 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc @@ -10,19 +10,13 @@ namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double, BFloat16 -#else -#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double -#endif - ONNX_OPERATOR_KERNEL_EX( BiasGeluGrad_dX, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T", BuildKernelDefConstraints()) .MayInplace(0, 0), BiasGeluGrad_dX); @@ -32,7 +26,7 @@ ONNX_OPERATOR_KERNEL_EX( 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T", BuildKernelDefConstraints()) .MayInplace(0, 0), BiasGeluGrad_dX); @@ -76,7 +70,7 @@ Status BiasGeluGrad_dX::ComputeInternal(OpKernelContext* co const auto input_size = input_shape.Size(), bias_size = bias_shape.Size(); - utils::MLTypeCallDispatcher dispatcher{X->GetElementType()}; + utils::MLTypeCallDispatcher dispatcher{X->GetElementType()}; dispatcher.Invoke(Stream(), input_size, bias_size, *dY, *X, *B, *dX); return Status::OK(); diff --git a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu index d6fae84ca4..67c95872a9 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu @@ -88,10 +88,8 @@ SPECIALIZED_BIAS_GELU_GRAD_IMPL(half, gelu_computation_mode::Approximation); SPECIALIZED_BIAS_GELU_GRAD_IMPL(float, gelu_computation_mode::Approximation); SPECIALIZED_BIAS_GELU_GRAD_IMPL(double, gelu_computation_mode::Approximation); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_BIAS_GELU_GRAD_IMPL(nv_bfloat16, gelu_computation_mode::Default); -SPECIALIZED_BIAS_GELU_GRAD_IMPL(nv_bfloat16, gelu_computation_mode::Approximation); -#endif +SPECIALIZED_BIAS_GELU_GRAD_IMPL(BFloat16, gelu_computation_mode::Default); +SPECIALIZED_BIAS_GELU_GRAD_IMPL(BFloat16, gelu_computation_mode::Approximation); #undef SPECIALIZED_BIAS_GELU_GRAD_IMPL diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 8d0882f2ee..7e01dc923e 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -57,12 +57,16 @@ class ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider class ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, 12, float, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, int64_t, SoftmaxCrossEntropyLoss); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, SoftmaxGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, SoftmaxGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SoftmaxGrad); @@ -167,7 +171,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_float, BatchNormInternal); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_float, BatchNormInternal); -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 // Adam class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float_BFloat16, AdamOptimizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer); @@ -200,7 +203,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, ReduceAllL2); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2); -#endif #if defined(ORT_USE_NCCL) || defined(USE_MPI) // P2P communication operators. @@ -299,12 +301,16 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -385,7 +391,6 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 // Adam BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -418,7 +423,6 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, -#endif // P2P communication operators. #if defined(ORT_USE_NCCL) || defined(USE_MPI) diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc index 1302481b12..6d8a522210 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc @@ -308,8 +308,10 @@ INSTANTIATE_VERSIONED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, float, int64_t, kO INSTANTIATE_VERSIONED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, MLFloat16, int64_t, kOnnxDomain, 12, 12) INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, float, int64_t, kOnnxDomain, 13) INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, MLFloat16, int64_t, kOnnxDomain, 13) +INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, BFloat16, int64_t, kOnnxDomain, 13) INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, float, int64_t, kMSDomain, 1) INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, kMSDomain, 1) +INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, BFloat16, int64_t, kMSDomain, 1) #define REGISTER_KERNEL_INTERNAL_TYPED(OpName, ClassName, T, Tin, CpuInputIndex) \ ONNX_OPERATOR_TWO_TYPED_KERNEL_EX(OpName, kMSDomain, 1, T, Tin, kCudaExecutionProvider, \ @@ -322,8 +324,10 @@ INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, kMSD REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, float, int64_t, 3) REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, MLFloat16, int64_t, 3) +REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, BFloat16, int64_t, 3) REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, float, int64_t, 4) REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, 4) +REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, BFloat16, int64_t, 4) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu index 12b2e0f66e..c40a656f08 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu @@ -106,6 +106,7 @@ void SoftmaxCrossEntropyLossImpl( INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(float, float, int32_t) INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(float, float, int64_t) INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(half, float, int64_t) +INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(BFloat16, float, int64_t) template __global__ void _WeightedSoftmaxCrossEntropyLossGrad( @@ -221,6 +222,7 @@ void SoftmaxCrossEntropyLossGradImpl( INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(float, float, int32_t) INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(float, float, int64_t) INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(half, float, int64_t) +INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(BFloat16, float, int64_t) #define INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(T, Tin) \ template void ComputeWeightsSoftmaxCrossEntropyImpl( \ @@ -235,6 +237,7 @@ INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(half, float, int64_t) INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(float, int32_t) INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(float, int64_t) INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(half, int64_t) +INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(BFloat16, int64_t) } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc index af9dd95ed1..636cb2d8d5 100644 --- a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc +++ b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc @@ -8,27 +8,13 @@ using namespace onnxruntime::common; namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#else -#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() -#endif - -#define REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(SrcT) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - MixedPrecisionScale, \ - kMSDomain, \ - 1, \ - SrcT, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("SrcT", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("ScaleT", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("DstT", ALL_IEEE_FLOAT_TENSOR_TYPES), \ +#define REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(SrcT) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MixedPrecisionScale, kMSDomain, 1, SrcT, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("SrcT", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("ScaleT", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("DstT", BuildKernelDefConstraints()), \ MixedPrecisionScale); Status BytesPerElement(ONNX_NAMESPACE::TensorProto_DataType to, size_t& bytes_per_elem) { @@ -42,11 +28,9 @@ Status BytesPerElement(ONNX_NAMESPACE::TensorProto_DataType to, size_t& bytes_pe case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: bytes_per_elem = sizeof(MLFloat16); break; -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: bytes_per_elem = sizeof(BFloat16); break; -#endif default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unexpected 'to' argument value: ", to); } @@ -118,9 +102,7 @@ Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) cons switch (to_) { CASE(TensorProto_DataType_FLOAT16, MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 CASE(TensorProto_DataType_BFLOAT16, BFloat16) -#endif CASE(TensorProto_DataType_FLOAT, float) default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unexpected 'to' argument value: ", to_); @@ -132,14 +114,11 @@ Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) cons REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(MLFloat16) REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(float) +REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16) template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16) template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; -#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu index b86641c091..cc38593376 100644 --- a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu +++ b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu @@ -50,14 +50,11 @@ SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, half) SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, float) SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, half) SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, float) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, nv_bfloat16) -SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, float) -SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, nv_bfloat16) -SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, half) -SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, nv_bfloat16) -#endif +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(BFloat16, BFloat16) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(BFloat16, float) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, BFloat16) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(BFloat16, half) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc index 7ce8f0184a..4d4c4bdc76 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc @@ -62,33 +62,27 @@ Status SoftMaxGradComputeHelper( return Status::OK(); } -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 // cudnnSoftmaxForward/Backward doesn't support BFloat16. -#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \ - template <> \ - Status SoftMaxGradComputeHelper( \ - cudaStream_t stream, \ - const BFloat16* dY, \ - const TensorShape& input_shape, \ - const BFloat16* Y, \ - BFloat16* dX, \ - cudnnHandle_t, \ - int64_t axis) { \ - typedef typename ToCudaType::MappedType CudaT; \ - 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(dY); \ - auto Y_data = reinterpret_cast(Y); \ - auto dX_data = reinterpret_cast(dX); \ - dispatch_softmax_backward, is_log_softmax>( \ - stream, dX_data, dY_data, Y_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ - return Status::OK(); \ +#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \ + template <> \ + Status SoftMaxGradComputeHelper(cudaStream_t stream, const BFloat16* dY, \ + const TensorShape& input_shape, const BFloat16* Y, \ + BFloat16* dX, cudnnHandle_t, int64_t axis) { \ + typedef typename ToCudaType::MappedType CudaT; \ + 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(dY); \ + auto Y_data = reinterpret_cast(Y); \ + auto dX_data = reinterpret_cast(dX); \ + dispatch_softmax_backward, is_log_softmax>( \ + stream, dX_data, dY_data, Y_data, gsl::narrow_cast(D), gsl::narrow_cast(D), \ + gsl::narrow_cast(N)); \ + return Status::OK(); \ } SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true) - SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false) -#endif +SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false) #define REGISTER_GRADIENT_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ @@ -211,9 +205,7 @@ SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true) SPECIALIZED_GRADIENT(float) SPECIALIZED_GRADIENT(double) SPECIALIZED_GRADIENT(MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 SPECIALIZED_GRADIENT(BFloat16) -#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu index 5abcdec918..17da25cc9d 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu @@ -200,10 +200,7 @@ template void dispatch_softmax_backward(cudaStre SPECIALIZED_SOFTMAX_GRAD_IMPL(float, float, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(half, half, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(double, double, double) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_SOFTMAX_GRAD_IMPL(nv_bfloat16, nv_bfloat16, float) -#endif +SPECIALIZED_SOFTMAX_GRAD_IMPL(BFloat16, BFloat16, float) } } \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc index 3ce65f6a4d..3bbe9f5acc 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc @@ -9,20 +9,13 @@ namespace onnxruntime { namespace cuda { -#define REGISTER_GRADIENT_KERNEL(OpName) \ - ONNX_OPERATOR_KERNEL_EX( \ - OpName, \ - kMSDomain, \ - 1, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \ - .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 2), \ - DropoutGrad); - -REGISTER_GRADIENT_KERNEL(DropoutGrad) +ONNX_OPERATOR_KERNEL_EX(DropoutGrad, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T1", BuildKernelDefConstraints()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 2), + DropoutGrad); template struct DropoutGradComputeImpl { @@ -53,13 +46,13 @@ Status DropoutGrad::ComputeInternal(OpKernelContext* context) const { float ratio_data = default_ratio_; auto ratio = context->Input(2); if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); t_disp.Invoke(ratio, ratio_data); } auto dX = context->Output(0, shape); - utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); + utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); t_disp.Invoke(Stream(), N, *dY, mask_data, ratio_data, *dX); return Status::OK(); diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu index bcd89738e5..e01c7fbf74 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu @@ -123,9 +123,7 @@ void DropoutGradientKernelImpl( SPECIALIZED_DROPOUT_GRAD_IMPL(float) SPECIALIZED_DROPOUT_GRAD_IMPL(double) SPECIALIZED_DROPOUT_GRAD_IMPL(half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_DROPOUT_GRAD_IMPL(nv_bfloat16) -#endif +SPECIALIZED_DROPOUT_GRAD_IMPL(BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc index 5b0bcb684f..15bac8f7d7 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc @@ -44,9 +44,7 @@ namespace cuda { REGISTER_GRADIENT_KERNEL_TYPED(float, float) REGISTER_GRADIENT_KERNEL_TYPED(double, double) REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_GRADIENT_KERNEL_TYPED(BFloat16, float) -#endif template LayerNormGrad::LayerNormGrad(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu index 574e57ea62..dddbd12de4 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu @@ -606,10 +606,8 @@ LAYERNORMGRAD_IMPL(half, float, true) LAYERNORMGRAD_IMPL(float, float, false) LAYERNORMGRAD_IMPL(double, double, false) LAYERNORMGRAD_IMPL(half, float, false) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -LAYERNORMGRAD_IMPL(nv_bfloat16, float, true) -LAYERNORMGRAD_IMPL(nv_bfloat16, float, false) -#endif +LAYERNORMGRAD_IMPL(BFloat16, float, true) +LAYERNORMGRAD_IMPL(BFloat16, float, false) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc index cc96e8bdfd..d8e424a439 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc @@ -46,8 +46,6 @@ REGISTER_ADAM_KERNEL_TYPED(MLFloat16, int64_t, float, MLFloat16, MLFloat16, MLFl REGISTER_ADAM_KERNEL_TYPED(MLFloat16, int64_t, float, MLFloat16, MLFloat16, float, MLFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, MLFloat16, MLFloat16, MLFloat16, MLFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, MLFloat16, MLFloat16, float, MLFloat16) - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, BFloat16) REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, float, float, BFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, float, float, BFloat16) @@ -57,7 +55,6 @@ REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, BFloat16, BFloat1 REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, BFloat16, float, BFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, BFloat16, BFloat16, BFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, BFloat16, float, BFloat16) -#endif template Status AdamOptimizer::ComputeInternal(OpKernelContext* ctx) const { diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu index f3205d44b1..d72b9f924e 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu @@ -263,18 +263,15 @@ SPECIALIZED_AdamOptimizerImpl(half, int64_t, float, half, half, half, half) SPECIALIZED_AdamOptimizerImpl(half, int64_t, float, half, half, float, half) SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, half, half, half, half) SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, half, half, float, half) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, float, float, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, float, float, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, float, float, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, nv_bfloat16, nv_bfloat16, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, nv_bfloat16, float, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, nv_bfloat16, float, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16) -SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, nv_bfloat16, float, nv_bfloat16) -#endif +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, float, float, BFloat16) +SPECIALIZED_AdamOptimizerImpl(BFloat16, int64_t, float, BFloat16, float, float, BFloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, BFloat16, float, float, BFloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, BFloat16, BFloat16, BFloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, BFloat16, float, BFloat16) +SPECIALIZED_AdamOptimizerImpl(BFloat16, int64_t, float, BFloat16, BFloat16, BFloat16, BFloat16) +SPECIALIZED_AdamOptimizerImpl(BFloat16, int64_t, float, BFloat16, BFloat16, float, BFloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, BFloat16, BFloat16, BFloat16, BFloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, BFloat16, BFloat16, float, BFloat16) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc index ab116b20d7..5a7815ff60 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc @@ -28,11 +28,9 @@ REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, float) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, MLFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(MLFloat16, MLFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(MLFloat16, float) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, BFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(BFloat16, BFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(BFloat16, float) -#endif template Status ZeroGradient::ComputeInternal(OpKernelContext* ctx) const { diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu index 1d83bb166c..48295f5a0c 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu @@ -35,23 +35,17 @@ void InPlaceAccumulatorImpl( N); } -#define SPECIALIZED_IMPL_InPlaceAccumulator(T, T_GRAD) \ - template void InPlaceAccumulatorImpl( \ - cudaStream_t stream, \ - const T* gradient_buffer, \ - const T_GRAD* gradient, \ - T* accumulated_gradient, \ - size_t count); +#define SPECIALIZED_IMPL_InPlaceAccumulator(T, T_GRAD) \ + template void InPlaceAccumulatorImpl(cudaStream_t stream, const T* gradient_buffer, const T_GRAD* gradient, \ + T* accumulated_gradient, size_t count); SPECIALIZED_IMPL_InPlaceAccumulator(float, float) SPECIALIZED_IMPL_InPlaceAccumulator(float, half) SPECIALIZED_IMPL_InPlaceAccumulator(half, half) SPECIALIZED_IMPL_InPlaceAccumulator(half, float) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_IMPL_InPlaceAccumulator(float, nv_bfloat16) -SPECIALIZED_IMPL_InPlaceAccumulator(nv_bfloat16, nv_bfloat16) -SPECIALIZED_IMPL_InPlaceAccumulator(nv_bfloat16, float) -#endif +SPECIALIZED_IMPL_InPlaceAccumulator(float, BFloat16) +SPECIALIZED_IMPL_InPlaceAccumulator(BFloat16, BFloat16) +SPECIALIZED_IMPL_InPlaceAccumulator(BFloat16, float) } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc index 59395bc623..17cff8db25 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc @@ -74,8 +74,6 @@ REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, MLFloat16, MLFloat16, ML REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, MLFloat16, float, MLFloat16) REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, float, MLFloat16, MLFloat16) REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, float, float, MLFloat16) - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_LAMB_KERNEL_TYPED(float, float, BFloat16, float, BFloat16, BFloat16) REGISTER_LAMB_KERNEL_TYPED(float, float, BFloat16, float, float, BFloat16) REGISTER_LAMB_KERNEL_TYPED(float, float, float, float, float, BFloat16) @@ -84,7 +82,6 @@ REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, BFloat16, BFloat16, BFloat REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, BFloat16, float, BFloat16) REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, float, BFloat16, BFloat16) REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, float, float, BFloat16) -#endif void check_inputs_and_outputs( const Tensor* w, diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu index ed289497f5..c7c9936ed1 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu @@ -177,13 +177,10 @@ SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, half, half) SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, half, float) SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, float, half) SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, float, float) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, nv_bfloat16, nv_bfloat16) -SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, nv_bfloat16, float) -SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, float, nv_bfloat16) -SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, float, float) -#endif +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, BFloat16, BFloat16, BFloat16) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, BFloat16, BFloat16, float) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, BFloat16, float, BFloat16) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, BFloat16, float, float) template __device__ __forceinline__ void _LambUpdateRule( @@ -305,13 +302,10 @@ INSTANTIATE_LAMB_UPDATE(float, float, float, half) INSTANTIATE_LAMB_UPDATE(double, double, double, half) INSTANTIATE_LAMB_UPDATE(half, float, half, half) INSTANTIATE_LAMB_UPDATE(float, float, half, half) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_LAMB_UPDATE(float, float, float, nv_bfloat16) -INSTANTIATE_LAMB_UPDATE(double, double, double, nv_bfloat16) -INSTANTIATE_LAMB_UPDATE(nv_bfloat16, float, nv_bfloat16, nv_bfloat16) -INSTANTIATE_LAMB_UPDATE(float, float, nv_bfloat16, nv_bfloat16) -#endif +INSTANTIATE_LAMB_UPDATE(float, float, float, BFloat16) +INSTANTIATE_LAMB_UPDATE(double, double, double, BFloat16) +INSTANTIATE_LAMB_UPDATE(BFloat16, float, BFloat16, BFloat16) +INSTANTIATE_LAMB_UPDATE(float, float, BFloat16, BFloat16) template __global__ void LambMultiTensorComputeDirectionImpl( @@ -406,13 +400,10 @@ INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, half, half) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, half, float) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, float, half) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, float, float) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, nv_bfloat16, nv_bfloat16) -INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, nv_bfloat16, float) -INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, float, nv_bfloat16) -INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, float, float) -#endif +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, BFloat16, BFloat16, BFloat16) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, BFloat16, BFloat16, float) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, BFloat16, float, BFloat16) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, BFloat16, float, float) template __global__ void LambMultiTensorUpdateImpl( @@ -477,13 +468,10 @@ INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, float, half) INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(double, double, double, half) INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(half, float, half, half) INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, half, half) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, float, nv_bfloat16) -INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(double, double, double, nv_bfloat16) -INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(nv_bfloat16, float, nv_bfloat16, nv_bfloat16) -INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, nv_bfloat16, nv_bfloat16) -#endif +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, float, BFloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(double, double, double, BFloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(BFloat16, float, BFloat16, BFloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, BFloat16, BFloat16) // w_buffer[i], d_buffer[i] is used to store the squared sum of all elements processed by the i-th block. // sync_range_and_lock is used for a well ordered reduction over blocks spanning the same tensor @@ -655,12 +643,9 @@ INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(double, double, double, double, INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, half, float, half, float) INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, half, float, float, float) INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(half, half, half, half, float) - -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, nv_bfloat16, float, nv_bfloat16, float) -INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, nv_bfloat16, float, float, float) -INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16, float) -#endif +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, BFloat16, float, BFloat16, float) +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, BFloat16, float, float, float) +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(BFloat16, BFloat16, BFloat16, BFloat16, float) } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc index 71d04393e6..529206ad16 100644 --- a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc +++ b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc @@ -103,11 +103,9 @@ REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, float) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, MLFloat16, float) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, MLFloat16) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, MLFloat16, MLFloat16) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, float) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, BFloat16) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, BFloat16) -#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu index dcb21e4b8c..7ac6a2cdbd 100644 --- a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu +++ b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu @@ -24,10 +24,8 @@ void ScalarSqrt(cudaStream_t stream, Tin* input, Tout* output) { template void ScalarSqrt(cudaStream_t stream, float* input, float* output); template void ScalarSqrt(cudaStream_t stream, half* input, half* output); template void ScalarSqrt(cudaStream_t stream, float* input, half* output); -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -template void ScalarSqrt(cudaStream_t stream, nv_bfloat16* input, nv_bfloat16* output); -template void ScalarSqrt(cudaStream_t stream, float* input, nv_bfloat16* output); -#endif +template void ScalarSqrt(cudaStream_t stream, BFloat16* input, BFloat16* output); +template void ScalarSqrt(cudaStream_t stream, float* input, BFloat16* output); template __launch_bounds__(ChunkGroup<1>::thread_count_per_block) @@ -114,11 +112,9 @@ INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, float) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, float) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, half) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(nv_bfloat16, float) -INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, nv_bfloat16) -INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(nv_bfloat16, nv_bfloat16) -#endif +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(BFloat16, float) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, BFloat16) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(BFloat16, BFloat16) } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc index 865443d207..d6932c8c09 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc @@ -10,29 +10,14 @@ namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#else -#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() -#endif - -ONNX_OPERATOR_KERNEL_EX( - GatherGrad, - kMSDomain, - 1, - kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .InputMemoryType(OrtMemTypeCPUInput, 0) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) - .TypeConstraint("Tind", std::vector{ - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}), - GatherGrad); +ONNX_OPERATOR_KERNEL_EX(GatherGrad, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 0) + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("Tind", std::vector{DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + GatherGrad); namespace { template @@ -96,11 +81,9 @@ Status DispatchToGatherGradImpl( } else if (utils::IsPrimitiveDataType(t_data_type)) { return DispatchToGatherGradImplByTindex( stream, tindex_data_type, allocator, num_gathered_per_index, gather_dimension_size, num_batches, dY, gathered_indices, dX); -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 } else if (utils::IsPrimitiveDataType(t_data_type)) { return DispatchToGatherGradImplByTindex( stream, tindex_data_type, allocator, num_gathered_per_index, gather_dimension_size, num_batches, dY, gathered_indices, dX); -#endif } return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GatherGrad unsupported T type: ", t_data_type); diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu index 930d71bf43..539f29630c 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu @@ -565,9 +565,7 @@ void GatherGradImpl( SPECIALIZED_WITH_IDX(float) SPECIALIZED_WITH_IDX(half) -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_WITH_IDX(nv_bfloat16) -#endif +SPECIALIZED_WITH_IDX(BFloat16) #undef SPECIALIZED_WITH_IDX #undef SPECIALIZED diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc index 285c7e66f1..b872916cdf 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc @@ -8,26 +8,14 @@ namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_TENSOR_TYPES \ - { DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType() } -#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double, BFloat16 -#else -#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() -#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double -#endif - -#define REGISTER_KERNEL_TYPED_GATHER_ND_GRAD(TIndex) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatherNDGrad, \ - kMSDomain, \ - 1, \ - TIndex, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES).TypeConstraint("Tind", DataTypeImpl::GetTensorType()).TypeConstraint("T1", DataTypeImpl::GetTensorType()).InputMemoryType(OrtMemTypeCPUInput, 0), \ +#define REGISTER_KERNEL_TYPED_GATHER_ND_GRAD(TIndex) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatherNDGrad, kMSDomain, 1, TIndex, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", BuildKernelDefConstraints()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 0), \ GatherNDGrad); REGISTER_KERNEL_TYPED_GATHER_ND_GRAD(int64_t) @@ -94,7 +82,7 @@ Status GatherNDGrad::ComputeInternal(OpKernelContext* context) const { const void* const kernel_input_data = update_tensor->DataRaw(); void* const kernel_output_data = output_tensor->MutableDataRaw(); - utils::MLTypeCallDispatcher t_disp(update_tensor->GetElementType()); + utils::MLTypeCallDispatcher t_disp(update_tensor->GetElementType()); t_disp.Invoke( Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get()); diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu index 8eddbd21a8..df1e8ebe17 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu @@ -44,9 +44,7 @@ SPECIALIZED_GRAD_IMPL(float); #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600 SPECIALIZED_GRAD_IMPL(half); SPECIALIZED_GRAD_IMPL(double); -#endif -#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -SPECIALIZED_GRAD_IMPL(nv_bfloat16); +SPECIALIZED_GRAD_IMPL(BFloat16); #endif } // namespace cuda