Ориентируясь на AVX2, каков самый быстрый способ транспонировать матрицу 8x8, содержащую 64-битные целые числа (или числа с двойной точностью)?

Я искал этот сайт и нашел несколько способов транспонирования 8x8, но в основном для 32-битных чисел с плавающей запятой. Поэтому я в основном спрашиваю, потому что не уверен, легко ли переводятся принципы, которые сделали эти алгоритмы быстрыми, в 64-битные и, во-вторых, очевидно, что AVX2 имеет только 16 регистров, поэтому только загрузка всех значений займет все регистры.

Один из способов сделать это - вызвать 2x2 _MM_TRANSPOSE4_PD, но мне было интересно, оптимально ли это:

  #define _MM_TRANSPOSE4_PD(row0,row1,row2,row3)                \
        {                                                       \
            __m256d tmp3, tmp2, tmp1, tmp0;                     \
                                                                \
            tmp0 = _mm256_shuffle_pd((row0),(row1), 0x0);       \
            tmp2 = _mm256_shuffle_pd((row0),(row1), 0xF);       \
            tmp1 = _mm256_shuffle_pd((row2),(row3), 0x0);       \
            tmp3 = _mm256_shuffle_pd((row2),(row3), 0xF);       \
                                                                \
            (row0) = _mm256_permute2f128_pd(tmp0, tmp1, 0x20);  \
            (row1) = _mm256_permute2f128_pd(tmp2, tmp3, 0x20);  \
            (row2) = _mm256_permute2f128_pd(tmp0, tmp1, 0x31);  \
            (row3) = _mm256_permute2f128_pd(tmp2, tmp3, 0x31);  \
        }

Все еще предполагая AVX2, транспонирование double[8][8] и int64_t[8][8] в принципе одинаково?

PS: И просто из любопытства, наличие AVX512 существенно изменит ситуацию, верно?

1
Ecir Hana 23 Мар 2021 в 18:47

2 ответа

Лучший ответ

После некоторых размышлений и обсуждения в комментариях я думаю, что это наиболее эффективная версия, по крайней мере, когда исходные и целевые данные находятся в ОЗУ. Не требует AVX2, достаточно AVX1.

Основная идея заключается в том, что современные процессоры могут выполнять в два раза больше микроопераций загрузки по сравнению с хранилищами, а на многих процессорах загрузка материала в более высокую половину векторов с помощью vinsertf128 имеет ту же стоимость, что и обычная 16-байтовая загрузка. По сравнению с вашим макросом этой версии больше не нужны эти относительно дорогие (3 цикла задержки на большинстве процессоров) vperm2f128 перемешивания.

struct Matrix4x4
{
    __m256d r0, r1, r2, r3;
};

inline void loadTransposed( Matrix4x4& mat, const double* rsi, size_t stride = 8 )
{
    // Load top half of the matrix into low half of 4 registers
    __m256d t0 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi ) );     // 00, 01
    __m256d t1 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi + 2 ) ); // 02, 03
    rsi += stride;
    __m256d t2 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi ) );     // 10, 11
    __m256d t3 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi + 2 ) ); // 12, 13
    rsi += stride;
    // Load bottom half of the matrix into high half of these registers
    t0 = _mm256_insertf128_pd( t0, _mm_loadu_pd( rsi ), 1 );    // 00, 01, 20, 21
    t1 = _mm256_insertf128_pd( t1, _mm_loadu_pd( rsi + 2 ), 1 );// 02, 03, 22, 23
    rsi += stride;
    t2 = _mm256_insertf128_pd( t2, _mm_loadu_pd( rsi ), 1 );    // 10, 11, 30, 31
    t3 = _mm256_insertf128_pd( t3, _mm_loadu_pd( rsi + 2 ), 1 );// 12, 13, 32, 33

    // Transpose 2x2 blocks in registers.
    // Due to the tricky way we loaded stuff, that's enough to transpose the complete 4x4 matrix.
    mat.r0 = _mm256_unpacklo_pd( t0, t2 ); // 00, 10, 20, 30
    mat.r1 = _mm256_unpackhi_pd( t0, t2 ); // 01, 11, 21, 31
    mat.r2 = _mm256_unpacklo_pd( t1, t3 ); // 02, 12, 22, 32
    mat.r3 = _mm256_unpackhi_pd( t1, t3 ); // 03, 13, 23, 33
}

