Adding TLoader method for Vector512 for EqualAndAscii

Author: khushal1996

src/libraries/System.Private.CoreLib/src/System/Text/Ascii.Equality.cs

@@ -61,42 +61,34 @@ private static bool Equals<TLeft, TRight, TLoader>(ref TLeft left, ref TRight ri
                     }
                 }
             }
-            else if (Vector512.IsHardwareAccelerated && length >= (uint)Vector512<TRight>.Count)
+            else if (Vector512.IsHardwareAccelerated && length >= (uint)Vector512<TLeft>.Count)
             {
                 ref TLeft currentLeftSearchSpace = ref left;
-                ref TLeft oneVectorAwayFromLeftEnd = ref Unsafe.Add(ref currentLeftSearchSpace, length - TLoader.Count512);
                 ref TRight currentRightSearchSpace = ref right;
-                ref TRight oneVectorAwayFromRightEnd = ref Unsafe.Add(ref currentRightSearchSpace, length - (uint)Vector512<TRight>.Count);
-
-                Vector512<TRight> leftValues;
-                Vector512<TRight> rightValues;
+                // Add Vector512<TLeft>.Count because TLeft == TRight
+                // Or we are in the Widen case where we iterate 2 * TRight.Count which is the same as TLeft.Count
+                Debug.Assert(Vector512<TLeft>.Count == Vector512<TRight>.Count
+                    || (typeof(TLoader) == typeof(WideningLoader) && Vector512<TLeft>.Count == Vector512<TRight>.Count * 2));
+                ref TRight oneVectorAwayFromRightEnd = ref Unsafe.Add(ref currentRightSearchSpace, length - (uint)Vector512<TLeft>.Count);
 
                 // Loop until either we've finished all elements or there's less than a vector's-worth remaining.
                 do
                 {
-                    leftValues = TLoader.Load512(ref currentLeftSearchSpace);
-                    rightValues = Vector512.LoadUnsafe(ref currentRightSearchSpace);
-
-                    if (leftValues != rightValues || !AllCharsInVectorAreAscii(leftValues | rightValues))
+                    if (!TLoader.EqualAndAscii512(ref currentLeftSearchSpace, ref currentRightSearchSpace))
                     {
                         return false;
                     }
 
-                    currentRightSearchSpace = ref Unsafe.Add(ref currentRightSearchSpace, Vector512<TRight>.Count);
-                    currentLeftSearchSpace = ref Unsafe.Add(ref currentLeftSearchSpace, TLoader.Count512);
+                    currentRightSearchSpace = ref Unsafe.Add(ref currentRightSearchSpace, Vector512<TLeft>.Count);
+                    currentLeftSearchSpace = ref Unsafe.Add(ref currentLeftSearchSpace, Vector512<TLeft>.Count);
                 }
                 while (!Unsafe.IsAddressGreaterThan(ref currentRightSearchSpace, ref oneVectorAwayFromRightEnd));
 
                 // If any elements remain, process the last vector in the search space.
-                if (length % (uint)Vector512<TRight>.Count != 0)
+                if (length % (uint)Vector512<TLeft>.Count != 0)
                 {
-                    leftValues = TLoader.Load512(ref oneVectorAwayFromLeftEnd);
-                    rightValues = Vector512.LoadUnsafe(ref oneVectorAwayFromRightEnd);
-
-                    if (leftValues != rightValues || !AllCharsInVectorAreAscii(leftValues | rightValues))
-                    {
-                        return false;
-                    }
+                    ref TLeft oneVectorAwayFromLeftEnd = ref Unsafe.Add(ref left, length - (uint)Vector512<TLeft>.Count);
+                    return TLoader.EqualAndAscii512(ref oneVectorAwayFromLeftEnd, ref oneVectorAwayFromRightEnd);
                 }
             }
             else if (Avx.IsSupported && length >= (uint)Vector256<TLeft>.Count)
