diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp --- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp +++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp @@ -861,6 +861,8 @@ return reduceOperationWidth(MI, TypeIdx, NarrowTy); case TargetOpcode::G_ADD: case TargetOpcode::G_SUB: + case TargetOpcode::G_UADDO: + case TargetOpcode::G_USUBO: return narrowScalarAddSub(MI, TypeIdx, NarrowTy); case TargetOpcode::G_MUL: case TargetOpcode::G_UMULH: @@ -4466,10 +4468,12 @@ unsigned OpO, OpE; switch (MI.getOpcode()) { + case TargetOpcode::G_UADDO: case TargetOpcode::G_ADD: OpO = TargetOpcode::G_UADDO; OpE = TargetOpcode::G_UADDE; break; + case TargetOpcode::G_USUBO: case TargetOpcode::G_SUB: OpO = TargetOpcode::G_USUBO; OpE = TargetOpcode::G_USUBE; @@ -4478,14 +4482,25 @@ llvm_unreachable("Unexpected add/sub opcode!"); } + // 1 for a plain add/sub, 2 if this is an operation with a carry-out. + unsigned NumDefs = MI.getNumExplicitDefs(); + Register Src1 = MI.getOperand(NumDefs).getReg(); + Register Src2 = MI.getOperand(NumDefs + 1).getReg(); + Register CarryDst; + if (NumDefs == 2) + CarryDst = MI.getOperand(1).getReg(); + SmallVector Src1Regs, Src2Regs, DstRegs; - extractParts(MI.getOperand(1).getReg(), NarrowTy, NumParts, Src1Regs); - extractParts(MI.getOperand(2).getReg(), NarrowTy, NumParts, Src2Regs); + extractParts(Src1, NarrowTy, NumParts, Src1Regs); + extractParts(Src2, NarrowTy, NumParts, Src2Regs); Register CarryIn; for (int i = 0; i < NumParts; ++i) { Register DstReg = MRI.createGenericVirtualRegister(NarrowTy); Register CarryOut = MRI.createGenericVirtualRegister(LLT::scalar(1)); + // Forward the final carry-out to the destination register + if (i == NumParts - 1 && CarryDst) + CarryOut = CarryDst; if (i == 0) MIRBuilder.buildInstr(OpO, {DstReg, CarryOut}, diff --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp --- a/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp +++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp @@ -880,6 +880,80 @@ EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF; } +TEST_F(AArch64GISelMITest, NarrowUADDO) { + setUp(); + if (!TM) + return; + + LLT S1 = LLT::scalar(1); + LLT S32 = LLT::scalar(32); + LLT S96 = LLT::scalar(96); + DefineLegalizerInfo(A, { + getActionDefinitionsBuilder({G_UADDO, G_UADDE}) + .legalFor({{LLT::scalar(32), LLT::scalar(1)}}); + }); + + auto Op0 = B.buildUndef(S96); + auto Op1 = B.buildUndef(S96); + auto UADDO = B.buildUAddo(S96, S1, Op0, Op1); + + AInfo Info(MF->getSubtarget()); + DummyGISelObserver Observer; + LegalizerHelper Helper(*MF, Info, Observer, B); + EXPECT_EQ(Helper.narrowScalar(*UADDO, 0, S32), + LegalizerHelper::LegalizeResult::Legalized); + + const char *CheckStr = R"( + CHECK: [[IMP_DEF0:%[0-9]+]]:_(s96) = G_IMPLICIT_DEF + CHECK: [[IMP_DEF1:%[0-9]+]]:_(s96) = G_IMPLICIT_DEF + CHECK: [[OP0_0:%[0-9]+]]:_(s32), [[OP0_1:%[0-9]+]]:_(s32), [[OP0_2:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[IMP_DEF0]] + CHECK: [[OP1_0:%[0-9]+]]:_(s32), [[OP1_1:%[0-9]+]]:_(s32), [[OP1_2:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[IMP_DEF1]] + CHECK: [[UADDO0:%[0-9]+]]:_(s32), [[CARRY0:%[0-9]+]]:_(s1) = G_UADDO [[OP0_0]]:_, [[OP1_0]]:_ + CHECK: [[UADDO1:%[0-9]+]]:_(s32), [[CARRY1:%[0-9]+]]:_(s1) = G_UADDE [[OP0_1]]:_, [[OP1_1]]:_, [[CARRY0]]:_ + CHECK: [[UADDO2:%[0-9]+]]:_(s32), [[CARRY2:%[0-9]+]]:_(s1) = G_UADDE [[OP0_2]]:_, [[OP1_2]]:_, [[CARRY1]]:_ + CHECK: [[UADDO:%[0-9]+]]:_(s96) = G_MERGE_VALUES [[UADDO0]]:_(s32), [[UADDO1]]:_(s32), [[UADDO2]]:_(s32) + )"; + + EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF; +} + +TEST_F(AArch64GISelMITest, NarrowUSUBO) { + setUp(); + if (!TM) + return; + + LLT S1 = LLT::scalar(1); + LLT S32 = LLT::scalar(32); + LLT S96 = LLT::scalar(96); + DefineLegalizerInfo(A, { + getActionDefinitionsBuilder({G_USUBO, G_USUBE}) + .legalFor({{LLT::scalar(32), LLT::scalar(1)}}); + }); + + auto Op0 = B.buildUndef(S96); + auto Op1 = B.buildUndef(S96); + auto USUBO = B.buildUSubo(S96, S1, Op0, Op1); + + AInfo Info(MF->getSubtarget()); + DummyGISelObserver Observer; + LegalizerHelper Helper(*MF, Info, Observer, B); + EXPECT_EQ(Helper.narrowScalar(*USUBO, 0, S32), + LegalizerHelper::LegalizeResult::Legalized); + + const char *CheckStr = R"( + CHECK: [[IMP_DEF0:%[0-9]+]]:_(s96) = G_IMPLICIT_DEF + CHECK: [[IMP_DEF1:%[0-9]+]]:_(s96) = G_IMPLICIT_DEF + CHECK: [[OP0_0:%[0-9]+]]:_(s32), [[OP0_1:%[0-9]+]]:_(s32), [[OP0_2:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[IMP_DEF0]] + CHECK: [[OP1_0:%[0-9]+]]:_(s32), [[OP1_1:%[0-9]+]]:_(s32), [[OP1_2:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[IMP_DEF1]] + CHECK: [[USUBO0:%[0-9]+]]:_(s32), [[CARRY0:%[0-9]+]]:_(s1) = G_USUBO [[OP0_0]]:_, [[OP1_0]]:_ + CHECK: [[USUBO1:%[0-9]+]]:_(s32), [[CARRY1:%[0-9]+]]:_(s1) = G_USUBE [[OP0_1]]:_, [[OP1_1]]:_, [[CARRY0]]:_ + CHECK: [[USUBO2:%[0-9]+]]:_(s32), [[CARRY2:%[0-9]+]]:_(s1) = G_USUBE [[OP0_2]]:_, [[OP1_2]]:_, [[CARRY1]]:_ + CHECK: [[USUBO:%[0-9]+]]:_(s96) = G_MERGE_VALUES [[USUBO0]]:_(s32), [[USUBO1]]:_(s32), [[USUBO2]]:_(s32) + )"; + + EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF; +} + TEST_F(AArch64GISelMITest, FewerElementsAnd) { setUp(); if (!TM)