[GlobalISel] Implement narrowScalar for UADDO/USUBO

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

4488

I don't like using Optional for Register, a default constructed Register behaves like none

[porglezomp]

Change Optional<Register> to Register

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

4485

Might as well hard code the 2 here, treating it as unknown is slightly more confusing

4489

This must be exactly ==, > adds ambiguity

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

898

Should have a dedicated buildUAddo

935

We should have a dedicated buildUSubo

[paquette]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

4485

I think this is supposed to be shared between the overflow ops/non-overflow ops, so it makes sense to not hardcode a 2.

I think a comment saying that this should be 1 or 2 would be handy though.

[porglezomp]

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

898

I wrote it this way to match the style of the surrounding tests. Which would we prefer here, better style, or more consistency? I'm already changing some other things in favor of better style in the new tests, so I wouldn't be opposed.

[arsenm]

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

898

Raw buildInstr is always uglier and only appropriate if you don't know the exact opcode you're using

[porglezomp]

Add comment about number of defs, change buildInstr

[porglezomp]

Change some ASSERT_TRUE to ASSERT_EQ

[arsenm]

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

903–904

It's a gtest-ism that these should be swapped to get the correct error message as EXPECT_EQ(expected, actual)

[porglezomp]

Flip argument order of EXPECT_EQ for better failure messages

[porglezomp]

ping on this one, since it's the first in the stack and blocking the others

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

case TargetOpcode::G_TRUNC: {
     MI.eraseFromParent();
     return Legalized;
   }
 
   case TargetOpcode::G_FREEZE:
     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:
     return narrowScalarMul(MI, NarrowTy);
   case TargetOpcode::G_EXTRACT:
     return narrowScalarExtract(MI, TypeIdx, NarrowTy);
   case TargetOpcode::G_INSERT:
     return narrowScalarInsert(MI, TypeIdx, NarrowTy);
LegalizerHelper::narrowScalarAddSub(MachineInstr &MI, unsigned TypeIdx,
   if (SizeOp0 % NarrowSize != 0)
     return UnableToLegalize;
 
   // Expand in terms of carry-setting/consuming G_<Op>E instructions.
   int NumParts = SizeOp0 / NarrowTy.getSizeInBits();
 
   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;
     break;
   default:
     llvm_unreachable("Unexpected add/sub opcode!");
   }
 
+  unsigned NumDefs = MI.getNumExplicitDefs();
+  Register Src1 = MI.getOperand(NumDefs).getReg();
+  Register Src2 = MI.getOperand(NumDefs + 1).getReg();
+  Optional<Register> CarryDst = None;
+  if (NumDefs >= 2)
+    CarryDst = MI.getOperand(1).getReg();
+
   SmallVector<Register, 2> 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},
                             {Src1Regs[i], Src2Regs[i]});
     else {
       MIRBuilder.buildInstr(OpE, {DstReg, CarryOut},
                             {Src1Regs[i], Src2Regs[i], CarryIn});
     }

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

TEST_F(AArch64GISelMITest, WidenSSUBE) {
   CHECK: G_ICMP intpred(ne), [[SSUBE]]:_(s16), [[SEXT]]:_
   CHECK: G_TRUNC [[SSUBE]]
   )";
 
   // Check
   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.buildInstr(TargetOpcode::G_UADDO, {S96, S1}, {Op0, Op1});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+  EXPECT_TRUE(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.buildInstr(TargetOpcode::G_USUBO, {S96, S1}, {Op0, Op1});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+  EXPECT_TRUE(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)
     return;
 
   const LLT V2S32 = LLT::vector(2, 32);
   const LLT V5S32 = LLT::vector(5, 32);