Skip to content

Commit

Permalink
[Generic][AIE2] Combiner for shufflevectors that use build vector
Browse files Browse the repository at this point in the history
Transforms a shufflevector that uses a build vector or undefined
into just a build vector. This can be done is because a shuffle
vector lowering is an unmerge and then merge. Since build is a
merge, the merge and unmerge cancel each other out and we can just
merge the vector directly.

Example:
```
    %1:_(s32) = COPY $r0
    %3:_(<8 x s32>) = G_IMPLICIT_DEF
    %5:_(s32) = G_IMPLICIT_DEF
    %2:_(<8 x s32>) = G_BUILD_VECTOR %1(s32), %5(s32), %5(s32), %5(s32), %5(s32), %5(s32), %5(s32), %5(s32)
    %0:_(<8 x s32>) = G_SHUFFLE_VECTOR %2(<8 x s32>), %3, shufflemask(0, 0, 0, 0, 0, 0, 0, 0)
    ===>
    %2:_(<8 x s32>) = G_BUILD_VECTOR %1(s32), %1(s32), %1(s32), %1(s32), %1(s32), %1(s32), %1(s32), %1(s32)
```
  • Loading branch information
ValentijnvdBeek committed Jul 19, 2024
1 parent 5654047 commit e0a5aae
Show file tree
Hide file tree
Showing 5 changed files with 427 additions and 4 deletions.
9 changes: 9 additions & 0 deletions llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,15 @@ class CombinerHelper {
applyCombineUnmergeMergeToPlainValues(MachineInstr &MI,
SmallVectorImpl<Register> &Operands);

/// Transform <ty, ...> G_SHUFFLE_VECTOR(G_MERGE ty X Y Z) -> G_MERGE
bool
matchCombineShuffleVectorBuildVector(MachineInstr &MI,
SmallVectorImpl<Register> &Operands);

void
applyCombineShuffleVectorBuildVector(MachineInstr &MI,
SmallVectorImpl<Register> &Operands);

/// Transform G_UNMERGE Constant -> Constant1, Constant2, ...
bool matchCombineUnmergeConstant(MachineInstr &MI,
SmallVectorImpl<APInt> &Csts);
Expand Down
15 changes: 12 additions & 3 deletions llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
Original file line number Diff line number Diff line change
Expand Up @@ -225,9 +225,9 @@ class GUnmerge : public GenericMachineInstr {
}
};

/// Represents G_BUILD_VECTOR, G_CONCAT_VECTORS or G_MERGE_VALUES.
/// All these have the common property of generating a single value from
/// multiple sources.
/// Represents G_BUILD_VECTOR, G_CONCAT_VECTORS, G_SHUFFLE_VECTOR or
/// G_MERGE_VALUES. All these have the common property of generating a single
/// value from multiple sources.
class GMergeLikeInstr : public GenericMachineInstr {
public:
/// Returns the number of source registers.
Expand All @@ -240,13 +240,22 @@ class GMergeLikeInstr : public GenericMachineInstr {
case TargetOpcode::G_MERGE_VALUES:
case TargetOpcode::G_CONCAT_VECTORS:
case TargetOpcode::G_BUILD_VECTOR:
case TargetOpcode::G_SHUFFLE_VECTOR:
return true;
default:
return false;
}
}
};

/// Represents a G_SHUFFLE_VECTOR
class GShuffleVector : public GMergeLikeInstr {
public:
static bool classof(const MachineInstr *MI) {
return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
}
};

/// Represents a G_MERGE_VALUES.
class GMerge : public GMergeLikeInstr {
public:
Expand Down
12 changes: 11 additions & 1 deletion llvm/include/llvm/Target/GlobalISel/Combine.td
Original file line number Diff line number Diff line change
Expand Up @@ -758,6 +758,16 @@ def unmerge_merge : GICombineRule<
(apply [{ Helper.applyCombineUnmergeMergeToPlainValues(*${d}, ${info}); }])
>;

// Fold (unmerge(merge x, y, z)) -> z, y, z.
def shufflevector_merge_matchinfo : GIDefMatchData<"SmallVector<Register, 8>">;
def shufflevector_merge : GICombineRule<
(defs root:$d, shufflevector_merge_matchinfo:$info),
(match (wip_match_opcode G_SHUFFLE_VECTOR): $d,
[{ return Helper.matchCombineShuffleVectorBuildVector(*${d}, ${info}); }]),
(apply [{ Helper.applyCombineShuffleVectorBuildVector(*${d}, ${info}); }])
>;


// Fold merge(unmerge).
def merge_unmerge : GICombineRule<
(defs root:$d, register_matchinfo:$matchinfo),
Expand Down Expand Up @@ -1324,7 +1334,7 @@ def all_combines : GICombineGroup<[trivial_combines, insert_vec_elt_combines,
intdiv_combines, mulh_combines, redundant_neg_operands,
and_or_disjoint_mask, fma_combines, fold_binop_into_select,
sub_add_reg, select_to_minmax, redundant_binop_in_equality,
fsub_to_fneg, commute_constant_to_rhs]>;
fsub_to_fneg, commute_constant_to_rhs, shufflevector_merge]>;

