SIMD를 이용한 깁스 샘플링 연산 최적화(SSE2, AVX)

지난 몇년간 CPU의 성능은 반도체 공정의 향상과 그로 인한 클럭 상승에 따라서 줄곧 높아져왔습니다. 무어의 법칙이 잘 작동하던 2000년대 초반까지는 프로그램 코드를 고치지 않아도 CPU의 클럭 상승으로 자동으로 프로그램의 성능도 높아져올 수 있었습니다. 그러나! 공정 미세화가 차츰 둔해지고, 집적화로 인한 발열을 감당할 수 없게 되자 CPU업계는 CPU의 발전 방향을 약간 틀었는데요, 1) 단일 코어의 성능을 향상시키기보다, 코어의 수를 늘리는 쪽으로, 2) 명령어 수준의 병렬성을 높여 최대한 효율적으로 연산을 처리하는 쪽으로 말이죠.

1번은 멀티 스레딩 프로그래밍을 촉발시켰습니다. 기존의 알고리즘들을 병렬화하는 방향으로 코드를 수정함으로써 프로그램의 성능을 높일수 있게된 것입니다. 깁스 샘플링 알고리즘을 병렬화하는 건 앞선 포스팅([토픽 모델링] 깁스 샘플링의 병렬화 & GPU 위에서 돌리기)에서 소개한 바 있습니다.

반면 2번은 명령어의 다양화를 촉발시켰습니다. CPU업계는 기존의 x86 명령어에 SSE나 AVX 같이 여러 데이터를 동시에 처리할 수 있는 명령어(이런 명령어를 Single Instruction Multiple Data, SIMD라고 합니다)를 추가하였고, 기존의 명령어 여러 개를 1~2개의 SSE나 AVX 명령어로 통합시킬수 있다면 하나의 코어로 훨씬 효율적인 작업을 수행할 수 있겠죠. 최신 CPU들은 SSE는 대부분 지원하고, AVX도 지원하는 CPU가 늘고있는 추세입니다. 문제는 컴파일러인데요, Auto Vectorization라는, 반복되는 연산을 캐치하여 SSE나 AVX 명령어로 변환해주는 기능을 제공하고 있지만, 이렇게 자동으로 SIMD화 해주는 기능이 생각보다는 아직 빈약하다는 것입니다. 결국 최적화를 위해서는 코드 상에서 직접 SIMD 명령어를 지정해줄 필요가 있겠지요.

본 포스팅에서는 기존 LDA의 깁스 샘플링 코드를 SIMD화하여 성능을 끌어올리는 과정 및 결과를 소개합니다.

SIMD 명령어

앞서 말했듯 SIMD 명령어는 여러 데이터에 대한 연산을 하나로 줄일수 있는데, 조금더 구체적으로 살펴보자면 다음과 같습니다.

(a, b, c, d) + (e, f, g, h) = (a+e, b+f, c+g, d+h)

위와 같은 벡터 연산을 생각해봅시다. 두 벡터의 합은 각각의 벡터의 원소끼리의 합으로 정의됩니다. SISD(Single Instruction Single Data, 하나의 명령어로 하나의 데이터를 처리하는 기존의 명령어)에서는 각각의 원소를 따로 더해서 결과를 얻어야 하므로 총 4번의 연산이 필요합니다. 반면 SIMD에서는 위와 같이 4개씩 묶인 벡터의 덧셈을 한번에 처리할 수 있게 지원해주는 명령어가 있습니다. 이를 이용하면 이론상 1번의 연산만으로 덧셈 결과를 구할수 있습니다.

벡터 간의 연산을 비롯하여, 행렬 연산, 기타 수치 연산들을 근본적으로 비슷한 작업을 복수의 데이터에 대해서 반복하는 경우가 많으므로 SIMD를 잘 활용한다면 성능을 끌어올릴수 있겠죠?

SSE에서는 128비트의 레지스터를 지원하여, 총 4개의 float(단정밀도 실수)나 2개의 double(배정밀도 실수)를 동시에 처리할 수 있구요, AVX에서는 256비트의 레지스터를 지원해, 총 8개의 float이나 4개의 double을 동시에 처리할 수 있습니다. SSE나 AVX의 버전에 따라 지원가능한 연산에는 약간의 차이가 있는데 이에 대해서는 https://software.intel.com/sites/landingpage/IntrinsicsGuide/ 에 잘 정리되어 있으니 참고하시길 바랍니다. 

Intrinsic 함수

SIMD 명령어들은 CPU 명령어이므로 어셈블리어가 아닌 이상 이 명령어들을 직접 사용하는 것은 어렵습니다. 다행히도 C나 C++에서는 이런 저수준 명령어를 바로 사용할 수 있도록 intrinsic 함수들을 제공합니다. 이 함수들은 컴파일시에 CPU 명령어와 일대일로 대응되어 실제 CPU 명령어로 변환됩니다. SSE나 AVX intrinsic 함수를 사용하기 위해서는 해당 명령어를 지원하는 컴파일러 상에서 다음 헤더를 include하면 되겠습니다.

