From c5866851402a30c90dc0ce968d895d8eee88cf66 Mon Sep 17 00:00:00 2001
From: Will Smith <lol.tihan@gmail.com>
Date: Thu, 16 Mar 2023 16:59:39 -0700
Subject: [PATCH] [JIT] ARM64 - Optimize `x < 0` and `x >= 0` (#83176)

* Initial work that allows optimizing x < 0

* Few tweaks

* Optimized x < 0 and x >= 0 for conditionals

* Removed old commented code. Switched tbz and tbnz

* Simplified

* Fix build

* Check for possible interference

* Reverse

* Check invariant-range

* Cleanup

* Cleanup

* Handle new JCMP flags in gtReverseCond

* Comment update

* Remove extra flags

* Simplify further

* Fix warning

* Update src/coreclr/jit/lower.cpp

Co-authored-by: Jakob Botsch Nielsen <Jakob.botsch.nielsen@gmail.com>

* Fix warning

* Fixing assert

---------

Co-authored-by: Jakob Botsch Nielsen <Jakob.botsch.nielsen@gmail.com>
---
 src/coreclr/jit/codegenarm64.cpp | 15 +++++++++++++--
 src/coreclr/jit/gentree.cpp      |  3 +--
 src/coreclr/jit/lower.cpp        |  8 ++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/coreclr/jit/codegenarm64.cpp b/src/coreclr/jit/codegenarm64.cpp
index 24e46849dec..505a2cc2409 100644
--- a/src/coreclr/jit/codegenarm64.cpp
+++ b/src/coreclr/jit/codegenarm64.cpp
@@ -4532,7 +4532,18 @@ void CodeGen::genCodeForCompare(GenTreeOp* tree)
         if (op2->isContainedIntOrIImmed())
         {
             GenTreeIntConCommon* intConst = op2->AsIntConCommon();
-            emit->emitIns_R_I(ins, cmpSize, op1->GetRegNum(), intConst->IconValue());
+
+            regNumber op1Reg = op1->GetRegNum();
+
+            if (compiler->opts.OptimizationEnabled() && (ins == INS_cmp) && (targetReg != REG_NA) &&
+                tree->OperIs(GT_LT) && intConst->IsIntegralConst(0) && ((cmpSize == EA_4BYTE) || (cmpSize == EA_8BYTE)))
+            {
+                emit->emitIns_R_R_I(INS_lsr, cmpSize, targetReg, op1Reg, (int)cmpSize * 8 - 1);
+                genProduceReg(tree);
+                return;
+            }
+
+            emit->emitIns_R_I(ins, cmpSize, op1Reg, intConst->IconValue());
         }
         else
         {
@@ -4725,7 +4736,7 @@ void CodeGen::genCodeForJumpCompare(GenTreeOp* tree)
     {
         ssize_t compareImm = op2->AsIntCon()->IconValue();
 
-        assert(isPow2(compareImm));
+        assert(isPow2(((size_t)compareImm)));
 
         instruction ins = (tree->gtFlags & GTF_JCMP_EQ) ? INS_tbz : INS_tbnz;
         int         imm = genLog2((size_t)compareImm);
diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp
index baef29856f3..f811a72061f 100644
--- a/src/coreclr/jit/gentree.cpp
+++ b/src/coreclr/jit/gentree.cpp
@@ -3481,8 +3481,7 @@ GenTree* Compiler::gtReverseCond(GenTree* tree)
     else if (tree->OperIs(GT_JCMP))
     {
         // Flip the GTF_JCMP_EQ
-        //
-        // This causes switching
+        // On ARM64, this causes switching
         //     cbz <=> cbnz
         //     tbz <=> tbnz
         tree->gtFlags ^= GTF_JCMP_EQ;
diff --git a/src/coreclr/jit/lower.cpp b/src/coreclr/jit/lower.cpp
index 614ca071c7e..6894843d8f7 100644
--- a/src/coreclr/jit/lower.cpp
+++ b/src/coreclr/jit/lower.cpp
@@ -3382,6 +3382,14 @@ GenTree* Lowering::LowerJTrue(GenTreeOp* jtrue)
             flags   = cond->OperIs(GT_EQ) ? GTF_JCMP_EQ : GTF_EMPTY;
             useJCMP = true;
         }
+        else if (cond->OperIs(GT_LT, GT_GE) && !cond->IsUnsigned() && relopOp2->IsIntegralConst(0))
+        {
+            // Codegen will use tbnz or tbz in codegen which do not affect the flag register
+            flags   = GTF_JCMP_TST | (cond->OperIs(GT_LT) ? GTF_EMPTY : GTF_JCMP_EQ);
+            useJCMP = true;
+            relopOp2->AsIntConCommon()->SetIntegralValue(
+                (static_cast<INT64>(1) << (8 * genTypeSize(genActualType(relopOp1)) - 1)));
+        }
         else if (cond->OperIs(GT_TEST_EQ, GT_TEST_NE) && isPow2(relopOp2->AsIntCon()->IconValue()))
         {
             // Codegen will use tbz or tbnz in codegen which do not affect the flag register
-- 
2.34.1