[X86] Enable load folding in the test shrinking code

[craig.topper]

Summary

This patch adds load folding support to the test shrinking code. This was noticed missing in the review for D52669

[RKSimon]

LGTM

lib/Target/X86/X86ISelDAGToDAG.cpp

if (N0.getOpcode() == ISD::AND &&
         N0.getValueType() != MVT::i8 &&
         X86::isZeroNode(N1)) {
       ConstantSDNode *C = dyn_cast<ConstantSDNode>(N0.getOperand(1));
       if (!C) break;
       uint64_t Mask = C->getZExtValue();
 
       MVT VT;
       int SubRegOp;
-      unsigned Op;
+      unsigned ROpc, MOpc;
 
       // For each of these checks we need to be careful if the sign flag is
       // being used. It is only safe to use the sign flag in two conditions,
       // either the sign bit in the shrunken mask is zero or the final test
       // size is equal to the original compare size.
 
       if (isUInt<8>(Mask) &&
           (!(Mask & 0x80) || CmpVT == MVT::i8 ||
            hasNoSignedComparisonUses(Node))) {
         // For example, convert "testl %eax, $8" to "testb %al, $8"
         VT = MVT::i8;
         SubRegOp = X86::sub_8bit;
-        Op = X86::TEST8ri;
+        ROpc = X86::TEST8ri;
+        MOpc = X86::TEST8mi;
       } else if (OptForMinSize && isUInt<16>(Mask) &&
                  (!(Mask & 0x8000) || CmpVT == MVT::i16 ||
                   hasNoSignedComparisonUses(Node))) {
         // For example, "testl %eax, $32776" to "testw %ax, $32776".
         // NOTE: We only want to form TESTW instructions if optimizing for
         // min size. Otherwise we only save one byte and possibly get a length
         // changing prefix penalty in the decoders.
         VT = MVT::i16;
         SubRegOp = X86::sub_16bit;
-        Op = X86::TEST16ri;
+        ROpc = X86::TEST16ri;
+        MOpc = X86::TEST16mi;
       } else if (isUInt<32>(Mask) && N0.getValueType() != MVT::i16 &&
                  (!(Mask & 0x80000000) || CmpVT == MVT::i32 ||
                   hasNoSignedComparisonUses(Node))) {
         // For example, "testq %rax, $268468232" to "testl %eax, $268468232".
         // NOTE: We only want to run that transform if N0 is 32 or 64 bits.
         // Otherwize, we find ourselves in a position where we have to do
         // promotion. If previous passes did not promote the and, we assume
         // they had a good reason not to and do not promote here.
         VT = MVT::i32;
         SubRegOp = X86::sub_32bit;
-        Op = X86::TEST32ri;
+        ROpc = X86::TEST32ri;
+        MOpc = X86::TEST32mi;
       } else {
         // No eligible transformation was found.
         break;
       }
 
       // FIXME: We should be able to fold loads here.
 
       SDValue Imm = CurDAG->getTargetConstant(Mask, dl, VT);
       SDValue Reg = N0.getOperand(0);
 
+      // Emit a testl or testw.
+      MachineSDNode *NewNode;
+      SDValue Tmp0, Tmp1, Tmp2, Tmp3, Tmp4;
+      if (tryFoldLoad(Node, N0.getNode(), Reg, Tmp0, Tmp1, Tmp2, Tmp3, Tmp4)) {
+        SDValue Ops[] = { Tmp0, Tmp1, Tmp2, Tmp3, Tmp4, Imm,
+                          Reg.getOperand(0) };
+        NewNode = CurDAG->getMachineNode(MOpc, dl, MVT::i32, MVT::Other, Ops);
+        // Update the chain.
+        ReplaceUses(Reg.getValue(1), SDValue(NewNode, 1));
+        // Record the mem-refs
+        CurDAG->setNodeMemRefs(NewNode,
+                               {cast<LoadSDNode>(Reg)->getMemOperand()});
+      } else {
         // Extract the subregister if necessary.
         if (N0.getValueType() != VT)
           Reg = CurDAG->getTargetExtractSubreg(SubRegOp, dl, VT, Reg);
 
-      // Emit a testl or testw.
-      SDNode *NewNode = CurDAG->getMachineNode(Op, dl, MVT::i32, Reg, Imm);
+        NewNode = CurDAG->getMachineNode(ROpc, dl, MVT::i32, Reg, Imm);
+      }
       // Replace CMP with TEST.
       ReplaceNode(Node, NewNode);
       return;
     }
     break;
   }
   case X86ISD::PCMPISTR: {
     if (!Subtarget->hasSSE42())

test/CodeGen/X86/test-shrink.ll

 ; CHECK-WIN32-64-NEXT:  # %bb.1: # %yes
 ; CHECK-WIN32-64-NEXT:    callq bar
 ; CHECK-WIN32-64-NEXT:  .LBB14_2: # %no
 ; CHECK-WIN32-64-NEXT:    addq $40, %rsp
 ; CHECK-WIN32-64-NEXT:    retq
 ;
 ; CHECK-X86-LABEL: and32_trunc_8_sign:
 ; CHECK-X86:       # %bb.0:
-; CHECK-X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-X86-NEXT:    testb $-128, %al
+; CHECK-X86-NEXT:    testb $-128, {{[0-9]+}}(%esp)
 ; CHECK-X86-NEXT:    jg .LBB14_2
 ; CHECK-X86-NEXT:  # %bb.1: # %yes
 ; CHECK-X86-NEXT:    calll bar
 ; CHECK-X86-NEXT:  .LBB14_2: # %no
 ; CHECK-X86-NEXT:    retl
   %t = and i32 %x, 128
   %r = trunc i32 %t to i8
   %s = icmp sle i8 %r, 0

 ; CHECK-WIN32-64-NEXT:  # %bb.1: # %yes
 ; CHECK-WIN32-64-NEXT:    callq bar
 ; CHECK-WIN32-64-NEXT:  .LBB15_2: # %no
 ; CHECK-WIN32-64-NEXT:    addq $40, %rsp
 ; CHECK-WIN32-64-NEXT:    retq
 ;
 ; CHECK-X86-LABEL: and64_trunc_8_sign:
 ; CHECK-X86:       # %bb.0:
-; CHECK-X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-X86-NEXT:    testb $-128, %al
+; CHECK-X86-NEXT:    testb $-128, {{[0-9]+}}(%esp)
 ; CHECK-X86-NEXT:    jg .LBB15_2
 ; CHECK-X86-NEXT:  # %bb.1: # %yes
 ; CHECK-X86-NEXT:    calll bar
 ; CHECK-X86-NEXT:  .LBB15_2: # %no
 ; CHECK-X86-NEXT:    retl
   %t = and i64 %x, 128
   %r = trunc i64 %t to i8
   %s = icmp sle i8 %r, 0

 ; CHECK-WIN32-64-NEXT:  # %bb.1: # %yes
 ; CHECK-WIN32-64-NEXT:    callq bar
 ; CHECK-WIN32-64-NEXT:  .LBB16_2: # %no
 ; CHECK-WIN32-64-NEXT:    addq $40, %rsp
 ; CHECK-WIN32-64-NEXT:    retq
 ;
 ; CHECK-X86-LABEL: and32_trunc_16_sign:
 ; CHECK-X86:       # %bb.0:
-; CHECK-X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-X86-NEXT:    testw $-32768, %ax # imm = 0x8000
+; CHECK-X86-NEXT:    testw $-32768, {{[0-9]+}}(%esp) # imm = 0x8000
 ; CHECK-X86-NEXT:    jg .LBB16_2
 ; CHECK-X86-NEXT:  # %bb.1: # %yes
 ; CHECK-X86-NEXT:    calll bar
 ; CHECK-X86-NEXT:  .LBB16_2: # %no
 ; CHECK-X86-NEXT:    retl
   %t = and i32 %x, 32768
   %r = trunc i32 %t to i16
   %s = icmp sle i16 %r, 0

 ; CHECK-WIN32-64-NEXT:  # %bb.1: # %yes
 ; CHECK-WIN32-64-NEXT:    callq bar
 ; CHECK-WIN32-64-NEXT:  .LBB17_2: # %no
 ; CHECK-WIN32-64-NEXT:    addq $40, %rsp
 ; CHECK-WIN32-64-NEXT:    retq
 ;
 ; CHECK-X86-LABEL: and64_trunc_32_sign:
 ; CHECK-X86:       # %bb.0:
-; CHECK-X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-X86-NEXT:    testw $-32768, %ax # imm = 0x8000
+; CHECK-X86-NEXT:    testw $-32768, {{[0-9]+}}(%esp) # imm = 0x8000
 ; CHECK-X86-NEXT:    jg .LBB17_2
 ; CHECK-X86-NEXT:  # %bb.1: # %yes
 ; CHECK-X86-NEXT:    calll bar
 ; CHECK-X86-NEXT:  .LBB17_2: # %no
 ; CHECK-X86-NEXT:    retl
   %t = and i64 %x, 32768
   %r = trunc i64 %t to i16
   %s = icmp sle i16 %r, 0