[turbofan] ARM64: Match 64 bit compare with zero and branch

src/compiler/arm64/instruction-selector-arm64.cc

Index: src/compiler/arm64/instruction-selector-arm64.cc
diff --git a/src/compiler/arm64/instruction-selector-arm64.cc b/src/compiler/arm64/instruction-selector-arm64.cc
index 06dce93feebcfc46bb7fbf678018a2a64badfc0c..ad877421990ce42a256672a12e96d7713bd05b01 100644
--- a/src/compiler/arm64/instruction-selector-arm64.cc
+++ b/src/compiler/arm64/instruction-selector-arm64.cc
@@ -1860,6 +1860,23 @@ void VisitWord64Test(InstructionSelector* selector, Node* node,
   VisitWordTest(selector, node, kArm64Tst, cont);
 }
 
+template <typename Matcher, ArchOpcode kOpcode>
+bool TryEmitTestAndBranch(InstructionSelector* selector, Node* node,
+                          FlagsContinuation* cont) {
+  Arm64OperandGenerator g(selector);
+  Matcher m(node);
+  if (cont->IsBranch() && m.right().HasValue() &&
+      (base::bits::CountPopulation(m.right().Value()) == 1)) {
+    // If the mask has only one bit set, we can use tbz/tbnz.
+    DCHECK((cont->condition() == kEqual) || (cont->condition() == kNotEqual));
+    selector->Emit(
+        cont->Encode(kOpcode), g.NoOutput(), g.UseRegister(m.left().node()),
+        g.TempImmediate(base::bits::CountTrailingZeros(m.right().Value())),
+        g.Label(cont->true_block()), g.Label(cont->false_block()));
+    return true;
+  }
+  return false;
+}
 
 // Shared routine for multiple float32 compare operations.
 void VisitFloat32Compare(InstructionSelector* selector, Node* node,
@@ -1904,6 +1921,8 @@ void VisitWordCompareZero(InstructionSelector* selector, Node* user,
   while (selector->CanCover(user, value)) {
     switch (value->opcode()) {
       case IrOpcode::kWord32Equal: {
+        // Combine with comparisons against 0 by simply inverting the
+        // continuation.
         Int32BinopMatcher m(value);
         if (m.right().Is(0)) {
           user = value;
@@ -1926,10 +1945,33 @@ void VisitWordCompareZero(InstructionSelector* selector, Node* user,
       case IrOpcode::kUint32LessThanOrEqual:
         cont->OverwriteAndNegateIfEqual(kUnsignedLessThanOrEqual);
         return VisitWord32Compare(selector, value, cont);
-      case IrOpcode::kWord64Equal:
+      case IrOpcode::kWord64Equal: {
         cont->OverwriteAndNegateIfEqual(kEqual);
+        Int64BinopMatcher m(value);
+        if (m.right().Is(0)) {
+          Node* const left = m.left().node();
+          if (selector->CanCover(value, left) &&
+              left->opcode() == IrOpcode::kWord64And) {
+            // Attempt to merge the Word64Equal(Word64And(x, y), 0) comparison
+            // into a tbz/tbnz instruction.
+            if (TryEmitTestAndBranch<Uint64BinopMatcher, kArm64TestAndBranch>(
+                    selector, left, cont)) {
+              return;
+            }
+            return VisitWordCompare(selector, left, kArm64Tst, cont, true,
+                                    kLogical64Imm);
+          }
+          // Merge the Word64Equal(x, 0) comparison into a cbz instruction.
+          if (cont->IsBranch()) {
+            selector->Emit(cont->Encode(kArm64CompareAndBranch), g.NoOutput(),
+                           g.UseRegister(left), g.Label(cont->true_block()),
+                           g.Label(cont->false_block()));
+            return;
+          }
+        }
         return VisitWordCompare(selector, value, kArm64Cmp, cont, false,
                                 kArithmeticImm);
+      }
       case IrOpcode::kInt64LessThan:
         cont->OverwriteAndNegateIfEqual(kSignedLessThan);
         return VisitWordCompare(selector, value, kArm64Cmp, cont, false,
@@ -2004,42 +2046,20 @@ void VisitWordCompareZero(InstructionSelector* selector, Node* user,
                                 kArithmeticImm);
       case IrOpcode::kInt32Sub:
         return VisitWord32Compare(selector, value, cont);
-      case IrOpcode::kWord32And: {
-        Int32BinopMatcher m(value);
-        if (cont->IsBranch() && m.right().HasValue() &&
-            (base::bits::CountPopulation32(m.right().Value()) == 1)) {
-          // If the mask has only one bit set, we can use tbz/tbnz.
-          DCHECK((cont->condition() == kEqual) ||
-                 (cont->condition() == kNotEqual));
-          selector->Emit(
-              cont->Encode(kArm64TestAndBranch32), g.NoOutput(),
-              g.UseRegister(m.left().node()),
-              g.TempImmediate(
-                  base::bits::CountTrailingZeros32(m.right().Value())),
-              g.Label(cont->true_block()), g.Label(cont->false_block()));
+      case IrOpcode::kWord32And:
+        if (TryEmitTestAndBranch<Uint32BinopMatcher, kArm64TestAndBranch32>(
+                selector, value, cont)) {
           return;
         }
         return VisitWordCompare(selector, value, kArm64Tst32, cont, true,
                                 kLogical32Imm);
-      }
-      case IrOpcode::kWord64And: {
-        Int64BinopMatcher m(value);
-        if (cont->IsBranch() && m.right().HasValue() &&
-            (base::bits::CountPopulation64(m.right().Value()) == 1)) {
-          // If the mask has only one bit set, we can use tbz/tbnz.
-          DCHECK((cont->condition() == kEqual) ||
-                 (cont->condition() == kNotEqual));
-          selector->Emit(
-              cont->Encode(kArm64TestAndBranch), g.NoOutput(),
-              g.UseRegister(m.left().node()),
-              g.TempImmediate(
-                  base::bits::CountTrailingZeros64(m.right().Value())),
-              g.Label(cont->true_block()), g.Label(cont->false_block()));
+      case IrOpcode::kWord64And:
+        if (TryEmitTestAndBranch<Uint64BinopMatcher, kArm64TestAndBranch>(
+                selector, value, cont)) {
           return;
         }
         return VisitWordCompare(selector, value, kArm64Tst, cont, true,
                                 kLogical64Imm);
-      }
       default:
         break;
     }