Skip to content

Commit

Permalink
[AIE2][AIE2P] Legalize G_SELECT for 512 bits only.
Browse files Browse the repository at this point in the history
  • Loading branch information
@SagarMaheshwari99
SagarMaheshwari99 committed Jan 30, 2025
1 parent 512188c commit 379c797
Show file tree
Hide file tree
Showing 8 changed files with 295 additions and 374 deletions.
34 changes: 30 additions & 4 deletions llvm/lib/Target/AIE/AIE2LegalizerInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
//
//===----------------------------------------------------------------------===//
/// \file
Expand Down Expand Up @@ -73,6 +73,13 @@ static LegalityPredicate isValidVectorAIE2(const unsigned TypeIdx) {
};
}

static LegalityPredicate vectorSmallerThan(unsigned TypeIdx, unsigned Size) {
return [=](const LegalityQuery &Query) {
const LLT QueryTy = Query.Types[TypeIdx];
return QueryTy.isVector() && QueryTy.getSizeInBits() < Size;
};
}

LegalityPredicate
negatePredicate(const std::function<bool(const LegalityQuery &)> &Func) {
return [=](const LegalityQuery &Query) { return !Func(Query); };
Expand Down Expand Up @@ -236,10 +243,24 @@ AIE2LegalizerInfo::AIE2LegalizerInfo(const AIE2Subtarget &ST) : AIEHelper(ST) {

getActionDefinitionsBuilder(G_SELECT)
.legalFor({{S32, S32}, {P0, S32}})
.clampScalar(1, S32, S32)
// AIE ISA supports only 512-bit vector select
.legalFor({V16S32, V32S16, V64S8})
// For scalar types >= 256, bitcast to a vector type and use existing
// selection patterns
.bitcastIf(
[=](const LegalityQuery &Query) {
const LLT &ResTy = Query.Types[0];
return ResTy.isScalar() && ResTy.getSizeInBits() >= 256;
},
[=](const LegalityQuery &Query) {
const LLT Ty = Query.Types[0];
const unsigned Size = Ty.getSizeInBits();
return std::pair(0, LLT::fixed_vector(Size / 32, LLT::scalar(32)));
})
.widenScalarToNextPow2(0)
.clampScalar(0, S32, S32)
.clampScalar(1, S32, S32)
.legalFor(AIE2VectorTypes)
.customIf(vectorSmallerThan(0, 512))
// We support G_SELECT only on the vector register bank
// Mapping the G_SELECT operands to the vector register bank
// during register bank selection introduces the proper cross-bank
Expand All @@ -248,7 +269,10 @@ AIE2LegalizerInfo::AIE2LegalizerInfo(const AIE2Subtarget &ST) : AIEHelper(ST) {
// type patterns in C++. Introducing bitcasts during legalization allows
// to re-use the existing code for register bank selection and ISEL
// patterns.
.bitcastIf(typeInSet(0, AIE2AccumulatorTypes), bitcastAccToVectorType(0));
.bitcastIf(typeInSet(0, AIE2AccumulatorTypes), bitcastAccToVectorType(0))
.clampMaxNumElements(0, S8, 64)
.clampMaxNumElements(0, S16, 32)
.clampMaxNumElements(0, S32, 16);

getActionDefinitionsBuilder({G_ADD, G_SUB})
.legalFor({S32})
Expand Down Expand Up @@ -546,6 +570,8 @@ bool AIE2LegalizerInfo::legalizeCustom(
return AIEHelper.legalizeG_SEXT_INREG(Helper, MI);
case TargetOpcode::G_BITCAST:
return AIEHelper.legalizeG_BITCAST(Helper, MI);
case TargetOpcode::G_SELECT:
return AIEHelper.legalizeG_SELECTLessThan512Bit(Helper, MI);
}

llvm_unreachable("Un-expected custom legalization");
Expand Down
68 changes: 64 additions & 4 deletions llvm/lib/Target/AIE/AIELegalizerHelper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
//
//===----------------------------------------------------------------------===//
/// \file
Expand Down Expand Up @@ -1186,19 +1186,79 @@ bool AIELegalizerHelper::legalizeLoopDecrement(LegalizerHelper &Helper,
return true;
}

// Legalize 2048-bit G_SELECT
bool AIELegalizerHelper::legalizeG_SELECT(LegalizerHelper &Helper,
MachineInstr &MI) const {
// Legalize < 512-bit G_SELECT
// Expand the source vectors to 512-bits by padding it with undefs.
// SrcReg0 controls the selection, effectively targeting element 0.
bool AIELegalizerHelper::legalizeG_SELECTLessThan512Bit(
LegalizerHelper &Helper, MachineInstr &MI) const {
MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
MachineRegisterInfo &MRI = *MIRBuilder.getMRI();

const Register DstReg = MI.getOperand(0).getReg();
const LLT DstTy = MRI.getType(DstReg);
const unsigned DstVecSize = DstTy.getSizeInBits();

assert(DstTy.isVector() && DstVecSize < 512 &&
"Expected to legalize < 512-bit vector G_SELECT");
assert(!(512 % DstVecSize) && "Vector size should be a factor of 512");

const Register SrcReg0 = MI.getOperand(1).getReg(); // Scalar
const Register SrcReg1 = MI.getOperand(2).getReg();
const Register SrcReg2 = MI.getOperand(3).getReg();

const LLT NewVecTy = LLT::fixed_vector(
512 / DstTy.getElementType().getSizeInBits(), DstTy.getElementType());

const Register UndefReg = MRI.createGenericVirtualRegister(DstTy);
MIRBuilder.buildUndef(UndefReg);

const unsigned NumPadElts = (512 / DstVecSize) - 1;
auto buildMergeInstr = [&](const Register SrcReg) -> Register {
SmallVector<Register, 4> Regs;
Regs.push_back(SrcReg);
for (unsigned I = 0; I < NumPadElts; I++)
Regs.push_back(UndefReg);
const Register NewSrcReg = MRI.createGenericVirtualRegister(NewVecTy);
MIRBuilder.buildMergeLikeInstr(NewSrcReg, Regs);
return NewSrcReg;
};

const Register NewSrcReg1 = buildMergeInstr(SrcReg1);
const Register NewSrcReg2 = buildMergeInstr(SrcReg2);

const Register NewDstReg = MRI.createGenericVirtualRegister(NewVecTy);
MIRBuilder.buildInstr(MI.getOpcode(), {NewDstReg},
{SrcReg0, NewSrcReg1, NewSrcReg2}, MI.getFlags());

SmallVector<Register, 4> Regs;
Regs.push_back(DstReg);
for (unsigned I = 0; I < NumPadElts; ++I)
Regs.push_back(MRI.createGenericVirtualRegister(DstTy));
MIRBuilder.buildUnmerge(Regs, NewDstReg);

MI.eraseFromParent();
return true;
}

// Legalize 256-bit / 2048-bit G_SELECT
bool AIELegalizerHelper::legalizeG_SELECT(LegalizerHelper &Helper,
MachineInstr &MI) const {
MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
MachineRegisterInfo &MRI = *MIRBuilder.getMRI();

const Register DstReg = MI.getOperand(0).getReg();
const LLT DstTy = MRI.getType(DstReg);

if (DstTy.isVector() && DstTy.getSizeInBits() < 512)
return legalizeG_SELECTLessThan512Bit(Helper, MI);

assert(DstTy.isVector() && DstTy.getSizeInBits() == 2048 &&
"Expected to legalize 2048-bit vector G_SELECT");

const Register SrcReg0 = MI.getOperand(1).getReg(); // Scalar
const Register SrcReg1 = MI.getOperand(2).getReg();
const Register SrcReg2 = MI.getOperand(3).getReg();

const LLT DstVecEltTy = DstTy.getElementType();
const unsigned ElTySize = DstVecEltTy.getSizeInBits();
const LLT ACC1024 = LLT::fixed_vector(1024 / ElTySize, ElTySize);
Expand Down
4 changes: 3 additions & 1 deletion llvm/lib/Target/AIE/AIELegalizerHelper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
//
//===----------------------------------------------------------------------===//
/// \file
Expand Down Expand Up @@ -57,6 +57,8 @@ class AIELegalizerHelper {
bool legalizeG_FABS(LegalizerHelper &Helper, MachineInstr &MI) const;
bool legalizeG_FADDSUB(LegalizerHelper &Helper, MachineInstr &MI) const;
bool legalizeG_SELECT(LegalizerHelper &Helper, MachineInstr &MI) const;
bool legalizeG_SELECTLessThan512Bit(LegalizerHelper &Helper,
MachineInstr &MI) const;
bool legalizeG_BITCAST(LegalizerHelper &Helper, MachineInstr &MI) const;
bool legalizeLoopDecrement(LegalizerHelper &Helper, MachineInstr &MI) const;
bool legalizeG_CONCAT_VECTORS(LegalizerHelper &Helper,
Expand Down
15 changes: 9 additions & 6 deletions llvm/lib/Target/AIE/aie2p/AIE2PLegalizerInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
// (c) Copyright 2024-2025 Advanced Micro Devices, Inc. or its affiliates
//
//===----------------------------------------------------------------------===//
/// \file
Expand Down Expand Up @@ -271,15 +271,14 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
getActionDefinitionsBuilder(G_SELECT)
.legalFor({{S32, S32}, {P0, S32}})
.clampScalar(1, S32, S32)
.legalFor(AIE2PVectorTypes)
// AIE ISA supports only 512-bit vector select
.legalFor({V16S32, V32S16, V64S8})
// For scalar types >= 256, bitcast to a vector type and use existing
// selection patterns
.bitcastIf(
[=](const LegalityQuery &Query) {
const LLT &ResTy = Query.Types[0];
if (ResTy.isVector())
return false;
return ResTy.getSizeInBits() >= 256;
return ResTy.isScalar() && ResTy.getSizeInBits() >= 256;
},
[=](const LegalityQuery &Query) {
const LLT Ty = Query.Types[0];
Expand All @@ -288,6 +287,7 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
})
.widenScalarToNextPow2(0)
.clampScalar(0, S32, S32)
.customIf(vectorSmallerThan(0, 512))
// We support G_SELECT only on the vector register bank
// Mapping the G_SELECT operands to the vector register bank
// during register bank selection introduces the proper cross-bank
Expand All @@ -298,7 +298,10 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
// patterns.
.bitcastIf(typeInSet(0, {AccV4S64, AccV8S64, AccV16S64}),
bitcastAccToVectorType(0))
.customFor({{AccV64S32, S32}});
.customFor({{AccV64S32, S32}})
.clampMaxNumElements(0, S8, 64)
.clampMaxNumElements(0, S16, 32)
.clampMaxNumElements(0, S32, 16);

getActionDefinitionsBuilder({G_ADD, G_SUB, G_XOR})
.legalFor({S32})
Expand Down
Loading

0 comments on commit 379c797

Please sign in to comment.