[C++, Eigen] Eigen cast함수 SIMD로 벡터화하기

Eigen은 C++기반의 선형 대수 연산 라이브러리입니다. 이 라이브러리의 특징은 c++의 템플릿을 쥐어짜서 컴파일 시간에 행렬 간의 연산식을 분석하고 최적의 연산 순서를 결정해서 연산을 수행한다는 것입니다. 예를 들어

Eigen::MatrixXf a, b, c, d; a = b * c + d;

와 같은 식이 있다면, 일반적인 c++객체에서는 b*c를 연산한 뒤 그 리턴값으로 임시 객체가 생성이 되고, 이 임시객체와 d를 더한뒤 임시 객체를 생성하고, 최종적으로 이 임시객체가 a에 대입되는 식으로 연산이 진행될 겁니다. 하지만 Eigen에서는 Expression Template이라는 템플릿의 응용기법을 이용해 이런 불필요한 연산을 회피하고, b*c를 바로 a에 대입한뒤 거기에 d를 더 더해서 임시객체 생성을 최소화하며 성능은 최대화하는 식으로 작동됩니다.

Eigen은 또한 행렬 연산 말고도 대량의 벡터 연산을 최적화하는데에도 유용합니다. 아주 긴 두 배열을 각각의 원소끼리 더하기 위해서는 반복문을 써야하지만, Eigen의 Array를 사용하면 반복문을 사용하지 않고 깔끔하게 덧셈을 수행할 수 있습니다. 덤으로 SSE나 AVX와 같은 명령어를 활용해서 최신 CPU에서는 최소 2배에서 8배 정도까지 빠른 속도를 낼 수도 있구요.

float a[1024], b[1024], c[1024]; // a = b + c를 하기 위해서는 for(int i = 0; i < 1024; ++i) { a[i] = b[i] + c[i]; } //Eigen을 사용하면 Eigen::VectorXf a(1024), b(1024), c(1024); a.array() = b.array() + c.array();

이러한 강력한 점 때문에 Eigen은 다양한 곳에서 쓰이고 있고, 저도 개인적으로 개발하는 토픽 모델링 라이브러리에 Eigen을 적극 활용하고 있습니다. 그런데 몇 가지 작업을 하던 도중 Eigen의 array가 제대로 벡터화되지 않는 문제를 발견했습니다.

Eigen::VectorXf a, b, c, d; a.array() = (b.array() + 0.1f) / (c.array() + 0.2f) * (d.array() + 0.1f); // 벡터화 잘됨 Eigen::VectorXf a; Eigen::VectorXi b, c, d; a.array() = (b.array().template cast<float>() + 0.1f) / (c.array().template cast<float>() + 0.2f) * (d.array().template cast<float>() + 0.1f); // 벡터화 안됨

위의 식의 경우 벡터화가 아주 잘 됩니다. AVX 옵션을 주고 컴파일하면 한번에 8개의 float를 묶어서 처리하기 때문에 개별로 연산할때보다 6~8배정도 빠르게 수행이 됩니다. 그런데 아래의 식의 경우 벡터화가 되질 않습니다. 즉 SSE나 AVX 옵션을 주고 컴파일해도 거의 성능 향상이 없습니다. 그 이유에 대해 조사를 하던 중 cast함수가 원인이라는 걸 알게 되었죠.

2019년 최신 버전인 Eigen 3.3.7에서 객체들은 cast가 호출될 경우 벡터화를 적용하지 않고 개별로 연산됩니다. 인터넷을 전전하며 원인을 찾아본 결과, 캐스팅을 수행하면서 자료형의 크기가 변경되는 경우(int8_t -> float로 간다던지, float에서 double로 간다던지 등등) 내부에서 처리하는 패킷의 크기가 달라져야 하기 때문에 이를 아직 적용하지 못하고 있는 것으로 보입니다. 아마 차기 버전에서는 멀티 패킷 관련 기능이 추가되면서 해당 문제가 해결될 가능성도 있어보입니다.

문제는 당장 제가 써야하는데, Eigen의 차기버전을 마냥 기다리고만 있을 수는 없다는 것이죠. 그리고 제 문제의 경우 int32_t -> float이므로 자료형의 크기에는 변화가 없어 패킷 크기도 동일하게 유지되기 때문에 크기 불일치도 문제가 되지 않구요. 어떻게든 방법이 있을것 같아서 하루종일 삽질한 결과 다음과 같이 템플릿 클래스 특수화를 통해 문제를 해결했습니다. 손 봐야하는 템플릿 클래스는 Eigen::internal::scalar_cast_op과 Eigen::internal::unary_evaluator 둘 입니다.

