[llvm] [VectorCombine] Fold Deinterleave/Interleave Pairs (PR #211022)

[Graham Hunter via llvm-commits]

================
@@ -5882,6 +5883,265 @@ bool VectorCombine::foldInsExtVectorToShuffle(Instruction &I) {
   return true;
 }
 
+namespace {
+
+class InterleavedElementwiseChain {
+  struct ElementwiseStep {
+    SmallVector<Instruction *, 8> Insts;
+    unsigned ChainOperand;
+
+    ElementwiseStep(ArrayRef<Instruction *> Insts, unsigned ChainOperand)
+        : Insts(Insts.begin(), Insts.end()), ChainOperand(ChainOperand) {}
+  };
+
+  IRBuilderBase &Builder;
+  Value *Root;
+  SmallVector<ElementwiseStep, 4> Steps;
+
+  static unsigned getNumDataOperands(Instruction *Inst) {
+    if (auto *II = dyn_cast<IntrinsicInst>(Inst))
+      return II->arg_size();
+    return Inst->getNumOperands();
+  }
+
+  static Value *getSplatOrScalar(Value *V) {
+    return isa<VectorType>(V->getType()) ? getSplatValue(V) : V;
+  }
+
+  static bool isSupportedElementwise(Instruction *Inst) {
+    auto *ResultTy = dyn_cast<VectorType>(Inst->getType());
+    if (!ResultTy)
+      return false;
+
+    if (auto *II = dyn_cast<IntrinsicInst>(Inst)) {
+      if (II->hasOperandBundles() ||
+          !isTriviallyVectorizable(II->getIntrinsicID()))
+        return false;
+    } else if (!isa<BinaryOperator, UnaryOperator, CastInst, CmpInst,
+                    SelectInst, FreezeInst>(Inst)) {
+      return false;
+    }
+
+    // Reject operations that change the element-count.
+    for (unsigned Op = 0, E = getNumDataOperands(Inst); Op != E; ++Op) {
+      auto *OperandTy = dyn_cast<VectorType>(Inst->getOperand(Op)->getType());
+      if (OperandTy &&
+          OperandTy->getElementCount() != ResultTy->getElementCount())
+        return false;
+    }
+
+    return true;
+  }
+
+  Value *createWideInstruction(Instruction *NarrowInst,
+                               ArrayRef<Value *> NewOperands,
+                               VectorType *WideResultTy) {
+    if (isa<BinaryOperator, UnaryOperator>(NarrowInst))
+      return Builder.CreateNAryOp(NarrowInst->getOpcode(), NewOperands);
+    if (auto *Cast = dyn_cast<CastInst>(NarrowInst))
+      return Builder.CreateCast(Cast->getOpcode(), NewOperands[0],
+                                WideResultTy);
+    if (auto *Cmp = dyn_cast<CmpInst>(NarrowInst))
+      return Builder.CreateCmp(Cmp->getPredicate(), NewOperands[0],
+                               NewOperands[1]);
+    if (isa<SelectInst>(NarrowInst))
+      return Builder.CreateSelect(NewOperands[0], NewOperands[1],
+                                  NewOperands[2]);
+    if (isa<FreezeInst>(NarrowInst))
+      return Builder.CreateFreeze(NewOperands[0]);
+    if (auto *II = dyn_cast<IntrinsicInst>(NarrowInst))
+      return Builder.CreateIntrinsic(WideResultTy, II->getIntrinsicID(),
+                                     NewOperands);
+    llvm_unreachable("Unsupported instruction");
+  }
+
+public:
+  InterleavedElementwiseChain(IRBuilderBase &Builder, Value *Root)
+      : Builder(Builder), Root(Root) {}
+
+  /// Visit a list  of instructions and check if they can be rewritten as a
+  /// single wider instruction. If so, the list is appened to \p Steps.
+  bool visitInstLevel(ArrayRef<Instruction *> Insts,
+                      ArrayRef<unsigned> OperandNumbers) {
+    if (Insts.empty() || Insts.size() != OperandNumbers.size())
+      return false;
+
+    Instruction *FirstInst = Insts.front();
+    unsigned ChainOperand = OperandNumbers.front();
+
+    if (!isSupportedElementwise(FirstInst))
+      return false;
+
+    for (unsigned Index = 1; Index != Insts.size(); ++Index)
+      if (OperandNumbers[Index] != ChainOperand ||
+          !FirstInst->isSameOperationAs(Insts[Index]))
+        return false;
+
+    // Non-chain operands must be either the same scalar or splats of that
+    // scalar. This intentionally rejects differing poison/undef or non-splat
+    // vector operands between chains.
+    for (unsigned Op = 0, E = getNumDataOperands(FirstInst); Op != E; ++Op) {
+      if (Op == ChainOperand)
+        continue;
+
+      Value *CommonValue = getSplatOrScalar(FirstInst->getOperand(Op));
+      if (!CommonValue || any_of(drop_begin(Insts), [&](Instruction *Inst) {
+            return getSplatOrScalar(Inst->getOperand(Op)) != CommonValue;
+          }))
+        return false;
+    }
+
+    Steps.emplace_back(Insts, ChainOperand);
+    return true;
+  }
+
+  /// Go through all the collected \p Steps and rewrite each of them as a wide
+  /// instruction.
+  Value *widenInstructions() {
+    Value *WideValue = Root;
+    ElementCount WideEC =
+        cast<VectorType>(WideValue->getType())->getElementCount();
+
+    for (const ElementwiseStep &Step : Steps) {
+      Instruction *NarrowInst = Step.Insts.front();
+
+      Builder.SetInsertPoint(NarrowInst);
+      Builder.SetCurrentDebugLocation(NarrowInst->getDebugLoc());
+
+      unsigned NumOperands = getNumDataOperands(NarrowInst);
+      SmallVector<Value *, 4> NewOperands;
+      NewOperands.reserve(NumOperands);
+
+      for (unsigned Op = 0; Op != NumOperands; ++Op) {
+        Value *Operand = NarrowInst->getOperand(Op);
+
+        if (Op == Step.ChainOperand)
+          Operand = WideValue;
+        else if (isa<VectorType>(Operand->getType()))
+          Operand = Builder.CreateVectorSplat(WideEC, getSplatValue(Operand));
+        NewOperands.push_back(Operand);
+      }
+
+      auto *WideResultTy =
+          VectorType::get(NarrowInst->getType()->getScalarType(), WideEC);
+      Value *NewValue =
+          createWideInstruction(NarrowInst, NewOperands, WideResultTy);
+
+      SmallVector<Value *, 8> NarrowInsts(Step.Insts.begin(), Step.Insts.end());
+      propagateIRFlags(NewValue, NarrowInsts);
+
+      if (auto *NewInst = dyn_cast<Instruction>(NewValue))
+        propagateMetadata(NewInst, NarrowInsts);
+
+      WideValue = NewValue;
+    }
+
+    return WideValue;
+  }
+};
+
+} // namespace
+
+/// Fold away a matched pair of vector.deinterleave/interleave intrinsics
+/// with a chain of elementwise operations on each between the
+/// deinterleave and interleave.
+///
+/// For example:
+///  ```
+///  %d = call { <2 x i16>, <2 x i16> } @deinterleave2.v4i16(<4 x i16> %v)
+///  %f0 = extractvalue { <2 x i16>, <2 x i16> } %d, 0
+///  %f1 = extractvalue { <2 x i16>, <2 x i16> } %d, 1
+///
+///  %u0 = add <2 x i16> %f0, splat (i16 3)
+///  %u1 = add <2 x i16> %f1, splat (i16 3)
+///
+///  %r = call <4 x i16> @interleave2.v4i16(<2 x i16> %u0, <2 x i16> %u1)
+///  ```
+/// Folds to:
+///  ```
+///  %r = add <4 x i16> %v, splat (i16 3)
+///  ```
+bool VectorCombine::foldDeinterleaveInterleavePair(Instruction &I) {
+  auto *Deinterleave = dyn_cast<IntrinsicInst>(&I);
+  if (!Deinterleave)
+    return false;
+
+  unsigned Factor =
+      getDeinterleaveIntrinsicFactor(Deinterleave->getIntrinsicID());
+  if (!Factor || Deinterleave->hasOperandBundles() ||
+      !Deinterleave->hasNUndroppableUses(Factor))
+    return false;
+
+  const Intrinsic::ID InterleaveIID =
+      Intrinsic::getInterleaveIntrinsicID(Factor);
+
+  // Collect one extract for each deinterleaved field.
+  SmallVector<Instruction *, 8> CurrentInsts(Factor, nullptr);
+  for (Use &U : Deinterleave->uses()) {
+    if (U.getUser()->isDroppable())
+      continue;
+
+    auto *Extract = dyn_cast<ExtractValueInst>(U.getUser());
+    if (!Extract || Extract->getNumIndices() != 1)
+      return false;
+
+    unsigned Index = *Extract->idx_begin();
+    if (Index >= Factor || CurrentInsts[Index])
+      return false;
+
+    CurrentInsts[Index] = Extract;
+  }
+
+  InterleavedElementwiseChain Chain(Builder, Deinterleave->getArgOperand(0));
+  IntrinsicInst *Interleave = nullptr;
+  unsigned NumVisited = 0;
+
+  // Traverse the Factor use chains with a breadth-first search.
+  // At each level, expect every chain to perform the same operation with the
+  // preceding chain value at the same operand position, until they all reach
+  // the matching interleave.
+  SmallVector<unsigned, 8> OperandNumbers(CurrentInsts.size());
+  while (NumVisited + Factor <= MaxInstrsToScan) {
+    NumVisited += Factor;
+
+    for (auto [Current, OpNumber] : zip_equal(CurrentInsts, OperandNumbers)) {
+      Use *U = Current->getSingleUndroppableUse();
+      auto *Next = U ? dyn_cast<Instruction>(U->getUser()) : nullptr;
----------------
huntergr-arm wrote:

```suggestion
      auto *Next = dyn_cast_if_present<Instruction>(U->getUser());
```

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