X86ISelDAGToDAG: Transform TEST + MOV64ri to SHR + TEST

[MatzeB]

Summary

Optimize a pattern where a sequence of 8/16 or 32 bits is tested for zero: LLVM normalizes this towards and AND with mask which is usually good, but does not work well on X86 when the mask does not fit into a 64bit register. This DagToDAG peephole transforms sequences like:

movabsq $562941363486720, %rax # imm = 0x1FFFE00000000
testq %rax, %rdi

to

shrq $33, %rdi
testw %di, %di

the result has a shorter encoding and saves a register if the tested value isn't used otherwise.

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

5620

Is this code already doing something similar? Sanjay had another patch to this code D121147 recently.

[MatzeB]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

5620

Interesting. I believe this case covers only the cases where the mask covers the highest bit in the register and so allows us to get rid of the test completely, while my case covers cases of the mask being "in-between" and not covering the highest bit so we still need the test...

That said let me dig deeper whether there is code to be shared/reorganized/merged...

[MatzeB]

Changed the code to hook into the existing similar case in X86DAGToDAGISel::Select.

[MatzeB]

Linted and added some comments.

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

5631

Are we emitting TEST+SHR rather than SHR by itself.

5636

Same question.

5663

Initialize TestOpcode to X86::TEST64rr instead of DELETED_NODE. Then this becomes

if (SubRegIdx != 0)
  Shift = CurDAG->getTargetExtractSubreg(SubRegIdx, dl, SubRegVT, Shift);
Test = CurDAG->getMachineNode(TestOpcode, dl, MVT::i32, Shift, Shift);
ReplaceNode(Node, Test)

[MatzeB]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

5631

It seems we emit a TEST here but a separate transformation will identify the TEST as redundant and re-use the flags of the shift instruction.

5636

as above.

[MatzeB]

Address review feedback

[spatel]

There was a question about the effect of extra uses within this pattern, so I added some tests to show the effect on the existing code:
ca808e89242f

Can you add some tests similar to those, so we'll have more visibility into what this change will do?

[MatzeB]

rebase; add extra hasOneUseCheck for the constant.

[spatel]

LGTM

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

