서열 정렬은 여러 생물정보학 애플리케이션과 파이프라인에서 중요한 알고리즘입니다. 정렬의 목표는 서열 간의 생물학적 관계에 대한 통찰을 얻는 것입니다. 특히 각 서열에서 일어난 점 돌연변이나 삽입 및 결실 같은 진화적 사건을 통해 서열들이 공통 조상으로부터 어떻게 분기했는지에 관심을 둡니다. 그러나 이 문제는 이차 복잡도를 가지므로, 이를 최적화하면 많은 애플리케이션에서 큰 이점을 얻을 수 있습니다. 많은 생물정보학 문제는 길이가 150~300 기호인 짧은 서열 조각 수백만 개를 정렬하는 것에서 시작하므로, SIMD 벡터를 사용하면 성능을 크게 향상할 수 있습니다. 이 실습에서는 데이터 레이아웃을 변환하고 SIMD 연산을 드러내어 알고리즘을 개선하는 방법을 배웁니다.

작성자: @rrahn.

힌트

힌트 1:

perf를 사용하거나 코드를 읽어 보면 compute_alignment의 루프가 병목임을 알 수 있습니다. 먼저 두 행렬 sequences1과 sequences2의 데이터 접근 패턴을 이해하고, 연산에서 어떤 패턴이 보이는지 살펴보십시오.


힌트 2:
  1. 각 서열 행렬의 차원은 sequence_count_v x sequence_size_v, 즉 16 x 200입니다.
  2. 16개의 각 행에 대해 다음을 수행합니다.
    • sequences1과 sequences2에서 i번째 행을 추출합니다.
    • sequences2[i] 행의 j번째 원소(또는 열)를 고정한 채로 sequences1[i]의 각 열을 순회하며 필요한 연산을 수행합니다.
    • sequences2[i]의 j개 열을 순회하면서 위 과정을 반복합니다.

여기서 주목해야 할 핵심은 완전히 동일한 연산을 여러 번 수행하고 있다는 점입니다. 이는 이러한 연산을 병렬로 수행할 수 있는 여지가 있음을 나타냅니다. 메모리 레이아웃(즉, 행렬 차원)을 컴파일러가 벡터화하기 쉽도록 하려면 행렬을 어떻게 재구성할 수 있겠습니까?


힌트 3:

행렬을 전치해야 합니다. 그러면 200 x 16이 되며, 이 16개의 열에는 점수 데이터(결과 타입은 16비트)가 담깁니다. 즉 행렬의 폭은 16 * 16비트 = 256비트가 되며, 이는 정확히 YMM/AVX 레지스터의 크기입니다. 이는 컴파일러가 자동 벡터화된 어셈블리 코드를 생성하도록 돕는 데 큰 진전입니다.

이제 행렬이 전치되었으므로, 루프에서 더 이상 스칼라 점수 누산기(score_column, horizontal_gap_column, last_vertical_gap 등)를 다루는 것이 아니라 폭이 sequence_count_v = 16인 벡터화된 값을 다루게 됩니다. 이 타입들을 크기 16의 std::array(또는 그에 상응하는 것)로 바꾸고, 각 원소를 갱신하도록 원소별 for 루프로 감싸야 합니다. 이렇게 작성하면 컴파일러가 벡터화된 코드를 더 쉽게 생성할 수 있습니다.

또한 변경된 과정을 반영하도록 기존 루프도 반드시 수정하십시오. 일부 루프는 더 이상 필요하지 않을 수 있습니다.


풀이:

배경:

이 실습에는 길이가 같은 시퀀스들의 두 집합을 나타내는 두 개의 행렬이 있습니다. 이 알고리즘은 DNA 및 RNA 분석과 관련된 서열 정렬(sequence alignment)이라는 분석의 일부로 일반적으로 수행되는 계산을 나타냅니다.

코드 프로파일링부터 시작하면 Linux perf에서 다음과 같은 성능 출력을 얻습니다:

----------------------------------------------------------------------------------
Benchmark                                        Time             CPU   Iterations
----------------------------------------------------------------------------------
bench_compute_alignment/iterations:2000       1252 us         1252 us         2000

