JIT: enable complementary jtrue assertions for cross-block local ap

AndyAyersMS · AndyAyersMS · commit 6f5287b61bf7 · 2023-11-14T15:39:11.000-08:00
Now that we can propagate assertions across block boundaries we can generate assertions for true/false branch conditions and propagate them along the corresponding edges. Contributes to dotnet#93246.
diff --git a/src/coreclr/jit/assertionprop.cpp b/src/coreclr/jit/assertionprop.cpp
@@ -590,12 +590,19 @@ void Compiler::optAssertionInit(bool isLocalProp)
         //
         optAssertionDep =
             new (this, CMK_AssertionProp) JitExpandArray<ASSERT_TP>(getAllocator(CMK_AssertionProp), max(1, lvaCount));
+
+        if (optCrossBlockLocalAssertionProp)
+        {
+            optComplementaryAssertionMap = new (this, CMK_AssertionProp)
+                AssertionIndex[optMaxAssertionCount + 1](); // zero-inited (NO_ASSERTION_INDEX)
+        }
     }
     else
     {
         // General assertion prop.
         //
-        optLocalAssertionProp = false;
+        optLocalAssertionProp           = false;
+        optCrossBlockLocalAssertionProp = 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
@@ -615,7 +622,6 @@ void Compiler::optAssertionInit(bool isLocalProp)
     }
 
     optAssertionTabPrivate = new (this, CMK_AssertionProp) AssertionDsc[optMaxAssertionCount];
-
     optAssertionTraitsInit(optMaxAssertionCount);
 
     optAssertionCount      = 0;
