[InstCombine] recognizeBSwapOrBitReverseIdiom - recognise zext(bswap(trunc(x))) patterns (PR39793)

[RKSimon]

Summary

PR39793 demonstrated an issue where we fail to recognize 'partial' bswap patterns of the lower bytes of an integer source.

In fact, most of this is already in place collectBitParts suitably tags zero bits, so we just need to correctly handle this case by finding the zero'd upper bits and reducing the bswap pattern just to the active demanded bits.

[RKSimon]

llvm/test/Transforms/InstCombine/bswap.ll

297

This looks like a minor separate issue

[RKSimon]

llvm/test/Transforms/InstCombine/bswap.ll

297

D88429

[RKSimon]

rebase

[lebedev.ri]

llvm/lib/Transforms/Utils/Local.cpp

2992–2994

Why is this restricted to powers of two?

[RKSimon]

llvm/lib/Transforms/Utils/Local.cpp

2992–2994

Good catch - my solution was too focused on the original test cases - I'm looking at refactoring this now.

[lebedev.ri]

llvm/lib/Transforms/Utils/Local.cpp

2992–2994

I think this might suggest that there's second half of the puzzle missing,
i.e. should this also know how to widen/pad the inputs so that it is applicable for bswap?

[RKSimon]

llvm/lib/Transforms/Utils/Local.cpp

2992–2994

yes - shifted bswaps would be possible as a future development - a lot of the existing bitTransformIsCorrectForBSwap code would need to be refactored for that though.

[RKSimon]

I've rewritten the leading zero count code to support any type width and added a 'PR39793_bswap_u50_as_u16' test case

[lebedev.ri]

Looks good to me. @spatel ?

llvm/lib/Transforms/Utils/Local.cpp

2991–3002

Suggestion:

while(!BitProvenance.empty() && BitProvenance.back() < 0)
  BitProvenance = BitProvenance.drop_back();
if(BitProvenance.empty())
  ???
DemandedTy = IntegerType::get(I->getContext(), BitProvenance.size());

[spatel]

llvm/lib/Transforms/Utils/Local.cpp

2992

Should this check and the later one be for == BitPart::Unset rather than <0?
Can we assert that BitProvenance is not <0 other than Unset?

[lebedev.ri]

llvm/lib/Transforms/Utils/Local.cpp

3005

