devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support BFloat16 in Eigen

Browse Source
This commit is contained in:
Teng Lu
2020-06-20 19:16:24 +00:00
committed by Rasmus Munk Larsen
parent 6b9c92fe7e
commit 386d809bde
19 changed files with 1893 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
unsupported/Eigen/CXX11/src/Tensor/TensorRandom.h
View File

@@ -101,6 +101,17 @@ Eigen::half RandomToTypeUniform<Eigen::half>(uint64_t* state, uint64_t stream) {
return result - Eigen::half(1.0f);
}
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
Eigen::bfloat16 RandomToTypeUniform<Eigen::bfloat16>(uint64_t* state, uint64_t stream) {
Eigen::bfloat16 result;
// Generate 7 random bits for the mantissa
unsigned rnd = PCG_XSH_RS_generator(state, stream);
result.value = static_cast<uint16_t>(rnd & 0x7fu);
// Set the exponent
result.value |= (static_cast<uint16_t>(127) << 7);
// Return the final result
return result - Eigen::bfloat16(1.0f);
}
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
float RandomToTypeUniform<float>(uint64_t* state, uint64_t stream) {

2
unsupported/Eigen/SpecialFunctions
View File

@@ -62,6 +62,7 @@ namespace Eigen {
#include "src/SpecialFunctions/BesselFunctionsImpl.h"
#include "src/SpecialFunctions/BesselFunctionsPacketMath.h"
#include "src/SpecialFunctions/BesselFunctionsBFloat16.h"
#include "src/SpecialFunctions/BesselFunctionsHalf.h"
#include "src/SpecialFunctions/BesselFunctionsFunctors.h"
#include "src/SpecialFunctions/BesselFunctionsArrayAPI.h"
@@ -70,6 +71,7 @@ namespace Eigen {
#include "src/SpecialFunctions/HipVectorCompatibility.h"
#endif
#include "src/SpecialFunctions/SpecialFunctionsPacketMath.h"
#include "src/SpecialFunctions/SpecialFunctionsBFloat16.h"
#include "src/SpecialFunctions/SpecialFunctionsHalf.h"
#include "src/SpecialFunctions/SpecialFunctionsFunctors.h"
#include "src/SpecialFunctions/SpecialFunctionsArrayAPI.h"

68
unsupported/Eigen/src/SpecialFunctions/BesselFunctionsBFloat16.h Normal file
View File

@@ -0,0 +1,68 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef EIGEN_BESSELFUNCTIONS_BFLOAT16_H
#define EIGEN_BESSELFUNCTIONS_BFLOAT16_H
namespace Eigen {
namespace numext {
#if EIGEN_HAS_C99_MATH
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_i0(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_i0(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_i0e(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_i0e(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_i1(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_i1(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_i1e(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_i1e(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_j0(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_j0(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_j1(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_j1(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_y0(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_y0(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_y1(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_y1(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_k0(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_k0(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_k0e(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_k0e(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_k1(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_k1(static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bessel_k1e(const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::bessel_k1e(static_cast<float>(x)));
}
#endif
} // end namespace numext
} // end namespace Eigen
#endif // EIGEN_BESSELFUNCTIONS_BFLOAT16_H

58
unsupported/Eigen/src/SpecialFunctions/SpecialFunctionsBFloat16.h Normal file
View File

@@ -0,0 +1,58 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef EIGEN_SPECIALFUNCTIONS_BFLOAT16_H
#define EIGEN_SPECIALFUNCTIONS_BFLOAT16_H
namespace Eigen {
namespace numext {
#if EIGEN_HAS_C99_MATH
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 lgamma(const Eigen::bfloat16& a) {
return Eigen::bfloat16(Eigen::numext::lgamma(static_cast<float>(a)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 digamma(const Eigen::bfloat16& a) {
return Eigen::bfloat16(Eigen::numext::digamma(static_cast<float>(a)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 zeta(const Eigen::bfloat16& x, const Eigen::bfloat16& q) {
return Eigen::bfloat16(Eigen::numext::zeta(static_cast<float>(x), static_cast<float>(q)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 polygamma(const Eigen::bfloat16& n, const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::polygamma(static_cast<float>(n), static_cast<float>(x)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 erf(const Eigen::bfloat16& a) {
return Eigen::bfloat16(Eigen::numext::erf(static_cast<float>(a)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 erfc(const Eigen::bfloat16& a) {
return Eigen::bfloat16(Eigen::numext::erfc(static_cast<float>(a)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 ndtri(const Eigen::bfloat16& a) {
return Eigen::bfloat16(Eigen::numext::ndtri(static_cast<float>(a)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 igamma(const Eigen::bfloat16& a, const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::igamma(static_cast<float>(a), static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 igamma_der_a(const Eigen::bfloat16& a, const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::igamma_der_a(static_cast<float>(a), static_cast<float>(x)));
}
template <>
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 gamma_sample_der_alpha(const Eigen::bfloat16& alpha, const Eigen::bfloat16& sample) {
return Eigen::bfloat16(Eigen::numext::gamma_sample_der_alpha(static_cast<float>(alpha), static_cast<float>(sample)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 igammac(const Eigen::bfloat16& a, const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::igammac(static_cast<float>(a), static_cast<float>(x)));
}
template<> EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 betainc(const Eigen::bfloat16& a, const Eigen::bfloat16& b, const Eigen::bfloat16& x) {
return Eigen::bfloat16(Eigen::numext::betainc(static_cast<float>(a), static_cast<float>(b), static_cast<float>(x)));
}
#endif
} // end namespace numext
} // end namespace Eigen
#endif // EIGEN_SPECIALFUNCTIONS_BFLOAT16_H

1
unsupported/test/cxx11_tensor_reduction.cpp
View File

@@ -511,6 +511,7 @@ EIGEN_DECLARE_TEST(cxx11_tensor_reduction) {
CALL_SUBTEST(( test_simple_reductions<float,ColMajor>() ));
CALL_SUBTEST(( test_simple_reductions<float,RowMajor>() ));
CALL_SUBTEST(( test_simple_reductions<Eigen::half,ColMajor>() ));
CALL_SUBTEST(( test_simple_reductions<Eigen::bfloat16,ColMajor>() ));
CALL_SUBTEST(test_reductions_in_expr<ColMajor>());
CALL_SUBTEST(test_reductions_in_expr<RowMajor>());
CALL_SUBTEST(test_full_reductions<ColMajor>());