🥕

SIMD로 행렬연산 10배이상 빠르게 구현하기(SSE, AVX)

간단소개

팔만코딩경 컨트리뷰터

ContributorNotionAccount

주제 / 분류

Assembly

Graphics

Scrap

태그

9 more properties

개요

MiniRT에서 쓰기 위한 행렬연산을 SIMD로 구현했다. 전치행렬, 벡터x행렬, 행렬x행렬 을 구현했다.

인라인 어셈으로 일단 짜고 나중에 인트린직으로 포팅하려 한다.

SIMD 결정

SSE, AVX를 쓰겠다. arm neon은 ISA 특성상 성능향상이 드라마틱하지 않고, AVX512는 인라인 어셈을 지원하지 않을 뿐더러 클러스터 아이맥에서 지원하지 않아 보류한다.

최적화도 일단 짜놓고 나중에 포팅한 후 어셈블리 뽑히는거 보면서 조정할 예정 아마 자동으로 vectorization이 되겠지만 네이티브보다 2~10배 정도는 성능이 잘 나올 것이다.

Transpose

보통 행렬연산은 Transpose해서 쓰는게 편하다.

코드는 아래와 같이 짰다. AVX로 한 번에 4x4 행렬을 뒤집는다.

void transpose(mat4_t* mat)
{
    __asm {
        mov             eax, mat

        vmovups         ymm2, [eax];                    ymm2 = { x0, y0, z0, w0, x1, y1, z1, w1 }
        vmovups         ymm3, [eax + 32];               ymm3 = { x2, y2, z2, w2, x3, y3, z3, w3 }
        vperm2i128      ymm0, ymm2, ymm3, 00100000b;    ymm0 = { x0, y0, z0, w0, x2, y2, z2, w2 }
        vperm2i128      ymm1, ymm2, ymm3, 00110001b;    ymm1 = { x1, y1, z1, w1, x3, y3, z3, w3 }

        ; shuffle to each ymm contain the x,z or y,w elements of every row
        vshufps         ymm2, ymm0, ymm1, 10001000b;    ymm2 = { x0, z0, x1, z1, x2, z2, x3, z3 }
        vshufps         ymm3, ymm0, ymm1, 11011101b;    ymm3 = { y0, w0, y1, w1, y2, w2, y3, w3 }

        vinsertf128     ymm4, ymm2, xmm3, 1;            ymm4 = { x0, z0, x1, z1, y0, w0, y1, w1 }
        vperm2f128      ymm5, ymm2, ymm3, 00110001b;    ymm5 = { x2, z2, x3, z3, y2, w2, y3, w3 }

        vshufps         ymm6, ymm4, ymm5, 10001000b;    ymm6 = { x0, x1, x2, x3, y0, y1, y2, y3 }
        vshufps         ymm7, ymm4, ymm5, 11011101b;    ymm7 = { z0, z1, z2, z3, w0, w1, w2, w3 }

        vmovups         [eax], ymm6
        vmovups         [eax + 32], ymm7
    }
}
C
복사

주요 아이디어는 한 ymm이 모든 row의 특정 원소를 가지도록 하여 셔플을 한 번에 할 수 있도록 하는것이다. 다이어그램으로 과정을 나타내면 다음과 같다.

graph TD

subgraph Transpose Process
    Memory[mat: x0, y0, z0, w0, x1, y1, z1, w1, x2, y2, z2, w2, x3, y3, z3, w3]
    Memory -->|vmovups| B[ymm2: x0, y0, z0, w0, x1, y1, z1, w1]
    Memory -->|vmovups| C[ymm3: x2, y2, z2, w2, x3, y3, z3, w3]
    B --> |vperm2il128 | D[ymm0: x0, y0, z0, w0, x2, y2, z2, w2]
    C --> |vperm2il128| D[ymm0: x0, y0, z0, w0, x2, y2, z2, w2]
    B --> |vperm2il128| E[ymm1: x1, y1, z1, w1, x3, y3, z3, w3]
    C --> |vperm2il128| E[ymm1: x1, y1, z1, w1, x3, y3, z3, w3]
    D --> |vshufps| F[ymm2: x0, z0, x1, z1, x2, z2, x3, z3]
    E --> |vshufps| F[ymm2: x0, z0, x1, z1, x2, z2, x3, z3]
    D --> |vshufps| G[ymm3: y0, w0, y1, w1, y2, w2, y3, w3]
    E --> |vshufps| G[ymm3: y0, w0, y1, w1, y2, w2, y3, w3]
    F --> |vinsertf128| H[ymm4: x0, z0, x1, z1, y0, w0, y1, w1]
    G --> |vinsertf128| H[ymm4: x0, z0, x1, z1, y0, w0, y1, w1]
    F --> |vperm2f128| I[ymm5: x2, z2, x3, z3, y2, w2, y3, w3]
    G --> |vperm2f128| I[ymm5: x2, z2, x3, z3, y2, w2, y3, w3]
    H --> |vshufps| J[ymm6: x0, x1, x2, x3, y0, y1, y2, y3]
    I --> |vshufps| J[ymm6: x0, x1, x2, x3, y0, y1, y2, y3]
    H --> |vshufps| K[ymm7: z0, z1, z2, z3, w0, w1, w2, w3]
    I --> |vshufps| K[ymm7: z0, z1, z2, z3, w0, w1, w2, w3]
    J --> |vmovups| eax[mat: x0, x1, x2, x3, y0, y1, y2, y3]
    K --> |vmovups| eax+32[mat + 32: z0, z1, z2, z3, w0, w1, w2, w3];
