Skip to content

Commit 52f5c4b

Browse files
agorenstein-nvidiaagorenstein-nvidia
committed
Restore old tests, reframe optimization as extension of existing one.
1 parent 70d3c01 commit 52f5c4b

File tree

6 files changed

+59
-47
lines changed

6 files changed

+59
-47
lines changed

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

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -542,34 +542,40 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
542542
}
543543
}
544544
} else if (auto *SVI = dyn_cast<ShuffleVectorInst>(I)) {
545-
// extractelt (shufflevector %v1, %v2, splat-mask) idx ->
546-
// extractelt %v1, splat-mask[0]
547-
auto Mask = SVI->getShuffleMask();
548-
if (Mask[0] != PoisonMaskElem && all_equal(Mask))
549-
return ExtractElementInst::Create(SVI->getOperand(0),
550-
Builder.getInt64(Mask[0]));
551-
552-
// If this is extracting an element from a shufflevector, figure out where
553-
// it came from and extract from the appropriate input element instead.
554-
// Restrict the following transformation to fixed-length vector.
555-
if (isa<FixedVectorType>(SVI->getType()) && isa<ConstantInt>(Index)) {
556-
int SrcIdx =
557-
SVI->getMaskValue(cast<ConstantInt>(Index)->getZExtValue());
558-
Value *Src;
559-
unsigned LHSWidth = cast<FixedVectorType>(SVI->getOperand(0)->getType())
560-
->getNumElements();
561-
562-
if (SrcIdx < 0)
563-
return replaceInstUsesWith(EI, PoisonValue::get(EI.getType()));
564-
if (SrcIdx < (int)LHSWidth)
565-
Src = SVI->getOperand(0);
566-
else {
567-
SrcIdx -= LHSWidth;
568-
Src = SVI->getOperand(1);
545+
int SplatIndex = getSplatIndex(SVI->getShuffleMask());
546+
// We know such a splat must be reading from the first operand, even
547+
// in the case of scalable vectors (vscale is always > 0).
548+
if (SplatIndex == 0)
549+
return ExtractElementInst::Create(SVI->getOperand(0), Builder.getInt64(0));
550+
// Restrict the non-zero index case to fixed-length vectors
551+
if (isa<FixedVectorType>(SVI->getType())) {
552+
553+
// getSplatIndex doesn't distinguish between the all-poison splat and
554+
// a non-splat mask. However, if Index is -1, we still want to propagate
555+
// that poison value.
556+
int SrcIdx = -2;
557+
if (SplatIndex != PoisonMaskElem)
558+
SrcIdx = SplatIndex;
559+
else if (ConstantInt* CI = dyn_cast<ConstantInt>(Index))
560+
SrcIdx = SVI->getMaskValue(CI->getZExtValue());
561+
562+
if (SrcIdx != -2) {
563+
Value *Src;
564+
unsigned LHSWidth = cast<FixedVectorType>(SVI->getOperand(0)->getType())
565+
->getNumElements();
566+
567+
if (SrcIdx < 0)
568+
return replaceInstUsesWith(EI, PoisonValue::get(EI.getType()));
569+
if (SrcIdx < (int)LHSWidth)
570+
Src = SVI->getOperand(0);
571+
else {
572+
SrcIdx -= LHSWidth;
573+
Src = SVI->getOperand(1);
574+
}
575+
Type *Int64Ty = Type::getInt64Ty(EI.getContext());
576+
return ExtractElementInst::Create(
577+
Src, ConstantInt::get(Int64Ty, SrcIdx, false));
569578
}
570-
Type *Int64Ty = Type::getInt64Ty(EI.getContext());
571-
return ExtractElementInst::Create(
572-
Src, ConstantInt::get(Int64Ty, SrcIdx, false));
573579
}
574580
} else if (auto *CI = dyn_cast<CastInst>(I)) {
575581
// Canonicalize extractelement(cast) -> cast(extractelement).

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,23 @@ define float @extract_from_general_scalable_splat(<vscale x 2 x float> %1, i64 %
3030
%4 = extractelement <vscale x 4 x float> %3, i64 %idx
3131
ret float %4
3232
}
33+
34+
define float @extract_from_splat_with_poison_0(<2 x float> %1, i64 %idx) {
35+
; CHECK-LABEL: @extract_from_splat_with_poison_0(
36+
; CHECK-NEXT: [[TMP2:%.*]] = extractelement <2 x float> [[TMP1:%.*]], i64 1
37+
; CHECK-NEXT: ret float [[TMP2]]
38+
;
39+
%3 = shufflevector <2 x float> %1, <2 x float> poison, <4 x i32> <i32 poison, i32 1, i32 1, i32 1>
40+
%4 = extractelement <4 x float> %3, i64 %idx
41+
ret float %4
42+
}
43+
44+
define float @extract_from_splat_with_poison_1(<2 x float> %1, i64 %idx) {
45+
; CHECK-LABEL: @extract_from_splat_with_poison_1(
46+
; CHECK-NEXT: [[TMP2:%.*]] = extractelement <2 x float> [[TMP1:%.*]], i64 1
47+
; CHECK-NEXT: ret float [[TMP2]]
48+
;
49+
%3 = shufflevector <2 x float> %1, <2 x float> poison, <4 x i32> <i32 1, i32 poison, i32 1, i32 1>
50+
%4 = extractelement <4 x float> %3, i64 %idx
51+
ret float %4
52+
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,7 @@ define float @test6(<4 x float> %X) {
6161

6262
define float @testvscale6(<vscale x 4 x float> %X) {
6363
; CHECK-LABEL: @testvscale6(
64-
; CHECK-NEXT: [[T2:%.*]] = shufflevector <vscale x 4 x float> [[X:%.*]], <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
65-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x float> [[T2]], i64 0
64+
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x float> [[X:%.*]], i64 0
6665
; CHECK-NEXT: ret float [[R]]
6766
;
6867
%X1 = bitcast <vscale x 4 x float> %X to <vscale x 4 x i32>

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,7 @@ define float @test6(<4 x float> %X) {
6767

6868
define float @testvscale6(<vscale x 4 x float> %X) {
6969
; CHECK-LABEL: @testvscale6(
70-
; CHECK-NEXT: [[T2:%.*]] = shufflevector <vscale x 4 x float> [[X:%.*]], <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
71-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x float> [[T2]], i64 0
70+
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x float> [[X:%.*]], i64 0
7271
; CHECK-NEXT: ret float [[R]]
7372
;
7473
%X1 = bitcast <vscale x 4 x float> %X to <vscale x 4 x i32>

‎llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,7 @@ define i8 @extractelement_bitcast_insert_extra_use_bitcast(<vscale x 2 x i32> %a
9191

9292
define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {
9393
; CHECK-LABEL: @extractelement_shuffle_maybe_out_of_range(
94-
; CHECK-NEXT: [[IN:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[V:%.*]], i64 0
95-
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
96-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i64 4
97-
; CHECK-NEXT: ret i32 [[R]]
94+
; CHECK-NEXT: ret i32 [[V:%.*]]
9895
;
9996
%in = insertelement <vscale x 4 x i32> poison, i32 %v, i32 0
10097
%splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
@@ -104,10 +101,7 @@ define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {
104101

105102
define i32 @extractelement_shuffle_invalid_index(i32 %v) {
106103
; CHECK-LABEL: @extractelement_shuffle_invalid_index(
107-
; CHECK-NEXT: [[IN:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[V:%.*]], i64 0
108-
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
109-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i64 4294967295
110-
; CHECK-NEXT: ret i32 [[R]]
104+
; CHECK-NEXT: ret i32 [[V:%.*]]
111105
;
112106
%in = insertelement <vscale x 4 x i32> poison, i32 %v, i32 0
113107
%splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,7 @@ define i8 @extractelement_bitcast_useless_insert(<vscale x 2 x i32> %a, i32 %x)
5555

5656
define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {
5757
; CHECK-LABEL: @extractelement_shuffle_maybe_out_of_range(
58-
; CHECK-NEXT: [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i64 0
59-
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
60-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i64 4
61-
; CHECK-NEXT: ret i32 [[R]]
58+
; CHECK-NEXT: ret i32 [[V:%.*]]
6259
;
6360
%in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
6461
%splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
@@ -68,10 +65,7 @@ define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {
6865

6966
define i32 @extractelement_shuffle_invalid_index(i32 %v) {
7067
; CHECK-LABEL: @extractelement_shuffle_invalid_index(
71-
; CHECK-NEXT: [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i64 0
72-
; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
73-
; CHECK-NEXT: [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i64 4294967295
74-
; CHECK-NEXT: ret i32 [[R]]
68+
; CHECK-NEXT: ret i32 [[V:%.*]]
7569
;
7670
%in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
7771
%splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer

0 commit comments

Comments
 (0)