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

CUDA BFloat16 Refactor (#10085)

Browse source
This commit is contained in:
Vincent Wang 2022-01-14 19:38:56 +08:00 committed by GitHub
parent e38e51ea8e
commit 44e2db9397
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
80 changed files with 710 additions and 767 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/ContribOperators.md
View file

@ -486,7 +486,7 @@ This version of the operator has been available since version 1 of the 'com.micr
#### Type Constraints
<dl>
<dt><tt>T</tt> : tensor(float16), tensor(float), tensor(double)</dt>
<dt><tt>T</tt> : tensor(float16), tensor(float), tensor(double), tensor(bfloat16)</dt>
<dd>Constrain input and output types to float tensors.</dd>
</dl>

2
docs/OperatorKernels.md
View file

@ -729,7 +729,7 @@ Do not modify directly.*
|**Operator Domain:** *com.microsoft*||||
|Attention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask_index:**M**<br> *in* past:**T**<br> *in* extra_add:**T**<br> *out* output:**T**<br> *out* present:**T**|1+|**T** = tensor(float), tensor(float16)|
|BiasDropout|*in* data:**T**<br> *in* bias:**T**<br> *in* residual:**T**<br> *in* ratio:**T1**<br> *in* training_mode:**T2**<br> *out* output:**T**<br> *out* mask:**T2**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)<br/> **T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)<br/> **T2** = tensor(bool)|
|BiasGelu|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|BiasGelu|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|BiasSoftmax|*in* data:**T**<br> *in* bias:**T**<br> *out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|ComplexMul|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(float), tensor(float16)|
|ComplexMulConj|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|1+|**T** = tensor(float), tensor(float16)|

87
include/onnxruntime/core/framework/float16.h
View file

@ -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<char*>(&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<const __nv_bfloat16*>(&val));
#else
float result;
char* const first = reinterpret_cast<char*>(&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<const unsigned short*>(&value); }
explicit ORT_HOST_DEVICE operator __nv_bfloat16() const { return *reinterpret_cast<const __nv_bfloat16*>(&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;
}
}
} // namespace onnxruntime

2
onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc
View file

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

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

@ -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 <unsigned TPB>
__global__ void FastGeluKernel2(const nv_bfloat162 a, const nv_bfloat162 b, const nv_bfloat162 c,
int input_length, int bias_length,
const nv_bfloat162* input, const nv_bfloat162* bias, nv_bfloat162* output) {
__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<const nv_bfloat162*>(input);
const nv_bfloat162* bias2 = reinterpret_cast<const nv_bfloat162*>(bias);
nv_bfloat162* output2 = reinterpret_cast<nv_bfloat162*>(output);
FastGeluKernel2<blockSize><<<gridSize, blockSize, 0, stream>>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2);
FastGeluKernel2<blockSize>
<<<gridSize, blockSize, 0, stream>>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2);
} else {
#endif
const int gridSize = (input_length + blockSize - 1) / blockSize;
FastGeluKernel<nv_bfloat16, blockSize><<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
FastGeluKernel<BFloat16, blockSize>
<<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
}
#endif
return CUDA_CALL(cudaPeekAtLastError());
}
#endif
} // namespace cuda
} // namespace contrib

8
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -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<void>() {
@ -112,6 +110,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, BiasGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, BiasGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, BiasGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, BiasGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, TransposeMatMul)>, // backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, TransposeMatMul)>, // backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, TransposeMatMul)>, // backward compatibility
@ -180,14 +179,11 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int8_t, QAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Trilu)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FastGelu)>,
// TransposedMatMul is still here for backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul)>, // backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float, LayerNormalization)>,
#endif
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FusedConv)>,
};

2
onnxruntime/contrib_ops/cuda/layer_norm.cc
View file

@ -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 <typename T, typename U, bool simplified>
LayerNorm<T, U, simplified>::LayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {

6
onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
View file

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

8
onnxruntime/contrib_ops/cuda/math/bias_dropout.cc
View file

@ -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<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T1", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>())
.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<Tensor>(3);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
@ -117,7 +117,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const {
const fast_divmod fdm_dim(gsl::narrow_cast<int>(dim));
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(X->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(X->GetElementType());
return t_disp.InvokeRet<Status, BiasDropoutComputeImpl>(
GetDeviceProp(), Stream(), N, fdm_dim, ratio_data, generator, *X, *bias, residual, *Y, mask_data, has_same_shape_bias);
}

4
onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu
View file

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

3
onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops.cc
View file

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

3
onnxruntime/contrib_ops/cuda/math/binary_elementwise_ops_impl.cu
View file

@ -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) \

5
onnxruntime/contrib_ops/cuda/math/fused_matmul.cc
View file

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

18
onnxruntime/contrib_ops/cuda/math/isfinite.cuh
View file

@ -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<float>(value));
}
template <>
__device__ __forceinline__ bool IsInfScalar(const nv_bfloat16 value) {
return __hisinf(value);
__device__ __forceinline__ bool IsInfScalar(const BFloat16 value) {
return isinf(static_cast<float>(value));
}
template <>
__device__ __forceinline__ bool IsNaNScalar(const nv_bfloat16 value) {
return __hisnan(value);
__device__ __forceinline__ bool IsNaNScalar(const BFloat16 value) {
return isnan(static_cast<float>(value));
}
#endif
} // namespace cuda
} // namespace onnxruntime

2
onnxruntime/core/framework/tensorprotoutils.cc
View file

@ -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<uint16_t>(v));
p_data[i] = BFloat16(static_cast<uint16_t>(v), BFloat16::FromBits());
}
return Status::OK();

2
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

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

14
onnxruntime/core/providers/cuda/activation/activations.cc
View file

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

