[RISCV] Legalize integer compares to RV64 containers

1 files changed, 151 insertions, 4 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 &&