[InstCombine] Catch more bswap cases missed due to zext and truncs.

[mcrosier]

Summary

This patch addresses PR27824: https://llvm.org/bugs/show_bug.cgi?id=27824

On AArch64 the i16 type is not legal type, so we end up with IR like what is included in the test cases.

Please take a look,
Chad

[gberry]

test/Transforms/InstCombine/bswap.ll

106

Shouldn't this be an i32 lshr of %conv and the second zext below deleted if you want the IR to look like it would for a target without i16 as a legal type?

[jmolloy]

LGTM!

[mcrosier]

test/Transforms/InstCombine/bswap.ll

106

With the exception of relabeling the variables and dropping irrelevant nsw/nuws, this is the code that is generated from this C test case when targeting AArch64:

unsigned short test16(unsigned short a) {
  unsigned short b = (a & 0xff00) >> 8;
  unsigned short c = (a & 0x00ff) << 8;
  return b | c;
}

Produces:

define i16 @test16(i16 %a) #0 {
entry:
  %conv = zext i16 %a to i32
  %shr11 = lshr i16 %a, 8
  %and3 = shl nuw nsw i32 %conv, 8
  %conv5 = zext i16 %shr11 to i32
  %or = or i32 %conv5, %and3
  %conv7 = trunc i32 %or to i16
  ret i16 %conv7
}

Regardless, I'm happy to add a test case similar to what you're suggesting.

[mcrosier]

Committed in r270853.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombiner::visitAnd(BinaryOperator &I) {
   }
 
   return Changed ? &I : nullptr;
 }
 
 /// Given an OR instruction, check to see if this is a bswap or bitreverse
 /// idiom. If so, insert the new intrinsic and return it.
 Instruction *InstCombiner::MatchBSwapOrBitReverse(BinaryOperator &I) {
+  Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
+
+  // Look through zero extends.
+  if (Instruction *Ext = dyn_cast<ZExtInst>(Op0))
+    Op0 = Ext->getOperand(0);
+
+  if (Instruction *Ext = dyn_cast<ZExtInst>(Op1))
+    Op1 = Ext->getOperand(0);
+
+  // (A | B) | C  and  A | (B | C)                  -> bswap if possible.
+  bool OrOfOrs = match(Op0, m_Or(m_Value(), m_Value())) ||
+                 match(Op1, m_Or(m_Value(), m_Value()));
+
+  // (A >> B) | (C << D)  and  (A << B) | (B >> C)  -> bswap if possible.
+  bool OrOfShifts = match(Op0, m_LogicalShift(m_Value(), m_Value())) &&
+                    match(Op1, m_LogicalShift(m_Value(), m_Value()));
+
+  // (A & B) | (C & D)                              -> bswap if possible.
+  bool OrOfAnds = match(Op0, m_And(m_Value(), m_Value())) &&
+                  match(Op1, m_And(m_Value(), m_Value()));
+
+  if (!OrOfOrs && !OrOfShifts && !OrOfAnds)
+    return nullptr;
+
   SmallVector<Instruction*, 4> Insts;
   if (!recognizeBitReverseOrBSwapIdiom(&I, true, false, Insts))
     return nullptr;
   Instruction *LastInst = Insts.pop_back_val();
   LastInst->removeFromParent();
 
   for (auto *Inst : Insts)
     Worklist.Add(Inst);
if (SelectInst *SI = dyn_cast<SelectInst>(Op0))
       if (Instruction *R = FoldOpIntoSelect(I, SI))
         return R;
 
     if (isa<PHINode>(Op0))
       if (Instruction *NV = FoldOpIntoPhi(I))
         return NV;
   }
 
-  Value *A = nullptr, *B = nullptr;
-  ConstantInt *C1 = nullptr, *C2 = nullptr;
-
-  // (A | B) | C  and  A | (B | C)                  -> bswap if possible.
-  bool OrOfOrs = match(Op0, m_Or(m_Value(), m_Value())) ||
-                 match(Op1, m_Or(m_Value(), m_Value()));
-  // (A >> B) | (C << D)  and  (A << B) | (B >> C)  -> bswap if possible.
-  bool OrOfShifts = match(Op0, m_LogicalShift(m_Value(), m_Value())) &&
-                    match(Op1, m_LogicalShift(m_Value(), m_Value()));
-  // (A & B) | (C & D)                              -> bswap if possible.
-  bool OrOfAnds = match(Op0, m_And(m_Value(), m_Value())) &&
-                  match(Op1, m_And(m_Value(), m_Value()));
-
-  if (OrOfOrs || OrOfShifts || OrOfAnds)
-    if (Instruction *BSwap = MatchBSwapOrBitReverse(I))
-      return BSwap;
+  // Given an OR instruction, check to see if this is a bswap.
 
