Оптимизация расчета скалярного произведения векторов int16 в Java с использованием Vector API ⇐ JAVA

[Anonymous]

TL;DR: оптимизация умножения 16-битных целочисленных массивов без переполнения с помощью Java Vector API.
Я пытаюсь оптимизировать критичный к производительности цикл, который применим функцию активации и вычисляет скалярное произведение двух массивов int16 с использованием векторного API Java (инкубирующего). Вот моя текущая скалярная реализация:

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

for (int i = 0; i < HIDDEN_SIZE; i++)
{
result += screlu(us.values[i]) * network.L1Weights[i]
+ screlu(them.values[i]) * network.L1Weights[i + HIDDEN_SIZE];
}

где

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

private static int screlu(short i)
{
int v = Math.max(0, Math.min(i, QA));
return v * v;
}

Я пытался векторизовать это следующим образом:

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

int[] usValues = new int[HIDDEN_SIZE];
int[] themValues = new int[HIDDEN_SIZE];

for (int i = 0; i < HIDDEN_SIZE; i++)
{
usValues[i] = (int) us.values[i];
themValues[i] = (int) them.values[i];
}

IntVector sum = IntVector.zero(INT_SPECIES);

for (; i < upperBound; i += INT_SPECIES.length())
{
IntVector va = IntVector.fromArray(INT_SPECIES, usValues, i);
IntVector vb = IntVector.fromArray(INT_SPECIES, themValues, i);
IntVector vc = IntVector.fromArray(INT_SPECIES, network.L1Weights, i);
IntVector vd = IntVector.fromArray(INT_SPECIES, network.L1Weights, i + HIDDEN_SIZE);

va = va.max(0).min(QA);
va = va.mul(va).mul(vc);

vb = vb.max(0).min(QA);
vb = vb.mul(vb).mul(vd);

sum = sum.add(va).add(vb);
}

int result = sum.reduceLanes(VectorOperators.ADD);

Из-за переполнения мне пришлось прибегнуть к использованию 32-битных полос, что снизило пропускную способность вдвое. В результате производительность лишь немного улучшилась. После некоторых исследований я обнаружил, что встроенная функция, такая как _mm256_madd_epi16, решает именно мою проблему, но я не смог найти никакой информации об этом в документации. Существует ли какая-либо эквивалентная операция в Vector API, а если нет, то есть ли другие решения этой проблемы?


Подробнее здесь: https://stackoverflow.com/questions/783 ... -using-vec