From 41a8e3963f2f94bc19565553433a3f48c8af052f Mon Sep 17 00:00:00 2001
From: Tanner Gooding <tagoo@outlook.com>
Date: Tue, 19 Sep 2023 15:09:19 -0700
Subject: [PATCH] Ensure VN handles both forms of the xarch shift instructions
 for SIMD (#91601)

---
 src/coreclr/jit/valuenum.cpp                  | 60 +++++++++++++++++++
 .../JitBlue/Runtime_91175/Runtime_91175.cs    | 30 ++++++++++
 .../Runtime_91175/Runtime_91175.csproj        |  8 +++
 3 files changed, 98 insertions(+)
 create mode 100644 src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.cs
 create mode 100644 src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.csproj

diff --git a/src/coreclr/jit/valuenum.cpp b/src/coreclr/jit/valuenum.cpp
index b4d7985cd97c5..79da0ca4601bd 100644
--- a/src/coreclr/jit/valuenum.cpp
+++ b/src/coreclr/jit/valuenum.cpp
@@ -7432,6 +7432,26 @@ ValueNum ValueNumStore::EvalHWIntrinsicFunBinary(var_types      type,
             case NI_AVX512BW_ShiftLeftLogical:
 #endif
             {
+#ifdef TARGET_XARCH
+                if (TypeOfVN(arg1VN) == TYP_SIMD16)
+                {
+                    // The xarch shift instructions support taking the shift amount as
+                    // a simd16, in which case they take the shift amount from the lower
+                    // 64-bits.
+
+                    uint64_t shiftAmount = GetConstantSimd16(arg1VN).u64[0];
+
+                    if (genTypeSize(baseType) != 8)
+                    {
+                        arg1VN = VNForIntCon(static_cast<int32_t>(shiftAmount));
+                    }
+                    else
+                    {
+                        arg1VN = VNForLongCon(static_cast<int64_t>(shiftAmount));
+                    }
+                }
+#endif // TARGET_XARCH
+
                 return EvaluateBinarySimd(this, GT_LSH, /* scalar */ false, type, baseType, arg0VN, arg1VN);
             }
 
@@ -7445,6 +7465,26 @@ ValueNum ValueNumStore::EvalHWIntrinsicFunBinary(var_types      type,
             case NI_AVX512BW_ShiftRightArithmetic:
 #endif
             {
+#ifdef TARGET_XARCH
+                if (TypeOfVN(arg1VN) == TYP_SIMD16)
+                {
+                    // The xarch shift instructions support taking the shift amount as
+                    // a simd16, in which case they take the shift amount from the lower
+                    // 64-bits.
+
+                    uint64_t shiftAmount = GetConstantSimd16(arg1VN).u64[0];
+
+                    if (genTypeSize(baseType) != 8)
+                    {
+                        arg1VN = VNForIntCon(static_cast<int32_t>(shiftAmount));
+                    }
+                    else
+                    {
+                        arg1VN = VNForLongCon(static_cast<int64_t>(shiftAmount));
+                    }
+                }
+#endif // TARGET_XARCH
+
                 return EvaluateBinarySimd(this, GT_RSH, /* scalar */ false, type, baseType, arg0VN, arg1VN);
             }
 
@@ -7457,6 +7497,26 @@ ValueNum ValueNumStore::EvalHWIntrinsicFunBinary(var_types      type,
             case NI_AVX512BW_ShiftRightLogical:
 #endif
             {
+#ifdef TARGET_XARCH
+                if (TypeOfVN(arg1VN) == TYP_SIMD16)
+                {
+                    // The xarch shift instructions support taking the shift amount as
+                    // a simd16, in which case they take the shift amount from the lower
+                    // 64-bits.
+
+                    uint64_t shiftAmount = GetConstantSimd16(arg1VN).u64[0];
+
+                    if (genTypeSize(baseType) != 8)
+                    {
+                        arg1VN = VNForIntCon(static_cast<int32_t>(shiftAmount));
+                    }
+                    else
+                    {
+                        arg1VN = VNForLongCon(static_cast<int64_t>(shiftAmount));
+                    }
+                }
+#endif // TARGET_XARCH
+
                 return EvaluateBinarySimd(this, GT_RSZ, /* scalar */ false, type, baseType, arg0VN, arg1VN);
             }
 
diff --git a/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.cs b/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.cs
new file mode 100644
index 0000000000000..6e1f8fe9fd525
--- /dev/null
+++ b/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.cs
@@ -0,0 +1,30 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.aa
+
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+using System.Runtime.Intrinsics.X86;
+using Xunit;
+
+public class TestClass
+{
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static Vector256<int> Method0() => Avx2.ShiftRightArithmetic(Vector256<int>.AllBitsSet, Vector128<int>.AllBitsSet);
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static Vector128<int> Method1() => Sse2.ShiftRightArithmetic(Vector128<int>.AllBitsSet, Vector128<int>.AllBitsSet);
+
+    [Fact]
+    public static void TestEntryPoint()
+    {
+        if (Avx2.IsSupported)
+        {
+            _ = Method0();
+        }
+
+        if (Sse2.IsSupported)
+        {
+            _ = Method1();
+        }
+    }
+}
diff --git a/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.csproj b/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.csproj
new file mode 100644
index 0000000000000..de6d5e08882e8
--- /dev/null
+++ b/src/tests/JIT/Regression/JitBlue/Runtime_91175/Runtime_91175.csproj
@@ -0,0 +1,8 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>