Improve compare-and-branch sequences produced by Emitter (#111797)

* Improve compare-and-branch sequences produced by Emitter

Introduces an overload of `genJumpToThrowHlpBlk` that allows you to pass
a function that generates the branch code, before it creates an inline
throw block. This allows us to choose compare-and-branch sequences such
as `cbz` on ARM64 for checks added in the emitter, such as divide-by-zero.

* Optimize bounds checks for zero index on ARM64

Emit `cbz` instead of `cmp+b.ls` when checking bounds for an access to the
0 index of an array. It can only throw when `arrayLength == 0`.

Fixes #42514

* Move `genGetThrowHelper` to codegenarm64.cpp

* Use `GetRegNum` instead of `GetReg`

Author: snickolls-arm

src/coreclr/jit/codegen.h

@@ -230,6 +230,61 @@ class CodeGen final : public CodeGenInterface
 
     void genExitCode(BasicBlock* block);
 
+#if defined(TARGET_ARM64)
+    BasicBlock* genGetThrowHelper(SpecialCodeKind codeKind);
+
+    // genEmitInlineThrow: Generate code for an inline exception.
+    void genEmitInlineThrow(SpecialCodeKind codeKind)
+    {
+        genEmitHelperCall(compiler->acdHelper(codeKind), 0, EA_UNKNOWN);
+    }
+
+    // throwCodeFn callback follows concept -> void(*)(BasicBlock* target, bool isInline)
+    //
+    // For conditional jumps:
+    //     If `isInline`, invert the condition for throw and fall into the exception block.
+    //     Otherwise emit compare and jump with the normal throw condition.
+    // For unconditional jumps:
+    //     Only emit the unconditional jump when `isInline == false`.
+    //     When `isInline == true` the code will fallthrough to throw without any jump added.
+    //
+    // Parameter `target` gives a label to jump to, which is the throw block if
+    // `isInline == false`, else the continuation.
+    template <typename throwCodeFn>
+    void genJumpToThrowHlpBlk(SpecialCodeKind codeKind, throwCodeFn emitJumpCode, BasicBlock* throwBlock = nullptr)
+    {
+        if (!throwBlock)
+        {
+            // If caller didn't supply a target block, then try to find a helper block.
+            throwBlock = genGetThrowHelper(codeKind);
+        }
+
+        if (throwBlock)
+        {
+            // check:
+            //   if (checkPassed)
+            //     goto throw;
+            //   ...
+            // throw:
+            //   throw();
+            emitJumpCode(throwBlock, false);
+        }
+        else
+        {
+            // check:
+            //   if (!checkPassed)
+            //     goto continue;
+            //   throw();
+            // continue:
+            //   ...
+            BasicBlock* over = genCreateTempLabel();
+            emitJumpCode(over, true);
+            genEmitInlineThrow(codeKind);
+            genDefineTempLabel(over);
+        }
+    }
+#endif
+
     void genJumpToThrowHlpBlk(emitJumpKind jumpKind, SpecialCodeKind codeKind, BasicBlock* failBlk = nullptr);
 
 #if defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
@@ -1285,6 +1340,8 @@ class CodeGen final : public CodeGenInterface
 #endif
 #if defined(TARGET_ARM64)
     void genCodeForJumpCompare(GenTreeOpCC* tree);
+    void genCompareImmAndJump(
+        GenCondition::Code cond, regNumber reg, ssize_t compareImm, emitAttr size, BasicBlock* target);
     void genCodeForBfiz(GenTreeOp* tree);
 #endif // TARGET_ARM64
 

src/coreclr/jit/codegenarm64.cpp

@@ -3622,9 +3622,10 @@ void CodeGen::genCodeForDivMod(GenTreeOp* tree)
             }
             else
             {
-                // Check if the divisor is zero throw a DivideByZeroException
-                emit->emitIns_R_I(INS_cmp, size, divisorReg, 0);
-                genJumpToThrowHlpBlk(EJ_eq, SCK_DIV_BY_ZERO);
+                genJumpToThrowHlpBlk(SCK_DIV_BY_ZERO, [&](BasicBlock* target, bool invert) {
+                    GenCondition::Code cond = invert ? GenCondition::NE : GenCondition::EQ;
+                    genCompareImmAndJump(cond, divisorReg, 0, size, target);
+                });
             }
         }
 
