[RISCV] Legalize integer compares to RV64 containers

2 files changed, 204 insertions, 20 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index bc437c0..b856476 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -683,6 +683,147 @@ lowerAtomicCompareExchangeWithSuccess(MachineInstr &Instruction, MachineFunction
   return rewriteUnary(Instruction, GenericOpcode::COPY, *OldValue, *AtomicResult);
 }
 
+[[nodiscard]] Expected<Register> widenIntegerCompareSource(MachineInstr &Instruction,
+                                                           MachineFunction &Function,
+                                                           MachineIRBuilder &Builder,
+                                                           Register Source, LLT WideType) {
+  const auto SourceInfo = Function.registerInfo().virtualRegisterInfo(Source);
+  if (!SourceInfo || !SourceInfo->get().Type)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V integer compare source has no scalar type");
+  const LLT SourceType = *SourceInfo->get().Type;
+  if (SourceType.scalarBitWidth() == WideType.scalarBitWidth())
+    return Source;
+
+  const MachineOperand *Definition = Function.registerInfo().getUniqueVRegDef(Source);
+  MachineInstr *Def =
+      Definition == nullptr ? nullptr : const_cast<MachineInstr *>(Definition->parent());
+  if (Def != nullptr && Def->opcode() == opcode(GenericOpcode::G_CONSTANT) &&
+      singleRegisterUse(Function, Source) == &Instruction) {
+    auto Widened = widenConstantToRV64(*Def, Function);
+    if (!Widened)
+      return std::unexpected(Widened.error());
+    return Source;
+  }
+
+  if (Def != nullptr && Def->opcode() == opcode(GenericOpcode::G_TRUNC) &&
+      Def->operandCount() == 2) {
+    const MachineOperand *WideOperand = Def->operandAt(1);
+    const Register *WideSource = WideOperand == nullptr ? nullptr : WideOperand->getIf<Register>();
+    const std::optional<LLT> ExistingWideType =
+        WideOperand == nullptr ? std::nullopt : registerType(Function, *WideOperand);
+    if (WideSource != nullptr && ExistingWideType && *ExistingWideType == WideType) {
+      auto Result = Function.registerInfo().createVirtualRegister(WideType);
+      if (!Result)
+        return std::unexpected(Result.error());
+      auto Point = Builder.setInsertPointBefore(Instruction);
+      if (!Point)
+        return std::unexpected(Point.error());
+      auto Pending = Builder.buildInstr(opcode(GenericOpcode::G_SEXT_INREG));
+      if (!Pending)
+        return std::unexpected(Pending.error());
+      if (auto Added = Pending->addDef(*Result); !Added)
+        return std::unexpected(Added.error());
+      if (auto Added = Pending->addUse(*WideSource); !Added)
+        return std::unexpected(Added.error());
+      if (auto Added =
+              Pending->addImmediate(static_cast<std::int64_t>(SourceType.scalarBitWidth()));
+          !Added) {
+        return std::unexpected(Added.error());
+      }
+      auto Built = Pending->build();
+      if (!Built)
+        return std::unexpected(Built.error());
+      return *Result;
+    }
+  }
+  return buildUnaryBefore(Builder, Instruction, GenericOpcode::G_SEXT, WideType, Source);
+}
+
+[[nodiscard]] Expected<void> lowerIntegerCompare(MachineInstr &Instruction,
+                                                 MachineFunction &Function,
+                                                 const TargetDescription &Target) {
+  const MachineOperand *DestinationOperand = Instruction.operandAt(0);
+  const MachineOperand *LeftOperand = Instruction.operandAt(2);
+  const MachineOperand *RightOperand = Instruction.operandAt(3);
+  const Register *Destination =
+      DestinationOperand == nullptr ? nullptr : DestinationOperand->getIf<Register>();
+  const Register *Left = LeftOperand == nullptr ? nullptr : LeftOperand->getIf<Register>();
+  const Register *Right = RightOperand == nullptr ? nullptr : RightOperand->getIf<Register>();
+  if (Destination == nullptr || Left == nullptr || Right == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V integer compare has invalid register operands");
+  const Register OriginalDestination = *Destination;
+  const Register OriginalLeft = *Left;
+  const Register OriginalRight = *Right;
+
+  auto WideType = LLT::makeAnyScalar(64);
+  if (!WideType)
+    return std::unexpected(WideType.error());
+  MachineIRBuilder Builder(Function, Target);
+  auto WideLeft =
+      widenIntegerCompareSource(Instruction, Function, Builder, OriginalLeft, *WideType);
+  if (!WideLeft)
+    return std::unexpected(WideLeft.error());
+  auto WideRight =
+      widenIntegerCompareSource(Instruction, Function, Builder, OriginalRight, *WideType);
+  if (!WideRight)
+    return std::unexpected(WideRight.error());
+  if (auto Replaced = Instruction.operandAt(2)->setRegister(*WideLeft); !Replaced)
+    return std::unexpected(Replaced.error());
+  if (auto Replaced = Instruction.operandAt(3)->setRegister(*WideRight); !Replaced)
+    return std::unexpected(Replaced.error());
+
+  const auto DestinationInfo = Function.registerInfo().virtualRegisterInfo(OriginalDestination);
+  if (!DestinationInfo || !DestinationInfo->get().Type)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V integer compare destination has no scalar type");
+  if (DestinationInfo->get().Type->scalarBitWidth() == 64)
+    return {};
+
+  MachineInstr *FoldedExtension = singleRegisterUse(Function, OriginalDestination);
+  if (FoldedExtension != nullptr &&
+      (FoldedExtension->opcode() == opcode(GenericOpcode::G_ANYEXT) ||
+       FoldedExtension->opcode() == opcode(GenericOpcode::G_SEXT) ||
+       FoldedExtension->opcode() == opcode(GenericOpcode::G_ZEXT)) &&
+      operandType(Function, *FoldedExtension, 0)
+          .transform([](LLT Type) { return Type.scalarBitWidth() == 64; })
+          .value_or(false)) {
+    const MachineOperand *ExtendedOperand = FoldedExtension->operandAt(0);
+    const Register *Extended =
+        ExtendedOperand == nullptr ? nullptr : ExtendedOperand->getIf<Register>();
+    if (Extended == nullptr)
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V integer compare extension has no destination");
+    const Register WideDestination = *Extended;
+    if (auto Replaced = Instruction.operandAt(0)->setRegister(WideDestination); !Replaced)
+      return std::unexpected(Replaced.error());
+    return eraseInstruction(*FoldedExtension);
+  }
+
+  auto WideDestination = Function.registerInfo().createVirtualRegister(*WideType);
+  if (!WideDestination)
+    return std::unexpected(WideDestination.error());
+  if (auto Replaced = Instruction.operandAt(0)->setRegister(*WideDestination); !Replaced)
+    return std::unexpected(Replaced.error());
+  if (Function.registerInfo().useEmpty(OriginalDestination))
+    return {};
+  auto Point = Builder.setInsertPointAfter(Instruction);
+  if (!Point)
+    return std::unexpected(Point.error());
+  auto Truncation = Builder.buildInstr(opcode(GenericOpcode::G_TRUNC));
+  if (!Truncation)
+    return std::unexpected(Truncation.error());
+  if (auto Added = Truncation->addDef(OriginalDestination); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = Truncation->addUse(*WideDestination); !Added)
+    return std::unexpected(Added.error());
+  auto Built = Truncation->build();
+  if (!Built)
+    return std::unexpected(Built.error());
+  return {};
+}
+
 [[nodiscard]] Expected<void> foldNarrowAnyExtBinary(MachineInstr &Extension,
                                                      MachineFunction &Function,
                                                      const TargetDescription &Target) {
@@ -2953,11 +3094,14 @@ public:
           Predicate == nullptr ? nullptr : Predicate->getIf<PredicateOperand>();
       const std::optional<LLT> Left = operandType(Function, Instruction, 2);
       const std::optional<LLT> Right = operandType(Function, Instruction, 3);
-      return Type && isScalarGPRType(*Type) && IntegerPredicate != nullptr &&
-                     IntegerPredicate->integerPredicate() != nullptr && Left && Right &&
-                     *Left == *Right && isSupportedIntegerType(*Left)
+      if (!Type || !isSupportedIntegerScalar(*Type) || IntegerPredicate == nullptr ||
+          IntegerPredicate->integerPredicate() == nullptr || !Left || !Right || *Left != *Right ||
+          !isSupportedIntegerType(*Left)) {
+        return LegalizationAction::Unsupported;
+      }
+      return Type->scalarBitWidth() == 64 && Left->scalarBitWidth() == 64
                  ? LegalizationAction::Legal
-                 : LegalizationAction::Unsupported;
+                 : LegalizationAction::Custom;
     }
     case GenericOpcode::G_FCMP: {
       const MachineOperand *Predicate = Instruction.operandAt(1);
@@ -3302,6 +3446,9 @@ public:
         Instruction.opcode() == opcode(GenericOpcode::G_PHI))
       return lowerNarrowPhi(Instruction, Function, Target);
     if (Action == LegalizationAction::Custom &&
+        Instruction.opcode() == opcode(GenericOpcode::G_ICMP))
+      return lowerIntegerCompare(Instruction, Function, Target);
+    if (Action == LegalizationAction::Custom &&
         Instruction.opcode() == opcode(GenericOpcode::G_ATOMIC_CMPXCHG))
       return lowerAtomicCompareExchange(Instruction, Function, Target);
     if (Action == LegalizationAction::Custom &&
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 12e59ba..7c0ec33 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -452,10 +452,21 @@ void testNarrowIntegerSelection(TestContext &Context) {
   if (Compare && Compare->firstBlock()->firstInstruction() != nullptr) {
     rcc::MachineInstr *Instruction = Compare->firstBlock()->firstInstruction();
     auto Action = rcc::riscv::legalizerInfo().classify(*Instruction, *Compare, *Target);
-    auto Selection = rcc::riscv::instructionSelector().select(*Instruction, *Compare, *Target);
-    Context.expect(Action && *Action == rcc::LegalizationAction::Legal && Selection,
-                   "signed i7 comparisons must normalize and select");
-    Context.expect(Selection && validatesSelectedFunction(*Target, *Compare),
+    auto Lowered = Action
+                       ? rcc::riscv::legalizerInfo().lower(*Instruction, *Compare, *Target, *Action)
+                       : rcc::Expected<void>(std::unexpected(Action.error()));
+    auto LegalAction = rcc::riscv::legalizerInfo().classify(*Instruction, *Compare, *Target);
+    bool Selected = Lowered && LegalAction && *LegalAction == rcc::LegalizationAction::Legal &&
+                    constrainVirtualRegisters(*Compare);
+    const auto Instructions = Compare->allInstructions();
+    for (rcc::MachineInstr *Candidate : Instructions) {
+      if (Selected && rcc::isGenericOpcode(Candidate->opcode()))
+        Selected =
+            rcc::riscv::instructionSelector().select(*Candidate, *Compare, *Target).has_value();
+    }
+    Context.expect(Action && *Action == rcc::LegalizationAction::Custom && Selected,
+                   "signed i7 comparisons must legalize to RV64 containers and select");
+    Context.expect(Selected && validatesSelectedFunction(*Target, *Compare),
                    "selected i7 comparison must satisfy target descriptors");
   }
 
@@ -608,26 +619,52 @@ void testComparisonsAndControl(TestContext &Context) {
                                          rcc::FloatingPredicate::AlwaysFalse);
   if (Equal) {
     auto *Instruction = Equal->firstBlock()->firstInstruction();
+    const rcc::Register *NarrowResult = Instruction->operandAt(0)->getIf<rcc::Register>();
+    auto WideResult = Equal->registerInfo().createVirtualRegister(*I64);
+    rcc::MachineIRBuilder Builder(*Equal, *Target);
+    if (NarrowResult == nullptr || !WideResult || !constrainRegister(*Equal, *WideResult, *I64) ||
+        !Builder.setInsertPoint(*Equal->firstBlock()))
+      return;
+    auto Extension = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ANYEXT));
+    if (!Extension || !Extension->addDef(*WideResult) || !Extension->addUse(*NarrowResult))
+      return;
+    auto BuiltExtension = Extension->build();
+    if (!BuiltExtension)
+      return;
+    auto Action = rcc::riscv::legalizerInfo().classify(*Instruction, *Equal, *Target);
+    auto Lowered = Action
+                       ? rcc::riscv::legalizerInfo().lower(*Instruction, *Equal, *Target, *Action)
+                       : rcc::Expected<void>(std::unexpected(Action.error()));
+    const rcc::Register *LegalResult = Instruction->operandAt(0)->getIf<rcc::Register>();
+    Context.expect(Lowered && (*BuiltExtension)->parent() == nullptr && LegalResult != nullptr &&
+                       *LegalResult == *WideResult,
+                   "RV64 integer comparison legalization must fold a single result extension");
     auto Selection = rcc::riscv::instructionSelector().select(*Instruction, *Equal, *Target);
-    Context.expect(Selection.has_value(), "integer equality must select");
-    Context.expect(Equal->allInstructions().size() == 2,
+    Context.expect(Action && *Action == rcc::LegalizationAction::Custom && Lowered && Selection,
+                   "integer equality must legalize its result container and select");
+    rcc::MachineInstr *FirstSelected = Equal->firstBlock()->firstInstruction();
+    rcc::MachineInstr *SecondSelected = FirstSelected == nullptr ? nullptr : FirstSelected->next();
+    Context.expect(SecondSelected != nullptr && SecondSelected->next() == nullptr,
                    "integer equality must emit XOR and zero-test instructions");
-    if (Equal->allInstructions().size() == 2) {
-      auto It = Equal->allInstructions().begin();
-      Context.expect((*It)->opcode() == rcc::Opcode::fromValue(static_cast<std::uint32_t>(
-                                            rcc::riscv::InstructionOpcode::XOR)),
+    if (SecondSelected != nullptr && SecondSelected->next() == nullptr) {
+      Context.expect(FirstSelected->opcode() == rcc::Opcode::fromValue(static_cast<std::uint32_t>(
+                                                    rcc::riscv::InstructionOpcode::XOR)),
                      "integer equality must start with XOR");
-      ++It;
-      Context.expect((*It)->opcode() == rcc::Opcode::fromValue(static_cast<std::uint32_t>(
-                                            rcc::riscv::InstructionOpcode::SLTIU)),
+      Context.expect(SecondSelected->opcode() == rcc::Opcode::fromValue(static_cast<std::uint32_t>(
+                                                     rcc::riscv::InstructionOpcode::SLTIU)),
                      "integer equality must finish with SLTIU");
     }
   }
   if (LessEqual) {
     auto *Instruction = LessEqual->firstBlock()->firstInstruction();
-    Context.expect(
-        rcc::riscv::instructionSelector().select(*Instruction, *LessEqual, *Target).has_value(),
-        "signed less-equal must select");
+    auto Action = rcc::riscv::legalizerInfo().classify(*Instruction, *LessEqual, *Target);
+    auto Lowered =
+        Action ? rcc::riscv::legalizerInfo().lower(*Instruction, *LessEqual, *Target, *Action)
+               : rcc::Expected<void>(std::unexpected(Action.error()));
+    Context.expect(Action && *Action == rcc::LegalizationAction::Custom && Lowered &&
+                       constrainVirtualRegisters(*LessEqual) &&
+                       rcc::riscv::instructionSelector().select(*Instruction, *LessEqual, *Target),
+                   "signed less-equal must legalize its result container and select");
     Context.expect(LessEqual->allInstructions().size() == 2,
                    "signed less-equal must emit compare and inversion");
   }