[llvm] [InstCombine] Fold zext into de-interleaving (factor=2) instructions (PR #195330)

[Min-Yih Hsu via llvm-commits]

================
@@ -3296,3 +3296,168 @@ Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
 
   return MadeChange ? &SVI : nullptr;
 }
+
+/// Given the following de-interleaving shufflevectors and the consuming zexts:
+/// ```
+/// %f0 = shufflevector <8 x i32> %v, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
+/// %f1 = shufflevector <8 x i32> %v, <4 x i32> <i32 1, i32 3, i32 5, i32 7>
+/// %z0 = zext <4 x i32> %f0 to <4 x i64>
+/// %z1 = zext <4 x i32> %f1 to <4 x i64>
+/// ```
+/// We can actually bitcast the input value, `%v` first before replacing zexts
+/// with simple arithmetics on this new bitcast:
+/// ```
+/// %bc = bitcast <8 x i32> %v to <4 x i64>
+//  %z0 = and <4 x i64> %bc, splat (i64 4294967295)
+//  %z1 = lshr <4 x i64> %bc, splat (i64 32)
+/// ```
+/// This transformation is almost always benefitial as shufflevector is more
+/// expensive than normal arithmetics.
+Instruction *
+InstCombinerImpl::foldExtractionOfVectorDeinterleave(ZExtInst &RootZExt) {
+  // This pattern involves bitcast that is not compatible with big endian.
+  if (DL.isBigEndian())
+    return nullptr;
+
+  // The actual value that got de-interleaved.
+  Value *DIV;
+
+  using namespace PatternMatch;
+  Value *SVI = nullptr, *DI = nullptr;
+  if (!match(&RootZExt,
+             m_ZExt(m_CombineOr(
+                 m_ExtractValue(m_Value(DI, m_Deinterleave2(m_Value(DIV)))),
+                 m_Value(SVI, m_Shuffle(m_Value(), m_Value()))))))
+    return nullptr;
+
+  auto isDeinterleaveShuffle =
+      [](Instruction *I) -> std::pair<Value *, unsigned> {
+    Value *V;
+    ArrayRef<int> ShuffleMask;
+    unsigned Index;
+    if (match(I, m_Shuffle(m_Value(V), m_Undef(), m_Mask(ShuffleMask))) &&
+        isa<FixedVectorType>(V->getType())) {
+      unsigned NumInputElements =
+          cast<VectorType>(V->getType())->getElementCount().getFixedValue();
+      if (ShuffleVectorInst::isDeInterleaveMaskOfFactor(ShuffleMask, 2,
+                                                        Index) &&
+          Index < 2 &&
+          ShuffleVectorInst::isSingleSourceMask(ShuffleMask,
+                                                NumInputElements) &&
+          ShuffleMask.size() * 2 == NumInputElements)
+        return {V, Index};
+    }
+    return {nullptr, UINT_MAX};
+  };
+
+  // Validate either the shufflevector or the vector.deinterleave2 and obtain
+  // the value they're de-interleaving.
+  if (SVI) {
+    // We will find other shufflevectors later.
+    DIV = isDeinterleaveShuffle(cast<Instruction>(SVI)).first;
+    if (!DIV)
+      return nullptr;
+  } else {
+    // We should already capture the value that got de-interleaved (i.e. DIV).
+    assert(DI && DIV);
+    if (!all_of(DI->users(), [](User *Usr) -> bool {
+          auto *EV = dyn_cast<ExtractValueInst>(Usr);
+          return EV && EV->getNumIndices() == 1;
+        }))
+      return nullptr;
+  }
+
+  auto *InputVecTy = dyn_cast<VectorType>(DIV->getType());
+  if (!InputVecTy)
+    return nullptr;
+  auto *InElementTy = dyn_cast<IntegerType>(InputVecTy->getElementType());
+  if (!InElementTy)
+    return nullptr;
+  if (!InputVecTy->getElementCount().isKnownEven())
+    return nullptr;
+
+  // {Field instruction, Field index}
+  SmallVector<std::pair<Instruction *, unsigned>, 4> Fields;
+  if (SVI) {
+    for (auto *Usr : DIV->users()) {
+      auto *FieldI = dyn_cast<Instruction>(Usr);
+      if (!FieldI)
+        continue;
+      auto [V, Index] = isDeinterleaveShuffle(FieldI);
+      if (V != DIV)
+        continue;
+      assert(Index < 2);
+      Fields.push_back({FieldI, Index});
+    }
+  } else {
+    // llvm.vector.deinterleave2.
+    for (User *Field : DI->users()) {
+      auto *FieldI = cast<ExtractValueInst>(Field);
+      unsigned FieldIdx = *FieldI->idx_begin();
+      assert(FieldIdx < 2);
+      Fields.push_back({FieldI, FieldIdx});
+    }
+  }
+
+  // We commit the transformation only if all the field users can be replaced,
+  // otherwise the primary de-interleaving construction, regardless of
+  // llvm.vector.deinterleave2 or shufflevectors, will still be there.
+  SmallVector<ZExtInst *, 4> Field0Replacements;
+  SmallVector<ZExtInst *, 4> Field1Replacements;
+  for (auto [Field, FieldIdx] : Fields) {
+    for (User *FieldUsr : Field->users()) {
+      auto *ZExt = dyn_cast<ZExtInst>(FieldUsr);
+      if (!ZExt)
+        return nullptr;
+      // Only if it's doubling the element size.
+      if (ZExt->getDestTy() != ZExt->getSrcTy()->getExtendedType())
+        return nullptr;
+      if (FieldIdx)
+        Field1Replacements.push_back(ZExt);
+      else
+        Field0Replacements.push_back(ZExt);
+    }
+  }
+
+  // Double the element size but half the vector length.
+  auto *BitcastedTy = VectorType::getExtendedElementVectorType(InputVecTy);
+  BitcastedTy = VectorType::getHalfElementsVectorType(BitcastedTy);
+  // Since we're going to "merge" lanes via bitcast, we need to freeze any
+  // potential poison lanes first.
+  Value *Freeze = Builder.CreateFreeze(DIV);
+  Value *Bitcast = Builder.CreateBitCast(Freeze, BitcastedTy);
+  unsigned InElementBitWidth = InElementTy->getBitWidth();
+  auto Mask = APInt::getLowBitsSet(InElementBitWidth * 2, InElementBitWidth);
+  Value *NewField0 = Builder.CreateAnd(Bitcast, Mask);
+  Value *NewField1 = Builder.CreateLShr(Bitcast, InElementBitWidth);
+
+  Value *RootNewField = nullptr;
+  for (ZExtInst *I : Field0Replacements) {
+    // The root ZExt will be replaced by the InstCombiner upon
+    // returning from this function.
+    if (I == &RootZExt) {
+      assert(!RootNewField);
+      RootNewField = NewField0;
+      continue;
+    }
+    replaceInstUsesWith(*I, NewField0);
+    // Make sure non-root ZExt are in the worklist so that they
+    // can be removed in the following iterations.
+    addToWorklist(I);
+  }
+  for (ZExtInst *I : Field1Replacements) {
+    if (I == &RootZExt) {
+      assert(!RootNewField);
+      RootNewField = NewField1;
+      continue;
+    }
+    replaceInstUsesWith(*I, NewField1);
+    addToWorklist(I);
+  }
+
+  assert(RootNewField && "cannot find replacement for the root zext");
+  auto *NewRootInst = cast<Instruction>(RootNewField);
+  // The InstCombiner expects a parent-less instruction.
----------------
mshockwave wrote:

Yes that works. The code is now updated.

https://github.com/llvm/llvm-project/pull/195330