diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp index 2a4b709858ec..367b480c2114 100644 --- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp +++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp @@ -4160,6 +4160,31 @@ static SDValue foldCONCAT_VECTORS(const SDLoc &DL, EVT VT, if (llvm::all_of(Ops, [](SDValue Op) { return Op.isUndef(); })) return DAG.getUNDEF(VT); + // Scan the operands and look for extract operations from a single source + // that correspond to insertion at the same location via this concatenation: + // concat (extract X, 0*subvec_elts), (extract X, 1*subvec_elts), ... + SDValue IdentitySrc; + bool IsIdentity = true; + for (unsigned i = 0, e = Ops.size(); i != e; ++i) { + SDValue Op = Ops[i]; + unsigned IdentityIndex = i * Op.getValueType().getVectorNumElements(); + if (Op.getOpcode() != ISD::EXTRACT_SUBVECTOR || + Op.getOperand(0).getValueType() != VT || + (IdentitySrc && Op.getOperand(0) != IdentitySrc) || + !isa(Op.getOperand(1)) || + Op.getConstantOperandVal(1) != IdentityIndex) { + IsIdentity = false; + break; + } + assert((!IdentitySrc || IdentitySrc == Op.getOperand(0)) && + "Unexpected identity source vector for concat of extracts"); + IdentitySrc = Op.getOperand(0); + } + if (IsIdentity) { + assert(IdentitySrc && "Failed to set source vector of extracts"); + return IdentitySrc; + } + // A CONCAT_VECTOR with all UNDEF/BUILD_VECTOR operands can be // simplified to one big BUILD_VECTOR. // FIXME: Add support for SCALAR_TO_VECTOR as well.