[llvm] [InstCombine] Refactor fixed and scalable binop shuffle combine. NFCI (PR #141287)
Luke Lau via llvm-commits
llvm-commits at lists.llvm.org
Sat May 24 06:07:47 PDT 2025
https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/141287
>From 00ccf64a0e130ea1f02f97d9ffe7262b96e83662 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/5] [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 2a96f4beb5f57..4acf7d2921d1b 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()))) {
// Lanes of NewC not used by the shuffle will be poison which will cause
// UB for div/rem. Mask them with a safe constant.
if (Inst.isIntDivRem())
@@ -2270,27 +2279,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 bbbac687be2e97d6f4fb9ebf0c5ebafcb7ffd937 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/5] Fix calling getSafeVectorConstant with scalable vectors
---
.../InstCombine/InstructionCombining.cpp | 6 +--
.../Transforms/InstCombine/vec_shuffle.ll | 44 +++++++++++++++++++
2 files changed, 47 insertions(+), 3 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 4acf7d2921d1b..0117951b92a99 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2266,9 +2266,9 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
bool ConstOp1 = isa<Constant>(RHS);
if (Constant *NewC =
unshuffleConstant(Mask, C, cast<VectorType>(V1->getType()))) {
- // Lanes of NewC not used by the shuffle will be poison which will cause
- // UB for div/rem. Mask them with a safe constant.
- if (Inst.isIntDivRem())
+ // For fixed vectors, lanes of NewC not used by the shuffle will be poison
+ // which will cause UB for div/rem. Mask them with a safe constant.
+ if (isa<FixedVectorType>(V1->getType()) && Inst.isIntDivRem())
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 90ad88089054e..a268956cc59e4 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:%.*]]
@@ -1048,6 +1070,28 @@ define <2 x i32> @udiv_splat_constant1(<2 x i32> %x) {
ret <2 x i32> %r
}
+define <vscale x 2 x i32> @udiv_splat_constant0_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @udiv_splat_constant0_scalable(
+; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 2 x i32> [[X:%.*]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: [[R:%.*]] = udiv <vscale x 2 x i32> splat (i32 42), [[SPLAT]]
+; 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 = udiv <vscale x 2 x i32> splat (i32 42), %splat
+ ret <vscale x 2 x i32> %r
+}
+
+define <vscale x 2 x i32> @udiv_splat_constant1_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @udiv_splat_constant1_scalable(
+; CHECK-NEXT: [[TMP1:%.*]] = udiv <vscale x 2 x i32> [[X:%.*]], splat (i32 42)
+; 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 = udiv <vscale x 2 x i32> %splat, splat (i32 42)
+ ret <vscale x 2 x i32> %r
+}
+
define <2 x i32> @sdiv_splat_constant0(<2 x i32> %x) {
; CHECK-LABEL: @sdiv_splat_constant0(
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <2 x i32> [[X:%.*]], <2 x i32> poison, <2 x i32> zeroinitializer
>From dd1e8e8b0e29e8ad53a6a7675692241db224f02c Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sat, 24 May 2025 14:04:17 +0100
Subject: [PATCH 3/5] Fix using old CTy instead of NewCTy for scalable splat
---
.../Transforms/InstCombine/InstructionCombining.cpp | 3 +--
llvm/test/Transforms/InstCombine/vec_shuffle.ll | 11 +++++++++++
2 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 0117951b92a99..07b634ef99604 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2106,8 +2106,7 @@ static Constant *unshuffleConstant(ArrayRef<int> ShMask, Constant *C,
Constant *Splat = C->getSplatValue();
if (!Splat)
return nullptr;
- return ConstantVector::getSplat(
- cast<VectorType>(C->getType())->getElementCount(), Splat);
+ return ConstantVector::getSplat(NewCTy->getElementCount(), Splat);
}
if (cast<FixedVectorType>(NewCTy)->getNumElements() >
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index a268956cc59e4..90959f517eb36 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -652,6 +652,17 @@ define <4 x i8> @widening_shuffle_add_1(<2 x i8> %x) {
ret <4 x i8> %r
}
+define <vscale x 4 x i8> @widening_shuffle_add_1_scalable(<vscale x 2 x i8> %x) {
+; CHECK-LABEL: @widening_shuffle_add_1_scalable(
+; CHECK-NEXT: [[TMP1:%.*]] = add <vscale x 2 x i8> [[X:%.*]], splat (i8 42)
+; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i8> [[TMP1]], <vscale x 2 x i8> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: ret <vscale x 4 x i8> [[R]]
+;
+ %widex = shufflevector <vscale x 2 x i8> %x, <vscale x 2 x i8> undef, <vscale x 4 x i32> zeroinitializer
+ %r = add <vscale x 4 x i8> %widex, splat (i8 42)
+ ret <vscale x 4 x i8> %r
+}
+
; Reduce the width of the binop by moving it ahead of a shuffle.
define <4 x i8> @widening_shuffle_add_2(<2 x i8> %x) {
>From 928ec026f6d7179803a8186352d28cf73fbede5d Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sat, 24 May 2025 14:04:39 +0100
Subject: [PATCH 4/5] Use /// for header comments
---
.../Transforms/InstCombine/InstructionCombining.cpp | 12 ++++++------
1 file changed, 6 insertions(+), 6 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 07b634ef99604..9ddcef0396e39 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2094,12 +2094,12 @@ 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>
+/// Find a 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)) {
>From 4fc211cd325858d26891ba143161e087904526e4 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sat, 24 May 2025 14:06:08 +0100
Subject: [PATCH 5/5] Use poison instead of undef in splats
---
llvm/test/Transforms/InstCombine/vec_shuffle.ll | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index 90959f517eb36..83919e743d384 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -658,7 +658,7 @@ define <vscale x 4 x i8> @widening_shuffle_add_1_scalable(<vscale x 2 x i8> %x)
; CHECK-NEXT: [[R:%.*]] = shufflevector <vscale x 2 x i8> [[TMP1]], <vscale x 2 x i8> poison, <vscale x 4 x i32> zeroinitializer
; CHECK-NEXT: ret <vscale x 4 x i8> [[R]]
;
- %widex = shufflevector <vscale x 2 x i8> %x, <vscale x 2 x i8> undef, <vscale x 4 x i32> zeroinitializer
+ %widex = shufflevector <vscale x 2 x i8> %x, <vscale x 2 x i8> poison, <vscale x 4 x i32> zeroinitializer
%r = add <vscale x 4 x i8> %widex, splat (i8 42)
ret <vscale x 4 x i8> %r
}
@@ -955,7 +955,7 @@ define <vscale x 2 x i32> @shl_splat_constant0_scalable(<vscale x 2 x i32> %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
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
%r = shl <vscale x 2 x i32> splat (i32 5), %splat
ret <vscale x 2 x i32> %r
}
@@ -966,7 +966,7 @@ define <vscale x 2 x i32> @shl_splat_constant1_scalable(<vscale x 2 x i32> %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
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
%r = shl <vscale x 2 x i32> %splat, splat (i32 5)
ret <vscale x 2 x i32> %r
}
@@ -1087,7 +1087,7 @@ define <vscale x 2 x i32> @udiv_splat_constant0_scalable(<vscale x 2 x i32> %x)
; CHECK-NEXT: [[R:%.*]] = udiv <vscale x 2 x i32> splat (i32 42), [[SPLAT]]
; 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
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
%r = udiv <vscale x 2 x i32> splat (i32 42), %splat
ret <vscale x 2 x i32> %r
}
@@ -1098,7 +1098,7 @@ define <vscale x 2 x i32> @udiv_splat_constant1_scalable(<vscale x 2 x i32> %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
+ %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
%r = udiv <vscale x 2 x i32> %splat, splat (i32 42)
ret <vscale x 2 x i32> %r
}
More information about the llvm-commits
mailing list