// A combine group used to for prelegalizer combiners at -O0. The combines in
// this group have been selected based on experiments to balance code size and
Expand Down
73 changes: 73 additions & 0 deletions llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
#include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
#include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
Expand All @@ -26,6 +27,7 @@
#include "llvm/CodeGen/MachineDominators.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineMemOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegisterBankInfo.h"
Expand Down Expand Up @@ -2147,6 +2149,77 @@ static Register peekThroughBitcast(Register Reg,
return Reg;
}

bool CombinerHelper::matchCombineShuffleVectorBuildVector(
MachineInstr &MI, SmallVectorImpl<Register> &Operands) {
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR &&
"Expected a shuffle vector");
auto &ShuffleVector = cast<GShuffleVector>(MI);
Register SrcReg1 = peekThroughBitcast(ShuffleVector.getSourceReg(0), MRI);
Register SrcReg2 = peekThroughBitcast(ShuffleVector.getSourceReg(1), MRI);

// Check if the Source registers are either merges or implicit definitions
auto *SrcInstr1 = getOpcodeDef<GBuildVector>(SrcReg1, MRI);
auto *SrcInstr2 = getOpcodeDef<GBuildVector>(SrcReg2, MRI);
auto *IsUndef1 = getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, SrcReg1, MRI);
auto *IsUndef2 = getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, SrcReg2, MRI);

// Our inputs need to be either be build vectors or undefined, register inputs
// break this optimization. You could maybe do something clever were you
// concatenate vectors to save half a build vector.
if ((SrcInstr1 == 0 && IsUndef1 == 0) || (SrcInstr2 == 0 && IsUndef2 == 0))
return false;

Register UndefReg;
if (SrcInstr1 || SrcInstr2)
UndefReg = MRI.createGenericVirtualRegister(MRI.getType(SrcReg1));

// Since our inputs to shufflevector must be of the same size, we can reuse
// the size of the defined register.
const unsigned NumElements = (SrcInstr1 != 0) ? SrcInstr1->getNumSources()
: SrcInstr2->getNumSources();
for (unsigned Idx = 0; Idx < NumElements; ++Idx) {
const Register Elt =
(SrcInstr1 != 0) ? SrcInstr1->getSourceReg(Idx) : UndefReg;
Operands.push_back(Elt);
}

for (unsigned Idx = 0; Idx < NumElements; ++Idx) {
const Register Elt =
(SrcInstr2 != 0) ? SrcInstr2->getSourceReg(Idx) : UndefReg;
Operands.push_back(Elt);
}

return true;
}

void CombinerHelper::applyCombineShuffleVectorBuildVector(
MachineInstr &MI, SmallVectorImpl<Register> &Operands) {
assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR &&
"Expected a shuffle vector");
auto &ShuffleVector = cast<GShuffleVector>(MI);
Register SrcReg1 = peekThroughBitcast(ShuffleVector.getSourceReg(0), MRI);
Register SrcReg2 = peekThroughBitcast(ShuffleVector.getSourceReg(1), MRI);

// Check if the Source registers are either merges or implicit definitions
auto *IsUndef1 = getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, SrcReg1, MRI);
auto *IsUndef2 = getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, SrcReg2, MRI);

// If they're both undefined, we will just return an undefined as well.
if (IsUndef1 && IsUndef2) {
Builder.buildUndef(ShuffleVector.getReg(0));
MI.eraseFromParent();
return;
}

const ArrayRef<int> ShiftMask = MI.getOperand(3).getShuffleMask();
SmallVector<Register, 8> Arguments;
for (int Index : ShiftMask)
Arguments.push_back(Operands[Index]);

Builder.buildBuildVector(ShuffleVector.getOperand(0), Arguments);
MI.eraseFromParent();
}

bool CombinerHelper::matchCombineUnmergeMergeToPlainValues(
MachineInstr &MI, SmallVectorImpl<Register> &Operands) {
assert(MI.getOpcode() == TargetOpcode::G_UNMERGE_VALUES &&
Expand Down
Loading

0 comments on commit e0a5aae

Please sign in to comment.