[llvm] [GlobalIsel] Combine extract vector element (PR #90339)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Apr 27 04:11:39 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Thorsten Schütt (tschuett)
<details>
<summary>Changes</summary>
look through shuffle vectors
---
Full diff: https://github.com/llvm/llvm-project/pull/90339.diff
6 Files Affected:
- (modified) llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h (+5)
- (modified) llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h (+12)
- (modified) llvm/include/llvm/Target/GlobalISel/Combine.td (+8)
- (modified) llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp (+107)
- (modified) llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir (+81)
- (modified) llvm/test/CodeGen/AArch64/extract-vector-elt.ll (+4-15)
``````````diff
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
index 76e8d1166ae0cd..4f1c9642e117c2 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
@@ -848,6 +848,11 @@ class CombinerHelper {
bool matchExtractVectorElementWithBuildVectorTrunc(const MachineOperand &MO,
BuildFnTy &MatchInfo);
+ /// Combine extract vector element with a shuffle vector on the vector
+ /// register.
+ bool matchExtractVectorElementWithShuffleVector(const MachineOperand &MO,
+ BuildFnTy &MatchInfo);
+
/// Combine extract vector element with a insert vector element on the vector
/// register and different indices.
bool matchExtractVectorElementWithDifferentIndices(const MachineOperand &MO,
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h b/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
index 25e47114e4a39a..47ad0c3d6430e8 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
@@ -294,6 +294,18 @@ class GBuildVectorTrunc : public GMergeLikeInstr {
}
};
+/// Respresents a G_SHUFFLE_VECTOR.
+class GShuffleVector : public GenericMachineInstr {
+public:
+ Register getSrc1Reg() const { return getOperand(1).getReg(); }
+ Register getSrc2Reg() const { return getOperand(2).getReg(); }
+ llvm::ArrayRef<int> getMask() const { return getOperand(3).getShuffleMask(); }
+
+ static bool classof(const MachineInstr *MI) {
+ return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
+ }
+};
+
/// Represents a G_PTR_ADD.
class GPtrAdd : public GenericMachineInstr {
public:
diff --git a/llvm/include/llvm/Target/GlobalISel/Combine.td b/llvm/include/llvm/Target/GlobalISel/Combine.td
index dbbb3abaa83047..55c17014e29808 100644
--- a/llvm/include/llvm/Target/GlobalISel/Combine.td
+++ b/llvm/include/llvm/Target/GlobalISel/Combine.td
@@ -1505,6 +1505,13 @@ def extract_vector_element_freeze : GICombineRule<
[{ return Helper.matchExtractVectorElementWithFreeze(${root}, ${matchinfo}); }]),
(apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
+def extract_vector_element_shuffle_vector : GICombineRule<
+ (defs root:$root, build_fn_matchinfo:$matchinfo),
+ (match (G_SHUFFLE_VECTOR $src, $src1, $src2, $mask),
+ (G_EXTRACT_VECTOR_ELT $root, $src, $idx),
+ [{ return Helper.matchExtractVectorElementWithShuffleVector(${root}, ${matchinfo}); }]),
+ (apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
+
// Combines concat operations
def concat_matchinfo : GIDefMatchData<"SmallVector<Register>">;
def combine_concat_vector : GICombineRule<
@@ -1582,6 +1589,7 @@ extract_vector_element_build_vector_trunc6,
extract_vector_element_build_vector_trunc7,
extract_vector_element_build_vector_trunc8,
extract_vector_element_freeze,
+extract_vector_element_shuffle_vector,
insert_vector_element_extract_vector_element
]>;
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
index fb33801a3a33fc..e1ba49ec2f0fa9 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
@@ -325,6 +325,113 @@ bool CombinerHelper::matchExtractVectorElementWithBuildVectorTrunc(
return true;
}
+bool CombinerHelper::matchExtractVectorElementWithShuffleVector(
+ const MachineOperand &MO, BuildFnTy &MatchInfo) {
+ GExtractVectorElement *Extract =
+ cast<GExtractVectorElement>(getDefIgnoringCopies(MO.getReg(), MRI));
+
+ //
+ // %zero:_(s64) = G_CONSTANT i64 0
+ // %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+ // shufflemask(0, 0, 0, 0)
+ // %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %zero(s64)
+ //
+ // -->
+ //
+ // %zero1:_(s64) = G_CONSTANT i64 0
+ // %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), %zero1(s64)
+ //
+ //
+ //
+ //
+ // %three:_(s64) = G_CONSTANT i64 3
+ // %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+ // shufflemask(0, 0, 0, -1)
+ // %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %three(s64)
+ //
+ // -->
+ //
+ // %extract:_(s32) = G_IMPLICIT_DEF
+ //
+ //
+ //
+ //
+ //
+ // %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+ // shufflemask(0, 0, 0, -1)
+ // %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
+ //
+ // -->
+ //
+ // %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+ // shufflemask(0, 0, 0, -1)
+ // %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
+ //
+ //
+
+ // We try to get the value of the Index register.
+ std::optional<ValueAndVReg> MaybeIndex =
+ getIConstantVRegValWithLookThrough(Extract->getIndexReg(), MRI);
+ if (!MaybeIndex)
+ return false;
+
+ GShuffleVector *Shuffle =
+ cast<GShuffleVector>(getDefIgnoringCopies(Extract->getVectorReg(), MRI));
+
+ llvm::ArrayRef<int> Mask = Shuffle->getMask();
+
+ unsigned Offset = MaybeIndex->Value.getZExtValue();
+ int SrcIdx = Mask[Offset];
+
+ LLT Src1Type = MRI.getType(Shuffle->getSrc1Reg());
+ // At the IR level a <1 x ty> shuffle vector is valid, but we want to extract
+ // from a vector.
+ assert(Src1Type.isVector() && "expected to extract from a vector");
+ unsigned LHSWidth = Src1Type.isVector() ? Src1Type.getNumElements() : 1;
+
+ // Note that there is no one use check.
+ Register Dst = Extract->getReg(0);
+ LLT DstTy = MRI.getType(Dst);
+
+ if (SrcIdx < 0 &&
+ isLegalOrBeforeLegalizer({TargetOpcode::G_IMPLICIT_DEF, {DstTy}})) {
+ MatchInfo = [=](MachineIRBuilder &B) { B.buildUndef(Dst); };
+ return true;
+ }
+
+ // If the legality check failed, then we still have to abort.
+ if (SrcIdx < 0)
+ return false;
+
+ Register NewVector;
+
+ // We check in which vector and at what offset to look through.
+ if (SrcIdx < (int)LHSWidth) {
+ NewVector = Shuffle->getSrc1Reg();
+ // SrcIdx unchanged
+ } else { // SrcIdx >= LHSWidth
+ NewVector = Shuffle->getSrc2Reg();
+ SrcIdx -= LHSWidth;
+ }
+
+ LLT IdxTy = MRI.getType(Extract->getIndexReg());
+ LLT NewVectorTy = MRI.getType(NewVector);
+
+ // We check the legality of the look through.
+ if (!isLegalOrBeforeLegalizer(
+ {TargetOpcode::G_EXTRACT_VECTOR_ELT, {DstTy, NewVectorTy, IdxTy}}) ||
+ !isConstantLegalOrBeforeLegalizer({IdxTy}))
+ return false;
+
+ // We look through the shuffle vector.
+ MatchInfo = [=](MachineIRBuilder &B) {
+ auto Idx = B.buildConstant(IdxTy, SrcIdx);
+ B.buildExtractVectorElement(Dst, NewVector, Idx);
+ };
+
+ return true;
+}
+
bool CombinerHelper::matchInsertVectorElementOOB(MachineInstr &MI,
BuildFnTy &MatchInfo) {
GInsertVectorElement *Insert = cast<GInsertVectorElement>(&MI);
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
index 587d53c300f83f..3b7515b407e37f 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
@@ -571,3 +571,84 @@ body: |
RET_ReallyLR implicit $x0
...
---
+name: extract_from_build_vector_shuffle_vector_undef
+alignment: 4
+liveins:
+ - { reg: '$x0' }
+ - { reg: '$x1' }
+frameInfo:
+ maxAlignment: 1
+body: |
+ bb.1:
+ liveins: $x0, $x1
+ ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_undef
+ ; CHECK: liveins: $x0, $x1
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %extract:_(s32) = G_IMPLICIT_DEF
+ ; CHECK-NEXT: $w0 = COPY %extract(s32)
+ ; CHECK-NEXT: RET_ReallyLR implicit $x0
+ %arg1:_(<4 x s32>) = COPY $q0
+ %arg2:_(<4 x s32>) = COPY $q1
+ %idx:_(s64) = G_CONSTANT i64 0
+ %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
+ %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+ $w0 = COPY %extract(s32)
+ RET_ReallyLR implicit $x0
+...
+---
+name: extract_from_build_vector_shuffle_vector_opaque
+alignment: 4
+liveins:
+ - { reg: '$x0' }
+ - { reg: '$x1' }
+frameInfo:
+ maxAlignment: 1
+body: |
+ bb.1:
+ liveins: $x0, $x1
+ ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_opaque
+ ; CHECK: liveins: $x0, $x1
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
+ ; CHECK-NEXT: %arg2:_(<4 x s32>) = COPY $q1
+ ; CHECK-NEXT: %idx:_(s64) = COPY $x1
+ ; CHECK-NEXT: %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2, shufflemask(undef, 0, 0, 0)
+ ; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+ ; CHECK-NEXT: $w0 = COPY %extract(s32)
+ ; CHECK-NEXT: RET_ReallyLR implicit $x0
+ %arg1:_(<4 x s32>) = COPY $q0
+ %arg2:_(<4 x s32>) = COPY $q1
+ %idx:_(s64) = COPY $x1
+ %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
+ %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+ $w0 = COPY %extract(s32)
+ RET_ReallyLR implicit $x0
+...
+---
+name: extract_from_build_vector_shuffle_vector_const
+alignment: 4
+liveins:
+ - { reg: '$x0' }
+ - { reg: '$x1' }
+frameInfo:
+ maxAlignment: 1
+body: |
+ bb.1:
+ liveins: $x0, $x1
+ ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_const
+ ; CHECK: liveins: $x0, $x1
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
+ ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 3
+ ; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), [[C]](s64)
+ ; CHECK-NEXT: $w0 = COPY %extract(s32)
+ ; CHECK-NEXT: RET_ReallyLR implicit $x0
+ %arg1:_(<4 x s32>) = COPY $q0
+ %arg2:_(<4 x s32>) = COPY $q1
+ %idx:_(s64) = G_CONSTANT i64 0
+ %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(3, 0, 0, 0)
+ %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+ $w0 = COPY %extract(s32)
+ RET_ReallyLR implicit $x0
+...
+---
diff --git a/llvm/test/CodeGen/AArch64/extract-vector-elt.ll b/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
index 504222e0036e22..0481d997d24faf 100644
--- a/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
+++ b/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
@@ -938,21 +938,10 @@ entry:
}
define i32 @extract_v4i32_shuffle_const(<4 x i32> %a, <4 x i32> %b, i32 %c) {
-; CHECK-SD-LABEL: extract_v4i32_shuffle_const:
-; CHECK-SD: // %bb.0: // %entry
-; CHECK-SD-NEXT: fmov w0, s1
-; CHECK-SD-NEXT: ret
-;
-; CHECK-GI-LABEL: extract_v4i32_shuffle_const:
-; CHECK-GI: // %bb.0: // %entry
-; CHECK-GI-NEXT: adrp x8, .LCPI36_0
-; CHECK-GI-NEXT: // kill: def $q0 killed $q0 killed $q0_q1 def $q0_q1
-; CHECK-GI-NEXT: ldr q2, [x8, :lo12:.LCPI36_0]
-; CHECK-GI-NEXT: // kill: def $q1 killed $q1 killed $q0_q1 def $q0_q1
-; CHECK-GI-NEXT: tbl v0.16b, { v0.16b, v1.16b }, v2.16b
-; CHECK-GI-NEXT: mov s0, v0.s[2]
-; CHECK-GI-NEXT: fmov w0, s0
-; CHECK-GI-NEXT: ret
+; CHECK-LABEL: extract_v4i32_shuffle_const:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: fmov w0, s1
+; CHECK-NEXT: ret
entry:
%vector = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 2, i32 4, i32 3>
%d = extractelement <4 x i32> %vector, i32 2
``````````
</details>
https://github.com/llvm/llvm-project/pull/90339
More information about the llvm-commits
mailing list