FastAvxFill.cs

/// <summary>
/// Fills the contents of this span with the given value.
/// </summary>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static unsafe void Fill<T>(this Span<T> @this, T value)
{
    if (@this.IsEmpty)
    {
        return;
    }

    // The branches based on this are stil properly folded, it is just more concise than writing 'Unsafe.SizeOf<T>()' everytime
    int size = Unsafe.SizeOf<T>();
    int len = @this.Length;

    Debug.Assert(size > 0 && len > 0);

    // This branch is either selected or elided by the JIT at JIT time
    if (RuntimeHelpers.IsReferenceOrContainsReferences<T>())
    {
        SoftwareFallback(@this, value);
        return;
    }

    int fullSize = len * size;

    // If the size is 1, which is a JIT time constant, the method just becomes initblk, which appears to perform better than the manual AVX or SSE pathways
    if (size == 1)
    {
        Unsafe.InitBlockUnaligned(ref Unsafe.As<T, byte>(ref MemoryMarshal.GetReference(@this)), Unsafe.As<T, byte>(ref value),
            (uint)len); // won't overflow, unless someone has reflected into the '_length' field and changed it to a negative value. Not particularly concerned with that scenario
        return;
    }

    // This and the SSE pathway cannot be fully elided by the JIT because they are dependent on the size of the data being filled,
    // but they will be elided if they are not supported or the size is not supported
    if (Avx.IsSupported
        && (fullSize >= 32) && ((size & (size - 1)) == 0 /* Is pow of 2 */) && (size <= 32))
    {
        Vector256<byte> vector;

        // Create the vector by filling it with {n} Ts, where {n} is 32 / sizeof(T)
        // This table is elided to a single branch at JIT time
        switch (size)
        {
            case 1:
                vector = Vector256.Create(Unsafe.As<T, byte>(ref value));
                break;
            case 2:
                vector = Vector256.Create(Unsafe.As<T, ushort>(ref value)).AsByte();
                break;
            case 4:
                vector = Vector256.Create(Unsafe.As<T, uint>(ref value)).AsByte();
                break;
            case 8:
                vector = Vector256.Create(Unsafe.As<T, ulong>(ref value)).AsByte();
                break;
            case 16:
                Vector128<byte> tmp = Unsafe.As<T, Vector128<byte>>(ref value);
                vector = Vector256.Create(tmp, tmp);
                break;
            case 32:
                vector = Unsafe.As<T, Vector256<byte>>(ref value);
                break;
            default:
                return; // unreachable, necessary
        }

        // We verified the span was not empty at the start, so the check from GetPinnableReference is not necessary
        // As this T is not constrained to be 'unmanaged' (even though we have confirmed it is), we must cast to a 'ref byte' first
        fixed (byte* p = &Unsafe.As<T, byte>(ref MemoryMarshal.GetReference(@this)))
        {
            byte* pAliasedVector = p; // a copy is required to find the alignment difference

            Avx.Store(pAliasedVector, vector); // initial, unaligned store

            if (fullSize == 32)
            {
                return; // previous store was all needed
            }

            pAliasedVector = (byte*)RoundUp(pAliasedVector, 32); // round up pointer to next 32 bytes to allow aligned stores
            Debug.Assert((ulong)pAliasedVector % 32 == 0);

            // This block rotates the vector to accomodate for the fact it has been offset by rounding up the pointer for alignment
            byte* pool = stackalloc byte[32 * 2];
            Avx.Store(pool, vector);
            Avx.Store(pool + 32, vector);

            var diff = (int)(pAliasedVector - p);
            fullSize -= diff;
            Vector256<byte> cpy = vector; // we make a copy here that we use for the final, unaligned store
            vector = Avx.LoadVector256(pool + diff);

            for (var i = 0; i < (fullSize & ~31U); i += 32)
            {
                Avx.Store(pAliasedVector + i, vector); // These stores are aligned (the assertion above confirms that), but we use the non aligned 
                                                       // instruction anyway for VEX encoding
            }

            if (fullSize % 32 == 0)
            {
                return; // no need for final unaligned store
            }

            Avx.Store((pAliasedVector + fullSize) - 32, cpy); // A final unaligned store, for up to the last 31 bytes, using the original vector
        }
    }
    else if (Sse2.IsSupported
    && (fullSize >= 16) && ((size & (size - 1)) == 0 /* Is pow of 2 */) && (size <= 16))
    {

        Vector128<byte> vector;

        // Create the vector by filling it with {n} Ts, where {n} is 16 / sizeof(T)
        // This table is elided to a single branch at JIT time
        switch (size)
        {
            case 1:
                vector = Vector128.Create(Unsafe.As<T, byte>(ref value));
                break;
            case 2:
                vector = Vector128.Create(Unsafe.As<T, ushort>(ref value)).AsByte();
                break;
            case 4:
                vector = Vector128.Create(Unsafe.As<T, uint>(ref value)).AsByte();
                break;
            case 8:
                vector = Vector128.Create(Unsafe.As<T, ulong>(ref value)).AsByte();
                break;
            case 16:
                vector = Unsafe.As<T, Vector128<byte>>(ref value);
                break;
            default:
                return; // unreachable, necessary
        }

        // We verified the span was not empty at the start, so the check from GetPinnableReference is not necessary
        // As this T is not constrained to be 'unmanaged' (even though we have confirmed it is), we must cast to a 'ref byte' first
        fixed (byte* p = &Unsafe.As<T, byte>(ref MemoryMarshal.GetReference(@this)))
        {
            byte* pAliasedVector = p; // a copy is required to find the alignment difference

            Sse2.Store(pAliasedVector, vector); // initial, unaligned store

            if (fullSize == 16)
            {
                return; // previous store was all needed
            }

            pAliasedVector = (byte*)RoundUp(pAliasedVector, 16); // round up pointer to next 16 bytes to allow aligned stores
            Debug.Assert((ulong)pAliasedVector % 16 == 0);

            // This block rotates the vector to accomodate for the fact it has been offset by rounding up the pointer for alignment
            byte* pool = stackalloc byte[16 * 2];
            Sse2.Store(pool, vector);
            Sse2.Store(pool + 16, vector);

            var diff = (int)(pAliasedVector - p);
            fullSize -= diff;
            Vector128<byte> cpy = vector; // we make a copy here that we use for the final, unaligned store
            vector = Sse2.LoadVector128(pool + diff);

            for (var i = 0; i < (fullSize & ~15U); i += 16)
            {
                Sse2.Store(pAliasedVector + i, vector); // These stores are aligned (the assertion above confirms that), but we use the non aligned 
                                                        // instruction anyway for VEX encoding
            }

            if (fullSize % 16 == 0)
            {
                return; // no need for final unaligned store
            }

            Sse2.Store((pAliasedVector + fullSize) - 16, cpy); // A final unaligned store, for up to the last 31 bytes, using the original vector
        }
    }
    else
    {
        SoftwareFallback(@this, value);
    }


    static void SoftwareFallback(Span<T> span, T value)
    {
        for (var i = 0; i < span.Length; i++)
        {
            span[i] = value;
        }
    }

    static void* RoundUp(void* p, uint alignment) => (void*)(((ulong)p + (alignment - 1UL)) & ~(alignment - 1UL));
}