[NativeAOT-LLVM] Stop caching GC_EXPOSED_NO results in LSSA #3138 from SingleAccretion/LSSA-DefStatusFix

[NativeAOT-LLVM] Stop caching `GC_EXPOSED_NO` results in LSSA

Author: pavelsavara

src/coreclr/jit/llvmlssa.cpp

@@ -242,7 +242,6 @@ class ShadowStackAllocator
     class LssaDomTreeVisitor : public DomTreeVisitor<LssaDomTreeVisitor>
     {
         class NotExposedReason;
-        struct NotExposedReasonLink;
 
         // Cannot use raw node pointers as their values influence hash table iteration order.
         struct DeterministicNodeHashInfo : public HashTableInfo<DeterministicNodeHashInfo>
@@ -329,7 +328,7 @@ class ShadowStackAllocator
         ArrayStack<BitVec> m_candidateBitSetPool;
 
 #ifdef DEBUG
-        JitHashTable<LclSsaVarDsc*, JitPtrKeyFuncs<LclSsaVarDsc>, NotExposedReasonLink*> m_notExposedReasonMap;
+        JitHashTable<LclSsaVarDsc*, JitPtrKeyFuncs<LclSsaVarDsc>, const char*> m_notExposedReasonMap;
 #endif // DEBUG
 
     public:
@@ -616,6 +615,16 @@ class ShadowStackAllocator
 
         void SpillSdsuValue(LIR::Range& blockRange, GenTree* defNode, unsigned* pSpillLclNum)
         {
+            // Like with candidates, the GC status of an SDSU can change during its lifetime, so we need to check it
+            // at every safepoint, in the general case. We could tighten this by filtering out more SDSUs early (like
+            // 'null' and such), but live-across-a-safepoint SDSUs are not that common, so we currently don't.
+            INDEBUG(NotExposedReason reason(this));
+            if (!IsGcExposedSdsuValue(defNode, DefStatus::Active DEBUGARG(&reason)))
+            {
+                JITDUMPEXEC(reason.Print("[%06u] is live, but not exposed: ", "\n", Compiler::dspTreeID(defNode)));
+                return;
+            }
+
             if (*pSpillLclNum != BAD_VAR_NUM)
             {
                 // We may have already spilled this def live across multiple safe points.
@@ -748,9 +757,7 @@ class ShadowStackAllocator
             unsigned status = GetGcExposedStatus(ssaDsc);
             if (status == GC_EXPOSED_UNKNOWN)
             {
-                bool isExposed = true;
                 GenTreeLclVarCommon* defNode = ssaDsc->GetDefNode();
-                INDEBUG(NotExposedReason defReason(this));
                 if (defNode == nullptr)
                 {
                     assert(ssaNum == SsaConfig::FIRST_SSA_NUM);
@@ -760,19 +767,25 @@ class ShadowStackAllocator
                     assert(initKind != ValueInitKind::None);
                     if (initKind != ValueInitKind::Param)
                     {
-                        INDEBUG(defReason.Add("implicit zero def"));
-                        isExposed = false;
+                        INDEBUG(pReason->Add("implicit zero def"));
+                        status = GC_EXPOSED_NO;
                     }
                 }
                 else if (defNode->OperIs(GT_STORE_LCL_VAR) &&
-                    !IsGcExposedValue(defNode->Data(), DefStatus::Unknown DEBUGARG(&defReason)))
+                    !IsGcExposedValue(defNode->Data(), DefStatus::Unknown DEBUGARG(pReason)))
                 {
-                    isExposed = false;
+                    status = GC_EXPOSED_NO;
                 }
 
-                // This caching is designed to prevent quadratic behavior.
-                status = isExposed ? GC_EXPOSED_YES : GC_EXPOSED_NO;
-                SetGcExposedStatus(ssaDsc, status DEBUGARG(varDsc) DEBUGARG(&defReason));
+                if (status != GC_EXPOSED_NO)
+                {
+                    // This caching is designed to prevent quadratic behavior. However, we can't cache "NO" results as
+                    // they depend on the current execution point (different points may yield different results as
+                    // shadow slots get overwritten with new values), so this is not a complete solution...
+                    // TODO-LLVM-LSSA: fix the above with a new LSSA algorithm.
+                    status = GC_EXPOSED_YES;
+                    SetGcExposedStatus(ssaDsc, status DEBUGARG(varDsc));
+                }
             }
             // Take care not to depend on potentially stale information. Consider:
             //
@@ -809,7 +822,6 @@ class ShadowStackAllocator
             }
 
             LclVarDsc* varDsc;
-            INDEBUG(NotExposedReason reason(this));
             if (IsCandidateLocalNode(node, &varDsc))
             {
                 JITDUMP(" -- Processing a candidate:\n");
@@ -836,8 +848,7 @@ class ShadowStackAllocator
                     IncrementActiveUseCount(varDsc->GetPerSsaData(ssaNum));
                 }
             }