@@ -1641,7 +1647,8 @@ AssertionIndex Compiler::optAddAssertion(AssertionDsc* newAssertion)
     //
     if (optLocalAssertionProp)
     {
-        assert(newAssertion->op1.kind == O1K_LCLVAR);
+        assert((newAssertion->op1.kind == O1K_LCLVAR) || (newAssertion->op1.kind == O1K_SUBTYPE) ||
+               (newAssertion->op1.kind == O1K_EXACT_TYPE));
 
         unsigned        lclNum = newAssertion->op1.lcl.lclNum;
         BitVecOps::Iter iter(apTraits, GetAssertionDep(lclNum));
@@ -1702,7 +1709,8 @@ AssertionIndex Compiler::optAddAssertion(AssertionDsc* newAssertion)
     // Assertion mask bits are [index + 1].
     if (optLocalAssertionProp)
     {
-        assert(newAssertion->op1.kind == O1K_LCLVAR);
+        assert((newAssertion->op1.kind == O1K_LCLVAR) || (newAssertion->op1.kind == O1K_SUBTYPE) ||
+               (newAssertion->op1.kind == O1K_EXACT_TYPE));
 
         // Mark the variables this index depends on
         unsigned lclNum = newAssertion->op1.lcl.lclNum;
@@ -1971,6 +1979,13 @@ AssertionIndex Compiler::optCreateJtrueAssertions(GenTree*                   op1
 
 AssertionInfo Compiler::optCreateJTrueBoundsAssertion(GenTree* tree)
 {
+    // These assertions are VN based, so not relevant for local prop
+    //
+    if (optLocalAssertionProp)
+    {
+        return NO_ASSERTION_INDEX;
+    }
+
     GenTree* relop = tree->gtGetOp1();
     if (!relop->OperIsCompare())
     {
@@ -2138,13 +2153,7 @@ AssertionInfo Compiler::optCreateJTrueBoundsAssertion(GenTree* tree)
  */
 AssertionInfo Compiler::optAssertionGenJtrue(GenTree* tree)
 {
-    // Only create assertions for JTRUE when we are in the global phase
-    if (optLocalAssertionProp)
-    {
-        return NO_ASSERTION_INDEX;
-    }
-
-    GenTree* relop = tree->AsOp()->gtOp1;
+    GenTree* const relop = tree->AsOp()->gtOp1;
     if (!relop->OperIsCompare())
     {
         return NO_ASSERTION_INDEX;
@@ -2158,6 +2167,11 @@ AssertionInfo Compiler::optAssertionGenJtrue(GenTree* tree)
         return info;
     }
 
+    if (optLocalAssertionProp && !optCrossBlockLocalAssertionProp)
+    {
+        return NO_ASSERTION_INDEX;
+    }
+
     // Find assertion kind.
     switch (relop->gtOper)
     {
@@ -2185,53 +2199,57 @@ AssertionInfo Compiler::optAssertionGenJtrue(GenTree* tree)
         std::swap(op1, op2);
     }
 
-    ValueNum op1VN = vnStore->VNConservativeNormalValue(op1->gtVNPair);
-    ValueNum op2VN = vnStore->VNConservativeNormalValue(op2->gtVNPair);
     // If op1 is lcl and op2 is const or lcl, create assertion.
     if ((op1->gtOper == GT_LCL_VAR) && (op2->OperIsConst() || (op2->gtOper == GT_LCL_VAR))) // Fix for Dev10 851483
     {
         return optCreateJtrueAssertions(op1, op2, assertionKind);
     }
-    else if (vnStore->IsVNCheckedBound(op1VN) && vnStore->IsVNInt32Constant(op2VN))
+    else if (!optLocalAssertionProp)
     {
-        assert(relop->OperIs(GT_EQ, GT_NE));
+        ValueNum op1VN = vnStore->VNConservativeNormalValue(op1->gtVNPair);
+        ValueNum op2VN = vnStore->VNConservativeNormalValue(op2->gtVNPair);
 
-        int con = vnStore->ConstantValue<int>(op2VN);
-        if (con >= 0)
+        if (vnStore->IsVNCheckedBound(op1VN) && vnStore->IsVNInt32Constant(op2VN))
         {
-            AssertionDsc dsc;
+            assert(relop->OperIs(GT_EQ, GT_NE));
 
-            // For arr.Length != 0, we know that 0 is a valid index
-            // For arr.Length == con, we know that con - 1 is the greatest valid index
-            if (con == 0)
+            int con = vnStore->ConstantValue<int>(op2VN);
+            if (con >= 0)
             {
-                dsc.assertionKind = OAK_NOT_EQUAL;
-                dsc.op1.bnd.vnIdx = vnStore->VNForIntCon(0);
-            }
-            else
-            {
-                dsc.assertionKind = OAK_EQUAL;
-                dsc.op1.bnd.vnIdx = vnStore->VNForIntCon(con - 1);
-            }
+                AssertionDsc dsc;
 
-            dsc.op1.vn         = op1VN;
-            dsc.op1.kind       = O1K_ARR_BND;
-            dsc.op1.bnd.vnLen  = op1VN;
-            dsc.op2.vn         = vnStore->VNConservativeNormalValue(op2->gtVNPair);
-            dsc.op2.kind       = O2K_CONST_INT;
-            dsc.op2.u1.iconVal = 0;
-            dsc.op2.SetIconFlag(GTF_EMPTY);
-
-            // when con is not zero, create an assertion on the arr.Length == con edge
-            // when con is zero, create an assertion on the arr.Length != 0 edge
-            AssertionIndex index = optAddAssertion(&dsc);
-            if (relop->OperIs(GT_NE) != (con == 0))
-            {
-                return AssertionInfo::ForNextEdge(index);
-            }
-            else
-            {
-                return index;
+                // For arr.Length != 0, we know that 0 is a valid index
+                // For arr.Length == con, we know that con - 1 is the greatest valid index
+                if (con == 0)
+                {
+                    dsc.assertionKind = OAK_NOT_EQUAL;
+                    dsc.op1.bnd.vnIdx = vnStore->VNForIntCon(0);
+                }
+                else
+                {
+                    dsc.assertionKind = OAK_EQUAL;
+                    dsc.op1.bnd.vnIdx = vnStore->VNForIntCon(con - 1);
+                }
+
+                dsc.op1.vn         = op1VN;
+                dsc.op1.kind       = O1K_ARR_BND;
+                dsc.op1.bnd.vnLen  = op1VN;
+                dsc.op2.vn         = vnStore->VNConservativeNormalValue(op2->gtVNPair);
+                dsc.op2.kind       = O2K_CONST_INT;
+                dsc.op2.u1.iconVal = 0;
+                dsc.op2.SetIconFlag(GTF_EMPTY);
+
+                // when con is not zero, create an assertion on the arr.Length == con edge
+                // when con is zero, create an assertion on the arr.Length != 0 edge
+                AssertionIndex index = optAddAssertion(&dsc);
+                if (relop->OperIs(GT_NE) != (con == 0))
+                {
+                    return AssertionInfo::ForNextEdge(index);
+                }
+                else
+                {
+                    return index;
+                }
             }
         }
     }
@@ -2260,7 +2278,7 @@ AssertionInfo Compiler::optAssertionGenJtrue(GenTree* tree)
         return NO_ASSERTION_INDEX;
     }
 
-    GenTreeCall* call = op1->AsCall();
+    GenTreeCall* const call = op1->AsCall();
 
     // Note CORINFO_HELP_READYTORUN_ISINSTANCEOF does not have the same argument pattern.
     // In particular, it is not possible to deduce what class is being tested from its args.
diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
@@ -7356,6 +7356,7 @@ class Compiler
     BitVecTraits* apTraits;
     ASSERT_TP     apFull;
     ASSERT_TP     apLocal;
+    ASSERT_TP     apLocalIfTrue;
 
     enum optAssertionKind
     {
diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp
@@ -12903,22 +12903,21 @@ void Compiler::fgAssertionGen(GenTree* tree)
     INDEBUG(unsigned oldAssertionCount = optAssertionCount;);
     optAssertionGen(tree);
 
-    if (tree->GeneratesAssertion())
-    {
-        AssertionIndex const apIndex = tree->GetAssertionInfo().GetAssertionIndex();
-        unsigned const       bvIndex = apIndex - 1;
-
+    // Helper to note when an existing assertion has been
+    // brought back to life.
+    //
+    auto announce = [&](AssertionIndex apIndex, const char* condition) {
 #ifdef DEBUG
         if (verbose)
         {
             if (oldAssertionCount == optAssertionCount)
             {
-                if (!BitVecOps::IsMember(apTraits, apLocal, bvIndex))
+                if (!BitVecOps::IsMember(apTraits, apLocal, apIndex - 1))
                 {
                     // This tree resurrected an existing assertion.
                     // We call that out here since assertion prop won't.
                     //
-                    printf("GenTreeNode creates assertion:\n");
+                    printf("GenTreeNode creates %sassertion:\n", condition);
                     gtDispTree(tree, nullptr, nullptr, true);
                     printf("In " FMT_BB " New Local ", compCurBB->bbNum);
                     optPrintAssertion(optGetAssertion(apIndex), apIndex);
@@ -12935,7 +12934,69 @@ void Compiler::fgAssertionGen(GenTree* tree)
             }
         }
 #endif
+    };
+
+    // For BBJ_COND nodes, we have two assertion out BVs.
+    // apLocal will be stored on bbAssertionOut and be used for false successors.
+    // apLocalIfTrue will be stored on bbAssertionGen and be used for true successors.
+    //
+    const bool doCondUpdates = tree->OperIs(GT_JTRUE) && compCurBB->KindIs(BBJ_COND) && (compCurBB->NumSucc() == 2);
+
+    // Intialize apLocalIfTrue if we might look for it later,
+    // even if it ends up identical to apLocal.
+    //
+    if (doCondUpdates)
+    {
+        apLocalIfTrue = BitVecOps::MakeCopy(apTraits, apLocal);
+    }
+
+    if (!tree->GeneratesAssertion())
+    {
+        return;
+    }
+
+    AssertionInfo info = tree->GetAssertionInfo();
+
+    if (doCondUpdates)
+    {
+        // Update apLocal and apIfTrue with suitable assertions
+        // from the JTRUE
+        //
+        assert(optCrossBlockLocalAssertionProp);
+
+        AssertionIndex ifFalseAssertionIndex;
+        AssertionIndex ifTrueAssertionIndex;
+
+        if (info.IsNextEdgeAssertion())
+        {
+            ifFalseAssertionIndex = info.GetAssertionIndex();
+            ifTrueAssertionIndex  = optFindComplementary(ifFalseAssertionIndex);
+        }
+        else
+        {
+            ifTrueAssertionIndex  = info.GetAssertionIndex();
+            ifFalseAssertionIndex = optFindComplementary(ifTrueAssertionIndex);
+        }
+
+        if (ifTrueAssertionIndex != NO_ASSERTION_INDEX)
+        {
+            announce(ifTrueAssertionIndex, " [if true]");
+            unsigned const bvIndex = ifTrueAssertionIndex - 1;
+            BitVecOps::AddElemD(apTraits, apLocalIfTrue, bvIndex);
+        }
 
+        if (ifFalseAssertionIndex != NO_ASSERTION_INDEX)
+        {
+            announce(ifFalseAssertionIndex, " [if false]");
+            unsigned const bvIndex = ifFalseAssertionIndex - 1;
+            BitVecOps::AddElemD(apTraits, apLocal, ifFalseAssertionIndex - 1);
+        }
+    }
+    else
+    {
+        AssertionIndex const apIndex = tree->GetAssertionInfo().GetAssertionIndex();
+        announce(apIndex, "");
+        unsigned const bvIndex = apIndex - 1;
         BitVecOps::AddElemD(apTraits, apLocal, bvIndex);
     }
 }
@@ -13833,17 +13894,35 @@ void Compiler::fgMorphBlock(BasicBlock* block, unsigned highestReachablePostorde
                         continue;
                     }
 
-                    // Yes, pred assertions are available. If this is the first pred, copy.
+                    // Yes, pred assertions are available.
+                    // If the pred is (a non-degenerate) BBJ_COND, fetch the appropriate out set.
+                    //
+                    ASSERT_TP assertionsOut = pred->bbAssertionOut;
+
+                    if (pred->KindIs(BBJ_COND) && (pred->NumSucc() == 2) && (block == pred->GetJumpDest()))
+                    {
+                        JITDUMP("Using `if true` assertions from pred " FMT_BB "\n", pred->bbNum);
+                        assertionsOut = pred->bbAssertionGen;
+                    }
+
+                    // If this is the first pred, copy (or share, when block is the only successor).
                     // If this is a subsequent pred, intersect.
                     //
                     if (!hasPredAssertions)
                     {
-                        apLocal           = BitVecOps::MakeCopy(apTraits, pred->bbAssertionOut);
+                        if (block->NumSucc() == 1)
+                        {
+                            apLocal = assertionsOut;
+                        }
+                        else
+                        {
+                            apLocal = BitVecOps::MakeCopy(apTraits, assertionsOut);
+                        }
                         hasPredAssertions = true;
                     }
                     else
                     {
-                        BitVecOps::IntersectionD(apTraits, apLocal, pred->bbAssertionOut);
+                        BitVecOps::IntersectionD(apTraits, apLocal, assertionsOut);
                     }
                 }
 
@@ -13885,6 +13964,13 @@ void Compiler::fgMorphBlock(BasicBlock* block, unsigned highestReachablePostorde
     {
         assert(optLocalAssertionProp);
         block->bbAssertionOut = BitVecOps::MakeCopy(apTraits, apLocal);
+
+        if (block->KindIs(BBJ_COND))
+        {
+            // We don't need to make a copy here as this BV
+            // was freshly copied in fgAssertionGen
+            block->bbAssertionGen = apLocalIfTrue;
+        }
     }
 
     compCurBB = nullptr;

Original file line number	Diff line number	Diff line change
`@@ -7356,6 +7356,7 @@ class Compiler`
`7356`	`7356`	`BitVecTraits* apTraits;`
`7357`	`7357`	`ASSERT_TP apFull;`
`7358`	`7358`	`ASSERT_TP apLocal;`
	`7359`	`+ ASSERT_TP apLocalIfTrue;`
`7359`	`7360`
`7360`	`7361`	`enum optAssertionKind`
`7361`	`7362`	`{`