[llvm] [GISel][AMDGPU] Fold ShuffleVec into ExtractSubvec, and custom lower ExtractSubvec (PR #124527)
Alan Li via llvm-commits
llvm-commits at lists.llvm.org
Tue Feb 11 13:14:54 PST 2025
https://github.com/lialan updated https://github.com/llvm/llvm-project/pull/124527
>From 7544ec3c7bc79517af4b6c55ac1d2b0ede2aec94 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 27 Jan 2025 18:26:19 +0800
Subject: [PATCH 1/6] First commit
---
.../llvm/CodeGen/GlobalISel/CombinerHelper.h | 4 ++
.../include/llvm/Target/GlobalISel/Combine.td | 9 ++-
.../lib/CodeGen/GlobalISel/CombinerHelper.cpp | 41 +++++++++++++
.../lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 59 +++++++++++++++++++
llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h | 2 +
5 files changed, 114 insertions(+), 1 deletion(-)
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
index 9b78342c8fc39..c1c303fd18e6b 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
@@ -264,6 +264,10 @@ class CombinerHelper {
void applyCombineShuffleConcat(MachineInstr &MI,
SmallVector<Register> &Ops) const;
+ /// Replace \p MI with a narrow extract_subvector.
+ bool matchCombineShuffleExtract(MachineInstr &MI, int64_t &IsFirst) const;
+ void applyCombineShuffleExtract(MachineInstr &MI, int64_t IsFirst) const;
+
/// Try to combine G_SHUFFLE_VECTOR into G_CONCAT_VECTORS.
/// Returns true if MI changed.
///
diff --git a/llvm/include/llvm/Target/GlobalISel/Combine.td b/llvm/include/llvm/Target/GlobalISel/Combine.td
index 3590ab221ad44..30316305d9e4f 100644
--- a/llvm/include/llvm/Target/GlobalISel/Combine.td
+++ b/llvm/include/llvm/Target/GlobalISel/Combine.td
@@ -1560,6 +1560,13 @@ def combine_shuffle_concat : GICombineRule<
[{ return Helper.matchCombineShuffleConcat(*${root}, ${matchinfo}); }]),
(apply [{ Helper.applyCombineShuffleConcat(*${root}, ${matchinfo}); }])>;
+// Combines shuffles of vector into extract_subvector
+def combine_shuffle_vector : GICombineRule<
+ (defs root:$root, int64_matchinfo:$matchinfo),
+ (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
+ [{ return Helper.matchCombineShuffleExtract(*${root}, ${matchinfo}); }]),
+ (apply [{ Helper.applyCombineShuffleExtract(*${root}, ${matchinfo}); }])>;
+
def insert_vector_element_idx_undef : GICombineRule<
(defs root:$root),
(match (G_IMPLICIT_DEF $idx),
@@ -2026,7 +2033,7 @@ def all_combines : GICombineGroup<[integer_reassoc_combines, trivial_combines,
and_or_disjoint_mask, fma_combines, fold_binop_into_select,
sub_add_reg, select_to_minmax,
fsub_to_fneg, commute_constant_to_rhs, match_ands, match_ors,
- simplify_neg_minmax, combine_concat_vector,
+ simplify_neg_minmax, combine_concat_vector, combine_shuffle_vector,
sext_trunc, zext_trunc, prefer_sign_combines, shuffle_combines,
combine_use_vector_truncate, merge_combines, overflow_combines]>;
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 0dfbb91f2ac54..2e517304d527a 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -384,6 +384,47 @@ void CombinerHelper::applyCombineConcatVectors(
MI.eraseFromParent();
}
+bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx) const {
+ assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR &&
+ "Invalid instruction");
+ auto &Shuffle = cast<GShuffleVector>(MI);
+ const auto &TLI = getTargetLowering();
+
+ auto SrcVec1 = Shuffle.getSrc1Reg();
+ auto SrcVec2 = Shuffle.getSrc2Reg();
+ auto Mask = Shuffle.getMask();
+
+ int Width = MRI.getType(SrcVec1).getNumElements();
+
+ // Check if all elements are extracted from the same vector, or within single
+ // vector.
+ auto MaxValue = *std::max_element(Mask.begin(), Mask.end());
+ auto MinValue = *std::min_element(Mask.begin(), Mask.end());
+ if (MaxValue >= Width && MinValue < Width) {
+ return false;
+ }
+ // Check if the extractee's order is kept:
+ if (!std::is_sorted(Mask.begin(), Mask.end())) {
+ return false;
+ }
+
+ Idx = Mask.front();
+ return true;
+}
+
+void CombinerHelper::applyCombineShuffleExtract(MachineInstr &MI, int64_t Idx) const {
+ auto &Shuffle = cast<GShuffleVector>(MI);
+
+ auto SrcVec1 = Shuffle.getSrc1Reg();
+ auto SrcVec2 = Shuffle.getSrc2Reg();
+ int Width = MRI.getType(SrcVec1).getNumElements();
+
+ auto SrcVec = Idx < Width ? SrcVec1 : SrcVec2;
+
+ Builder.buildExtractSubvector(MI.getOperand(0).getReg(), SrcVec, Idx);
+ MI.eraseFromParent();
+}
+
bool CombinerHelper::matchCombineShuffleConcat(
MachineInstr &MI, SmallVector<Register> &Ops) const {
ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index e9e47eaadd557..68b0a8b5aecbf 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -29,6 +29,7 @@
#include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
#include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
#include "llvm/CodeGen/GlobalISel/Utils.h"
+#include "llvm/CodeGen/Register.h"
#include "llvm/CodeGen/TargetOpcodes.h"
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
@@ -1832,6 +1833,11 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
.lower();
}
+ getActionDefinitionsBuilder(G_EXTRACT_SUBVECTOR)
+ //.fewerElementsIf(isWideVec16(0), changeTo(0, V2S16))
+ .customFor({V8S16, V4S16})
+ .lower();
+
getActionDefinitionsBuilder(G_EXTRACT_VECTOR_ELT)
.unsupportedIf([=](const LegalityQuery &Query) {
const LLT &EltTy = Query.Types[1].getElementType();
@@ -2127,6 +2133,8 @@ bool AMDGPULegalizerInfo::legalizeCustom(
case TargetOpcode::G_FMINNUM_IEEE:
case TargetOpcode::G_FMAXNUM_IEEE:
return legalizeMinNumMaxNum(Helper, MI);
+ case TargetOpcode::G_EXTRACT_SUBVECTOR:
+ return legalizeExtractSubvector(MI, MRI, B);
case TargetOpcode::G_EXTRACT_VECTOR_ELT:
return legalizeExtractVectorElt(MI, MRI, B);
case TargetOpcode::G_INSERT_VECTOR_ELT:
@@ -2716,6 +2724,57 @@ bool AMDGPULegalizerInfo::legalizeMinNumMaxNum(LegalizerHelper &Helper,
return Helper.lowerFMinNumMaxNum(MI) == LegalizerHelper::Legalized;
}
+static auto buildExtractSubvector(MachineIRBuilder &B, SrcOp Src,
+ LLT DstTy, unsigned Start) {
+ SmallVector<Register, 8> Subvectors;
+ for (unsigned i = Start, e = Start + DstTy.getNumElements(); i != e; ++i) {
+ Subvectors.push_back(
+ B.buildExtractVectorElementConstant(DstTy.getElementType(), Src, i)
+ .getReg(0));
+ }
+ return B.buildBuildVector(DstTy, Subvectors);
+}
+
+bool AMDGPULegalizerInfo::legalizeExtractSubvector(
+ MachineInstr &MI, MachineRegisterInfo &MRI,
+ MachineIRBuilder &B) const {
+ const auto &Instr = llvm::cast<GExtractSubvector>(MI);
+ Register Src = Instr.getSrcVec();
+ Register Dst = MI.getOperand(0).getReg();
+ auto Start = Instr.getIndexImm();
+
+ LLT SrcTy = MRI.getType(Src);
+ LLT DstTy = MRI.getType(Dst);
+
+ LLT EltTy = SrcTy.getElementType();
+ assert(EltTy == DstTy.getElementType());
+ auto Count = DstTy.getNumElements();
+ assert(SrcTy.getNumElements() % 2 == 0 && Count % 2 == 0);
+
+ // Split vector size into legal sub vectors, and use build_vector
+ // to merge the result.
+ if (EltTy.getScalarSizeInBits() == 16 && Start % 2 == 0) {
+ bool UseScalar = Count == 2;
+ // Extract 32-bit registers at a time.
+ LLT NewSrcTy =
+ UseScalar ? S32 : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
+ auto Bitcasted = B.buildBitcast(NewSrcTy, Src).getReg(0);
+ LLT NewDstTy = LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
+
+ SmallVector<Register, 8> Subvectors;
+ for (unsigned i = Start / 2, e = (Start + Count) / 2; i != e; ++i) {
+ auto Subvec = B.buildExtractVectorElementConstant(S32, Bitcasted, i);
+ Subvectors.push_back(Subvec.getReg(0));
+ }
+
+ auto BuildVec = B.buildBuildVector(NewDstTy, Subvectors);
+ B.buildBitcast(Dst, BuildVec.getReg(0));
+ MI.eraseFromParent();
+ return true;
+ }
+ return false;
+}
+
bool AMDGPULegalizerInfo::legalizeExtractVectorElt(
MachineInstr &MI, MachineRegisterInfo &MRI,
MachineIRBuilder &B) const {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
index 86c15197805d2..7b55492afb982 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h
@@ -56,6 +56,8 @@ class AMDGPULegalizerInfo final : public LegalizerInfo {
bool legalizeFPTOI(MachineInstr &MI, MachineRegisterInfo &MRI,
MachineIRBuilder &B, bool Signed) const;
bool legalizeMinNumMaxNum(LegalizerHelper &Helper, MachineInstr &MI) const;
+ bool legalizeExtractSubvector(MachineInstr &MI, MachineRegisterInfo &MRI,
+ MachineIRBuilder &B) const;
bool legalizeExtractVectorElt(MachineInstr &MI, MachineRegisterInfo &MRI,
MachineIRBuilder &B) const;
bool legalizeInsertVectorElt(MachineInstr &MI, MachineRegisterInfo &MRI,
>From ffdddd603e63b976e2d1f5455b6ea10448a20d77 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 27 Jan 2025 19:17:12 +0800
Subject: [PATCH 2/6] update
---
llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp | 14 ++++++++++++--
llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 12 ++++--------
2 files changed, 16 insertions(+), 10 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 2e517304d527a..fb91d7446319a 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -388,13 +388,16 @@ bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx)
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR &&
"Invalid instruction");
auto &Shuffle = cast<GShuffleVector>(MI);
- const auto &TLI = getTargetLowering();
auto SrcVec1 = Shuffle.getSrc1Reg();
- auto SrcVec2 = Shuffle.getSrc2Reg();
+ int SrcVec2 = Shuffle.getSrc2Reg();
auto Mask = Shuffle.getMask();
int Width = MRI.getType(SrcVec1).getNumElements();
+ int Width2 = MRI.getType(SrcVec2).getNumElements();
+
+ if (!llvm::isPowerOf2_32(Width))
+ return false;
// Check if all elements are extracted from the same vector, or within single
// vector.
@@ -403,6 +406,13 @@ bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx)
if (MaxValue >= Width && MinValue < Width) {
return false;
}
+
+ // Check that the extractee length is power of 2.
+ if ((MaxValue < Width && !llvm::isPowerOf2_32(Width)) ||
+ (MinValue >= Width && !llvm::isPowerOf2_32(Width2))) {
+ return false;
+ }
+
// Check if the extractee's order is kept:
if (!std::is_sorted(Mask.begin(), Mask.end())) {
return false;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 68b0a8b5aecbf..75017c23bb502 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -2758,16 +2758,12 @@ bool AMDGPULegalizerInfo::legalizeExtractSubvector(
// Extract 32-bit registers at a time.
LLT NewSrcTy =
UseScalar ? S32 : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
- auto Bitcasted = B.buildBitcast(NewSrcTy, Src).getReg(0);
LLT NewDstTy = LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
+ auto Bitcasted = B.buildBitcast(NewSrcTy, Src);
- SmallVector<Register, 8> Subvectors;
- for (unsigned i = Start / 2, e = (Start + Count) / 2; i != e; ++i) {
- auto Subvec = B.buildExtractVectorElementConstant(S32, Bitcasted, i);
- Subvectors.push_back(Subvec.getReg(0));
- }
-
- auto BuildVec = B.buildBuildVector(NewDstTy, Subvectors);
+ auto BuildVec =
+ UseScalar ? Bitcasted
+ : buildExtractSubvector(B, Bitcasted, NewDstTy, Start / 2);
B.buildBitcast(Dst, BuildVec.getReg(0));
MI.eraseFromParent();
return true;
>From 967138e82c39c560c8b84f8909a58fdfc7d0fa0e Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 27 Jan 2025 22:15:15 +0800
Subject: [PATCH 3/6] small fix
---
llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp | 12 ++++++++++--
llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 4 ++--
2 files changed, 12 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index fb91d7446319a..88b0dfd46652a 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -407,14 +407,22 @@ bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx)
return false;
}
+ LLT ExtractTy = MaxValue < Width ? MRI.getType(SrcVec1) : MRI.getType(SrcVec2);
+
// Check that the extractee length is power of 2.
if ((MaxValue < Width && !llvm::isPowerOf2_32(Width)) ||
(MinValue >= Width && !llvm::isPowerOf2_32(Width2))) {
return false;
}
- // Check if the extractee's order is kept:
- if (!std::is_sorted(Mask.begin(), Mask.end())) {
+ // Check if the extractee's order is kept, and they should be conscecutive.
+ for (size_t i = 1; i < Mask.size(); ++i) {
+ if (Mask[i] != Mask[i - 1] + 1 || Mask[i] == -1) {
+ return false; // Not consecutive
+ }
+ }
+
+ if (!LI->isLegalOrCustom({TargetOpcode::G_EXTRACT_SUBVECTOR, {ExtractTy}})) {
return false;
}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 75017c23bb502..5873b7a56d6da 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -1834,8 +1834,8 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
}
getActionDefinitionsBuilder(G_EXTRACT_SUBVECTOR)
- //.fewerElementsIf(isWideVec16(0), changeTo(0, V2S16))
- .customFor({V8S16, V4S16})
+ .widenScalarOrEltToNextPow2(0)
+ .customFor(AllS16Vectors)
.lower();
getActionDefinitionsBuilder(G_EXTRACT_VECTOR_ELT)
>From 9c98f96659dbca047b0c22e7b188a2db5130f925 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 28 Jan 2025 12:44:48 +0800
Subject: [PATCH 4/6] Fix scalar corner issues.
---
llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 17 ++++++++++-------
1 file changed, 10 insertions(+), 7 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 5873b7a56d6da..33cc0c872c643 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -2727,6 +2727,9 @@ bool AMDGPULegalizerInfo::legalizeMinNumMaxNum(LegalizerHelper &Helper,
static auto buildExtractSubvector(MachineIRBuilder &B, SrcOp Src,
LLT DstTy, unsigned Start) {
SmallVector<Register, 8> Subvectors;
+ if (!DstTy.isVector()) {
+ return B.buildExtractVectorElementConstant(DstTy, Src, Start);
+ }
for (unsigned i = Start, e = Start + DstTy.getNumElements(); i != e; ++i) {
Subvectors.push_back(
B.buildExtractVectorElementConstant(DstTy.getElementType(), Src, i)
@@ -2754,16 +2757,16 @@ bool AMDGPULegalizerInfo::legalizeExtractSubvector(
// Split vector size into legal sub vectors, and use build_vector
// to merge the result.
if (EltTy.getScalarSizeInBits() == 16 && Start % 2 == 0) {
- bool UseScalar = Count == 2;
// Extract 32-bit registers at a time.
- LLT NewSrcTy =
- UseScalar ? S32 : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
- LLT NewDstTy = LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
+ LLT NewSrcTy = SrcTy.getNumElements() == 2
+ ? S32
+ : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
auto Bitcasted = B.buildBitcast(NewSrcTy, Src);
- auto BuildVec =
- UseScalar ? Bitcasted
- : buildExtractSubvector(B, Bitcasted, NewDstTy, Start / 2);
+ LLT NewDstTy = DstTy.getNumElements() == 2
+ ? S32
+ : LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
+ auto BuildVec = buildExtractSubvector(B, Bitcasted, NewDstTy, Start / 2);
B.buildBitcast(Dst, BuildVec.getReg(0));
MI.eraseFromParent();
return true;
>From a45ef34ad93772ca0cff5bdc43f5199a447cbe4b Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 11 Feb 2025 07:28:28 -0800
Subject: [PATCH 5/6] update
---
llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp | 13 ++++++++-----
1 file changed, 8 insertions(+), 5 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 88b0dfd46652a..e28f93711b215 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -37,6 +37,7 @@
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Target/TargetMachine.h"
+#include <algorithm>
#include <cmath>
#include <optional>
#include <tuple>
@@ -415,11 +416,13 @@ bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx)
return false;
}
- // Check if the extractee's order is kept, and they should be conscecutive.
- for (size_t i = 1; i < Mask.size(); ++i) {
- if (Mask[i] != Mask[i - 1] + 1 || Mask[i] == -1) {
- return false; // Not consecutive
- }
+ // Check if the extractee's order is kept, and they should be consecutive.
+ bool isConsecutive =
+ std::adjacent_find(Mask.begin(), Mask.end(), [](int a, int b) {
+ return b != a + 1 || b == -1;
+ }) == Mask.end();
+ if (!isConsecutive) {
+ return false;
}
if (!LI->isLegalOrCustom({TargetOpcode::G_EXTRACT_SUBVECTOR, {ExtractTy}})) {
>From 97da0ee0509125279f91d36c0b7f15fb8fd14b4b Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 11 Feb 2025 13:03:32 -0800
Subject: [PATCH 6/6] Checkup
---
.../llvm/CodeGen/GlobalISel/LegalizerHelper.h | 3 +
.../lib/CodeGen/GlobalISel/CombinerHelper.cpp | 14 ++---
.../CodeGen/GlobalISel/LegalizerHelper.cpp | 33 ++++++++++
.../lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp | 63 ++++++++++---------
4 files changed, 78 insertions(+), 35 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h
index 4e18f5cc913a7..ac2a763b22402 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h
@@ -347,6 +347,9 @@ class LegalizerHelper {
LegalizeResult narrowScalarShiftByConstant(MachineInstr &MI, const APInt &Amt,
LLT HalfTy, LLT ShiftAmtTy);
+ LegalizeResult fewerElementsExtractSubvector(MachineInstr &MI,
+ unsigned TypeIdx, LLT NarrowTy);
+
LegalizeResult fewerElementsVectorReductions(MachineInstr &MI,
unsigned TypeIdx, LLT NarrowTy);
LegalizeResult fewerElementsVectorSeqReductions(MachineInstr &MI,
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index e28f93711b215..baf1232e098c8 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -416,18 +416,18 @@ bool CombinerHelper::matchCombineShuffleExtract(MachineInstr &MI, int64_t &Idx)
return false;
}
- // Check if the extractee's order is kept, and they should be consecutive.
- bool isConsecutive =
- std::adjacent_find(Mask.begin(), Mask.end(), [](int a, int b) {
- return b != a + 1 || b == -1;
- }) == Mask.end();
- if (!isConsecutive) {
- return false;
+ // Check if the extractee's order is kept, and they should be conscecutive.
+ for (size_t i = 1; i < Mask.size(); ++i) {
+ if (Mask[i] != Mask[i - 1] + 1 || Mask[i] == -1) {
+ return false; // Not consecutive
+ }
}
+ /*
if (!LI->isLegalOrCustom({TargetOpcode::G_EXTRACT_SUBVECTOR, {ExtractTy}})) {
return false;
}
+ */
Idx = Mask.front();
return true;
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index d4cb224c35d74..b306ac786e456 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -5463,6 +5463,8 @@ LegalizerHelper::fewerElementsVector(MachineInstr &MI, unsigned TypeIdx,
return fewerElementsVectorSeqReductions(MI, TypeIdx, NarrowTy);
case G_SHUFFLE_VECTOR:
return fewerElementsVectorShuffle(MI, TypeIdx, NarrowTy);
+ case G_EXTRACT_SUBVECTOR:
+ return fewerElementsExtractSubvector(MI, TypeIdx, NarrowTy);
case G_FPOWI:
return fewerElementsVectorMultiEltType(GMI, NumElts, {2 /*pow*/});
case G_BITCAST:
@@ -5643,6 +5645,36 @@ LegalizerHelper::LegalizeResult LegalizerHelper::fewerElementsVectorShuffle(
return Legalized;
}
+LegalizerHelper::LegalizeResult LegalizerHelper::fewerElementsExtractSubvector(
+ MachineInstr &MI, unsigned int TypeIdx, LLT NarrowTy) {
+ assert(MI.getOpcode() == TargetOpcode::G_EXTRACT_SUBVECTOR);
+ if (TypeIdx != 0)
+ return UnableToLegalize;
+
+ auto [DstReg, DstTy, SrcReg, SrcTy] = MI.getFirst2RegLLTs();
+
+
+ if (!isPowerOf2_32(DstTy.getNumElements()))
+ return UnableToLegalize;
+
+ uint64_t SplitIdx = MI.getOperand(2).getImm();
+ unsigned NewElts = NarrowTy.isVector() ? NarrowTy.getNumElements() : 1;
+ unsigned SplitParts = DstTy.getNumElements() / NewElts;
+
+ // Split the sources into NarrowTy size pieces.
+ SmallVector<Register> SplitDstParts;
+
+ for (unsigned i = 0; i < SplitParts; i++) {
+ auto Part = MIRBuilder.buildExtractSubvector(NarrowTy, SrcReg,
+ i * NewElts + SplitIdx);
+ SplitDstParts.push_back(Part.getReg(0));
+ }
+
+ MIRBuilder.buildMergeLikeInstr(DstReg, SplitDstParts);
+ MI.eraseFromParent();
+ return Legalized;
+}
+
LegalizerHelper::LegalizeResult LegalizerHelper::fewerElementsVectorReductions(
MachineInstr &MI, unsigned int TypeIdx, LLT NarrowTy) {
auto &RdxMI = cast<GVecReduce>(MI);
@@ -6049,6 +6081,7 @@ LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
unsigned Opc = MI.getOpcode();
switch (Opc) {
case TargetOpcode::G_IMPLICIT_DEF:
+ case TargetOpcode::G_EXTRACT_SUBVECTOR:
case TargetOpcode::G_LOAD: {
if (TypeIdx != 0)
return UnableToLegalize;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
index 33cc0c872c643..8156bd326b25f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
@@ -108,6 +108,18 @@ static LegalizeMutation oneMoreElement(unsigned TypeIdx) {
};
}
+static LegalizeMutation fewerEltsToSize32Vector(unsigned TypeIdx) {
+ return [=](const LegalityQuery &Query) {
+ const LLT Ty = Query.Types[TypeIdx];
+ const LLT EltTy = Ty.getElementType();
+ unsigned Size = Ty.getSizeInBits();
+ unsigned Pieces = llvm::divideCeil(Size, 32);
+ unsigned NewNumElts = (Ty.getNumElements() + 1) / Pieces;
+ return std::pair(TypeIdx, LLT::scalarOrVector(
+ ElementCount::getFixed(NewNumElts), EltTy));
+ };
+}
+
static LegalizeMutation fewerEltsToSize64Vector(unsigned TypeIdx) {
return [=](const LegalityQuery &Query) {
const LLT Ty = Query.Types[TypeIdx];
@@ -1834,9 +1846,10 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
}
getActionDefinitionsBuilder(G_EXTRACT_SUBVECTOR)
- .widenScalarOrEltToNextPow2(0)
- .customFor(AllS16Vectors)
- .lower();
+ .widenScalarOrEltToNextPow2(0, 2)
+ .moreElementsIf(isSmallOddVector(0), oneMoreElement(0))
+ .fewerElementsIf(vectorWiderThan(0, 32), fewerEltsToSize32Vector(0))
+ .customIf(sizeIsMultipleOf32(0));
getActionDefinitionsBuilder(G_EXTRACT_VECTOR_ELT)
.unsupportedIf([=](const LegalityQuery &Query) {
@@ -2742,36 +2755,30 @@ bool AMDGPULegalizerInfo::legalizeExtractSubvector(
MachineInstr &MI, MachineRegisterInfo &MRI,
MachineIRBuilder &B) const {
const auto &Instr = llvm::cast<GExtractSubvector>(MI);
- Register Src = Instr.getSrcVec();
- Register Dst = MI.getOperand(0).getReg();
- auto Start = Instr.getIndexImm();
- LLT SrcTy = MRI.getType(Src);
- LLT DstTy = MRI.getType(Dst);
+ auto [DstReg, DstTy, SrcReg, SrcTy] = MI.getFirst2RegLLTs();
+ auto Start = Instr.getIndexImm();
LLT EltTy = SrcTy.getElementType();
assert(EltTy == DstTy.getElementType());
- auto Count = DstTy.getNumElements();
- assert(SrcTy.getNumElements() % 2 == 0 && Count % 2 == 0);
-
- // Split vector size into legal sub vectors, and use build_vector
- // to merge the result.
- if (EltTy.getScalarSizeInBits() == 16 && Start % 2 == 0) {
- // Extract 32-bit registers at a time.
- LLT NewSrcTy = SrcTy.getNumElements() == 2
- ? S32
- : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
- auto Bitcasted = B.buildBitcast(NewSrcTy, Src);
-
- LLT NewDstTy = DstTy.getNumElements() == 2
- ? S32
- : LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
- auto BuildVec = buildExtractSubvector(B, Bitcasted, NewDstTy, Start / 2);
- B.buildBitcast(Dst, BuildVec.getReg(0));
- MI.eraseFromParent();
- return true;
+
+ if (DstTy.getSizeInBits() % 32 != 0) {
+ return false;
}
- return false;
+
+ // Extract 32-bit registers at a time.
+ LLT NewSrcTy = SrcTy.getNumElements() == 2
+ ? S32
+ : LLT::fixed_vector(SrcTy.getNumElements() / 2, S32);
+ auto Bitcasted = B.buildBitcast(NewSrcTy, SrcReg);
+
+ LLT NewDstTy = DstTy.getNumElements() == 2
+ ? S32
+ : LLT::fixed_vector(DstTy.getNumElements() / 2, S32);
+ auto BuildVec = buildExtractSubvector(B, Bitcasted, NewDstTy, Start / 2);
+ B.buildBitcast(DstReg, BuildVec.getReg(0));
+ MI.eraseFromParent();
+ return true;
}
bool AMDGPULegalizerInfo::legalizeExtractVectorElt(
More information about the llvm-commits
mailing list