[InstCombine] collectBitParts - add fshl/fshr handling

Pulled from D87452, this is a fixed version of the collectBitParts fshl/fshr handling which as @nikic noticed wasn't checking for different providers or had correct bit ordering (which was hid by only testing shift amounts of bitwidth/2).

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

2 files changed, 37 insertions, 12 deletions

diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index 0861e99..61f4dff 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -2939,6 +2939,41 @@ collectBitParts(Value *V, bool MatchBSwaps, bool MatchBitReversals,
         Result->Provenance[i] = BitPart::Unset;
       return Result;
     }
+
+    // Handle intrinsic calls.
+    if (auto *II = dyn_cast<IntrinsicInst>(I)) {
+      Intrinsic::ID IntrinsicID = II->getIntrinsicID();
+
+      // Funnel 'double' shifts take 3 operands, 2 inputs and the shift
+      // amount (modulo).
+      // fshl(X,Y,Z): (X << (Z % BW)) | (Y >> (BW - (Z % BW)))
+      // fshr(X,Y,Z): (X << (BW - (Z % BW))) | (Y >> (Z % BW))
+      const APInt *Amt;
+      if ((IntrinsicID == Intrinsic::fshl || IntrinsicID == Intrinsic::fshr) &&
+          match(II->getArgOperand(2), m_APInt(Amt))) {
+
+        // We can treat fshr as a fshl by flipping the modulo amount.
+        unsigned ModAmt = Amt->urem(BitWidth);
+        if (IntrinsicID == Intrinsic::fshr)
+          ModAmt = BitWidth - ModAmt;
+
+        const auto &LHS = collectBitParts(II->getArgOperand(0), MatchBSwaps,
+                                          MatchBitReversals, BPS, Depth + 1);
+        const auto &RHS = collectBitParts(II->getArgOperand(1), MatchBSwaps,
+                                          MatchBitReversals, BPS, Depth + 1);
+
+        // Check we have both sources and they are from the same provider.
+        if (!LHS || !RHS || !LHS->Provider || LHS->Provider != RHS->Provider)
+          return Result;
+
+        Result = BitPart(LHS->Provider, BitWidth);
+        for (unsigned I = 0; I < (BitWidth - ModAmt); ++I)
+          Result->Provenance[I + ModAmt] = LHS->Provenance[I];
+        for (unsigned I = 0; I < ModAmt; ++I)
+          Result->Provenance[I] = RHS->Provenance[I + BitWidth - ModAmt];
+        return Result;
+      }
+    }
   }
 
   // Okay, we got to something that isn't a shift, 'or' or 'and'.  This must be
diff --git a/llvm/test/Transforms/InstCombine/bswap.ll b/llvm/test/Transforms/InstCombine/bswap.ll
index d81d4ac..3036ef2 100644
--- a/llvm/test/Transforms/InstCombine/bswap.ll
+++ b/llvm/test/Transforms/InstCombine/bswap.ll
@@ -431,11 +431,7 @@ declare i32 @llvm.fshr.i32(i32, i32, i32)
 
 define i32 @funnel_unary(i32 %abcd) {
 ; CHECK-LABEL: @funnel_unary(
-; CHECK-NEXT:    [[DABC:%.*]] = call i32 @llvm.fshl.i32(i32 [[ABCD:%.*]], i32 [[ABCD]], i32 24)
-; CHECK-NEXT:    [[BCDA:%.*]] = call i32 @llvm.fshl.i32(i32 [[ABCD]], i32 [[ABCD]], i32 8)
-; CHECK-NEXT:    [[DZBZ:%.*]] = and i32 [[DABC]], -16711936
-; CHECK-NEXT:    [[ZCZA:%.*]] = and i32 [[BCDA]], 16711935
-; CHECK-NEXT:    [[DCBA:%.*]] = or i32 [[DZBZ]], [[ZCZA]]
+; CHECK-NEXT:    [[DCBA:%.*]] = call i32 @llvm.bswap.i32(i32 [[ABCD:%.*]])
 ; CHECK-NEXT:    ret i32 [[DCBA]]
 ;
   %dabc = call i32 @llvm.fshl.i32(i32 %abcd, i32 %abcd, i32 24)
@@ -448,13 +444,7 @@ define i32 @funnel_unary(i32 %abcd) {
 
 define i32 @funnel_binary(i32 %abcd) {
 ; CHECK-LABEL: @funnel_binary(
-; CHECK-NEXT:    [[CDZZ:%.*]] = shl i32 [[ABCD:%.*]], 16
-; CHECK-NEXT:    [[DCDZ:%.*]] = call i32 @llvm.fshl.i32(i32 [[ABCD]], i32 [[CDZZ]], i32 24)
-; CHECK-NEXT:    [[ZZAB:%.*]] = lshr i32 [[ABCD]], 16
-; CHECK-NEXT:    [[ZABA:%.*]] = call i32 @llvm.fshl.i32(i32 [[ZZAB]], i32 [[ABCD]], i32 8)
-; CHECK-NEXT:    [[DCZZ:%.*]] = and i32 [[DCDZ]], -65536
-; CHECK-NEXT:    [[ZZBA:%.*]] = and i32 [[ZABA]], 65535
-; CHECK-NEXT:    [[DCBA:%.*]] = or i32 [[DCZZ]], [[ZZBA]]
+; CHECK-NEXT:    [[DCBA:%.*]] = call i32 @llvm.bswap.i32(i32 [[ABCD:%.*]])
 ; CHECK-NEXT:    ret i32 [[DCBA]]
 ;
   %cdzz = shl i32 %abcd, 16