[X86] Widen the 'AND' mask if doing so shrinks the encoding size

[majnemer]

Summary

We can set additional bits in a mask given that we know the other
operand of an AND already has some bits set to zero. This can be more
efficient if doing so allows us to use an instruction which implicitly
sign extends the immediate.

This fixes PR24085.

[majnemer]

• Small cleanup

[escha]

lib/Target/X86/X86ISelDAGToDAG.cpp

2291

Is this right? (int8_t)-384 would still pass this check, right?

[majnemer]

lib/Target/X86/X86ISelDAGToDAG.cpp

2291

We wouldn't get that far because it would fail the above Val > 0 check.

[chandlerc]

Cool implementation. Some thoughts inline about structuring the code a bit better.

lib/Target/X86/X86ISelDAGToDAG.cpp

2280–2284

At this point, I think it would be worth extracting this to a helper function. That will let you use early-exit to make everything easier to read IMO.

2305–2323

I would put all of this into a lambda that you can call with the number of bits, the MVT, and the opcode. It should be easier to read, and make the above tests able to just return the call of the lambda each time they find a valid opcode.

[majnemer]

• Small cleanup
• Address review comments

[chandlerc]

Ugh. The lambda didn't help as much as i had hoped because you need to only compute the known bits once. Unfortunate.

lib/Target/X86/X86ISelDAGToDAG.cpp

2145–2146

uint64_t(...) <= 0? Why not just == 0, or did you mean something else?

[majnemer]

In D11289#207674, @chandlerc wrote:

Ugh. The lambda didn't help as much as i had hoped because you need to only compute the known bits once. Unfortunate.

Is it still worth keeping the lambda?

lib/Target/X86/X86ISelDAGToDAG.cpp

2145–2146

It's a victim of refactoring. Val should be int64_t.

[majnemer]

• Fix a bug introduced during refactoring.

[chandlerc]

LGTM

I'd just leave the code as is.

If you can come up with tests that would tickle the int64_t and int8_t case, it'd be nice to add those.

lib/Target/X86/X86ISelDAGToDAG.cpp

 
 // Include the pieces autogenerated from the target description.
 #include "X86GenDAGISel.inc"
 
   private:
     SDNode *Select(SDNode *N) override;
     SDNode *SelectGather(SDNode *N, unsigned Opc);
     SDNode *SelectAtomicLoadArith(SDNode *Node, MVT NVT);
+    SDNode *SelectAndWithSExtImmediate(SDNode *Node, MVT NVT);
 
     bool FoldOffsetIntoAddress(uint64_t Offset, X86ISelAddressMode &AM);
     bool MatchLoadInAddress(LoadSDNode *N, X86ISelAddressMode &AM);
     bool MatchWrapper(SDValue N, X86ISelAddressMode &AM);
     bool MatchAddress(SDValue N, X86ISelAddressMode &AM);
     bool MatchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
                                  unsigned Depth);
     bool MatchAddressBase(SDValue N, X86ISelAddressMode &AM);
SDNode *X86DAGToDAGISel::SelectGather(SDNode *Node, unsigned Opc) {
   // ResNode has 3 outputs: VDst, VMask_wb, and MVT::Other.
   // We replace VDst of Node with VDst of ResNode, and Other of Node with Other
   // of ResNode.
   ReplaceUses(SDValue(Node, 0), SDValue(ResNode, 0));
   ReplaceUses(SDValue(Node, 1), SDValue(ResNode, 2));
   return ResNode;
 }
 