if (N0.getOpcode() == ISD::AND && N0.getNode()->hasOneUse() &&
       // We may have looked through a truncate so mask off any bits that
       // shouldn't be part of the compare.
       uint64_t Mask = MaskC->getZExtValue();
       Mask &= maskTrailingOnes<uint64_t>(CmpVT.getScalarSizeInBits());
 
       // Check if we can replace AND+IMM64 with a shift. This is possible for
       // masks like 0xFF000000 or 0x00FFFFFF and if we care only about the zero
       // flag.
       if (CmpVT == MVT::i64 && !isInt<32>(Mask) &&
           onlyUsesZeroFlag(SDValue(Node, 0))) {
         unsigned ShiftOpcode = ISD::DELETED_NODE;
         unsigned ShiftAmt;
-        if (isMask_64(~Mask)) {
+        unsigned SubRegIdx = 0;
+        unsigned TestOpcode = ISD::DELETED_NODE;
+        MVT SubRegVT;
+        if (isShiftedMask_64(Mask)) {
+          unsigned LeadingZeros = countLeadingZeros(Mask);
+          unsigned TrailingZeros = countTrailingZeros(Mask);
+          if (LeadingZeros == 0) {
             ShiftOpcode = X86::SHR64ri;
-          ShiftAmt = countTrailingZeros(Mask);
-        } else if (isMask_64(Mask)) {
+            ShiftAmt = TrailingZeros;
+          } else if (TrailingZeros == 0) {
             ShiftOpcode = X86::SHL64ri;
-          ShiftAmt = countLeadingZeros(Mask);
+            ShiftAmt = LeadingZeros;
+          } else {
+            unsigned PopCount = 64 - LeadingZeros - TrailingZeros;
+            if (PopCount == 8) {
+              ShiftOpcode = X86::SHR64ri;
+              TestOpcode = X86::TEST8rr;
+              SubRegIdx = X86::sub_8bit;
+              SubRegVT = MVT::i8;
+            } else if (PopCount == 16) {
+              ShiftOpcode = X86::SHR64ri;
+              TestOpcode = X86::TEST16rr;
+              SubRegIdx = X86::sub_16bit;
+              SubRegVT = MVT::i16;
+            } else if (PopCount == 32) {
+              ShiftOpcode = X86::SHR64ri;
+              TestOpcode = X86::TEST32rr;
+              SubRegIdx = X86::sub_32bit;
+              SubRegVT = MVT::i32;
+            }
+          }
         }
         if (ShiftOpcode != ISD::DELETED_NODE) {
           SDValue ShiftC = CurDAG->getTargetConstant(ShiftAmt, dl, MVT::i64);
           SDValue Shift = SDValue(
               CurDAG->getMachineNode(ShiftOpcode, dl, MVT::i64, MVT::i32,
                                      N0.getOperand(0), ShiftC),
               0);
-          MachineSDNode *Test =
-              CurDAG->getMachineNode(X86::TEST64rr, dl, MVT::i32, Shift, Shift);
+          MachineSDNode *Test;
+          if (SubRegIdx == 0) {
+            Test = CurDAG->getMachineNode(X86::TEST64rr, dl, MVT::i32, Shift,
+                                          Shift);
+          } else {
+            SDValue TestOp = CurDAG->getTargetExtractSubreg(SubRegIdx, dl,
+                                                            SubRegVT, Shift);
+            Test = CurDAG->getMachineNode(TestOpcode, dl, MVT::i32, TestOp,
+                                          TestOp);
+          }
           ReplaceNode(Node, Test);
           return;
         }
       }
 
       MVT VT;
       int SubRegOp;
       unsigned ROpc, MOpc;

llvm/test/CodeGen/X86/cmp.ll

; CHECK-NEXT:    retq # encoding: [0xc3]
   %cmp = icmp ne i32 %and, 0
   %ret = zext i1 %cmp to i32
   ret i32 %ret
 }
 
 define i1 @shifted_mask_testb(i64 %a) {
 ; CHECK-LABEL: shifted_mask_testb:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movabsq $287104476244869120, %rax # encoding: [0x48,0xb8,0x00,0x00,0x00,0x00,0x00,0x00,0xfc,0x03]
-; CHECK-NEXT:    # imm = 0x3FC000000000000
-; CHECK-NEXT:    testq %rax, %rdi # encoding: [0x48,0x85,0xc7]
+; CHECK-NEXT:    shrq $64, %rdi # encoding: [0x48,0xc1,0xef,0x40]
+; CHECK-NEXT:    testb %dil, %dil # encoding: [0x40,0x84,0xff]
 ; CHECK-NEXT:    setne %al # encoding: [0x0f,0x95,0xc0]
 ; CHECK-NEXT:    retq # encoding: [0xc3]
   %v0 = and i64 %a, 287104476244869120  ; 0xff << 50
   %v1 = icmp ne i64 %v0, 0
   ret i1 %v1
 }
 
 define i1 @shifted_mask_testw(i64 %a) {
 ; CHECK-LABEL: shifted_mask_testw:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movabsq $562941363486720, %rax # encoding: [0x48,0xb8,0x00,0x00,0x00,0x00,0xfe,0xff,0x01,0x00]
-; CHECK-NEXT:    # imm = 0x1FFFE00000000
-; CHECK-NEXT:    testq %rax, %rdi # encoding: [0x48,0x85,0xc7]
+; CHECK-NEXT:    shrq $64, %rdi # encoding: [0x48,0xc1,0xef,0x40]
+; CHECK-NEXT:    testw %di, %di # encoding: [0x66,0x85,0xff]
 ; CHECK-NEXT:    setne %al # encoding: [0x0f,0x95,0xc0]
 ; CHECK-NEXT:    retq # encoding: [0xc3]
   %v0 = and i64 %a, 562941363486720  ; 0xffff << 33
   %v1 = icmp ne i64 %v0, 0
   ret i1 %v1
 }
 
 define i1 @shifted_mask_testl(i64 %a) {
 ; CHECK-LABEL: shifted_mask_testl:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movabsq $549755813760, %rax # encoding: [0x48,0xb8,0x80,0xff,0xff,0xff,0x7f,0x00,0x00,0x00]
-; CHECK-NEXT:    # imm = 0x7FFFFFFF80
-; CHECK-NEXT:    testq %rax, %rdi # encoding: [0x48,0x85,0xc7]
+; CHECK-NEXT:    shrq $64, %rdi # encoding: [0x48,0xc1,0xef,0x40]
+; CHECK-NEXT:    testl %edi, %edi # encoding: [0x85,0xff]
 ; CHECK-NEXT:    sete %al # encoding: [0x0f,0x94,0xc0]
 ; CHECK-NEXT:    retq # encoding: [0xc3]
   %v0 = and i64 %a, 549755813760  ; 0xffffffff << 7
   %v1 = icmp eq i64 %v0, 0
   ret i1 %v1
 }
 
 define { i64, i64 } @pr39968(i64, i64, i32) {