[release/11.0] Preserve mask element granularity in SIMD transformations

src/coreclr/jit/gentree.cpp

@@ -23182,7 +23182,7 @@ bool GenTree::isEmbeddedMaskingCompatible(Compiler*  comp,
                 {
                     assert(broadcastOpIndex != nullptr);
 
-                    // If the contained broadcast is 4 bytes, we can change it to 8 bytes
+                    // A 4-byte broadcast may be widened if its scalar is constant (checked below).
                     supportsMaskBaseSize2Or4 = true;
                     *broadcastOpIndex        = 2;
                 }
@@ -23204,7 +23204,7 @@ bool GenTree::isEmbeddedMaskingCompatible(Compiler*  comp,
                 {
                     assert(broadcastOpIndex != nullptr);
 
-                    // If the contained broadcast is 4 bytes, we can change it to 8 bytes
+                    // A 4-byte broadcast may be widened if its scalar is constant (checked below).
                     supportsMaskBaseSize2Or4 = true;
                     *broadcastOpIndex        = 3;
                 }
@@ -23250,6 +23250,18 @@ bool GenTree::isEmbeddedMaskingCompatible(Compiler*  comp,
 
         if (supportsMaskBaseSize2Or4)
         {
+            if ((broadcastOpIndex != nullptr) && (*broadcastOpIndex != 0))
+            {
+                const GenTreeHWIntrinsic* broadcastNode = node->Op(*broadcastOpIndex)->AsHWIntrinsic();
+
+                // Only a constant can be duplicated into a wider broadcast without changing
+                // the original memory access or the bits broadcast into each element.
+                if (broadcastNode->OperIsMemoryLoad() || !broadcastNode->Op(1)->OperIsConst())
+                {
+                    return false;
+                }
+            }
+
             if (tgtMaskBaseSize == 2)
             {
                 if (varTypeIsFloating(simdBaseType))
@@ -35542,19 +35554,21 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
 
         assert(opCount == (size_t)HWIntrinsicInfo::lookupNumArgs(maskVariant));
 
-        // Check all operands are valid
-        bool canFold = true;
+        size_t firstVectorOperand = 1;
         if (ni == NI_Sve_ConditionalSelect)
         {
             assert(varTypeIsMask(op1));
-            canFold = (op2->OperIsConvertMaskToVector() && op3->OperIsConvertMaskToVector());
+            firstVectorOperand = 2;
         }
-        else
+
+        // Predicate bits are spaced according to element size. Reinterpreting the
+        // expanded vector does not reinterpret the predicate at the new granularity.
+        bool canFold = true;
+        for (size_t i = firstVectorOperand; (i <= opCount) && canFold; i++)
         {
-            for (size_t i = 1; i <= opCount && canFold; i++)
-            {
-                canFold &= tree->Op(i)->OperIsConvertMaskToVector();
-            }
+            GenTree* operand = tree->Op(i);
+            canFold          = operand->OperIsConvertMaskToVector() &&
+                      (genTypeSize(operand->AsHWIntrinsic()->GetSimdBaseType()) == genTypeSize(simdBaseType));
         }
 
         if (canFold)

src/coreclr/jit/lowerxarch.cpp

@@ -2618,7 +2618,7 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node)
 
             // If either of the value operands is const zero and the mask is either all
             // zeros or all ones per-element, we can optimize down to AND or AND_NOT.
-            if (op3->IsVectorPerElementMask(m_compiler, TYP_BYTE, simdSize) &&
+            if (op3->IsVectorPerElementMask(m_compiler, simdBaseType, simdSize) &&
                 (op1->IsVectorZero() || op2->IsVectorZero()))
             {
                 var_types simdType = node->TypeGet();
@@ -3014,7 +3014,7 @@ GenTree* Lowering::LowerHWIntrinsicCmpOp(GenTreeHWIntrinsic* node, genTreeOps cm
         // This works since the upper bits are implicitly zero and so by inverting matches also become
         // zero, which in turn means that `AllBitsSet` will become `Zero` and other cases become non-zero
 
-        if (varTypeIsMask(op1Msk) && op2->IsCnsVec())
+        if (varTypeIsMask(op1Msk) && (op2->IsVectorZero() || op2->IsVectorAllBitsSet()))
         {
             // We want to specially handle the common cases of `mask op Zero` and `mask op AllBitsSet`
             //
@@ -3180,31 +3180,37 @@ GenTree* Lowering::LowerHWIntrinsicCmpOp(GenTreeHWIntrinsic* node, genTreeOps cm
                                 (nestedIntrinId == NI_AVX2_BroadcastScalarToVector256) ||
                                 (nestedIntrinId == NI_AVX512_BroadcastScalarToVector512))
                             {
-                                // We need to rewrite the embedded broadcast back to a regular constant
-                                // so that the subsequent containment check for ptestm can determine
-                                // if the embedded broadcast is still relevant
+                                // Restore constant broadcasts so ptestm can choose its own broadcast
+                                // granularity. Runtime broadcasts must retain their original load size.
 
                                 GenTree* broadcastOp = nestedIntrin->Op(1);
+                                GenTree* scalarOp    = broadcastOp;
 
                                 if (broadcastOp->OperIsHWIntrinsic(NI_Vector_CreateScalarUnsafe) &&
                                     broadcastOp->TypeIs(TYP_SIMD16))
                                 {
-                                    BlockRange().Remove(broadcastOp);
-                                    broadcastOp = broadcastOp->AsHWIntrinsic()->Op(1);
+                                    scalarOp = broadcastOp->AsHWIntrinsic()->Op(1);
                                 }
 
-                                assert(broadcastOp->OperIsConst());
+                                if (!nestedIntrin->OperIsMemoryLoad() && scalarOp->OperIsConst())
+                                {
+                                    GenTree* vecCns =
+                                        m_compiler->gtNewSimdCreateBroadcastNode(simdType, scalarOp,
+                                                                                 nestedIntrin->GetSimdBaseType(),
+                                                                                 simdSize);
 
-                                GenTree* vecCns =
-                                    m_compiler->gtNewSimdCreateBroadcastNode(simdType, broadcastOp,
-                                                                             nestedIntrin->GetSimdBaseType(), simdSize);
+                                    assert(vecCns->IsCnsVec());
+                                    BlockRange().InsertAfter(scalarOp, vecCns);
+                                    nestedOp2 = vecCns;
 
-                                assert(vecCns->IsCnsVec());
-                                BlockRange().InsertAfter(broadcastOp, vecCns);
-                                nestedOp2 = vecCns;
+                                    if (scalarOp != broadcastOp)
+                                    {
+                                        BlockRange().Remove(broadcastOp);
+                                    }
 
-                                BlockRange().Remove(broadcastOp);
-                                BlockRange().Remove(nestedIntrin);
+                                    BlockRange().Remove(scalarOp);
+                                    BlockRange().Remove(nestedIntrin);
+                                }
                             }
                         }
 
@@ -3275,6 +3281,7 @@ GenTree* Lowering::LowerHWIntrinsicCmpOp(GenTreeHWIntrinsic* node, genTreeOps cm
             // so ensure that we track the base type as the one we'll be producing
             // via the vector comparison introduced here.
             maskBaseType = simdBaseType;
+            count        = simdSize / genTypeSize(maskBaseType);
 
             // We have `x == y` or `x != y` both of which where we want to find `AllBitsSet` in the mask since
             // we can directly do the relevant comparison. Given the above tables then when we have a full mask
@@ -10481,10 +10488,10 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
                                                 GenTreeHWIntrinsic* broadcastNode =
                                                     op2->AsHWIntrinsic()->Op(broadcastOpIndex)->AsHWIntrinsic();
                                                 GenTree* constNode = broadcastNode->Op(1);
-                                                int64_t  lval      = 0;
+                                                uint64_t lval      = 0;
 
                                                 assert(genTypeSize(constNode) == 4);
-                                                assert(tgtMaskSize == 2);
+                                                assert(tgtMaskSize == (simdSize / 8));
 
                                                 if (constNode->IsCnsFltOrDbl())
                                                 {

src/tests/JIT/Regression_ro_2/Runtime_133548.cs

@@ -0,0 +1,23 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+using System.Runtime.Intrinsics.X86;
+using Xunit;
+
+public class Runtime_133548
+{
+    [ConditionalFact(typeof(Sse41), nameof(Sse41.IsSupported))]
+    public static void TestEntryPoint()
+    {
+        Vector128<byte> right = Vector128.Create(0x01010000u, 0x00000101u, 0x01010000u, 0x00000101u).AsByte();
+        Vector128<float> value = Vector128.Create(1.1f);
+
+        Assert.Equal(Vector128.Create(0.0f, 1.1f, 0.0f, 1.1f), Blend(value, Vector128<byte>.Zero, right));
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector128<float> Blend(Vector128<float> value, Vector128<byte> left, Vector128<byte> right) =>
+        Sse41.BlendVariable(Vector128<float>.Zero, value, Sse2.CompareEqual(left, right).AsSingle());
+}

src/tests/JIT/Regression_ro_2/Runtime_133552.cs

@@ -0,0 +1,55 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+using Xunit;
+
+public class Runtime_133552
+{
+    [Fact]
+    public static void TestRuntimeBroadcast()
+    {
+        // A widened load must not read the different neighboring word.
+        float[] source = [-1.0f, 2.0f];
+        Vector128<float> value = Vector128.Create(0x12345678u).AsSingle();
+        Vector128<double> left = Vector128.Create(1.0, 2.0);
+        Vector128<double> right = Vector128.Create(1.0, 3.0);
+        Vector128<double> fallback = Vector128.Create(9.0);
+        Vector128<double> expected = Vector128.Create(0x1200000012000000ul, 0).AsDouble().WithElement(1, 9.0);
+
+        Assert.Equal(expected, SelectRuntime(value, ref source[0], left, right, fallback));
+    }
+
+    [Fact]
+    public static void TestWideConstantBroadcast()
+    {
+        Vector256<float> value = Vector256.Create(0x12345678u).AsSingle();
+        Vector256<double> left = Vector256.Create(1.0, 2.0, 1.0, 2.0);
+        Vector256<double> right = Vector256.Create(1.0, 3.0, 1.0, 3.0);
+        Vector256<double> fallback = Vector256.Create(9.0);
+        Vector256<double> expected = Vector256.Create(0xADB45678ADB45678ul).AsDouble().WithElement(1, 9.0).WithElement(3, 9.0);
+
+        Assert.Equal(expected, SelectConstant(value, left, right, fallback));
+    }
+
+    [Fact]
+    public static void TestBroadcastAndZero()
+    {
+        int[] source = [0x12345678, 0];
+        Assert.True(BroadcastAndZero(Vector512<int>.Zero, ref source[0]));
+        Assert.False(BroadcastAndZero(Vector512<int>.AllBitsSet, ref source[0]));
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector128<double> SelectRuntime(Vector128<float> value, ref float source, Vector128<double> left, Vector128<double> right, Vector128<double> fallback) =>
+        Vector128.ConditionalSelect(Vector128.Equals(left, right), (value & Vector128.Create(source)).AsDouble(), fallback);
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector256<double> SelectConstant(Vector256<float> value, Vector256<double> left, Vector256<double> right, Vector256<double> fallback) =>
+        Vector256.ConditionalSelect(Vector256.Equals(left, right), (value ^ Vector256.Create(-1.0f)).AsDouble(), fallback);
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static bool BroadcastAndZero(Vector512<int> value, ref int source) =>
+        (value & Vector512.Create(source)).AsInt64() == Vector512<long>.Zero;
+}

src/tests/JIT/Regression_ro_2/Runtime_133555.cs

@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+using Xunit;
+
+public class Runtime_133555
+{
+    [Fact]
+    public static void TestEntryPoint()
+    {
+        Assert.True(Equal(Vector128<byte>.AllBitsSet, Vector128<byte>.Zero, Vector128<long>.AllBitsSet));
+        Assert.False(Equal(Vector128<byte>.AllBitsSet, Vector128<byte>.Zero, Vector128<long>.Zero));
+    }
+
+    [Fact]
+    public static void TestConstantComparison()
+    {
+        Assert.True(EqualConstant(Vector128.Create(-1L, 0L).AsByte(), Vector128<byte>.Zero));
+        Assert.False(EqualConstant(Vector128<byte>.AllBitsSet, Vector128<byte>.Zero));
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static bool Equal(Vector128<byte> left, Vector128<byte> right, Vector128<long> other) =>
+        Vector128.GreaterThan(left, right).AsInt64() == other;
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static bool EqualConstant(Vector128<byte> left, Vector128<byte> right) =>
+        Vector128.GreaterThan(left, right).AsInt64() == Vector128.Create(-1L, 0L);
+}

src/tests/JIT/opt/SVE/PredicateInstructions.cs

@@ -14,6 +14,42 @@ public class PredicateInstructions
     private static readonly float[] s_floatValues = new float[64];
     private static readonly double[] s_doubleValues = new double[64];
 
+    [ConditionalFact(typeof(Sve), nameof(Sve.IsSupported))]
+    public static void TestMaskElementGranularity()
+    {
+        int[] integers = new int[Vector<int>.Count];
+        Array.Fill(integers, 1);
+        integers[0] = 0;
+
+        short[] shorts = new short[Vector<short>.Count];
+        Array.Fill(shorts, (short)1);
+        shorts[0] = 0;
+
+        Vector<int> right = new Vector<int>(integers);
+        Vector<byte> reversed = ReverseReinterpretedComparison(Vector<int>.Zero, right);
+        Vector<byte> anded = AndReinterpretedComparisons(Vector<int>.Zero, right, Vector<short>.Zero, new Vector<short>(shorts));
+        Vector<byte> selected = SelectReinterpretedComparisons(Vector<byte>.AllBitsSet, Vector<int>.Zero, right);
+
+        for (int i = 0; i < Vector<byte>.Count; i++)
+        {
+            Assert.Equal(i >= Vector<byte>.Count - sizeof(int) ? byte.MaxValue : (byte)0, reversed[i]);
+            Assert.Equal(i < sizeof(short) ? byte.MaxValue : (byte)0, anded[i]);
+            Assert.Equal(i < sizeof(int) ? byte.MaxValue : (byte)0, selected[i]);
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector<byte> ReverseReinterpretedComparison(Vector<int> left, Vector<int> right) =>
+        Sve.ReverseElement(Vector.AsVectorByte(Sve.CompareEqual(left, right)));
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector<byte> AndReinterpretedComparisons(Vector<int> left, Vector<int> right, Vector<short> leftShort, Vector<short> rightShort) =>
+        Sve.And(Vector.AsVectorByte(Sve.CompareEqual(left, right)), Vector.AsVectorByte(Sve.CompareEqual(leftShort, rightShort)));
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static Vector<byte> SelectReinterpretedComparisons(Vector<byte> condition, Vector<int> left, Vector<int> right) =>
+        Sve.ConditionalSelect(condition, Vector.AsVectorByte(Sve.CompareEqual(left, right)), Vector.AsVectorByte(Sve.CompareGreaterThan(left, right)));
+
     [MethodImpl(MethodImplOptions.NoInlining)]
     [Fact]
     public static void TestPredicateInstructions()