Safe vectorization is a lot slower than unsafe

[EgorBo]

Currently, when we do code vectorization in C#, we switch to raw pointers (either managed or unmanaged) and unsafe APIs resulting in lack of safety/bounds checks. A good example is a function that calculates a sum using SIMD, today we would write it like this:

ref int r = ref MemoryMarshal.GetReference(span);
for (int i = 0; i < span.Length - Vector128<int>.Count; i += Vector128<int>.Count)
{
    vSum += Vector128.LoadUnsafe(ref r, (nuint)i);
}

It would be nice to move away from this practice towards fully safe APIs:

-ref int r = ref MemoryMarshal.GetReference(span);
for (int i = 0; i < span.Length - Vector128<int>.Count; i += Vector128<int>.Count)
{
-    vSum += Vector128.LoadUnsafe(ref r, (nuint)i);
+    vSum += Vector128.Create(span[i..]);
}

Unfortunately, JIT is not able to eliminate safety checks produced by Slice despite those being redundant in this case. This makes the safer version up to 2x slower, see EgorBot/runtime-utils#226

Full C# impl we want to have 0 redundant bounds check in it:

public static int Sum(ReadOnlySpan<int> span)
{
    Vector128<int> vSum = default;

    // Main loop
    int i;
    for (i = 0; i < span.Length - Vector128<int>.Count; i += Vector128<int>.Count)
        vSum += Vector128.Create(span[i..]);

    // Horizontal sum
    int sum = Vector128.Sum(vSum);

    // Trailing elements
    for (; i < span.Length; i++)
        sum += span[i];

    return sum;
}

Current codegen

