cosine_similarity.cs

using System.Runtime.InteropServices;
using System.Runtime.Intrinsics;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Jobs;
using BenchmarkDotNet.Running;

#pragma warning disable CS8618

var summary = BenchmarkRunner.Run<MainTest>();

// = 1K
// | Method                     | Mean       | Error   | StdDev  | Ratio | Allocated | Alloc Ratio |
// |--------------------------- |-----------:|--------:|--------:|------:|----------:|------------:|
// | cosine_similarity_simple   | 1,299.0 ns | 3.23 ns | 3.02 ns |  1.00 |         - |          NA |
// | cosine_similarity_SIMD_256 |   171.0 ns | 0.50 ns | 0.44 ns |  0.13 |         - |          NA |
// | cosine_similarity_SIMD_512 |   130.4 ns | 0.35 ns | 0.32 ns |  0.10 |         - |          NA |

// = 512K
// | Method                     | Mean      | Error    | StdDev   | Ratio | Allocated | Alloc Ratio |
// |--------------------------- |----------:|---------:|---------:|------:|----------:|------------:|
// | cosine_similarity_simple   | 679.45 us | 3.438 us | 3.216 us |  1.00 |         - |          NA |
// | cosine_similarity_SIMD_256 |  88.13 us | 0.265 us | 0.248 us |  0.13 |         - |          NA |
// | cosine_similarity_SIMD_512 |  76.38 us | 0.283 us | 0.265 us |  0.11 |         - |          NA |

[SimpleJob(RuntimeMoniker.Net90, baseline: true)]
// [WarmupCount(2)]
// [IterationCount(2)]
[MemoryDiagnoser]
public class MainTest
{
    const    int      N    = 1024 + 3;
    readonly double[] arr1 = Enumerable.Range(0, N).Select(_ => Random.Shared.NextDouble()).ToArray();
    readonly double[] arr2 = Enumerable.Range(0, N).Select(_ => Random.Shared.NextDouble()).ToArray();

    [Benchmark(Baseline = true)]
    public double cosine_similarity_simple()
    {
        return cosine_similarity_simple(arr1, arr2);
    }

    double cosine_similarity_simple(double[] x, double[] y)
    {
        ArgumentNullException.ThrowIfNull(x);
        ArgumentNullException.ThrowIfNull(y);
        if (x.Length != y.Length)
            throw new ArgumentException("Arrays must have the same length");

        var sumScalar = 0.0;
        for (var i = 0; i < x.Length; i++)
            sumScalar += x[i] * y[i];

        var sumSquaredX = 0.0;
        var sumSquaredY = 0.0;
        for (var i = 0; i < x.Length; i++)
        {
            sumSquaredX += x[i] * x[i];
            sumSquaredY += y[i] * y[i];
        }

        var k = sumScalar / (Math.Sqrt(sumSquaredX) * Math.Sqrt(sumSquaredY));
        return k;
    }

    [Benchmark]
    public double cosine_similarity_SIMD_256()
    {
        return cosine_similarity_SIMD_256(arr1, arr2);
    }

    double cosine_similarity_SIMD_256(double[] x, double[] y)
    {
        ArgumentNullException.ThrowIfNull(x);
        ArgumentNullException.ThrowIfNull(y);
        if (x.Length != y.Length)
            throw new ArgumentException("Arrays must have the same length");

        var vx = MemoryMarshal.Cast<double, Vector256<double>>(x);
        var vy = MemoryMarshal.Cast<double, Vector256<double>>(y);

        var sumAccum    = Vector256<double>.Zero;
        var sumSquaredX = Vector256<double>.Zero;
        var sumSquaredY = Vector256<double>.Zero;
        for (var i = 0; i < vx.Length; i++)
        {
            var x1 = vx[i];
            var y1 = vy[i];
            sumAccum    += x1 * y1;
            sumSquaredX += x1 * y1;
            sumSquaredY += x1 * y1;
        }

        var sumScalar         = sumAccum[0]    + sumAccum[1]    + sumAccum[2]    + sumAccum[3];
        var sumSquaredXScalar = sumSquaredX[0] + sumSquaredX[1] + sumSquaredX[2] + sumSquaredX[3];
        var sumSquaredYScalar = sumSquaredY[0] + sumSquaredY[1] + sumSquaredY[2] + sumSquaredY[3];
        for (var i = vx.Length * 4; i < x.Length; i++)
        {
            sumScalar         += x[i] * y[i];
            sumSquaredXScalar += x[i] * x[i];
            sumSquaredYScalar += y[i] * y[i];
        }

        var k = sumScalar / (Math.Sqrt(sumSquaredXScalar) * Math.Sqrt(sumSquaredYScalar));
        return k;
    }

    [Benchmark]
    public double cosine_similarity_SIMD_512()
    {
        return cosine_similarity_SIMD_512(arr1, arr2);
    }

    double cosine_similarity_SIMD_512(double[] x, double[] y)
    {
        ArgumentNullException.ThrowIfNull(x);
        ArgumentNullException.ThrowIfNull(y);
        if (x.Length != y.Length)
            throw new ArgumentException("Arrays must have the same length");

        var vx = MemoryMarshal.Cast<double, Vector512<double>>(x);
        var vy = MemoryMarshal.Cast<double, Vector512<double>>(y);

        var sumAccum    = Vector512<double>.Zero;
        var sumSquaredX = Vector512<double>.Zero;
        var sumSquaredY = Vector512<double>.Zero;
        for (var i = 0; i < vx.Length; i++)
        {
            var x1 = vx[i];
            var y1 = vy[i];
            sumAccum    += x1 * y1;
            sumSquaredX += x1 * y1;
            sumSquaredY += x1 * y1;
        }

        var sumScalar         = sumAccum[0]    + sumAccum[1]    + sumAccum[2]    + sumAccum[3];
        var sumSquaredXScalar = sumSquaredX[0] + sumSquaredX[1] + sumSquaredX[2] + sumSquaredX[3];
        var sumSquaredYScalar = sumSquaredY[0] + sumSquaredY[1] + sumSquaredY[2] + sumSquaredY[3];
        for (var i = vx.Length * 8; i < x.Length; i++)
        {
            sumScalar         += x[i] * y[i];
            sumSquaredXScalar += x[i] * x[i];
            sumSquaredYScalar += y[i] * y[i];
        }

        var k = sumScalar / (Math.Sqrt(sumSquaredXScalar) * Math.Sqrt(sumSquaredYScalar));
        return k;
    }
}