JIT: cross-block local assertion prop in morph

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,85 @@ 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;
+
+        // See if we should disable cross-block local prop
+        //
+        if (optCrossBlockLocalAssertionProp)
+        {
+            // Disable via config
+            //
+            if (JitConfig.JitDoCrossBlockLocalAssertionProp() == 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 know lvaCount is less than the tracked limit.
+            // 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 table will be reset for each block, so can be smaller.
+            //
+            optMaxAssertionCount = 64;
+        }
+
+        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;
@@ -4379,12 +4428,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(JitDoCrossBlockLocalAssertionProp, W("JitDoCrossBlockLocalAssertionProp"), 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

@@ -13780,10 +13780,80 @@ 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)
+            {
+                bool hasPred = false;
+
+                for (BasicBlock* const pred : block->PredBlocks())
+                {
+                    // A higher postorder number means the block appears earlier in
+                    // the postorder. Use this to detect if a pred's assertion info is available.
+                    //
+                    if (pred->bbPostorderNum > block->bbPostorderNum)
+                    {
+                        // Yes, available. If this is the first pred, copy.
+                        // If this is a subsequent pred, intersect.
+                        //
+                        if (!hasPred)
+                        {
+                            apLocal = BitVecOps::MakeCopy(apTraits, pred->bbAssertionOut);
+                            hasPred = true;
+                        }
+                        else
+                        {
+                            BitVecOps::IntersectionD(apTraits, apLocal, pred->bbAssertionOut);
+                        }
+
+                        continue;
+                    }
+
+                    // No, not available.
+                    //
+                    JITDUMP(FMT_BB " pred " FMT_BB " not processed; clearing assertions in\n", block->bbNum,
+                            pred->bbNum);
+                    canUsePredAssertions = false;
+                }
+
+                // If there were no preds, there are no asserions in.
+                //
+                if (!hasPred)
+                {
+                    JITDUMP(FMT_BB " has no preds, so no assertions in\n", block->bbNum);
+                    canUsePredAssertions = false;
+                }
+
+                if (canUsePredAssertions)
+                {
+                    JITDUMPEXEC(optDumpAssertionIndices("Assertions in: ", apLocal));
+                }
+            }
+            else
+            {
+                JITDUMP(FMT_BB " ineligible for cross-block; clearing assertions in\n", block->bbNum);
+            }
+
+            if (!canUsePredAssertions)
+            {
+                apLocal = BitVecOps::MakeEmpty(apTraits);
+            }
+        }
     }
 
     // Make the current basic block address available globally.
@@ -13801,6 +13871,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;
 }
 
@@ -13822,7 +13900,8 @@ PhaseStatus Compiler::fgMorphBlocks()
 
     // Local assertion prop is enabled if we are optimized
     //
-    optLocalAssertionProp = opts.OptimizationEnabled();
+    optLocalAssertionProp           = opts.OptimizationEnabled();
+    optCrossBlockLocalAssertionProp = optLocalAssertionProp;
 
     if (optLocalAssertionProp)
     {

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" JitDoCrossBlockLocalAssertionProp="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" />