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 29, 2024
1 parent 5654047 commit cc0f633
Show file tree
Hide file tree
Showing 5 changed files with 415 additions and 2 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 ty X,Y,Z
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
19 changes: 19 additions & 0 deletions llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// Modifications (c) Copyright 2024 Advanced Micro Devices, Inc. or its
// affiliates
//
//===----------------------------------------------------------------------===//
/// \file
/// Declares convenience wrapper classes for interpreting MachineInstr instances
Expand Down Expand Up @@ -225,6 +228,22 @@ class GUnmerge : public GenericMachineInstr {
}
};

/// Represents a G_SHUFFLE_VECTOR.
class GShuffleVector : public GenericMachineInstr {
public:
/// Returns the number of source registers.
unsigned getNumSources() const { return getNumOperands() - 2; }
/// Returns the I'th source register.
Register getSourceReg(unsigned I) const {
assert(I + 1 <= getNumSources());
return getReg(I + 1);
}

static bool classof(const MachineInstr *MI) {
return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
}
};

/// 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.
Expand Down
14 changes: 12 additions & 2 deletions llvm/include/llvm/Target/GlobalISel/Combine.td
Original file line number Diff line number Diff line change
Expand Up @@ -749,7 +749,7 @@ def fneg_fneg_fold: GICombineRule <
(apply (GIReplaceReg $dst, $src))
>;

// Fold (unmerge(merge x, y, z)) -> z, y, z.
// Fold (unmerge(merge x, y, z)) -> x, y, z.
def unmerge_merge_matchinfo : GIDefMatchData<"SmallVector<Register, 8>">;
def unmerge_merge : GICombineRule<
(defs root:$d, unmerge_merge_matchinfo:$info),
Expand All @@ -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
76 changes: 76 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,80 @@ 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 && !IsUndef1) || (!SrcInstr2 && !IsUndef2))
return false;

if (IsUndef1 && IsUndef2)
return true;

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 cc0f633

Please sign in to comment.