[InstCombine] Fold or(zext(bswap(x)),shl(zext(bswap(y)),bw/2)) -> bswap(or(zext(x),shl(zext(y), bw/2))

[RKSimon]

Summary

This adds a general combine that can be used to fold:

or(zext(OP(x)), shl(zext(OP(y)),bw/2))
-->
OP(or(zext(x), shl(zext(y),bw/2)))

Allowing us to widen 'concat-able' style or+zext patterns - I've just set this up for BSWAP but we could use this for other similar ops (BITREVERSE for instance).

We already do something similar for binop(bswap(x),bswap(y)) --> bswap(binop(x,y))

Fixes PR45715

[RKSimon]

rebase

[lebedev.ri]

This looks good to me, but i'm not sure this will work more generally.
I.e., i will be surprized if the or tree reaches us in the proper form,
without reassociation into some other form we aren't expecting.
Especially for larger tree depths.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

static Instruction *matchRotate(Instruction &Or) {
 
   bool IsFshl = (!SubIsOnLHS && ShiftOpcode0 == BinaryOperator::Shl) ||
                 (SubIsOnLHS && ShiftOpcode1 == BinaryOperator::Shl);
   Intrinsic::ID IID = IsFshl ? Intrinsic::fshl : Intrinsic::fshr;
   Function *F = Intrinsic::getDeclaration(Or.getModule(), IID, Or.getType());
   return IntrinsicInst::Create(F, { ShVal, ShVal, ShAmt });
 }
 
+/// Attempt to combine or(zext(x),shl(zext(y),bw/2) concat packing patterns.
+static Instruction *matchOrConcat(Instruction &Or,
+                                  InstCombiner::BuilderTy &Builder) {
+  assert(Or.getOpcode() == Instruction::Or && "bswap requires an 'or'");
+  Value *Op0 = Or.getOperand(0), *Op1 = Or.getOperand(1);
+  Type *Ty = Or.getType();
+
+  unsigned Width = Ty->getScalarSizeInBits();
+  if ((Width & 1) != 0)
+    return nullptr;
+  unsigned HalfWidth = Width / 2;
+
+  // Canonicalize zext (lower half) to LHS.
+  if (!isa<ZExtInst>(Op0))
+    std::swap(Op0, Op1);
+
+  // Find lower/upper half.
+  Value *LowerSrc, *ShlVal, *UpperSrc;
+  const APInt *C;
+  if (!match(Op0, m_OneUse(m_ZExt(m_Value(LowerSrc)))) ||
+      !match(Op1, m_OneUse(m_Shl(m_Value(ShlVal), m_APInt(C)))) ||
+      !match(ShlVal, m_OneUse(m_ZExt(m_Value(UpperSrc)))))
+    return nullptr;
+  if (*C != HalfWidth || LowerSrc->getType() != UpperSrc->getType() ||
+      LowerSrc->getType()->getScalarSizeInBits() != HalfWidth)
+    return nullptr;
+
+  // Find matching bswap instructions.
+  // TODO: Add more patterns (bitreverse?)
+  Value *LowerBSwap, *UpperBSwap;
+  if (!match(LowerSrc, m_BSwap(m_Value(LowerBSwap))) ||
+      !match(UpperSrc, m_BSwap(m_Value(UpperBSwap))))
+    return nullptr;
+
+  // Push the concat down, swapping the lower/upper sources.
+  Value *NewLower = Builder.CreateZExt(UpperBSwap, Ty);
+  Value *NewUpper = Builder.CreateZExt(LowerBSwap, Ty);
+  NewUpper = Builder.CreateShl(NewUpper, HalfWidth);
+  Value *BinOp = Builder.CreateOr(NewLower, NewUpper);
+  Function *F = Intrinsic::getDeclaration(Or.getModule(), Intrinsic::bswap, Ty);
+  return Builder.CreateCall(F, BinOp);
+}
+
 /// If all elements of two constant vectors are 0/-1 and inverses, return true.
 static bool areInverseVectorBitmasks(Constant *C1, Constant *C2) {
   unsigned NumElts = cast<VectorType>(C1->getType())->getNumElements();
   for (unsigned i = 0; i != NumElts; ++i) {
     Constant *EltC1 = C1->getAggregateElement(i);
     Constant *EltC2 = C2->getAggregateElement(i);
     if (!EltC1 || !EltC2)
       return false;
if (Instruction *FoldedLogic = foldBinOpIntoSelectOrPhi(I))
     return FoldedLogic;
 
   if (Instruction *BSwap = matchBSwap(I))
     return BSwap;
 
   if (Instruction *Rotate = matchRotate(I))
     return Rotate;
 
+  if (Instruction *Concat = matchOrConcat(I, Builder))
+    return replaceInstUsesWith(I, Concat);
+
   Value *X, *Y;
   const APInt *CV;
   if (match(&I, m_c_Or(m_OneUse(m_Xor(m_Value(X), m_APInt(CV))), m_Value(Y))) &&
       !CV->isAllOnesValue() && MaskedValueIsZero(Y, *CV, 0, &I)) {
     // (X ^ C) | Y -> (X | Y) ^ C iff Y & C == 0
     // The check for a 'not' op is for efficiency (if Y is known zero --> ~X).
     Value *Or = Builder.CreateOr(X, Y);
     return BinaryOperator::CreateXor(Or, ConstantInt::get(I.getType(), *CV));

llvm/test/Transforms/InstCombine/or-concat.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -instcombine < %s | FileCheck %s
 
 ;
 ; Tests for combining concat-able ops:
 ; or(zext(OP(x)), shl(zext(OP(y)),bw/2))
 ; -->
 ; OP(or(zext(x), shl(zext(y),bw/2)))
 ;
 
 ; BSWAP
 
 ; PR45715
 define i64 @concat_bswap32_unary_split(i64 %a0) {
 ; CHECK-LABEL: @concat_bswap32_unary_split(
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[A0:%.*]], 32
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 [[TMP1]] to i32
-; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 [[A0]] to i32
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[TMP2]])
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[TMP3]])
-; CHECK-NEXT:    [[TMP6:%.*]] = zext i32 [[TMP4]] to i64
-; CHECK-NEXT:    [[TMP7:%.*]] = zext i32 [[TMP5]] to i64
-; CHECK-NEXT:    [[TMP8:%.*]] = shl nuw i64 [[TMP7]], 32
-; CHECK-NEXT:    [[TMP9:%.*]] = or i64 [[TMP8]], [[TMP6]]
-; CHECK-NEXT:    ret i64 [[TMP9]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @llvm.bswap.i64(i64 [[A0:%.*]])
+; CHECK-NEXT:    ret i64 [[TMP1]]
 ;
   %1 = lshr i64 %a0, 32
   %2 = trunc i64 %1 to i32
   %3 = trunc i64 %a0 to i32
   %4 = tail call i32 @llvm.bswap.i32(i32 %2)
   %5 = tail call i32 @llvm.bswap.i32(i32 %3)
   %6 = zext i32 %4 to i64
   %7 = zext i32 %5 to i64
   %8 = shl nuw i64 %7, 32
   %9 = or i64 %6, %8
   ret i64 %9
 }
 
 define i64 @concat_bswap32_unary_flip(i64 %a0) {
 ; CHECK-LABEL: @concat_bswap32_unary_flip(
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[A0:%.*]], 32
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 [[TMP1]] to i32
-; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 [[A0]] to i32
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[TMP2]])
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[TMP3]])
-; CHECK-NEXT:    [[TMP6:%.*]] = zext i32 [[TMP4]] to i64
-; CHECK-NEXT:    [[TMP7:%.*]] = zext i32 [[TMP5]] to i64
-; CHECK-NEXT:    [[TMP8:%.*]] = shl nuw i64 [[TMP6]], 32
-; CHECK-NEXT:    [[TMP9:%.*]] = or i64 [[TMP8]], [[TMP7]]
-; CHECK-NEXT:    ret i64 [[TMP9]]
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i64 [[A0]], 32
+; CHECK-NEXT:    [[TMP3:%.*]] = or i64 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = call i64 @llvm.bswap.i64(i64 [[TMP3]])
+; CHECK-NEXT:    ret i64 [[TMP4]]
 ;
   %1 = lshr i64 %a0, 32
   %2 = trunc i64 %1 to i32
   %3 = trunc i64 %a0 to i32
   %4 = tail call i32 @llvm.bswap.i32(i32 %2)
   %5 = tail call i32 @llvm.bswap.i32(i32 %3)
   %6 = zext i32 %4 to i64
   %7 = zext i32 %5 to i64
   %8 = shl nuw i64 %6, 32
   %9 = or i64 %7, %8
   ret i64 %9
 }
 
 define i64 @concat_bswap32_binary(i32 %a0, i32 %a1) {
 ; CHECK-LABEL: @concat_bswap32_binary(
-; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[A0:%.*]])
-; CHECK-NEXT:    [[TMP2:%.*]] = tail call i32 @llvm.bswap.i32(i32 [[A1:%.*]])
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP1]] to i64
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[TMP2]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = shl nuw i64 [[TMP4]], 32
-; CHECK-NEXT:    [[TMP6:%.*]] = or i64 [[TMP5]], [[TMP3]]
-; CHECK-NEXT:    ret i64 [[TMP6]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[A1:%.*]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i32 [[A0:%.*]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = shl nuw i64 [[TMP2]], 32
+; CHECK-NEXT:    [[TMP4:%.*]] = or i64 [[TMP3]], [[TMP1]]
+; CHECK-NEXT:    [[TMP5:%.*]] = call i64 @llvm.bswap.i64(i64 [[TMP4]])
+; CHECK-NEXT:    ret i64 [[TMP5]]
 ;
   %1 = tail call i32 @llvm.bswap.i32(i32 %a0)
   %2 = tail call i32 @llvm.bswap.i32(i32 %a1)
   %3 = zext i32 %1 to i64
   %4 = zext i32 %2 to i64
   %5 = shl nuw i64 %4, 32
   %6 = or i64 %3, %5
   ret i64 %6