-            else if (node->IsValue() && !node->IsUnusedValue() && IsGcExposedType(node) &&
-                IsGcExposedSdsuValue(node, DefStatus::Active DEBUGARG(&reason)))
+            else if (node->IsValue() && !node->IsUnusedValue() && IsGcExposedType(node))
             {
                 node->gtLIRFlags |= LIR::Flags::Mark;
                 m_liveSdsuGcDefs.AddOrUpdate(node, BAD_VAR_NUM);
@@ -852,17 +863,7 @@ class ShadowStackAllocator
             }
             if (node->TypeIs(TYP_STRUCT))
             {
-                // TODO-LLVM: delete this once we're up to date with upstream deleting "STORE_DYN_BLK".
-                if (node->OperIs(GT_IND))
-                {
-                    return false;
-                }
-                if (!node->GetLayout(m_compiler)->HasGCPtr())
-                {
-                    return false;
-                }
-
-                return true;
+                return node->GetLayout(m_compiler)->HasGCPtr();
             }
 
             return false;
@@ -1112,7 +1113,7 @@ class ShadowStackAllocator
         }
 
         void SetGcExposedStatus(
-            LclSsaVarDsc* ssaDsc, unsigned status DEBUGARG(LclVarDsc* varDsc) DEBUGARG(NotExposedReason* pReason))
+            LclSsaVarDsc* ssaDsc, unsigned status DEBUGARG(LclVarDsc* varDsc) DEBUGARG(NotExposedReason* pReason = nullptr))
         {
             unsigned oldStatus = GetGcExposedStatus(ssaDsc);
             assert((oldStatus == GC_EXPOSED_UNKNOWN) || ((oldStatus == GC_EXPOSED_YES) && (status == GC_EXPOSED_SPILL)));
@@ -1123,7 +1124,7 @@ class ShadowStackAllocator
 
             JITDUMP("V%02u/%u: %s -> %s\n", m_compiler->lvaGetLclNum(varDsc),
                 varDsc->GetSsaNumForSsaDef(ssaDsc), GcStatusToString(oldStatus), GcStatusToString(status));
-            INDEBUG(pReason->SaveForDef(ssaDsc));
+            DBEXEC(pReason != nullptr, pReason->SaveForDef(ssaDsc));
         }
 
         unsigned GetGcExposedStatus(LclSsaVarDsc* ssaDsc)
@@ -1353,12 +1354,6 @@ class ShadowStackAllocator
             }
         }
 
