Squared L2 Norm with AVX2

The L2 Norm and its Squared Form

The L2 norm, also known as the Euclidean norm, is the geometric length of a vector. For a vector $\mathbf{v}=(v_1,v_2,...,v_n)$, it is defined as:

In many applications, including the intermediate steps of computing Cosine Similarity, it's more efficient to work with the squared L2 norm. This avoids a computationally expensive square root operation until the very end. The squared L2 norm is simply:

Look closely at this formula. Does it remind you of anything? A sum of products. This is precisely the definition of a dot product, but where the vector is multiplied by itself:

This crucial insight means that all the work we did to optimize the dot product in the last lesson is directly applicable here. We don't need to reinvent the wheel; we just need to adapt our existing high-performance pattern.

A Basic SIMD Approach: Multiply and Add

Before we jump to the most optimized solution, let's look at a straightforward SIMD implementation. The core task is to square each element of the vector and sum the results. This can be done with a multiplication followed by an addition.

The following resource demonstrates exactly this. While the article's goal is to normalize a vector (dividing by the norm), the first part of its SIMD function calculates the sum of squares, which is precisely our squared L2 norm.

This approach works and is a correct vectorization of the formula. However, we can do better. As you learned in the previous lesson, combining a multiplication and an addition into a single Fused Multiply-Add (FMA) operation is more efficient, both in terms of throughput and numerical precision.

The Optimal FMA-based Implementation

The calculation $\text{sum}=\text{sum}+(v_i\cdot v_i)$ is a perfect use case for the _mm256_fmadd_ps intrinsic. The logic becomes identical to our unrolled dot product implementation, with the simple change that both vector inputs to the FMA operation are the same.

Let's revisit the high-performance dot product code from the last lesson and see how to adapt it.

With that modification in mind, a complete, high-performance l2_norm_squared function would look strikingly similar to the dotProductFma function.

Here is the structure for your own implementation. Note its near-identical form to the function from the last lesson:

1. Function Signature: It will take a single vector (e.g., const float* data, size_t size) as input.
2. Initialization: Initialize four __m256 accumulator registers (e.g., dot0, dot1, dot2, dot3) to zero using _mm256_setzero_ps().
3. Main Loop (Unrolled): The loop processes 32 floats (4 AVX registers) per iteration. Inside the loop, you will:
   • Load a 256-bit chunk of the vector: const __m256 vec = _mm256_loadu_ps(p);
   • Update an accumulator using FMA, passing the same vector register twice: acc = _mm256_fmadd_ps(vec, vec, acc);
   • Repeat this for all four accumulators, advancing the data pointer by 8 floats for each.
4. Remainder Handling: (Optional for now, but crucial for a robust library) Handle cases where the vector length is not a multiple of 32, using masked loads as shown in the updated Stack Overflow answer.
5. Reduction: Perform the exact same horizontal reduction sequence as in the dot product function to sum the four accumulator vectors into a single scalar float result.

This exercise demonstrates a key principle in performance engineering: recognizing and reusing optimized computational patterns. The dot product and squared L2 norm are mathematically related, and their low-level implementations can be nearly identical.

Conclusion

In this lesson, you've learned how to implement a SIMD-accelerated function for the squared L2 norm. By recognizing its mathematical relationship to the dot product, you were able to leverage the powerful FMA and loop unrolling techniques from our previous session to create a highly efficient implementation.

Key Takeaways:

• The squared L2 norm, $||\mathbf{v}|{|}_2^2$, is mathematically equivalent to the dot product of a vector with itself, $\mathbf{v}\cdot \mathbf{v}$.
• A performant SIMD implementation of the squared L2 norm reuses the same FMA-based, unrolled-loop structure as a dot product.
• The core operation in the loop becomes _mm256_fmadd_ps(vec, vec, accumulator), where vec is a chunk of the input vector.
• The horizontal reduction logic to obtain the final scalar result is identical to that of the dot product.

We now have optimized functions for both the dot product and the squared L2 norm. In our next and final lesson of this module, we will combine these pieces to assemble the complete, SIMD-accelerated Cosine Similarity function.