Skip to content

Commit b42f553

Browse files
author
Thorsten Schütt
authored
[GlobalIsel] Combine extract vector element (#90339)
look through shuffle vectors
1 parent 879245e commit b42f553

File tree

6 files changed

+198
-15
lines changed

6 files changed

+198
-15
lines changed

‎llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -848,6 +848,11 @@ class CombinerHelper {
848848
bool matchExtractVectorElementWithBuildVectorTrunc(const MachineOperand &MO,
849849
BuildFnTy &MatchInfo);
850850

851+
/// Combine extract vector element with a shuffle vector on the vector
852+
/// register.
853+
bool matchExtractVectorElementWithShuffleVector(const MachineOperand &MO,
854+
BuildFnTy &MatchInfo);
855+
851856
/// Combine extract vector element with a insert vector element on the vector
852857
/// register and different indices.
853858
bool matchExtractVectorElementWithDifferentIndices(const MachineOperand &MO,

‎llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,18 @@ class GBuildVectorTrunc : public GMergeLikeInstr {
294294
}
295295
};
296296

297+
/// Represents a G_SHUFFLE_VECTOR.
298+
class GShuffleVector : public GenericMachineInstr {
299+
public:
300+
Register getSrc1Reg() const { return getOperand(1).getReg(); }
301+
Register getSrc2Reg() const { return getOperand(2).getReg(); }
302+
ArrayRef<int> getMask() const { return getOperand(3).getShuffleMask(); }
303+
304+
static bool classof(const MachineInstr *MI) {
305+
return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
306+
}
307+
};
308+
297309
/// Represents a G_PTR_ADD.
298310
class GPtrAdd : public GenericMachineInstr {
299311
public:

‎llvm/include/llvm/Target/GlobalISel/Combine.td

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1501,6 +1501,13 @@ def extract_vector_element_freeze : GICombineRule<
15011501
[{ return Helper.matchExtractVectorElementWithFreeze(${root}, ${matchinfo}); }]),
15021502
(apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
15031503

1504+
def extract_vector_element_shuffle_vector : GICombineRule<
1505+
(defs root:$root, build_fn_matchinfo:$matchinfo),
1506+
(match (G_SHUFFLE_VECTOR $src, $src1, $src2, $mask),
1507+
(G_EXTRACT_VECTOR_ELT $root, $src, $idx),
1508+
[{ return Helper.matchExtractVectorElementWithShuffleVector(${root}, ${matchinfo}); }]),
1509+
(apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
1510+
15041511
// Combines concat operations
15051512
def concat_matchinfo : GIDefMatchData<"SmallVector<Register>">;
15061513
def combine_concat_vector : GICombineRule<
@@ -1578,6 +1585,7 @@ extract_vector_element_build_vector_trunc6,
15781585
extract_vector_element_build_vector_trunc7,
15791586
extract_vector_element_build_vector_trunc8,
15801587
extract_vector_element_freeze,
1588+
extract_vector_element_shuffle_vector,
15811589
insert_vector_element_extract_vector_element
15821590
]>;
15831591

‎llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,112 @@ bool CombinerHelper::matchExtractVectorElementWithBuildVectorTrunc(
325325
return true;
326326
}
327327

328+
bool CombinerHelper::matchExtractVectorElementWithShuffleVector(
329+
const MachineOperand &MO, BuildFnTy &MatchInfo) {
330+
GExtractVectorElement *Extract =
331+
cast<GExtractVectorElement>(getDefIgnoringCopies(MO.getReg(), MRI));
332+
333+
//
334+
// %zero:_(s64) = G_CONSTANT i64 0
335+
// %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
336+
// shufflemask(0, 0, 0, 0)
337+
// %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %zero(s64)
338+
//
339+
// -->
340+
//
341+
// %zero1:_(s64) = G_CONSTANT i64 0
342+
// %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), %zero1(s64)
343+
//
344+
//
345+
//
346+
//
347+
// %three:_(s64) = G_CONSTANT i64 3
348+
// %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
349+
// shufflemask(0, 0, 0, -1)
350+
// %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %three(s64)
351+
//
352+
// -->
353+
//
354+
// %extract:_(s32) = G_IMPLICIT_DEF
355+
//
356+
//
357+
//
358+
//
359+
//
360+
// %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
361+
// shufflemask(0, 0, 0, -1)
362+
// %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
363+
//
364+
// -->
365+
//
366+
// %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
367+
// shufflemask(0, 0, 0, -1)
368+
// %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
369+
//
370+
371+
// We try to get the value of the Index register.
372+
std::optional<ValueAndVReg> MaybeIndex =
373+
getIConstantVRegValWithLookThrough(Extract->getIndexReg(), MRI);
374+
if (!MaybeIndex)
375+
return false;
376+
377+
GShuffleVector *Shuffle =
378+
cast<GShuffleVector>(getDefIgnoringCopies(Extract->getVectorReg(), MRI));
379+
380+
ArrayRef<int> Mask = Shuffle->getMask();
381+
382+
unsigned Offset = MaybeIndex->Value.getZExtValue();
383+
int SrcIdx = Mask[Offset];
384+
385+
LLT Src1Type = MRI.getType(Shuffle->getSrc1Reg());
386+
// At the IR level a <1 x ty> shuffle vector is valid, but we want to extract
387+
// from a vector.
388+
assert(Src1Type.isVector() && "expected to extract from a vector");
389+
unsigned LHSWidth = Src1Type.isVector() ? Src1Type.getNumElements() : 1;
390+
391+
// Note that there is no one use check.
392+
Register Dst = Extract->getReg(0);
393+
LLT DstTy = MRI.getType(Dst);
394+
395+
if (SrcIdx < 0 &&
396+
isLegalOrBeforeLegalizer({TargetOpcode::G_IMPLICIT_DEF, {DstTy}})) {
397+
MatchInfo = [=](MachineIRBuilder &B) { B.buildUndef(Dst); };
398+
return true;
399+
}
400+
401+
// If the legality check failed, then we still have to abort.
402+
if (SrcIdx < 0)
403+
return false;
404+
405+
Register NewVector;
406+
407+
// We check in which vector and at what offset to look through.
408+
if (SrcIdx < (int)LHSWidth) {
409+
NewVector = Shuffle->getSrc1Reg();
410+
// SrcIdx unchanged
411+
} else { // SrcIdx >= LHSWidth
412+
NewVector = Shuffle->getSrc2Reg();
413+
SrcIdx -= LHSWidth;
414+
}
415+
416+
LLT IdxTy = MRI.getType(Extract->getIndexReg());
417+
LLT NewVectorTy = MRI.getType(NewVector);
418+
419+
// We check the legality of the look through.
420+
if (!isLegalOrBeforeLegalizer(
421+
{TargetOpcode::G_EXTRACT_VECTOR_ELT, {DstTy, NewVectorTy, IdxTy}}) ||
422+
!isConstantLegalOrBeforeLegalizer({IdxTy}))
423+
return false;
424+
425+
// We look through the shuffle vector.
426+
MatchInfo = [=](MachineIRBuilder &B) {
427+
auto Idx = B.buildConstant(IdxTy, SrcIdx);
428+
B.buildExtractVectorElement(Dst, NewVector, Idx);
429+
};
430+
431+
return true;
432+
}
433+
328434
bool CombinerHelper::matchInsertVectorElementOOB(MachineInstr &MI,
329435
BuildFnTy &MatchInfo) {
330436
GInsertVectorElement *Insert = cast<GInsertVectorElement>(&MI);

‎llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -571,3 +571,66 @@ body: |
571571
RET_ReallyLR implicit $x0
572572
...
573573
---
574+
name: extract_from_build_vector_shuffle_vector_undef
575+
body: |
576+
bb.1:
577+
liveins: $x0, $x1
578+
; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_undef
579+
; CHECK: liveins: $x0, $x1
580+
; CHECK-NEXT: {{ $}}
581+
; CHECK-NEXT: %extract:_(s32) = G_IMPLICIT_DEF
582+
; CHECK-NEXT: $w0 = COPY %extract(s32)
583+
; CHECK-NEXT: RET_ReallyLR implicit $x0
584+
%arg1:_(<4 x s32>) = COPY $q0
585+
%arg2:_(<4 x s32>) = COPY $q1
586+
%idx:_(s64) = G_CONSTANT i64 0
587+
%sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
588+
%extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
589+
$w0 = COPY %extract(s32)
590+
RET_ReallyLR implicit $x0
591+
...
592+
---
593+
name: extract_from_build_vector_shuffle_vector_opaque
594+
body: |
595+
bb.1:
596+
liveins: $x0, $x1
597+
; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_opaque
598+
; CHECK: liveins: $x0, $x1
599+
; CHECK-NEXT: {{ $}}
600+
; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
601+
; CHECK-NEXT: %arg2:_(<4 x s32>) = COPY $q1
602+
; CHECK-NEXT: %idx:_(s64) = COPY $x1
603+
; CHECK-NEXT: %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2, shufflemask(undef, 0, 0, 0)
604+
; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
605+
; CHECK-NEXT: $w0 = COPY %extract(s32)
606+
; CHECK-NEXT: RET_ReallyLR implicit $x0
607+
%arg1:_(<4 x s32>) = COPY $q0
608+
%arg2:_(<4 x s32>) = COPY $q1
609+
%idx:_(s64) = COPY $x1
610+
%sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
611+
%extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
612+
$w0 = COPY %extract(s32)
613+
RET_ReallyLR implicit $x0
614+
...
615+
---
616+
name: extract_from_build_vector_shuffle_vector_const
617+
body: |
618+
bb.1:
619+
liveins: $x0, $x1
620+
; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_const
621+
; CHECK: liveins: $x0, $x1
622+
; CHECK-NEXT: {{ $}}
623+
; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
624+
; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 3
625+
; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), [[C]](s64)
626+
; CHECK-NEXT: $w0 = COPY %extract(s32)
627+
; CHECK-NEXT: RET_ReallyLR implicit $x0
628+
%arg1:_(<4 x s32>) = COPY $q0
629+
%arg2:_(<4 x s32>) = COPY $q1
630+
%idx:_(s64) = G_CONSTANT i64 0
631+
%sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(3, 0, 0, 0)
632+
%extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
633+
$w0 = COPY %extract(s32)
634+
RET_ReallyLR implicit $x0
635+
...
636+
---

‎llvm/test/CodeGen/AArch64/extract-vector-elt.ll

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -938,21 +938,10 @@ entry:
938938
}
939939

940940
define i32 @extract_v4i32_shuffle_const(<4 x i32> %a, <4 x i32> %b, i32 %c) {
941-
; CHECK-SD-LABEL: extract_v4i32_shuffle_const:
942-
; CHECK-SD: // %bb.0: // %entry
943-
; CHECK-SD-NEXT: fmov w0, s1
944-
; CHECK-SD-NEXT: ret
945-
;
946-
; CHECK-GI-LABEL: extract_v4i32_shuffle_const:
947-
; CHECK-GI: // %bb.0: // %entry
948-
; CHECK-GI-NEXT: adrp x8, .LCPI36_0
949-
; CHECK-GI-NEXT: // kill: def $q0 killed $q0 killed $q0_q1 def $q0_q1
950-
; CHECK-GI-NEXT: ldr q2, [x8, :lo12:.LCPI36_0]
951-
; CHECK-GI-NEXT: // kill: def $q1 killed $q1 killed $q0_q1 def $q0_q1
952-
; CHECK-GI-NEXT: tbl v0.16b, { v0.16b, v1.16b }, v2.16b
953-
; CHECK-GI-NEXT: mov s0, v0.s[2]
954-
; CHECK-GI-NEXT: fmov w0, s0
955-
; CHECK-GI-NEXT: ret
941+
; CHECK-LABEL: extract_v4i32_shuffle_const:
942+
; CHECK: // %bb.0: // %entry
943+
; CHECK-NEXT: fmov w0, s1
944+
; CHECK-NEXT: ret
956945
entry:
957946
%vector = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 2, i32 4, i32 3>
958947
%d = extractelement <4 x i32> %vector, i32 2

0 commit comments

Comments
 (0)