Fix ARM GCStress hole with byref write barrier helper

src/jit/codegen.h

@@ -269,7 +269,7 @@ class CodeGen : public CodeGenInterface
     void genEmitHelperCall(unsigned helper, int argSize, emitAttr retSize);
 #endif
 
-    void genGCWriteBarrier(GenTreePtr tree, GCInfo::WriteBarrierForm wbf);
+    void genGCWriteBarrier(GenTreePtr tgt, GCInfo::WriteBarrierForm wbf);
 
     BasicBlock* genCreateTempLabel();
 

src/jit/codegenarm.cpp

@@ -873,8 +873,10 @@ void CodeGen::genCodeForCpObj(GenTreeObj* cpObjNode)
 #endif // DEBUG
 
     // Consume the operands and get them into the right registers.
-    // They may now contain gc pointers; genConsumeBlockOp will take care of that.
+    // They may now contain gc pointers (depending on their type; gcMarkRegPtrVal will "do the right thing").
     genConsumeBlockOp(cpObjNode, REG_WRITE_BARRIER_DST_BYREF, REG_WRITE_BARRIER_SRC_BYREF, REG_NA);
+    gcInfo.gcMarkRegPtrVal(REG_WRITE_BARRIER_SRC_BYREF, srcAddrType);
+    gcInfo.gcMarkRegPtrVal(REG_WRITE_BARRIER_DST_BYREF, dstAddr->TypeGet());
 
     // Temp register used to perform the sequence of loads and stores.
     regNumber tmpReg = cpObjNode->ExtractTempReg();

src/jit/codegencommon.cpp

@@ -600,8 +600,11 @@ void CodeGenInterface::genUpdateRegLife(const LclVarDsc* varDsc, bool isBorn, bo
     }
 }
 
-// Gets a register mask that represent the kill set for a helper call since
-// not all JIT Helper calls follow the standard ABI on the target architecture.
+//----------------------------------------------------------------------
+// compNoGCHelperCallKillSet:
+//
+// Gets a register mask that represents the kill set for a helper call.
+// Not all JIT Helper calls follow the standard ABI on the target architecture.
 //
 // TODO-CQ: Currently this list is incomplete (not all helpers calls are
 //          enumerated) and not 100% accurate (some killsets are bigger than
