[AArch64] Optimize some simple TBZ/TBNZ cases.

[gberry]

Summary

Add some AArch64 dag combines to optimize some simple TBZ/TBNZ cases:

(tbz (and x, m), b) -> (tbz x, b)
(tbz (shl x, c), b) -> (tbz x, b-c)
(tbz (shr x, c), b) -> (tbz x, b+c)
(tbz (xor x, -1), b) -> (tbnz x, b)

[mcrosier]

lib/Target/AArch64/AArch64ISelLowering.cpp

9528

Please remove the extra curly brackets.

9534

Mind adding a test that covers this case?

9543

I don't think SRA is covered by a test. If so, please add one.

[• chatur01]

Hi Geoff, looks good. :-)

I have some potentially silly questions inline, just curiosities.

lib/Target/AArch64/AArch64ISelLowering.cpp

9528

Are there interesting cases where m[b]==0, so that you an elide the test and generate a direct branch? Or is that handled somewhere else in the compiler (please ignore if this would zap time to figure out, I'm only curious)

9540

Same here, f we shl > Bits, then we know x[Bit] == 0.

[gberry]

Add more tests and handle ashr testing sign bit case.

[gberry]

Chad, Charlie,

Thanks for the review. I believe I have addressed all of your concerns.

Charlie, as to your question about folding opportunities, I don't think those should occur (they should have already been optimized). I did a test run to detect if any of these cases show up and was unable to find a single occurrence. I did however hit a few instances of a missing case of a tbz of an ashr that was essentially testing the carried-in sign bit (see the test case test18 for an example). I added some new code to handle that case in the latest version of the change.

[mcrosier]

LGTM

lib/Target/AArch64/AArch64ISelLowering.cpp

else
     BR = DAG.getNode(AArch64ISD::CBNZ, SDLoc(N), MVT::Other, Chain, LHS, Dest);
 
   // Do not add new nodes to DAG combiner worklist.
   DCI.CombineTo(N, BR, false);
 
   return SDValue();
 }
 
+// Optimize some simple tbz/tbnz cases.  Returns the new operand and bit to test
+// as well as whether the test should be inverted.  This code is required to
+// catch these cases (as opposed to standard dag combines) because
+// AArch64ISD::TBZ is matched during legalization.
+static SDValue getTestBitOperand(SDValue Op, unsigned &Bit, bool &Invert,
+                                 SelectionDAG &DAG) {
+
+  if (!Op->hasOneUse())
+    return Op;
+
+  // We don't handle undef/constant-fold cases below, as they should have
+  // already been taken care of (e.g. and of 0, test of undefined shifted bits,
+  // etc.)
+
+  // (tbz (trunc x), b) -> (tbz x, b)
+  // This case is just here to enable more of the below cases to be caught.
+  if (Op->getOpcode() == ISD::TRUNCATE &&
+      Bit < Op->getValueType(0).getSizeInBits()) {
+    return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+  }
+
+  if (Op->getNumOperands() != 2)
+    return Op;
+
+  auto *C = dyn_cast<ConstantSDNode>(Op->getOperand(1));
+  if (!C)
+    return Op;
+
+  switch (Op->getOpcode()) {
+  default:
+    return Op;
+
+  // (tbz (and x, m), b) -> (tbz x, b)
+  case ISD::AND:
+    if ((C->getZExtValue() >> Bit) & 1)
+      return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+    return Op;
+
+  // (tbz (shl x, c), b) -> (tbz x, b-c)
+  case ISD::SHL:
+    if (C->getZExtValue() <= Bit &&
+        (Bit - C->getZExtValue()) < Op->getValueType(0).getSizeInBits()) {
+      Bit = Bit - C->getZExtValue();
+      return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+    }
+    return Op;
+
+  // (tbz (sra x, c), b) -> (tbz x, b+c) or (tbz x, msb) if b+c is > # bits in x
+  case ISD::SRA:
+    Bit = Bit + C->getZExtValue();
+    if (Bit >= Op->getValueType(0).getSizeInBits())
+      Bit = Op->getValueType(0).getSizeInBits() - 1;
+    return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+
+  // (tbz (srl x, c), b) -> (tbz x, b+c)
+  case ISD::SRL:
+    if ((Bit + C->getZExtValue()) < Op->getValueType(0).getSizeInBits()) {
+      Bit = Bit + C->getZExtValue();
+      return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+    }
+    return Op;
+
+  // (tbz (xor x, -1), b) -> (tbnz x, b)
+  case ISD::XOR:
+    if ((C->getZExtValue() >> Bit) & 1)
+      Invert = !Invert;
+    return getTestBitOperand(Op->getOperand(0), Bit, Invert, DAG);
+  }
+}
+
+// Optimize test single bit zero/non-zero and branch.
+static SDValue performTBZCombine(SDNode *N,
+                                 TargetLowering::DAGCombinerInfo &DCI,
+                                 SelectionDAG &DAG) {
+  unsigned Bit = cast<ConstantSDNode>(N->getOperand(2))->getZExtValue();
+  bool Invert = false;
+  SDValue TestSrc = N->getOperand(1);
+  SDValue NewTestSrc = getTestBitOperand(TestSrc, Bit, Invert, DAG);
+
+  if (TestSrc == NewTestSrc)
+    return SDValue();
+
+  unsigned NewOpc = N->getOpcode();
+  if (Invert) {
+    if (NewOpc == AArch64ISD::TBZ)
+      NewOpc = AArch64ISD::TBNZ;
+    else {
+      assert(NewOpc == AArch64ISD::TBNZ);
+      NewOpc = AArch64ISD::TBZ;
+    }
+  }
+
+  SDLoc DL(N);
+  return DAG.getNode(NewOpc, DL, MVT::Other, N->getOperand(0), NewTestSrc,
+                     DAG.getConstant(Bit, DL, MVT::i64), N->getOperand(3));
+}
+
 // vselect (v1i1 setcc) ->
 //     vselect (v1iXX setcc)  (XX is the size of the compared operand type)
 // FIXME: Currently the type legalizer can't handle VSELECT having v1i1 as
 // condition. If it can legalize "VSELECT v1i1" correctly, no need to combine
 // such VSELECT.
 static SDValue performVSelectCombine(SDNode *N, SelectionDAG &DAG) {
   SDValue N0 = N->getOperand(0);
   EVT CCVT = N0.getValueType();
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case ISD::LOAD:
     if (performTBISimplification(N->getOperand(1), DCI, DAG))
       return SDValue(N, 0);
     break;
   case ISD::STORE:
     return performSTORECombine(N, DCI, DAG, Subtarget);
   case AArch64ISD::BRCOND:
     return performBRCONDCombine(N, DCI, DAG);
+  case AArch64ISD::TBNZ:
+  case AArch64ISD::TBZ:
+    return performTBZCombine(N, DCI, DAG);
   case AArch64ISD::CSEL:
     return performCONDCombine(N, DCI, DAG, 2, 3);
   case AArch64ISD::DUP:
     return performPostLD1Combine(N, DCI, false);
   case AArch64ISD::NVCAST:
     return performNVCASTCombine(N);
   case ISD::INSERT_VECTOR_ELT:
     return performPostLD1Combine(N, DCI, true);

test/CodeGen/AArch64/tbz-tbnz.ll

 
 if.then:
   call void @t()
   br label %if.end
 
 if.end:
   ret void
 }