; Method Bench:Sum(System.ReadOnlySpan`1[int]):int (FullOpts)
G_M000_IG01:                ;; offset=0x0000
       push     rbx
       sub      rsp, 32

G_M000_IG02:                ;; offset=0x0005
       vxorps   xmm0, xmm0, xmm0
       xor      eax, eax
       mov      edx, dword ptr [rcx+0x08]
       lea      r8d, [rdx-0x04]
       test     r8d, r8d
       jle      SHORT G_M000_IG05

G_M000_IG03:                ;; offset=0x0017
       mov      r10d, edx
       align    [6 bytes for IG04]

G_M000_IG04:                ;; offset=0x0020
       mov      r9d, edx
       sub      r9d, eax
       mov      r11d, eax
       mov      ebx, r9d
       add      rbx, r11
       cmp      rbx, r10
       ja       SHORT G_M000_IG11
       mov      rbx, bword ptr [rcx]
       lea      r11, bword ptr [rbx+4*r11]
       cmp      r9d, 4
       jl       SHORT G_M000_IG10
       vpaddd   xmm0, xmm0, xmmword ptr [r11]
       add      eax, 4
       cmp      eax, r8d
       jl       SHORT G_M000_IG04

G_M000_IG05:                ;; offset=0x004E
       vpsrldq  xmm1, xmm0, 8
       vpaddd   xmm0, xmm1, xmm0
       vpsrldq  xmm1, xmm0, 4
       vpaddd   xmm0, xmm1, xmm0
       vmovd    r8d, xmm0
       cmp      eax, edx
       jge      SHORT G_M000_IG08

G_M000_IG06:                ;; offset=0x0069
       mov      rbx, bword ptr [rcx]
       align    [0 bytes for IG07]

G_M000_IG07:                ;; offset=0x006C
       cmp      eax, edx
       jae      SHORT G_M000_IG12
       mov      ecx, eax
       add      r8d, dword ptr [rbx+4*rcx]
       inc      eax
       cmp      eax, edx
       jl       SHORT G_M000_IG07

G_M000_IG08:                ;; offset=0x007C
       mov      eax, r8d

G_M000_IG09:                ;; offset=0x007F
       add      rsp, 32
       pop      rbx
       ret      

G_M000_IG10:                ;; offset=0x0085
       mov      ecx, 6
       call     [System.ThrowHelper:ThrowArgumentOutOfRangeException(int)]
       int3     

G_M000_IG11:                ;; offset=0x0091
       call     [System.ThrowHelper:ThrowArgumentOutOfRangeException()]
       int3     

G_M000_IG12:                ;; offset=0x0098
       call     CORINFO_HELP_RNGCHKFAIL
       int3     
; Total bytes of code: 158

There are other loop shapes typically used with SIMD.

[dotnet-policy-service]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

[EgorBo]

@EgorBot -amd -arm

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

public class Bench
{
    private static int[] Testdata = new int[10000];

    [Benchmark]
    public int Sum_safe()
    {
        ReadOnlySpan<int> span = Testdata;

        Vector128<int> vSum = default;
        int i;
        for (i = 0; i < span.Length - Vector128<int>.Count; i += Vector128<int>.Count)
            vSum += Vector128.Create(span[i..]);
        int sum = Vector128.Sum(vSum);
        for (; i < span.Length; i++)
            sum += span[i];
        return sum;
    }

    [Benchmark(Baseline = true)]
    public int Sum_unsafe()
    {
        ReadOnlySpan<int> span = Testdata;

        Vector128<int> vSum = default;
        ref int r = ref MemoryMarshal.GetReference(span);
        int i;
        for (i = 0; i < span.Length - Vector128<int>.Count; i += Vector128<int>.Count)
            vSum += Vector128.LoadUnsafe(ref r, (nuint)i);
        int sum = Vector128.Sum(vSum);
        for (; i < span.Length; i++)
            sum += span[i];
        return sum;
    }
}

[neon-sunset]

There could also be a few helper methods provided by CoreLib as a middle ground for simpler vectorized loop implementations. For example, Rust exposes https://doc.rust-lang.org/std/primitive.slice.html#method.as_simd_mut which returns a tuple of scalar prefix, aligned SIMD-typed middle section and scalar suffix. It's not ideal if you can use masking/overlapping or other alternatives for handling short inputs and/or remainders, but it's quite nice when you know you just want to make the "long input" case go fast.

[EgorBo]

cc @AndyAyersMS this issue I was talking about

[tannergooding]

I'm not sure I'd call the pattern above necessarily representative of what the BCL would write, so I don't think it's a pattern that would allow us to move off of the unsafe logic.

Rather, it's more the baseline pattern that we'd want to support for allowing some code to move off. Which is generally that this shape of loop (given that Vector128<T>.Count is a constant):

for (int i = 0; i < (span.Length - constant); i += constant)

Should allow bounds checks to be elided for span[i] (inclusive) through span[i + constant] (exclusive), as well as allow the following branch to be elided:

if (span.Length < constant)
{
    ArgumentOutOfRangeException.Throw(...);
}

---

More generally the patterns that need to be supported are a bit more expansive and also include things like manual unrolling or tracking a remainder:

int remainder = span.Length;

while (remainder >= (Vector128<T>.Count * 4))
{
    vector1 = Vector128.Create(span);
    vector2 = Vector128.Create(span[Vector128<T>.Count..]);
    vector3 = Vector128.Create(span[(Vector128<T>.Count * 2)..]);
    vector4 = Vector128.Create(span[(Vector128<T>.Count * 3)..]);

    remainder -= Vector128<T>.Count * 4;
}

scenarios like adjusting the baseline to handle remaining elements in a single vector operation:

if (span.Length > constant)
{
    int i;
    for (i = 0; i < (span.Length - constant); i += constant)
    {
        M(Vector128.Create(span[i..]));
    }

    i = span.Length - constant;
    M(Vector128.Create(span[i..]));
}

A more complex example of the previous pattern that is done alongside unrolling is:

nuint endIndex = remainder;
remainder = (remainder + (uint)(Vector128<T>.Count - 1)) & (nuint)(-Vector128<T>.Count);
 
switch (remainder / (uint)Vector128<T>.Count)
{
    case 4:
    {
        Vector128<T> vector = M(Vector128.Create(x[remainder - (uint)(Vector128<T>.Count * 4..]));
        vector.CopyTo(d[remainder - (uint)(Vector128<T>.Count * 4)..]);
        goto case 3;
    }

    case 3: // ...
}

and so on; with the general goal likely just getting a better understanding of the valid ranges given a higher condition that was met, thus allowing bounds check elimination for safe arithmetic.

[tfenise]

Currently it is possible to do the following:

//Unsafeness encapsulated
readonly ref struct Vector128Span<T> where T : struct
{
    readonly Span<Vector128<T>> _spanVector;//Unaligned
    internal readonly Span<T> _spanRemainder;

    internal Vector128Span(Span<T> span)
    {
        if (!Vector128<T>.IsSupported) throw new NotSupportedException();

        _spanVector = MemoryMarshal.Cast<T, Vector128<T>>(span);
        _spanRemainder = span.Slice(_spanVector.Length * Vector128<T>.Count);
    }

    internal int VectorLength => _spanVector.Length;
    internal Vector128<T> ReadVector(int index) => Unsafe.ReadUnaligned<Vector128<T>>(ref Unsafe.As<Vector128<T>, byte>(ref _spanVector[index]));
    internal void WriteVector(int index, Vector128<T> v) => Unsafe.WriteUnaligned<Vector128<T>>(ref Unsafe.As<Vector128<T>, byte>(ref _spanVector[index]), v);
}

//Safe code
static void Xor(Span<byte> bits1, Span<byte> bits2)
{
    Vector128Span<byte> vs1 = new(bits1), vs2 = new(bits2);

    if (vs1.VectorLength != vs2.VectorLength || vs1._spanRemainder.Length != vs2._spanRemainder.Length)
        throw new ArgumentException();

    for (int i = 0; i < vs1.VectorLength; i++)
    {
        vs1.WriteVector(i, vs1.ReadVector(i) ^ vs2.ReadVector(i));
    }
    for (int i = 0; i < vs1._spanRemainder.Length; i++)
    {
        vs1._spanRemainder[i] = (byte)(vs1._spanRemainder[i] ^ vs2._spanRemainder[i]);
    }
}

The codegen is decent without extra range checks in the loop.

Disasm

; Assembly listing for method ConsoleApp1.SafeVec:Xor(System.Span`1[ubyte],System.Span`1[ubyte]) (Tier1)
; Emitting BLENDED_CODE for X64 with AVX512 - Windows
; Tier1 code
; optimized code
; optimized using Synthesized PGO
; rsp based frame
; fully interruptible
; with Synthesized PGO: fgCalledCount is 71
; 4 inlinees with PGO data; 10 single block inlinees; 2 inlinees without PGO data

G_M000_IG01:                ;; offset=0x0000
       push     rsi
       push     rbx
       sub      rsp, 40

G_M000_IG02:                ;; offset=0x0006
       mov      rax, bword ptr [rcx]
       mov      ecx, dword ptr [rcx+0x08]
       mov      r8d, ecx
       shr      r8d, 4
       mov      r10d, r8d
       shl      r10d, 4
       cmp      r10d, ecx
       ja       SHORT G_M000_IG08
       mov      r9d, r10d
       add      r9, rax
       sub      ecx, r10d
       mov      r10, bword ptr [rdx]
       mov      edx, dword ptr [rdx+0x08]
       mov      r11d, edx
       shr      r11d, 4
       mov      ebx, r11d
       shl      ebx, 4
       cmp      ebx, edx
       ja       SHORT G_M000_IG08
       mov      esi, ebx
       add      rsi, r10
       sub      edx, ebx
       cmp      r8d, r11d
       jne      SHORT G_M000_IG09
       cmp      ecx, edx
       jne      SHORT G_M000_IG09
       test     r8d, r8d
       jle      SHORT G_M000_IG05

G_M000_IG03:                ;; offset=0x0058
       xor      edx, edx
       align    [6 bytes for IG04]

G_M000_IG04:                ;; offset=0x0060
       vmovups  xmm0, xmmword ptr [rax+rdx]
       vpxor    xmm0, xmm0, xmmword ptr [r10+rdx]
       vmovups  xmmword ptr [rax+rdx], xmm0
       add      rdx, 16
       dec      r8d
       jne      SHORT G_M000_IG04

G_M000_IG05:                ;; offset=0x0079
       xor      eax, eax
       test     ecx, ecx
       jle      SHORT G_M000_IG07
       align    [1 bytes for IG06]

G_M000_IG06:                ;; offset=0x0080
       mov      edx, eax
       movzx    r8, byte  ptr [rsi+rdx]
       xor      byte  ptr [r9+rdx], r8b
       inc      eax
       cmp      eax, ecx
       jl       SHORT G_M000_IG06

G_M000_IG07:                ;; offset=0x0091
       add      rsp, 40
       pop      rbx
       pop      rsi
       ret

G_M000_IG08:                ;; offset=0x0098
       call     [System.ThrowHelper:ThrowArgumentOutOfRangeException()]
       int3

G_M000_IG09:                ;; offset=0x009F
       mov      rcx, 0x7FFB2733C080
       call     CORINFO_HELP_NEWSFAST
       mov      rbx, rax
       mov      rcx, rbx
       call     [System.ArgumentException:.ctor():this]
       mov      rcx, rbx
       call     CORINFO_HELP_THROW
       int3

; Total bytes of code 195

However, it only works with relatively simple operations, and it is a bit hacky because internally it keeps an unaligned Span<Vector128<T>>.

It could be practised in places like BitArray to make them look safer.

Side notes

I would expect the following to produce better codegen, as bits1.Length now must be equal to bits2.Length and the arithmetic around the lengths wouldn't need to be duplicated:

internal static void Xor(Span<byte> bits1, Span<byte> bits2)
{
    if (bits1.Length != bits2.Length)
        throw new ArgumentException();

    Vector128Span<byte> vs1 = new(bits1), vs2 = new(bits2);

    //if (vs1.VectorLength != vs2.VectorLength || vs1._spanRemainder.Length != vs2._spanRemainder.Length)
    //    throw new ArgumentException();

    for (int i = 0; i < vs1.VectorLength; i++)
    {
        vs1.WriteVector(i, vs1.ReadVector(i) ^ vs2.ReadVector(i));
    }
    for (int i = 0; i < vs1._spanRemainder.Length; i++)
    {
        vs1._spanRemainder[i] = (byte)(vs1._spanRemainder[i] ^ vs2._spanRemainder[i]);
    }
}

Apparently not. It's worse:

; Assembly listing for method ConsoleApp1.SafeVec:Xor(System.Span`1[ubyte],System.Span`1[ubyte]) (Tier1)
; Emitting BLENDED_CODE for X64 with AVX512 - Windows
; Tier1 code
; optimized code
; optimized using Synthesized PGO
; rsp based frame
; fully interruptible
; with Synthesized PGO: fgCalledCount is 70
; 4 inlinees with PGO data; 8 single block inlinees; 2 inlinees without PGO data

G_M000_IG01:                ;; offset=0x0000
       push     rdi
       push     rsi
       push     rbp
       push     rbx
       sub      rsp, 40

G_M000_IG02:                ;; offset=0x0008
       mov      eax, dword ptr [rcx+0x08]
       mov      r8d, dword ptr [rdx+0x08]
       cmp      eax, r8d
       jne      G_M000_IG12
       mov      rcx, bword ptr [rcx]
       mov      r10d, eax
       shr      r10d, 4
       mov      r9d, r10d
       shl      r9d, 4
       cmp      r9d, eax
       ja       G_M000_IG13
       mov      r11d, r9d
       add      r11, rcx
       sub      eax, r9d
       mov      rdx, bword ptr [rdx]
       mov      r9d, r8d
       shr      r9d, 4
       mov      ebx, r9d
       shl      ebx, 4
       cmp      ebx, r8d
       ja       G_M000_IG13
       mov      esi, ebx
       add      rsi, rdx
       sub      r8d, ebx
       xor      ebx, ebx
       test     r10d, r10d
       jle      SHORT G_M000_IG06

G_M000_IG03:                ;; offset=0x0063
       cmp      r10d, r9d
       jg       SHORT G_M000_IG10

G_M000_IG04:                ;; offset=0x0068
       xor      r9d, r9d
       align    [5 bytes for IG05]

G_M000_IG05:                ;; offset=0x0070
       vmovups  xmm0, xmmword ptr [rcx+r9]
       vpxor    xmm0, xmm0, xmmword ptr [rdx+r9]
       vmovups  xmmword ptr [rcx+r9], xmm0
       add      r9, 16
       dec      r10d
       jne      SHORT G_M000_IG05

G_M000_IG06:                ;; offset=0x008B
       xor      ecx, ecx
       test     eax, eax
       jle      SHORT G_M000_IG09

G_M000_IG07:                ;; offset=0x0091
       cmp      eax, r8d
       jg       SHORT G_M000_IG11
       align    [10 bytes for IG08]

G_M000_IG08:                ;; offset=0x00A0
       mov      r8d, ecx
       movzx    rdx, byte  ptr [rsi+r8]
       xor      byte  ptr [r11+r8], dl
       inc      ecx
       cmp      ecx, eax
       jl       SHORT G_M000_IG08

G_M000_IG09:                ;; offset=0x00B2
       add      rsp, 40
       pop      rbx
       pop      rbp
       pop      rsi
       pop      rdi
       ret

G_M000_IG10:                ;; offset=0x00BB
       mov      edi, ebx
       shl      rdi, 4
       lea      rbp, bword ptr [rcx+rdi]
       vmovups  xmm0, xmmword ptr [rbp]
       cmp      ebx, r9d
       jae      SHORT G_M000_IG14
       vpxor    xmm0, xmm0, xmmword ptr [rdx+rdi]
       vmovups  xmmword ptr [rbp], xmm0
       inc      ebx
       cmp      ebx, r10d
       jl       SHORT G_M000_IG10
       jmp      SHORT G_M000_IG06

G_M000_IG11:                ;; offset=0x00E2
       mov      edx, ecx
       mov      r10d, ecx
       movzx    r10, byte  ptr [r11+r10]
       cmp      ecx, r8d
       jae      SHORT G_M000_IG14
       mov      r9d, ecx
       movzx    r9, byte  ptr [rsi+r9]
       xor      r10d, r9d
       mov      byte  ptr [r11+rdx], r10b
       inc      ecx
       cmp      ecx, eax
       jl       SHORT G_M000_IG11
       jmp      SHORT G_M000_IG09

G_M000_IG12:                ;; offset=0x0108
       mov      rcx, 0x7FFB2735C080
       call     CORINFO_HELP_NEWSFAST
       mov      rbx, rax
       mov      rcx, rbx
       call     [System.ArgumentException:.ctor():this]
       mov      rcx, rbx
       call     CORINFO_HELP_THROW
       int3

G_M000_IG13:                ;; offset=0x012C
       call     [System.ThrowHelper:ThrowArgumentOutOfRangeException()]
       int3

G_M000_IG14:                ;; offset=0x0133
       call     CORINFO_HELP_RNGCHKFAIL
       int3

; Total bytes of code 313

[tannergooding]

_spanVector = MemoryMarshal.Cast<T, Vector128>(span);

This is dangerous for many reasons and is generally not something you'd want to do.

I'd likely argue it is more unsafe than the standard Vector128.LoadUnsafe(ref address, nuint index) and Vector128.StoreUnsafe(ref address, nuint index) pattern all while generating more complex and likely less optimal code.

The general point of this issue is that we have some simple and clear patterns for unsafely loading/storing the data today. But ideally users should just be able to pass in a Span<T> to get a Vector128<T>, slice it, and iterate along it until the end; that way it is safe and performant.