#include <xmmintrin.h> // SSE Intrinsic 사용 #include <immintrin.h> // AVX Intrinsic 사용

깁스 샘플링 코드를 보자

FLOAT* PLDAModel::getZLikelihoods(size_t threadId, const Document& doc, size_t vid) const { vector<float>& zLikelihood = ldaLocalData[threadId].zLikelihood; auto&& docState = static_pointer_cast_ptr<DocStateLDA>(doc.docState); float acc = 0; for (size_t k = 0; k < K; ++k) { zLikelihood[k] = acc += (docState->numByTopic[k] + alpha) * (ldaLocalData[threadId].numByWordTopic[V*k + vid] + beta) / (ldaLocalData[threadId].numByTopic[k] + V * beta); } return &zLikelihood[0]; }

대략적으로 이런 형태입니다. k = 0 ~ K-1의 경우에 대해 각각의 Likelihood를 계산하여 누적분포를 계산하는 코드이지요. K가 크다면, 같은 덧셈, 곱셈, 나눗셈을 반복하는 비중이 높아지므로, SIMD를 사용하기에 딱 제격입니다. 먼저 SSE를 사용해보도록 하죠. SSE를 위한 데이터 타입은 __m128입니다.

FLOAT* PLDAModel::getZLikelihoods(size_t threadId, const Document& doc, size_t vid) const { vector<float>& zLikelihood = ldaLocalData[threadId].zLikelihood; auto&& docState = static_pointer_cast_ptr<DocStateLDA>(doc.docState); float acc = 0; size_t k = 0; #ifdef USE_SSE2 __m128 mAlpha = _mm_set1_ps(alpha), mBeta = _mm_set1_ps(beta), mVBeta = _mm_set1_ps(V*beta); for (; k < (K >> 2) << 2; k += 4) { __m128i docNum = _mm_load_si128((__m128i*)&docState->numByTopic[k]), wordNum = _mm_set_epi32( ldaLocalData[threadId].numByWordTopic[V*(k + 3) + vid], ldaLocalData[threadId].numByWordTopic[V*(k + 2) + vid], ldaLocalData[threadId].numByWordTopic[V*(k + 1) + vid], ldaLocalData[threadId].numByWordTopic[V*k + vid]), allNum = _mm_load_si128((__m128i*)&ldaLocalData[threadId].numByTopic[k]); __m128 d = _mm_div_ps( _mm_mul_ps(_mm_add_ps(_mm_cvtepi32_ps(docNum), mAlpha), _mm_add_ps(_mm_cvtepi32_ps(wordNum), mBeta)), _mm_add_ps(_mm_cvtepi32_ps(allNum), mVBeta)); for (size_t t = 0; t < 4; ++t) zLikelihood[k + t] = acc += d.m128_f32[t]; } #endif for (; k < K; ++k) { zLikelihood[k] = acc += (docState->numByTopic[k] + alpha) * (ldaLocalData[threadId].numByWordTopic[V*k + vid] + beta) / (ldaLocalData[threadId].numByTopic[k] + V * beta); } return &zLikelihood[0]; }

_mm_set1_ps(a)는 4개의 float 값을 a로 세팅하는 함수입니다.

_mm_load_si128는 해당 포인터로부터 4개의 int 정수를 읽어서 레지스터에 넣는 함수, _mm_cvtepi32_ps는 4개의 int정수를 4개의 float 변수로 변환하는 함수입니다.

_mm_add_ps, _mm_mul_ps, _mm_div_ps는 설명안해도 아시겠죠?

SSE에서는 4개의 float을 하나로 묶어서 처리해야하기 때문에 반복문이 4개마다 묶여야 합니다. 그리고 4개로 묶이지 않는 나머지는 기존 방식처럼 하나씩 계산해주어야하구요.

이때 중요한 것은 load에 사용하는 포인터는 16바이트 단위로 정렬이 되어 있어야한다는 것입니다. 따라서 임의로 할당한 메모리는 작동을 보장할 수 없고, _mm_malloc, _mm_free로 정렬된 메모리 할당을 해야하죠. 이점만 지킨다면 연산에서 문제가 생길 부분은 없습니다.

이렇게 likelihood의 누적 분포를 구했으니, 이 중에서 랜덤하게 k를 뽑는 과정이 필요합니다. 이는 일반적으로 다음과 같이 구현됩니다.

template<class RealIt, class Random> inline size_t random_discrete_from_acc(RealIt begin, RealIt end, Random& rg) { auto r = std::generate_canonical<float, 32>(rg) * *(end - 1); size_t K = std::distance(begin, end); for (size_t z = 0; z < K; ++z) { if (r < begin[z]) return z; } return K - 1; }

[3, 7, 10, 12, 20]과 같이 누적 분포가 계산되어있을때 0~20 사이의 수를 랜덤하게 뽑아서 어느 구간이 선택되었는지를 찾는 코드에요. 만약 1이 뽑혔다면 0을, 11이 뽑혔다면 3을 반환해주겠죠.

반복문 안에 조건문이 들어있어서 병렬화하기가 어려워 보이지만, 조금만 머리를 굴리면 방법이 있습니다.

[3, 7, 10, 12, 20, 25, 26, 30, 32] 에서 22의 위치를 찾기

4개씩 묶여서 비교를 합니다.

[22, 22, 22, 22] < [3, 7, 10, 12] 는 [0, 0, 0, 0]. 전부 0이니깐 패스

[22, 22, 22, 22] < [20, 25, 26, 30] 는 [0, -1, -1, -1] (SSE에서 비교 연산이 참일 경우 -1의 값을 돌려줍니다.) 
0이 아닌 값이 있습니다. 0이 아닌 값의 개수를 세면 3개이므로 1번째 위치인 22 < 25가 발견되었다는 것을 알수 있습니다.

template<class RealIt, class Random> inline size_t random_discrete_from_acc(RealIt begin, RealIt end, Random& rg) { auto r = std::generate_canonical<float, 32>(rg) * *(end - 1); size_t K = std::distance(begin, end); size_t z = 0; #ifdef USE_SSE2 __m128 mr = _mm_set1_ps(r); for (; z < (K >> 2) << 2; z += 4) { __m128 mz = _mm_load_ps(&begin[z]); int mask = _mm_movemask_ps(_mm_cmplt_ps(mr, mz)); if (mask) return z + 4 - __popcnt(mask); } #endif for (; z < K; ++z) { if (r < begin[z]) return z; } return K - 1; }

_mm_movemask_ps는 4개의 int들의 최상위 비트만 모아서 4비트 마스크를 리턴해줍니다. popcnt는 1로 세트된 비트의 수를 돌려주고요. 따라서 둘을 활용하면 4개짜리 float 비교 결과를 한번에 알아낼 수 있습니다.

엄밀히 따져보면 이 경우는 굳이 병렬화하지 않아도 좋습니다. 사실 이 문제는 오름차순으로 정렬된 배열에서 특정한 원소이 위치를 찾는 것과 동치이므로 이진 탐색을 사용할 수 있기 때문이죠. 다만 실제로 구현해본 결과 K의 값이 특별히 크지 않은 이상 이진 탐색의 오버헤드가 생각보다 컸습니다. 아마 인접한 배열의 값을 참조할 경우 캐시에 적중할 가능성이 높기 때문에, 순차 탐색이 생각보다 느리지 않은 것으로 예상됩니다. 좀 더 정교한 구현이 되려면 K이 특정 threshold 이상일 경우 이진 탐색을 실시하다가, 탐색 구간이 작아지면 순차탐색을 진행하는 것이 좋을것 같습니다.

AVX를 사용한 경우도 유사하게 구현할 수 있습니다. 이에 대해서는 https://github.com/bab2min/tomoto 에서 확인하시면 되겠습니다.

성능 비교

그렇다면 이렇게 SSE 및 AVX로 최적화할 경우 성능 향상은 얼마나 될까요?

먼저 LDA(tomoto 구현)의 경우입니다. 5000개의 문헌, 14976종류의 어휘, 총 451119개의 단어로 이루어진 코퍼스를 대상, K=50으로 설정하여 테스트를 실시했습니다.

Tomoto LDA	32.900s	0%
Tomoto LDA With AVX	22.575s	+45.7%

AVX를 사용한 경우 약 45퍼센트의 성능 향상이 있었습니다. 

조금 더 연산 부하가 큰 DMR의 경우는 어떨까요? 같은 문헌 셋을 가지고 DMR 토픽 모델링을 실시해보았습니다.

Tomoto DMR	46.014s	0%
Tomoto DMR With SSE2	23.341s	+97.1%
Tomoto DMR With AVX	22.332s	+106.0%
Tomoto DMR With OpenCL	12.318s	+283.6%
Mallet DMR	33.123s	+38.9%

SSE2에 맞춰 최적화한 경우 약 97%, AVX에 맞춰 최적화한 경우 약 106%까지 성능 향상이 있었습니다. (물론 GPU를 사용한 경우를 이길수는 없었네요) 그리고 심심해서 Java로 작성된 토픽 모델링 라이브러리인 Mallet DMR로도 같은 테스트를 해보았는데 SSE2로 최적한 것보다 조금 느린 속도가 나왔습니다. 아마 Java에서 직접 SIMD를 사용할 방법은 없지만 컴파일시 Auto Vectorization이 어느 정도 진행되어 SIMD없는 C++보다 빠른 속도가 나온 것으로 보입니다.  Java도 생각보다 느리지 않네요