코드를 분석하면 다음과 같은 사실을 알 수 있습니다:

  1. 각 행렬의 차원은 sequence_count_v x sequence_size_v, 즉 16 x 200입니다.
  2. 16개의 각 행에 대해 다음을 수행합니다:
    • sequences1과 sequences2에서 i번째 행을 추출합니다.
    • sequences2[i] 행의 j번째 요소(또는 열)에 고정한 상태에서 sequences1[i]의 각 열을 순회하며 필요한 연산을 수행합니다.
    • sequences2[i]의 j개 열을 순회하면서 위 과정을 반복합니다.

여기서 주목해야 할 핵심은 완전히 동일한 연산을 여러 번 수행하고 있다는 점입니다. 이는 이러한 연산을 병렬로 수행할 여지가 있음을 의미합니다. 이를 달성하려면 이것이 가능하도록 문제를 리팩터링해야 합니다.

해결책

행렬의 차원이 16 x 200이었음을 떠올려 봅시다. 이 행렬을 전치하여 대신 200개의 열과 16개의 행을 갖도록 하면 어떨까요? 그렇게 하면 200번의 행 반복 각각에서 단 16개의 요소에 대한 연산만 처리하게 됩니다. 행렬에는 std::uint8_t가 저장되어 있으므로, 한 행에 16개가 있으면 128비트가 됩니다. 이는 벡터화된 데이터에 대한 SIMD 연산에 사용되는 XMM 레지스터의 크기와 같습니다. 점점 가능성이 보이기 시작합니다!

먼저 행렬을 전치할 수 있도록 유틸리티 메서드를 제공합시다:

using simd_score_t = std::array<int16_t, sequence_count_v>;
using simd_sequence_t = std::array<simd_score_t, sequence_size_v>;

simd_sequence_t transpose(const std::vector<sequence_t>& vec) {
    simd_sequence_t transposed{};

    for (size_t i = 0; i < transposed.size(); ++i) {
        for (size_t j = 0; j < vec.size(); ++j) {
            transposed[i][j] = vec[j][i];
        }
    }
    return transposed;
}

여기서는 전치된 행렬을 생성하며, 그 요소의 크기는 원래 루프의 누산기인 score_t와 같은 16비트 정수입니다. 이는 전치된 행렬의 행이 실제로는 128비트가 아니라 256비트가 된다는 의미입니다. 하지만 대신 YMM 레지스터를 사용하여 여전히 벡터화를 수행할 수 있으므로 문제없습니다.

행렬을 전치했으니 이제 모든 누산기를 새로운 너비에 맞게 넓혀야 합니다. 각 반복에서 단순히 std::uint16_t 크기의 스칼라 score_t를 처리하는 대신, 16개의 점수 누산기(전치된 입력 데이터의 각 열마다 하나씩)를 담은 벡터 점수 누산기를 처리하려는 것입니다. 기억하십시오. 지금 우리는 컴파일러가 벡터화된 코드를 생성할 수 있는 방식으로 알고리즘을 작성하려는 것이므로, 이 접근 방식이 바로 우리가 원하는 것입니다.

먼저 score_t를 16비트 정수 16개로 이루어진 배열로 바꿉시다:

/// From our transpose change above...
using simd_score_t = std::array<int16_t, sequence_count_v>;
...
using score_t = simd_score_t;

초기화 값도 std::array::fill을 사용하도록 변경해야 합니다:

score_t gap_open{};
gap_open.fill(-11);
score_t gap_extension{};
gap_extension.fill(-1);
score_t match{};
match.fill(6);
score_t mismatch{};
mismatch.fill(-4);

다른 변수들도 배열 너비가 되어야 합니다. 행렬의 첫 번째 열을 초기화하는 루프는 다음과 같이 수정합니다:

for (size_t i = 1; i < score_column.size(); ++i) {
    for (size_t k = 0; k < sequence_count_v; ++k) {
        score_column[i][k] = last_vertical_gap[k];
        horizontal_gap_column[i][k] = last_vertical_gap[k] + gap_open[k];
        last_vertical_gap[k] += gap_extension[k];
    }
}

주 재귀 루프의 시작 부분에서도 같은 작업을 수행합니다:

/*
* Compute the main recursion to fill the matrix.
*/
for (unsigned row2 = 1; row2 <= sequence2.size(); ++row2) {
  score_t last_diagonal_score =
      score_column[0]; // Cache last diagonal score to compute this cell.
  for (size_t k = 0; k < sequence_count_v; ++k) {
      score_column[0][k] = horizontal_gap_column[0][k];
      last_vertical_gap[k] = horizontal_gap_column[0][k] + gap_open[k];
      horizontal_gap_column[0][k] += gap_extension[k];
  }
...
}

최종 루프는 다음과 같아야 합니다:

for (unsigned row = 1; row <= sequence1.size(); ++row) {
   // Compute next score from diagonal direction with match/mismatch.
   score_t best_cell_score = last_diagonal_score;
   for (size_t k = 0; k < sequence_count_v; ++k) {
       best_cell_score[k] += sequence1[row - 1][k] == sequence2[row2 - 1][k] ? match[k] : mismatch[k];
   }
   for (size_t k = 0; k < sequence_count_v; ++k) {
       // Determine best score from diagonal, vertical, or horizontal
       // direction.
       best_cell_score[k] = std::max(best_cell_score[k], last_vertical_gap[k]);
       best_cell_score[k] = std::max(best_cell_score[k], horizontal_gap_column[row][k]);
       // Cache next diagonal value and store optimum in score_column.
       last_diagonal_score[k] = score_column[row][k];
       score_column[row][k] = best_cell_score[k];
       // Compute the next values for vertical and horizontal gap.
       best_cell_score[k] += gap_open[k];
       last_vertical_gap[k] += gap_extension[k];
       horizontal_gap_column[row][k] += gap_extension[k];
       // Store optimum between gap open and gap extension.
       last_vertical_gap[k] = std::max(last_vertical_gap[k], best_cell_score[k]);
       horizontal_gap_column[row][k] =
           std::max(horizontal_gap_column[row][k], best_cell_score[k]);
   }
}

best_cell_score에 대한 첫 번째 내부 루프는 인덱스가 같은데도 마지막 루프와 별도로 계산된다는 점에 주목하십시오. 왜 그럴까요? 이는 컴파일러가 벡터화된 코드를 생성하여 두 번째 루프에서 best_cell_score의 16개 열을 한 번에 모두 갱신할 수 있도록 하기 위한 필수 조건입니다. 두 루프를 합치면 반복마다 열을 하나씩만 채울 수 있으므로, 결국 원래의 스칼라 코드보다 나을 것이 없습니다.

마지막으로, 다음과 같이 결과를 준비합니다.

// Report the best score.
for (size_t k = 0; k < sequence_count_v; ++k) {
  result[k] = score_column.back()[k];
}

이 모든 변경 사항을 적용하면 다음과 같은 성능 보고서를 얻습니다.

----------------------------------------------------------------------------------
Benchmark                                        Time             CPU   Iterations
----------------------------------------------------------------------------------
bench_compute_alignment/iterations:2000        497 us          497 us         2000

이 기법을 통해 약 60%의 속도 향상을 달성했습니다.

perf report를 사용하여 컴파일러가 생성한 어셈블리를 살펴보면, 이론에서 예상한 대로 YMM 레지스터(256비트)를 사용하는 벡터화된 명령어가 많이 보입니다.

     │     for (unsigned row = 1; row <= sequence1.size(); ++row) {                                                                                                                                                    ▒
     │       nop                                                                                                                                                                                                       ▒
     │     // Compute next score from diagonal direction with match/mismatch.                                                                                                                                          ▒
     │     score_t best_cell_score = last_diagonal_score;                                                                                                                                                              ▒
 559 │800:   vmovdqa      ymm5,YMMWORD PTR [rsp+0xc0]                                                                                                                                                                  ▒
     │     for (size_t k = 0; k < sequence_count_v; ++k) {                                                                                                                                                             ▒
     │       xor          eax,eax                                                                                                                                                                                      ▒
     │     score_t best_cell_score = last_diagonal_score;                                                                                                                                                              ▒
 282 │       vmovdqa      YMMWORD PTR [rsp+0xe0],ymm5       

해결책 코드(링크):

솔루션 링크