JIT: cross-block local assertion prop in morph (#94363)

During global morph, allow assertions to propagate to a block from the block's
predecessors. Handle special cases where we can't allow this to happen:
* block has preds that have not yet been morphed
* block has no preds
* block is specially flagged as one that might gain new preds during morph
* block is an EH handler entry

Contributes to #93246.

When enabled, size the assertion table based on the number of locals, up to the tracked limit.

Disabled by default; use 'DOTNET_JitEnableCrossBlockLocalAssertionProp=1` to enable.

Author: AndyAyersMS

eng/pipelines/common/templates/runtimes/run-test-job.yml

@@ -585,7 +585,7 @@ jobs:
             - jitobjectstackallocation
             - jitphysicalpromotion_only
             - jitphysicalpromotion_full
-
+            - jitcrossblocklocalassertionprop
         ${{ if in(parameters.testGroup, 'jit-cfg') }}:
           scenarios:
           - jitcfg

src/coreclr/jit/assertionprop.cpp

@@ -541,36 +541,83 @@ void Compiler::optAssertionTraitsInit(AssertionIndex assertionCount)
 
 void Compiler::optAssertionInit(bool isLocalProp)
 {
-    // Use a function countFunc to determine a proper maximum assertion count for the
-    // method being compiled. The function is linear to the IL size for small and
-    // moderate methods. For large methods, considering throughput impact, we track no
-    // more than 64 assertions.
-    // Note this tracks at most only 256 assertions.
-    static const AssertionIndex countFunc[] = {64, 128, 256, 64};
-    static const unsigned       lowerBound  = 0;
-    static const unsigned       upperBound  = ArrLen(countFunc) - 1;
-    const unsigned              codeSize    = info.compILCodeSize / 512;
-    optMaxAssertionCount                    = countFunc[isLocalProp ? lowerBound : min(upperBound, codeSize)];
-
-    optLocalAssertionProp  = isLocalProp;
-    optAssertionTabPrivate = new (this, CMK_AssertionProp) AssertionDsc[optMaxAssertionCount];
     assert(NO_ASSERTION_INDEX == 0);
+    const unsigned maxTrackedLocals = (unsigned)JitConfig.JitMaxLocalsToTrack();
 
-    if (!isLocalProp)
+    // We initialize differently for local prop / global prop
+    //
+    if (isLocalProp)
     {
+        optLocalAssertionProp           = true;
+        optCrossBlockLocalAssertionProp = true;
+
+        // Disable via config
+        //
+        if (JitConfig.JitEnableCrossBlockLocalAssertionProp() == 0)
+        {
+            JITDUMP("Disabling cross-block assertion prop by config setting\n");
+            optCrossBlockLocalAssertionProp = false;
+        }
+
+        // Disable if too many locals
+        //
+        // The typical number of local assertions is roughly proportional
+        // to the number of locals. So when we have huge numbers of locals,
+        // just do within-block local assertion prop.
+        //
+        if (lvaCount > maxTrackedLocals)
+        {
+            JITDUMP("Disabling cross-block assertion prop: too many locals\n");
+            optCrossBlockLocalAssertionProp = false;
+        }
+
+        if (optCrossBlockLocalAssertionProp)
+        {
+            // We may need a fairly large table.
+            // Allow for roughly one assertion per local, up to the tracked limit.
+            // (empirical studies show about 0.6 asserions/local)
+            //
+            optMaxAssertionCount = (AssertionIndex)min(maxTrackedLocals, ((lvaCount / 64) + 1) * 64);
+        }
+        else
+        {
+            // The assertion table will be reset for each block, so it can be smaller.
+            //
+            optMaxAssertionCount = 64;
+        }
+
+        // Local assertion prop keeps mappings from each local var to the assertions about that var.
+        //
+        optAssertionDep =
+            new (this, CMK_AssertionProp) JitExpandArray<ASSERT_TP>(getAllocator(CMK_AssertionProp), max(1, lvaCount));
+    }
+    else
+    {
+        // General assertion prop.
+        //
+        optLocalAssertionProp = false;
+
+        // Use a function countFunc to determine a proper maximum assertion count for the
+        // method being compiled. The function is linear to the IL size for small and
+        // moderate methods. For large methods, considering throughput impact, we track no
+        // more than 64 assertions.
+        // Note this tracks at most only 256 assertions.
+        //
+        static const AssertionIndex countFunc[] = {64, 128, 256, 64};
+        static const unsigned       upperBound  = ArrLen(countFunc) - 1;
+        const unsigned              codeSize    = info.compILCodeSize / 512;
+        optMaxAssertionCount                    = countFunc[min(upperBound, codeSize)];
+
         optValueNumToAsserts =
             new (getAllocator(CMK_AssertionProp)) ValueNumToAssertsMap(getAllocator(CMK_AssertionProp));
         optComplementaryAssertionMap = new (this, CMK_AssertionProp)
             AssertionIndex[optMaxAssertionCount + 1](); // zero-inited (NO_ASSERTION_INDEX)
     }
 
-    if (optAssertionDep == nullptr)
-    {
-        optAssertionDep =
-            new (this, CMK_AssertionProp) JitExpandArray<ASSERT_TP>(getAllocator(CMK_AssertionProp), max(1, lvaCount));
-    }
+    optAssertionTabPrivate = new (this, CMK_AssertionProp) AssertionDsc[optMaxAssertionCount];
 
     optAssertionTraitsInit(optMaxAssertionCount);
+
     optAssertionCount      = 0;
     optAssertionOverflow   = 0;
     optAssertionPropagated = false;
@@ -4457,12 +4504,11 @@ AssertionIndex Compiler::optAssertionIsNonNullInternal(GenTree*         op,
         // Find live assertions related to lclNum
         //
         unsigned const lclNum      = op->AsLclVarCommon()->GetLclNum();
-        ASSERT_TP      apDependent = GetAssertionDep(lclNum);
-        BitVecOps::IntersectionD(apTraits, apDependent, apLocal);
+        ASSERT_TP      apDependent = BitVecOps::Intersection(apTraits, GetAssertionDep(lclNum), assertions);
 
         // Scan those looking for a suitable assertion
         //
-        BitVecOps::Iter iter(apTraits, assertions);
+        BitVecOps::Iter iter(apTraits, apDependent);
         unsigned        index = 0;
         while (iter.NextElem(&index))
         {

src/coreclr/jit/compiler.h

@@ -7640,6 +7640,7 @@ class Compiler
     AssertionDsc*  optAssertionTabPrivate;      // table that holds info about value assignments
     AssertionIndex optAssertionCount;           // total number of assertions in the assertion table
     AssertionIndex optMaxAssertionCount;
+    bool           optCrossBlockLocalAssertionProp;
     unsigned       optAssertionOverflow;
     bool           optCanPropLclVar;
     bool           optCanPropEqual;

src/coreclr/jit/jitconfigvalues.h

@@ -653,6 +653,9 @@ CONFIG_INTEGER(JitEnableHeadTailMerge, W("JitEnableHeadTailMerge"), 1)
 // Enable physical promotion
 CONFIG_INTEGER(JitEnablePhysicalPromotion, W("JitEnablePhysicalPromotion"), 1)
 
+// Enable cross-block local assertion prop
+CONFIG_INTEGER(JitEnableCrossBlockLocalAssertionProp, W("JitEnableCrossBlockLocalAssertionProp"), 0)
+
 #if defined(DEBUG)
 // JitFunctionFile: Name of a file that contains a list of functions. If the currently compiled function is in the
 // file, certain other JIT config variables will be active. If the currently compiled function is not in the file,

src/coreclr/jit/morph.cpp

@@ -13779,10 +13779,73 @@ void Compiler::fgMorphBlock(BasicBlock* block)
 
     if (optLocalAssertionProp)
     {
-        // For now, each block starts with an empty table, and no available assertions
-        //
-        optAssertionReset(0);
-        apLocal = BitVecOps::MakeEmpty(apTraits);
+        if (!optCrossBlockLocalAssertionProp)
+        {
+            // Each block starts with an empty table, and no available assertions
+            //
+            optAssertionReset(0);
+            apLocal = BitVecOps::MakeEmpty(apTraits);
+        }
+        else
+        {
+            // Determine if this block can leverage assertions from its pred blocks.
+            //
+            // Some blocks are ineligible.
+            //
+            bool canUsePredAssertions = ((block->bbFlags & BBF_CAN_ADD_PRED) == 0) && !bbIsHandlerBeg(block);
+
+            // Validate all preds have valid info
+            //
+            if (!canUsePredAssertions)
+            {
+                JITDUMP(FMT_BB " ineligible for cross-block\n", block->bbNum);
+            }
+            else
+            {
+                bool hasPredAssertions = false;
+
+                for (BasicBlock* const pred : block->PredBlocks())
+                {
+                    // A smaller pred postorder number means the pred appears later in the postorder.
+                    // An equal number means pred == block (block is a self-loop).
+                    // Either way the assertion info is not available, and we must assume the worst.
+                    //
+                    if (pred->bbPostorderNum <= block->bbPostorderNum)
+                    {
+                        JITDUMP(FMT_BB " pred " FMT_BB " not processed; clearing assertions in\n", block->bbNum,
+                                pred->bbNum);
+                        break;
+                    }
+
+                    // Yes, pred assertions are available. If this is the first pred, copy.
+                    // If this is a subsequent pred, intersect.
+                    //
+                    if (!hasPredAssertions)
+                    {
+                        apLocal           = BitVecOps::MakeCopy(apTraits, pred->bbAssertionOut);
+                        hasPredAssertions = true;
+                    }
+                    else
+                    {
+                        BitVecOps::IntersectionD(apTraits, apLocal, pred->bbAssertionOut);
+                    }
+                }
+
+                if (!hasPredAssertions)
+                {
+                    // Either no preds, or some preds w/o assertions.
+                    //
+                    canUsePredAssertions = false;
+                }
+            }
+
+            if (!canUsePredAssertions)
+            {
+                apLocal = BitVecOps::MakeEmpty(apTraits);
+            }
+
+            JITDUMPEXEC(optDumpAssertionIndices("Assertions in: ", apLocal));
+        }
     }
 
     // Make the current basic block address available globally.
@@ -13800,6 +13863,14 @@ void Compiler::fgMorphBlock(BasicBlock* block)
         }
     }
 
+    // Publish the live out state.
+    //
+    if (optCrossBlockLocalAssertionProp && (block->NumSucc() > 0))
+    {
+        assert(optLocalAssertionProp);
+        block->bbAssertionOut = BitVecOps::MakeCopy(apTraits, apLocal);
+    }
+
     compCurBB = nullptr;
 }
 
@@ -13819,15 +13890,18 @@ PhaseStatus Compiler::fgMorphBlocks()
     //
     fgGlobalMorph = true;
 
-    // Local assertion prop is enabled if we are optimized
-    //
-    optLocalAssertionProp = opts.OptimizationEnabled();
-
-    if (optLocalAssertionProp)
+    if (opts.OptimizationEnabled())
+    {
+        // Local assertion prop is enabled if we are optimizing.
+        //
+        optAssertionInit(/* isLocalProp*/ true);
+    }
+    else
     {
-        // Initialize for local assertion prop
+        // Not optimizing. No assertion prop.
         //
-        optAssertionInit(true);
+        optLocalAssertionProp           = false;
+        optCrossBlockLocalAssertionProp = false;
     }
 
     if (!compEnregLocals())
@@ -13854,10 +13928,6 @@ PhaseStatus Compiler::fgMorphBlocks()
     {
         // If we aren't optimizing, we just morph in normal bbNext order.
         //
-        // Note morph can add blocks downstream from the current block,
-        // and alter (but not null out) the current block's bbNext;
-        // this iterator ensures they all get visited.
-        //
         for (BasicBlock* block : Blocks())
         {
             fgMorphBlock(block);

src/tests/Common/testenvironment.proj

@@ -80,6 +80,7 @@
       RunningIlasmRoundTrip;
       DOTNET_JitSynthesizeCounts;
       DOTNET_JitCheckSynthesizedCounts
+      DOTNET_JitDoCrossBlockLocalAssertionProp
     </DOTNETVariables>
   </PropertyGroup>
   <ItemGroup>
@@ -220,6 +221,7 @@
     <TestEnvironment Include="jitobjectstackallocation" JitObjectStackAllocation="1" TieredCompilation="0" />
     <TestEnvironment Include="jitphysicalpromotion_only" JitStressModeNames="STRESS_NO_OLD_PROMOTION" TieredCompilation="0" />
     <TestEnvironment Include="jitphysicalpromotion_full" JitStressModeNames="STRESS_PHYSICAL_PROMOTION_COST STRESS_NO_OLD_PROMOTION" TieredCompilation="0" />
+    <TestEnvironment Include="jitcrossblocklocalassertionprop" JitEnableCrossBlockLocalAssertionProp="1" TieredCompilation="0" />
     <TestEnvironment Include="jitcfg" JitForceControlFlowGuard="1" />
     <TestEnvironment Include="jitcfg_dispatcher_always" JitForceControlFlowGuard="1" JitCFGUseDispatcher="1" />
     <TestEnvironment Include="jitcfg_dispatcher_never" JitForceControlFlowGuard="1" JitCFGUseDispatcher="0" />