[InstCombine] Added pattern for recognising the construction of packed integers.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -11,10 +11,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/FloatingPointPredicateUtils.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/ConstantRange.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
@@ -3592,6 +3595,141 @@ static Value *foldOrOfInversions(BinaryOperator &I,
   return nullptr;
 }
 
+/// Match "shufflevector -> bitcast" or "extractelement -> zext -> shl" patterns
+/// which extract vector elements and pack them in the same relative positions.
+static bool matchSubIntegerPackFromVector(Value *V, Value *&Vec,
+                                          uint64_t &VecOffset,
+                                          SmallBitVector &Mask) {
+  static const auto m_ConstShlOrSelf = [](const auto &Base, uint64_t &ShlAmt) {
+    ShlAmt = 0;
+    return m_CombineOr(m_Shl(Base, m_ConstantInt(ShlAmt)), Base);
+  };
+
+  // First try to match extractelement -> zext -> shl
+  uint64_t VecIdx, ShlAmt;
+  if (match(V, m_ConstShlOrSelf(m_ZExtOrSelf(m_ExtractElt(
+                                    m_Value(Vec), m_ConstantInt(VecIdx))),
+                                ShlAmt))) {
+    auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
+    if (!VecTy)
+      return false;
+    auto *EltTy = dyn_cast<IntegerType>(VecTy->getElementType());
+    if (!EltTy)
+      return false;
+
+    const unsigned EltBitWidth = EltTy->getBitWidth();
+    const unsigned TargetBitWidth = V->getType()->getIntegerBitWidth();
+    if (TargetBitWidth % EltBitWidth != 0 || ShlAmt % EltBitWidth != 0)
+      return false;
+    const unsigned ShlEltAmt = ShlAmt / EltBitWidth;
+
+    if (ShlEltAmt > VecIdx)
+      return false;
+    VecOffset = VecIdx - ShlEltAmt;
+    Mask.resize(V->getType()->getIntegerBitWidth() / EltBitWidth);
+    Mask.set(ShlEltAmt);
+    return true;
+  }
+
+  // Now try to match a bitcasted subvector.
+  Instruction *DstVecI;
+  if (!match(V, m_BitCast(m_Instruction(DstVecI))))
+    return false;
+
+  auto *DstTy = dyn_cast<FixedVectorType>(DstVecI->getType());
+  if (!DstTy)
+    return false;
+
+  Mask.resize(DstTy->getNumElements());
+
+  // First check for a subvector obtained from a shufflevector.
+  if (isa<ShuffleVectorInst>(DstVecI)) {
+    Constant *ConstVec;
+    ArrayRef<int> ShuffleMask;
+    if (!match(DstVecI, m_Shuffle(m_Value(Vec), m_Constant(ConstVec),
+                                  m_Mask(ShuffleMask))))
+      return false;
+
+    auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
+    if (!VecTy)
+      return false;
+
+    const unsigned NumVecElts = VecTy->getNumElements();
+    bool FoundVecOffset = false;
+    for (unsigned Idx = 0; Idx < ShuffleMask.size(); ++Idx) {
+      if (ShuffleMask[Idx] == PoisonMaskElem)
+        return false;
+      const unsigned ShuffleIdx = ShuffleMask[Idx];
+      if (ShuffleIdx >= NumVecElts) {
+        const unsigned ConstIdx = ShuffleIdx - NumVecElts;
+        auto *ConstElt =
+            dyn_cast<ConstantInt>(ConstVec->getAggregateElement(ConstIdx));
+        if (!ConstElt || !ConstElt->isNullValue())
+          return false;
+        continue;
+      }
+
+      if (FoundVecOffset) {
+        if (VecOffset + Idx != ShuffleIdx)
+          return false;
+      } else {
+        if (ShuffleIdx < Idx)
+          return false;
+        VecOffset = ShuffleIdx - Idx;
+        FoundVecOffset = true;
+      }
+      Mask.set(Idx);
+    }
+    return FoundVecOffset;
+  }
+
+  // Check for a subvector obtained as an (insertelement V, 0, idx)
+  uint64_t InsertIdx;
+  if (!match(DstVecI,
+             m_InsertElt(m_Value(Vec), m_Zero(), m_ConstantInt(InsertIdx))))
+    return false;
+
+  auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
+  if (!VecTy)
+    return false;
+  VecOffset = 0;
+  bool AlreadyInsertedMaskedElt = Mask.test(InsertIdx);
+  Mask.set();
+  if (!AlreadyInsertedMaskedElt)
+    Mask.reset(InsertIdx);
+  return true;
+}
+
+/// Try to fold the or of two scalar integers whose contents are packed elements
+/// of the same vector.
+Instruction *foldIntegerPackFromVector(Instruction &I,
+                                       InstCombiner::BuilderTy &Builder) {
+  assert(I.getOpcode() == Instruction::Or);
+  Value *LhsVec, *RhsVec;
+  uint64_t LhsVecOffset, RhsVecOffset;
+  SmallBitVector Mask;
+  if (!matchSubIntegerPackFromVector(I.getOperand(0), LhsVec, LhsVecOffset,
+                                     Mask))
+    return nullptr;
+  if (!matchSubIntegerPackFromVector(I.getOperand(1), RhsVec, RhsVecOffset,
+                                     Mask))
+    return nullptr;
+  if (LhsVec != RhsVec || LhsVecOffset != RhsVecOffset)
+    return nullptr;
+
+  // Convert into shufflevector -> bitcast;
+  const unsigned ZeroVecIdx =
+      cast<FixedVectorType>(LhsVec->getType())->getNumElements();
+  SmallVector<int> ShuffleMask(Mask.size(), ZeroVecIdx);
+  for (unsigned Idx : Mask.set_bits())
+    ShuffleMask[Idx] = LhsVecOffset + Idx;
+
+  Value *MaskedVec = Builder.CreateShuffleVector(
+      LhsVec, Constant::getNullValue(LhsVec->getType()), ShuffleMask,
+      I.getName() + ".v");
+  return CastInst::Create(Instruction::BitCast, MaskedVec, I.getType());
+}
+
 // A decomposition of ((X & Mask) * Factor). The NUW / NSW bools
 // track these properities for preservation. Note that we can decompose
 // equivalent select form of this expression (e.g. (!(X & Mask) ? 0 : Mask *