+// Try to shrink the encoding of an AND by setting additional bits in the mask.
+// It is only correct to do so if we know a priori that the other operand of the
+// AND already has those bits set to zero.
+SDNode *X86DAGToDAGISel::SelectAndWithSExtImmediate(SDNode *Node, MVT NVT) {
+  SDValue N0 = Node->getOperand(0);
+  SDValue N1 = Node->getOperand(1);
+  auto *Cst = dyn_cast<ConstantSDNode>(N1);
+  if (!Cst)
+    return nullptr;
+  int64_t Val = Cst->getSExtValue();
+  if (Val <= 0)
+    return nullptr;
+  unsigned Opc8 = 0, Opc32 = 0;
+  // Limit ourselves to constants which already have sign bits to save on
+  // compile time.
+  if ((int8_t)Val < 0) {
+    if (NVT.SimpleTy == MVT::i64)
+      Opc8 = X86::AND64ri8;
+    else if (NVT.SimpleTy == MVT::i32)
+      Opc8 = X86::AND32ri8;
+    else if (NVT.SimpleTy == MVT::i16)
+      Opc8 = X86::AND16ri8;
+  }
+  if ((int32_t)Val < 0)
+    if (NVT.SimpleTy == MVT::i64)
+      Opc32 = X86::AND64ri32;
+  // Only bother computing the zero bits in the non-constant operand if we
+  // have a suitable instruction.
+  if (!Opc8 && !Opc32)
+    return nullptr;
+  APInt Op0Zero, Op0One;
+  CurDAG->computeKnownBits(N0, Op0Zero, Op0One);
+  // Grow the mask using the known zero bits.
+  Op0Zero |= Val;
+  auto SelectAndInstruction = [&](unsigned Opc, unsigned NumBits) -> SDNode * {
+    // Do we have an instruction for this number of bits?
+    if (!Opc)
+      return nullptr;
+    // See if the mask can be efficiently encoded using at most NumBits.
+    if (!Op0Zero.isSignedIntN(NumBits))
+      return nullptr;
+    SDLoc DL(Node);
+    SDValue NewCst = CurDAG->getTargetConstant(Op0Zero.getSExtValue(), DL,
+                                               MVT::getIntegerVT(NumBits));
+    return CurDAG->getMachineNode(Opc, DL, NVT, N0, NewCst);
+  };
+  if (SDNode *NewInst = SelectAndInstruction(Opc8, 8))
+    return NewInst;
+  if (SDNode *NewInst = SelectAndInstruction(Opc32, 32))
+    return NewInst;
+  return nullptr;
+}
+
 SDNode *X86DAGToDAGISel::Select(SDNode *Node) {
   MVT NVT = Node->getSimpleValueType(0);
   unsigned Opc, MOpc;
   unsigned Opcode = Node->getOpcode();
   SDLoc dl(Node);
 
   DEBUG(dbgs() << "Selecting: "; Node->dump(CurDAG); dbgs() << '\n');
 
SDNode *X86DAGToDAGISel::Select(SDNode *Node) {
   case ISD::ATOMIC_LOAD_AND:
   case ISD::ATOMIC_LOAD_OR:
   case ISD::ATOMIC_LOAD_ADD: {
     SDNode *RetVal = SelectAtomicLoadArith(Node, NVT);
     if (RetVal)
       return RetVal;
     break;
   }
-  case ISD::AND:
+  case ISD::AND: {
+    // FALLTHROUGH
+    if (SDNode *NewNode = SelectAndWithSExtImmediate(Node, NVT)) {
+      ReplaceUses(SDValue(Node, 0), SDValue(NewNode, 0));
+      return nullptr;
+    }
+  }
   case ISD::OR:
   case ISD::XOR: {
     // For operations of the form (x << C1) op C2, check if we can use a smaller
     // encoding for C2 by transforming it into (x op (C2>>C1)) << C1.
     SDValue N0 = Node->getOperand(0);
     SDValue N1 = Node->getOperand(1);
 
     if (N0->getOpcode() != ISD::SHL || !N0->hasOneUse())
       break;
 
     // i8 is unshrinkable, i16 should be promoted to i32.
     if (NVT != MVT::i32 && NVT != MVT::i64)
       break;
 
     ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(N1);
     ConstantSDNode *ShlCst = dyn_cast<ConstantSDNode>(N0->getOperand(1));
     if (!Cst || !ShlCst)
       break;
 
     int64_t Val = Cst->getSExtValue();
     uint64_t ShlVal = ShlCst->getZExtValue();
 
     // Make sure that we don't change the operation by removing bits.
     // This only matters for OR and XOR, AND is unaffected.
     uint64_t RemovedBitsMask = (1ULL << ShlVal) - 1;
     if (Opcode != ISD::AND && (Val & RemovedBitsMask) != 0)
       break;
 
     unsigned ShlOp, AddOp, Op;
     MVT CstVT = NVT;
 
     // Check the minimum bitwidth for the new constant.
     // TODO: AND32ri is the same as AND64ri32 with zext imm.
     // TODO: MOV32ri+OR64r is cheaper than MOV64ri64+OR64rr
     // TODO: Using 16 and 8 bit operations is also possible for or32 & xor32.
     if (!isInt<8>(Val) && isInt<8>(Val >> ShlVal))
       CstVT = MVT::i8;
     else if (!isInt<32>(Val) && isInt<32>(Val >> ShlVal))
       CstVT = MVT::i32;
 
     // Bail if there is no smaller encoding.
     if (NVT == CstVT)
       break;
 
     switch (NVT.SimpleTy) {
     default: llvm_unreachable("Unsupported VT!");

test/CodeGen/X86/shift-pair.ll

 ; RUN: llc < %s -march=x86-64 | FileCheck %s
 
 define i64 @test(i64 %A) {
 ; CHECK: @test
 ; CHECK: shrq $54
-; CHECK: andl $1020
+; CHECK: andq $-4
 ; CHECK: ret
     %B = lshr i64 %A, 56
     %C = shl i64 %B, 2
     ret i64 %C
 }

test/CodeGen/X86/win64_frame.ll

define i32 @f8(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) "no-frame-pointer-elim"="true" {
   ; CHECK:        .seh_setframe 5, 128
 
   %alloca = alloca [300 x i8], align 64
   ; CHECK:        andq    $-64, %rsp
   ; CHECK:        movq    %rsp, %rbx
 
   alloca i32, i32 %a
   ; CHECK:        movl    %ecx, %eax
-  ; CHECK:        leaq    15(,%rax,4), %rcx
-  ; CHECK:        movabsq $34359738352, %rax
-  ; CHECK:        andq    %rcx, %rax
+  ; CHECK:        leaq    15(,%rax,4), %rax
+  ; CHECK:        andq    $-16, %rax
   ; CHECK:        callq   __chkstk
   ; CHECK:        subq    %rax, %rsp
 
   %gep = getelementptr [300 x i8], [300 x i8]* %alloca, i32 0, i32 0
   call void @external(i8* %gep)
   ; CHECK:        subq    $32, %rsp
   ; CHECK:        leaq    (%rbx), %rcx
   ; CHECK:        callq   external

test/CodeGen/X86/zext-fold.ll

 ; RUN: llc < %s -mcpu=generic -march=x86 -enable-misched=false | FileCheck %s
 
 ;; Simple case
 define i32 @test1(i8 %x) nounwind readnone {
   %A = and i8 %x, -32
   %B = zext i8 %A to i32
   ret i32 %B
 }
 ; CHECK: test1
 ; CHECK: movzbl
-; CHECK-NEXT: andl {{.*}}224
+; CHECK-NEXT: andl {{.*}}-32
 
 ;; Multiple uses of %x but easily extensible.
 define i32 @test2(i8 %x) nounwind readnone {
   %A = and i8 %x, -32
   %B = zext i8 %A to i32
   %C = or i8 %x, 63
   %D = zext i8 %C to i32
   %E = add i32 %B, %D
   ret i32 %E
 }
 ; CHECK: test2
 ; CHECK: movzbl
-; CHECK: andl $224
+; CHECK: andl $-32
 ; CHECK: orl $63
 
 declare void @use(i32, i8)
 
 ;; Multiple uses of %x where we shouldn't extend the load.
 define void @test3(i8 %x) nounwind readnone {
   %A = and i8 %x, -32
   %B = zext i8 %A to i32
   call void @use(i32 %B, i8 %x)
   ret void
 }
 ; CHECK: test3
 ; CHECK: movzbl {{[0-9]+}}(%esp), [[REGISTER:%e[a-z]{2}]]
 ; CHECK-NEXT: movl [[REGISTER]], 4(%esp)
-; CHECK-NEXT: andl $224, [[REGISTER]]
+; CHECK-NEXT: andl $-32, [[REGISTER]]
 ; CHECK-NEXT: movl [[REGISTER]], (%esp)
 ; CHECK-NEXT: call{{.*}}use