From ccb8fd2c4aceb0b076fe33c3e6a8356c463209e6 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
Date: Mon, 24 Jun 2024 20:20:30 +0200
Subject: [PATCH] JIT: Handle QMARK properly for LCL_ADDR propagation

The assertions created/used need to take into account that qmark arms
are only conditionally executed.
---
 src/coreclr/jit/lclmorph.cpp | 81 +++++++++++++++++++++++++++++++-----
 1 file changed, 70 insertions(+), 11 deletions(-)

diff --git a/src/coreclr/jit/lclmorph.cpp b/src/coreclr/jit/lclmorph.cpp
index 1774aea55c831..16b65294c6512 100644
--- a/src/coreclr/jit/lclmorph.cpp
+++ b/src/coreclr/jit/lclmorph.cpp
@@ -225,10 +225,11 @@ class LocalEqualsLocalAddrAssertions
     AssertionToIndexMap                       m_map;
     uint64_t*                                 m_lclAssertions;
     uint64_t*                                 m_outgoingAssertions;
-    uint64_t                                  m_currentAssertions = 0;
     BitVec                                    m_localsToExpose;
 
 public:
+    uint64_t CurrentAssertions = 0;
+
     LocalEqualsLocalAddrAssertions(Compiler* comp)
         : m_comp(comp)
         , m_assertions(comp->getAllocator(CMK_LocalAddressVisitor))
@@ -285,28 +286,28 @@ class LocalEqualsLocalAddrAssertions
     {
         if ((m_assertions.Height() == 0) || (block->bbPreds == nullptr) || m_comp->bbIsHandlerBeg(block))
         {
-            m_currentAssertions = 0;
+            CurrentAssertions = 0;
             return;
         }
 
-        m_currentAssertions = UINT64_MAX;
+        CurrentAssertions = UINT64_MAX;
         for (BasicBlock* pred : block->PredBlocks())
         {
             assert(m_comp->m_dfsTree->Contains(pred));
             if (pred->bbPostorderNum <= block->bbPostorderNum)
             {
-                m_currentAssertions = 0;
+                CurrentAssertions = 0;
                 break;
             }
 
-            m_currentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
+            CurrentAssertions &= m_outgoingAssertions[pred->bbPostorderNum];
         }
 
 #ifdef DEBUG
-        if (m_currentAssertions != 0)
+        if (CurrentAssertions != 0)
         {
             JITDUMP(FMT_BB " incoming assertions:\n", block->bbNum);
-            uint64_t assertions = m_currentAssertions;
+            uint64_t assertions = CurrentAssertions;
             do
             {
                 uint32_t index = BitOperations::BitScanForward(assertions);
@@ -328,7 +329,7 @@ class LocalEqualsLocalAddrAssertions
     //
     void EndBlock(BasicBlock* block)
     {
-        m_outgoingAssertions[block->bbPostorderNum] = m_currentAssertions;
+        m_outgoingAssertions[block->bbPostorderNum] = CurrentAssertions;
     }
 
     //-------------------------------------------------------------------
@@ -389,7 +390,7 @@ class LocalEqualsLocalAddrAssertions
             }
         }
 
-        m_currentAssertions |= uint64_t(1) << index;
+        CurrentAssertions |= uint64_t(1) << index;
     }
 
     //-------------------------------------------------------------------
@@ -400,7 +401,7 @@ class LocalEqualsLocalAddrAssertions
     //
     void Clear(unsigned dstLclNum)
     {
-        m_currentAssertions &= ~m_lclAssertions[dstLclNum];
+        CurrentAssertions &= ~m_lclAssertions[dstLclNum];
     }
 
     //-----------------------------------------------------------------------------------
@@ -415,7 +416,7 @@ class LocalEqualsLocalAddrAssertions
     //
     const LocalEqualsLocalAddrAssertion* GetCurrentAssertion(unsigned lclNum)
     {
-        uint64_t curAssertion = m_currentAssertions & m_lclAssertions[lclNum];
+        uint64_t curAssertion = CurrentAssertions & m_lclAssertions[lclNum];
         assert(genMaxOneBit(curAssertion));
         if (curAssertion == 0)
         {
@@ -797,6 +798,64 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
             }
             break;
 
+            case GT_QMARK:
+            {
+                // We have to inline the pre/postorder visit here to handle
+                // assertions properly.
+                GenTreeQmark* qmark = node->AsQmark();
+
+                assert(!node->IsReverseOp());
+                if (WalkTree(&qmark->gtOp1, qmark) == Compiler::WALK_ABORT)
+                {
+                    return Compiler::WALK_ABORT;
+                }
+
+                if (m_lclAddrAssertions != nullptr)
+                {
+                    uint64_t origAssertions = m_lclAddrAssertions->CurrentAssertions;
+
+                    if (WalkTree(&qmark->gtOp2->AsOp()->gtOp1, qmark->gtOp2) == Compiler::WALK_ABORT)
+                    {
+                        return Compiler::WALK_ABORT;
+                    }
+
+                    uint64_t op1Assertions                 = m_lclAddrAssertions->CurrentAssertions;
+                    m_lclAddrAssertions->CurrentAssertions = origAssertions;
+
+                    if (WalkTree(&qmark->gtOp2->AsOp()->gtOp2, qmark->gtOp2) == Compiler::WALK_ABORT)
+                    {
+                        return Compiler::WALK_ABORT;
+                    }
+
+                    uint64_t op2Assertions                 = m_lclAddrAssertions->CurrentAssertions;
+                    m_lclAddrAssertions->CurrentAssertions = op1Assertions & op2Assertions;
+                }
+                else
+                {
+                    if ((WalkTree(&qmark->gtOp2->AsOp()->gtOp1, qmark->gtOp2) == Compiler::WALK_ABORT) ||
+                        (WalkTree(&qmark->gtOp2->AsOp()->gtOp2, qmark->gtOp2) == Compiler::WALK_ABORT))
+                    {
+                        return Compiler::WALK_ABORT;
+                    }
+                }
+
+                assert(TopValue(0).Node() == qmark->gtGetOp2()->gtGetOp2());
+                assert(TopValue(1).Node() == qmark->gtGetOp2()->gtGetOp1());
+                assert(TopValue(2).Node() == qmark->gtGetOp1());
+
+                EscapeValue(TopValue(0), qmark->gtGetOp2());
+                PopValue();
+
+                EscapeValue(TopValue(0), qmark->gtGetOp2());
+                PopValue();
+
+                EscapeValue(TopValue(0), qmark);
+                PopValue();
+
+                PushValue(use);
+                return Compiler::WALK_SKIP_SUBTREES;
+            }
+
             default:
                 break;
         }