vectorize squaredNorm() for complex types

This commit is contained in:
Charles Schlosser 2024-08-21 10:54:17 +00:00 committed by Rasmus Munk Larsen
parent 32d95bb097
commit 87239e058a
2 changed files with 35 additions and 1 deletions

View File

@ -41,6 +41,20 @@ struct dot_nocheck<T, U, true> {
} }
}; };
template <typename Derived, typename Scalar = typename traits<Derived>::Scalar>
struct squared_norm_impl {
using Real = typename NumTraits<Scalar>::Real;
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Real run(const Derived& a) {
Scalar result = a.unaryExpr(squared_norm_functor<Scalar>()).sum();
return numext::real(result) + numext::imag(result);
}
};
template <typename Derived>
struct squared_norm_impl<Derived, bool> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool run(const Derived& a) { return a.any(); }
};
} // end namespace internal } // end namespace internal
/** \fn MatrixBase::dot /** \fn MatrixBase::dot
@ -85,7 +99,7 @@ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
template <typename Derived> template <typename Derived>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename NumTraits<typename internal::traits<Derived>::Scalar>::Real EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename NumTraits<typename internal::traits<Derived>::Scalar>::Real
MatrixBase<Derived>::squaredNorm() const { MatrixBase<Derived>::squaredNorm() const {
return numext::real((*this).cwiseAbs2().sum()); return internal::squared_norm_impl<Derived>::run(derived());
} }
/** \returns, for vectors, the \em l2 norm of \c *this, and for matrices the Frobenius norm. /** \returns, for vectors, the \em l2 norm of \c *this, and for matrices the Frobenius norm.

View File

@ -103,6 +103,26 @@ struct functor_traits<scalar_abs2_op<Scalar>> {
enum { Cost = NumTraits<Scalar>::MulCost, PacketAccess = packet_traits<Scalar>::HasAbs2 }; enum { Cost = NumTraits<Scalar>::MulCost, PacketAccess = packet_traits<Scalar>::HasAbs2 };
}; };
template <typename Scalar, bool IsComplex = NumTraits<Scalar>::IsComplex>
struct squared_norm_functor {
typedef Scalar result_type;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(const Scalar& a) const {
return Scalar(numext::real(a) * numext::real(a), numext::imag(a) * numext::imag(a));
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(const Packet& a) const {
return Packet(pmul(a.v, a.v));
}
};
template <typename Scalar>
struct squared_norm_functor<Scalar, false> : scalar_abs2_op<Scalar> {};
template <typename Scalar>
struct functor_traits<squared_norm_functor<Scalar>> {
using Real = typename NumTraits<Scalar>::Real;
enum { Cost = NumTraits<Real>::MulCost, PacketAccess = packet_traits<Real>::HasMul };
};
/** \internal /** \internal
* \brief Template functor to compute the conjugate of a complex value * \brief Template functor to compute the conjugate of a complex value
* *