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.
This commit is contained in:
Dmitri Smirnov 2018-12-21 12:40:59 -08:00 committed by GitHub
parent 4e74ffba91
commit 7af1887b33
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 322 additions and 74 deletions

2
cmake/external/onnx vendored

@ -1 +1 @@
Subproject commit 0c8d857bb162431912b255d5c0e773fb7c131a65
Subproject commit 0a7cc483eb0c34e15414437bbbb420f52df4d8c2

View file

@ -5,8 +5,9 @@
#include <string>
#include <stdint.h>
#include <unordered_map>
#include <type_traits>
#include <map>
#include <unordered_map>
#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<uint16_t*>(&v);
if (ort_endian::is_little()) {
val = dst[1];
} else {
val = dst[0];
}
}
float ToFloat() const {
float result;
uint16_t* dst = reinterpret_cast<uint16_t*>(&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<T, H, Tail...> {
template <typename T>
struct IsTensorContainedType : public IsAnyOf<T, float, uint8_t, int8_t, uint16_t, int16_t,
int32_t, int64_t, std::string, bool, MLFloat16,
double, uint32_t, uint64_t> {
double, uint32_t, uint64_t, BFloat16> {
};
/// This template's Get() returns a corresponding MLDataType
@ -203,7 +267,7 @@ struct SetMapTypes {
MLDataType dt = GetMLDataType<V, IsTensorContainedType<V>::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<T, IsTensorContainedType<T>::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<ELEM_TYPE>::Type() { \
static TensorType<ELEM_TYPE> tensor_type; \
@ -493,7 +557,7 @@ class NonOnnxType : public DataTypeImpl {
return TensorType<ELEM_TYPE>::Type(); \
}
#define ORT_REGISTER_MAP(TYPE) \
#define ORT_REGISTER_MAP(TYPE) \
template <> \
MLDataType MapType<TYPE>::Type() { \
static MapType<TYPE> map_type; \
@ -504,7 +568,7 @@ class NonOnnxType : public DataTypeImpl {
return MapType<TYPE>::Type(); \
}
#define ORT_REGISTER_SEQ(TYPE) \
#define ORT_REGISTER_SEQ(TYPE) \
template <> \
MLDataType SequenceType<TYPE>::Type() { \
static SequenceType<TYPE> sequence_type; \
@ -515,25 +579,25 @@ class NonOnnxType : public DataTypeImpl {
return SequenceType<TYPE>::Type(); \
}
#define ORT_REGISTER_NON_ONNX_TYPE(TYPE) \
template <> \
MLDataType NonOnnxType<TYPE>::Type() { \
static NonOnnxType<TYPE> non_onnx_type; \
return &non_onnx_type; \
} \
template <> \
MLDataType DataTypeImpl::GetType<TYPE>() { \
return NonOnnxType<TYPE>::Type(); \
#define ORT_REGISTER_NON_ONNX_TYPE(TYPE) \
template <> \
MLDataType NonOnnxType<TYPE>::Type() { \
static NonOnnxType<TYPE> non_onnx_type; \
return &non_onnx_type; \
} \
template <> \
MLDataType DataTypeImpl::GetType<TYPE>() { \
return NonOnnxType<TYPE>::Type(); \
}
#define ORT_REGISTER_OPAQUE_TYPE(CPPType, Domain, Name) \
template <> \
MLDataType OpaqueType<CPPType, Domain, Name>::Type() { \
static OpaqueType<CPPType, Domain, Name> opaque_type; \
return &opaque_type; \
} \
template <> \
MLDataType DataTypeImpl::GetType<CPPType>() { \
return OpaqueType<CPPType, Domain, Name>::Type(); \
#define ORT_REGISTER_OPAQUE_TYPE(CPPType, Domain, Name) \
template <> \
MLDataType OpaqueType<CPPType, Domain, Name>::Type() { \
static OpaqueType<CPPType, Domain, Name> opaque_type; \
return &opaque_type; \
} \
template <> \
MLDataType DataTypeImpl::GetType<CPPType>() { \
return OpaqueType<CPPType, Domain, Name>::Type(); \
}
} // namespace onnxruntime

View file

@ -85,6 +85,10 @@ template <>
constexpr ONNX_NAMESPACE::TensorProto_DataType ToTensorDataType<uint64_t>() {
return ONNX_NAMESPACE::TensorProto_DataType_UINT64;
};
template <>
constexpr ONNX_NAMESPACE::TensorProto_DataType ToTensorDataType<BFloat16>() {
return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16;
};
template <typename T>
struct TensorContainedTypeSetter<T> {
@ -123,6 +127,8 @@ template struct
TensorContainedTypeSetter<uint32_t>;
template struct
TensorContainedTypeSetter<uint64_t>;
template struct
TensorContainedTypeSetter<BFloat16>;
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<void(MLDataType)>& 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<uint64_t>();
case TensorProto_DataType_FLOAT16:
return DataTypeImpl::GetTensorType<MLFloat16>();
case TensorProto_DataType_BFLOAT16:
return DataTypeImpl::GetTensorType<BFloat16>();
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<MLDataType>& DataTypeImpl::AllFixedSizeTensorTypes() {
static std::vector<MLDataType> all_fixed_size_tensor_types =
@ -699,6 +710,7 @@ const std::vector<MLDataType>& DataTypeImpl::AllFixedSizeTensorTypes() {
DataTypeImpl::GetTensorType<int8_t>(),
DataTypeImpl::GetTensorType<uint8_t>(),
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<BFloat16>(),
DataTypeImpl::GetTensorType<bool>()};
return all_fixed_size_tensor_types;
@ -717,6 +729,7 @@ const std::vector<MLDataType>& DataTypeImpl::AllTensorTypes() {
DataTypeImpl::GetTensorType<int8_t>(),
DataTypeImpl::GetTensorType<uint8_t>(),
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<BFloat16>(),
DataTypeImpl::GetTensorType<bool>(),
DataTypeImpl::GetTensorType<std::string>()};

View file

@ -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<uint64_t>();
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
return DataTypeImpl::GetType<MLFloat16>();
case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16:
return DataTypeImpl::GetType<BFloat16>();
default:
ORT_NOT_IMPLEMENTED(__FUNCTION__, ":tensor type ", type, " is not supported");
}

View file

@ -10,6 +10,7 @@
#include <stdexcept>
#include <atomic>
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<MLFloat16>()) {
type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
} else if (cpp_type == onnxruntime::DataTypeImpl::GetType<BFloat16>()) {
type = ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16;
} else if (cpp_type == onnxruntime::DataTypeImpl::GetType<double>()) {
type = ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE;
} else if (cpp_type == onnxruntime::DataTypeImpl::GetType<uint32_t>()) {

View file

@ -138,6 +138,37 @@ common::Status GetTensorByTypeFromTensorProto<MLFloat16>(const TensorProto& tens
return common::Status::OK();
}
template <>
common::Status GetTensorByTypeFromTensorProto<BFloat16>(const TensorProto& tensor_proto,
const TensorShape& tensor_shape,
std::unique_ptr<Tensor>* 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<size_t>(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<BFloat16*>(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<Tensor>(DataTypeImpl::GetType<BFloat16>(),
tensor_shape,
static_cast<void*>(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<Tensor> 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<bool>())
dtype = TensorProto_DataType_BOOL;
else if (tensor_type == DataTypeImpl::GetType<MLFloat16>())
dtype = TensorProto_DataType_FLOAT16;
else if (tensor_type == DataTypeImpl::GetType<BFloat16>())
dtype = TensorProto_DataType_BFLOAT16;
return dtype;
}

View file

@ -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<const T*>(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<const T*>(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<uint16_t>::max();
for (int i = 0; i < static_cast<int>(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<uint16_t>(v));
}
return Status::OK();
}
#define CASE_PROTO_TRACE(X, Y) \
case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \
if (!IAllocator::CalcMemSizeForArrayWithAlignment<alignment>(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);

View file

@ -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<MLFloat16>(shape, shape_len, info, p_data, p_data_len, &tensor));
break;
case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
ORT_API_RETURN_IF_ERROR(CreateTensorImpl<BFloat16>(shape, shape_len, info, p_data, p_data_len, &tensor));
break;
case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
ORT_API_RETURN_IF_ERROR(CreateTensorImpl<double>(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<MLFloat16>(shape, shape_len, allocator, &tensor));
break;
case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
ORT_API_RETURN_IF_ERROR(CreateTensorImpl<BFloat16>(shape, shape_len, allocator, &tensor));
break;
case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
ORT_API_RETURN_IF_ERROR(CreateTensorImpl<double>(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";

View file

@ -351,6 +351,42 @@ TEST_F(DataTypeTest, VectorMapInt64ToFloatTest) {
EXPECT_FALSE(DataTypeImpl::GetType<VectorMapInt64ToFloat>()->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<uint64_t>()->IsCompatible(from_dt_proto));
}
// Test Tensor with bfloat16
{
const std::string tensor_uint64("tensor(bfloat16)");
const auto* ten_proto = DataTypeImpl::GetTensorType<BFloat16>()->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<BFloat16>()->IsCompatible(from_dt_proto));
}
// SparseTensor
// Currently test only with proto, no MLDataType yet.
{

View file

@ -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<bool, float, double, uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t, std::string, MLFloat16>(output_tensor.DataType(), expected_data, output_tensor, provider_type);
CheckDispatch<bool, float, double, uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t, std::string, MLFloat16, BFloat16>(output_tensor.DataType(), expected_data, output_tensor, provider_type);
}
// Check for non tensor types

View file

@ -83,6 +83,9 @@ constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType<std::string>() { r
template <>
constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType<MLFloat16>() { return ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; }
template <>
constexpr ONNX_NAMESPACE::TensorProto_DataType TypeToDataType<BFloat16>() { return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; }
template <typename T>
struct TTypeProto : ONNX_NAMESPACE::TypeProto {
TTypeProto(const std::vector<int64_t>* 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);

View file

@ -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<COMPARE_RESULT, std::string> CompareFloat16Result(const Tensor& outval
return std::make_pair(COMPARE_RESULT::SUCCESS, "");
}
std::pair<COMPARE_RESULT, std::string> 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<BFloat16>();
const BFloat16* real_output = outvalue.template Data<BFloat16>();
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<COMPARE_RESULT, std::string> CompareTwoTensors(const Tensor& outvalue, const Tensor& expected_tensor,
double per_sample_tolerance,
double relative_per_sample_tolerance,
@ -170,6 +193,9 @@ std::pair<COMPARE_RESULT, std::string> CompareTwoTensors(const Tensor& outvalue,
} else if (p1 == DataTypeImpl::GetType<MLFloat16>()) {
return CompareFloat16Result(outvalue, expected_tensor,
per_sample_tolerance, relative_per_sample_tolerance, post_processing);
} else if (p1 == DataTypeImpl::GetType<BFloat16>()) {
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<MLFloat16>()) {
return "tensor(MLFloat16)";
} else if (type == DataTypeImpl::GetType<BFloat16>()) {
return "tensor(bfloat16)";
} else {
return "unknown";
}

View file

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