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

Vectorize cast

Browse Source
This commit is contained in:
Charles Schlosser
2023-04-26 02:50:13 +00:00
committed by Rasmus Munk Larsen
parent 3918768be1
commit eb5ff1861a
8 changed files with 254 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
Eigen/src/Core/CoreEvaluators.h
View File

@@ -621,6 +621,110 @@ protected:
Data m_d;
};
// ----------------------- Casting ---------------------
template <typename SrcType, typename DstType, typename ArgType>
struct unary_evaluator<CwiseUnaryOp<scalar_cast_op<SrcType, DstType>, ArgType>, IndexBased> {
using CastOp = scalar_cast_op<SrcType, DstType>;
using XprType = CwiseUnaryOp<CastOp, ArgType>;
using SrcPacketType = typename packet_traits<SrcType>::type;
static constexpr int SrcPacketSize = packet_traits<SrcType>::size;
static constexpr int SrcPacketSizeBytes = SrcPacketSize * sizeof(SrcType);
enum {
CoeffReadCost = int(evaluator<ArgType>::CoeffReadCost) + int(functor_traits<CastOp>::Cost),
Flags = evaluator<ArgType>::Flags &
(HereditaryBits | LinearAccessBit | (functor_traits<CastOp>::PacketAccess ? PacketAccessBit : 0)),
Alignment = evaluator<ArgType>::Alignment
};
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE explicit unary_evaluator(const XprType& xpr)
: m_argImpl(xpr.nestedExpression()) {
EIGEN_INTERNAL_CHECK_COST_VALUE(functor_traits<CastOp>::Cost);
EIGEN_INTERNAL_CHECK_COST_VALUE(CoeffReadCost);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DstType coeff(Index row, Index col) const {
return CastOp()(m_argImpl.coeff(row, col));
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DstType coeff(Index index) const {
return CastOp()(m_argImpl.coeff(index));
}
template <int LoadMode>
EIGEN_ALWAYS_INLINE SrcPacketType srcPacket(Index row, Index col, Index offset) const {
EIGEN_STATIC_ASSERT((LoadMode & (LoadMode - 1)) == 0, LoadMode must be a power of two)
constexpr bool ArgIsRowMajor = evaluator<ArgType>::Flags & RowMajorBit;
return m_argImpl.template packet<LoadMode, SrcPacketType>(ArgIsRowMajor ? row : row + (offset * SrcPacketSize),
ArgIsRowMajor ? col + (offset * SrcPacketSize) : col);
}
template <int LoadMode>
EIGEN_ALWAYS_INLINE SrcPacketType srcPacket(Index index, Index offset) const {
EIGEN_STATIC_ASSERT((LoadMode & (LoadMode - 1)) == 0, LoadMode must be a power of two)
return m_argImpl.template packet<LoadMode, SrcPacketType>(index + (offset * SrcPacketSize));
}
template <typename DstPacketType>
using SrcPacketArgs1 = std::enable_if_t<unpacket_traits<DstPacketType>::size <= (1 * SrcPacketSize), bool>;
template <typename DstPacketType>
using SrcPacketArgs2 = std::enable_if_t<unpacket_traits<DstPacketType>::size == (2 * SrcPacketSize), bool>;
template <typename DstPacketType>
using SrcPacketArgs4 = std::enable_if_t<unpacket_traits<DstPacketType>::size == (4 * SrcPacketSize), bool>;
template <int LoadMode, typename DstPacketType, SrcPacketArgs1<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index row, Index col) const {
constexpr int DstPacketSize = unpacket_traits<DstPacketType>::size;
constexpr int SrcIncrementBytes = DstPacketSize * sizeof(SrcType);
constexpr int SrcLoadMode = plain_enum_min(SrcIncrementBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(srcPacket<SrcLoadMode>(row, col, 0));
}
template <int LoadMode, typename DstPacketType, SrcPacketArgs2<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index row, Index col) const {
constexpr int SrcLoadMode0 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode1 = plain_enum_min(1 * SrcPacketSizeBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(srcPacket<SrcLoadMode0>(row, col, 0),
srcPacket<SrcLoadMode1>(row, col, 1));
}
template <int LoadMode, typename DstPacketType, SrcPacketArgs4<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index row, Index col) const {
constexpr int SrcLoadMode0 = plain_enum_min(4 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode1 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode2 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode3 = plain_enum_min(1 * SrcPacketSizeBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(
srcPacket<SrcLoadMode0>(row, col, 0), srcPacket<SrcLoadMode1>(row, col, 1),
srcPacket<SrcLoadMode2>(row, col, 2), srcPacket<SrcLoadMode3>(row, col, 3));
}
template <int LoadMode, typename DstPacketType, SrcPacketArgs1<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index index) const {
constexpr int DstPacketSize = unpacket_traits<DstPacketType>::size;
constexpr int SrcIncrementBytes = DstPacketSize * sizeof(SrcType);
constexpr int SrcLoadMode = plain_enum_min(SrcIncrementBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(srcPacket<SrcLoadMode>(index, 0));
}
template <int LoadMode, typename DstPacketType, SrcPacketArgs2<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index index) const {
constexpr int SrcLoadMode0 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode1 = plain_enum_min(1 * SrcPacketSizeBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(srcPacket<SrcLoadMode0>(index, 0),
srcPacket<SrcLoadMode1>(index, 1));
}
template <int LoadMode, typename DstPacketType, SrcPacketArgs4<DstPacketType> = true>
EIGEN_STRONG_INLINE DstPacketType packet(Index index) const {
constexpr int SrcLoadMode0 = plain_enum_min(4 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode1 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode2 = plain_enum_min(2 * SrcPacketSizeBytes, LoadMode);
constexpr int SrcLoadMode3 = plain_enum_min(1 * SrcPacketSizeBytes, LoadMode);
return CastOp().template packetOp<DstPacketType>(
srcPacket<SrcLoadMode0>(index, 0), srcPacket<SrcLoadMode1>(index, 1),
srcPacket<SrcLoadMode2>(index, 2), srcPacket<SrcLoadMode3>(index, 3));
}
protected:
const evaluator<ArgType> m_argImpl;
};
// -------------------- CwiseTernaryOp --------------------
// this is a ternary expression

6
Eigen/src/Core/MathFunctions.h
View File

@@ -430,6 +430,12 @@ struct cast_impl
}
};
template <typename OldType>
struct cast_impl<OldType, bool> {
EIGEN_DEVICE_FUNC
static inline bool run(const OldType& x) { return x != OldType(0); }
};
// Casting from S -> Complex<T> leads to an implicit conversion from S to T,
// generating warnings on clang. Here we explicitly cast the real component.
template<typename OldType, typename NewType>

13
Eigen/src/Core/arch/AVX/TypeCasting.h
View File

@@ -62,6 +62,15 @@ struct type_casting_traits<float, bool> {
TgtCoeffRatio = 1
};
};
template <>
struct type_casting_traits<float, double> {
enum {
VectorizedCast = 1,
SrcCoeffRatio = 1,
TgtCoeffRatio = 2
};
};
#endif // EIGEN_VECTORIZE_AVX512
template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet8f, Packet8i>(const Packet8f& a) {
@@ -80,6 +89,10 @@ template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet4d, Packet8i>(const Packet4d
return _mm256_set_m128i(_mm256_cvttpd_epi32(b), _mm256_cvttpd_epi32(a));
}
template<> EIGEN_STRONG_INLINE Packet4d pcast<Packet8f, Packet4d>(const Packet8f& a) {
return _mm256_cvtps_pd(_mm256_castps256_ps128(a));
}
template <>
EIGEN_STRONG_INLINE Packet16b pcast<Packet8f, Packet16b>(const Packet8f& a,
const Packet8f& b) {

45
Eigen/src/Core/arch/Default/Half.h
View File

@@ -1014,4 +1014,49 @@ struct hash<Eigen::half> {
} // end namespace std
#endif
namespace Eigen {
namespace internal {
template <>
struct cast_impl<float, half> {
EIGEN_DEVICE_FUNC
static inline half run(const float& a) {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __float2half(a);
#else
return Eigen::half(a);
#endif
}
};
template <>
struct cast_impl<int, half> {
EIGEN_DEVICE_FUNC
static inline half run(const int& a) {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __float2half(static_cast<float>(a));
#else
return half(static_cast<float>(a));
#endif
}
};
template <>
struct cast_impl<half, float> {
EIGEN_DEVICE_FUNC
static inline float run(const half& a) {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __half2float(a);
#else
return static_cast<float>(a);
#endif
}
};
} // namespace internal
} // namespace Eigen
#endif // EIGEN_HALF_H

116
Eigen/src/Core/arch/Default/TypeCasting.h
View File

@@ -1,116 +0,0 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2016 Benoit Steiner <benoit.steiner.goog@gmail.com>
// Copyright (C) 2019 Rasmus Munk Larsen <rmlarsen@google.com>
//
// 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_GENERIC_TYPE_CASTING_H
#define EIGEN_GENERIC_TYPE_CASTING_H
#include "../../InternalHeaderCheck.h"
namespace Eigen {
namespace internal {
template<>
struct scalar_cast_op<float, Eigen::half> {
typedef Eigen::half result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::half operator() (const float& a) const {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __float2half(a);
#else
return Eigen::half(a);
#endif
}
};
template<>
struct functor_traits<scalar_cast_op<float, Eigen::half> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
template<>
struct scalar_cast_op<int, Eigen::half> {
typedef Eigen::half result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::half operator() (const int& a) const {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __float2half(static_cast<float>(a));
#else
return Eigen::half(static_cast<float>(a));
#endif
}
};
template<>
struct functor_traits<scalar_cast_op<int, Eigen::half> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
template<>
struct scalar_cast_op<Eigen::half, float> {
typedef float result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float operator() (const Eigen::half& a) const {
#if (defined(EIGEN_HAS_CUDA_FP16) && defined(EIGEN_CUDA_ARCH) && EIGEN_CUDA_ARCH >= 300) || \
(defined(EIGEN_HAS_HIP_FP16) && defined(EIGEN_HIP_DEVICE_COMPILE))
return __half2float(a);
#else
return static_cast<float>(a);
#endif
}
};
template<>
struct functor_traits<scalar_cast_op<Eigen::half, float> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
template<>
struct scalar_cast_op<float, Eigen::bfloat16> {
typedef Eigen::bfloat16 result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::bfloat16 operator() (const float& a) const {
return Eigen::bfloat16(a);
}
};
template<>
struct functor_traits<scalar_cast_op<float, Eigen::bfloat16> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
template<>
struct scalar_cast_op<int, Eigen::bfloat16> {
typedef Eigen::bfloat16 result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::bfloat16 operator() (const int& a) const {
return Eigen::bfloat16(static_cast<float>(a));
}
};
template<>
struct functor_traits<scalar_cast_op<int, Eigen::bfloat16> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
template<>
struct scalar_cast_op<Eigen::bfloat16, float> {
typedef float result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float operator() (const Eigen::bfloat16& a) const {
return static_cast<float>(a);
}
};
template<>
struct functor_traits<scalar_cast_op<Eigen::bfloat16, float> >
{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
}
}
#endif // EIGEN_GENERIC_TYPE_CASTING_H

40
Eigen/src/Core/functors/UnaryFunctors.h
View File

@@ -173,22 +173,40 @@ struct functor_traits<scalar_carg_op<Scalar>> {
*
* \sa class CwiseUnaryOp, MatrixBase::cast()
*/
template<typename Scalar, typename NewType>
template <typename SrcType, typename DstType>
struct scalar_cast_op {
typedef NewType result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const NewType operator() (const Scalar& a) const { return cast<Scalar, NewType>(a); }
using result_type = DstType;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DstType operator()(const SrcType& a) const {
return cast<SrcType, DstType>(a);
}
using SrcPacket = typename packet_traits<SrcType>::type;
template <typename DstPacket>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DstPacket packetOp(const SrcPacket& a) const {
return pcast<SrcPacket, DstPacket>(a);
}
template <typename DstPacket>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DstPacket packetOp(const SrcPacket& a, const SrcPacket& b) const {
return pcast<SrcPacket, DstPacket>(a, b);
}
template <typename DstPacket>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DstPacket packetOp(const SrcPacket& a, const SrcPacket& b,
const SrcPacket& c, const SrcPacket& d) const {
return pcast<SrcPacket, DstPacket>(a, b, c, d);
}
};
template <typename Scalar>
struct scalar_cast_op<Scalar, bool> {
typedef bool result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator()(const Scalar& a) const { return a != Scalar(0); }
template <typename SrcType, typename DstType>
struct functor_traits<scalar_cast_op<SrcType, DstType>> {
enum {
Cost = is_same<SrcType, DstType>::value ? 0 : NumTraits<DstType>::AddCost,
PacketAccess = (type_casting_traits<SrcType, DstType>::VectorizedCast != 0) &&
(type_casting_traits<SrcType, DstType>::SrcCoeffRatio <= 4)
};
};
template<typename Scalar, typename NewType>
struct functor_traits<scalar_cast_op<Scalar,NewType> >
{ enum { Cost = is_same<Scalar, NewType>::value ? 0 : NumTraits<NewType>::AddCost, PacketAccess = false }; };
/** \internal
* \brief Template functor to arithmetically shift a scalar right by a number of bits
*