JIT: build pred lists before patchpoint expansion

src/coreclr/jit/compiler.cpp

@@ -1910,8 +1910,9 @@ void Compiler::compInit(ArenaAllocator*       pAlloc,
     compGeneratingProlog = false;
     compGeneratingEpilog = false;
 
-    compLSRADone       = false;
-    compRationalIRForm = false;
+    compPostImportationCleanupDone = false;
+    compLSRADone                   = false;
+    compRationalIRForm             = false;
 
 #ifdef DEBUG
     compCodeGenDone        = false;
@@ -4431,19 +4432,6 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         DoPhase(this, PHASE_IBCINSTR, &Compiler::fgInstrumentMethod);
     }
 
-    // Expand any patchpoints
-    //
-    DoPhase(this, PHASE_PATCHPOINTS, &Compiler::fgTransformPatchpoints);
-
-    // Transform indirect calls that require control flow expansion.
-    //
-    DoPhase(this, PHASE_INDXCALL, &Compiler::fgTransformIndirectCalls);
-
-    // Cleanup un-imported BBs, cleanup un-imported or
-    // partially imported try regions, add OSR step blocks.
-    //
-    DoPhase(this, PHASE_POST_IMPORT, &Compiler::fgPostImportationCleanup);
-
     // Compute bbNum, bbRefs and bbPreds
     //
     // This is the first time full (not cheap) preds will be computed.
@@ -4461,6 +4449,19 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     };
     DoPhase(this, PHASE_COMPUTE_PREDS, computePredsPhase);
 