end
Mermaid
복사

명령어는 다음을 참조하자. https://www.felixcloutier.com/x86/movups https://www.felixcloutier.com/x86/vperm2i128 https://www.felixcloutier.com/x86/shufps https://www.felixcloutier.com/x86/vinsertf128:vinsertf32x4:vinsertf64x2:vinsertf32x8:vinsertf64x4

Transform (Vector x Matrix)

벡터x행렬은 다음과 같다. dpps(dot product packed single precision)를 사용하여 간단히 해결했다. (다 짜보고나니 transpose 안하는게 쉬울지도)

void transform(vec4_t* dst, const vec4_t* src, const mat4_t* mat_tr)
{
    __asm {
        mov             edi, dst
        mov             esi, src
        mov             eax, mat_tr

        vbroadcastf128  ymm0, [esi];                    ymm0 = { x, y, z, w, x, y, z, w }
        vmovups         ymm1, [eax];                    ymm1 = { x0, y0, z0, w0, x1, y1, z1, w1 }
        vmovups         ymm2, [eax + 32];               ymm2 = { x2, y2, z2, w2, x3, y3, z3, w3 }

        vdpps           ymm1, ymm0, ymm1, 11110001b;    ymm1 = { s0, 0, 0, 0, s1, 0, 0, 0 }
        vdpps           ymm2, ymm0, ymm2, 11110001b;    ymm2 = { s2, 0, 0, 0, s3, 0, 0, 0 }

        vpermq          ymm1, ymm1, 11111000b;          xmm1 = { s0, 0, s1, 0 }
        vpermq          ymm2, ymm2, 11111000b;          xmm2 = { s2, 0, s3, 0 }
        shufps          xmm1, xmm2, 10001000b;          xmm1 = { s0, s1, s2, s3 }

        movups          [edi], xmm1
    }
} 
C
복사

그래프 그리기 어려워서 그래프는 생략한다.