@@ -5071,6 +5072,34 @@ void CodeGen::genCodeForJumpCompare(GenTreeOpCC* tree)
     }
 }
 
+void CodeGen::genCompareImmAndJump(
+    GenCondition::Code cond, regNumber reg, ssize_t compareImm, emitAttr size, BasicBlock* target)
+{
+    // For ARM64 we only expect equality comparisons.
+    assert((cond == GenCondition::EQ) || (cond == GenCondition::NE));
+
+    if (compareImm == 0)
+    {
+        // We can use cbz/cbnz
+        instruction ins = (cond == GenCondition::EQ) ? INS_cbz : INS_cbnz;
+        GetEmitter()->emitIns_J_R(ins, size, target, reg);
+    }
+    else if (isPow2(compareImm))
+    {
+        // We can use tbz/tbnz
+        instruction ins = (cond == GenCondition::EQ) ? INS_tbz : INS_tbnz;
+        int         imm = genLog2((size_t)compareImm);
+        GetEmitter()->emitIns_J_R_I(ins, size, target, reg, imm);
+    }
+    else
+    {
+        // Emit compare and branch pair default.
+        emitJumpKind jumpKind = (cond == GenCondition::EQ) ? EJ_eq : EJ_ne;
+        GetEmitter()->emitIns_R_I(INS_cmp, size, reg, compareImm);
+        inst_JMP(jumpKind, target);
+    }
+}
+
 //---------------------------------------------------------------------
 // genSPtoFPdelta - return offset from the stack pointer (Initial-SP) to the frame pointer. The frame pointer
 // will point to the saved frame pointer slot (i.e., there will be frame pointer chaining).
@@ -5939,4 +5968,27 @@ insOpts CodeGen::ShiftOpToInsOpts(genTreeOps shiftOp)
     }
 }
 
+//---------------------------------------------------------------------------------
+// genGetThrowHelper: Search for the throw helper for the exception kind `codeKind`
+BasicBlock* CodeGen::genGetThrowHelper(SpecialCodeKind codeKind)
+{
+    BasicBlock* excpRaisingBlock = nullptr;
+    if (compiler->fgUseThrowHelperBlocks())
+    {
+        // For code with throw helper blocks, find and use the helper block for
+        // raising the exception. The block may be shared by other trees too.
+        Compiler::AddCodeDsc* add = compiler->fgFindExcptnTarget(codeKind, compiler->compCurBB);
+        PREFIX_ASSUME_MSG((add != nullptr), ("ERROR: failed to find exception throw block"));
+        assert(add->acdUsed);
+        excpRaisingBlock = add->acdDstBlk;
+#if !FEATURE_FIXED_OUT_ARGS
+        assert(add->acdStkLvlInit || isFramePointerUsed());
+#endif // !FEATURE_FIXED_OUT_ARGS
+
+        noway_assert(excpRaisingBlock != nullptr);
+    }
+
+    return excpRaisingBlock;
+}
+
 #endif // TARGET_ARM64

src/coreclr/jit/codegenarmarch.cpp

@@ -1474,6 +1474,23 @@ void CodeGen::genRangeCheck(GenTree* oper)
         src1    = arrLen;
         src2    = arrIndex;
         jmpKind = EJ_ls;
+
+#if defined(TARGET_ARM64)
+        if (arrIndex->IsIntegralConst(0))
+        {
+            assert(!arrLen->isContained());
+            // For (index == 0), we can just test if (length == 0) as this is the only case that would throw.
+            // This may lead to an optimization by using cbz/tbnz.
+            genJumpToThrowHlpBlk(
+                bndsChk->gtThrowKind,
+                [&](BasicBlock* target, bool isInline) {
+                genCompareImmAndJump(isInline ? GenCondition::NE : GenCondition::EQ, arrLen->GetRegNum(), 0,
+                                     emitActualTypeSize(arrLen), target);
+            },
+                bndsChk->gtIndRngFailBB);
+            return;
+        }
+#endif
     }
     else
     {