inline void store( const Matrix4x4& mat, double* rdi, size_t stride = 8 )
{
    _mm256_storeu_pd( rdi, mat.r0 );
    _mm256_storeu_pd( rdi + stride, mat.r1 );
    _mm256_storeu_pd( rdi + stride * 2, mat.r2 );
    _mm256_storeu_pd( rdi + stride * 3, mat.r3 );
}

// Transpose 8x8 matrix of double values
void transpose8x8( double* rdi, const double* rsi )
{
    Matrix4x4 block;
    // Top-left corner
    loadTransposed( block, rsi );
    store( block, rdi );

#if 1
    // Using another instance of the block to support in-place transpose
    Matrix4x4 block2;
    loadTransposed( block, rsi + 4 );       // top right block
    loadTransposed( block2, rsi + 8 * 4 ); // bottom left block

    store( block2, rdi + 4 );
    store( block, rdi + 8 * 4 );
#else
    // Flip the #if if you can guarantee ( rsi != rdi )
    // Performance is about the same, but this version uses 4 less vector registers,
    // slightly more efficient when some registers need to be backed up / restored.
    assert( rsi != rdi );
    loadTransposed( block, rsi + 4 );
    store( block, rdi + 8 * 4 );

    loadTransposed( block, rsi + 8 * 4 );
    store( block, rdi + 4 );
#endif
    // Bottom-right corner
    loadTransposed( block, rsi + 8 * 4 + 4 );
    store( block, rdi + 8 * 4 + 4 );
}

Для полноты, вот версия, которая использует код, очень похожий на ваш макрос, выполняет вдвое меньше загрузок, такое же количество хранений и больше перемешиваний. Не тестировал, но я ожидал, что он будет немного медленнее.

struct Matrix4x4
{
    __m256d r0, r1, r2, r3;
};

inline void load( Matrix4x4& mat, const double* rsi, size_t stride = 8 )
{
    mat.r0 = _mm256_loadu_pd( rsi );
    mat.r1 = _mm256_loadu_pd( rsi + stride );
    mat.r2 = _mm256_loadu_pd( rsi + stride * 2 );
    mat.r3 = _mm256_loadu_pd( rsi + stride * 3 );
}

inline void store( const Matrix4x4& mat, double* rdi, size_t stride = 8 )
{
    _mm256_storeu_pd( rdi, mat.r0 );
    _mm256_storeu_pd( rdi + stride, mat.r1 );
    _mm256_storeu_pd( rdi + stride * 2, mat.r2 );
    _mm256_storeu_pd( rdi + stride * 3, mat.r3 );
}

inline void transpose( Matrix4x4& m4 )
{
    // These unpack instructions transpose lanes within 2x2 blocks of the matrix
    const __m256d t0 = _mm256_unpacklo_pd( m4.r0, m4.r1 );
    const __m256d t1 = _mm256_unpacklo_pd( m4.r2, m4.r3 );
    const __m256d t2 = _mm256_unpackhi_pd( m4.r0, m4.r1 );
    const __m256d t3 = _mm256_unpackhi_pd( m4.r2, m4.r3 );
    // Produce the transposed matrix by combining these blocks
    m4.r0 = _mm256_permute2f128_pd( t0, t1, 0x20 );
    m4.r1 = _mm256_permute2f128_pd( t2, t3, 0x20 );
    m4.r2 = _mm256_permute2f128_pd( t0, t1, 0x31 );
    m4.r3 = _mm256_permute2f128_pd( t2, t3, 0x31 );
}

// Transpose 8x8 matrix with double values
void transpose8x8( double* rdi, const double* rsi )
{
    Matrix4x4 block;
    // Top-left corner
    load( block, rsi );
    transpose( block );
    store( block, rdi );

    // Using another instance of the block to support in-place transpose, with very small overhead
    Matrix4x4 block2;
    load( block, rsi + 4 );     // top right block
    load( block2, rsi + 8 * 4 ); // bottom left block

    transpose( block2 );
    store( block2, rdi + 4 );
    transpose( block );
    store( block, rdi + 8 * 4 );

    // Bottom-right corner
    load( block, rsi + 8 * 4 + 4 );
    transpose( block );
    store( block, rdi + 8 * 4 + 4 );
}
2
Soonts 24 Мар 2021 в 05:19

