[DAGCombiner] Improve MatchBswapHword logic (PR31357)

The existing code only looks at half of the tree when matching bswap + rol patterns ending in an OR tree (as opposed to a cascade).

Patch originally introduced by Jim Lewis.

Submitted on the behalf of Dinar Temirbulatov.

Differential Revision: https://reviews.llvm.org/D32039

llvm-svn: 301907

2 files changed, 26 insertions, 54 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 71b96daff13..c2200742add 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -3819,7 +3819,7 @@ SDValue DAGCombiner::MatchBSwapHWordLow(SDNode *N, SDValue N0, SDValue N1,
   EVT VT = N->getValueType(0);
   if (VT != MVT::i64 && VT != MVT::i32 && VT != MVT::i16)
     return SDValue();
-  if (!TLI.isOperationLegal(ISD::BSWAP, VT))
+  if (!TLI.isOperationLegalOrCustom(ISD::BSWAP, VT))
     return SDValue();
 
   // Recognize (and (shl a, 8), 0xff), (and (srl a, 8), 0xff00)
@@ -3933,8 +3933,15 @@ static bool isBSwapHWordElement(SDValue N, MutableArrayRef<SDNode *> Parts) {
 
   SDValue N0 = N.getOperand(0);
   unsigned Opc0 = N0.getOpcode();
+  if (Opc0 != ISD::AND && Opc0 != ISD::SHL && Opc0 != ISD::SRL)
+    return false;
 
-  ConstantSDNode *N1C = dyn_cast<ConstantSDNode>(N.getOperand(1));
+  ConstantSDNode *N1C = nullptr;
+  // SHL or SRL: look upstream for AND mask operand
+  if (Opc == ISD::AND)
+    N1C = dyn_cast<ConstantSDNode>(N.getOperand(1));
+  else if (Opc0 == ISD::AND)
+    N1C = dyn_cast<ConstantSDNode>(N0.getOperand(1));
   if (!N1C)
     return false;
 
@@ -4005,7 +4012,7 @@ SDValue DAGCombiner::MatchBSwapHWord(SDNode *N, SDValue N0, SDValue N1) {
   EVT VT = N->getValueType(0);
   if (VT != MVT::i32)
     return SDValue();
-  if (!TLI.isOperationLegal(ISD::BSWAP, VT))
+  if (!TLI.isOperationLegalOrCustom(ISD::BSWAP, VT))
     return SDValue();
 
   // Look for either
@@ -4020,18 +4027,16 @@ SDValue DAGCombiner::MatchBSwapHWord(SDNode *N, SDValue N0, SDValue N1) {
   if (N1.getOpcode() == ISD::OR &&
       N00.getNumOperands() == 2 && N01.getNumOperands() == 2) {
     // (or (or (and), (and)), (or (and), (and)))
-    SDValue N000 = N00.getOperand(0);
-    if (!isBSwapHWordElement(N000, Parts))
+    if (!isBSwapHWordElement(N00, Parts))
       return SDValue();
 
-    SDValue N001 = N00.getOperand(1);
-    if (!isBSwapHWordElement(N001, Parts))
+    if (!isBSwapHWordElement(N01, Parts))
       return SDValue();
-    SDValue N010 = N01.getOperand(0);
-    if (!isBSwapHWordElement(N010, Parts))
+    SDValue N10 = N1.getOperand(0);
+    if (!isBSwapHWordElement(N10, Parts))
       return SDValue();
-    SDValue N011 = N01.getOperand(1);
-    if (!isBSwapHWordElement(N011, Parts))
+    SDValue N11 = N1.getOperand(1);
+    if (!isBSwapHWordElement(N11, Parts))
       return SDValue();
   } else {
     // (or (or (or (and), (and)), (and)), (and))
diff --git a/llvm/test/CodeGen/X86/bswap_tree.ll b/llvm/test/CodeGen/X86/bswap_tree.ll
index 35a28af8557..c217879d438 100644
--- a/llvm/test/CodeGen/X86/bswap_tree.ll
+++ b/llvm/test/CodeGen/X86/bswap_tree.ll
@@ -13,32 +13,16 @@
 define i32 @test1(i32 %x) nounwind {
 ; CHECK-LABEL: test1:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; CHECK-NEXT:    movl %ecx, %edx
-; CHECK-NEXT:    andl $16711680, %edx # imm = 0xFF0000
-; CHECK-NEXT:    movl %ecx, %eax
-; CHECK-NEXT:    andl $-16777216, %eax # imm = 0xFF000000
-; CHECK-NEXT:    shll $8, %edx
-; CHECK-NEXT:    shrl $8, %eax
-; CHECK-NEXT:    bswapl %ecx
-; CHECK-NEXT:    shrl $16, %ecx
-; CHECK-NEXT:    orl %edx, %eax
-; CHECK-NEXT:    orl %ecx, %eax
+; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; CHECK-NEXT:    bswapl %eax
+; CHECK-NEXT:    roll $16, %eax
 ; CHECK-NEXT:    retl
 ;
 ; CHECK64-LABEL: test1:
 ; CHECK64:       # BB#0:
-; CHECK64-NEXT:    # kill: %EDI<def> %EDI<kill> %RDI<def>
-; CHECK64-NEXT:    movl %edi, %eax
-; CHECK64-NEXT:    andl $16711680, %eax # imm = 0xFF0000
-; CHECK64-NEXT:    movl %edi, %ecx
-; CHECK64-NEXT:    andl $-16777216, %ecx # imm = 0xFF000000
-; CHECK64-NEXT:    shll $8, %eax
-; CHECK64-NEXT:    shrl $8, %ecx
 ; CHECK64-NEXT:    bswapl %edi
-; CHECK64-NEXT:    shrl $16, %edi
-; CHECK64-NEXT:    orl %eax, %ecx
-; CHECK64-NEXT:    leal (%rcx,%rdi), %eax
+; CHECK64-NEXT:    roll $16, %edi
+; CHECK64-NEXT:    movl %edi, %eax
 ; CHECK64-NEXT:    retq
   %byte0 = and i32 %x, 255        ; 0x000000ff
   %byte1 = and i32 %x, 65280      ; 0x0000ff00
@@ -62,33 +46,16 @@ define i32 @test1(i32 %x) nounwind {
 define i32 @test2(i32 %x) nounwind {
 ; CHECK-LABEL: test2:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    pushl %esi
 ; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-NEXT:    movl %eax, %ecx
-; CHECK-NEXT:    shll $8, %ecx
-; CHECK-NEXT:    shrl $8, %eax
-; CHECK-NEXT:    movzwl %cx, %edx
-; CHECK-NEXT:    movzbl %al, %esi
-; CHECK-NEXT:    andl $-16777216, %ecx # imm = 0xFF000000
-; CHECK-NEXT:    andl $16711680, %eax # imm = 0xFF0000
-; CHECK-NEXT:    orl %edx, %esi
-; CHECK-NEXT:    orl %ecx, %eax
-; CHECK-NEXT:    orl %esi, %eax
-; CHECK-NEXT:    popl %esi
+; CHECK-NEXT:    bswapl %eax
+; CHECK-NEXT:    roll $16, %eax
 ; CHECK-NEXT:    retl
 ;
 ; CHECK64-LABEL: test2:
 ; CHECK64:       # BB#0:
-; CHECK64-NEXT:    movl %edi, %ecx
-; CHECK64-NEXT:    shll $8, %ecx
-; CHECK64-NEXT:    shrl $8, %edi
-; CHECK64-NEXT:    movzwl %cx, %edx
-; CHECK64-NEXT:    movzbl %dil, %eax
-; CHECK64-NEXT:    andl $-16777216, %ecx # imm = 0xFF000000
-; CHECK64-NEXT:    andl $16711680, %edi # imm = 0xFF0000
-; CHECK64-NEXT:    orl %edx, %eax
-; CHECK64-NEXT:    orl %ecx, %edi
-; CHECK64-NEXT:    orl %edi, %eax
+; CHECK64-NEXT:    bswapl %edi
+; CHECK64-NEXT:    roll $16, %edi
+; CHECK64-NEXT:    movl %edi, %eax
 ; CHECK64-NEXT:    retq
   %byte1 = shl  i32 %x, 8
   %byte0 = lshr i32 %x, 8