From 7af1887b330fc8b768c638488441ed5e956c1f2c Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Fri, 21 Dec 2018 12:40:59 -0800 Subject: [PATCH] Introduce basic BFloat16 runtime support (#235) * Add basic support for BFloat16 type. * Advance onnx submodule for bfloat16 support. * Update install_deps for linux. * Address review comments. --- cmake/external/onnx | 2 +- .../onnxruntime/core/framework/data_types.h | 122 +++++++++++++----- onnxruntime/core/framework/data_types.cc | 13 ++ .../core/framework/onnxruntime_typeinfo.cc | 4 + .../core/framework/tensor_type_and_shape.cc | 3 + .../core/framework/tensorprotoutils.cc | 36 ++++++ onnxruntime/core/framework/tensorutils.cc | 105 ++++++++++----- onnxruntime/core/session/onnxruntime_c_api.cc | 8 +- onnxruntime/test/framework/data_types_test.cc | 50 +++++++ .../test/providers/provider_test_utils.cc | 2 +- .../test/providers/provider_test_utils.h | 9 +- onnxruntime/test/util/compare_mlvalue.cc | 38 +++++- .../linux/docker/scripts/install_deps.sh | 4 +- 13 files changed, 322 insertions(+), 74 deletions(-) diff --git a/cmake/external/onnx b/cmake/external/onnx index 0c8d857bb1..0a7cc483eb 160000 --- a/cmake/external/onnx +++ b/cmake/external/onnx @@ -1 +1 @@ -Subproject commit 0c8d857bb162431912b255d5c0e773fb7c131a65 +Subproject commit 0a7cc483eb0c34e15414437bbbb420f52df4d8c2 diff --git a/include/onnxruntime/core/framework/data_types.h b/include/onnxruntime/core/framework/data_types.h index 023bc2b362..e2c93126e4 100644 --- a/include/onnxruntime/core/framework/data_types.h +++ b/include/onnxruntime/core/framework/data_types.h @@ -5,8 +5,9 @@ #include #include -#include +#include #include +#include #include "core/common/common.h" #include "core/common/exceptions.h" @@ -45,13 +46,76 @@ union MLFloat16 { MLFloat16() : val(0) {} }; -inline bool operator==(const MLFloat16& left, const MLFloat16& right) -{ +inline bool operator==(const MLFloat16& left, const MLFloat16& right) { return left.val == right.val; } -inline bool operator!=(const MLFloat16& left, const MLFloat16& right) -{ +inline bool operator!=(const MLFloat16& left, const MLFloat16& right) { + return left.val != right.val; +} + +struct ort_endian { + union q { + uint16_t v_; + uint8_t b_[2]; + constexpr q(uint16_t v) noexcept : v_(v) {} + }; + static constexpr bool is_little() { + return q(0x200).b_[0] == 0x0; + } + static constexpr bool is_big() { + return q(0x200).b_[0] == 0x2; + } +}; + +//BFloat16 +struct BFloat16 { + uint16_t val; + explicit BFloat16() : val(0) {} + explicit BFloat16(uint16_t v) : val(v) {} + explicit BFloat16(float v) { + uint16_t* dst = reinterpret_cast(&v); + if (ort_endian::is_little()) { + val = dst[1]; + } else { + val = dst[0]; + } + } + float ToFloat() const { + float result; + uint16_t* dst = reinterpret_cast(&result); + if (ort_endian::is_little()) { + dst[1] = val; + dst[0] = 0; + } else { + dst[0] = val; + dst[1] = 0; + } + return result; + } +}; + +inline void BFloat16ToFloat(const BFloat16* blf, float* flt, size_t size) { + auto src = blf; + auto d = flt; + for (; size != 0; ++src, ++d, --size) { + *d = src->ToFloat(); + } +} + +inline void FloatToBFloat16(const float* flt, BFloat16* blf, size_t size) { + auto src = flt; + auto d = blf; + for (; size != 0; ++src, ++d, --size) { + new (d) BFloat16(*src); + } +} + +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; } @@ -165,7 +229,7 @@ struct IsAnyOf { template struct IsTensorContainedType : public IsAnyOf { + double, uint32_t, uint64_t, BFloat16> { }; /// This template's Get() returns a corresponding MLDataType @@ -203,7 +267,7 @@ struct SetMapTypes { MLDataType dt = GetMLDataType::value>::Get(); const auto* value_proto = dt->GetTypeProto(); ORT_ENFORCE(value_proto != nullptr, typeid(V).name(), - " expected to be a registered ONNX type"); + " expected to be a registered ONNX type"); CopyMutableMapValue(*value_proto, proto); } }; @@ -220,7 +284,7 @@ struct SetSequenceType { MLDataType dt = GetMLDataType::value>::Get(); const auto* elem_proto = dt->GetTypeProto(); ORT_ENFORCE(elem_proto != nullptr, typeid(T).name(), - " expected to be a registered ONNX type"); + " expected to be a registered ONNX type"); CopyMutableSeqElement(*elem_proto, proto); } }; @@ -482,7 +546,7 @@ class NonOnnxType : public DataTypeImpl { // thus a simple way to pre-instantiate a given template // at a registration time does not currently work and the macro // is needed. -#define ORT_REGISTER_TENSOR_TYPE(ELEM_TYPE) \ +#define ORT_REGISTER_TENSOR_TYPE(ELEM_TYPE) \ template <> \ MLDataType TensorType::Type() { \ static TensorType tensor_type; \ @@ -493,7 +557,7 @@ class NonOnnxType : public DataTypeImpl { return TensorType::Type(); \ } -#define ORT_REGISTER_MAP(TYPE) \ +#define ORT_REGISTER_MAP(TYPE) \ template <> \ MLDataType MapType::Type() { \ static MapType map_type; \ @@ -504,7 +568,7 @@ class NonOnnxType : public DataTypeImpl { return MapType::Type(); \ } -#define ORT_REGISTER_SEQ(TYPE) \ +#define ORT_REGISTER_SEQ(TYPE) \ template <> \ MLDataType SequenceType::Type() { \ static SequenceType sequence_type; \ @@ -515,25 +579,25 @@ class NonOnnxType : public DataTypeImpl { return SequenceType::Type(); \ } -#define ORT_REGISTER_NON_ONNX_TYPE(TYPE) \ - template <> \ - MLDataType NonOnnxType::Type() { \ - static NonOnnxType non_onnx_type; \ - return &non_onnx_type; \ - } \ - template <> \ - MLDataType DataTypeImpl::GetType() { \ - return NonOnnxType::Type(); \ +#define ORT_REGISTER_NON_ONNX_TYPE(TYPE) \ + template <> \ + MLDataType NonOnnxType::Type() { \ + static NonOnnxType non_onnx_type; \ + return &non_onnx_type; \ + } \ + template <> \ + MLDataType DataTypeImpl::GetType() { \ + return NonOnnxType::Type(); \ } -#define ORT_REGISTER_OPAQUE_TYPE(CPPType, Domain, Name) \ - template <> \ - MLDataType OpaqueType::Type() { \ - static OpaqueType opaque_type; \ - return &opaque_type; \ - } \ - template <> \ - MLDataType DataTypeImpl::GetType() { \ - return OpaqueType::Type(); \ +#define ORT_REGISTER_OPAQUE_TYPE(CPPType, Domain, Name) \ + template <> \ + MLDataType OpaqueType::Type() { \ + static OpaqueType opaque_type; \ + return &opaque_type; \ + } \ + template <> \ + MLDataType DataTypeImpl::GetType() { \ + return OpaqueType::Type(); \ } } // namespace onnxruntime diff --git a/onnxruntime/core/framework/data_types.cc b/onnxruntime/core/framework/data_types.cc index 1ebb69b086..73d1db82d6 100644 --- a/onnxruntime/core/framework/data_types.cc +++ b/onnxruntime/core/framework/data_types.cc @@ -85,6 +85,10 @@ template <> constexpr ONNX_NAMESPACE::TensorProto_DataType ToTensorDataType() { return ONNX_NAMESPACE::TensorProto_DataType_UINT64; }; +template <> +constexpr ONNX_NAMESPACE::TensorProto_DataType ToTensorDataType() { + return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; +}; template struct TensorContainedTypeSetter { @@ -123,6 +127,8 @@ template struct TensorContainedTypeSetter; template struct TensorContainedTypeSetter; +template struct + TensorContainedTypeSetter; void CopyMutableMapValue(const ONNX_NAMESPACE::TypeProto& value_proto, ONNX_NAMESPACE::TypeProto& map_proto) { @@ -435,6 +441,7 @@ ORT_REGISTER_TENSOR_TYPE(double); ORT_REGISTER_TENSOR_TYPE(uint32_t); ORT_REGISTER_TENSOR_TYPE(uint64_t); ORT_REGISTER_TENSOR_TYPE(MLFloat16); +ORT_REGISTER_TENSOR_TYPE(BFloat16); ORT_REGISTER_MAP(MapStringToString); ORT_REGISTER_MAP(MapStringToInt64); @@ -482,6 +489,7 @@ void RegisterAllProtos(const std::function& reg_fn) { REGISTER_TENSOR_PROTO(uint32_t, reg_fn); REGISTER_TENSOR_PROTO(uint64_t, reg_fn); REGISTER_TENSOR_PROTO(MLFloat16, reg_fn); + REGISTER_TENSOR_PROTO(BFloat16, reg_fn); REGISTER_ONNX_PROTO(MapStringToString, reg_fn); REGISTER_ONNX_PROTO(MapStringToInt64, reg_fn); @@ -540,6 +548,8 @@ MLDataType DataTypeImpl::TypeFromProto(const ONNX_NAMESPACE::TypeProto& proto) { return DataTypeImpl::GetTensorType(); case TensorProto_DataType_FLOAT16: return DataTypeImpl::GetTensorType(); + case TensorProto_DataType_BFLOAT16: + return DataTypeImpl::GetTensorType(); default: ORT_NOT_IMPLEMENTED("tensor type ", tensor_type.elem_type(), " is not supported"); } @@ -685,6 +695,7 @@ ORT_REGISTER_NON_ONNX_TYPE(double); ORT_REGISTER_NON_ONNX_TYPE(uint32_t); ORT_REGISTER_NON_ONNX_TYPE(uint64_t); ORT_REGISTER_NON_ONNX_TYPE(MLFloat16); +ORT_REGISTER_NON_ONNX_TYPE(BFloat16); const std::vector& DataTypeImpl::AllFixedSizeTensorTypes() { static std::vector all_fixed_size_tensor_types = @@ -699,6 +710,7 @@ const std::vector& DataTypeImpl::AllFixedSizeTensorTypes() { DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}; return all_fixed_size_tensor_types; @@ -717,6 +729,7 @@ const std::vector& DataTypeImpl::AllTensorTypes() { DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}; diff --git a/onnxruntime/core/framework/onnxruntime_typeinfo.cc b/onnxruntime/core/framework/onnxruntime_typeinfo.cc index 655bc840fc..0c45349982 100644 --- a/onnxruntime/core/framework/onnxruntime_typeinfo.cc +++ b/onnxruntime/core/framework/onnxruntime_typeinfo.cc @@ -8,6 +8,7 @@ #include "core/framework/tensor.h" #include "core/graph/onnx_protobuf.h" +using onnxruntime::BFloat16; using onnxruntime::DataTypeImpl; using onnxruntime::MLFloat16; using onnxruntime::Tensor; @@ -80,6 +81,9 @@ const DataTypeImpl* ElementTypeFromProto(int type) { return DataTypeImpl::GetType(); case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: return DataTypeImpl::GetType(); + case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: + return DataTypeImpl::GetType(); + default: ORT_NOT_IMPLEMENTED(__FUNCTION__, ":tensor type ", type, " is not supported"); } diff --git a/onnxruntime/core/framework/tensor_type_and_shape.cc b/onnxruntime/core/framework/tensor_type_and_shape.cc index eb22ef015c..da355705ec 100644 --- a/onnxruntime/core/framework/tensor_type_and_shape.cc +++ b/onnxruntime/core/framework/tensor_type_and_shape.cc @@ -10,6 +10,7 @@ #include #include +using onnxruntime::BFloat16; using onnxruntime::DataTypeImpl; using onnxruntime::MLFloat16; using onnxruntime::Tensor; @@ -102,6 +103,8 @@ inline ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType( type = ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; } else if (cpp_type == onnxruntime::DataTypeImpl::GetType()) { type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; + } else if (cpp_type == onnxruntime::DataTypeImpl::GetType()) { + type = ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16; } else if (cpp_type == onnxruntime::DataTypeImpl::GetType()) { type = ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; } else if (cpp_type == onnxruntime::DataTypeImpl::GetType()) { diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index b36c72266a..a0b49a2938 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -138,6 +138,37 @@ common::Status GetTensorByTypeFromTensorProto(const TensorProto& tens return common::Status::OK(); } +template <> +common::Status GetTensorByTypeFromTensorProto(const TensorProto& tensor_proto, + const TensorShape& tensor_shape, + std::unique_ptr* p_tensor, + AllocatorPtr alloc, + void* preallocated, + size_t preallocated_size) { + int64_t tensor_size = tensor_shape.Size(); + if (tensor_size < 0) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Tensor shape cannot contain any negative value"); + } + static_assert(sizeof(BFloat16) == sizeof(uint16_t), "BFloat16 must has 16 bit size"); + size_t size_to_allocate; + if (!IAllocator::CalcMemSizeForArrayWithAlignment<256>(static_cast(tensor_size), sizeof(BFloat16), &size_to_allocate)) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "size overflow"); + } + + if (preallocated && preallocated_size != size_to_allocate) + return Status(ONNXRUNTIME, FAIL, "The buffer planner is not consistent with tensor buffer size"); + + BFloat16* p_data = static_cast(preallocated ? preallocated : alloc->Alloc(size_to_allocate)); + ORT_RETURN_IF_ERROR(::onnxruntime::utils::TensorUtils::UnpackTensor(tensor_proto, p_data, tensor_size)); + *p_tensor = std::make_unique(DataTypeImpl::GetType(), + tensor_shape, + static_cast(p_data), + alloc->Info(), + preallocated ? nullptr : alloc); // no deleter for preallocated + + return common::Status::OK(); +} + Status TensorProtoToMLValue(const ONNX_NAMESPACE::TensorProto& input, AllocatorPtr allocator, void* preallocated, size_t preallocated_size, MLValue& value) { std::unique_ptr p_tensor; @@ -174,6 +205,7 @@ common::Status GetTensorFromTensorProto(const TensorProto& tensor_proto, CASE_PROTO(UINT64, uint64_t); CASE_PROTO(STRING, std::string); CASE_PROTO(FLOAT16, MLFloat16); + CASE_PROTO(BFLOAT16, BFloat16); default: { std::ostringstream ostr; ostr << "Initialized tensor with unexpected type: " << tensor_proto.data_type(); @@ -208,6 +240,10 @@ TensorProto::DataType GetTensorProtoType(const Tensor& tensor) { dtype = TensorProto_DataType_UINT64; else if (tensor_type == DataTypeImpl::GetType()) dtype = TensorProto_DataType_BOOL; + else if (tensor_type == DataTypeImpl::GetType()) + dtype = TensorProto_DataType_FLOAT16; + else if (tensor_type == DataTypeImpl::GetType()) + dtype = TensorProto_DataType_BFLOAT16; return dtype; } diff --git a/onnxruntime/core/framework/tensorutils.cc b/onnxruntime/core/framework/tensorutils.cc index 7723e12168..b2b800fd60 100644 --- a/onnxruntime/core/framework/tensorutils.cc +++ b/onnxruntime/core/framework/tensorutils.cc @@ -45,38 +45,38 @@ static void UnpackTensorWithRawData(const ONNX_NAMESPACE::TensorProto& tensor, / namespace onnxruntime { namespace utils { -#define DEFINE_UNPACK_TENSOR(T, Type, field_name, field_size) \ - template <> \ - Status TensorUtils::UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, /*out*/ T* p_data, int64_t expected_size) { \ - if (nullptr == p_data) { \ - const size_t size = tensor.has_raw_data() ? tensor.raw_data().size() : tensor.field_size(); \ - if (size == 0) \ - return Status::OK(); \ - else \ - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); \ - } \ - if (nullptr == p_data || Type != tensor.data_type()) { \ - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); \ - } \ - if (tensor.has_raw_data()) { \ - size_t expected_size_in_bytes; \ - if (!IAllocator::CalcMemSizeForArray(expected_size, sizeof(T), &expected_size_in_bytes)) { \ - return Status(common::ONNXRUNTIME, common::FAIL, "size overflow"); \ - } \ - if (tensor.raw_data().size() != expected_size_in_bytes) \ - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, \ - "UnpackTensor: the pre-allocated size does not match the raw data size, expected ", \ - expected_size_in_bytes, ", got ", tensor.raw_data().size()); \ - UnpackTensorWithRawData(tensor, p_data); \ - return Status::OK(); \ - } \ - if (tensor.field_size() != expected_size) \ - return Status(common::ONNXRUNTIME, common::FAIL, \ - "UnpackTensor: the pre-allocated size does not match the size in proto"); \ - auto& data = tensor.field_name(); \ - for (auto data_iter = data.cbegin(); data_iter != data.cend(); ++data_iter) \ - *p_data++ = *reinterpret_cast(data_iter); \ - return Status::OK(); \ +#define DEFINE_UNPACK_TENSOR(T, Type, field_name, field_size) \ + template <> \ + Status TensorUtils::UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, /*out*/ T* p_data, int64_t expected_size) { \ + if (nullptr == p_data) { \ + const size_t size = tensor.has_raw_data() ? tensor.raw_data().size() : tensor.field_size(); \ + if (size == 0) \ + return Status::OK(); \ + else \ + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); \ + } \ + if (nullptr == p_data || Type != tensor.data_type()) { \ + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); \ + } \ + if (tensor.has_raw_data()) { \ + size_t expected_size_in_bytes; \ + if (!IAllocator::CalcMemSizeForArray(expected_size, sizeof(T), &expected_size_in_bytes)) { \ + return Status(common::ONNXRUNTIME, common::FAIL, "size overflow"); \ + } \ + if (tensor.raw_data().size() != expected_size_in_bytes) \ + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, \ + "UnpackTensor: the pre-allocated size does not match the raw data size, expected ", \ + expected_size_in_bytes, ", got ", tensor.raw_data().size()); \ + UnpackTensorWithRawData(tensor, p_data); \ + return Status::OK(); \ + } \ + if (tensor.field_size() != expected_size) \ + return Status(common::ONNXRUNTIME, common::FAIL, \ + "UnpackTensor: the pre-allocated size does not match the size in proto"); \ + auto& data = tensor.field_name(); \ + for (auto data_iter = data.cbegin(); data_iter != data.cend(); ++data_iter) \ + *p_data++ = *reinterpret_cast(data_iter); \ + return Status::OK(); \ } //TODO: uint32 uint64 complex64 complex128 @@ -190,6 +190,46 @@ Status TensorUtils::UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, return Status::OK(); } +template <> +Status TensorUtils::UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, + /*out*/ BFloat16* p_data, + int64_t expected_size) { + if (nullptr == p_data) { + const size_t size = tensor.has_raw_data() ? tensor.raw_data().size() : tensor.int32_data_size(); + if (size == 0) + return Status::OK(); + else + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); + } + if (ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16 != tensor.data_type()) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); + } + + if (tensor.has_raw_data()) { + if (tensor.raw_data().size() != (expected_size) * sizeof(uint16_t)) + return Status(common::ONNXRUNTIME, common::FAIL, + "UnpackTensor: the pre-allocate size does not match the raw data size"); + + UnpackTensorWithRawData(tensor, p_data); + return Status::OK(); + } + + if (tensor.int32_data_size() != expected_size) + return Status(common::ONNXRUNTIME, common::FAIL, + "UnpackTensor: the pre-allocate size does not match the size in proto"); + + const int max_value = std::numeric_limits::max(); + for (int i = 0; i < static_cast(expected_size); i++) { + int v = tensor.int32_data()[i]; + if (v < 0 || v > max_value) { + return Status(common::ONNXRUNTIME, common::FAIL, "data overflow"); + } + p_data[i] = BFloat16(static_cast(v)); + } + + return Status::OK(); +} + #define CASE_PROTO_TRACE(X, Y) \ case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \ if (!IAllocator::CalcMemSizeForArrayWithAlignment(size, sizeof(Y), out)) { \ @@ -222,6 +262,7 @@ common::Status GetSizeInBytesFromTensorProto(const ONNX_NAMESPACE::TensorProto& CASE_PROTO_TRACE(UINT32, uint32_t); CASE_PROTO_TRACE(UINT64, uint64_t); CASE_PROTO_TRACE(FLOAT16, MLFloat16); + CASE_PROTO_TRACE(BFLOAT16, BFloat16); case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_STRING: default: return common::Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED); diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 627821cb09..07ddb8e837 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -26,6 +26,7 @@ #include "abi_session_options_impl.h" using namespace onnxruntime::logging; +using onnxruntime::BFloat16; using onnxruntime::DataTypeImpl; using onnxruntime::Environment; using onnxruntime::IAllocator; @@ -259,6 +260,9 @@ ORT_API_STATUS_IMPL(OrtCreateTensorWithDataAsOrtValue, _In_ const OrtAllocatorIn case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, info, p_data, p_data_len, &tensor)); break; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: + ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, info, p_data, p_data_len, &tensor)); + break; case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, info, p_data, p_data_len, &tensor)); break; @@ -322,6 +326,9 @@ ORT_API_STATUS_IMPL(OrtCreateTensorAsOrtValue, _Inout_ OrtAllocator* allocator, case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, allocator, &tensor)); break; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: + ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, allocator, &tensor)); + break; case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: ORT_API_RETURN_IF_ERROR(CreateTensorImpl(shape, shape_len, allocator, &tensor)); break; @@ -333,7 +340,6 @@ ORT_API_STATUS_IMPL(OrtCreateTensorAsOrtValue, _Inout_ OrtAllocator* allocator, break; case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64: case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128: - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: default: { std::ostringstream oss; oss << "type " << type << " is not supported in this function"; diff --git a/onnxruntime/test/framework/data_types_test.cc b/onnxruntime/test/framework/data_types_test.cc index 2ebda7c4bb..6f6b13142e 100644 --- a/onnxruntime/test/framework/data_types_test.cc +++ b/onnxruntime/test/framework/data_types_test.cc @@ -351,6 +351,42 @@ TEST_F(DataTypeTest, VectorMapInt64ToFloatTest) { EXPECT_FALSE(DataTypeImpl::GetType()->IsCompatible(tensor_type)); } +TEST_F(DataTypeTest, BFloat16Test) { + // Test data type + { + const float sample = 1.0f; + BFloat16 flt16(sample); + auto int_rep = flt16.val; + BFloat16 flt_from_int(int_rep); + const double diff = fabs(sample - flt_from_int.ToFloat()); + if (diff > FLT_EPSILON || (std::isnan(diff) && !std::isnan(sample))) { + EXPECT_TRUE(false); + } + } + // Test bulk conversion + { + float sample[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f}; + BFloat16 converted[sizeof(sample) / sizeof(float)]; + static_assert(sizeof(sample) / sizeof(float) == sizeof(converted) / sizeof(BFloat16), "Must have the same count"); + FloatToBFloat16(sample, converted, sizeof(sample) / sizeof(float)); + for (size_t i = 0; i < sizeof(sample) / sizeof(float); ++i) { + const double diff = fabs(sample[i] - converted[i].ToFloat()); + if (diff > FLT_EPSILON || (std::isnan(diff) && !std::isnan(sample[i]))) { + EXPECT_TRUE(false); + } + } + + float back_converted[sizeof(sample) / sizeof(float)]; + BFloat16ToFloat(converted, back_converted, sizeof(sample) / sizeof(float)); + for (size_t i = 0; i < sizeof(sample) / sizeof(float); ++i) { + const double diff = fabs(sample[i] - back_converted[i]); + if (diff > FLT_EPSILON || (std::isnan(diff) && !std::isnan(sample[i]))) { + EXPECT_TRUE(false); + } + } + } +} + TEST_F(DataTypeTest, DataUtilsTest) { using namespace ONNX_NAMESPACE::Utils; // Test Tensor @@ -367,6 +403,20 @@ TEST_F(DataTypeTest, DataUtilsTest) { const auto& from_dt_proto = DataTypeUtils::ToTypeProto(ten_dt); EXPECT_TRUE(DataTypeImpl::GetTensorType()->IsCompatible(from_dt_proto)); } + // Test Tensor with bfloat16 + { + const std::string tensor_uint64("tensor(bfloat16)"); + const auto* ten_proto = DataTypeImpl::GetTensorType()->GetTypeProto(); + EXPECT_NE(ten_proto, nullptr); + DataType ten_dt = DataTypeUtils::ToType(*ten_proto); + EXPECT_NE(ten_dt, nullptr); + EXPECT_EQ(tensor_uint64, *ten_dt); + DataType ten_from_str = DataTypeUtils::ToType(*ten_dt); + // Expect internalized strings + EXPECT_EQ(ten_dt, ten_from_str); + const auto& from_dt_proto = DataTypeUtils::ToTypeProto(ten_dt); + EXPECT_TRUE(DataTypeImpl::GetTensorType()->IsCompatible(from_dt_proto)); + } // SparseTensor // Currently test only with proto, no MLDataType yet. { diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index 98e368f0fd..a4e5080f5e 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -90,7 +90,7 @@ void Check(const OpTester::Data& expected_data, const Tensor& output_tensor, con "] did not match run output shape [" + output_tensor.Shape().ToString() + "] for " + expected_data.def_.Name()); - CheckDispatch(output_tensor.DataType(), expected_data, output_tensor, provider_type); + CheckDispatch(output_tensor.DataType(), expected_data, output_tensor, provider_type); } // Check for non tensor types diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index e9287c6cfe..b2056b856a 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -83,6 +83,9 @@ constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType() { r template <> constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType() { return ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; } +template <> +constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType() { return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; } + template struct TTypeProto : ONNX_NAMESPACE::TypeProto { TTypeProto(const std::vector* shape = nullptr) { @@ -149,8 +152,8 @@ class OpTester { // Set whether the NodeArg created by AddInput/AddOutput should include shape information // for Tensor types. If not added, shape inferencing should resolve. If added, shape inferencing - // should validate. Default is to not add. - // Additionally a symbolic dimension will be added if symbolic_dim matches a dimension in the input. + // should validate. Default is to not add. + // Additionally a symbolic dimension will be added if symbolic_dim matches a dimension in the input. OpTester& AddShapeToTensorData(bool add_shape = true, int symbolic_dim = -1) { add_shape_to_tensor_data_ = add_shape; add_symbolic_dim_to_tensor_data_ = symbolic_dim; @@ -271,7 +274,7 @@ class OpTester { try { TensorShape shape{dims}; ORT_ENFORCE(shape.Size() == values_count, values_count, - " input values doesn't match tensor size of ", shape.Size()); + " input values doesn't match tensor size of ", shape.Size()); auto allocator = test::AllocatorManager::Instance().GetAllocator(CPU); auto size_in_bytes = values_count * sizeof(T); diff --git a/onnxruntime/test/util/compare_mlvalue.cc b/onnxruntime/test/util/compare_mlvalue.cc index 24459e5e80..6a2dc1f7d6 100644 --- a/onnxruntime/test/util/compare_mlvalue.cc +++ b/onnxruntime/test/util/compare_mlvalue.cc @@ -13,11 +13,11 @@ using namespace onnxruntime; #if (!EIGEN_VERSION_AT_LEAST(3, 3, 6)) - namespace Eigen { - namespace half_impl { - using __half_raw = ::Eigen::half_impl::__half; - } - } +namespace Eigen { +namespace half_impl { +using __half_raw = ::Eigen::half_impl::__half; +} +} // namespace Eigen #endif #define CASE_TYPE(X) \ @@ -131,6 +131,29 @@ std::pair CompareFloat16Result(const Tensor& outval return std::make_pair(COMPARE_RESULT::SUCCESS, ""); } +std::pair CompareBFloat16Result(const Tensor& outvalue, const Tensor& expected_value, + double per_sample_tolerance, + double relative_per_sample_tolerance, + bool post_processing) { + const size_t size1 = expected_value.Shape().Size(); + const BFloat16* expected_output = expected_value.template Data(); + const BFloat16* real_output = outvalue.template Data(); + for (size_t di = 0; di != size1; ++di) { + float expected = expected_output[di].ToFloat(); + float real = real_output[di].ToFloat(); + real = post_processing ? std::max(0.0f, std::min(255.0f, real)) : real; + const double diff = fabs(expected - real); + const double rtol = per_sample_tolerance + relative_per_sample_tolerance * fabs(expected); + if (diff > rtol || (std::isnan(diff) && !std::isnan(expected))) { + std::ostringstream oss; + oss << "expected " << expected << ", got " << real << ", diff: " << diff << ", tol=" << rtol; + + return std::make_pair(COMPARE_RESULT::RESULT_DIFFERS, oss.str()); + } + } + return std::make_pair(COMPARE_RESULT::SUCCESS, ""); +} + std::pair CompareTwoTensors(const Tensor& outvalue, const Tensor& expected_tensor, double per_sample_tolerance, double relative_per_sample_tolerance, @@ -170,6 +193,9 @@ std::pair CompareTwoTensors(const Tensor& outvalue, } else if (p1 == DataTypeImpl::GetType()) { return CompareFloat16Result(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, post_processing); + } else if (p1 == DataTypeImpl::GetType()) { + return CompareBFloat16Result(outvalue, expected_tensor, + per_sample_tolerance, relative_per_sample_tolerance, post_processing); } else { return std::make_pair(COMPARE_RESULT::NOT_SUPPORT, ""); } @@ -259,6 +285,8 @@ const char* ElementTypeToString(MLDataType type) { else if (type == DataTypeImpl::GetType()) { return "tensor(MLFloat16)"; + } else if (type == DataTypeImpl::GetType()) { + return "tensor(bfloat16)"; } else { return "unknown"; } diff --git a/tools/ci_build/github/linux/docker/scripts/install_deps.sh b/tools/ci_build/github/linux/docker/scripts/install_deps.sh index d092283f78..31f989d86b 100755 --- a/tools/ci_build/github/linux/docker/scripts/install_deps.sh +++ b/tools/ci_build/github/linux/docker/scripts/install_deps.sh @@ -33,8 +33,8 @@ else #Install ONNX #5af210ca8a1c73aa6bae8754c9346ec54d0a756e is v1.2.3 #bae6333e149a59a3faa9c4d9c44974373dcf5256 is v1.3.0 - #0c8d857bb162431912b255d5c0e773fb7c131a65 is v1.3.0 latest - for onnx_version in "5af210ca8a1c73aa6bae8754c9346ec54d0a756e" "bae6333e149a59a3faa9c4d9c44974373dcf5256" "0c8d857bb162431912b255d5c0e773fb7c131a65"; do + #0a7cc483eb0c34e15414437bbbb420f52df4d8c2 is v1.3.0 latest + for onnx_version in "5af210ca8a1c73aa6bae8754c9346ec54d0a756e" "bae6333e149a59a3faa9c4d9c44974373dcf5256" "0a7cc483eb0c34e15414437bbbb420f52df4d8c2"; do if [ -z ${lastest_onnx_version+x} ]; then echo "first pass"; else