Optimize Span<T>.Fill implementation

src/libraries/System.Memory/tests/Span/Fill.cs

@@ -1,6 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Generic;
+using System.Linq;
+using System.Reflection;
 using System.Runtime.InteropServices;
 using Xunit;
 using static System.TestHelpers;
@@ -147,5 +150,53 @@ public static unsafe void FillNativeBytes()
                 Marshal.FreeHGlobal(new IntPtr(ptr));
             }
         }
+
+        public static IEnumerable<object[]> FillData
+        {
+            get
+            {
+                yield return new object[] { typeof(sbyte), (sbyte)0x20 };
+                yield return new object[] { typeof(byte), (byte)0x20 };
+                yield return new object[] { typeof(bool), true };
+                yield return new object[] { typeof(short), (short)0x1234 };
+                yield return new object[] { typeof(ushort), (ushort)0x1234 };
+                yield return new object[] { typeof(char), 'x' };
+                yield return new object[] { typeof(int), (int)0x12345678 };
+                yield return new object[] { typeof(uint), (uint)0x12345678 };
+                yield return new object[] { typeof(long), (long)0x0123456789abcdef };
+                yield return new object[] { typeof(ulong), (ulong)0x0123456789abcdef };
+                yield return new object[] { typeof(nint), unchecked((nint)0x0123456789abcdef) };
+                yield return new object[] { typeof(nuint), unchecked((nuint)0x0123456789abcdef) };
+                yield return new object[] { typeof(float), 1.0f };
+                yield return new object[] { typeof(double), 1.0d };
+                yield return new object[] { typeof(string), "Hello world!" }; // shouldn't go down SIMD path
+                yield return new object[] { typeof(decimal), 1.0m }; // shouldn't go down SIMD path
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(FillData))]
+        public static void FillWithRecognizedType(Type expectedType, object value)
+        {
+            Assert.IsType(expectedType, value);
+            typeof(SpanTests).GetMethod(nameof(FillWithRecognizedType_RunTest), BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static)
+                .MakeGenericMethod(expectedType)
+                .Invoke(null, BindingFlags.DoNotWrapExceptions, null, new object[] { value }, null);
+        }
+
+        private static void FillWithRecognizedType_RunTest<T>(T value)
+        {
+            T[] arr = new T[128];
+
+            // Run tests for lengths := 0 to 64, ensuring we don't overrun our buffer
+
+            for (int i = 0; i <= 64; i++)
+            {
+                arr.AsSpan(0, i).Fill(value);
+                Assert.Equal(Enumerable.Repeat(value, i), arr.Take(i)); // first i entries should've been populated with 'value'
+                Assert.Equal(Enumerable.Repeat(default(T), arr.Length - i), arr.Skip(i)); // remaining entries should contain default(T)
+                Array.Clear(arr, 0, arr.Length);
+            }
+        }
     }
 }

src/libraries/System.Private.CoreLib/src/System/Span.cs

@@ -5,7 +5,6 @@
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Runtime.Versioning;
-using System.Text;
 using EditorBrowsableAttribute = System.ComponentModel.EditorBrowsableAttribute;
 using EditorBrowsableState = System.ComponentModel.EditorBrowsableState;
 using Internal.Runtime.CompilerServices;
