Как эффективно загружать и хранить данные для новых инструкций AVX-VNNI и Arm Neon MMLA? ⇐ C++

[Anonymous]

Каков правильный способ загрузки данных для последних инструкций AVX-VNNI и Arm Neon MMLA?
Например, описание SMMLA:

8-битная целочисленная матрица со знаком, умножение-накопление. Эта инструкция умножает матрицу 2x8 8-битных целых чисел со знаком в первом исходном векторе на матрицу 8x2 8-битных целых чисел со знаком во втором исходном векторе. Полученная 32-битная целочисленная матрица размером 2x2 [...]

Аналогично, описание для _mm256_dpbusd_epi32:

Умножить группы из 4 соседних пар беззнаковых 8-битных целых чисел в a на соответствующие 8-битные целые числа со знаком в b, получив 4 промежуточных знаковых числа. 16-битные результаты. Суммируйте эти 4 результата с соответствующим 32-битным целым числом в src и сохраните упакованные 32-битные результаты в dst.

Похоже, что все они требуют входы вида 2[4]x8 и 8x[4]2. и выдать выходные данные в форме 2[4]x[4]2. Как я могу эффективно загружать и хранить данные для этих функций?
Я вижу три широкие возможности использования этих инструкций, ни одна из которых не является привлекательной:

• [Разделение и объединение] Я загружаю два последовательных 128-битных вектора, а затем разделяю их. Аналогично, для AVX я бы загрузил 4 128 или 256 векторов, а затем разделил их. Хранение также «сложно», поскольку мне нужно извлечь соответствующие части матрицы 2[4]x[4]2 перед ее сохранением. Мой код перегружен инструкциями по разделению/объединению.
• [Меньшие векторы] Альтернативно я мог бы загружать меньшие части, но это тоже кажется неэффективным.
• [Изменение порядка входных данных] Конечно, я мог бы изменить порядок входных данных так, чтобы векторизованные нагрузки уже охватывали несколько строк или столбцов. Должно ли это быть предполагаемое использование?

Пример кода для внутреннего цикла (сокращение по K) небольшой входной матрицы A размером 4xK (основной по строкам) и матрица B Kx4 (основной столбец) выглядит следующим образом:

Код: Выделить всё

for (size_t k = 0; k < 64; k += 8) {
uint8x8_t low = vld1_u8(row0);
uint8x8_t high = vld1_u8(row1);
uint8x16_t row01x01234567 = vcombine_u8(low, high);
row0 += 8;
row1 += 8;
low = vld1_u8(row2);
high = vld1_u8(row3);
uint8x16_t row23x01234567 = vcombine_u8(low, high);
row2 += 8;
row3 += 8;
low = vld1_u8(col0);
high = vld1_u8(col1);
uint8x16_t col01x01234567 = vcombine_u8(low, high);
col0 += 8;
col1 += 8;
low = vld1_u8(col2);
high = vld1_u8(col3);
uint8x16_t col23x01234567 = vcombine_u8(low, high);
col2 += 8;
col3 += 8;
out01x01 = vmmlaq_u32(out01x01, row01x01234567, col01x01234567);
out01x23 = vmmlaq_u32(out01x23, row01x01234567, col23x01234567);

out23x01 = vmmlaq_u32(out23x01, row23x01234567, col01x01234567);
out23x23 = vmmlaq_u32(out23x23, row23x01234567, col23x01234567);
}

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


Подробнее здесь: https://stackoverflow.com/questions/792 ... ructions-e