Skip to content

Commit 8c9fcd9

Browse files
RKSimonRKSimon
committed
[InstCombine] foldVecExtTruncToExtElt - extend to handle trunc(lshr(extractelement(x,c1),c2)) -> extractelement(bitcast(x),c3) patterns.
Fixes #107404
1 parent 729626b commit 8c9fcd9

File tree

3 files changed

+100
-45
lines changed

3 files changed

+100
-45
lines changed

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

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -460,7 +460,11 @@ static Instruction *foldVecExtTruncToExtElt(TruncInst &Trunc,
460460

461461
Value *VecOp;
462462
ConstantInt *Cst;
463-
if (!match(Src, m_OneUse(m_ExtractElt(m_Value(VecOp), m_ConstantInt(Cst)))))
463+
const APInt *ShiftAmount = nullptr;
464+
if (!match(Src, m_OneUse(m_ExtractElt(m_Value(VecOp), m_ConstantInt(Cst)))) &&
465+
!match(Src,
466+
m_OneUse(m_LShr(m_ExtractElt(m_Value(VecOp), m_ConstantInt(Cst)),
467+
m_APInt(ShiftAmount)))))
464468
return nullptr;
465469

466470
auto *VecOpTy = cast<VectorType>(VecOp->getType());
@@ -469,10 +473,23 @@ static Instruction *foldVecExtTruncToExtElt(TruncInst &Trunc,
469473
uint64_t BitCastNumElts = VecElts.getKnownMinValue() * TruncRatio;
470474
uint64_t VecOpIdx = Cst->getZExtValue();
471475
uint64_t NewIdx = IC.getDataLayout().isBigEndian()
472-
? (VecOpIdx + 1) * TruncRatio - 1
476+
? (VecOpIdx * TruncRatio) + (TruncRatio - 1)
473477
: VecOpIdx * TruncRatio;
478+
479+
// Adjust index by the whole number of truncated elements.
480+
if (ShiftAmount) {
481+
// Check shift amount is in range and shifts a whole number of truncated
482+
// elements.
483+
if (ShiftAmount->uge(SrcBits) || ShiftAmount->urem(DstBits) != 0)
484+
return nullptr;
485+
486+
uint64_t IdxOfs = ShiftAmount->udiv(DstBits).getZExtValue();
487+
NewIdx = IC.getDataLayout().isBigEndian() ? (NewIdx - IdxOfs)
488+
: (NewIdx + IdxOfs);
489+
}
490+
474491
assert(BitCastNumElts <= std::numeric_limits<uint32_t>::max() &&
475-
"overflow 32-bits");
492+
NewIdx <= std::numeric_limits<uint32_t>::max() && "overflow 32-bits");
476493

477494
auto *BitCastTo =
478495
VectorType::get(DstType, BitCastNumElts, VecElts.isScalable());

‎llvm/test/Transforms/InstCombine/trunc-extractelement-inseltpoison.ll‎

Lines changed: 40 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,15 @@ define i32 @shrinkExtractElt_i64_to_i32_0(<3 x i64> %x) {
1919
}
2020

2121
define i32 @shrinkShiftExtractElt_i64_to_i32_0(<3 x i64> %x) {
22-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
23-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 0
24-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 32
25-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i32
26-
; ANY-NEXT: ret i32 [[T]]
22+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
23+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <6 x i32>
24+
; LE-NEXT: [[T:%.*]] = extractelement <6 x i32> [[TMP1]], i64 1
25+
; LE-NEXT: ret i32 [[T]]
26+
;
27+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
28+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <6 x i32>
29+
; BE-NEXT: [[T:%.*]] = extractelement <6 x i32> [[TMP1]], i64 0
30+
; BE-NEXT: ret i32 [[T]]
2731
;
2832
%e = extractelement <3 x i64> %x, i32 0
2933
%s = lshr i64 %e, 32
@@ -48,11 +52,15 @@ define i32 @vscale_shrinkExtractElt_i64_to_i32_0(<vscale x 3 x i64> %x) {
4852
}
4953

5054
define i32 @vscale_shrinkShiftExtractElt_i64_to_i32_0(<vscale x 3 x i64> %x) {
51-
; ANY-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
52-
; ANY-NEXT: [[E:%.*]] = extractelement <vscale x 3 x i64> [[X:%.*]], i64 0
53-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 32
54-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i32
55-
; ANY-NEXT: ret i32 [[T]]
55+
; LE-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
56+
; LE-NEXT: [[TMP1:%.*]] = bitcast <vscale x 3 x i64> [[X:%.*]] to <vscale x 6 x i32>
57+
; LE-NEXT: [[T:%.*]] = extractelement <vscale x 6 x i32> [[TMP1]], i64 1
58+
; LE-NEXT: ret i32 [[T]]
59+
;
60+
; BE-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
61+
; BE-NEXT: [[TMP1:%.*]] = bitcast <vscale x 3 x i64> [[X:%.*]] to <vscale x 6 x i32>
62+
; BE-NEXT: [[T:%.*]] = extractelement <vscale x 6 x i32> [[TMP1]], i64 0
63+
; BE-NEXT: ret i32 [[T]]
5664
;
5765
%e = extractelement <vscale x 3 x i64> %x, i32 0
5866
%s = lshr i64 %e, 32
@@ -109,11 +117,15 @@ define i16 @shrinkExtractElt_i64_to_i16_0(<3 x i64> %x) {
109117
}
110118

111119
define i16 @shrinkShiftExtractElt_i64_to_i16_0(<3 x i64> %x) {
112-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
113-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 0
114-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 48
115-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i16
116-
; ANY-NEXT: ret i16 [[T]]
120+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
121+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <12 x i16>
122+
; LE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 3
123+
; LE-NEXT: ret i16 [[T]]
124+
;
125+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
126+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <12 x i16>
127+
; BE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 0
128+
; BE-NEXT: ret i16 [[T]]
117129
;
118130
%e = extractelement <3 x i64> %x, i16 0
119131
%s = ashr i64 %e, 48
@@ -242,12 +254,19 @@ define i16 @shrinkShiftExtractElt_i64_to_i16_2_extra_shift_use(<3 x i64> %x) {
242254

243255
; OK to reuse the extract if we remove the shift+trunc.
244256
define i16 @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(<3 x i64> %x) {
245-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
246-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
247-
; ANY-NEXT: call void @use(i64 [[E]])
248-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 48
249-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i16
250-
; ANY-NEXT: ret i16 [[T]]
257+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
258+
; LE-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
259+
; LE-NEXT: call void @use(i64 [[E]])
260+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X]] to <12 x i16>
261+
; LE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 11
262+
; LE-NEXT: ret i16 [[T]]
263+
;
264+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
265+
; BE-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
266+
; BE-NEXT: call void @use(i64 [[E]])
267+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X]] to <12 x i16>
268+
; BE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 8
269+
; BE-NEXT: ret i16 [[T]]
251270
;
252271
%e = extractelement <3 x i64> %x, i64 2
253272
call void @use(i64 %e)

‎llvm/test/Transforms/InstCombine/trunc-extractelement.ll‎

Lines changed: 40 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,15 @@ define i32 @shrinkExtractElt_i64_to_i32_0(<3 x i64> %x) {
1919
}
2020

2121
define i32 @shrinkShiftExtractElt_i64_to_i32_0(<3 x i64> %x) {
22-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
23-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 0
24-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 32
25-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i32
26-
; ANY-NEXT: ret i32 [[T]]
22+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
23+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <6 x i32>
24+
; LE-NEXT: [[T:%.*]] = extractelement <6 x i32> [[TMP1]], i64 1
25+
; LE-NEXT: ret i32 [[T]]
26+
;
27+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i32_0(
28+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <6 x i32>
29+
; BE-NEXT: [[T:%.*]] = extractelement <6 x i32> [[TMP1]], i64 0
30+
; BE-NEXT: ret i32 [[T]]
2731
;
2832
%e = extractelement <3 x i64> %x, i32 0
2933
%s = lshr i64 %e, 32
@@ -48,11 +52,15 @@ define i32 @vscale_shrinkExtractElt_i64_to_i32_0(<vscale x 3 x i64> %x) {
4852
}
4953

5054
define i32 @vscale_shrinkShiftExtractElt_i64_to_i32_0(<vscale x 3 x i64> %x) {
51-
; ANY-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
52-
; ANY-NEXT: [[E:%.*]] = extractelement <vscale x 3 x i64> [[X:%.*]], i64 0
53-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 32
54-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i32
55-
; ANY-NEXT: ret i32 [[T]]
55+
; LE-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
56+
; LE-NEXT: [[TMP1:%.*]] = bitcast <vscale x 3 x i64> [[X:%.*]] to <vscale x 6 x i32>
57+
; LE-NEXT: [[T:%.*]] = extractelement <vscale x 6 x i32> [[TMP1]], i64 1
58+
; LE-NEXT: ret i32 [[T]]
59+
;
60+
; BE-LABEL: @vscale_shrinkShiftExtractElt_i64_to_i32_0(
61+
; BE-NEXT: [[TMP1:%.*]] = bitcast <vscale x 3 x i64> [[X:%.*]] to <vscale x 6 x i32>
62+
; BE-NEXT: [[T:%.*]] = extractelement <vscale x 6 x i32> [[TMP1]], i64 0
63+
; BE-NEXT: ret i32 [[T]]
5664
;
5765
%e = extractelement <vscale x 3 x i64> %x, i32 0
5866
%s = lshr i64 %e, 32
@@ -109,11 +117,15 @@ define i16 @shrinkExtractElt_i64_to_i16_0(<3 x i64> %x) {
109117
}
110118

111119
define i16 @shrinkShiftExtractElt_i64_to_i16_0(<3 x i64> %x) {
112-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
113-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 0
114-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 48
115-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i16
116-
; ANY-NEXT: ret i16 [[T]]
120+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
121+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <12 x i16>
122+
; LE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 3
123+
; LE-NEXT: ret i16 [[T]]
124+
;
125+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i16_0(
126+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X:%.*]] to <12 x i16>
127+
; BE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 0
128+
; BE-NEXT: ret i16 [[T]]
117129
;
118130
%e = extractelement <3 x i64> %x, i16 0
119131
%s = ashr i64 %e, 48
@@ -242,12 +254,19 @@ define i16 @shrinkShiftExtractElt_i64_to_i16_2_extra_shift_use(<3 x i64> %x) {
242254

243255
; OK to reuse the extract if we remove the shift+trunc.
244256
define i16 @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(<3 x i64> %x) {
245-
; ANY-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
246-
; ANY-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
247-
; ANY-NEXT: call void @use(i64 [[E]])
248-
; ANY-NEXT: [[S:%.*]] = lshr i64 [[E]], 48
249-
; ANY-NEXT: [[T:%.*]] = trunc nuw i64 [[S]] to i16
250-
; ANY-NEXT: ret i16 [[T]]
257+
; LE-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
258+
; LE-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
259+
; LE-NEXT: call void @use(i64 [[E]])
260+
; LE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X]] to <12 x i16>
261+
; LE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 11
262+
; LE-NEXT: ret i16 [[T]]
263+
;
264+
; BE-LABEL: @shrinkShiftExtractElt_i64_to_i16_2_extra_extract_use(
265+
; BE-NEXT: [[E:%.*]] = extractelement <3 x i64> [[X:%.*]], i64 2
266+
; BE-NEXT: call void @use(i64 [[E]])
267+
; BE-NEXT: [[TMP1:%.*]] = bitcast <3 x i64> [[X]] to <12 x i16>
268+
; BE-NEXT: [[T:%.*]] = extractelement <12 x i16> [[TMP1]], i64 8
269+
; BE-NEXT: ret i16 [[T]]
251270
;
252271
%e = extractelement <3 x i64> %x, i64 2
253272
call void @use(i64 %e)

0 commit comments

Comments
 (0)