@@ -280,53 +279,21 @@ public unsafe void Clear()
         /// <summary>
         /// Fills the contents of this span with the given value.
         /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void Fill(T value)
         {
             if (Unsafe.SizeOf<T>() == 1)
             {
-                uint length = (uint)_length;
-                if (length == 0)
-                    return;
-
-                T tmp = value; // Avoid taking address of the "value" argument. It would regress performance of the loop below.
-                Unsafe.InitBlockUnaligned(ref Unsafe.As<T, byte>(ref _pointer.Value), Unsafe.As<T, byte>(ref tmp), length);
+                // Special-case single-byte types like byte / sbyte / bool.
+                // The runtime eventually calls memset, which can efficiently support large buffers.
+                // We don't need to check IsReferenceOrContainsReferences because no references
+                // can ever be stored in types this small.
+                Unsafe.InitBlockUnaligned(ref Unsafe.As<T, byte>(ref _pointer.Value), Unsafe.As<T, byte>(ref value), (uint)_length);
             }
             else
             {
-                // Do all math as nuint to avoid unnecessary 64->32->64 bit integer truncations
-                nuint length = (uint)_length;
-                if (length == 0)
-                    return;
-
-                ref T r = ref _pointer.Value;
-
-                // TODO: Create block fill for value types of power of two sizes e.g. 2,4,8,16
-
-                nuint elementSize = (uint)Unsafe.SizeOf<T>();
-                nuint i = 0;
-                for (; i < (length & ~(nuint)7); i += 8)
-                {
-                    Unsafe.AddByteOffset<T>(ref r, (i + 0) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 1) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 2) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 3) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 4) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 5) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 6) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 7) * elementSize) = value;
-                }
-                if (i < (length & ~(nuint)3))
-                {
-                    Unsafe.AddByteOffset<T>(ref r, (i + 0) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 1) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 2) * elementSize) = value;
-                    Unsafe.AddByteOffset<T>(ref r, (i + 3) * elementSize) = value;
-                    i += 4;
-                }
-                for (; i < length; i++)
-                {
-                    Unsafe.AddByteOffset<T>(ref r, i * elementSize) = value;
-                }
+                // Call our optimized workhorse method for all other types.
+                SpanHelpers.Fill(ref _pointer.Value, (uint)_length, value);
             }
         }
 

src/libraries/System.Private.CoreLib/src/System/SpanHelpers.T.cs