+    // Expand any patchpoints
+    //
+    DoPhase(this, PHASE_PATCHPOINTS, &Compiler::fgTransformPatchpoints);
+
+    // Transform indirect calls that require control flow expansion.
+    //
+    DoPhase(this, PHASE_INDXCALL, &Compiler::fgTransformIndirectCalls);
+
+    // Cleanup un-imported BBs, cleanup un-imported or
+    // partially imported try regions, add OSR step blocks.
+    //
+    DoPhase(this, PHASE_POST_IMPORT, &Compiler::fgPostImportationCleanup);
+
     // If we're importing for inlining, we're done.
     if (compIsForInlining())
     {

src/coreclr/jit/compiler.h

@@ -9106,6 +9106,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     bool fgLocalVarLivenessChanged;
     bool fgIsDoingEarlyLiveness;
     bool fgDidEarlyLiveness;
+    bool compPostImportationCleanupDone;
     bool compLSRADone;
     bool compRationalIRForm;
 

src/coreclr/jit/fgdiagnostic.cpp

@@ -2578,6 +2578,14 @@ bool BBPredsChecker::CheckEhTryDsc(BasicBlock* block, BasicBlock* blockPred, EHb
         return true;
     }
 
+    // If this is an OSR method and we haven't run post-importation cleanup, we may see a branch
+    // from fgFirstBB to the middle of a try. Those get fixed during cleanup. Tolerate.
+    //
+    if (comp->opts.IsOSR() && !comp->compPostImportationCleanupDone && (blockPred == comp->fgFirstBB))
+    {
+        return true;
+    }
+
     printf("Jump into the middle of try region: " FMT_BB " branches to " FMT_BB "\n", blockPred->bbNum, block->bbNum);
     assert(!"Jump into middle of try region");
     return false;
@@ -2655,6 +2663,15 @@ bool BBPredsChecker::CheckJump(BasicBlock* blockPred, BasicBlock* block)
             assert(!"SWITCH in the predecessor list with no jump label to BLOCK!");
             break;
 
+        case BBJ_LEAVE:
+            // We may see BBJ_LEAVE preds in unimported blocks if we haven't done cleanup yet.
+            if (!comp->compPostImportationCleanupDone && ((blockPred->bbFlags & BBF_IMPORTED) == 0))
+            {
+                return true;
+            }
+            assert(!"Unexpected BBJ_LEAVE predecessor");
+            break;
+
         default:
             assert(!"Unexpected bbJumpKind");
             break;
@@ -2820,16 +2837,22 @@ void Compiler::fgDebugCheckBBlist(bool checkBBNum /* = false */, bool checkBBRef
         // If the block is a BBJ_COND, a BBJ_SWITCH or a
         // lowered GT_SWITCH_TABLE node then make sure it
         // ends with a conditional jump or a GT_SWITCH
-
-        if (block->bbJumpKind == BBJ_COND)
-        {
-            assert((!allNodesLinked || (block->lastNode()->gtNext == nullptr)) &&
-                   block->lastNode()->OperIsConditionalJump());
-        }
-        else if (block->bbJumpKind == BBJ_SWITCH)
+        //
+        // This may not be true for unimported blocks, if
+        // we haven't run post-importation cleanup yet.
+        //
+        if (compPostImportationCleanupDone || ((block->bbFlags & BBF_IMPORTED) != 0))
         {
-            assert((!allNodesLinked || (block->lastNode()->gtNext == nullptr)) &&
-                   (block->lastNode()->gtOper == GT_SWITCH || block->lastNode()->gtOper == GT_SWITCH_TABLE));
+            if (block->bbJumpKind == BBJ_COND)
+            {
+                assert((!allNodesLinked || (block->lastNode()->gtNext == nullptr)) &&
+                       block->lastNode()->OperIsConditionalJump());
+            }
+            else if (block->bbJumpKind == BBJ_SWITCH)
+            {
+                assert((!allNodesLinked || (block->lastNode()->gtNext == nullptr)) &&
+                       (block->lastNode()->gtOper == GT_SWITCH || block->lastNode()->gtOper == GT_SWITCH_TABLE));
+            }
         }
 
         if (block->bbCatchTyp == BBCT_FILTER)
@@ -2884,7 +2907,7 @@ void Compiler::fgDebugCheckBBlist(bool checkBBNum /* = false */, bool checkBBRef
 
         // Under OSR, if we also are keeping the original method entry around,
         // mark that as implicitly referenced as well.
-        if (opts.IsOSR() && (block == fgEntryBB))
+        if (opts.IsOSR() && (block == fgEntryBB) && ((block->bbFlags & BBF_IMPORTED) != 0))
         {
             blockRefs += 1;
         }

src/coreclr/jit/fgopt.cpp

@@ -1452,6 +1452,17 @@ PhaseStatus Compiler::fgPostImportationCleanup()
             {
                 JITDUMP(FMT_BB " was not imported, marking as removed (%d)\n", cur->bbNum, removedBlks);
 
+                // Notify all successors that cur is no longer a pred.
+                //
+                // This may not be necessary once we have pred lists built before importation.
+                // When we alter flow in the importer branch opts, we should be able to make
+                // suitable updates there for blocks that we plan to keep.
+                //
+                for (BasicBlock* succ : cur->Succs(this))
+                {
+                    fgRemoveAllRefPreds(succ, cur);
+                }
+
                 cur->bbFlags |= BBF_REMOVED;
                 removedBlks++;
 
@@ -1592,6 +1603,7 @@ PhaseStatus Compiler::fgPostImportationCleanup()
                         // here as the oldTryEntry is no longer in the main bb list.
                         newTryEntry = bbNewBasicBlock(BBJ_NONE);
                         newTryEntry->bbFlags |= (BBF_IMPORTED | BBF_INTERNAL);
+                        newTryEntry->bbRefs = 0;
 
                         // Set the right EH region indices on this new block.
                         //
@@ -1611,6 +1623,10 @@ PhaseStatus Compiler::fgPostImportationCleanup()
                         {
                             newTryEntry->bbJumpKind = BBJ_THROW;
                         }
+                        else
+                        {
+                            fgAddRefPred(newTryEntry->bbNext, newTryEntry);
+                        }
 
                         JITDUMP("OSR: changing start of try region #%u from " FMT_BB " to new " FMT_BB "\n",
                                 XTnum + delCnt, oldTryEntry->bbNum, newTryEntry->bbNum);
@@ -1846,6 +1862,11 @@ PhaseStatus Compiler::fgPostImportationCleanup()
     //
     const bool madeChanges = madeFlowChanges || addedTemps;
 
+    // Note that we have now run post importation cleanup,
+    // so we can enable more stringent checking.
+    //
+    compPostImportationCleanupDone = true;
+
     return madeChanges ? PhaseStatus::MODIFIED_EVERYTHING : PhaseStatus::MODIFIED_NOTHING;
 }
 

src/coreclr/jit/indirectcalltransformer.cpp

@@ -279,9 +279,27 @@ class IndirectCallTransformer
         //
         virtual void ChainFlow()
         {
-            assert(!compiler->fgComputePredsDone);
+            assert(compiler->fgComputePredsDone);
+
+            // currBlock
+            compiler->fgRemoveRefPred(remainderBlock, currBlock);
+
+            if (checkBlock != currBlock)
+            {
+                compiler->fgAddRefPred(checkBlock, currBlock);
+            }
+
+            // checkBlock
             checkBlock->bbJumpDest = elseBlock;
-            thenBlock->bbJumpDest  = remainderBlock;
+            compiler->fgAddRefPred(elseBlock, checkBlock);
+            compiler->fgAddRefPred(thenBlock, checkBlock);
+
+            // thenBlock
+            thenBlock->bbJumpDest = remainderBlock;
+            compiler->fgAddRefPred(remainderBlock, thenBlock);
+
+            // elseBlock
+            compiler->fgAddRefPred(remainderBlock, elseBlock);
         }
 
         Compiler*    compiler;
@@ -747,7 +765,6 @@ class IndirectCallTransformer
         {
             thenBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, checkBlock);
             thenBlock->bbFlags |= currBlock->bbFlags & BBF_SPLIT_GAINED;
-
             InlineCandidateInfo* inlineInfo = origCall->gtInlineCandidateInfo;
             CORINFO_CLASS_HANDLE clsHnd     = inlineInfo->guardedClassHandle;
 
@@ -1002,8 +1019,10 @@ class IndirectCallTransformer
             // Finally, rewire the cold block to jump to the else block,
             // not fall through to the check block.
             //
+            compiler->fgRemoveRefPred(checkBlock, coldBlock);
             coldBlock->bbJumpKind = BBJ_ALWAYS;
             coldBlock->bbJumpDest = elseBlock;
+            compiler->fgAddRefPred(elseBlock, coldBlock);
         }
 
         // When the current candidate hads sufficiently high likelihood, scan
@@ -1358,10 +1377,28 @@ class IndirectCallTransformer
         //
         virtual void ChainFlow() override
         {
-            assert(!compiler->fgComputePredsDone);
-            checkBlock->bbJumpDest  = elseBlock;
+            assert(compiler->fgComputePredsDone);
+
+            // currBlock
+            compiler->fgRemoveRefPred(remainderBlock, currBlock);
+            compiler->fgAddRefPred(checkBlock, currBlock);
+
+            // checkBlock
+            checkBlock->bbJumpDest = elseBlock;
+            compiler->fgAddRefPred(elseBlock, checkBlock);
+            compiler->fgAddRefPred(checkBlock2, checkBlock);
+
+            // checkBlock2
             checkBlock2->bbJumpDest = elseBlock;
-            thenBlock->bbJumpDest   = remainderBlock;
+            compiler->fgAddRefPred(elseBlock, checkBlock2);
+            compiler->fgAddRefPred(thenBlock, checkBlock2);
+
+            // thenBlock
+            thenBlock->bbJumpDest = remainderBlock;
+            compiler->fgAddRefPred(remainderBlock, thenBlock);
+
+            // elseBlock
+            compiler->fgAddRefPred(remainderBlock, elseBlock);
         }
 
     private:

src/coreclr/jit/patchpoint.cpp

@@ -147,9 +147,13 @@ class PatchpointTransformer
         // Update flow and flags
         block->bbJumpKind = BBJ_COND;
         block->bbJumpDest = remainderBlock;
-        helperBlock->bbFlags |= BBF_BACKWARD_JUMP;
         block->bbFlags |= BBF_INTERNAL;
 
+        helperBlock->bbFlags |= BBF_BACKWARD_JUMP;
+
+        compiler->fgAddRefPred(helperBlock, block);
+        compiler->fgAddRefPred(remainderBlock, helperBlock);
+
         // Update weights
         remainderBlock->inheritWeight(block);
         helperBlock->inheritWeightPercentage(block, 100 - HIGH_PROBABILITY);