[release/10.0] JIT: fix switch peeling stealing the wrong operand of its compare

src/coreclr/jit/fgopt.cpp

@@ -2792,14 +2792,17 @@ void Compiler::fgPeelSwitch(BasicBlock* block)
 
     // Set up a compare in the upstream block, "stealing" the switch value tree.
     //
-    GenTree* const   dominantCaseCompare = gtNewOperNode(GT_EQ, TYP_INT, switchValue, gtNewIconNode(dominantCase));
-    GenTree* const   jmpTree             = gtNewOperNode(GT_JTRUE, TYP_VOID, dominantCaseCompare);
-    Statement* const jmpStmt             = fgNewStmtFromTree(jmpTree, switchStmt->GetDebugInfo());
-    fgInsertStmtAtEnd(block, jmpStmt);
+    GenTree* const dominantCaseCompare = gtNewOperNode(GT_EQ, TYP_INT, switchValue, gtNewIconNode(dominantCase));
+    GenTree* const jmpTree             = gtNewOperNode(GT_JTRUE, TYP_VOID, dominantCaseCompare);
 
     // Reattach switch value to the switch. This may introduce a comma
     // in the upstream compare tree, if the switch value expression is complex.
     //
+    // Note this must happen before the compare is put into a statement below: creating the
+    // statement sequences the tree via gtSetEvalOrder, which is allowed to swap the operands
+    // of the compare (and does so when the switch value is a constant). After that point
+    // "gtOp1" is no longer guaranteed to be the switch value.
+    //
     switchTree->AsOp()->gtOp1 = fgMakeMultiUse(&dominantCaseCompare->AsOp()->gtOp1);
 
     // Update flags
@@ -2809,6 +2812,9 @@ void Compiler::fgPeelSwitch(BasicBlock* block)
     jmpTree->gtFlags |= dominantCaseCompare->gtFlags & GTF_ALL_EFFECT;
     dominantCaseCompare->gtFlags |= GTF_RELOP_JMP_USED | GTF_DONT_CSE;
 
+    Statement* const jmpStmt = fgNewStmtFromTree(jmpTree, switchStmt->GetDebugInfo());
+    fgInsertStmtAtEnd(block, jmpStmt);
+
     // Wire up the new control flow.
     //
     FlowEdge* const blockToTargetEdge   = fgAddRefPred(dominantTarget, block, dominantEdge);
@@ -2851,11 +2857,8 @@ void Compiler::fgPeelSwitch(BasicBlock* block)
         gtSetStmtInfo(switchStmt);
         fgSetStmtSeq(switchStmt);
 
-        // fgNewStmtFromTree() already threaded the tree, but calling fgMakeMultiUse() might have
-        // added new nodes if a COMMA was introduced.
-        JITDUMP("Rethreading " FMT_STMT "\n", jmpStmt->GetID());
-        gtSetStmtInfo(jmpStmt);
-        fgSetStmtSeq(jmpStmt);
+        // Note the compare does not need rethreading here: it was fully built (including any
+        // nodes fgMakeMultiUse() added) before fgNewStmtFromTree() sequenced it.
     }
 }
 

src/tests/JIT/Regression/JitBlue/Runtime_132370/Runtime_132370.cs

@@ -0,0 +1,95 @@
+// 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.Threading;
+using Xunit;
+
+// When a loop containing a switch is unrolled, the switch value becomes a constant in each
+// unrolled copy. Switch peeling would then "steal" the wrong operand of the compare it creates
+// (sequencing the compare can swap its operands when both are constants), so every copy of the
+// switch ended up dispatching on the dominant case value instead of on its own value.
+
+public class Runtime_132370
+{
+    private readonly struct V3
+    {
+        public readonly long X;
+        public readonly long Y;
+        public readonly long Z;
+
+        public V3(long x, long y, long z)
+        {
+            X = x;
+            Y = y;
+            Z = z;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static long Get(V3 value, int axis) => axis switch
+        {
+            0 => value.X,
+            1 => value.Y,
+            2 => value.Z,
+            _ => throw new ArgumentOutOfRangeException(nameof(axis)),
+        };
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static bool IsZero(V3 value)
+    {
+        for (int axis = 0; axis < 3; axis++)
+        {
+            if (V3.Get(value, axis) != 0)
+            {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static int CountZero(V3[] values)
+    {
+        int zero = 0;
+
+        for (int i = 0; i < values.Length; i++)
+        {
+            if (IsZero(values[i]))
+            {
+                zero++;
+            }
+        }
+
+        return zero;
+    }
+
+    [Fact]
+    public static void TestEntryPoint()
+    {
+        // A mixture of values is needed so that case 0 is the dominant case in the profile.
+        V3[] values = new V3[256];
+        for (int i = 0; i < values.Length - 1; i++)
+        {
+            values[i] = new V3(1, 0, 0);
+        }
+
+        values[^1] = new V3(0, 0, 1);
+
+        // None of the vectors is the zero vector, so this must always be 0. The bad code only
+        // showed up once IsZero reached tier1 with profile data, so keep calling it, yielding in
+        // between so that the background compilation has a chance to happen. In practice this
+        // fails on the ~7th outer iteration when the JIT is broken.
+        for (int i = 0; i < 50; i++)
+        {
+            for (int j = 0; j < 1000; j++)
+            {
+                Assert.Equal(0, CountZero(values));
+            }
+
+            Thread.Sleep(16);
+        }
+    }
+}

src/tests/JIT/Regression/JitBlue/Runtime_132370/Runtime_132370.csproj

@@ -0,0 +1,14 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+    <CLRTestPriority>1</CLRTestPriority>
+    <!-- Needed for CLRTestEnvironmentVariable -->
+    <RequiresProcessIsolation>true</RequiresProcessIsolation>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+    <!-- The switch needs profile data for a dominant case to be peeled -->
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredCompilation" Value="1" />
+    <CLRTestEnvironmentVariable Include="DOTNET_TieredPGO" Value="1" />
+  </ItemGroup>
+</Project>