#pragma once #include <type_traits> #include <Eigen/Dense> #ifdef EIGEN_VECTORIZE_AVX #include <immintrin.h> // AVX를 사용하는 경우 namespace Eigen { namespace internal { // to_int_packet은 float 타입의 패킷을 int타입의 패킷으로 바꿔줍니다. template<typename PacketType> struct to_int_packet { typedef PacketType type; }; template<> struct to_int_packet<Packet8f> { typedef Packet8i type; }; template<> struct to_int_packet<Packet4f> { typedef Packet4i type; }; EIGEN_STRONG_INLINE Packet8f p_to_f32(const Packet8i& a) { // __m256i를 __m256로 캐스팅합니다. return _mm256_cvtepi32_ps(a); } EIGEN_STRONG_INLINE Packet4f p_to_f32(const Packet4i& a) { // __m128i를 __m128로 캐스팅합니다. return _mm_cvtepi32_ps(a); } } } #elif defined(EIGEN_VECTORIZE_SSE2) // SSE를 사용하는 경우 #include <xmmintrin.h> namespace Eigen { namespace internal { template<typename PacketType> struct to_int_packet { typedef PacketType type; }; template<> struct to_int_packet<Packet4f> { typedef Packet4i type; }; EIGEN_STRONG_INLINE Packet4f p_to_f32(const Packet4i& a) { // __m128i를 __m128로 캐스팅합니다. return _mm_cvtepi32_ps(a); } } } #else // 벡터화가 없는 경우.. template<typename PacketType> struct to_int_packet { typedef PacketType type; }; #endif namespace Eigen { namespace internal { // Eigen::internal::scalar_cast_op이 캐스팅 연산을 실제 수행하는 클래스입니다. // 얘를 int32_t -> float에 대해서만 특수화해서 아래와 같이 사용하겠습니다. template<> struct scalar_cast_op<int32_t, float> { EIGEN_EMPTY_STRUCT_CTOR(scalar_cast_op) typedef float result_type; // 개별로 처리되는 캐스팅 함수 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const float operator() (const int32_t& a) const { return cast<int32_t, float>(a); } // 패킷으로 처리되는(벡터화된) 캐스팅 함수 template<typename Packet> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(const typename to_int_packet<typename std::remove_const<Packet>::type>::type& a) const { return p_to_f32(a); } }; template<> struct functor_traits<scalar_cast_op<int32_t, float> > { enum { Cost = NumTraits<float>::AddCost, PacketAccess = 1 }; }; // scalar_cast_op가 수행될때 그 값을 실제로 평가하는 클래스 // 이 역시 다음과 같이 특수화를 진행하였습니다. template<typename ArgType> struct unary_evaluator<CwiseUnaryOp<scalar_cast_op<int32_t, float>, ArgType>, IndexBased > : evaluator_base<CwiseUnaryOp<scalar_cast_op<int32_t, float>, ArgType> > { typedef CwiseUnaryOp<scalar_cast_op<int32_t, float>, ArgType> XprType; enum { CoeffReadCost = evaluator<ArgType>::CoeffReadCost + functor_traits<scalar_cast_op<int32_t, float>>::Cost, Flags = evaluator<ArgType>::Flags & (HereditaryBits | LinearAccessBit | (functor_traits<scalar_cast_op<int32_t, float>>::PacketAccess ? PacketAccessBit : 0)), Alignment = evaluator<ArgType>::Alignment }; EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE explicit unary_evaluator(const XprType& op) : m_functor(op.functor()), m_argImpl(op.nestedExpression()) { EIGEN_INTERNAL_CHECK_COST_VALUE(NumTraits<float>::AddCost); EIGEN_INTERNAL_CHECK_COST_VALUE(CoeffReadCost); } typedef typename XprType::CoeffReturnType CoeffReturnType; EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index row, Index col) const { return m_functor(m_argImpl.coeff(row, col)); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const { return m_functor(m_argImpl.coeff(index)); } // to_int_packet이 들어간 부분을 고쳤습니다. template<int LoadMode, typename PacketType> EIGEN_STRONG_INLINE PacketType packet(Index row, Index col) const { return m_functor.packetOp<PacketType>(m_argImpl.template packet<LoadMode, typename to_int_packet<PacketType>::type>(row, col)); } // 여기도 마찬가지. template<int LoadMode, typename PacketType> EIGEN_STRONG_INLINE PacketType packet(Index index) const { return m_functor.packetOp<PacketType>(m_argImpl.template packet<LoadMode, typename to_int_packet<PacketType>::type>(index)); } protected: const scalar_cast_op<int32_t, float> m_functor; evaluator<ArgType> m_argImpl; }; } }

이제 <Eigen/Dense> 대신 위에서 만든 헤더파일을 include하여 사용하면 됩니다. 간단하죠? (불행히도 x86나 x86_64가 아닌 arm 등에서는 작동이 안된다는 점 양해바랍니다. 다른 아키텍쳐를 쓰시는 분들은 #ifdef 부분을 수정하여 아키텍쳐에 맞는 명령어로 고쳐 쓰시면 되겠네요.)

이렇게 헤더파일만 고쳐서 AVX 옵션으로 새로 컴파일해본 결과 단순한 연산이 포함된 식에서는 116.13초에서 92.64초로 약 1.3배 속도 향상, 로그 및 로그 감마 함수가 포함된 복잡한 식에서는 718.71초에서 211.82초로 약 3.5배 속도 향상이 있었습니다!! 초월함수가 포함된 식에서 성능 향상이 극적이었는데, 이는 초월함수가 SIMD 명령어를 타지 않고 개별적으로 계산될 경우 처리 비용이 너무 크기 때문인 걸로 보입니다. 사칙 연산 정도야 SIMD 안 써도 빠른 클럭으로 비벼볼 수 있지만, log, sin, cos 같은 함수들은 그렇지 않지요.

어쩌다 Eigen 헤더를 뜯어서 분석한 것을 기회로 삼아 더 복잡한 초월함수들도 SIMD로 구현하고 있습니다. lgamma나 digamma 같은거요. 살짝만 구현했는데도 속도 향상이 느껴져서 구현할 맛이 나네요. 오랜만에 intrinsic 명령어들 다루니 어셈블리하는 느낌도 들고요. (확실히 C++은 완전 low-level 밑바닥에서부터 모래알들을 잘 쌓아올려, 편하게 쓸 수 있는 거대한 구조들을 만들어 낸다는 점에서 매력이 있는것 같습니다.) Eigen 가지고 고생하시는 분들에게 작은 도움이 되었길 바랍니다.