@@ -624,19 +627,24 @@ void CodeGenInterface::genUpdateRegLife(const LclVarDsc* varDsc, bool isBorn, bo
 //
 //         The interim solution is to only add known helper calls that don't
 //         follow the AMD64 ABI and actually trash registers that are supposed to be non-volatile.
+//
+// Arguments:
+//   helper - The helper being inquired about
+//
+// Return Value:
+//   Mask of register kills -- registers whose value is no longer guaranteed to be the same.
+//
 regMaskTP Compiler::compHelperCallKillSet(CorInfoHelpFunc helper)
 {
     switch (helper)
     {
         case CORINFO_HELP_ASSIGN_BYREF:
 #if defined(_TARGET_AMD64_)
-            return RBM_RSI | RBM_RDI | RBM_CALLEE_TRASH;
-#elif defined(_TARGET_ARM64_)
+            return RBM_RSI | RBM_RDI | RBM_CALLEE_TRASH_NOGC;
+#elif defined(_TARGET_ARMARCH_)
             return RBM_WRITE_BARRIER_SRC_BYREF | RBM_WRITE_BARRIER_DST_BYREF | RBM_CALLEE_TRASH_NOGC;
 #elif defined(_TARGET_X86_)
             return RBM_ESI | RBM_EDI | RBM_ECX;
-#elif defined(_TARGET_ARM_)
-            return RBM_ARG_1 | RBM_ARG_0 | RBM_CALLEE_TRASH_NOGC;
 #else
             NYI("Model kill set for CORINFO_HELP_ASSIGN_BYREF on target arch");
             return RBM_CALLEE_TRASH;
@@ -663,6 +671,29 @@ regMaskTP Compiler::compHelperCallKillSet(CorInfoHelpFunc helper)
             NYI("Model kill set for CORINFO_HELP_PROF_FCN_TAILCALL on target arch");
 #endif
 
+#ifdef _TARGET_X86_
+        case CORINFO_HELP_ASSIGN_REF_EAX:
+        case CORINFO_HELP_ASSIGN_REF_ECX:
+        case CORINFO_HELP_ASSIGN_REF_EBX:
+        case CORINFO_HELP_ASSIGN_REF_EBP:
+        case CORINFO_HELP_ASSIGN_REF_ESI:
+        case CORINFO_HELP_ASSIGN_REF_EDI:
+
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_EAX:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_ECX:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_EBX:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_EBP:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_ESI:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF_EDI:
+            return RBM_EDX;
+
+#ifdef FEATURE_USE_ASM_GC_WRITE_BARRIERS
+        case CORINFO_HELP_ASSIGN_REF:
+        case CORINFO_HELP_CHECKED_ASSIGN_REF:
+            return RBM_EAX | RBM_EDX;
+#endif // FEATURE_USE_ASM_GC_WRITE_BARRIERS
+#endif
+
         case CORINFO_HELP_STOP_FOR_GC:
             return RBM_STOP_FOR_GC_TRASH;
 
@@ -674,19 +705,30 @@ regMaskTP Compiler::compHelperCallKillSet(CorInfoHelpFunc helper)
     }
 }
 
+//----------------------------------------------------------------------
+// compNoGCHelperCallKillSet: Gets a register mask that represents the set of registers that no longer
+// contain GC or byref pointers, for "NO GC" helper calls. This is used by the emitter when determining
+// what registers to remove from the current live GC/byref sets (and thus what to report as dead in the
+// GC info). Note that for the CORINFO_HELP_ASSIGN_BYREF helper, in particular, the kill set reported by
+// compHelperCallKillSet() doesn't match this kill set. compHelperCallKillSet() reports the dst/src
+// address registers as killed for liveness purposes, since their values change. However, they still are
+// valid byref pointers after the call, so the dst/src address registers are NOT reported as killed here.
 //
-// Gets a register mask that represents the kill set for "NO GC" helper calls since
-// not all JIT Helper calls follow the standard ABI on the target architecture.
+// Note: This list may not be complete and defaults to the default RBM_CALLEE_TRASH_NOGC registers.
 //
-// Note: This list may not be complete and defaults to the default NOGC registers.
+// Arguments:
+//   helper - The helper being inquired about
+//
+// Return Value:
+//   Mask of GC register kills
 //
 regMaskTP Compiler::compNoGCHelperCallKillSet(CorInfoHelpFunc helper)
 {
     assert(emitter::emitNoGChelper(helper));
 
     switch (helper)
     {
-#if defined(_TARGET_AMD64_) || defined(_TARGET_X86_)
+#if defined(_TARGET_XARCH_)
         case CORINFO_HELP_PROF_FCN_ENTER:
             return RBM_PROFILER_ENTER_TRASH;
 
@@ -695,18 +737,15 @@ regMaskTP Compiler::compNoGCHelperCallKillSet(CorInfoHelpFunc helper)
 
         case CORINFO_HELP_PROF_FCN_TAILCALL:
             return RBM_PROFILER_TAILCALL_TRASH;
-#endif // defined(_TARGET_AMD64_) || defined(_TARGET_X86_)
+#endif // defined(_TARGET_XARCH_)
 
         case CORINFO_HELP_ASSIGN_BYREF:
 #if defined(_TARGET_AMD64_)
-            // this helper doesn't trash RSI and RDI
-            return RBM_CALLEE_TRASH_NOGC & ~(RBM_RSI | RBM_RDI);
+            return RBM_CALLEE_TRASH_NOGC;
 #elif defined(_TARGET_X86_)
             // This helper only trashes ECX.
             return RBM_ECX;
-#elif defined(_TARGET_ARM64_)
-            return RBM_CALLEE_TRASH_NOGC & ~(RBM_WRITE_BARRIER_SRC_BYREF | RBM_WRITE_BARRIER_DST_BYREF);
-#else
+#elif defined(_TARGET_ARMARCH_)
             return RBM_CALLEE_TRASH_NOGC;
 #endif // defined(_TARGET_AMD64_)
 
@@ -3912,16 +3951,86 @@ void CodeGen::genReportEH()
     assert(XTnum == EHCount);
 }
 
-void CodeGen::genGCWriteBarrier(GenTreePtr tgt, GCInfo::WriteBarrierForm wbf)
+//----------------------------------------------------------------------
+// genUseOptimizedWriteBarriers: Determine if an optimized write barrier
+// helper should be used.
+//
+// Arguments:
+//   wbf - The WriteBarrierForm of the write (GT_STOREIND) that is happening.
+//
+// Return Value:
+//   true if an optimized write barrier helper should be used, false otherwise.
+//   Note: only x86 implements (register-specific source) optimized write
+//   barriers currently).
+//
+bool CodeGenInterface::genUseOptimizedWriteBarriers(GCInfo::WriteBarrierForm wbf)
+{
+#if defined(_TARGET_X86_) && NOGC_WRITE_BARRIERS
+#ifdef DEBUG
+    return (wbf != GCInfo::WBF_NoBarrier_CheckNotHeapInDebug); // This one is always a call to a C++ method.
+#else
+    return true;
+#endif
+#else
+    return false;
+#endif
+}
+
+//----------------------------------------------------------------------
+// genUseOptimizedWriteBarriers: Determine if an optimized write barrier
+// helper should be used.
+//
+// This has the same functionality as the version of
+// genUseOptimizedWriteBarriers that takes a WriteBarrierForm, but avoids
+// determining what the required write barrier form is, if possible.
+//
+// Arguments:
+//   tgt - target tree of write (e.g., GT_STOREIND)
+//   assignVal - tree with value to write
+//
+// Return Value:
+//   true if an optimized write barrier helper should be used, false otherwise.
+//   Note: only x86 implements (register-specific source) optimized write
+//   barriers currently).
+//
+bool CodeGenInterface::genUseOptimizedWriteBarriers(GenTree* tgt, GenTree* assignVal)
+{
+#if defined(_TARGET_X86_) && NOGC_WRITE_BARRIERS
+#ifdef DEBUG
+    GCInfo::WriteBarrierForm wbf = compiler->codeGen->gcInfo.gcIsWriteBarrierCandidate(tgt, assignVal);
+    return (wbf != GCInfo::WBF_NoBarrier_CheckNotHeapInDebug); // This one is always a call to a C++ method.
+#else
+    return true;
+#endif
+#else
+    return false;
+#endif
+}
+
+//----------------------------------------------------------------------
+// genWriteBarrierHelperForWriteBarrierForm: Given a write node requiring a write
+// barrier, and the write barrier form required, determine the helper to call.
+//
+// Arguments:
+//   tgt - target tree of write (e.g., GT_STOREIND)
+//   wbf - already computed write barrier form to use
+//
+// Return Value:
+//   Write barrier helper to use.
+//
+// Note: do not call this function to get an optimized write barrier helper (e.g.,
+// for x86).
+//
+CorInfoHelpFunc CodeGenInterface::genWriteBarrierHelperForWriteBarrierForm(GenTree* tgt, GCInfo::WriteBarrierForm wbf)
 {
 #ifndef LEGACY_BACKEND
     noway_assert(tgt->gtOper == GT_STOREIND);
 #else  // LEGACY_BACKEND
     noway_assert(tgt->gtOper == GT_IND || tgt->gtOper == GT_CLS_VAR); // enforced by gcIsWriteBarrierCandidate
 #endif // LEGACY_BACKEND
 
-    /* Call the proper vm helper */
-    int helper = CORINFO_HELP_ASSIGN_REF;
+    CorInfoHelpFunc helper = CORINFO_HELP_ASSIGN_REF;
+
 #ifdef DEBUG
     if (wbf == GCInfo::WBF_NoBarrier_CheckNotHeapInDebug)
     {
@@ -3949,6 +4058,20 @@ void CodeGen::genGCWriteBarrier(GenTreePtr tgt, GCInfo::WriteBarrierForm wbf)
            ((helper == CORINFO_HELP_ASSIGN_REF) &&
             (wbf == GCInfo::WBF_BarrierUnchecked || wbf == GCInfo::WBF_BarrierUnknown)));
 
+    return helper;
+}
+
+//----------------------------------------------------------------------
+// genGCWriteBarrier: Generate a write barrier for a node.
+//
+// Arguments:
+//   tgt - target tree of write (e.g., GT_STOREIND)
+//   wbf - already computed write barrier form to use
+//
+void CodeGen::genGCWriteBarrier(GenTreePtr tgt, GCInfo::WriteBarrierForm wbf)
+{
+    CorInfoHelpFunc helper = genWriteBarrierHelperForWriteBarrierForm(tgt, wbf);
+
 #ifdef FEATURE_COUNT_GC_WRITE_BARRIERS
     // We classify the "tgt" trees as follows:
     // If "tgt" is of the form (where [ x ] indicates an optional x, and { x1, ..., xn } means "one of the x_i forms"):

src/jit/codegeninterface.h

@@ -151,6 +151,11 @@ class CodeGenInterface
     regMaskTP genLiveMask(VARSET_VALARG_TP liveSet);
 #endif
 
+public:
+    bool genUseOptimizedWriteBarriers(GCInfo::WriteBarrierForm wbf);
+    bool genUseOptimizedWriteBarriers(GenTree* tgt, GenTree* assignVal);
+    CorInfoHelpFunc genWriteBarrierHelperForWriteBarrierForm(GenTree* tgt, GCInfo::WriteBarrierForm wbf);
+
     // The following property indicates whether the current method sets up
     // an explicit stack frame or not.
 private:

src/jit/codegenxarch.cpp

@@ -4987,30 +4987,34 @@ bool CodeGen::genEmitOptimizedGCWriteBarrier(GCInfo::WriteBarrierForm writeBarri
     assert(writeBarrierForm != GCInfo::WBF_NoBarrier);
 
 #if defined(_TARGET_X86_) && NOGC_WRITE_BARRIERS
-    bool useOptimizedWriteBarriers = true;
-
-#ifdef DEBUG
-    useOptimizedWriteBarriers =
-        (writeBarrierForm != GCInfo::WBF_NoBarrier_CheckNotHeapInDebug); // This one is always a call to a C++ method.
-#endif
-
-    if (!useOptimizedWriteBarriers)
+    if (!genUseOptimizedWriteBarriers(writeBarrierForm))
     {
         return false;
     }
 
     const static int regToHelper[2][8] = {
         // If the target is known to be in managed memory
         {
-            CORINFO_HELP_ASSIGN_REF_EAX, CORINFO_HELP_ASSIGN_REF_ECX, -1, CORINFO_HELP_ASSIGN_REF_EBX, -1,
-            CORINFO_HELP_ASSIGN_REF_EBP, CORINFO_HELP_ASSIGN_REF_ESI, CORINFO_HELP_ASSIGN_REF_EDI,
+            CORINFO_HELP_ASSIGN_REF_EAX, // EAX
+            CORINFO_HELP_ASSIGN_REF_ECX, // ECX
+            -1,                          // EDX (always the target address)
+            CORINFO_HELP_ASSIGN_REF_EBX, // EBX
+            -1,                          // ESP
+            CORINFO_HELP_ASSIGN_REF_EBP, // EBP
+            CORINFO_HELP_ASSIGN_REF_ESI, // ESI
+            CORINFO_HELP_ASSIGN_REF_EDI, // EDI
         },
 
         // Don't know if the target is in managed memory
         {
-            CORINFO_HELP_CHECKED_ASSIGN_REF_EAX, CORINFO_HELP_CHECKED_ASSIGN_REF_ECX, -1,
-            CORINFO_HELP_CHECKED_ASSIGN_REF_EBX, -1, CORINFO_HELP_CHECKED_ASSIGN_REF_EBP,
-            CORINFO_HELP_CHECKED_ASSIGN_REF_ESI, CORINFO_HELP_CHECKED_ASSIGN_REF_EDI,
+            CORINFO_HELP_CHECKED_ASSIGN_REF_EAX, // EAX
+            CORINFO_HELP_CHECKED_ASSIGN_REF_ECX, // ECX
+            -1,                                  // EDX (always the target address)
+            CORINFO_HELP_CHECKED_ASSIGN_REF_EBX, // EBX
+            -1,                                  // ESP
+            CORINFO_HELP_CHECKED_ASSIGN_REF_EBP, // EBP
+            CORINFO_HELP_CHECKED_ASSIGN_REF_ESI, // ESI
+            CORINFO_HELP_CHECKED_ASSIGN_REF_EDI, // EDI
         },
     };
 

src/jit/emitarm.cpp

@@ -4476,6 +4476,16 @@ void emitter::emitIns_Call(EmitCallType          callType,
         {
             savedSet |= RBM_PROFILER_RET_SCRATCH;
         }
+
+#ifdef DEBUG
+        if (emitComp->verbose)
+        {
+            printf("NOGC Call: savedSet=");
+            printRegMaskInt(savedSet);
+            emitDispRegSet(savedSet);
+            printf("\n");
+        }
+#endif
     }
     else
     {

src/jit/lsra.cpp

@@ -2786,6 +2786,47 @@ void LinearScan::addRefsForPhysRegMask(regMaskTP mask, LsraLocation currentLoc,
     }
 }
 
+//------------------------------------------------------------------------
+// getKillSetForStoreInd: Determine the liveness kill set for a GT_STOREIND node.
+// If the GT_STOREIND will generate a write barrier, determine the specific kill
+// set required by the case-specific, platform-specific write barrier. If no
+// write barrier is required, the kill set will be RBM_NONE.
+//
+// Arguments:
+//    tree - the GT_STOREIND node
+//
+// Return Value:    a register mask of the registers killed
+//
+regMaskTP LinearScan::getKillSetForStoreInd(GenTreeStoreInd* tree)
+{
+    assert(tree->OperIs(GT_STOREIND));
+
+    regMaskTP killMask = RBM_NONE;
+
+    GenTree* data = tree->Data();
+
+    GCInfo::WriteBarrierForm writeBarrierForm = compiler->codeGen->gcInfo.gcIsWriteBarrierCandidate(tree, data);
+    if (writeBarrierForm != GCInfo::WBF_NoBarrier)
+    {
+        if (compiler->codeGen->genUseOptimizedWriteBarriers(writeBarrierForm))
+        {
+            // We can't determine the exact helper to be used at this point, because it depends on
+            // the allocated register for the `data` operand. However, all the (x86) optimized
+            // helpers have the same kill set: EDX.
+            killMask = RBM_CALLEE_TRASH_NOGC;
+        }
+        else
+        {
+            // Figure out which helper we're going to use, and then get the kill set for that helper.
+            CorInfoHelpFunc helper =
+                compiler->codeGen->genWriteBarrierHelperForWriteBarrierForm(tree, writeBarrierForm);
+            killMask = compiler->compHelperCallKillSet(helper);
+        }
+    }
+
+    return killMask;
+}
+
 //------------------------------------------------------------------------
 // getKillSetForNode:   Return the registers killed by the given tree node.
 //
@@ -2936,10 +2977,7 @@ regMaskTP LinearScan::getKillSetForNode(GenTree* tree)
             }
             break;
         case GT_STOREIND:
-            if (compiler->codeGen->gcInfo.gcIsWriteBarrierAsgNode(tree))
-            {
-                killMask = RBM_CALLEE_TRASH_NOGC;
-            }
+            killMask = getKillSetForStoreInd(tree->AsStoreInd());
             break;
 
 #if defined(PROFILING_SUPPORTED)