Для небольших матриц, где более 1 строки может поместиться в один вектор SIMD, AVX-512 имеет очень удобную тасовку с пересечением полос с 2 входами с 32-битной или 64-битной степенью детализации с векторным управлением. (В отличие от _mm512_unpacklo_pd, который представляет собой 4 отдельных 128-битных перемешивания.)

Матрица 4x4 double - это "всего" 128 байтов, два вектора ZMM __m512d, поэтому вам нужно только два vpermt2ps (_mm512_permutex2var_pd) для создания обоих выходных векторов: по одному на каждый выходной вектор в случайном порядке, с обоими грузы и магазины во всю ширину. Однако вам нужны константы управляющих векторов.

Использование 512-битных векторных инструкций имеет некоторые недостатки (тактовая частота и пропускная способность порта выполнения), но если ваша программа может проводить много времени в коде, который использует 512-битные векторы, вероятно, есть значительный выигрыш в пропускной способности от передачи большего количества данных с каждым инструкция, и более мощные перемешивания.

С 256-битными векторами vpermt2pd ymm, вероятно, не будет полезен для 4x4, потому что для каждой выходной строки __m256d каждый из 4 элементов, которые вы хотите, поступает из другой входной строки. Таким образом, одно перемешивание с двумя входами не может дать желаемого результата.

Я думаю, что перемешивание с пересечением полос с детализацией менее 128 бит бесполезно, если ваша матрица не достаточно мала для размещения нескольких строк в одном векторе SIMD. См. Как транспонировать матрицу 16x16 с помощью инструкций SIMD? для рассуждений об алгоритмической сложности 32-битных элементов - xpose 8x8 32-битных элементов с AVX1 примерно такой же, как 8x8 64-битных элементов с AVX-512, где каждый вектор SIMD содержит ровно одну целую строку.

Таким образом, нет необходимости в векторных константах, просто мгновенная перетасовка 128-битных фрагментов и unpacklo/hi


Транспонирование 8x8 с 512-битными векторами (8 двойников) будет иметь ту же проблему: каждая выходная строка из 8 двойников требует по 1 двойнику из каждого из 8 входных векторов. Итак, в конечном счете, я думаю, вам нужна стратегия, аналогичная ответу Soonts AVX, начиная с _mm512_insertf64x4(v, load, 1) в качестве первого шага для преобразования первой половины двух входных строк в один вектор.

(Если вас интересует KNL / Xeon Phi, другой ответ @ ZBoson на Как транспонировать матрицу 16x16 с помощью инструкций SIMD? показаны некоторые интересные идеи использования маскирования слияния с перетасовкой с одним входом, например vpermpd или vpermq вместо двух входов. перемешивается как vunpcklpd или vpermt2pd)

Использование более широких векторов означает меньшее количество загрузок и сохранений и, возможно, даже меньшее количество перетасовок, потому что каждый из них объединяет больше данных. Но у вас также есть дополнительная работа по перетасовке, чтобы собрать все 8 элементов строки в один вектор, вместо того, чтобы просто загружать и сохранять в разных местах кусками половинного размера строки. Неочевидно, что лучше; Я обновлю этот ответ, если дойду до написания кода.

Обратите внимание, что Ice Lake (первый потребительский процессор с AVX-512) может выполнять 2 загрузки и 2 сохранения за такт. У него лучшая пропускная способность перемешивания, чем у Skylake-X для некоторых перемешивания, но не для тех, которые полезны для этого или ответа Сунта. (Все vperm2f128, vunpcklpd и vpermt2pd работают только на порту 5 для версий ymm и zmm. https://uops.info/ . vinsertf64x4 zmm, mem, 1 - это 2 мопа для внешнего интерфейса, и ему нужен порт загрузки и моп для p0 / p5. (Не p1, потому что это 512-битный uop, а также см. инструкции SIMD, снижающие частоту процессора).)

1
Peter Cordes 24 Мар 2021 в 17:23