+
+define void @test14(i1 %cond) {
+; CHECK-LABEL: @test14
+  br i1 %cond, label %if.end, label %if.then
+
+; CHECK-NOT: and
+; CHECK: tbnz w0, #0
+
+if.then:
+  call void @t()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+define void @test15(i1 %cond) {
+; CHECK-LABEL: @test15
+  %cond1 = xor i1 %cond, -1
+  br i1 %cond1, label %if.then, label %if.end
+
+; CHECK-NOT: movn
+; CHECK: tbnz w0, #0
+
+if.then:
+  call void @t()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+define void @test16(i64 %in) {
+; CHECK-LABEL: @test16
+  %shl = shl i64 %in, 3
+  %and = and i64 %shl, 32
+  %cond = icmp eq i64 %and, 0
+  br i1 %cond, label %then, label %end
+
+; CHECK-NOT: lsl
+; CHECK: tbnz w0, #2
+
+then:
+  call void @t()
+  br label %end
+
+end:
+  ret void
+}
+
+define void @test17(i64 %in) {
+; CHECK-LABEL: @test17
+  %shr = ashr i64 %in, 3
+  %and = and i64 %shr, 1
+  %cond = icmp eq i64 %and, 0
+  br i1 %cond, label %then, label %end
+
+; CHECK-NOT: lsr
+; CHECK: tbnz w0, #3
+
+then:
+  call void @t()
+  br label %end
+
+end:
+  ret void
+}
+
+define void @test18(i32 %in) {
+; CHECK-LABEL: @test18
+  %shr = ashr i32 %in, 2
+  %cond = icmp sge i32 %shr, 0
+  br i1 %cond, label %then, label %end
+
+; CHECK-NOT: asr
+; CHECK: tbnz w0, #31
+
+then:
+  call void @t()
+  br label %end
+
+end:
+  ret void
+}
+
+define void @test19(i64 %in) {
+; CHECK-LABEL: @test19
+  %shl = lshr i64 %in, 3
+  %trunc = trunc i64 %shl to i32
+  %and = and i32 %trunc, 1
+  %cond = icmp eq i32 %and, 0
+  br i1 %cond, label %then, label %end
+
+; CHECK-NOT: ubfx
+; CHECK: tbnz w0, #3
+
+then:
+  call void @t()
+  br label %end
+
+end:
+  ret void
+}