+  Value *A = nullptr, *B = nullptr;
+  ConstantInt *C1 = nullptr, *C2 = nullptr;
+
   // (X^C)|Y -> (X|Y)^C iff Y&C == 0
   if (Op0->hasOneUse() &&
       match(Op0, m_Xor(m_Value(A), m_ConstantInt(C1))) &&
       MaskedValueIsZero(Op1, C1->getValue(), 0, &I)) {
     Value *NOr = Builder->CreateOr(A, Op1);
     NOr->takeName(Op0);
     return BinaryOperator::CreateXor(NOr, C1);
   }

lib/Transforms/Utils/Local.cpp

if (I->getOpcode() == Instruction::And &&
       if (!Res)
         return Result;
       Result = Res;
 
       for (unsigned i = 0; i < BitWidth; ++i, Bit <<= 1)
         // If the AndMask is zero for this bit, clear the bit.
         if ((AndMask & Bit) == 0)
           Result->Provenance[i] = BitPart::Unset;
+      return Result;
+    }
 
+    // If this is a zext instruction zero extend the result.
+    if (I->getOpcode() == Instruction::ZExt) {
+      auto &Res = collectBitParts(I->getOperand(0), MatchBSwaps,
+                                  MatchBitReversals, BPS);
+      if (!Res)
+        return Result;
+      Result = Res;
+
+      // Try and merge the two together.
+      Result = BitPart(Res->Provider, BitWidth);
+
+      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;
     }
   }
 
   // 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)
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())) {
+      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);
   if (!Res)
     return false;
   auto &BitProvenance = Res->Provenance;
 
   // Now, is the bit permutation correct for a bswap or a bitreverse? We can
   // only byteswap values with an even number of bytes.
-  bool OKForBSwap = BW % 16 == 0, OKForBitReverse = true;
-  for (unsigned i = 0; i < BW; ++i) {
-    OKForBSwap &= bitTransformIsCorrectForBSwap(BitProvenance[i], i, BW);
+  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, BW);
+        bitTransformIsCorrectForBitReverse(BitProvenance[i], i, DemandedBW);
   }
 
   Intrinsic::ID Intrin;
   if (OKForBSwap && MatchBSwaps)
     Intrin = Intrinsic::bswap;
   else if (OKForBitReverse && MatchBitReversals)
     Intrin = Intrinsic::bitreverse;
   else
     return false;
 
+  if (ITy != DemandedTy) {
+    Function *F = Intrinsic::getDeclaration(I->getModule(), Intrin, DemandedTy);
+    Value *Provider = Res->Provider;
+    IntegerType *ProviderTy = cast<IntegerType>(Provider->getType());
+    // We may need to truncate the provider.
+    if (DemandedTy != ProviderTy) {
+      auto *Trunc = CastInst::Create(Instruction::Trunc, Provider, DemandedTy, "trunc", I);
+      InsertedInsts.push_back(Trunc);
+      Provider = Trunc;
+    }
+    auto *CI = CallInst::Create(F, Provider, "rev", I);
+    InsertedInsts.push_back(CI);
+    auto *ExtInst = CastInst::Create(Instruction::ZExt, CI, ITy, "zext", I);
+    InsertedInsts.push_back(ExtInst);
+    return true;
+  }
+
   Function *F = Intrinsic::getDeclaration(I->getModule(), Intrin, ITy);
   InsertedInsts.push_back(CallInst::Create(F, Res->Provider, "rev", I));
   return true;
 }

test/Transforms/InstCombine/bswap.ll

define i32 @test7(i32 %x) {
   %or = or i32 %shl, %shr
   %and2 = shl i32 %or, 8
   %shl3 = and i32 %and2, -16711936
   %and4 = lshr i32 %or, 8
   %shr5 = and i32 %and4, 16711935
   %or6 = or i32 %shl3, %shr5
   ret i32 %or6
 }
+
+; CHECK-LABEL: @test8
+; CHECK: call i16 @llvm.bswap.i16(i16 %a)
+define i16 @test8(i16 %a) {
+entry:
+  %conv = zext i16 %a to i32
+  %shr = lshr i16 %a, 8
+  %shl = shl i32 %conv, 8
+  %conv1 = zext i16 %shr to i32
+  %or = or i32 %conv1, %shl
+  %conv2 = trunc i32 %or to i16
+  ret i16 %conv2
+}
+
+; CHECK-LABEL: @test9
+; CHECK: trunc i32 %a to i16
+; CHECK: call i16 @llvm.bswap.i16(i16 %trunc)
+define i16 @test9(i32 %a) {
+  %shr1 = lshr i32 %a, 8
+  %and1 = and i32 %shr1, 255
+  %and2 = shl i32 %a, 8
+  %shl1 = and i32 %and2, 65280
+  %or = or i32 %and1, %shl1
+  %conv = trunc i32 %or to i16
+  ret i16 %conv
+}