Fix handling of CreateScalarUnsafe for embedded broadcast (#87134)

* Adding a regression test for #87116 * Ensure IsContainableHWIntrinsicOp takes into account whether CreateScalarUnsafe is coming from memory for embedded broadcast
dotnet · Jun 7, 2023 · eb09ba9 · eb09ba9
1 parent 54dab73
commit eb09ba9
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 5 deletions.
diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp
@@ -25059,6 +25059,37 @@ bool GenTreeHWIntrinsic::OperIsBroadcastScalar() const
 #endif
 }
 
+//------------------------------------------------------------------------
+// OperIsCreateScalarUnsafe: Is this HWIntrinsic a CreateScalarUnsafe node.
+//
+// Return Value:
+//    Whether "this" is a CreateScalarUnsafe node.
+//
+bool GenTreeHWIntrinsic::OperIsCreateScalarUnsafe() const
+{
+    NamedIntrinsic intrinsicId = GetHWIntrinsicId();
+
+    switch (intrinsicId)
+    {
+#if defined(TARGET_ARM64)
+        case NI_Vector64_CreateScalarUnsafe:
+#endif // TARGET_ARM64
+        case NI_Vector128_CreateScalarUnsafe:
+#if defined(TARGET_XARCH)
+        case NI_Vector256_CreateScalarUnsafe:
+        case NI_Vector512_CreateScalarUnsafe:
+#endif // TARGET_XARCH
+        {
+            return true;
+        }
+
+        default:
+        {
+            return false;
+        }
+    }
+}
+
 //------------------------------------------------------------------------------
 // OperRequiresAsgFlag : Check whether the operation requires GTF_ASG flag regardless
 //                       of the children's flags.

diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h
@@ -6245,6 +6245,7 @@ struct GenTreeHWIntrinsic : public GenTreeJitIntrinsic
     bool OperIsMemoryStoreOrBarrier() const;
     bool OperIsEmbBroadcastCompatible() const;
     bool OperIsBroadcastScalar() const;
+    bool OperIsCreateScalarUnsafe() const;
 
     bool OperRequiresAsgFlag() const;
     bool OperRequiresCallFlag() const;

diff --git a/src/coreclr/jit/lowerxarch.cpp b/src/coreclr/jit/lowerxarch.cpp
@@ -7599,7 +7599,7 @@ bool Lowering::IsContainableHWIntrinsicOp(GenTreeHWIntrinsic* parentNode, GenTre
 
             GenTree* op1 = hwintrinsic->Op(1);
 
-            if (IsSafeToContainMem(parentNode, hwintrinsic, op1))
+            if (IsInvariantInRange(op1, parentNode, hwintrinsic))
             {
                 if (op1->isContained())
                 {
@@ -7706,14 +7706,29 @@ bool Lowering::IsContainableHWIntrinsicOp(GenTreeHWIntrinsic* parentNode, GenTre
             if (intrinsicId == NI_SSE3_MoveAndDuplicate)
             {
                 // NI_SSE3_MoveAndDuplicate is for Vector128<double> only.
-                assert(childNode->AsHWIntrinsic()->GetSimdBaseType() == TYP_DOUBLE);
+                assert(hwintrinsic->GetSimdBaseType() == TYP_DOUBLE);
             }
+
             if (comp->compOpportunisticallyDependsOn(InstructionSet_AVX512F_VL) &&
                 parentNode->OperIsEmbBroadcastCompatible())
             {
-                GenTree* broadcastOperand = childNode->AsHWIntrinsic()->Op(1);
-                bool     childSupportsRegOptional;
-                if (IsContainableHWIntrinsicOp(childNode->AsHWIntrinsic(), broadcastOperand, &childSupportsRegOptional))
+                GenTree* broadcastOperand = hwintrinsic->Op(1);
+
+                if (broadcastOperand->OperIsHWIntrinsic())
+                {
+                    GenTreeHWIntrinsic* hwintrinsicOperand = broadcastOperand->AsHWIntrinsic();
+
+                    if (hwintrinsicOperand->OperIsCreateScalarUnsafe())
+                    {
+                        // CreateScalarUnsafe can contain non-memory operands such as enregistered
+                        // locals, so we want to check if its operand is containable instead. This
+                        // will result in such enregistered locals returning `false`.
+                        broadcastOperand = hwintrinsicOperand->Op(1);
+                    }
+                }
+
+                bool childSupportsRegOptional;
+                if (IsContainableHWIntrinsicOp(hwintrinsic, broadcastOperand, &childSupportsRegOptional))
                 {
                     return true;
                 }

diff --git a/src/tests/JIT/Regression/JitBlue/Runtime_87116/Runtime_87116.cs b/src/tests/JIT/Regression/JitBlue/Runtime_87116/Runtime_87116.cs
@@ -0,0 +1,26 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+using Xunit;
+
+public class Runtime_87116
+{
+    [Fact]
+    public static int Test()
+    {
+        return TryVectorAdd(1, 2, 1 + 2) ? 100 : 0;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public static bool TryVectorAdd(float a, float b, float c)
+    {
+        Vector128<float> A = Vector128.Create(a);
+        Vector128<float> B = Vector128.Create(b);
+
+        Vector128<float> C = A + B;
+        return C == Vector128.Create(c);
+    }
+}
diff --git a/src/tests/JIT/Regression/JitBlue/Runtime_87116/Runtime_87116.csproj b/src/tests/JIT/Regression/JitBlue/Runtime_87116/Runtime_87116.csproj
@@ -0,0 +1,8 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>