15
onnxruntime/core/providers/cuda/activation/activations_impl.cu
View file

@ -91,20 +91,15 @@ struct OP_ThresholdedRelu : public CtxThresholdedRelu {
count); \
}
#define SPECIALIZED_UNARY_ACTIVATION_IMPL(name, T) \
template void Impl_##name<T>(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<T>(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); \

21
onnxruntime/core/providers/cuda/atomic/common.cuh
View file

@ -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<unsigned int*>(reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
__device__ __forceinline__ void atomic_add(BFloat16* address, BFloat16 value) {
unsigned int* base_address =
reinterpret_cast<unsigned int*>(reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
unsigned int old = *base_address;
unsigned int assumed;
unsigned short x;
BFloat16 bsum;
do {
assumed = old;
x = reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
x = __bfloat16_as_short(__float2bfloat16(__bfloat162float(*reinterpret_cast<const __nv_bfloat16*>(&x)) + __bfloat162float(value)));
old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (x << 16) : (old & 0xffff0000) | x;
bsum.val = reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
bsum = bsum + value;
old = reinterpret_cast<size_t>(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

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

@ -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<float>(a) + static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator-(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) - static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator*(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) * static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator/(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) / static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator-(const BFloat16& a) { return -static_cast<float>(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<float>(a) + b; }
__device__ __forceinline__ float operator-(BFloat16 a, float b) { return static_cast<float>(a) - b; }
__device__ __forceinline__ float operator*(BFloat16 a, float b) { return static_cast<float>(a) * b; }
__device__ __forceinline__ float operator/(BFloat16 a, float b) { return static_cast<float>(a) / b; }
__device__ __forceinline__ float operator+(float a, BFloat16 b) { return a + static_cast<float>(b); }
__device__ __forceinline__ float operator-(float a, BFloat16 b) { return a - static_cast<float>(b); }
__device__ __forceinline__ float operator*(float a, BFloat16 b) { return a * static_cast<float>(b); }
__device__ __forceinline__ float operator/(float a, BFloat16 b) { return a / static_cast<float>(b); }
__device__ __forceinline__ float& operator+=(float& a, const BFloat16& b) { return a += static_cast<float>(b); }
__device__ __forceinline__ float& operator-=(float& a, const BFloat16& b) { return a -= static_cast<float>(b); }
__device__ __forceinline__ float& operator*=(float& a, const BFloat16& b) { return a *= static_cast<float>(b); }
__device__ __forceinline__ float& operator/=(float& a, const BFloat16& b) { return a /= static_cast<float>(b); }
/// Arithmetic with doubles
__device__ __forceinline__ double operator+(BFloat16 a, double b) { return static_cast<double>(a) + b; }
__device__ __forceinline__ double operator-(BFloat16 a, double b) { return static_cast<double>(a) - b; }
__device__ __forceinline__ double operator*(BFloat16 a, double b) { return static_cast<double>(a) * b; }
__device__ __forceinline__ double operator/(BFloat16 a, double b) { return static_cast<double>(a) / b; }
__device__ __forceinline__ double operator+(double a, BFloat16 b) { return a + static_cast<double>(b); }
__device__ __forceinline__ double operator-(double a, BFloat16 b) { return a - static_cast<double>(b); }
__device__ __forceinline__ double operator*(double a, BFloat16 b) { return a * static_cast<double>(b); }
__device__ __forceinline__ double operator/(double a, BFloat16 b) { return a / static_cast<double>(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 <typename T>
__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 <typename T>
__device__ __inline__ T _Gelu(T a) {
return a * _Normcdf(a);
}
template <>
__device__ __inline__ BFloat16 _Sqrt(BFloat16 a) { return sqrtf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Exp(BFloat16 a) { return expf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Log(BFloat16 a) { return logf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Tanh(BFloat16 a) { return tanhf(static_cast<float>(a)); }
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template <>
__device__ __inline__ nv_bfloat16 _Sqrt(nv_bfloat16 a) { return nv_bfloat16(sqrtf(static_cast<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Exp(nv_bfloat16 a) { return nv_bfloat16(expf(static_cast<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Log(nv_bfloat16 a) { return nv_bfloat16(logf(static_cast<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Tanh(nv_bfloat16 a) { return nv_bfloat16(tanhf(static_cast<float>(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<float>(a))); }
#endif
__device__ __inline__ BFloat16 _Normcdf(BFloat16 a) { return normcdff(static_cast<float>(a)); }
template <typename T>
__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.

12
onnxruntime/core/providers/cuda/cuda_common.h
View file

@ -69,18 +69,6 @@ class ToCudaType<MLFloat16> {
}
};
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template <>
class ToCudaType<BFloat16> {
public:
typedef nv_bfloat16 MappedType;
static MappedType FromFloat(float f) {
uint16_t h = BFloat16(f).val;
return *reinterpret_cast<MappedType*>(&h);
}
};
#endif
inline bool CalculateFdmStrides(gsl::span<fast_divmod> p, const std::vector<int64_t>& dims) {
int stride = 1;
if (dims.empty() || p.size() < dims.size())

12
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, bool, Pad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, SpaceToDepth)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, DepthToSpace)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Add)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Sub)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Mul)>,
@ -1980,7 +1972,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum)>,
#endif
// OpSet 14
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, float, Relu)>,
@ -2035,13 +2027,11 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, int8_t, ReduceMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, uint8_t, ReduceMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, int64_t, ReduceMin)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Add)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Sub)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Mul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Div)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Relu)>,
#endif
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, Trilu)>,
// OpSet 15

10
onnxruntime/core/providers/cuda/cuda_execution_provider.h
View file

@ -143,13 +143,11 @@ class CUDAExecutionProvider : public IExecutionProvider {
constant_ones_half_ = cuda::CreateConstantOnes<half>();
}
return reinterpret_cast<const T*>(constant_ones_half_->GetBuffer(stream_, count));
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
} else if (std::is_same<T, nv_bfloat16>::value) {
} else if (std::is_same<T, BFloat16>::value) {
if (!constant_ones_bfloat16_) {
constant_ones_bfloat16_ = cuda::CreateConstantOnes<nv_bfloat16>();
constant_ones_bfloat16_ = cuda::CreateConstantOnes<BFloat16>();
}
return reinterpret_cast<const T*>(constant_ones_bfloat16_->GetBuffer(stream_, count));
#endif
} else {
return nullptr;
}
@ -172,9 +170,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
std::unique_ptr<cuda::IConstantBuffer<float>> constant_ones_float_;
std::unique_ptr<cuda::IConstantBuffer<double>> constant_ones_double_;
std::unique_ptr<cuda::IConstantBuffer<half>> constant_ones_half_;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
std::unique_ptr<cuda::IConstantBuffer<nv_bfloat16>> constant_ones_bfloat16_;
#endif
std::unique_ptr<cuda::IConstantBuffer<BFloat16>> constant_ones_bfloat16_;
AllocatorPtr allocator_;
};

8
onnxruntime/core/providers/cuda/cuda_utils.cu
View file

@ -71,9 +71,7 @@ std::unique_ptr<IConstantBuffer<T>> CreateConstantOnes() {
template std::unique_ptr<IConstantBuffer<float>> CreateConstantOnes<float>();
template std::unique_ptr<IConstantBuffer<double>> CreateConstantOnes<double>();
template std::unique_ptr<IConstantBuffer<half>> CreateConstantOnes<half>();
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template std::unique_ptr<IConstantBuffer<nv_bfloat16>> CreateConstantOnes<nv_bfloat16>();
#endif
template std::unique_ptr<IConstantBuffer<BFloat16>> CreateConstantOnes<BFloat16>();
#define SPECIALIZED_FILL(T) \
template void Fill<T>(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

6
onnxruntime/core/providers/cuda/cudnn_common.cc
View file

@ -161,10 +161,8 @@ const float Consts<half>::Zero = 0;
const float Consts<half>::One = 1;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
const float Consts<nv_bfloat16>::Zero = 0;
const float Consts<nv_bfloat16>::One = 1;
#endif
const float Consts<BFloat16>::Zero = 0;
const float Consts<BFloat16>::One = 1;
template <>
const int8_t Consts<int8_t>::Zero = 0;

4
onnxruntime/core/providers/cuda/cudnn_common.h
View file

@ -124,13 +124,11 @@ struct Consts<half> {
static const float One;
};
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template <>
struct Consts<nv_bfloat16> {
struct Consts<BFloat16> {
static const float Zero;
static const float One;
};
#endif
inline double ClampCudnnBatchNormEpsilon(double epsilon) {
if (epsilon < CUDNN_BN_MIN_EPSILON) {

12
onnxruntime/core/providers/cuda/fpgeneric.cu
View file

@ -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<<<dimGrid, dimBlock, 0, stream>>>(x, incx, y, incy, n);
return CUBLAS_STATUS_SUCCESS;
}
#endif

40
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc
View file

@ -189,22 +189,6 @@ Status BinaryElementwise<ShouldBroadcast>::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<ShouldBroadcast>::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<ShouldBroadcast>::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<ShouldBroadcast>::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) \

20
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu
View file

@ -79,23 +79,15 @@ namespace cuda {
const TArray<int64_t>* rhs_padded_strides, const T2* rhs_data, \
const TArray<fast_divmod>* 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)

2
onnxruntime/core/providers/cuda/math/gemm.cc
View file

@ -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 <typename T>
Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {

2
onnxruntime/core/providers/cuda/math/matmul.cc
View file

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

34
onnxruntime/core/providers/cuda/math/softmax.cc
View file

@ -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<BFloat16, is_log_softmax>( \
cudaStream_t stream, \
const BFloat16* X, \
const TensorShape& input_shape, \
BFloat16* Y, \
int64_t axis) { \
typedef typename ToCudaType<BFloat16>::MappedType CudaT; \
int64_t N = input_shape.SizeToDimension(axis); \
int64_t D = input_shape.SizeFromDimension(axis); \
auto Y_data = reinterpret_cast<CudaT*>(Y); \
auto X_data = reinterpret_cast<const CudaT*>(X); \
dispatch_warpwise_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
stream, Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(N)); \
return Status::OK(); \
#define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \
template <> \
Status SoftMaxComputeHelper<BFloat16, is_log_softmax>(cudaStream_t stream, const BFloat16* X, \
const TensorShape& input_shape, BFloat16* Y, int64_t axis) { \
typedef typename ToCudaType<BFloat16>::MappedType CudaT; \
int64_t N = input_shape.SizeToDimension(axis); \
int64_t D = input_shape.SizeFromDimension(axis); \
auto Y_data = reinterpret_cast<CudaT*>(Y); \
auto X_data = reinterpret_cast<const CudaT*>(X); \
dispatch_warpwise_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
stream, Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(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

20
onnxruntime/core/providers/cuda/math/softmax_impl.cu
View file

@ -109,9 +109,7 @@ template void dispatch_warpwise_softmax_forward<input_t, output_t, acc_t, true>(
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 <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
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<input_t, output_t, acc_t, false>(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<input_t, output_t, acc_t, true>(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<input_t, output_t, acc_t, false>( \
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<input_t, output_t, acc_t, true>( \
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)
}
}

26
onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu
View file

@ -94,24 +94,18 @@ struct ViaTypeMap<half> {
typedef float ViaT;
};
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template <>
struct ViaTypeMap<nv_bfloat16> {
struct ViaTypeMap<BFloat16> {
typedef float ViaT;
};
#endif
template <typename InT, typename OutT>
struct OP_Cast {
__device__ __inline__ OutT operator()(const InT& a) const {
const bool any_float16 = std::is_same<half, InT>::value || std::is_same<half, OutT>::value;
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
const bool any_bf16 = std::is_same<nv_bfloat16, InT>::value || std::is_same<nv_bfloat16, OutT>::value;
typedef typename std::conditional<any_bf16, nv_bfloat16, OutT>::type T1;
const bool any_bf16 = std::is_same<BFloat16, InT>::value || std::is_same<BFloat16, OutT>::value;
typedef typename std::conditional<any_bf16, BFloat16, OutT>::type T1;
typedef typename std::conditional<any_float16, half, T1>::type T;
#else
typedef typename std::conditional<any_float16, half, OutT>::type T;
#endif
typedef typename ViaTypeMap<T>::ViaT ViaT;
return (OutT)((ViaT)a);
}
@ -133,15 +127,8 @@ void Impl_Cast(
#define SPECIALIZED_CAST_IMPL2(InT, OutT) \
template void Impl_Cast<InT, OutT>(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

26
onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc
View file

@ -228,26 +228,18 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
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<variadic_elementwise_ops::Sum, MLFloat16, float, double, BFloat16>;
using SumOp = VariadicElementwiseOp<
variadic_elementwise_ops::Sum,
ALL_IEEE_FLOAT_DATA_TYPES>;
using MinOp = VariadicElementwiseOp<variadic_elementwise_ops::Min, uint32_t, uint64_t, int32_t, int64_t, MLFloat16,
float, double, BFloat16>;
using MinOp = VariadicElementwiseOp<
variadic_elementwise_ops::Min,
uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>;
using MaxOp = VariadicElementwiseOp<variadic_elementwise_ops::Max, uint32_t, uint64_t, int32_t, int64_t, MLFloat16,
float, double, BFloat16>;
using MaxOp = VariadicElementwiseOp<
variadic_elementwise_ops::Max,
uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>;
const DeleteOnUnloadPtr<std::vector<MLDataType>> k_uzilhfd_datatypes = new std::vector<MLDataType>(BuildKernelDefConstraints<uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>());
const DeleteOnUnloadPtr<std::vector<MLDataType>> k_hfd_datatypes = new std::vector<MLDataType>(BuildKernelDefConstraints<ALL_IEEE_FLOAT_DATA_TYPES>());
const DeleteOnUnloadPtr<std::vector<MLDataType>> k_uzilhfd_datatypes = new std::vector<MLDataType>(
BuildKernelDefConstraints<uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double, BFloat16>());
const DeleteOnUnloadPtr<std::vector<MLDataType>> k_hfd_datatypes =
new std::vector<MLDataType>(BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>());
} // namespace

10
onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu
View file

@ -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) \

8
onnxruntime/core/providers/cuda/nn/dropout.cc
View file

@ -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<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T1", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>())
.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<Tensor>(1);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
@ -80,7 +80,7 @@ Status Dropout::ComputeInternal(OpKernelContext* context) const {
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(X->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(X->GetElementType());
t_disp.Invoke<DropoutComputeImpl>(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data);
return Status::OK();

12
onnxruntime/core/providers/cuda/nn/dropout.h
View file

@ -10,18 +10,6 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>() }
#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 <typename T>
struct GetRatioDataImpl {
void operator()(const Tensor* ratio, float& ratio_data) const {

4
onnxruntime/core/providers/cuda/nn/dropout_impl.cu
View file

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

18
onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
View file

@ -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 <typename TIn, typename TOut>
@ -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 <typename TIn, typename TOut>
@ -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

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

@ -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<true>::ComputeImpl<BFloat16, CUDNN_REDUCE_TENSOR_NO_INDICES>(
@ -940,7 +939,6 @@ Status ReduceKernel<true>::ComputeImpl<BFloat16, CUDNN_REDUCE_TENSOR_NO_INDICES>
return Status::OK();
}
#endif
namespace ReductionOps {
@ -997,17 +995,11 @@ template std::unique_ptr<Tensor> ReduceCompute<MLFloat16, CUDNN_REDUCE_TENSOR_NO
} // namespace ReductionOps
#if defined(CUDA_VERSION) && CUDA_VERSION >= 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)

5
onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h
View file

@ -4,6 +4,7 @@
#pragma once
#include <cuda_fp16.h>
#include "core/framework/float16.h"
namespace onnxruntime {
namespace cuda {
@ -17,10 +18,8 @@ template <>
struct AccumulationType<float> { using type = float; };
template <>
struct AccumulationType<double> { using type = double; };
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template <>
struct AccumulationType<nv_bfloat16> { using type = float; };
#endif
struct AccumulationType<BFloat16> { using type = float; };
template <typename T>
using AccumulationType_t = typename AccumulationType<T>::type;

54
onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h
View file

@ -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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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);

6
onnxruntime/core/providers/cuda/tensor/cast_op.cc
View file

@ -11,9 +11,7 @@ namespace cuda {
const DeleteOnUnloadPtr<std::vector<MLDataType>> castOpTypeConstraints = new std::vector<MLDataType> {
DataTypeImpl::GetTensorType<MLFloat16>(),
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
DataTypeImpl::GetTensorType<BFloat16>(),
#endif
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int8_t>(),
@ -81,9 +79,7 @@ Status Cast<SrcT>::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

41
onnxruntime/core/providers/cuda/tensor/gather_nd.cc
View file

@ -112,35 +112,12 @@ Status GatherNDBase::PrepareCompute(
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()), \
GatherND<TIndex>);
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define GATHER_ND_T_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>(), \
DataTypeImpl::GetTensorType<bool>(), \
DataTypeImpl::GetTensorType<int64_t>() }
#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16, bool
#else
#define GATHER_ND_T_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<bool>(), \
DataTypeImpl::GetTensorType<int64_t>() }
#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<TIndex>()), \
#define REGISTER_KERNEL_TYPED_GATHER_ND(TIndex, ver) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
GatherND, kOnnxDomain, ver, TIndex, kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", BuildKernelDefConstraints<float, MLFloat16, double, int64_t, BFloat16, bool>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()), \
GatherND<TIndex>);
// TODO: decprecate GatherND-1 after updating training models to opset-12
@ -211,9 +188,9 @@ Status GatherND<TIndex>::ComputeInternal(OpKernelContext* context) const {
const void* const kernel_input_data = input_tensor->DataRaw();
void* const kernel_output_data = output_tensor->MutableDataRaw();
utils::MLTypeCallDispatcher<GATHER_ND_T_DATA_TYPES> t_disp(input_tensor->GetElementType());
t_disp.Invoke<GatherNDComputeImpl>(
Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get());
utils::MLTypeCallDispatcher<float, MLFloat16, double, int64_t, BFloat16, bool> t_disp(input_tensor->GetElementType());
t_disp.Invoke<GatherNDComputeImpl>(Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data,
input_slice_offsets_buffer.get());
return Status::OK();
}

4
onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu
View file

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

18
onnxruntime/core/providers/rocm/atomic/common.cuh
View file

@ -4,6 +4,7 @@
#pragma once
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#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<unsigned int*>(reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
unsigned int old = *base_address;
unsigned int assumed;
BFloat16 bsum;
do {
assumed = old;
bsum.val = reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
bsum = bsum + value;
old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (bsum.val << 16) : (old & 0xffff0000) | bsum.val;
old = atomicCAS(base_address, assumed, old);
} while (assumed != old);
}
} // namespace rocm
} // namespace onnxruntime
} // namespace onnxruntime

91
onnxruntime/core/providers/rocm/cu_inc/common.cuh
View file

@ -14,6 +14,82 @@
namespace onnxruntime {
namespace rocm {
/// Arithmetic for BFloat16
__device__ __forceinline__ BFloat16 operator+(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) + static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator-(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) - static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator*(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) * static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator/(const BFloat16& a, const BFloat16& b) {
return static_cast<float>(a) / static_cast<float>(b);
}
__device__ __forceinline__ BFloat16 operator-(const BFloat16& a) { return -static_cast<float>(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<float>(a) + b; }
__device__ __forceinline__ float operator-(BFloat16 a, float b) { return static_cast<float>(a) - b; }
__device__ __forceinline__ float operator*(BFloat16 a, float b) { return static_cast<float>(a) * b; }
__device__ __forceinline__ float operator/(BFloat16 a, float b) { return static_cast<float>(a) / b; }
__device__ __forceinline__ float operator+(float a, BFloat16 b) { return a + static_cast<float>(b); }
__device__ __forceinline__ float operator-(float a, BFloat16 b) { return a - static_cast<float>(b); }
__device__ __forceinline__ float operator*(float a, BFloat16 b) { return a * static_cast<float>(b); }
__device__ __forceinline__ float operator/(float a, BFloat16 b) { return a / static_cast<float>(b); }
__device__ __forceinline__ float& operator+=(float& a, const BFloat16& b) { return a += static_cast<float>(b); }
__device__ __forceinline__ float& operator-=(float& a, const BFloat16& b) { return a -= static_cast<float>(b); }
__device__ __forceinline__ float& operator*=(float& a, const BFloat16& b) { return a *= static_cast<float>(b); }
__device__ __forceinline__ float& operator/=(float& a, const BFloat16& b) { return a /= static_cast<float>(b); }
/// Arithmetic with doubles
__device__ __forceinline__ double operator+(BFloat16 a, double b) { return static_cast<double>(a) + b; }
__device__ __forceinline__ double operator-(BFloat16 a, double b) { return static_cast<double>(a) - b; }
__device__ __forceinline__ double operator*(BFloat16 a, double b) { return static_cast<double>(a) * b; }
__device__ __forceinline__ double operator/(BFloat16 a, double b) { return static_cast<double>(a) / b; }
__device__ __forceinline__ double operator+(double a, BFloat16 b) { return a + static_cast<double>(b); }
__device__ __forceinline__ double operator-(double a, BFloat16 b) { return a - static_cast<double>(b); }
__device__ __forceinline__ double operator*(double a, BFloat16 b) { return a * static_cast<double>(b); }
__device__ __forceinline__ double operator/(double a, BFloat16 b) { return a / static_cast<double>(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 <typename T>
__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<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Exp(BFloat16 a) { return expf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Log(BFloat16 a) { return logf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Tanh(BFloat16 a) { return tanhf(static_cast<float>(a)); }
template <>
__device__ __inline__ BFloat16 _Normcdf(BFloat16 a) { return normcdff(static_cast<float>(a)); }
template <typename T>
__device__ __inline__ T _Gelu(T a) {
return a * _Normcdf(a);

2
onnxruntime/test/framework/data_types_test.cc
View file

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

209
onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc
View file

@ -14,12 +14,37 @@ namespace test {
std::vector<MLFloat16> MakeMLFloat16(const std::initializer_list<float>& input) {
std::vector<MLFloat16> 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<int64_t>& lhs_dim,
const std::initializer_list<float>& lhs_values, const std::vector<int64_t>& rhs_dim,
const std::initializer_list<float>& rhs_values, const std::vector<int64_t>& out_dim,
const std::initializer_list<float>& out_values) {
{
OpTester tester(op_name, 14);
tester.AddInput<MLFloat16>("A", lhs_dim, MakeMLFloat16(lhs_values));
tester.AddInput<MLFloat16>("B", rhs_dim, MakeMLFloat16(rhs_values));
tester.AddOutput<MLFloat16>("C", out_dim, MakeMLFloat16(out_values));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
{
OpTester tester(op_name, 14);
tester.AddInput<BFloat16>("A", lhs_dim, MakeBFloat16(lhs_values));
tester.AddInput<BFloat16>("B", rhs_dim, MakeBFloat16(rhs_values));
tester.AddOutput<BFloat16>("C", out_dim, MakeBFloat16(out_values));
std::vector<std::unique_ptr<IExecutionProvider>> 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<int64_t> dims{3, 3};
test.AddInput<float>("A", dims,
{1.0f, 2.0f, -1.0f,
0.0f, 1.5f, -100.0f,
-5.4f, 9.3f, -10000.0f});
test.AddInput<float>("B", dims,
{-1.0f, 4.4f, 432.3f,
0.0f, 3.5f, 64.0f,
-5.4f, 9.3f, 10000.0f});
test.AddOutput<float>("C", dims,
{0.0f, 6.4f, 431.3f,
0.0f, 5.0f, -36.0f,
-10.8f, 18.6f, 0.0f});
std::initializer_list<float> lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.4f, 9.3f, -10000.0f};
std::initializer_list<float> rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.3f, 10000.0f};
std::initializer_list<float> out_values{0.0f, 6.4f, 431.3f, 0.0f, 5.0f, -36.0f, -10.8f, 18.6f, 0.0f};
test.AddInput<float>("A", dims, lhs_values);
test.AddInput<float>("B", dims, rhs_values);
test.AddOutput<float>("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<int64_t> dims{3, 3};
test.AddInput<float>("A", dims,
{1.0f, 2.0f, 3.0f,
4.0f, 5.0f, 6.0f,
7.0f, 8.0f, 9.0f});
test.AddInput<float>("B", {3, 1},
{3.0f,
2.0f,
1.0f});
test.AddOutput<float>("C", dims,
{4.0f, 5.0f, 6.0f,
6.0f, 7.0f, 8.0f,
8.0f, 9.0f, 10.0f});
std::initializer_list<float> lhs_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
std::initializer_list<float> rhs_values{3.0f, 2.0f, 1.0f};
std::initializer_list<float> out_values{4.0f, 5.0f, 6.0f, 6.0f, 7.0f, 8.0f, 8.0f, 9.0f, 10.0f};
test.AddInput<float>("A", dims, lhs_values);
test.AddInput<float>("B", {3, 1}, rhs_values);
test.AddOutput<float>("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<float>("A", {3, 1},
{3.0f,
2.0f,
1.0f});
test.AddInput<float>("B", {3},
{1.0f, 2.0f, 3.0f});
test.AddOutput<float>("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<float> lhs_values{3.0f, 2.0f, 1.0f};
std::initializer_list<float> rhs_values{1.0f, 2.0f, 3.0f};
std::initializer_list<float> out_values{4.0f, 5.0f, 6.0f, 3.0f, 4.0f, 5.0f, 2.0f, 3.0f, 4.0f};
test.AddInput<float>("A", {3, 1}, lhs_values);
test.AddInput<float>("B", {3}, rhs_values);
test.AddOutput<float>("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<float>("A", {3},
{1.0f, 2.0f, 3.0f});
test.AddInput<float>("B", {3, 1},
{3.0f,
2.0f,
1.0f});
test.AddOutput<float>("C", {3, 3},
{4.0f, 5.0f, 6.0f,
3.0f, 4.0f, 5.0f,
2.0f, 3.0f, 4.0f});
std::initializer_list<float> lhs_values{1.0f, 2.0f, 3.0f};
std::initializer_list<float> rhs_values{3.0f, 2.0f, 1.0f};
std::initializer_list<float> out_values{4.0f, 5.0f, 6.0f, 3.0f, 4.0f, 5.0f, 2.0f, 3.0f, 4.0f};
test.AddInput<float>("A", {3}, lhs_values);
test.AddInput<float>("B", {3, 1}, rhs_values);
test.AddOutput<float>("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<int64_t> dims{3, 3};
test.AddInput<float>("A", dims,
{1.0f, 2.0f, -1.0f,
0.0f, 1.5f, -100.0f,
-5.4f, 9.3f, -10000.0f});
test.AddInput<float>("B", dims,
{-1.0f, 4.4f, 432.3f,
0.0f, 3.5f, 64.0f,
-5.4f, 9.3f, 10000.0f});
test.AddOutput<float>("C", dims,
{2.0f, -2.4f, -433.3f,
0.0f, -2.0f, -164.0f,
0.0f, 0.0f, -20000.0f});
std::initializer_list<float> lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.4f, 9.3f, -10000.0f};
std::initializer_list<float> rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.3f, 10000.0f};
std::initializer_list<float> out_values{2.0f, -2.4f, -433.3f, 0.0f, -2.0f, -164.0f, 0.0f, 0.0f, -20000.0f};
test.AddInput<float>("A", dims, lhs_values);
test.AddInput<float>("B", dims, rhs_values);
test.AddOutput<float>("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<int64_t> dims{3, 3};
test.AddInput<float>("A", dims,
{1.0f, 2.0f, -1.0f,
0.0f, 1.5f, -100.0f, -5.4f,
9.30f, -10000.0f});
test.AddInput<float>("B", dims,
{-1.0f, 4.4f, 432.3f,
0.0f, 3.5f, 64.0f, -5.4f,
9.30f, 10000.0f});
test.AddOutput<float>("C", dims,
{-1.0f, 8.8f, -432.3f,
0.0f, 5.25f, -6400.0f,
29.16f, 86.49f, -100000000.0f});
std::initializer_list<float> lhs_values{1.0f, 2.0f, -1.0f, 0.0f, 1.5f, -100.0f, -5.0f, 9.30f, -10000.0f};
std::initializer_list<float> rhs_values{-1.0f, 4.4f, 432.3f, 0.0f, 3.5f, 64.0f, -5.4f, 9.0f, 10000.0f};
std::initializer_list<float> out_values{-1.0f, 8.8f, -432.3f, 0.0f, 5.25f, -6400.0f, 27.0f, 83.7f, -100000000.0f};
test.AddInput<float>("A", dims, lhs_values);
test.AddInput<float>("B", dims, rhs_values);
test.AddOutput<float>("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<int64_t> dims{2, 3};
test.AddInput<float>("A", dims,
{1000.0f, 1.0f, 6.0f,
0.0f, -10.0f, -1.0f});
test.AddInput<float>("B", dims,
{1000.0f, 2.0f, 3.0f,
1.0f, -1.0f, 4.0f});
test.AddOutput<float>("C", dims,
{1.0f, 0.5f, 2.0f,
0.0f, 10.0f, -0.25f});
std::initializer_list<float> lhs_values{1000.0f, 1.0f, 6.0f, 0.0f, -10.0f, -1.0f};
std::initializer_list<float> rhs_values{1000.0f, 2.0f, 3.0f, 1.0f, -1.0f, 4.0f};
std::initializer_list<float> out_values{1.0f, 0.5f, 2.0f, 0.0f, 10.0f, -0.25f};
test.AddInput<float>("A", dims, lhs_values);
test.AddInput<float>("B", dims, rhs_values);
test.AddOutput<float>("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) {

23
onnxruntime/test/providers/cpu/math/gemm_test.cc
View file

@ -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<BFloat16>("A", {2, 4}, MakeBFloat16({1.0f, 2.0f, 3.0f, 4.0f, -1.0f, -2.0f, -3.0f, -4.0f}));
test.AddInput<BFloat16>("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<BFloat16>("C", {2, 3}, MakeBFloat16({1.f, 1.f, 1.f, 1.f, 1.f, 1.f}));
test.AddOutput<BFloat16>("Y", {2, 3}, MakeBFloat16({11.0f, 11.0f, 11.0f, -9.0f, -9.0f, -9.0f}));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
#endif
template <typename T>
void TestGemmBroadcast() {
auto run_test = [](bool b_is_initializer, bool c_is_initializer) {

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

@ -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<BFloat16>("data", {3, 2, 2},
MakeBFloat16({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f}));
test.AddInput<int64_t>("axes", {2}, std::vector<int64_t>{0, 1});
test.AddOutput<BFloat16>("reduced", {2}, MakeBFloat16({36.0f, 42.0f}));
std::vector<std::unique_ptr<IExecutionProvider>> 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);

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

@ -329,6 +329,9 @@ struct TensorCheck<BFloat16> {
/// XXX: May need to adjust threshold as BFloat is coarse
float threshold = 0.001f;
#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING) || defined(USE_CUDA) || defined(USE_ROCM)
threshold = 0.008f;
#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;

6
onnxruntime/test/providers/provider_test_utils.h
View file

@ -1088,6 +1088,12 @@ inline std::vector<MLFloat16> FloatsToMLFloat16s(const std::vector<float>& f) {
return m;
}
inline std::vector<BFloat16> MakeBFloat16(const std::initializer_list<float>& input) {
std::vector<BFloat16> 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_};
}

2
orttraining/orttraining/core/graph/training_op_defs.cc
View file

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

12
orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc
View file

@ -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<ALL_IEEE_FLOAT_DATA_TYPES>())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.MayInplace(0, 0),
BiasGeluGrad_dX<gelu_computation_mode::Default>);
@ -32,7 +26,7 @@ ONNX_OPERATOR_KERNEL_EX(
1,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraints<ALL_IEEE_FLOAT_DATA_TYPES>())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.MayInplace(0, 0),
BiasGeluGrad_dX<gelu_computation_mode::Approximation>);
@ -76,7 +70,7 @@ Status BiasGeluGrad_dX<GeluComputationMode>::ComputeInternal(OpKernelContext* co
const auto input_size = input_shape.Size(), bias_size = bias_shape.Size();
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> dispatcher{X->GetElementType()};
utils::MLTypeCallDispatcher<MLFloat16, float, double, BFloat16> dispatcher{X->GetElementType()};
dispatcher.Invoke<KernelLaunchDispatcher>(Stream(), input_size, bias_size, *dY, *X, *B, *dX);
return Status::OK();

6
orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu
View file

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

12
orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc
View file

@ -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<ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, 12, float, int64_t, SoftmaxCrossEntropyLoss)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, int64_t, SoftmaxCrossEntropyLoss)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, int64_t, SoftmaxCrossEntropyLoss)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, int64_t, SoftmaxCrossEntropyLoss)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float, BatchNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double, BatchNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_MLFloat16, BatchNormalizationGrad)>,
@ -385,7 +391,6 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_float, BatchNormInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float_float, BatchNormInternal)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
// Adam
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer)>,
@ -418,7 +423,6 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2)>,
#endif
// P2P communication operators.
#if defined(ORT_USE_NCCL) || defined(USE_MPI)

4
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc
View file

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

3
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu
View file

@ -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 <typename T, typename TAcc, typename Tin>
__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

37
orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc
View file

@ -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<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>() }
#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<SrcT>()) \
.TypeConstraint("ScaleT", DataTypeImpl::GetTensorType<float>()) \
.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<SrcT>()) \
.TypeConstraint("ScaleT", DataTypeImpl::GetTensorType<float>()) \
.TypeConstraint("DstT", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>()), \
MixedPrecisionScale<SrcT>);
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<SrcT>::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<SrcT>::ComputeInternal(OpKernelContext* context) cons
REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(MLFloat16)
REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(float)
REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16)
template Status MixedPrecisionScale<MLFloat16>::ComputeInternal(OpKernelContext* context) const;
template Status MixedPrecisionScale<float>::ComputeInternal(OpKernelContext* context) const;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16)
template Status MixedPrecisionScale<BFloat16>::ComputeInternal(OpKernelContext* context) const;
#endif
} // namespace cuda
} // namespace onnxruntime

13
orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu
View file

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

42
orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc
View file

@ -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<BFloat16, is_log_softmax>( \
cudaStream_t stream, \
const BFloat16* dY, \
const TensorShape& input_shape, \
const BFloat16* Y, \
BFloat16* dX, \
cudnnHandle_t, \
int64_t axis) { \
typedef typename ToCudaType<BFloat16>::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<const CudaT*>(dY); \
auto Y_data = reinterpret_cast<const CudaT*>(Y); \
auto dX_data = reinterpret_cast<CudaT*>(dX); \
dispatch_softmax_backward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
stream, dX_data, dY_data, Y_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(N)); \
return Status::OK(); \
#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \
template <> \
Status SoftMaxGradComputeHelper<BFloat16, is_log_softmax>(cudaStream_t stream, const BFloat16* dY, \
const TensorShape& input_shape, const BFloat16* Y, \
BFloat16* dX, cudnnHandle_t, int64_t axis) { \
typedef typename ToCudaType<BFloat16>::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<const CudaT*>(dY); \
auto Y_data = reinterpret_cast<const CudaT*>(Y); \
auto dX_data = reinterpret_cast<CudaT*>(dX); \
dispatch_softmax_backward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
stream, dX_data, dY_data, Y_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), \
gsl::narrow_cast<int>(N)); \
return Status::OK(); \
}
SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true)
SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false)
#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

5
orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu
View file

@ -200,10 +200,7 @@ template void dispatch_softmax_backward<input_t, output_t, acc_t, true>(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)
}
}

25
orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc
View file

@ -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<bool>()) \
.InputMemoryType(OrtMemTypeCPUInput, 2), \
DropoutGrad);
REGISTER_GRADIENT_KERNEL(DropoutGrad)
ONNX_OPERATOR_KERNEL_EX(DropoutGrad, kMSDomain, 1, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T1", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>())
.InputMemoryType(OrtMemTypeCPUInput, 2),
DropoutGrad);
template <typename T>
struct DropoutGradComputeImpl {
@ -53,13 +46,13 @@ Status DropoutGrad::ComputeInternal(OpKernelContext* context) const {
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(2);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
auto dX = context->Output(0, shape);
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(dY->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(dY->GetElementType());
t_disp.Invoke<DropoutGradComputeImpl>(Stream(), N, *dY, mask_data, ratio_data, *dX);
return Status::OK();

4
orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu
View file

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

2
orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc
View file

@ -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 <typename T, typename U, bool simplified>
LayerNormGrad<T, U, simplified>::LayerNormGrad(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {

6
orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
View file

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

3
orttraining/orttraining/training_ops/cuda/optimizer/adam.cc
View file

@ -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 <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM, typename T_MIXED_PRECISION_FP>
Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>::ComputeInternal(OpKernelContext* ctx) const {

21
orttraining/orttraining/training_ops/cuda/optimizer/adam.cu
View file

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

2
orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc
View file

@ -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 <typename T>
Status ZeroGradient<T>::ComputeInternal(OpKernelContext* ctx) const {

18
orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu
View file

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

3
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc
View file

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

53
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu
View file

@ -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 <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
__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 <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
__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 <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
__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

2
orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc
View file

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

14
orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu
View file

@ -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 <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
__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

33
orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
View file

@ -10,29 +10,14 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>() }
#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<int64_t>())
.TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES)
.TypeConstraint("Tind", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
GatherGrad);
ONNX_OPERATOR_KERNEL_EX(GatherGrad, kMSDomain, 1, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 0)
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>())
.TypeConstraint("Tind", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
GatherGrad);
namespace {
template <typename T, typename TIndex>
@ -96,11 +81,9 @@ Status DispatchToGatherGradImpl(
} else if (utils::IsPrimitiveDataType<MLFloat16>(t_data_type)) {
return DispatchToGatherGradImplByTindex<MLFloat16>(
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<BFloat16>(t_data_type)) {
return DispatchToGatherGradImplByTindex<BFloat16>(
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);

4
orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu
View file

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

30
orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc
View file

@ -8,26 +8,14 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>() }
#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<TIndex>()).TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()).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<MLFloat16, float, double, BFloat16>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()) \
.InputMemoryType(OrtMemTypeCPUInput, 0), \
GatherNDGrad<TIndex>);
REGISTER_KERNEL_TYPED_GATHER_ND_GRAD(int64_t)
@ -94,7 +82,7 @@ Status GatherNDGrad<TIndex>::ComputeInternal(OpKernelContext* context) const {
const void* const kernel_input_data = update_tensor->DataRaw();
void* const kernel_output_data = output_tensor->MutableDataRaw();
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(update_tensor->GetElementType());
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(update_tensor->GetElementType());
t_disp.Invoke<GatherNDGradComputeImpl>(
Stream(), num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get());

4
orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu
View file

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