perf(avx512f): vectorize 16x16 float / 8x8 double transpose

[DiamonDinoia]

Without a 32/64-bit avx512f transpose specialization, batch/ (and int32/int64) transpose fell through to the generic common path.

Add native AVX-512 transpose networks:

• float 16x16: unpcklo/hi_ps (32-bit) -> shuffle_ps (64-bit) -> two shuffle_f32x4 (128-bit lane) stages = 64 shuffle-domain ops.
• double 8x8: unpcklo/hi_pd (64-bit) -> two shuffle_f64x2 (128-bit lane) stages = 24 shuffle-domain ops. Both use only AVX512F intrinsics. int32/uint32 alias to float, int64/uint64 to double.

[serge-sans-paille]

Can you share some benchmark numbers? I've looked at the generated assembly and it's super dense ^^!

[DiamonDinoia]

Yes, the asm is a lot of instructions. Give me a moment to gather the numbers. I have to run this in a server as my laptop does not have avx512.

[serge-sans-paille]

Thanks! Mine does not have those either :-/

[DiamonDinoia]

#include <chrono>
#include <cstdint>
#include <cstdio>
#include <vector>
#include <xsimd/xsimd.hpp>

using namespace std::chrono;

template <class T>
static void do_not_optimize(T& v) { asm volatile("" : "+x,m"(v) : : "memory"); }

template <class T>
static double bench(const char* name, std::size_t nmat, std::size_t reps)
{
    using arch = xsimd::default_arch;
    using batch = xsimd::batch<T, arch>;
    constexpr std::size_t N = batch::size; // matrix is N x N (N batches)

    std::vector<batch> data(nmat * N);
    for (std::size_t i = 0; i < data.size(); ++i)
    {
        alignas(64) T buf[N];
        for (std::size_t j = 0; j < N; ++j)
            buf[j] = static_cast<T>(i * N + j);
        data[i] = batch::load_aligned(buf);
    }

    // warmup
    for (std::size_t m = 0; m < nmat; ++m)
        xsimd::transpose(&data[m * N], &data[m * N] + N);

    auto t0 = steady_clock::now();
    for (std::size_t r = 0; r < reps; ++r)
        for (std::size_t m = 0; m < nmat; ++m)
        {
            xsimd::transpose(&data[m * N], &data[m * N] + N);
            do_not_optimize(data[m * N]);
        }
    auto t1 = steady_clock::now();

    double ns = duration_cast<nanoseconds>(t1 - t0).count();
    double per = ns / double(reps * nmat);
    std::printf("  %-10s %2zux%-2zu  %8.2f ns/transpose\n", name, N, N, per);
    return per;
}

int main()
{
    std::printf("arch = %s\n", xsimd::default_arch::name());
    const std::size_t nmat = 4096, reps = 2000;
    bench<float>("float", nmat, reps);
    bench<double>("double", nmat, reps);
    bench<int32_t>("int32", nmat, reps);
    bench<int64_t>("int64", nmat, reps);
    return 0;
}

===== MASTER =====
arch = avx512vnni+avx512vbmi2
  float      16x16     77.66 ns/transpose
  double      8x8      19.71 ns/transpose
  int32      16x16     78.70 ns/transpose
  int64       8x8      24.46 ns/transpose

===== PR #1371 =====
arch = avx512vnni+avx512vbmi2
  float      16x16     30.34 ns/transpose
  double      8x8       9.32 ns/transpose
  int32      16x16     31.01 ns/transpose
  int64       8x8      10.17 ns/transpose

[serge-sans-paille]

Great! It would be interesting to compare with some HPC routines to check how far we are from state of the art, but that's definitely a follow-up.