JIT: Use faster write barriers when we know the target address is on the heap

src/coreclr/jit/assertionprop.cpp

@@ -4599,13 +4599,16 @@ GenTree* Compiler::optAssertionProp_Ind(ASSERT_VALARG_TP assertions, GenTree* tr
 {
     assert(tree->OperIsIndir());
 
-    bool didNonNullProp = optNonNullAssertionProp_Ind(assertions, tree);
-    bool didNonHeapProp = optNonHeapAssertionProp_Ind(assertions, tree);
-    if (didNonNullProp || didNonHeapProp)
+    bool updated = optNonNullAssertionProp_Ind(assertions, tree);
+    if (tree->OperIs(GT_STOREIND))
     {
-        return optAssertionProp_Update(tree, tree, stmt);
+        updated |= optWriteBarrierAssertionProp_StoreInd(assertions, tree->AsStoreInd());
     }
 
+    if (updated)
+    {
+        return optAssertionProp_Update(tree, tree, stmt);
+    }
     return nullptr;
 }
 
@@ -4866,94 +4869,135 @@ bool Compiler::optNonNullAssertionProp_Ind(ASSERT_VALARG_TP assertions, GenTree*
 }
 
 //------------------------------------------------------------------------
-// optAssertionIsNonHeap: see if we can prove a tree's value will be not GC-tracked
-//   based on assertions
+// GetWriteBarrierForm: Determinate the exact type of write barrier required for the
+//    given address.
 //
 // Arguments:
-//   op         - tree to check
-//   assertions - set of live assertions
-//   pVnBased   - [out] set to true if value numbers were used
-//   pIndex     - [out] the assertion used in the proof
+//    vnStore - ValueNumStore object
+//    vn      - VN of the address
 //
 // Return Value:
-//   true if the tree's value will be not GC-tracked
+//    Exact type of write barrier required for the given address.
 //
-// Notes:
-//   Sets "pVnBased" if the assertion is value number based. If no matching
-//    assertions are found from the table, then returns "NO_ASSERTION_INDEX."
-//
-//   If both VN and assertion table yield a matching assertion, "pVnBased"
-//   is only set and the return value is "NO_ASSERTION_INDEX."
-//
-bool Compiler::optAssertionIsNonHeap(GenTree*         op,
-                                     ASSERT_VALARG_TP assertions DEBUGARG(bool* pVnBased)
-                                         DEBUGARG(AssertionIndex* pIndex))
+static GCInfo::WriteBarrierForm GetWriteBarrierForm(ValueNumStore* vnStore, ValueNum vn)
 {
-    bool vnBased = (!optLocalAssertionProp && vnStore->IsVNObjHandle(op->gtVNPair.GetConservative()));
-#ifdef DEBUG
-    *pIndex   = NO_ASSERTION_INDEX;
-    *pVnBased = vnBased;
-#endif
-
-    if (vnBased)
+    const var_types type = vnStore->TypeOfVN(vn);
+    if (type == TYP_REF)
     {
-        return true;
+        return GCInfo::WriteBarrierForm::WBF_BarrierUnchecked;
+    }
+    if (type != TYP_BYREF)
+    {
+        return GCInfo::WriteBarrierForm::WBF_BarrierUnknown;
     }
 
-    if (op->IsIconHandle(GTF_ICON_OBJ_HDL))
+    VNFuncApp funcApp;
+    if (vnStore->GetVNFunc(vnStore->VNNormalValue(vn), &funcApp))
     {
-        return true;
+        if (funcApp.m_func == VNF_PtrToArrElem)
+        {
+            // Arrays are always on the heap
+            return GCInfo::WriteBarrierForm::WBF_BarrierUnchecked;
+        }
+        if (funcApp.m_func == VNF_PtrToLoc)
+        {
+            // Pointer to a local
+            return GCInfo::WriteBarrierForm::WBF_NoBarrier;
+        }
+        if (funcApp.m_func == VNFunc(GT_ADD))
+        {
+            // Check arguments of the GT_ADD
+            // To make it conservative, we require one of the arguments to be a constant, e.g.:
+            //
+            //   addressOfLocal + cns    -> NoBarrier
+            //   cns + addressWithinHeap -> BarrierUnchecked
+            //
+            // Because "addressOfLocal + nativeIntVariable" could be in fact a pointer to the heap.
+            // if "nativeIntVariable == addressWithinHeap - addressOfLocal".
+            //
+            if (vnStore->IsVNConstantNonHandle(funcApp.m_args[0]))
+            {
+                return GetWriteBarrierForm(vnStore, funcApp.m_args[1]);
+            }
+            if (vnStore->IsVNConstantNonHandle(funcApp.m_args[1]))
+            {
+                return GetWriteBarrierForm(vnStore, funcApp.m_args[0]);
+            }
+        }
     }
 
-    return false;
+    // TODO:
+    // * addr is ByRefLike - NoBarrier (https://github.com/dotnet/runtime/issues/9512)
+    //
+    return GCInfo::WriteBarrierForm::WBF_BarrierUnknown;
 }
 
 //------------------------------------------------------------------------
-// optNonHeapAssertionProp_Ind: Possibly prove an indirection not GC-tracked.
+// optWriteBarrierAssertionProp_StoreInd: This function assists gcIsWriteBarrierCandidate with help of
+//    assertions and VNs since CSE may "hide" addresses/values under locals, making it impossible for
+//    gcIsWriteBarrierCandidate to determine the exact type of write barrier required
+//    (it's too late for it to rely on VNs).
+//
+//    There are three cases we handle here:
+//     * Target is not on the heap - no write barrier is required
+//     * Target could be on the heap, but the value being stored doesn't require any write barrier
+//     * Target is definitely on the heap - checked (slower) write barrier is not required
 //
 // Arguments:
 //    assertions - Active assertions
-//    indir      - The indirection
+//    indir      - The STOREIND node
 //
 // Return Value:
-//    Whether the indirection was found to be not GC-tracked and marked as such.
+//    Whether the exact type of write barrier was determined and marked on the STOREIND node.
 //
-bool Compiler::optNonHeapAssertionProp_Ind(ASSERT_VALARG_TP assertions, GenTree* indir)
+bool Compiler::optWriteBarrierAssertionProp_StoreInd(ASSERT_VALARG_TP assertions, GenTreeStoreInd* indir)
 {
-    assert(indir->OperIsIndir());
+    const GenTree* value = indir->AsIndir()->Data();
+    const GenTree* addr  = indir->AsIndir()->Addr();
 
-    if (!indir->OperIs(GT_STOREIND) || (indir->gtFlags & GTF_IND_TGT_NOT_HEAP) != 0)
+    if (optLocalAssertionProp || !indir->TypeIs(TYP_REF) || !value->TypeIs(TYP_REF) ||
+        ((indir->gtFlags & GTF_IND_TGT_NOT_HEAP) != 0))
     {
         return false;
     }
 
-// We like CSE to happen for handles, as the codegen for loading a 64-bit constant can be pretty heavy
-// and this is particularly true on platforms with a fixed-width instruction encoding. However, this
-// pessimizes stores as we can no longer optimize around some object handles that would allow us to
-// bypass the write barrier.
-//
-// In order to handle that, we'll propagate the IND_TGT_NOT_HEAP flag onto the store if the handle is
-// directly or if the underlying value number is an applicable object handle.
+    GCInfo::WriteBarrierForm barrierType = GCInfo::WriteBarrierForm::WBF_BarrierUnknown;
 
-#ifdef DEBUG
-    bool           vnBased = false;
-    AssertionIndex index   = NO_ASSERTION_INDEX;
-#endif
-    if (optAssertionIsNonHeap(indir->AsIndir()->Data(), assertions DEBUGARG(&vnBased) DEBUGARG(&index)))
+    // First, analyze the value being stored
+    if (value->IsIntegralConst(0) || (value->gtVNPair.GetConservative() == ValueNumStore::VNForNull()))
     {
-#ifdef DEBUG
-        if (verbose)
-        {
-            (vnBased) ? printf("\nVN based non-heap prop in " FMT_BB ":\n", compCurBB->bbNum)
-                      : printf("\nNon-heap prop for index #%02u in " FMT_BB ":\n", index, compCurBB->bbNum);
-            gtDispTree(indir, nullptr, nullptr, true);
-        }
-#endif
-        indir->gtFlags |= GTF_IND_TGT_NOT_HEAP;
+        // The value being stored is null
+        barrierType = GCInfo::WriteBarrierForm::WBF_NoBarrier;
+    }
+    else if (value->IsIconHandle(GTF_ICON_OBJ_HDL) || vnStore->IsVNObjHandle(value->gtVNPair.GetConservative()))
+    {
+        // The value being stored is a handle
+        barrierType = GCInfo::WriteBarrierForm::WBF_NoBarrier;
+    }
+    else if ((indir->gtFlags & GTF_IND_TGT_HEAP) == 0)
+    {
+        // Next, analyze the address if we haven't already determined the barrier type from the value
+        //
+        // NOTE: we might want to inspect indirs with GTF_IND_TGT_HEAP flag as well - what if we can prove
+        // that they actually need no barrier? But that comes with a TP regression.
+        barrierType = GetWriteBarrierForm(vnStore, addr->gtVNPair.GetConservative());
+    }
 
+    JITDUMP("Trying to determine the exact type of write barrier for STOREIND [%d06]: ", dspTreeID(indir));
+    if (barrierType == GCInfo::WriteBarrierForm::WBF_NoBarrier)
+    {
+        JITDUMP("is not needed at all.\n");
+        indir->gtFlags |= GTF_IND_TGT_NOT_HEAP;
+        return true;
+    }
+    if (barrierType == GCInfo::WriteBarrierForm::WBF_BarrierUnchecked)
+    {
+        JITDUMP("unchecked is fine.\n");
+        indir->gtFlags |= GTF_IND_TGT_HEAP;
         return true;
     }
 
+    JITDUMP("unknown (checked).\n");
     return false;
 }
 

src/coreclr/jit/compiler.h

@@ -7784,8 +7784,6 @@ class Compiler
     AssertionIndex optAssertionIsNonNullInternal(GenTree* op, ASSERT_VALARG_TP assertions DEBUGARG(bool* pVnBased));
     bool optAssertionIsNonNull(GenTree*         op,
                                ASSERT_VALARG_TP assertions DEBUGARG(bool* pVnBased) DEBUGARG(AssertionIndex* pIndex));
-    bool optAssertionIsNonHeap(GenTree*         op,
-                               ASSERT_VALARG_TP assertions DEBUGARG(bool* pVnBased) DEBUGARG(AssertionIndex* pIndex));
 
     AssertionIndex optGlobalAssertionIsEqualOrNotEqual(ASSERT_VALARG_TP assertions, GenTree* op1, GenTree* op2);
     AssertionIndex optGlobalAssertionIsEqualOrNotEqualZero(ASSERT_VALARG_TP assertions, GenTree* op1);
@@ -7822,7 +7820,7 @@ class Compiler
     GenTree* optAssertionProp_Update(GenTree* newTree, GenTree* tree, Statement* stmt);
     GenTree* optNonNullAssertionProp_Call(ASSERT_VALARG_TP assertions, GenTreeCall* call);
     bool optNonNullAssertionProp_Ind(ASSERT_VALARG_TP assertions, GenTree* indir);
-    bool optNonHeapAssertionProp_Ind(ASSERT_VALARG_TP assertions, GenTree* indir);
+    bool optWriteBarrierAssertionProp_StoreInd(ASSERT_VALARG_TP assertions, GenTreeStoreInd* indir);
 
     // Implied assertion functions.
     void optImpliedAssertions(AssertionIndex assertionIndex, ASSERT_TP& activeAssertions);