wehub-resource-sync
246 lines
8.4 KiB
C++
246 lines
8.4 KiB
C++
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_F16_H_
|
|
#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_F16_H_
|
|
|
|
#include <cstdint>
|
|
#include <type_traits>
|
|
|
|
#include "fp16.h" // from @FP16 // IWYU pragma: keep, used with no builtin float16
|
|
|
|
// Use __FLT16_MAX__ to determine whether _Float16 is builtin
|
|
#if defined(__FLT16_MAX__) && !SHLO_REF_EMULATE_F16
|
|
#define SHLO_REF_HAS_BUILTIN_FLOAT16 1
|
|
#endif
|
|
|
|
namespace shlo_ref {
|
|
|
|
class alignas(uint16_t) F16 {
|
|
public:
|
|
F16() = default;
|
|
|
|
template <typename T,
|
|
typename = std::enable_if_t<std::is_convertible_v<T, float>>>
|
|
// Allow implicit conversions from types convertible to float.
|
|
// NOLINTNEXTLINE(google-explicit-constructor)
|
|
F16(T x);
|
|
|
|
// Tagged constructor to allow construction from bits
|
|
struct bitcast_construct_t {};
|
|
explicit F16(bitcast_construct_t, uint16_t bits) : bits_(bits) {}
|
|
|
|
// Allow implicit conversions to float.
|
|
// NOLINTNEXTLINE(google-explicit-constructor)
|
|
operator float() const;
|
|
|
|
explicit operator bool() const;
|
|
|
|
F16& operator=(float x) { return *this = static_cast<F16>(x); }
|
|
|
|
#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
|
|
#define SHLO_REF_DEFINE_BINARY_OP(OP) \
|
|
friend F16 operator OP(F16 x, F16 y) { return x.native_ OP y.native_; } \
|
|
\
|
|
template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
|
|
friend auto operator OP(F16 x, T y) { \
|
|
return x.native_ OP y; \
|
|
} \
|
|
\
|
|
template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
|
|
friend auto operator OP(T x, F16 y) { \
|
|
return x OP y.native_; \
|
|
}
|
|
|
|
#define SHLO_REF_DEFINE_BINARY_ASSIGN_OP(OP) \
|
|
SHLO_REF_DEFINE_BINARY_OP(OP); \
|
|
friend F16& operator OP##=(F16 & x, F16 y) { \
|
|
x.native_ OP## = y.native_; \
|
|
return x; \
|
|
} \
|
|
\
|
|
template <class T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
|
|
friend F16& operator OP##=(F16 & x, T y) { \
|
|
x.native_ OP## = y; \
|
|
return x; \
|
|
}
|
|
|
|
#else // !SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
|
|
#define SHLO_REF_DEFINE_BINARY_OP(OP) \
|
|
friend F16 operator OP(F16 x, F16 y) { \
|
|
return F16(static_cast<float>(x) OP static_cast<float>(y)); \
|
|
} \
|
|
\
|
|
template <typename T, typename C = std::common_type_t<F16, T>> \
|
|
friend C operator OP(F16 x, T y) { \
|
|
return static_cast<C>(static_cast<C>(x) OP static_cast<C>(y)); \
|
|
} \
|
|
\
|
|
template <typename T, typename C = std::common_type_t<F16, T>> \
|
|
friend C operator OP(T x, F16 y) { \
|
|
return static_cast<C>(static_cast<C>(x) OP static_cast<C>(y)); \
|
|
}
|
|
|
|
#define SHLO_REF_DEFINE_BINARY_ASSIGN_OP(OP) \
|
|
SHLO_REF_DEFINE_BINARY_OP(OP); \
|
|
friend F16& operator OP##=(F16 & x, F16 y) { \
|
|
return x = F16(static_cast<float>(x) OP static_cast<float>(y)); \
|
|
} \
|
|
\
|
|
template <class T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
|
|
friend F16& operator OP##=(F16 & x, T y) { \
|
|
return x = static_cast<float>(x) OP y; \
|
|
}
|
|
|
|
#endif // SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
|
|
friend F16 operator+(F16 x);
|
|
friend F16 operator-(F16 x);
|
|
friend F16& operator++(F16& x);
|
|
friend F16 operator++(F16& x, int);
|
|
friend F16& operator--(F16& x);
|
|
friend F16 operator--(F16& x, int);
|
|
SHLO_REF_DEFINE_BINARY_ASSIGN_OP(+);
|
|
SHLO_REF_DEFINE_BINARY_ASSIGN_OP(-);
|
|
SHLO_REF_DEFINE_BINARY_ASSIGN_OP(*);
|
|
SHLO_REF_DEFINE_BINARY_ASSIGN_OP(/);
|
|
SHLO_REF_DEFINE_BINARY_OP(<);
|
|
SHLO_REF_DEFINE_BINARY_OP(<=);
|
|
SHLO_REF_DEFINE_BINARY_OP(>);
|
|
SHLO_REF_DEFINE_BINARY_OP(>=);
|
|
SHLO_REF_DEFINE_BINARY_OP(==);
|
|
SHLO_REF_DEFINE_BINARY_OP(!=);
|
|
#undef SHLO_REF_DEFINE_BINARY_ASSIGN_OP
|
|
#undef SHLO_REF_DEFINE_BINARY_OP
|
|
|
|
private:
|
|
union {
|
|
#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
_Float16 native_;
|
|
#endif
|
|
|
|
uint16_t bits_;
|
|
};
|
|
};
|
|
|
|
namespace detail {
|
|
|
|
template <class T, class SFINAE = void>
|
|
struct F16CommonType {};
|
|
|
|
template <class T>
|
|
struct F16CommonType<T, std::enable_if_t<std::is_integral_v<T>>> {
|
|
using type = F16;
|
|
};
|
|
|
|
template <class T>
|
|
struct F16CommonType<T, std::enable_if_t<std::is_floating_point_v<T>>> {
|
|
using type = T;
|
|
};
|
|
|
|
template <class T>
|
|
struct F16CommonType<T, std::enable_if_t<!std::is_arithmetic_v<T> &&
|
|
std::is_convertible_v<T, float>>> {
|
|
using type = float;
|
|
};
|
|
|
|
template <class T>
|
|
struct F16CommonType<T, std::enable_if_t<!std::is_arithmetic_v<T> &&
|
|
!std::is_convertible_v<T, float> &&
|
|
std::is_convertible_v<T, double>>> {
|
|
using type = double;
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
} // namespace shlo_ref
|
|
|
|
namespace std {
|
|
|
|
template <>
|
|
struct common_type<shlo_ref::F16, shlo_ref::F16> {
|
|
using type = shlo_ref::F16;
|
|
};
|
|
|
|
template <typename T>
|
|
struct common_type<shlo_ref::F16, T> : shlo_ref::detail::F16CommonType<T> {};
|
|
|
|
template <typename T>
|
|
struct common_type<T, shlo_ref::F16> : shlo_ref::detail::F16CommonType<T> {};
|
|
|
|
} // namespace std
|
|
|
|
namespace shlo_ref {
|
|
static_assert(sizeof(F16) == sizeof(uint16_t));
|
|
static_assert(alignof(F16) == alignof(uint16_t));
|
|
|
|
#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
|
|
template <typename T, typename _>
|
|
F16::F16(T x) : native_(static_cast<_Float16>(static_cast<float>(x))) {}
|
|
|
|
inline F16::operator float() const { return native_; }
|
|
|
|
inline F16::operator bool() const { return native_; }
|
|
|
|
inline F16 operator+(F16 x) { return x.native_; }
|
|
|
|
inline F16 operator-(F16 x) { return -x.native_; }
|
|
|
|
inline F16& operator++(F16& x) { return x += 1; }
|
|
|
|
inline F16 operator++(F16& x, int) { return x.native_++; }
|
|
|
|
inline F16& operator--(F16& x) { return x -= 1; }
|
|
|
|
inline F16 operator--(F16& x, int) { return x.native_--; }
|
|
|
|
#else // !SHLO_REF_HAS_BUILTIN_FLOAT16
|
|
|
|
template <typename T, typename _>
|
|
inline F16::F16(T x)
|
|
: bits_(fp16_ieee_from_fp32_value(static_cast<float>(x))) {}
|
|
|
|
inline F16::operator float() const { return fp16_ieee_to_fp32_value(bits_); }
|
|
|
|
inline F16::operator bool() const { return bits_; }
|
|
|
|
inline F16 operator-(F16 x) { return F16(-static_cast<float>(x)); }
|
|
inline F16 operator+(F16 x) { return F16(static_cast<float>(x)); }
|
|
|
|
inline F16& operator++(F16& x) { return x += 1; }
|
|
|
|
inline F16 operator++(F16& x, int) {
|
|
const F16 y = x;
|
|
++x;
|
|
return y;
|
|
}
|
|
|
|
inline F16& operator--(F16& x) { return x -= 1; }
|
|
|
|
inline F16 operator--(F16& x, int) {
|
|
const F16 y = x;
|
|
--x;
|
|
return y;
|
|
}
|
|
|
|
#endif
|
|
|
|
} // namespace shlo_ref
|
|
|
|
#endif // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_F16_H_
|