GlobalISel: Implement MoreElements for select of vector conditions

This commit is contained in:
Matt Arsenault 2022-04-11 21:31:15 -04:00
parent 6009122250
commit 3754f60112
2 changed files with 68 additions and 4 deletions

View File

@ -4910,10 +4910,26 @@ LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
moreElementsVectorDst(MI, MoreTy, 0);
Observer.changedInstr(MI);
return Legalized;
case TargetOpcode::G_SELECT:
if (TypeIdx != 0)
case TargetOpcode::G_SELECT: {
Register DstReg = MI.getOperand(0).getReg();
Register CondReg = MI.getOperand(1).getReg();
LLT DstTy = MRI.getType(DstReg);
LLT CondTy = MRI.getType(CondReg);
if (TypeIdx == 1) {
if (!CondTy.isScalar() ||
DstTy.getElementCount() != MoreTy.getElementCount())
return UnableToLegalize;
if (MRI.getType(MI.getOperand(1).getReg()).isVector())
// This is turning a scalar select of vectors into a vector
// select. Broadcast the select condition.
auto ShufSplat = MIRBuilder.buildShuffleSplat(MoreTy, CondReg);
Observer.changingInstr(MI);
MI.getOperand(1).setReg(ShufSplat.getReg(0));
Observer.changedInstr(MI);
return Legalized;
}
if (CondTy.isVector())
return UnableToLegalize;
Observer.changingInstr(MI);
@ -4922,6 +4938,7 @@ LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
moreElementsVectorDst(MI, MoreTy, 0);
Observer.changedInstr(MI);
return Legalized;
}
case TargetOpcode::G_UNMERGE_VALUES:
return UnableToLegalize;
case TargetOpcode::G_PHI:

View File

@ -4127,4 +4127,51 @@ TEST_F(AArch64GISelMITest, narrowScalarShiftByConstant) {
EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
}
TEST_F(AArch64GISelMITest, MoreElementsSelect) {
setUp();
if (!TM)
return;
LLT s1 = LLT::scalar(1);
LLT s64 = LLT::scalar(64);
LLT v2s1 = LLT::fixed_vector(2, 1);
LLT v2s32 = LLT::fixed_vector(2, 32);
LegalizerInfo LI;
DummyGISelObserver Observer;
LegalizerHelper Helper(*MF, LI, Observer, B);
B.setInsertPt(*EntryMBB, EntryMBB->end());
auto Val0 = B.buildBitcast(v2s32, Copies[0]);
auto Val1 = B.buildBitcast(v2s32, Copies[1]);
// Build select of vectors with scalar condition.
auto Zero = B.buildConstant(s64, 0);
auto Cond = B.buildICmp(CmpInst::ICMP_EQ, s1, Copies[2], Zero);
auto Select = B.buildSelect(v2s32, Cond, Val0, Val1);
// Splat the condition into a vector select
B.setInstr(*Select);
EXPECT_EQ(LegalizerHelper::LegalizeResult::UnableToLegalize,
Helper.moreElementsVector(*Select, 1, LLT::fixed_vector(3, 1)));
EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
Helper.moreElementsVector(*Select, 1, v2s1));
auto CheckStr = R"(
CHECK: [[BITCAST0:%[0-9]+]]:_(<2 x s32>) = G_BITCAST
CHECK: [[BITCAST1:%[0-9]+]]:_(<2 x s32>) = G_BITCAST
CHECK: [[ZERO0:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
CHECK: [[CMP:%[0-9]+]]:_(s1) = G_ICMP intpred(eq), %{{[0-9]+}}:_(s64), [[ZERO0]]
CHECK: [[IMPDEF:%[0-9]+]]:_(<2 x s1>) = G_IMPLICIT_DEF
CHECK: [[ZERO1:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
CHECK: [[INSERT:%[0-9]+]]:_(<2 x s1>) = G_INSERT_VECTOR_ELT [[IMPDEF]]:_, [[CMP]]:_(s1), [[ZERO1]]
CHECK: [[SHUFFLE:%[0-9]+]]:_(<2 x s1>) = G_SHUFFLE_VECTOR [[INSERT]]:_(<2 x s1>), [[IMPDEF]]:_, shufflemask(0, 0)
CHECK: [[SELECT:%[0-9]+]]:_(<2 x s32>) = G_SELECT [[SHUFFLE]]:_(<2 x s1>), [[BITCAST0]]:_, [[BITCAST1]]:_
)";
EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
}
} // namespace