@@ -112,7 +104,7 @@ private static bool Equals<TLeft, TRight, TLoader>(ref TLeft left, ref TRight ri
                 // Loop until either we've finished all elements or there's less than a vector's-worth remaining.
                 do
                 {
-                    if (!TLoader.EqualAndAscii(ref currentLeftSearchSpace, ref currentRightSearchSpace))
+                    if (!TLoader.EqualAndAscii256(ref currentLeftSearchSpace, ref currentRightSearchSpace))
                     {
                         return false;
                     }
@@ -126,7 +118,7 @@ private static bool Equals<TLeft, TRight, TLoader>(ref TLeft left, ref TRight ri
                 if (length % (uint)Vector256<TLeft>.Count != 0)
                 {
                     ref TLeft oneVectorAwayFromLeftEnd = ref Unsafe.Add(ref left, length - (uint)Vector256<TLeft>.Count);
-                    return TLoader.EqualAndAscii(ref oneVectorAwayFromLeftEnd, ref oneVectorAwayFromRightEnd);
+                    return TLoader.EqualAndAscii256(ref oneVectorAwayFromLeftEnd, ref oneVectorAwayFromRightEnd);
                 }
             }
             else
@@ -255,7 +247,6 @@ private static bool EqualsIgnoreCase<TLeft, TRight, TLoader>(ref TLeft left, ref
                 {
                     leftValues = TLoader.Load512(ref currentLeftSearchSpace);
                     rightValues = Vector512.LoadUnsafe(ref currentRightSearchSpace);
-
                     if (!AllCharsInVectorAreAscii(leftValues | rightValues))
                     {
                         return false;
@@ -467,7 +458,8 @@ private interface ILoader<TLeft, TRight>
             static abstract Vector128<TRight> Load128(ref TLeft ptr);
             static abstract Vector256<TRight> Load256(ref TLeft ptr);
             static abstract Vector512<TRight> Load512(ref TLeft ptr);
-            static abstract bool EqualAndAscii(ref TLeft left, ref TRight right);
+            static abstract bool EqualAndAscii256(ref TLeft left, ref TRight right);
+            static abstract bool EqualAndAscii512(ref TLeft left, ref TRight right);
         }
 
         private readonly struct PlainLoader<T> : ILoader<T, T> where T : unmanaged, INumberBase<T>
@@ -477,10 +469,11 @@ private interface ILoader<TLeft, TRight>
             public static nuint Count512 => (uint)Vector512<T>.Count;
             public static Vector128<T> Load128(ref T ptr) => Vector128.LoadUnsafe(ref ptr);
             public static Vector256<T> Load256(ref T ptr) => Vector256.LoadUnsafe(ref ptr);
+            public static Vector512<T> Load512(ref T ptr) => Vector512.LoadUnsafe(ref ptr);
 
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
             [CompExactlyDependsOn(typeof(Avx))]
-            public static bool EqualAndAscii(ref T left, ref T right)
+            public static bool EqualAndAscii256(ref T left, ref T right)
             {
                 Vector256<T> leftValues = Vector256.LoadUnsafe(ref left);
                 Vector256<T> rightValues = Vector256.LoadUnsafe(ref right);
@@ -493,7 +486,19 @@ public static bool EqualAndAscii(ref T left, ref T right)
                 return true;
             }
 
-            public static Vector512<T> Load512(ref T ptr) => Vector512.LoadUnsafe(ref ptr);
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public static bool EqualAndAscii512(ref T left, ref T right)
+            {
+                Vector512<T> leftValues = Vector512.LoadUnsafe(ref left);
+                Vector512<T> rightValues = Vector512.LoadUnsafe(ref right);
+
+                if (leftValues != rightValues || !AllCharsInVectorAreAscii(leftValues))
+                {
+                    return false;
+                }
+
+                return true;
+            }
         }
 
         private readonly struct WideningLoader : ILoader<byte, ushort>
@@ -528,9 +533,16 @@ public static Vector256<ushort> Load256(ref byte ptr)
                 return Vector256.Create(lower, upper);
             }
 
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public static Vector512<ushort> Load512(ref byte ptr)
+            {
+                (Vector256<ushort> lower, Vector256<ushort> upper) = Vector256.Widen(Vector256.LoadUnsafe(ref ptr));
+                return Vector512.Create(lower, upper);
+            }
+
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
             [CompExactlyDependsOn(typeof(Avx))]
-            public static bool EqualAndAscii(ref byte utf8, ref ushort utf16)
+            public static bool EqualAndAscii256(ref byte utf8, ref ushort utf16)
             {
                 // We widen the utf8 param so we can compare it to utf16, this doubles how much of the utf16 vector we search
                 Debug.Assert(Vector256<byte>.Count == Vector256<ushort>.Count * 2);
@@ -554,10 +566,29 @@ public static bool EqualAndAscii(ref byte utf8, ref ushort utf16)
                 return true;
             }
 
-            public static Vector512<ushort> Load512(ref byte ptr)
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public static bool EqualAndAscii512(ref byte utf8, ref ushort utf16)
             {
-                (Vector256<ushort> lower, Vector256<ushort> upper) = Vector256.Widen(Vector256.LoadUnsafe(ref ptr));
-                return Vector512.Create(lower, upper);
+                // We widen the utf8 param so we can compare it to utf16, this doubles how much of the utf16 vector we search
+                Debug.Assert(Vector512<byte>.Count == Vector512<ushort>.Count * 2);
+
+                Vector512<byte> leftNotWidened = Vector512.LoadUnsafe(ref utf8);
+                if (!AllCharsInVectorAreAscii(leftNotWidened))
+                {
+                    return false;
+                }
+
+                (Vector512<ushort> leftLower, Vector512<ushort> leftUpper) = Vector512.Widen(leftNotWidened);
+                Vector512<ushort> right = Vector512.LoadUnsafe(ref utf16);
+                Vector512<ushort> rightNext = Vector512.LoadUnsafe(ref utf16, (uint)Vector512<ushort>.Count);
+
+                // A branchless version of "leftLower != right || leftUpper != rightNext"
+                if (((leftLower ^ right) | (leftUpper ^ rightNext)) != Vector512<ushort>.Zero)
+                {
+                    return false;
+                }
+
+                return true;
             }
         }
     }