SIMD 명령어는 reference를 무조건 잘 봐야한다. (https://www.felixcloutier.com/x86/dpps)

Concatenate (Matrix x Matrix)

행렬x행렬은 그냥 Vector x Matrix 를 4번 반복하면 된다.

void concatenate(mat4_t* dst, const mat4_t* m0, const mat4_t* m1_tr)
{
    __asm {
        mov             edi, dst
        mov             esi, m0
        mov             eax, m1_tr

        // load m1_tr
        vmovups         ymm2, [eax];                    ymm2 = { x0, y0, z0, w0, x1, y1, z1, w1 }
        vmovups         ymm3, [eax + 32];               ymm3 = { x2, y2, z2, w2, x3, y3, z3, w3 }

        // m0_r0
        vbroadcastf128  ymm0, [esi]
        vdpps           ymm4, ymm0, ymm2, 11110001b;    ymm4 = { x0, 0, 0, 0, y0, 0, 0, 0 }
        vdpps           ymm5, ymm0, ymm3, 11110001b;    ymm5 = { z0, 0, 0, 0, w0, 0, 0, 0 }
        vpermq          ymm4, ymm4, 11111000b;          xmm4 = { x0, 0, y0, 0 }
        vpermq          ymm5, ymm5, 11111000b;          xmm5 = { z0, 0, w0, 0 }
        shufps          xmm4, xmm5, 10001000b;          xmm4 = { x0, y0, z0, w0 }
        movups          [edi], xmm4

        // m0_r1
        vbroadcastf128  ymm0, [esi + 16]
        vdpps           ymm4, ymm0, ymm2, 11110001b;    ymm4 = { x1, 0, 0, 0, y1, 0, 0, 0 }
        vdpps           ymm5, ymm0, ymm3, 11110001b;    ymm5 = { z1, 0, 0, 0, w1, 0, 0, 0 }
        vpermq          ymm4, ymm4, 11111000b;          xmm4 = { x1, 0, y1, 0 }
        vpermq          ymm5, ymm5, 11111000b;          xmm5 = { z1, 0, w1, 0 }
        shufps          xmm4, xmm5, 10001000b;          xmm4 = { x1, y1, z1, w1 }
        movups          [edi + 16], xmm4

        // m0_r2
        vbroadcastf128  ymm0, [esi + 32]
        vdpps           ymm4, ymm0, ymm2, 11110001b;    ymm4 = { x2, 0, 0, 0, y2, 0, 0, 0 }
        vdpps           ymm5, ymm0, ymm3, 11110001b;    ymm5 = { z2, 0, 0, 0, w2, 0, 0, 0 }
        vpermq          ymm4, ymm4, 11111000b;          xmm4 = { x2, 0, y2, 0 }
        vpermq          ymm5, ymm5, 11111000b;          xmm5 = { z2, 0, w2, 0 }
        shufps          xmm4, xmm5, 10001000b;          xmm4 = { x2, y2, z2, w2 }
        movups          [edi + 32], xmm4

        // m0_r3
        vbroadcastf128  ymm0, [esi + 48]
        vdpps           ymm4, ymm0, ymm2, 11110001b;    ymm4 = { x3, 0, 0, 0, y3, 0, 0, 0 }
        vdpps           ymm5, ymm0, ymm3, 11110001b;    ymm5 = { z3, 0, 0, 0, w3, 0, 0, 0 }
        vpermq          ymm4, ymm4, 11111000b;          xmm4 = { x3, 0, y3, 0 }
        vpermq          ymm5, ymm5, 11111000b;          xmm5 = { z3, 0, w3, 0 }
        shufps          xmm4, xmm5, 10001000b;          xmm4 = { x3, y3, z3, w3 }
        movups          [edi + 48], xmm4
    }
}
C
복사

사용

이제 다른 코드에 붙여서 사용하면 된다. 끝!

[Directx 샘플 : HongLab Graphics]

사용 예시

// copyright. POCU Labs inc.

void run(vec4_t* world_pts, const vec4_t* local_pts, const size_t count, const vec3_t* scale, const vec3_t* rotation, const vec3_t* translation)
{
    // make transformation matrices
    mat4_t m_scale;
    mat_scale(&m_scale, scale->x, scale->y, scale->z);
    
    mat4_t m_rotation_x;
    mat4_t m_rotation_y;
    mat4_t m_rotation_z;
    mat_rotation_x(&m_rotation_x, rotation->x);
    mat_rotation_y(&m_rotation_y, rotation->y);
    mat_rotation_z(&m_rotation_z, rotation->z);

    mat4_t m_translation;
    mat_translation(&m_translation, translation->x, translation->y, translation->z);

    // concatenate tranformation matrices
    mat4_t m_concat;
    
    mat4_t m_scale_tr = m_scale;
    mat4_t m_rot_x_tr = m_rotation_x;
    mat4_t m_rot_y_tr = m_rotation_y;
    mat4_t m_rot_z_tr = m_rotation_z;
    mat4_t m_trans_tr = m_translation;
    transpose(&m_scale_tr);
    transpose(&m_rot_x_tr);
    transpose(&m_rot_y_tr);
    transpose(&m_rot_z_tr);
    transpose(&m_trans_tr);

    concatenate(&m_concat, &m_scale, &m_rot_x_tr);
    concatenate(&m_concat, &m_concat, &m_rot_y_tr);
    concatenate(&m_concat, &m_concat, &m_rot_z_tr);
    concatenate(&m_concat, &m_concat, &m_trans_tr);

    mat4_t m_concat_tr = m_concat;
    transpose(&m_concat_tr);

    // transform
    for (size_t i = 0; i < count; i++) {
        transform(&world_pts[i], &local_pts[i], &m_concat_tr);
    }
}
C++
복사

여담

어떤 데이터와 어떤 상황에서 사용하냐에 따라서 vmovups가 아닌 vmovntps를 사용할 수도 있다. (글쓴이의 non-temporal hint 관련 글 참고)