Skip to content

Commit 8f9549c

Browse files
lukel97lukel97
authored
[InstCombine] Refactor fixed and scalable binop shuffle combine. NFCI (#141287)
This extracts the logic that works out the "unshuffled" constant when pulling shuffle vectors out of binary ops, so the same combine can be generic over fixed and scalable vectors. The plan is to reuse this helper to do the same canonicalization on intrinsics too.
1 parent 8d374f1 commit 8f9549c

File tree

2 files changed

+106
-64
lines changed

2 files changed

+106
-64
lines changed

‎llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Lines changed: 51 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -2094,6 +2094,49 @@ static bool shouldMergeGEPs(GEPOperator &GEP, GEPOperator &Src) {
20942094
return true;
20952095
}
20962096

2097+
/// Find a constant NewC that has property:
2098+
/// shuffle(NewC, ShMask) = C
2099+
/// Returns nullptr if such a constant does not exist e.g. ShMask=<0,0> C=<1,2>
2100+
///
2101+
/// A 1-to-1 mapping is not required. Example:
2102+
/// ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <poison,5,6,poison>
2103+
static Constant *unshuffleConstant(ArrayRef<int> ShMask, Constant *C,
2104+
VectorType *NewCTy) {
2105+
if (isa<ScalableVectorType>(NewCTy)) {
2106+
Constant *Splat = C->getSplatValue();
2107+
if (!Splat)
2108+
return nullptr;
2109+
return ConstantVector::getSplat(NewCTy->getElementCount(), Splat);
2110+
}
2111+
2112+
if (cast<FixedVectorType>(NewCTy)->getNumElements() >
2113+
cast<FixedVectorType>(C->getType())->getNumElements())
2114+
return nullptr;
2115+
2116+
unsigned NewCNumElts = cast<FixedVectorType>(NewCTy)->getNumElements();
2117+
PoisonValue *PoisonScalar = PoisonValue::get(C->getType()->getScalarType());
2118+
SmallVector<Constant *, 16> NewVecC(NewCNumElts, PoisonScalar);
2119+
unsigned NumElts = cast<FixedVectorType>(C->getType())->getNumElements();
2120+
for (unsigned I = 0; I < NumElts; ++I) {
2121+
Constant *CElt = C->getAggregateElement(I);
2122+
if (ShMask[I] >= 0) {
2123+
assert(ShMask[I] < (int)NumElts && "Not expecting narrowing shuffle");
2124+
Constant *NewCElt = NewVecC[ShMask[I]];
2125+
// Bail out if:
2126+
// 1. The constant vector contains a constant expression.
2127+
// 2. The shuffle needs an element of the constant vector that can't
2128+
// be mapped to a new constant vector.
2129+
// 3. This is a widening shuffle that copies elements of V1 into the
2130+
// extended elements (extending with poison is allowed).
2131+
if (!CElt || (!isa<PoisonValue>(NewCElt) && NewCElt != CElt) ||
2132+
I >= NewCNumElts)
2133+
return nullptr;
2134+
NewVecC[ShMask[I]] = CElt;
2135+
}
2136+
}
2137+
return ConstantVector::get(NewVecC);
2138+
}
2139+
20972140
Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
20982141
if (!isa<VectorType>(Inst.getType()))
20992142
return nullptr;
@@ -2213,53 +2256,18 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
22132256
// other binops, so they can be folded. It may also enable demanded elements
22142257
// transforms.
22152258
Constant *C;
2216-
auto *InstVTy = dyn_cast<FixedVectorType>(Inst.getType());
2217-
if (InstVTy &&
2218-
match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
2259+
if (match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
22192260
m_Mask(Mask))),
2220-
m_ImmConstant(C))) &&
2221-
cast<FixedVectorType>(V1->getType())->getNumElements() <=
2222-
InstVTy->getNumElements()) {
2223-
assert(InstVTy->getScalarType() == V1->getType()->getScalarType() &&
2261+
m_ImmConstant(C)))) {
2262+
assert(Inst.getType()->getScalarType() == V1->getType()->getScalarType() &&
22242263
"Shuffle should not change scalar type");
22252264

2226-
// Find constant NewC that has property:
2227-
// shuffle(NewC, ShMask) = C
2228-
// If such constant does not exist (example: ShMask=<0,0> and C=<1,2>)
2229-
// reorder is not possible. A 1-to-1 mapping is not required. Example:
2230-
// ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <undef,5,6,undef>
22312265
bool ConstOp1 = isa<Constant>(RHS);
2232-
ArrayRef<int> ShMask = Mask;
2233-
unsigned SrcVecNumElts =
2234-
cast<FixedVectorType>(V1->getType())->getNumElements();
2235-
PoisonValue *PoisonScalar = PoisonValue::get(C->getType()->getScalarType());
2236-
SmallVector<Constant *, 16> NewVecC(SrcVecNumElts, PoisonScalar);
2237-
bool MayChange = true;
2238-
unsigned NumElts = InstVTy->getNumElements();
2239-
for (unsigned I = 0; I < NumElts; ++I) {
2240-
Constant *CElt = C->getAggregateElement(I);
2241-
if (ShMask[I] >= 0) {
2242-
assert(ShMask[I] < (int)NumElts && "Not expecting narrowing shuffle");
2243-
Constant *NewCElt = NewVecC[ShMask[I]];
2244-
// Bail out if:
2245-
// 1. The constant vector contains a constant expression.
2246-
// 2. The shuffle needs an element of the constant vector that can't
2247-
// be mapped to a new constant vector.
2248-
// 3. This is a widening shuffle that copies elements of V1 into the
2249-
// extended elements (extending with poison is allowed).
2250-
if (!CElt || (!isa<PoisonValue>(NewCElt) && NewCElt != CElt) ||
2251-
I >= SrcVecNumElts) {
2252-
MayChange = false;
2253-
break;
2254-
}
2255-
NewVecC[ShMask[I]] = CElt;
2256-
}
2257-
}
2258-
if (MayChange) {
2259-
Constant *NewC = ConstantVector::get(NewVecC);
2260-
// Lanes of NewC not used by the shuffle will be poison which will cause
2261-
// UB for div/rem. Mask them with a safe constant.
2262-
if (Inst.isIntDivRem())
2266+
if (Constant *NewC =
2267+
unshuffleConstant(Mask, C, cast<VectorType>(V1->getType()))) {
2268+
// For fixed vectors, lanes of NewC not used by the shuffle will be poison
2269+
// which will cause UB for div/rem. Mask them with a safe constant.
2270+
if (isa<FixedVectorType>(V1->getType()) && Inst.isIntDivRem())
22632271
NewC = getSafeVectorConstantForBinop(Opcode, NewC, ConstOp1);
22642272

22652273
// Op(shuffle(V1, Mask), C) -> shuffle(Op(V1, NewC), Mask)
@@ -2270,27 +2278,6 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
22702278
}
22712279
}
22722280

2273-
// Similar to the combine above, but handles the case for scalable vectors
2274-
// where both shuffle(V1, 0) and C are splats.
2275-
//
2276-
// Op(shuffle(V1, 0), (splat C)) -> shuffle(Op(V1, (splat C)), 0)
2277-
if (isa<ScalableVectorType>(Inst.getType()) &&
2278-
match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
2279-
m_ZeroMask())),
2280-
m_ImmConstant(C)))) {
2281-
if (Constant *Splat = C->getSplatValue()) {
2282-
bool ConstOp1 = isa<Constant>(RHS);
2283-
VectorType *V1Ty = cast<VectorType>(V1->getType());
2284-
Constant *NewC = ConstantVector::getSplat(V1Ty->getElementCount(), Splat);
2285-
2286-
Value *NewLHS = ConstOp1 ? V1 : NewC;
2287-
Value *NewRHS = ConstOp1 ? NewC : V1;
2288-
VectorType *VTy = cast<VectorType>(Inst.getType());
2289-
SmallVector<int> Mask(VTy->getElementCount().getKnownMinValue(), 0);
2290-
return createBinOpShuffle(NewLHS, NewRHS, Mask);
2291-
}
2292-
}
2293-
22942281
// Try to reassociate to sink a splat shuffle after a binary operation.
22952282
if (Inst.isAssociative() && Inst.isCommutative()) {
22962283
// Canonicalize shuffle operand as LHS.

‎llvm/test/Transforms/InstCombine/vec_shuffle.ll

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -652,6 +652,17 @@ define <4 x i8> @widening_shuffle_add_1(<2 x i8> %x) {
652652
ret <4 x i8> %r
653653
}
654654

655+
define <vscale x 4 x i8> @widening_shuffle_add_1_scalable(<vscale x 2 x i8> %x) {
656+
; CHECK-LABEL: @widening_shuffle_add_1_scalable(
657+
; CHECK-NEXT: [[TMP1:%.*]] = add <vscale x 2 x i8> [[X:%.*]], splat (i8 42)
658+
; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i8> [[TMP1]], <vscale x 2 x i8> poison, <vscale x 4 x i32> zeroinitializer
659+
; CHECK-NEXT: ret <vscale x 4 x i8> [[R]]
660+
;
661+
%widex = shufflevector <vscale x 2 x i8> %x, <vscale x 2 x i8> poison, <vscale x 4 x i32> zeroinitializer
662+
%r = add <vscale x 4 x i8> %widex, splat (i8 42)
663+
ret <vscale x 4 x i8> %r
664+
}
665+
655666
; Reduce the width of the binop by moving it ahead of a shuffle.
656667

657668
define <4 x i8> @widening_shuffle_add_2(<2 x i8> %x) {
@@ -938,6 +949,28 @@ define <2 x i32> @shl_splat_constant1(<2 x i32> %x) {
938949
ret <2 x i32> %r
939950
}
940951

952+
define <vscale x 2 x i32> @shl_splat_constant0_scalable(<vscale x 2 x i32> %x) {
953+
; CHECK-LABEL: @shl_splat_constant0_scalable(
954+
; CHECK-NEXT: [[TMP1:%.*]] = shl <vscale x 2 x i32> splat (i32 5), [[X:%.*]]
955+
; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
956+
; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
957+
;
958+
%splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
959+
%r = shl <vscale x 2 x i32> splat (i32 5), %splat
960+
ret <vscale x 2 x i32> %r
961+
}
962+
963+
define <vscale x 2 x i32> @shl_splat_constant1_scalable(<vscale x 2 x i32> %x) {
964+
; CHECK-LABEL: @shl_splat_constant1_scalable(
965+
; CHECK-NEXT: [[TMP1:%.*]] = shl <vscale x 2 x i32> [[X:%.*]], splat (i32 5)
966+
; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
967+
; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
968+
;
969+
%splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
970+
%r = shl <vscale x 2 x i32> %splat, splat (i32 5)
971+
ret <vscale x 2 x i32> %r
972+
}
973+
941974
define <2 x i32> @ashr_splat_constant0(<2 x i32> %x) {
942975
; CHECK-LABEL: @ashr_splat_constant0(
943976
; CHECK-NEXT: [[TMP1:%.*]] = lshr <2 x i32> <i32 5, i32 poison>, [[X:%.*]]
@@ -1048,6 +1081,28 @@ define <2 x i32> @udiv_splat_constant1(<2 x i32> %x) {
10481081
ret <2 x i32> %r
10491082
}
10501083

1084+
define <vscale x 2 x i32> @udiv_splat_constant0_scalable(<vscale x 2 x i32> %x) {
1085+
; CHECK-LABEL: @udiv_splat_constant0_scalable(
1086+
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 2 x i32> [[X:%.*]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
1087+
; CHECK-NEXT: [[R:%.*]] = udiv <vscale x 2 x i32> splat (i32 42), [[SPLAT]]
1088+
; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
1089+
;
1090+
%splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
1091+
%r = udiv <vscale x 2 x i32> splat (i32 42), %splat
1092+
ret <vscale x 2 x i32> %r
1093+
}
1094+
1095+
define <vscale x 2 x i32> @udiv_splat_constant1_scalable(<vscale x 2 x i32> %x) {
1096+
; CHECK-LABEL: @udiv_splat_constant1_scalable(
1097+
; CHECK-NEXT: [[TMP1:%.*]] = udiv <vscale x 2 x i32> [[X:%.*]], splat (i32 42)
1098+
; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
1099+
; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
1100+
;
1101+
%splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
1102+
%r = udiv <vscale x 2 x i32> %splat, splat (i32 42)
1103+
ret <vscale x 2 x i32> %r
1104+
}
1105+
10511106
define <2 x i32> @sdiv_splat_constant0(<2 x i32> %x) {
10521107
; CHECK-LABEL: @sdiv_splat_constant0(
10531108
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <2 x i32> [[X:%.*]], <2 x i32> poison, <2 x i32> zeroinitializer

0 commit comments

Comments
 (0)