vector matrix multiplication in nnue

[octopus-prime]

I still try to understand affine transformation in nnue.
I focus on vector matrix multiplication.
No bias, no sparse - to keep it most simple.

    constexpr auto N = 64;
    constexpr auto M = 16;
    alignas(64) int8_t input[N] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64};
    alignas(64) int8_t weight[M][N] = {
        1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
    };
    alignas(64) int32_t output[M];

// vector matrix multiplication here

    std::println("output = {},{},{},{},{},{},{},{},{},{},{},{},{},{},{},{}", output[0], output[1], output[2], output[3], output[4], output[5], output[6], output[7], output[8], output[9], output[10], output[11], output[12], output[13], output[14], output[15]);

Is this design correct (weight[M][N] vs weight[N][M]) ??
What is the expected outcome? Is it output = 1,2,3,4,0,0,0,0,0,0,0,0,0,0,0,0 ?

[MinetaS]

Check no SIMD implementation:

    std::memcpy(output, biases, sizeof(OutputType) * OutputDimensions);

    // Traverse weights in transpose order to take advantage of input sparsity
    for (IndexType i = 0; i < InputDimensions; ++i)
    {
        const InputType in = input[i];
        if (in)
        {
            const WeightType* w = &weights[i];
            for (IndexType j = 0; j < OutputDimensions; ++j)
                output[j] += w[j * PaddedInputDimensions] * in;
        }
    }

It doesn't matter whether the array is defined as N*M or M*N, it depends on how you store numbers (transpose or not).
Weights should be M*N matrix so the result vector is M*1 after multiplication with input vector which is N*1.

If the top-leftmost 4x4 of weight matrix is identical and bias is zero, the result is {1, 2, 3, 4, 0, 0, ...}.

[octopus-prime]

Thanks a lot. I read your answer as:
Is this design correct (weight[M][N] vs weight[N][M]) ?? -> yes
What is the expected outcome? Is it output = 1,2,3,4,0,0,0,0,0,0,0,0,0,0,0,0 ? -> yes

Ok, since design is correct it's time to implement the mutliplication.

Notice this implementation gives the expected output output = 1,2,3,4,0,0,0,0,0,0,0,0,0,0,0,0

    for (auto&& [output_element, weight_column] : std::views::zip(output, weight)) {
        __m256i sum = _mm256_setzero_si256();
        for (auto&& [input_elements, weight_elements] : std::views::zip(to_simd(input), to_simd(weight_column)))
            sum = m256_add_dpbusd_epi32(sum, input_elements, weight_elements);
        output_element = m256_hadd(sum);
    }

And this impl uses m256_add_dpbusd_epi32 like your impl.
But this impl uses m256_hadd too - your impl does not (ignoring special case M = 1).
But this impl iterates over input in the inner loop - your impl iterates over input in the outer loop.
So i did not get your impl.

And if i try to get your impl i end up with

    auto input_span = std::span<int32_t>{( int32_t*) input, N / 4};
    __m256i sums[M / 8] = {_mm256_setzero_si256(), _mm256_setzero_si256()};
    for (auto&& [input_element, weight_column] : std::views::zip(input_span, weight))
        for (auto&& [sum, weight_elements] : std::views::zip(sums, to_simd(weight_column)))
            sum = m256_add_dpbusd_epi32(sum, _mm256_set1_epi32(input_element), weight_elements);
    std::ranges::copy(sums, (__m256i*) output);

But this produces nonsence output = 34,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0

So i still missed something very important.

[MinetaS]

m256_hadd is required if the weight matrix is not transposed respectively, because _mm256_dpbusd_epi32 calculates sum of dot products of four 8-bit elements per each 32bit packed integer.
Check #3663 for more information.

Stockfish/src/nnue/layers/affine_transform.h

Lines 153 to 156 in 9fb5832

	static constexpr IndexType get_weight_index_scrambled(IndexType i) {
	return (i / 4) % (PaddedInputDimensions / 4) * OutputDimensions * 4
	+ i / PaddedInputDimensions * 4 + i % 4;
	}

[octopus-prime]

Thanks a lot.
This is still like voodoo for me.

constexpr auto get_weight_index_scrambled = [](auto i) { 
    return (i / 4) % (N / 4) * M * 4 + i / N * 4 + i % 4; 
};

Is there an explaination for this?

Looked at the scambled matirx too - looks just wiered :-)

But it did help!

[octopus-prime]

So i have a working toy-impl now.

I noticed that the size of the output is now limited by simd size.
For avx2 -> 32
For avx512 -> 64

Correct?

void affine_transform(auto&& input, auto&& weight, auto&& bias, auto&& output) noexcept {
    using simd_t = __m256i;
    constexpr auto set1 = [](auto&& chunk) -> simd_t { return _mm256_set1_epi32(*((int32_t*) chunk.data())); };
    constexpr auto load = [](auto&& chunk) -> simd_t { return _mm256_load_si256((simd_t*) chunk.data()); };
    constexpr auto dpbusd = [](simd_t acc, simd_t a, simd_t b) -> simd_t { return m256_add_dpbusd_epi32(acc, a, b); };

    constexpr auto ChunkSize = sizeof(int32_t) / sizeof(int8_t);
    constexpr auto Simd32Size = sizeof(simd_t) / sizeof(int32_t);
    constexpr auto NumRegs = std::size(output) / Simd32Size;
    constexpr auto NumCols = sizeof(simd_t) / std::size(output);

    auto input_range = input | std::views::chunk(ChunkSize) | std::views::transform(set1);
    auto weight_range = weight | std::views::chunk(ChunkSize);
    auto bias_range = bias | std::views::chunk(Simd32Size) | std::views::transform(load);
    auto output_range = (simd_t*) output;

    simd_t sums[NumRegs];
    std::ranges::copy(bias_range, sums);
    for (auto&& [input_element, weight_column] : std::views::zip(input_range, weight_range)) {
        auto column_range = weight_column | std::views::chunk(NumCols) | std::views::transform(load);
        for (auto&& [sum, column_elements] : std::views::zip(sums, column_range))
            sum = dpbusd(sum, input_element, column_elements);
    }
    std::ranges::copy(sums, output_range);
}

This impl is dense.
For N = 3072 and M = 32 it takes about 4µs.

So i want a sparse impl too.
So here my naive approch (using std::simd for the all_of() function)

    simd_t sums[NumRegs];
    std::ranges::copy(bias_range, sums);
    for (auto&& [input_element, weight_column] : std::views::zip(input_range, weight_range)) {
        if (std::experimental::all_of((std::experimental::native_simd<int32_t>) input_element == 0))
            continue;
        auto column_range = weight_column | std::views::chunk(NumCols) | std::views::transform(load);
        for (auto&& [sum, column_elements] : std::views::zip(sums, column_range))
            sum = dpbusd(sum, input_element, column_elements);
    }
    std::ranges::copy(sums, output_range);

As you can see there is no
std::uint16_t nnz[NumChunks];
and no
find_nnz<NumChunks>(input32, nnz, count);

Nevertheless the time dropped from 4µs to 1µs (very sparse input data).
Is this the expected perf gain? Factor ~ 4 ?!

If my impl does well, why you have this 'find_nnz' stuff and more important the nnz array and the 2 loops (1x input + 1x nnz)?

Did i still miss something?

[octopus-prime]

Yes I did. I should just copy/paste.