[InstCombine] collectBitParts - add fshl/fshr handling

[RKSimon]

Summary

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).

[nikic]

This LGTM now, thanks!

llvm/lib/Transforms/Utils/Local.cpp

2947

I would if (auto *II = dyn_cast<IntrinsicInst>(I)) and then II->getInrinsicID() here. Doesn't seem necessary to explicitly go through the Function.

2950

modulo -> modulo the bitwidth?

llvm/lib/Transforms/Utils/Local.cpp

if (I->getOpcode() == Instruction::ZExt) {
       auto NarrowBitWidth =
           cast<IntegerType>(cast<ZExtInst>(I)->getSrcTy())->getBitWidth();
       for (unsigned i = 0; i < NarrowBitWidth; ++i)
         Result->Provenance[i] = Res->Provenance[i];
       for (unsigned i = NarrowBitWidth; i < BitWidth; ++i)
         Result->Provenance[i] = BitPart::Unset;
       return Result;
     }
+
+    if (auto *CI = dyn_cast<CallInst>(I)) {
+      Value *Callee = CI->getCalledOperand();
+      // Handle intrinsic calls.
+      if (auto *F = dyn_cast<Function>(Callee)) {
+        Intrinsic::ID IntrinsicID = F->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(CI->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(CI->getArgOperand(0), MatchBSwaps,
+                                            MatchBitReversals, BPS, Depth + 1);
+          const auto &RHS = collectBitParts(CI->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
   // the input value to the bswap/bitreverse.
   Result = BitPart(V, BitWidth);
   for (unsigned i = 0; i < BitWidth; ++i)
     Result->Provenance[i] = i;
   return Result;

llvm/test/Transforms/InstCombine/bswap.ll

;
   ret i32 %cast
 }
 
 declare i32 @llvm.fshl.i32(i32, i32, i32)
 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)
   %bcda = call i32 @llvm.fshr.i32(i32 %abcd, i32 %abcd, i32 24)
   %dzbz = and i32 %dabc, -16711936
   %zcza = and i32 %bcda,  16711935
   %dcba = or i32 %dzbz, %zcza
   ret i32 %dcba
 }
 
 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
   %dcdz = call i32 @llvm.fshl.i32(i32 %abcd, i32 %cdzz, i32 24)
   %zzab = lshr i32 %abcd, 16
   %zaba = call i32 @llvm.fshr.i32(i32 %zzab, i32 %abcd, i32 24)
   %dczz = and i32 %dcdz, -65536
   %zzba = and i32 %zaba,  65535
   %dcba = or i32 %dczz, %zzba
   ret i32 %dcba
 }