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

Use select ternary op in tensor select evaulator

Browse Source
This commit is contained in:
Charles Schlosser
2023-04-18 20:52:16 +00:00
committed by Rasmus Munk Larsen
parent 2b954be663
commit e2bbf496f6
3 changed files with 72 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
unsupported/Eigen/CXX11/src/Tensor/TensorEvaluator.h
View File

@@ -817,14 +817,21 @@ struct TensorEvaluator<const TensorSelectOp<IfArgType, ThenArgType, ElseArgType>
{
typedef TensorSelectOp<IfArgType, ThenArgType, ElseArgType> XprType;
typedef typename XprType::Scalar Scalar;
using TernarySelectOp = internal::scalar_boolean_select_op<typename internal::traits<ThenArgType>::Scalar,
typename internal::traits<ElseArgType>::Scalar,
typename internal::traits<IfArgType>::Scalar>;
static constexpr bool TernaryPacketAccess =
TensorEvaluator<ThenArgType, Device>::PacketAccess && TensorEvaluator<ElseArgType, Device>::PacketAccess &&
TensorEvaluator<IfArgType, Device>::PacketAccess && internal::functor_traits<TernarySelectOp>::PacketAccess;
static constexpr int Layout = TensorEvaluator<IfArgType, Device>::Layout;
enum {
IsAligned = TensorEvaluator<ThenArgType, Device>::IsAligned &
TensorEvaluator<ElseArgType, Device>::IsAligned,
PacketAccess = TensorEvaluator<ThenArgType, Device>::PacketAccess &
TensorEvaluator<ElseArgType, Device>::PacketAccess &
PacketType<Scalar, Device>::HasBlend,
PacketAccess = (TensorEvaluator<ThenArgType, Device>::PacketAccess &&
TensorEvaluator<ElseArgType, Device>::PacketAccess &&
PacketType<Scalar, Device>::HasBlend) || TernaryPacketAccess,
BlockAccess = TensorEvaluator<IfArgType, Device>::BlockAccess &&
TensorEvaluator<ThenArgType, Device>::BlockAccess &&
TensorEvaluator<ElseArgType, Device>::BlockAccess,
@@ -922,7 +929,9 @@ struct TensorEvaluator<const TensorSelectOp<IfArgType, ThenArgType, ElseArgType>
{
return m_condImpl.coeff(index) ? m_thenImpl.coeff(index) : m_elseImpl.coeff(index);
}
template<int LoadMode>
template<int LoadMode, bool UseTernary = TernaryPacketAccess,
std::enable_if_t<!UseTernary, bool> = true>
EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const
{
internal::Selector<PacketSize> select;
@@ -936,6 +945,14 @@ struct TensorEvaluator<const TensorSelectOp<IfArgType, ThenArgType, ElseArgType>
}
template <int LoadMode, bool UseTernary = TernaryPacketAccess,
std::enable_if_t<UseTernary, bool> = true>
EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
return TernarySelectOp().template packetOp<PacketReturnType>(m_thenImpl.template packet<LoadMode>(index),
m_elseImpl.template packet<LoadMode>(index),
m_condImpl.template packet<LoadMode>(index));
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost
costPerCoeff(bool vectorized) const {
return m_condImpl.costPerCoeff(vectorized) +