Super minor: this then should be if (BitProvenance.back() == BitPart::Unset) {

llvm/lib/Transforms/Utils/Local.cpp

bool llvm::recognizeBSwapOrBitReverseIdiom(
     SmallVectorImpl<Instruction *> &InsertedInsts) {
   if (Operator::getOpcode(I) != Instruction::Or)
     return false;
   if (!MatchBSwaps && !MatchBitReversals)
     return false;
   IntegerType *ITy = dyn_cast<IntegerType>(I->getType());
   if (!ITy || ITy->getBitWidth() > 128)
     return false;   // Can't do vectors or integers > 128 bits.
-  unsigned BW = ITy->getBitWidth();
 
-  unsigned DemandedBW = BW;
   IntegerType *DemandedTy = ITy;
-  if (I->hasOneUse()) {
-    if (TruncInst *Trunc = dyn_cast<TruncInst>(I->user_back())) {
+  if (I->hasOneUse())
+    if (auto *Trunc = dyn_cast<TruncInst>(I->user_back()))
       DemandedTy = cast<IntegerType>(Trunc->getType());
-      DemandedBW = DemandedTy->getBitWidth();
-    }
-  }
 
   // Try to find all the pieces corresponding to the bswap.
   std::map<Value *, Optional<BitPart>> BPS;
   auto Res = collectBitParts(I, MatchBSwaps, MatchBitReversals, BPS, 0);
   if (!Res)
     return false;
-  auto &BitProvenance = Res->Provenance;
+  ArrayRef<int8_t> BitProvenance = Res->Provenance;
+
+  // If the upper half is zero, then attempt to perform as a truncated op.
+  if (isPowerOf2_32(BitProvenance.size())) {
+    // Keep halving the demanded width while the upper half is zero.
+    while (BitProvenance.size() > 2) {
+      unsigned HalfBW = BitProvenance.size() / 2;
+      ArrayRef<int8_t> UpperBits = BitProvenance.slice(HalfBW);
+      if (!llvm::all_of(UpperBits, [](int8_t x) { return x < 0; }))
+        break;
+      BitProvenance = BitProvenance.slice(0, HalfBW);
+    }
+    DemandedTy = IntegerType::get(I->getContext(), BitProvenance.size());
+  }
 
   // Now, is the bit permutation correct for a bswap or a bitreverse? We can
   // only byteswap values with an even number of bytes.
+  unsigned DemandedBW = DemandedTy->getBitWidth();
   bool OKForBSwap = DemandedBW % 16 == 0, OKForBitReverse = true;
   for (unsigned i = 0; i < DemandedBW; ++i) {
     OKForBSwap &=
         bitTransformIsCorrectForBSwap(BitProvenance[i], i, DemandedBW);
     OKForBitReverse &=
         bitTransformIsCorrectForBitReverse(BitProvenance[i], i, DemandedBW);
   }
 

llvm/test/Transforms/InstCombine/bswap.ll

;
   %shl1 = and i32 %and2, 65280
   %or = or i32 %and1, %shl1
   %conv = trunc i32 %or to i16
   ret i16 %conv
 }
 
 define i64 @PR39793_bswap_u64_as_u32(i64 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u64_as_u32(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP0:%.*]], 24
-; CHECK-NEXT:    [[TMP3:%.*]] = and i64 [[TMP2]], 255
-; CHECK-NEXT:    [[TMP4:%.*]] = lshr i64 [[TMP0]], 8
-; CHECK-NEXT:    [[TMP5:%.*]] = and i64 [[TMP4]], 65280
-; CHECK-NEXT:    [[TMP6:%.*]] = or i64 [[TMP3]], [[TMP5]]
-; CHECK-NEXT:    [[TMP7:%.*]] = shl i64 [[TMP0]], 8
-; CHECK-NEXT:    [[TMP8:%.*]] = and i64 [[TMP7]], 16711680
-; CHECK-NEXT:    [[TMP9:%.*]] = or i64 [[TMP6]], [[TMP8]]
-; CHECK-NEXT:    [[TMP10:%.*]] = shl i64 [[TMP0]], 24
-; CHECK-NEXT:    [[TMP11:%.*]] = and i64 [[TMP10]], 4278190080
-; CHECK-NEXT:    [[TMP12:%.*]] = or i64 [[TMP9]], [[TMP11]]
-; CHECK-NEXT:    ret i64 [[TMP12]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i64 [[TMP0:%.*]] to i32
+; CHECK-NEXT:    [[REV:%.*]] = call i32 @llvm.bswap.i32(i32 [[TRUNC]])
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i32 [[REV]] to i64
+; CHECK-NEXT:    ret i64 [[TMP2]]
 ;
   %2 = lshr i64 %0, 24
   %3 = and i64 %2, 255
   %4 = lshr i64 %0, 8
   %5 = and i64 %4, 65280
   %6 = or i64 %3, %5
   %7 = shl i64 %0, 8
   %8 = and i64 %7, 16711680
   %9 = or i64 %6, %8
   %10 = shl i64 %0, 24
   %11 = and i64 %10, 4278190080
   %12 = or i64 %9, %11
   ret i64 %12
 }
 
 define i16 @PR39793_bswap_u64_as_u32_trunc(i64 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u64_as_u32_trunc(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP0:%.*]], 24
-; CHECK-NEXT:    [[TMP3:%.*]] = and i64 [[TMP2]], 255
-; CHECK-NEXT:    [[TMP4:%.*]] = lshr i64 [[TMP0]], 8
-; CHECK-NEXT:    [[TMP5:%.*]] = and i64 [[TMP4]], 65280
-; CHECK-NEXT:    [[TMP6:%.*]] = or i64 [[TMP3]], [[TMP5]]
-; CHECK-NEXT:    [[TMP7:%.*]] = trunc i64 [[TMP6]] to i16
-; CHECK-NEXT:    ret i16 [[TMP7]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i64 [[TMP0:%.*]] to i32
+; CHECK-NEXT:    [[REV:%.*]] = call i32 @llvm.bswap.i32(i32 [[TRUNC]])
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i32 [[REV]] to i16
+; CHECK-NEXT:    ret i16 [[TMP2]]
 ;
   %2 = lshr i64 %0, 24
   %3 = and i64 %2, 255
   %4 = lshr i64 %0, 8
   %5 = and i64 %4, 65280
   %6 = or i64 %3, %5
   %7 = shl i64 %0, 8
   %8 = and i64 %7, 16711680
   %9 = or i64 %6, %8
   %10 = shl i64 %0, 24
   %11 = and i64 %10, 4278190080
   %12 = or i64 %9, %11
   %13 = trunc i64 %12 to i16
   ret i16 %13
 }
 
 define i64 @PR39793_bswap_u64_as_u16(i64 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u64_as_u16(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[TMP3:%.*]] = and i64 [[TMP2]], 255
-; CHECK-NEXT:    [[TMP4:%.*]] = shl i64 [[TMP0]], 8
-; CHECK-NEXT:    [[TMP5:%.*]] = and i64 [[TMP4]], 65280
-; CHECK-NEXT:    [[TMP6:%.*]] = or i64 [[TMP3]], [[TMP5]]
-; CHECK-NEXT:    ret i64 [[TMP6]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i64 [[TMP0:%.*]] to i16
+; CHECK-NEXT:    [[REV:%.*]] = call i16 @llvm.bswap.i16(i16 [[TRUNC]])
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[REV]] to i64
+; CHECK-NEXT:    ret i64 [[TMP2]]
 ;
   %2 = lshr i64 %0, 8
   %3 = and i64 %2, 255
   %4 = shl i64 %0, 8
   %5 = and i64 %4, 65280
   %6 = or i64 %3, %5
   ret i64 %6
 }
 
 define i8 @PR39793_bswap_u64_as_u16_trunc(i64 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u64_as_u16_trunc(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 [[TMP2]] to i8
-; CHECK-NEXT:    ret i8 [[TMP3]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i64 [[TMP0:%.*]] to i16
+; CHECK-NEXT:    [[REV:%.*]] = lshr i16 [[TRUNC]], 8
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i16 [[REV]] to i8
+; CHECK-NEXT:    ret i8 [[TMP2]]
 ;
   %2 = lshr i64 %0, 8
   %3 = and i64 %2, 255
   %4 = shl i64 %0, 8
   %5 = and i64 %4, 65280
   %6 = or i64 %3, %5
   %7 = trunc i64 %6 to i8
   ret i8 %7
 }
 
 define i32 @PR39793_bswap_u32_as_u16(i32 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u32_as_u16(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i32 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], 255
-; CHECK-NEXT:    [[TMP4:%.*]] = shl i32 [[TMP0]], 8
-; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP4]], 65280
-; CHECK-NEXT:    [[TMP6:%.*]] = or i32 [[TMP3]], [[TMP5]]
-; CHECK-NEXT:    ret i32 [[TMP6]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[TMP0:%.*]] to i16
+; CHECK-NEXT:    [[REV:%.*]] = call i16 @llvm.bswap.i16(i16 [[TRUNC]])
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[REV]] to i32
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %2 = lshr i32 %0, 8
   %3 = and i32 %2, 255
   %4 = shl i32 %0, 8
   %5 = and i32 %4, 65280
   %6 = or i32 %3, %5
   ret i32 %6
 }
 
 define i8 @PR39793_bswap_u32_as_u16_trunc(i32 %0) {
 ; CHECK-LABEL: @PR39793_bswap_u32_as_u16_trunc(
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i32 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[TMP3:%.*]] = trunc i32 [[TMP2]] to i8
-; CHECK-NEXT:    ret i8 [[TMP3]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[TMP0:%.*]] to i16
+; CHECK-NEXT:    [[REV:%.*]] = lshr i16 [[TRUNC]], 8
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i16 [[REV]] to i8
+; CHECK-NEXT:    ret i8 [[TMP2]]
 ;
   %2 = lshr i32 %0, 8
   %3 = and i32 %2, 255
   %4 = shl i32 %0, 8
   %5 = and i32 %4, 65280
   %6 = or i32 %3, %5
   %7 = trunc i32 %6 to i8
   ret i8 %7