Fix AVX-512 AudioChannelMixer tail over-read; guard-page buffers in kernel tests

Common/Cpp/Containers/GuardPageBuffer.cpp

@@ -0,0 +1,130 @@
+/*  Guard Page Buffer
+ *
+ *  From: https://github.com/PokemonAutomation/
+ *
+ */
+
+#include <new>
+#include "Common/Cpp/Exceptions.h"
+#include "GuardPageBuffer.h"
+
+#ifdef _WIN32
+#include <Windows.h>
+#else
+#include <sys/mman.h>
+#include <unistd.h>
+#endif
+
+namespace PokemonAutomation{
+
+
+
+//  Return the OS page size in bytes: the unit in which the OS maps memory and
+//  sets access rights. It is 4 KB on x86 Windows and Linux, and 16 KB on
+//  Apple Silicon macOS.
+static size_t page_size(){
+#ifdef _WIN32
+    SYSTEM_INFO info;
+    GetSystemInfo(&info);
+    return info.dwPageSize;
+#else
+    return (size_t)sysconf(_SC_PAGESIZE);
+#endif
+}
+
+
+
+GuardPageMemory::~GuardPageMemory(){
+    release();
+}
+
+//  Take over the pages of "x". Afterwards "x" owns no memory, so its
+//  destructor does nothing.
+GuardPageMemory::GuardPageMemory(GuardPageMemory&& x) noexcept
+    : m_base(x.m_base)
+    , m_mapped_bytes(x.m_mapped_bytes)
+    , m_data(x.m_data)
+{
+    x.m_base = nullptr;
+    x.m_mapped_bytes = 0;
+    x.m_data = nullptr;
+}
+void GuardPageMemory::operator=(GuardPageMemory&& x) noexcept{
+    if (this == &x){
+        return;
+    }
+    release();
+    m_base = x.m_base;
+    m_mapped_bytes = x.m_mapped_bytes;
+    m_data = x.m_data;
+    x.m_base = nullptr;
+    x.m_mapped_bytes = 0;
+    x.m_data = nullptr;
+}
+
+//  Unmap the pages if this object still owns them. The destructor calls it,
+//  and so does the move assignment before it takes the pages of another
+//  object.
+void GuardPageMemory::release() noexcept{
+    if (m_base == nullptr){
+        return;
+    }
+#ifdef _WIN32
+    VirtualFree(m_base, 0, MEM_RELEASE);
+#else
+    munmap(m_base, m_mapped_bytes);
+#endif
+    m_base = nullptr;
+}
+
+//  Map the data pages and the guard page, then place the data so that it ends
+//  at the guard page:
+//  1. Round the data size up to whole pages. Add alignment - 1 bytes first,
+//     so that rounding the start down to "alignment" stays inside the mapping.
+//  2. Map all the pages as read-write (mmap() or VirtualAlloc()). New pages
+//     from the OS are all zero, so the buffer starts as zero bytes.
+//  3. Make the last page no-access (mprotect() with PROT_NONE, or
+//     VirtualProtect() with PAGE_NOACCESS). This is the guard page.
+//  4. Put the start of the data at guard - bytes, rounded down to "alignment".
+//  Example with 4096-byte pages, bytes = 72 and alignment = 4: it maps 1 data
+//  page and 1 guard page, and the data starts at byte 4024 of the first page,
+//  so the data ends at byte 4096, where the guard page starts.
+//  If a step fails, it unmaps what it mapped and throws std::bad_alloc.
+GuardPageMemory::GuardPageMemory(size_t bytes, size_t alignment){
+    if (alignment == 0 || (alignment & (alignment - 1)) != 0){
+        throw InternalProgramError(nullptr, PA_CURRENT_FUNCTION, "Alignment must be a power of two.");
+    }
+
+    size_t page = page_size();
+    size_t data_pages = (bytes + alignment - 1 + page - 1) / page;
+    m_mapped_bytes = (data_pages + 1) * page;
+
+#ifdef _WIN32
+    m_base = VirtualAlloc(nullptr, m_mapped_bytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    if (m_base == nullptr){
+        throw std::bad_alloc();
+    }
+    char* guard = (char*)m_base + data_pages * page;
+    DWORD old_protect;
+    if (!VirtualProtect(guard, page, PAGE_NOACCESS, &old_protect)){
+        VirtualFree(m_base, 0, MEM_RELEASE);
+        throw std::bad_alloc();
+    }
+#else
+    m_base = mmap(nullptr, m_mapped_bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    if (m_base == MAP_FAILED){
+        throw std::bad_alloc();
+    }
+    char* guard = (char*)m_base + data_pages * page;
+    if (mprotect(guard, page, PROT_NONE) != 0){
+        munmap(m_base, m_mapped_bytes);
+        throw std::bad_alloc();
+    }
+#endif
+
+    m_data = (void*)((size_t)(guard - bytes) & ~(alignment - 1));
+}
+
+
+
+}

Common/Cpp/Containers/GuardPageBuffer.h

@@ -0,0 +1,137 @@
+/*  Guard Page Buffer
+ *
+ *  From: https://github.com/PokemonAutomation/
+ *
+ *  A test-only array that crashes on any access past its end. Unit tests of
+ *  SIMD kernels use it to catch reads and writes past the end of a buffer.
+ *
+ */
+
+#ifndef PokemonAutomation_GuardPageBuffer_H
+#define PokemonAutomation_GuardPageBuffer_H
+
+#include <stddef.h>
+#include <type_traits>
+
+namespace PokemonAutomation{
+
+
+//  The memory behind `GuardPageBuffer`, without an element type.
+//
+//  It maps enough whole pages from the OS for "bytes" bytes (plus the room to
+//  round the start down to "alignment"), and then one more page that it marks
+//  as no-access: the guard page. `data()` points to the first byte of the
+//  data, placed so that the data ends right before the guard page.
+//
+//  The memory starts as all zero bytes, and the destructor unmaps it. It can
+//  be moved but not copied. A moved-from object owns no memory.
+//
+//  Throws std::bad_alloc if the OS cannot map or protect the pages, and
+//  InternalProgramError if "alignment" is not a power of two.
+class GuardPageMemory{
+public:
+    ~GuardPageMemory();
+    GuardPageMemory(GuardPageMemory&& x) noexcept;
+    void operator=(GuardPageMemory&& x) noexcept;
+    GuardPageMemory(const GuardPageMemory&) = delete;
+    void operator=(const GuardPageMemory&) = delete;
+
+public:
+    GuardPageMemory(size_t bytes, size_t alignment = 1);
+
+    //  First byte of the data. The data ends right before the guard page
+    //  when "bytes" is a multiple of "alignment".
+    void* data() const{ return m_data; }
+
+private:
+    void release() noexcept;
+
+private:
+    void* m_base;           //  Start of the mapped pages, or nullptr after a move.
+    size_t m_mapped_bytes;  //  Size of the mapping, guard page included.
+    void* m_data;           //  What `data()` returns.
+};
+
+
+
+//  A fixed-size array for unit tests whose last element sits right before a
+//  guard page: a page of memory that the program is not allowed to read or
+//  write. Any access just past the end of the array crashes the test at once.
+//
+//  Why it exists: SIMD kernels work on whole vectors (for example, 16 floats
+//  with AVX-512) and need special code for the last few elements, the "tail".
+//  A tail that loads or stores a full vector touches memory past the end of
+//  the buffer. With a normal heap buffer like `AlignedVector`, that memory
+//  usually belongs to the heap, so nothing happens. The test only crashes on
+//  the rare run where the buffer happens to end next to an unmapped page.
+//  That is how an over-read in the AVX-512 `AudioChannelMixer` tail crashed
+//  CI only now and then. With a GuardPageBuffer, every such run crashes.
+//
+//  How it works: `GuardPageMemory` maps whole pages and makes the last page
+//  no-access. The array is placed so that its end touches that page.
+//  Example with 4 KB pages and `GuardPageBuffer<float> buffer(18)`:
+//  1. The data needs 18 * 4 = 72 bytes, so it maps 1 data page and 1 guard
+//     page (8 KB in total).
+//  2. `buffer.data()` points 72 bytes before the guard page.
+//  3. `buffer[17]` is the last valid element. Reading `buffer[18]`, or a
+//     16-float vector load at `buffer.data() + 16`, crashes.
+//
+//  Where it is used: the kernel unit tests use it for the input and output
+//  buffers of the kernels under test, for example in
+//  `Kernels_AudioChannelMixer_TestInfra.cpp`. `GuardPageImageRGB32` in
+//  `Tests/GuardPageImage.h` uses it for images.
+//
+//  Limits:
+//  - Only trivial types like float and uint32_t. The elements are not
+//    constructed or destroyed, and they start as zero bytes.
+//  - Only accesses after the end are caught, not accesses before the start.
+//  - By default, `data()` is only aligned to alignof(Object). For a kernel
+//    that needs aligned pointers, pass a larger "alignment", for example
+//    PA_ALIGNMENT (64 bytes) for `fft_abs()`. The start is then rounded
+//    down, which can leave a gap of up to alignment - 1 bytes before the
+//    guard page. Accesses in that gap are not caught, so use sizes that are
+//    a multiple of the alignment when you can.
+//  - Each buffer maps at least 2 pages with a system call, so it is slower
+//    than a heap buffer and only meant for tests.
+//
+//  Throws std::bad_alloc if the OS cannot map or protect the pages, and
+//  InternalProgramError if "alignment" is not a power of two.
+template <typename Object>
+class GuardPageBuffer{
+    static_assert(
+        std::is_trivially_default_constructible_v<Object> &&
+        std::is_trivially_destructible_v<Object>,
+        "GuardPageBuffer only holds trivial types. The elements start as zero bytes."
+    );
+
+public:
+    //  Map "size" elements, all zero bytes. "alignment" is the alignment of
+    //  `data()` in bytes and must be a power of two. Example:
+    //  `GuardPageBuffer<float> real(length, PA_ALIGNMENT)` gives `fft_abs()`
+    //  the 64-byte aligned input that it requires.
+    GuardPageBuffer(size_t size, size_t alignment = alignof(Object))
+        : m_memory(size * sizeof(Object), alignment)
+        , m_size(size)
+    {}
+
+    //  Number of elements, as given to the constructor.
+    size_t size() const{ return m_size; }
+
+    //  Element access without a bounds check. An index at or just past
+    //  size() reads the guard page and crashes, which is what tests want.
+    const Object& operator[](size_t index) const{ return data()[index]; }
+          Object& operator[](size_t index)      { return data()[index]; }
+
+    //  Pointer to the first element, to give to the kernel under test.
+    const Object* data() const{ return (const Object*)m_memory.data(); }
+          Object* data()      { return (Object*)m_memory.data(); }
+
+private:
+    GuardPageMemory m_memory;
+    size_t m_size;
+};
+
+
+
+}
+#endif

SerialPrograms/Source/Kernels/AbsFFT/Kernels_AbsFFT_TestInfra.cpp

@@ -11,7 +11,7 @@
 #include <cstdio>
 #include <random>
 #include "Common/Compiler.h"
-#include "Common/Cpp/Containers/AlignedVector.tpp"
+#include "Common/Cpp/Containers/GuardPageBuffer.h"
 #include "Kernels_AbsFFT_TestInfra.h"
 
 namespace PokemonAutomation{
@@ -96,11 +96,16 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{
     //  Random input with a fixed seed. fft_abs() is destructive on "real",
     //  so each run gets its own copy of the same data.
     //
-    AlignedVector<float> real0(length);
-    AlignedVector<float> real1(length);
-    AlignedVector<float> abs0(half_length);
-    AlignedVector<float> abs1(half_length);
-    AlignedVector<float> weights(length);
+    //  Every buffer ends right before a guard page, so a read or write past
+    //  the end faults on every run. "real" and "abs" must be aligned to 64
+    //  bytes. Their sizes are powers of two, so for all but the smallest
+    //  transforms they still end exactly at the guard page.
+    //
+    GuardPageBuffer<float> real0(length, PA_ALIGNMENT);
+    GuardPageBuffer<float> real1(length, PA_ALIGNMENT);
+    GuardPageBuffer<float> abs0(half_length, PA_ALIGNMENT);
+    GuardPageBuffer<float> abs1(half_length, PA_ALIGNMENT);
+    GuardPageBuffer<float> weights(length);
 
     std::mt19937 generator(42);
     std::uniform_int_distribution<int> sample_distribution(-10000, 10000);

SerialPrograms/Source/Kernels/AudioChannelInterleave/Kernels_AudioChannelInterleave_TestInfra.cpp

@@ -5,7 +5,7 @@
  */
 
 #include <random>
-#include "Common/Cpp/Containers/AlignedVector.tpp"
+#include "Common/Cpp/Containers/GuardPageBuffer.h"
 #include "Kernels_AudioChannelInterleave_TestInfra.h"
 
 namespace PokemonAutomation{
@@ -32,9 +32,11 @@ Test::Test(
 UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{
     size_t total = m_channels * m_samples;
 
-    AlignedVector<float> interleaved0(total);
-    AlignedVector<float> interleaved1(total);
-    AlignedVector<float> strided(total);
+    //  Every buffer ends right before a guard page, so a kernel that reads or
+    //  writes past the end faults on every run.
+    GuardPageBuffer<float> interleaved0(total);
+    GuardPageBuffer<float> interleaved1(total);
+    GuardPageBuffer<float> strided(total);
 
     std::mt19937 generator(42);
 

SerialPrograms/Source/Kernels/AudioChannelMixer/Kernels_AudioChannelMixer_TestInfra.cpp

@@ -5,7 +5,7 @@
  */
 
 #include <random>
-#include "Common/Cpp/Containers/AlignedVector.tpp"
+#include "Common/Cpp/Containers/GuardPageBuffer.h"
 #include "Kernels_AudioChannelMixer_TestInfra.h"
 
 namespace PokemonAutomation{
@@ -35,9 +35,13 @@ Test::Test(
 
 UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{
     std::vector<float> matrix(m_in_channels * m_out_channels);
-    AlignedVector<float> in(m_samples * m_in_channels);
-    AlignedVector<float> out0(m_samples * m_out_channels);
-    AlignedVector<float> out1(m_samples * m_out_channels);
+
+    //  The input and both outputs end right before a guard page. A tail that
+    //  reads or writes past the end then faults on every run, not just when
+    //  the heap happens to end there.
+    GuardPageBuffer<float> in(m_samples * m_in_channels);
+    GuardPageBuffer<float> out0(m_samples * m_out_channels);
+    GuardPageBuffer<float> out1(m_samples * m_out_channels);
 
     std::mt19937 generator(42);
 
@@ -63,7 +67,7 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{
     );
 
     //  Check for errors.
-    for (size_t c = 0; c < m_samples; c++){
+    for (size_t c = 0; c < m_samples * m_out_channels; c++){
         if (out0[c] != out1[c]){
             return "Mismatch at " + std::to_string(c) + ": " +
                 std::to_string(out0[c]) + " != " + std::to_string(out1[c]);

SerialPrograms/Source/Kernels/AudioChannelMixer/Kernels_AudioChannelMixer_x64_AVX512.cpp

@@ -126,8 +126,8 @@ void mix_channels_x64_AVX512<2, 1>(
     }
     {
         __mmask16 mask = ((uint32_t)1 << samples) - 1;
-        __m512 out0 = _mm512_mul_ps(i0o0, _mm512_maskz_loadu_ps(mask, in + in_stride*0));
-        out0 = _mm512_fmadd_ps(i1o0, _mm512_loadu_ps(in + in_stride * 1), out0);
+        __m512 out0 = _mm512_mul_ps(i0o0, _mm512_maskz_loadu_ps(mask, in + in_stride * 0));
+        out0 = _mm512_fmadd_ps(i1o0, _mm512_maskz_loadu_ps(mask, in + in_stride * 1), out0);
         _mm512_mask_storeu_ps(out, mask, out0);
     }
 }

SerialPrograms/Source/Kernels/AudioPipeline/Kernels_AudioPipeline_RoundTrip_TestInfra.cpp

@@ -10,6 +10,7 @@
 #include "Common/Cpp/Exceptions.h"
 #include "Common/Cpp/PrettyPrint.h"
 #include "Common/Cpp/Containers/AlignedVector.tpp"
+#include "Common/Cpp/Containers/GuardPageBuffer.h"
 #include "Kernels_AudioPipeline_RoundTrip_TestInfra.h"
 #include "Kernels_RawAudio2WorkingAudio.h"
 #include "Kernels_WorkingAudio2RawAudio.h"
@@ -82,8 +83,9 @@ UnitTestResult run_round_trip(
     size_t samples,
     float multiplier
 ){
-    //  Generate random input.
-    AlignedVector<Type> input(samples * channels);
+    //  Generate random input. It ends right before a guard page, so a read
+    //  past the end faults.
+    GuardPageBuffer<Type> input(samples * channels);
     std::mt19937 generator(42);
     if constexpr (std::is_same_v<Type, float>){
         //  Random integer values, all exactly representable.