Simplify JIT label handling (#51208)

* Simplify JIT label handling

Remove the BBF_JMP_TARGET flag that was set early and attempted
to be maintained all through compilation. Instead, use
BBF_USE_LABEL to indicate to the emitter where we need labels.

Also, stop setting and maintaining BBF_USE_LABEL early.

Add a pass over the blocks when preparing for codegen that sets
most of the necessary BBF_USE_LABEL flags. This flag will never
be set before codegen. A few places set the flag after codegen starts,
namely `genCreateTempLabel` and BBJ_COND handling for alignment.
Note that this flag must be set before the block is processed for
codegen (and an insGroup is created).

Together, these changes make it easier to work with the flow graph
without worrying about maintaining these bits of information through
various optimizations.

Add a few more details about alignment processing to the dump.

There are a few code asm diffs due to alignment processing not previously
creating a label to help compute how large a loop is.

There are a lot of textual asm diffs due to there being (mostly)
fewer labels, plus some additional insGroup output. This can happen if
a block was labeled with `BBF_JMP_TARGET` or `BBF_USE_LABEL` before,
but didn't need to be, perhaps after some optimizations. Now, the flag is
never added in the first place.

There are a large number of GC info diffs. Labels are where GC info state
changes are recorded between codegen and the emitter. If we eliminate an
unnecessary emitter label, then we also eliminate a capture of the current
codegen GC state. Since the emitter is lazy at marking GC deaths, this
means that we see a lot of lengthened GC lifetimes -- until the next
label, or some other cause of GC kill. Often, you see a register kill
followed by register birth just disappear, and the register is maintained
alive across the interim.

* Remove loop align flag if we decide a loop is no longer a loop

src/coreclr/jit/block.cpp

@@ -419,10 +419,6 @@ void BasicBlock::dspFlags()
     {
         printf("label ");
     }
-    if (bbFlags & BBF_JMP_TARGET)
-    {
-        printf("target ");
-    }
     if (bbFlags & BBF_HAS_JMP)
     {
         printf("jmp ");

src/coreclr/jit/block.h

@@ -411,7 +411,7 @@ struct BasicBlock : private LIR::Range
 #define BBF_LOOP_CALL1          MAKE_BBFLAG(15) // BB starts a loop that will always     call
 
 #define BBF_HAS_LABEL           MAKE_BBFLAG(16) // BB needs a label
-#define BBF_JMP_TARGET          MAKE_BBFLAG(17) // BB is a target of an implicit/explicit jump
+// Unused                       MAKE_BBFLAG(17)
 #define BBF_HAS_JMP             MAKE_BBFLAG(18) // BB executes a JMP instruction (instead of return)
 #define BBF_GC_SAFE_POINT       MAKE_BBFLAG(19) // BB has a GC safe point (a call).  More abstractly, BB does not require a
                                                 // (further) poll -- this may be because this BB has a call, or, in some
@@ -501,9 +501,8 @@ struct BasicBlock : private LIR::Range
 // TODO: Should BBF_RUN_RARELY be added to BBF_SPLIT_GAINED ?
 
 #define BBF_SPLIT_GAINED                                                                                               \
-    (BBF_DONT_REMOVE | BBF_HAS_LABEL | BBF_HAS_JMP | BBF_BACKWARD_JUMP | BBF_HAS_IDX_LEN | BBF_HAS_NEWARRAY |          \
-     BBF_PROF_WEIGHT | BBF_HAS_NEWOBJ | BBF_KEEP_BBJ_ALWAYS | BBF_CLONED_FINALLY_END | BBF_HAS_NULLCHECK |             \
-     BBF_HAS_CLASS_PROFILE)
+    (BBF_DONT_REMOVE | BBF_HAS_JMP | BBF_BACKWARD_JUMP | BBF_HAS_IDX_LEN | BBF_HAS_NEWARRAY | BBF_PROF_WEIGHT |        \
+     BBF_HAS_NEWOBJ | BBF_KEEP_BBJ_ALWAYS | BBF_CLONED_FINALLY_END | BBF_HAS_NULLCHECK | BBF_HAS_CLASS_PROFILE)
 
 #ifndef __GNUC__ // GCC doesn't like C_ASSERT at global scope
     static_assert_no_msg((BBF_SPLIT_NONEXIST & BBF_SPLIT_LOST) == 0);

src/coreclr/jit/codegen.h

@@ -88,7 +88,7 @@ class CodeGen final : public CodeGenInterface
 
     void genPrepForCompiler();
 
-    void genPrepForEHCodegen();
+    void genMarkLabelsForCodegen();
 
     inline RegState* regStateForType(var_types t)
     {

src/coreclr/jit/codegenarm.cpp

@@ -129,7 +129,6 @@ BasicBlock* CodeGen::genCallFinally(BasicBlock* block)
     assert(block->bbNext->bbJumpDest->bbFlags & BBF_FINALLY_TARGET);
 
     bbFinallyRet = block->bbNext->bbJumpDest;
-    bbFinallyRet->bbFlags |= BBF_JMP_TARGET;
 
     // Load the address where the finally funclet should return into LR.
     // The funclet prolog/epilog will do "push {lr}" / "pop {pc}" to do the return.
@@ -633,7 +632,7 @@ void CodeGen::genJumpTable(GenTree* treeNode)
     for (unsigned i = 0; i < jumpCount; i++)
     {
         BasicBlock* target = *jumpTable++;
-        noway_assert(target->bbFlags & BBF_JMP_TARGET);
+        noway_assert(target->bbFlags & BBF_HAS_LABEL);
 
         JITDUMP("            DD      L_M%03u_" FMT_BB "\n", compiler->compMethodID, target->bbNum);
 

src/coreclr/jit/codegenarm64.cpp

@@ -2743,7 +2743,6 @@ void CodeGen::genTableBasedSwitch(GenTree* treeNode)
     GetEmitter()->emitIns_R_R_R(INS_ldr, EA_4BYTE, baseReg, baseReg, idxReg, INS_OPTS_LSL);
 
     // add it to the absolute address of fgFirstBB
-    compiler->fgFirstBB->bbFlags |= BBF_JMP_TARGET;
     GetEmitter()->emitIns_R_L(INS_adr, EA_PTRSIZE, compiler->fgFirstBB, tmpReg);
     GetEmitter()->emitIns_R_R_R(INS_add, EA_PTRSIZE, baseReg, baseReg, tmpReg);
 
@@ -2771,7 +2770,7 @@ void CodeGen::genJumpTable(GenTree* treeNode)
     for (unsigned i = 0; i < jumpCount; i++)
     {
         BasicBlock* target = *jumpTable++;
-        noway_assert(target->bbFlags & BBF_JMP_TARGET);
+        noway_assert(target->bbFlags & BBF_HAS_LABEL);
 
         JITDUMP("            DD      L_M%03u_" FMT_BB "\n", compiler->compMethodID, target->bbNum);
 

src/coreclr/jit/codegencommon.cpp

@@ -318,84 +318,165 @@ void CodeGen::genPrepForCompiler()
 #endif
 }
 
-/*****************************************************************************
- *  To report exception handling information to the VM, we need the size of the exception
- *  handling regions. To compute that, we need to emit labels for the beginning block of
- *  an EH region, and the block that immediately follows a region. Go through the EH
- *  table and mark all these blocks with BBF_HAS_LABEL to make this happen.
- *
- *  The beginning blocks of the EH regions already should have this flag set.
- *
- *  No blocks should be added or removed after this.
- *
- *  This code is closely couple with genReportEH() in the sense that any block
- *  that this procedure has determined it needs to have a label has to be selected
- *  using the same logic both here and in genReportEH(), so basically any time there is
- *  a change in the way we handle EH reporting, we have to keep the logic of these two
- *  methods 'in sync'.
- */
-
-void CodeGen::genPrepForEHCodegen()
+//------------------------------------------------------------------------
+// genMarkLabelsForCodegen: Mark labels required for codegen.
+//
+// Mark all blocks that require a label with BBF_HAS_LABEL. These are either blocks that are:
+// 1. the target of jumps (fall-through flow doesn't require a label),
+// 2. referenced labels such as for "switch" codegen,
+// 3. needed to denote the range of EH regions to the VM.
+// 4. needed to denote the range of code for alignment processing.
+//
+// No labels will be in the IR before now, but future codegen might annotate additional blocks
+// with this flag, such as "switch" codegen, or codegen-created blocks from genCreateTempLabel().
+// Also, the alignment processing code marks BBJ_COND fall-through labels elsewhere.
+//
+// To report exception handling information to the VM, we need the size of the exception
+// handling regions. To compute that, we need to emit labels for the beginning block of
+// an EH region, and the block that immediately follows a region. Go through the EH
+// table and mark all these blocks with BBF_HAS_LABEL to make this happen.
+//
+// This code is closely couple with genReportEH() in the sense that any block
+// that this procedure has determined it needs to have a label has to be selected
+// using the same logic both here and in genReportEH(), so basically any time there is
+// a change in the way we handle EH reporting, we have to keep the logic of these two
+// methods 'in sync'.
+//
+// No blocks should be added or removed after this.
+//
+void CodeGen::genMarkLabelsForCodegen()
 {
     assert(!compiler->fgSafeBasicBlockCreation);
 
+    JITDUMP("Mark labels for codegen\n");
+
+#ifdef DEBUG
+    // No label flags should be set before this.
+    for (BasicBlock* block = compiler->fgFirstBB; block != nullptr; block = block->bbNext)
+    {
+        assert((block->bbFlags & BBF_HAS_LABEL) == 0);
+    }
+#endif // DEBUG
+
+    // The first block is special; it always needs a label. This is to properly set up GC info.
+    JITDUMP("  " FMT_BB " : first block\n", compiler->fgFirstBB->bbNum);
+    compiler->fgFirstBB->bbFlags |= BBF_HAS_LABEL;
+
+    // The current implementation of switch tables requires the first block to have a label so it
+    // can generate offsets to the switch label targets.
+    // (This is duplicative with the fact we always set the first block with a label above.)
+    // TODO-CQ: remove this when switches have been re-implemented to not use this.
+    if (compiler->fgHasSwitch)
+    {
+        JITDUMP("  " FMT_BB " : function has switch; mark first block\n", compiler->fgFirstBB->bbNum);
+        compiler->fgFirstBB->bbFlags |= BBF_HAS_LABEL;
+    }
+
+    for (BasicBlock* block = compiler->fgFirstBB; block != nullptr; block = block->bbNext)
+    {
+        switch (block->bbJumpKind)
+        {
+            case BBJ_ALWAYS: // This will also handle the BBJ_ALWAYS of a BBJ_CALLFINALLY/BBJ_ALWAYS pair.
+            case BBJ_COND:
+            case BBJ_EHCATCHRET:
+                JITDUMP("  " FMT_BB " : branch target\n", block->bbJumpDest->bbNum);
+                block->bbJumpDest->bbFlags |= BBF_HAS_LABEL;
+                break;
+
+            case BBJ_SWITCH:
+                unsigned jumpCnt;
+                jumpCnt = block->bbJumpSwt->bbsCount;
+                BasicBlock** jumpTab;
+                jumpTab = block->bbJumpSwt->bbsDstTab;
+                do
+                {
+                    JITDUMP("  " FMT_BB " : branch target\n", (*jumpTab)->bbNum);
+                    (*jumpTab)->bbFlags |= BBF_HAS_LABEL;
+                } while (++jumpTab, --jumpCnt);
+                break;
+
+            case BBJ_CALLFINALLY:
+                // The finally target itself will get marked by walking the EH table, below, and marking
+                // all handler begins.
+                CLANG_FORMAT_COMMENT_ANCHOR;
+
+#if FEATURE_EH_CALLFINALLY_THUNKS
+                {
+                    // For callfinally thunks, we need to mark the block following the callfinally/always pair,
+                    // as that's needed for identifying the range of the "duplicate finally" region in EH data.
+                    BasicBlock* bbToLabel = block->bbNext;
+                    if (block->isBBCallAlwaysPair())
+                    {
+                        bbToLabel = bbToLabel->bbNext; // skip the BBJ_ALWAYS
+                    }
+                    if (bbToLabel != nullptr)
+                    {
+                        JITDUMP("  " FMT_BB " : callfinally thunk region end\n", bbToLabel->bbNum);
+                        bbToLabel->bbFlags |= BBF_HAS_LABEL;
+                    }
+                }
+#endif // FEATURE_EH_CALLFINALLY_THUNKS
+
+                break;
+
+            case BBJ_EHFINALLYRET:
+            case BBJ_EHFILTERRET:
+            case BBJ_RETURN:
+            case BBJ_THROW:
+            case BBJ_NONE:
+                break;
+
+            default:
+                noway_assert(!"Unexpected bbJumpKind");
+                break;
+        }
+    }
+
+    // Walk all the exceptional code blocks and mark them, since they don't appear in the normal flow graph.
+    for (Compiler::AddCodeDsc* add = compiler->fgAddCodeList; add; add = add->acdNext)
+    {
+        JITDUMP("  " FMT_BB " : throw helper block\n", add->acdDstBlk->bbNum);
+        add->acdDstBlk->bbFlags |= BBF_HAS_LABEL;
+    }
+
     EHblkDsc* HBtab;
     EHblkDsc* HBtabEnd;
 
-    bool anyFinallys = false;
-
     for (HBtab = compiler->compHndBBtab, HBtabEnd = compiler->compHndBBtab + compiler->compHndBBtabCount;
          HBtab < HBtabEnd; HBtab++)
     {
-        assert(HBtab->ebdTryBeg->bbFlags & BBF_HAS_LABEL);
-        assert(HBtab->ebdHndBeg->bbFlags & BBF_HAS_LABEL);
+        HBtab->ebdTryBeg->bbFlags |= BBF_HAS_LABEL;
+        HBtab->ebdHndBeg->bbFlags |= BBF_HAS_LABEL;
+
+        JITDUMP("  " FMT_BB " : try begin\n", HBtab->ebdTryBeg->bbNum);
+        JITDUMP("  " FMT_BB " : hnd begin\n", HBtab->ebdHndBeg->bbNum);
 
         if (HBtab->ebdTryLast->bbNext != nullptr)
         {
             HBtab->ebdTryLast->bbNext->bbFlags |= BBF_HAS_LABEL;
+            JITDUMP("  " FMT_BB " : try end\n", HBtab->ebdTryLast->bbNext->bbNum);
         }
 
         if (HBtab->ebdHndLast->bbNext != nullptr)
         {
             HBtab->ebdHndLast->bbNext->bbFlags |= BBF_HAS_LABEL;
+            JITDUMP("  " FMT_BB " : hnd end\n", HBtab->ebdHndLast->bbNext->bbNum);
         }
 
         if (HBtab->HasFilter())
         {
-            assert(HBtab->ebdFilter->bbFlags & BBF_HAS_LABEL);
-            // The block after the last block of the filter is
-            // the handler begin block, which we already asserted
-            // has BBF_HAS_LABEL set.
+            HBtab->ebdFilter->bbFlags |= BBF_HAS_LABEL;
+            JITDUMP("  " FMT_BB " : filter begin\n", HBtab->ebdFilter->bbNum);
         }
-
-#if FEATURE_EH_CALLFINALLY_THUNKS
-        if (HBtab->HasFinallyHandler())
-        {
-            anyFinallys = true;
-        }
-#endif // FEATURE_EH_CALLFINALLY_THUNKS
     }
 
-#if FEATURE_EH_CALLFINALLY_THUNKS
-    if (anyFinallys)
+#ifdef DEBUG
+    if (compiler->verbose)
     {
-        for (BasicBlock* block = compiler->fgFirstBB; block != nullptr; block = block->bbNext)
-        {
-            if (block->bbJumpKind == BBJ_CALLFINALLY)
-            {
-                BasicBlock* bbToLabel = block->bbNext;
-                if (block->isBBCallAlwaysPair())
-                {
-                    bbToLabel = bbToLabel->bbNext; // skip the BBJ_ALWAYS
-                }
-                if (bbToLabel != nullptr)
-                {
-                    bbToLabel->bbFlags |= BBF_HAS_LABEL;
-                }
-            } // block is BBJ_CALLFINALLY
-        }     // for each block
-    }         // if (anyFinallys)
-#endif        // FEATURE_EH_CALLFINALLY_THUNKS
+        printf("*************** After genMarkLabelsForCodegen()\n");
+        compiler->fgDispBasicBlocks();
+    }
+#endif // DEBUG
 }
 
 void CodeGenInterface::genUpdateLife(GenTree* tree)
@@ -954,7 +1035,8 @@ BasicBlock* CodeGen::genCreateTempLabel()
     compiler->fgSafeBasicBlockCreation = false;
 #endif
 
-    block->bbFlags |= BBF_JMP_TARGET | BBF_HAS_LABEL;
+    JITDUMP("Mark " FMT_BB " as label: codegen temp block\n", block->bbNum);
+    block->bbFlags |= BBF_HAS_LABEL;
 
     // Use coldness of current block, as this label will
     // be contained in it.
@@ -1067,7 +1149,7 @@ void CodeGen::genAdjustStackLevel(BasicBlock* block)
 
     if (!isFramePointerUsed() && compiler->fgIsThrowHlpBlk(block))
     {
-        noway_assert(block->bbFlags & BBF_JMP_TARGET);
+        noway_assert(block->bbFlags & BBF_HAS_LABEL);
 
         SetStackLevel(compiler->fgThrowHlpBlkStkLevel(block) * sizeof(int));
 
@@ -1979,7 +2061,7 @@ void CodeGen::genInsertNopForUnwinder(BasicBlock* block)
     // calls the funclet during non-exceptional control flow.
     if (block->bbFlags & BBF_FINALLY_TARGET)
     {
-        assert(block->bbFlags & BBF_JMP_TARGET);
+        assert(block->bbFlags & BBF_HAS_LABEL);
 
 #ifdef DEBUG
         if (compiler->verbose)

src/coreclr/jit/codegenlinear.cpp

@@ -93,14 +93,6 @@ void CodeGen::genInitialize()
     initializeVariableLiveKeeper();
 #endif //  USING_VARIABLE_LIVE_RANGE
 
-    // The current implementation of switch tables requires the first block to have a label so it
-    // can generate offsets to the switch label targets.
-    // TODO-CQ: remove this when switches have been re-implemented to not use this.
-    if (compiler->fgHasSwitch)
-    {
-        compiler->fgFirstBB->bbFlags |= BBF_JMP_TARGET;
-    }
-
     genPendingCallLabel = nullptr;
 
     // Initialize the pointer tracking code
@@ -167,8 +159,7 @@ void CodeGen::genCodeForBBlist()
 
 #endif // defined(DEBUG) && defined(TARGET_XARCH)
 
-    // Prepare the blocks for exception handling codegen: mark the blocks that needs labels.
-    genPrepForEHCodegen();
+    genMarkLabelsForCodegen();
 
     assert(!compiler->fgFirstBBScratch ||
            compiler->fgFirstBB == compiler->fgFirstBBScratch); // compiler->fgFirstBBScratch has to be first.
@@ -321,7 +312,7 @@ void CodeGen::genCodeForBBlist()
 
         // If this block is a jump target or it requires a label then set 'needLabel' to true,
         //
-        bool needLabel = (block->bbFlags & (BBF_JMP_TARGET | BBF_HAS_LABEL)) != 0;
+        bool needLabel = (block->bbFlags & BBF_HAS_LABEL) != 0;
 
         if (block == compiler->fgFirstColdBlock)
         {
@@ -758,12 +749,24 @@ void CodeGen::genCodeForBBlist()
                 //
                 // During emitter, this information will be used to calculate the loop size.
                 // Depending on the loop size, decision of whether to align a loop or not will be taken.
+                //
+                // In the emitter, we need to calculate the loop size from `block->bbJumpDest` through
+                // `block` (inclusive). Thus, we need to ensure there is a label on the lexical fall-through
+                // block, even if one is not otherwise needed, to be able to calculate the size of this
+                // loop (loop size is calculated by walking the instruction groups; see emitter::getLoopSize()).
 
                 if (block->bbJumpDest->isLoopAlign())
                 {
                     GetEmitter()->emitSetLoopBackEdge(block->bbJumpDest);
+
+                    if (block->bbNext != nullptr)
+                    {
+                        JITDUMP("Mark " FMT_BB " as label: alignment end-of-loop\n", block->bbNext->bbNum);
+                        block->bbNext->bbFlags |= BBF_HAS_LABEL;
+                    }
                 }
-#endif
+#endif // FEATURE_LOOP_ALIGN
+
                 break;
 
             default:

src/coreclr/jit/codegenxarch.cpp

@@ -3531,7 +3531,6 @@ void CodeGen::genTableBasedSwitch(GenTree* treeNode)
     GetEmitter()->emitIns_R_ARX(INS_mov, EA_4BYTE, baseReg, baseReg, idxReg, 4, 0);
 
     // add it to the absolute address of fgFirstBB
-    compiler->fgFirstBB->bbFlags |= BBF_JMP_TARGET;
     GetEmitter()->emitIns_R_L(INS_lea, EA_PTR_DSP_RELOC, compiler->fgFirstBB, tmpReg);
     GetEmitter()->emitIns_R_R(INS_add, EA_PTRSIZE, baseReg, tmpReg);
     // jmp baseReg
@@ -3558,7 +3557,7 @@ void CodeGen::genJumpTable(GenTree* treeNode)
     for (unsigned i = 0; i < jumpCount; i++)
     {
         BasicBlock* target = *jumpTable++;
-        noway_assert(target->bbFlags & BBF_JMP_TARGET);
+        noway_assert(target->bbFlags & BBF_HAS_LABEL);
 
         JITDUMP("            DD      L_M%03u_" FMT_BB "\n", compiler->compMethodID, target->bbNum);