@@ -3688,6 +3826,9 @@ Instruction *InstCombinerImpl::visitOr(BinaryOperator &I) {
   if (Instruction *X = foldComplexAndOrPatterns(I, Builder))
     return X;
 
+  if (Instruction *X = foldIntegerPackFromVector(I, Builder))
+    return X;
+
   // (A & B) | (C & D) -> A ^ D where A == ~C && B == ~D
   // (A & B) | (C & D) -> A ^ C where A == ~D && B == ~C
   if (Value *V = foldOrOfInversions(I, Builder))

llvm/test/Transforms/InstCombine/or-packed-int-vecs.ll

@@ -0,0 +1,109 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=instcombine %s | FileCheck %s
+
+define i32 @bitcast.v2i(<4 x i8> %v) {
+; CHECK-LABEL: define i32 @bitcast.v2i(
+; CHECK-SAME: <4 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[X_3:%.*]] = bitcast <4 x i8> [[V]] to i32
+; CHECK-NEXT:    ret i32 [[X_3]]
+;
+  %v.0 = extractelement <4 x i8> %v, i32 0
+  %z.0 = zext i8 %v.0 to i32
+
+  %v.1 = extractelement <4 x i8> %v, i32 1
+  %z.1 = zext i8 %v.1 to i32
+  %s.1 = shl i32 %z.1, 8
+  %x.1 = or i32 %z.0, %s.1
+
+  %v.2 = extractelement <4 x i8> %v, i32 2
+  %z.2 = zext i8 %v.2 to i32
+  %s.2 = shl i32 %z.2, 16
+  %x.2 = or i32 %x.1, %s.2
+
+  %v.3 = extractelement <4 x i8> %v, i32 3
+  %z.3 = zext i8 %v.3 to i32
+  %s.3 = shl i32 %z.3, 24
+  %x.3 = or i32 %x.2, %s.3
+
+  ret i32 %x.3
+}
+
+define i32 @bitcast.v2i.tree(<4 x i8> %v) {
+; CHECK-LABEL: define i32 @bitcast.v2i.tree(
+; CHECK-SAME: <4 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[X:%.*]] = bitcast <4 x i8> [[V]] to i32
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %v.0 = extractelement <4 x i8> %v, i32 0
+  %z.0 = zext i8 %v.0 to i32
+
+  %v.1 = extractelement <4 x i8> %v, i32 1
+  %z.1 = zext i8 %v.1 to i32
+  %s.1 = shl i32 %z.1, 8
+  %x.1 = or i32 %z.0, %s.1
+
+  %v.2 = extractelement <4 x i8> %v, i32 2
+  %z.2 = zext i8 %v.2 to i32
+  %s.2 = shl i32 %z.2, 16
+
+  %v.3 = extractelement <4 x i8> %v, i32 3
+  %z.3 = zext i8 %v.3 to i32
+  %s.3 = shl i32 %z.3, 24
+  %x.3 = or i32 %s.2, %s.3
+
+  %x = or i32 %x.1, %x.3
+
+  ret i32 %x
+}
+
+define i32 @extract.i32(<8 x i8> %v) {
+; CHECK-LABEL: define i32 @extract.i32(
+; CHECK-SAME: <8 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[X_3_V_EXTRACT:%.*]] = shufflevector <8 x i8> [[V]], <8 x i8> poison, <8 x i32> <i32 3, i32 4, i32 5, i32 6, i32 poison, i32 poison, i32 poison, i32 poison>
+; CHECK-NEXT:    [[X_3_V_BC:%.*]] = bitcast <8 x i8> [[X_3_V_EXTRACT]] to <2 x i32>
+; CHECK-NEXT:    [[X_3_V_EXTRACT1:%.*]] = extractelement <2 x i32> [[X_3_V_BC]], i64 0
+; CHECK-NEXT:    ret i32 [[X_3_V_EXTRACT1]]
+;
+  %v.0 = extractelement <8 x i8> %v, i32 3
+  %z.0 = zext i8 %v.0 to i32
+
+  %v.1 = extractelement <8 x i8> %v, i32 4
+  %z.1 = zext i8 %v.1 to i32
+  %s.1 = shl i32 %z.1, 8
+  %x.1 = or i32 %z.0, %s.1
+
+  %v.2 = extractelement <8 x i8> %v, i32 5
+  %z.2 = zext i8 %v.2 to i32
+  %s.2 = shl i32 %z.2, 16
+  %x.2 = or i32 %x.1, %s.2
+
+  %v.3 = extractelement <8 x i8> %v, i32 6
+  %z.3 = zext i8 %v.3 to i32
+  %s.3 = shl i32 %z.3, 24
+  %x.3 = or i32 %x.2, %s.3
+
+  ret i32 %x.3
+}
+
+define i32 @partial(<4 x i8> %v) {
+; CHECK-LABEL: define i32 @partial(
+; CHECK-SAME: <4 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[X_3_V1:%.*]] = insertelement <4 x i8> [[V]], i8 0, i64 2
+; CHECK-NEXT:    [[X_3:%.*]] = bitcast <4 x i8> [[X_3_V1]] to i32
+; CHECK-NEXT:    ret i32 [[X_3]]
+;
+  %v.0 = extractelement <4 x i8> %v, i32 0
+  %z.0 = zext i8 %v.0 to i32
+
+  %v.1 = extractelement <4 x i8> %v, i32 1
+  %z.1 = zext i8 %v.1 to i32
+  %s.1 = shl i32 %z.1, 8
+  %x.1 = or i32 %z.0, %s.1
+
+  %v.3 = extractelement <4 x i8> %v, i32 3
+  %z.3 = zext i8 %v.3 to i32
+  %s.3 = shl i32 %z.3, 24
+  %x.3 = or i32 %x.1, %s.3
+
+  ret i32 %x.3
+}