From de7cee24b6fe97cc11225b20853ea97532a1c2e9 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Thu, 28 Apr 2022 16:01:33 +0100
Subject: [PATCH] [X86] getBT - attempt to peek through aext(and(trunc(x),c))
 mask/modulo

Ideally we'd fold this with generic DAGCombiner, but that only works for !isTruncateFree cases - we might be able to adapt IsDesirableToPromoteOp to find truncated src ops in the future, but for now just use this peephole.

Noticed in Issue #55138
---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 15 +++++++++++++--
 llvm/test/CodeGen/X86/setcc.ll          | 12 +++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index f66194d..c23c832 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -22975,8 +22975,19 @@ static SDValue getBT(SDValue Src, SDValue BitNo, const SDLoc &DL, SelectionDAG &
 
   // If the operand types disagree, extend the shift amount to match.  Since
   // BT ignores high bits (like shifts) we can use anyextend.
-  if (Src.getValueType() != BitNo.getValueType())
-    BitNo = DAG.getNode(ISD::ANY_EXTEND, DL, Src.getValueType(), BitNo);
+  if (Src.getValueType() != BitNo.getValueType()) {
+    // Peek through a mask/modulo operation.
+    // TODO: DAGCombine fails to do this as it just checks isTruncateFree, but
+    // we probably need a better IsDesirableToPromoteOp to handle this as well.
+    if (BitNo.getOpcode() == ISD::AND && BitNo->hasOneUse())
+      BitNo = DAG.getNode(ISD::AND, DL, Src.getValueType(),
+                          DAG.getNode(ISD::ANY_EXTEND, DL, Src.getValueType(),
+                                      BitNo.getOperand(0)),
+                          DAG.getNode(ISD::ANY_EXTEND, DL, Src.getValueType(),
+                                      BitNo.getOperand(1)));
+    else
+      BitNo = DAG.getNode(ISD::ANY_EXTEND, DL, Src.getValueType(), BitNo);
+  }
 
   return DAG.getNode(X86ISD::BT, DL, MVT::i32, Src, BitNo);
 }
diff --git a/llvm/test/CodeGen/X86/setcc.ll b/llvm/test/CodeGen/X86/setcc.ll
index d600aa6..9bac47c 100644
--- a/llvm/test/CodeGen/X86/setcc.ll
+++ b/llvm/test/CodeGen/X86/setcc.ll
@@ -311,9 +311,8 @@ define i16 @shift_and(i16 %a) {
 define i32 @PR55138(i32 %x) {
 ; X86-LABEL: PR55138:
 ; X86:       ## %bb.0:
-; X86-NEXT:    movb {{[0-9]+}}(%esp), %al
-; X86-NEXT:    andb $15, %al
-; X86-NEXT:    movzbl %al, %ecx
+; X86-NEXT:    movzbl {{[0-9]+}}(%esp), %ecx
+; X86-NEXT:    andl $15, %ecx
 ; X86-NEXT:    movl $27030, %edx ## imm = 0x6996
 ; X86-NEXT:    xorl %eax, %eax
 ; X86-NEXT:    btl %ecx, %edx
@@ -322,11 +321,10 @@ define i32 @PR55138(i32 %x) {
 ;
 ; X64-LABEL: PR55138:
 ; X64:       ## %bb.0:
-; X64-NEXT:    andb $15, %dil
-; X64-NEXT:    movzbl %dil, %ecx
-; X64-NEXT:    movl $27030, %edx ## imm = 0x6996
+; X64-NEXT:    andl $15, %edi
+; X64-NEXT:    movl $27030, %ecx ## imm = 0x6996
 ; X64-NEXT:    xorl %eax, %eax
-; X64-NEXT:    btl %ecx, %edx
+; X64-NEXT:    btl %edi, %ecx
 ; X64-NEXT:    setb %al
 ; X64-NEXT:    retq
   %urem = and i32 %x, 15
-- 
2.7.4