-        struct NotExposedReasonLink
-        {
-            const char* Reason;
-            NotExposedReasonLink* Next;
-        };
-
         class NotExposedReason
         {
             LssaDomTreeVisitor* m_visitor;
@@ -1393,12 +1388,11 @@ class ShadowStackAllocator
             {
                 if (m_enabled)
                 {
+                    const char* reason;
                     LclSsaVarDsc* ssaDsc = varDsc->GetPerSsaData(ssaNum);
-                    NotExposedReasonLink* link = m_visitor->m_notExposedReasonMap[ssaDsc];
-                    while (link != nullptr)
+                    if (m_visitor->m_notExposedReasonMap.Lookup(ssaDsc, &reason))
                     {
-                        Add(link->Reason);
-                        link = link->Next;
+                        Add(reason);
                     }
 
                     unsigned lclNum = m_visitor->m_compiler->lvaGetLclNum(varDsc);
@@ -1429,26 +1423,11 @@ class ShadowStackAllocator
 
             void SaveForDef(LclSsaVarDsc* ssaDsc)
             {
-                CompAllocator alloc = m_visitor->m_compiler->getAllocator(CMK_DebugOnly);
-                NotExposedReasonLink* links = nullptr;
-                NotExposedReasonLink* last = nullptr;
-                for (int i = 0; i < m_chain.Height(); i++)
-                {
-                    NotExposedReasonLink* link = new (alloc) NotExposedReasonLink();
-                    link->Reason = m_chain.Bottom(i);
-                    if (last == nullptr)
-                    {
-                        links = link;
-                    }
-                    else
-                    {
-                        last->Next = link;
-                    }
-                    last = link;
-                }
-                if (links != nullptr)
+                if (m_enabled)
                 {
-                    m_visitor->m_notExposedReasonMap.Set(ssaDsc, links);
+                    assert(m_chain.Height() == 1); // Full support currently not needed.
+                    const char* reason = m_chain.Top();
+                    m_visitor->m_notExposedReasonMap.Set(ssaDsc, reason);
                 }
             }
 

src/tests/nativeaot/SmokeTests/HelloWasm/HelloWasm.cs

@@ -430,6 +430,8 @@ private static unsafe int Main(string[] args)
 
         LSSATests.Run();
 
+        FunctionalLSSATests.Run();
+
         TestThreadStaticAlignment();
 
         TestLiveInFirstBlockForTracked(new object());

src/tests/nativeaot/SmokeTests/HelloWasm/LSSATests.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 
 namespace System.Runtime.JitTesting;
 
@@ -630,9 +631,83 @@ private static void YetAnotherSafePoint() { }
 
     // Roslyn likes to eliminate locals we write. Force it to abstain with this function.
     private static void ForceILLocal(void* x) { }
+}
+
+public unsafe class FunctionalLSSATests
+{
+    private static bool s_failed;
+    private static GCHandle s_chkObjWeakHandle;
+
+    // These tests, unlike the ones above, test the generated code indirectly, by creating objects and checking they're
+    // kept alive as expected. They are ideally suited for testing correctness, where we need something to be visible
+    // to the GC but do not care particularly about how precisely that is achieved.
+    public static void Run()
+    {
+        s_chkObjWeakHandle = GCHandle.Alloc(null, GCHandleType.Weak);
+
+        Program.StartTest("FunctionalLSSATests");
+
+        SpilledCandidateSlotInvalidated();
+        if (!FunctionalTestSucceeded(nameof(SpilledCandidateSlotInvalidated)))
+            return;
+
+        Program.PassTest();
+    }
+
+    private static bool FunctionalTestSucceeded(string name)
+    {
+        if (s_failed)
+        {
+            Program.FailTest($"  {name} failed");
+            return false;
+        }
+        return true;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static ref int SpilledCandidateSlotInvalidated()
+    {
+        // Here we are testing that 'derivedRef' keeps 'chkObj' alive as expected,
+        // despite the fact 'chkObj' is spilled at the point 'derivedRef' is formed.
+        ClassWithField chkObj = NewCheckedObject();
+        LogicalSafePoint(chkObj, chkObj);
+        LogicalSafePoint(chkObj, chkObj);
+
+        ref int derivedRef = ref chkObj.Field;
+        chkObj = GetNullOpaquely(); // Invalidate the pin.
+        LogicalSafePoint(chkObj, chkObj);
+        LogicalSafePoint(chkObj, chkObj);
+        ValidateCheckedObjectIsAlive();
+
+        return ref derivedRef;
+    }
 
-    class ClassWithField
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static ClassWithField NewCheckedObject()
     {
-        public int Field;
+        ClassWithField obj = new();
+        s_chkObjWeakHandle.Target = obj;
+        return obj;
     }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static void ValidateCheckedObjectIsAlive()
+    {
+        GC.Collect();
+        if (s_chkObjWeakHandle.Target is null)
+        {
+            s_failed = true;
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static void LogicalSafePoint(object x = null, object y = null) { }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static ClassWithField GetNullOpaquely() => null;
+}
+
+class ClassWithField
+{
+    public int Field;
 }