[llvm] [InstCombine] Refactor fixed and scalable binop shuffle combine. NFCI (PR #141287)
Luke Lau via llvm-commits
llvm-commits at lists.llvm.org
Fri May 23 16:26:17 PDT 2025
https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/141287
>From 689b76298568f70857391e26d5f05aa9df88c1d1 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 23 May 2025 20:37:14 +0100
Subject: [PATCH 1/2] [InstCombine] Refactor fixed and scalable binop shuffle
combine. NFCI
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.
---
.../InstCombine/InstructionCombining.cpp | 110 ++++++++----------
1 file changed, 49 insertions(+), 61 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index c1608b1866a5d..eb5352524853f 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2094,6 +2094,50 @@ static bool shouldMergeGEPs(GEPOperator &GEP, GEPOperator &Src) {
return true;
}
+// Find constant NewC that has property:
+// shuffle(NewC, ShMask) = C
+// Returns nullptr if such a constant does not exist e.g. ShMask=<0,0> C=<1,2>
+//
+// A 1-to-1 mapping is not required. Example:
+// ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <poison,5,6,poison>
+static Constant *unshuffleConstant(ArrayRef<int> ShMask, Constant *C,
+ VectorType *NewCTy) {
+ if (isa<ScalableVectorType>(NewCTy)) {
+ Constant *Splat = C->getSplatValue();
+ if (!Splat)
+ return nullptr;
+ return ConstantVector::getSplat(
+ cast<VectorType>(C->getType())->getElementCount(), Splat);
+ }
+
+ if (cast<FixedVectorType>(NewCTy)->getNumElements() >
+ cast<FixedVectorType>(C->getType())->getNumElements())
+ return nullptr;
+
+ unsigned NewCNumElts = cast<FixedVectorType>(NewCTy)->getNumElements();
+ PoisonValue *PoisonScalar = PoisonValue::get(C->getType()->getScalarType());
+ SmallVector<Constant *, 16> NewVecC(NewCNumElts, PoisonScalar);
+ unsigned NumElts = cast<FixedVectorType>(C->getType())->getNumElements();
+ for (unsigned I = 0; I < NumElts; ++I) {
+ Constant *CElt = C->getAggregateElement(I);
+ if (ShMask[I] >= 0) {
+ assert(ShMask[I] < (int)NumElts && "Not expecting narrowing shuffle");
+ Constant *NewCElt = NewVecC[ShMask[I]];
+ // Bail out if:
+ // 1. The constant vector contains a constant expression.
+ // 2. The shuffle needs an element of the constant vector that can't
+ // be mapped to a new constant vector.
+ // 3. This is a widening shuffle that copies elements of V1 into the
+ // extended elements (extending with poison is allowed).
+ if (!CElt || (!isa<PoisonValue>(NewCElt) && NewCElt != CElt) ||
+ I >= NewCNumElts)
+ return nullptr;
+ NewVecC[ShMask[I]] = CElt;
+ }
+ }
+ return ConstantVector::get(NewVecC);
+}
+
Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
if (!isa<VectorType>(Inst.getType()))
return nullptr;
@@ -2213,50 +2257,15 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
// other binops, so they can be folded. It may also enable demanded elements
// transforms.
Constant *C;
- auto *InstVTy = dyn_cast<FixedVectorType>(Inst.getType());
- if (InstVTy &&
- match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
+ if (match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
m_Mask(Mask))),
- m_ImmConstant(C))) &&
- cast<FixedVectorType>(V1->getType())->getNumElements() <=
- InstVTy->getNumElements()) {
- assert(InstVTy->getScalarType() == V1->getType()->getScalarType() &&
+ m_ImmConstant(C)))) {
+ assert(Inst.getType()->getScalarType() == V1->getType()->getScalarType() &&
"Shuffle should not change scalar type");
- // Find constant NewC that has property:
- // shuffle(NewC, ShMask) = C
- // If such constant does not exist (example: ShMask=<0,0> and C=<1,2>)
- // reorder is not possible. A 1-to-1 mapping is not required. Example:
- // ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <undef,5,6,undef>
bool ConstOp1 = isa<Constant>(RHS);
- ArrayRef<int> ShMask = Mask;
- unsigned SrcVecNumElts =
- cast<FixedVectorType>(V1->getType())->getNumElements();
- PoisonValue *PoisonScalar = PoisonValue::get(C->getType()->getScalarType());
- SmallVector<Constant *, 16> NewVecC(SrcVecNumElts, PoisonScalar);
- bool MayChange = true;
- unsigned NumElts = InstVTy->getNumElements();
- for (unsigned I = 0; I < NumElts; ++I) {
- Constant *CElt = C->getAggregateElement(I);
- if (ShMask[I] >= 0) {
- assert(ShMask[I] < (int)NumElts && "Not expecting narrowing shuffle");
- Constant *NewCElt = NewVecC[ShMask[I]];
- // Bail out if:
- // 1. The constant vector contains a constant expression.
- // 2. The shuffle needs an element of the constant vector that can't
- // be mapped to a new constant vector.
- // 3. This is a widening shuffle that copies elements of V1 into the
- // extended elements (extending with poison is allowed).
- if (!CElt || (!isa<PoisonValue>(NewCElt) && NewCElt != CElt) ||
- I >= SrcVecNumElts) {
- MayChange = false;
- break;
- }
- NewVecC[ShMask[I]] = CElt;
- }
- }
- if (MayChange) {
- Constant *NewC = ConstantVector::get(NewVecC);
+ if (Constant *NewC =
+ unshuffleConstant(Mask, C, cast<VectorType>(V1->getType()))) {
// It may not be safe to execute a binop on a vector with poison elements
// because the entire instruction can be folded to undef or create poison
// that did not exist in the original code.
@@ -2272,27 +2281,6 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
}
}
- // Similar to the combine above, but handles the case for scalable vectors
- // where both shuffle(V1, 0) and C are splats.
- //
- // Op(shuffle(V1, 0), (splat C)) -> shuffle(Op(V1, (splat C)), 0)
- if (isa<ScalableVectorType>(Inst.getType()) &&
- match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(),
- m_ZeroMask())),
- m_ImmConstant(C)))) {
- if (Constant *Splat = C->getSplatValue()) {
- bool ConstOp1 = isa<Constant>(RHS);
- VectorType *V1Ty = cast<VectorType>(V1->getType());
- Constant *NewC = ConstantVector::getSplat(V1Ty->getElementCount(), Splat);
-
- Value *NewLHS = ConstOp1 ? V1 : NewC;
- Value *NewRHS = ConstOp1 ? NewC : V1;
- VectorType *VTy = cast<VectorType>(Inst.getType());
- SmallVector<int> Mask(VTy->getElementCount().getKnownMinValue(), 0);
- return createBinOpShuffle(NewLHS, NewRHS, Mask);
- }
- }
-
// Try to reassociate to sink a splat shuffle after a binary operation.
if (Inst.isAssociative() && Inst.isCommutative()) {
// Canonicalize shuffle operand as LHS.
>From 660475b076e78f08439088078d4a4454590beb79 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sat, 24 May 2025 00:25:47 +0100
Subject: [PATCH 2/2] Fix shl crash
---
.../InstCombine/InstructionCombining.cpp | 4 +++-
.../Transforms/InstCombine/vec_shuffle.ll | 22 +++++++++++++++++++
2 files changed, 25 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index eb5352524853f..8e98c520a35be 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2269,8 +2269,10 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
// It may not be safe to execute a binop on a vector with poison elements
// because the entire instruction can be folded to undef or create poison
// that did not exist in the original code.
+ // This isn't needed for scalable vectors since they're always splats.
// TODO: The shift case should not be necessary.
- if (Inst.isIntDivRem() || (Inst.isShift() && ConstOp1))
+ if (isa<FixedVectorType>(V1->getType()) &&
+ (Inst.isIntDivRem() || (Inst.isShift() && ConstOp1)))
NewC = getSafeVectorConstantForBinop(Opcode, NewC, ConstOp1);
// Op(shuffle(V1, Mask), C) -> shuffle(Op(V1, NewC), Mask)
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index eb88fbadab956..998be273aa445 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -938,6 +938,28 @@ define <2 x i32> @shl_splat_constant1(<2 x i32> %x) {
ret <2 x i32> %r
}
+define <vscale x 2 x i32> @shl_splat_constant0_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @shl_splat_constant0_scalable(
+; CHECK-NEXT: [[TMP1:%.*]] = shl <vscale x 2 x i32> splat (i32 5), [[X:%.*]]
+; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
+;
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+ %r = shl <vscale x 2 x i32> splat (i32 5), %splat
+ ret <vscale x 2 x i32> %r
+}
+
+define <vscale x 2 x i32> @shl_splat_constant1_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @shl_splat_constant1_scalable(
+; CHECK-NEXT: [[TMP1:%.*]] = shl <vscale x 2 x i32> [[X:%.*]], splat (i32 5)
+; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: ret <vscale x 2 x i32> [[R]]
+;
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+ %r = shl <vscale x 2 x i32> %splat, splat (i32 5)
+ ret <vscale x 2 x i32> %r
+}
+
define <2 x i32> @ashr_splat_constant0(<2 x i32> %x) {
; CHECK-LABEL: @ashr_splat_constant0(
; CHECK-NEXT: [[TMP1:%.*]] = lshr <2 x i32> <i32 5, i32 poison>, [[X:%.*]]
More information about the llvm-commits
mailing list