@@ -2,13 +2,160 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Diagnostics;
-
+using System.Numerics;
+using System.Runtime.CompilerServices;
 using Internal.Runtime.CompilerServices;
 
 namespace System
 {
     internal static partial class SpanHelpers // .T
     {
+        public static void Fill<T>(ref T refData, uint numElements, T value)
+        {
+            // n.b. If Fill is ever adapted to work for lengths beyond uint.MaxValue, double-check
+            // where it's used in the arithmetic operations below to ensure no integer overflow is possible.
+
+            if (numElements == 0)
+            {
+                return; // nothing to do
+            }
+
+            if (typeof(T) != typeof(float) && typeof(T) != typeof(double) && !RuntimeHelpers.IsBitwiseEquatable<T>())
+            {
+                goto CannotVectorize; // it might not be valid to SIMD-spray this value into the backing span
+            }
+
+            if (Vector.IsHardwareAccelerated)
+            {
+                if (numElements < (uint)(Vector<byte>.Count / Unsafe.SizeOf<T>()))
+                {
+                    goto CannotVectorize; // not enough data for even a single run through the vectorized loop
+                }
+
+                Vector<byte> vector;
+
+                if (typeof(T) == typeof(float))
+                {
+                    vector = (Vector<byte>)(new Vector<float>((float)(object)value!));
+                }
+                else if (typeof(T) == typeof(double))
+                {
+                    vector = (Vector<byte>)(new Vector<double>((double)(object)value!));
+                }
+                else
+                {
+                    T tmp = value; // Avoid taking address of the "value" argument. It would regress performance of the loops below.
+                    if (Unsafe.SizeOf<T>() == 1)
+                    {
+                        vector = new Vector<byte>(Unsafe.As<T, byte>(ref tmp));
+                    }
+                    else if (Unsafe.SizeOf<T>() == 2)
+                    {
+                        vector = (Vector<byte>)(new Vector<ushort>(Unsafe.As<T, ushort>(ref tmp)));
+                    }
+                    else if (Unsafe.SizeOf<T>() == 4)
+                    {
+                        vector = (Vector<byte>)(new Vector<uint>(Unsafe.As<T, uint>(ref tmp)));
+                    }
+                    else if (Unsafe.SizeOf<T>() == 8)
+                    {
+                        vector = (Vector<byte>)(new Vector<ulong>(Unsafe.As<T, ulong>(ref tmp)));
+                    }
+                    else
+                    {
+                        Debug.Fail("This should never happen. Did we miss special-casing a bitwise equatable type?");
+                        goto CannotVectorize;
+                    }
+                }
+
+                nuint totalByteLength = numElements * (nuint)Unsafe.SizeOf<T>(); // get this calculation ready ahead of time
+                nuint stopLoopAtOffset = totalByteLength & (nuint)(-Vector<byte>.Count * 2);
+                nuint offset = 0;
+
+                // Loop, writing 2 vectors at a time.
+
+                if (numElements >= 2 * (uint)(Vector<byte>.Count / Unsafe.SizeOf<T>()))
+                {
+                    do
+                    {
+                        Unsafe.WriteUnaligned(ref Unsafe.AddByteOffset(ref Unsafe.As<T, byte>(ref refData), offset), vector);
+                        Unsafe.WriteUnaligned(ref Unsafe.AddByteOffset(ref Unsafe.As<T, byte>(ref refData), offset + (uint)Vector<byte>.Count), vector);
+                        offset += 2 * (uint)Vector<byte>.Count;
+                    } while (offset < stopLoopAtOffset);
+                }
+
+                // There are [ 0, 2 * sizeof(Vector) ) bytes remaining.
+                // If there are >= sizeof(Vector) bytes remaining, write one vector now.
+
+                if ((totalByteLength & (nuint)Vector<byte>.Count) != 0)
+                {
+                    Unsafe.WriteUnaligned(ref Unsafe.AddByteOffset(ref Unsafe.As<T, byte>(ref refData), offset), vector);
+                }
+
+                // If there's any remaining space that won't fill a full vector, write a vector at
+                // the very end of the destination buffer. This will result in overwriting a previous
+                // entry, but that's ok since we're splatting the same value for all entries, so this
+                // won't result in data corruption.
+
+                if ((totalByteLength & (nuint)(Vector<byte>.Count - 1)) != 0)
+                {
+                    Unsafe.WriteUnaligned(ref Unsafe.AddByteOffset(ref Unsafe.As<T, byte>(ref refData), totalByteLength - (uint)Vector<byte>.Count), vector);
+                }
+
+                // And we're done!
+
+                return;
+            }
+
+        CannotVectorize:
+
+            {
+                nuint i = 0;
+                nuint stopLoopAtOffset = numElements & ~(uint)7;
+
+                // Write 8 elements at a time
+
+                for (; i < stopLoopAtOffset; i += 8)
+                {
+                    Unsafe.Add(ref refData, (nint)i + 0) = value;
+                    Unsafe.Add(ref refData, (nint)i + 1) = value;
+                    Unsafe.Add(ref refData, (nint)i + 2) = value;
+                    Unsafe.Add(ref refData, (nint)i + 3) = value;
+                    Unsafe.Add(ref refData, (nint)i + 4) = value;
+                    Unsafe.Add(ref refData, (nint)i + 5) = value;
+                    Unsafe.Add(ref refData, (nint)i + 6) = value;
+                    Unsafe.Add(ref refData, (nint)i + 7) = value;
+                }
+
+                // Write next 4 elements if needed
+
+                if ((numElements & 4) != 0)
+                {
+                    Unsafe.Add(ref refData, (nint)i + 0) = value;
+                    Unsafe.Add(ref refData, (nint)i + 1) = value;
+                    Unsafe.Add(ref refData, (nint)i + 2) = value;
+                    Unsafe.Add(ref refData, (nint)i + 3) = value;
+                    i += 4;
+                }
+
+                // Write next 2 elements if needed
+
+                if ((numElements & 2) != 0)
+                {
+                    Unsafe.Add(ref refData, (nint)i + 0) = value;
+                    Unsafe.Add(ref refData, (nint)i + 1) = value;
+                    i += 2;
+                }
+
+                // Write final element if needed
+
+                if ((numElements & 1) != 0)
+                {
+                    Unsafe.Add(ref refData, (nint)i) = value;
+                }
+            }
+        }
+
         public static int IndexOf<T>(ref T searchSpace, int searchSpaceLength, ref T value, int valueLength) where T : IEquatable<T>
         {
             Debug.Assert(searchSpaceLength >= 0);