Optimizing Vector Storage with Aligned Memory

The Problem with Tightly Packed Data

In our current k-NN implementation, you are likely allocating one large block of memory for all your vectors, something like this:

float* vectors = malloc(num_vectors * dimensions * sizeof(float));

While simple, this has two potential performance issues that we must address:

1. The block returned by malloc is not guaranteed to start on a 64-byte cache line boundary.
2. Even if it were, the starting address of each subsequent vector (vectors[1], vectors[2], etc.) is almost certainly not aligned to a 64-byte boundary unless the number of dimensions happens to result in a byte size that is a perfect multiple of 64.

The first problem is what we solved with posix_memalign in the last lesson. The second problem is more subtle but equally important. When we introduce SIMD (Single Instruction, Multiple Data) operations, the CPU will load data into wide 256-bit or 512-bit registers. These operations are fastest when the memory address they load from is aligned to the register size (e.g., 32 or 64 bytes). An unaligned access can force a slower code path or, worse, require loading two cache lines to fill one register.

The goal is to move from a simple, tightly packed array to a layout where each vector resides in its own aligned memory segment.

Structuring Data for Performance

The key to solving this is to treat each vector as a record and ensure each record starts on an aligned boundary. We achieve this by calculating a stride: the total number of bytes for each vector, including any padding needed to reach the next alignment boundary.

The principle is clearly stated in the following resource.

Let's apply these rules.

1. Starting Address: We'll use our safe_aligned_alloc function to ensure the entire block of memory starts on a 64-byte boundary.
2. Size of "Class": We will calculate a padded size, or stride, for each vector so that its total size is a multiple of 64 bytes. This guarantees that vector[i] and vector[i+1] do not share a cache line in a way that could cause a split access, and ensures vector[i+1] also starts on an aligned boundary relative to the start of the whole block.

For example, if our vectors have 100 dimensions (400 bytes) and we want 64-byte alignment, the stride for each vector would be 448 bytes (the next multiple of 64). The last 48 bytes would be unused padding.

The following video explains the importance of alignment specifically in the context of the AVX intrinsics we will be using later.

Your Task: Refactor the Vector Storage

It's time to modify your project. You will replace your simple malloc with a more sophisticated allocation strategy that incorporates padding to ensure each vector is aligned.

Here are the steps to follow, along with code snippets to guide you. Assume you have your safe_aligned_alloc function from the previous lesson available.

1. Define Alignment and Calculate Stride:
   In your data storage file (e.g., database.c), define the alignment you want. 64 bytes is an excellent choice as it matches the cache line size on most modern CPUs. Then, calculate the stride.

   
   
   
   #define VECTOR_ALIGNMENT 64
   
   // Inside your function that allocates memory for vectors...
   // Let's say you have 'db_header->num_vectors' and 'db_header->dimensions'
   
   // The actual size of one vector's data
   size_t vector_data_size = db_header->dimensions * sizeof(float);
   
   // Calculate the padded size (stride). This rounds up vector_data_size to the nearest multiple of VECTOR_ALIGNMENT.
   size_t vector_stride_bytes = (vector_data_size + VECTOR_ALIGNMENT - 1) / VECTOR_ALIGNMENT * VECTOR_ALIGNMENT;
   
   // We can also store the stride in floats for easier indexing later
   size_t vector_stride_floats = vector_stride_bytes / sizeof(float);
   
   // Store this stride in your database header/struct for later use.
   db->vector_stride = vector_stride_floats; // Or in bytes, be consistent!
   

2. Allocate Aligned and Padded Memory:
   Calculate the total size needed based on the stride and allocate the memory using your aligned allocator.

   // Total size is based on the stride, not the raw data size
   size_t total_alloc_size = db_header->num_vectors * vector_stride_bytes;
   
   // Allocate the main block, aligned to 64 bytes
   db->vectors = safe_aligned_alloc(VECTOR_ALIGNMENT, total_alloc_size);
   if (db->vectors == NULL) {
       // Handle allocation failure
       return -1;
   }
   

3. Create an Accessor Function:
   Since the vectors are no longer tightly packed, you can no longer use simple pointer arithmetic like db->vectors + i * dimensions. You must use the stride. It's best practice to create a helper function for this.

   /**
    * @brief Gets a pointer to the beginning of a specific vector in the database.
    *
    * @param db The database instance.
    * @param index The index of the vector to retrieve.
    * @return A pointer to the first float of the requested vector.
    */
   static inline float* get_vector(const Database* db, size_t index) {
       // Here we assume db->vector_stride is in units of floats
       return db->vectors + index * db->vector_stride;
   }
   

4. Update Your Code:
   Go through your k-NN search function and any other code that accesses vectors. Replace all direct calculations of a vector's address with a call to your new get_vector() helper function.

   For example, a loop in your brute-force search might change from this:

   // Before
   for (size_t i = 0; i < db->header.num_vectors; ++i) {
       float* db_vector = db->vectors + i * db->header.dimensions;
       // ... calculate distance ...
   }
   

   To this:

   // After
   for (size_t i = 0; i < db->header.num_vectors; ++i) {
       float* db_vector = get_vector(db, i);
       // ... calculate distance ...
   }
   

This refactoring is a significant step towards a truly high-performance system. You are now explicitly controlling data layout in a way that "respects the silicon," paving the way for massive speedups with SIMD.

Conclusion

In this lesson, you have taken a critical step from academic understanding of memory alignment to practical, low-level implementation. By refactoring your vector storage, you have done more than just fix a potential performance issue; you have fundamentally changed your data layout to prepare for advanced instruction-level parallelism.

Key Takeaways:

• Optimal memory access requires not only that the start of a data block is aligned, but also that individual elements within it are aligned, especially when preparing for SIMD.
• We can enforce alignment for each vector by calculating a stride—the size of the vector's data padded out to the next multiple of our desired alignment (e.g., 64 bytes).
• Accessing data in a strided array requires a helper function that uses the stride for pointer arithmetic, abstracting the memory layout from the logic that uses it.

In the next lesson, we will shift focus slightly to improve the efficiency of the search algorithm itself. You will use a bounded max-heap data structure to efficiently track the top 'k' candidates during a search. This will reduce the computational overhead of sorting candidates at the end of the search, which becomes increasingly important as we start to process candidates faster and faster.