[RISCV] Match LLVM register bank mappings

3 files changed, 266 insertions, 28 deletions

diff --git a/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp b/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
index 3d5ea2a..a7160b4 100644
--- a/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
+++ b/backend/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
@@ -139,7 +139,7 @@ public:
   [[nodiscard]] Expected<void>
   constrainSelectedInstruction(MachineInstr &Instruction, MachineFunction &Function,
                                const TargetDescription &Target) const override {
-    return RegisterBanks.constrainInstruction(Instruction, Function, Target);
+    return RegisterBanks.constrainSelectedInstruction(Instruction, Function, Target);
   }
 
   [[nodiscard]] Expected<void> canSelect(const MachineInstr &Instruction,
diff --git a/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.cpp b/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.cpp
index 6e38d52..7c60219 100644
--- a/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.cpp
+++ b/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.cpp
@@ -2,13 +2,18 @@
 
 #include "RISCVRegisterBankInfo.h"
 
+#include "rcc/CodeGen/GenericOpcodes.h"
 #include "rcc/CodeGen/MachineFunction.h"
 #include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/MachineRegisterInfo.h"
 #include "rcc/CodeGen/TargetDescription.h"
 #include "rcc/Target/RISCV/RISCVRegisters.h"
 
 #include <cstdint>
+#include <optional>
 #include <ranges>
+#include <string_view>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -25,6 +30,181 @@ namespace {
           Type.scalarBitWidth() <= 64);
 }
 
+[[nodiscard]] constexpr RegisterBankID gprBank() noexcept {
+  return RegisterBankID::fromValue(static_cast<std::uint32_t>(RegisterBank::GPR));
+}
+
+[[nodiscard]] constexpr RegisterBankID fprBank() noexcept {
+  return RegisterBankID::fromValue(static_cast<std::uint32_t>(RegisterBank::FPR));
+}
+
+[[nodiscard]] std::optional<std::size_t> operandIndex(const MachineInstr &Instruction,
+                                                      const MachineOperand &Operand) noexcept {
+  for (std::size_t Index = 0; Index < Instruction.operandCount(); ++Index) {
+    if (Instruction.operandAt(Index) == &Operand)
+      return Index;
+  }
+  return std::nullopt;
+}
+
+[[nodiscard]] std::optional<RegisterBankID>
+physicalRegisterBank(Register Reg, const TargetDescription &Target) noexcept {
+  const std::optional<PhysicalRegister> Physical = Reg.physicalRegister();
+  if (!Physical)
+    return std::nullopt;
+  for (const RegisterBankDesc &Bank : Target.registers().registerBanks()) {
+    if (Target.registers().bankContainsRegister(Bank.ID, *Physical))
+      return Bank.ID;
+  }
+  return std::nullopt;
+}
+
+[[nodiscard]] bool isFloatingGenericOpcode(GenericOpcode OpcodeValue) noexcept {
+  switch (OpcodeValue) {
+  case GenericOpcode::G_FRAME_INDEX:
+  case GenericOpcode::G_FREEZE:
+  case GenericOpcode::G_FENCE:
+  case GenericOpcode::G_FSHL:
+  case GenericOpcode::G_FSHR:
+    return false;
+  default:
+    return getGenericOpcodeName(OpcodeValue).starts_with("G_F");
+  }
+}
+
+[[nodiscard]] std::optional<RegisterBankID>
+requiredBankForOperand(const MachineInstr &Instruction, const MachineOperand &Operand,
+                       const TargetDescription &Target) {
+  const std::optional<std::size_t> Index = operandIndex(Instruction, Operand);
+  if (!Index)
+    return std::nullopt;
+
+  if (const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction.opcode());
+      Descriptor != nullptr && *Index < Descriptor->Operands.size() &&
+      Descriptor->Operands[*Index].RegBank) {
+    return Descriptor->Operands[*Index].RegBank;
+  }
+
+  if (Instruction.opcode() == opcode(GenericOpcode::COPY)) {
+    if (!Operand.isDef())
+      return std::nullopt;
+    for (const MachineOperand &Other : Instruction.operands()) {
+      if (&Other == &Operand)
+        continue;
+      const Register *OtherReg = Other.getIf<Register>();
+      if (OtherReg != nullptr && Other.isUse() && OtherReg->isPhysical())
+        return physicalRegisterBank(*OtherReg, Target);
+    }
+    return std::nullopt;
+  }
+
+  if (!isGenericOpcode(Instruction.opcode()))
+    return std::nullopt;
+  const GenericOpcode Generic = static_cast<GenericOpcode>(Instruction.opcode().value());
+  switch (Generic) {
+  case GenericOpcode::G_FPEXTLOAD:
+  case GenericOpcode::G_FPTRUNCSTORE:
+    return *Index == 0 ? fprBank() : gprBank();
+  case GenericOpcode::G_FPTOSI:
+  case GenericOpcode::G_FPTOUI:
+  case GenericOpcode::G_FPTOSI_SAT:
+  case GenericOpcode::G_FPTOUI_SAT:
+  case GenericOpcode::G_LROUND:
+  case GenericOpcode::G_LLROUND:
+    return *Index == 0 ? gprBank() : fprBank();
+  case GenericOpcode::G_SITOFP:
+  case GenericOpcode::G_UITOFP:
+    return *Index == 0 ? fprBank() : gprBank();
+  case GenericOpcode::G_FCMP:
+  case GenericOpcode::G_IS_FPCLASS:
+    return *Index == 0 ? gprBank() : fprBank();
+  case GenericOpcode::G_FPOWI:
+  case GenericOpcode::G_FLDEXP:
+    return *Index < 2 ? fprBank() : gprBank();
+  case GenericOpcode::G_FFREXP:
+    return *Index == 1 ? gprBank() : fprBank();
+  case GenericOpcode::G_LOAD:
+  case GenericOpcode::G_SEXTLOAD:
+  case GenericOpcode::G_ZEXTLOAD:
+    return *Index == 1 ? std::optional<RegisterBankID>(gprBank()) : std::nullopt;
+  case GenericOpcode::G_STORE:
+    return *Index == 1 ? std::optional<RegisterBankID>(gprBank()) : std::nullopt;
+  case GenericOpcode::G_SELECT:
+    return *Index == 1 ? std::optional<RegisterBankID>(gprBank()) : std::nullopt;
+  case GenericOpcode::G_IMPLICIT_DEF:
+  case GenericOpcode::G_PHI:
+    return std::nullopt;
+  default:
+    return isFloatingGenericOpcode(Generic) ? fprBank() : gprBank();
+  }
+}
+
+void appendConnectedRegisters(const MachineInstr &Instruction, const MachineOperand &Operand,
+                              std::vector<Register> &Pending) {
+  const std::optional<std::size_t> Index = operandIndex(Instruction, Operand);
+  if (!Index)
+    return;
+  const bool IsPhi = Instruction.opcode() == opcode(GenericOpcode::G_PHI);
+  if (!IsPhi)
+    return;
+  for (std::size_t OtherIndex = 0; OtherIndex < Instruction.operandCount(); ++OtherIndex) {
+    if (OtherIndex == *Index)
+      continue;
+    const MachineOperand *Other = Instruction.operandAt(OtherIndex);
+    const Register *OtherReg = Other == nullptr ? nullptr : Other->getIf<Register>();
+    if (OtherReg != nullptr && OtherReg->isVirtual())
+      Pending.push_back(*OtherReg);
+  }
+}
+
+[[nodiscard]] Expected<std::optional<RegisterBankID>>
+requiredBankForRegister(const MachineFunction &Function, Register Reg,
+                        const TargetDescription &Target) {
+  std::vector<Register> Pending{Reg};
+  std::unordered_set<Register> Visited;
+  std::optional<RegisterBankID> Required;
+  while (!Pending.empty()) {
+    const Register Current = Pending.back();
+    Pending.pop_back();
+    if (!Visited.insert(Current).second)
+      continue;
+    for (const MachineOperand &Operand : Function.registerInfo().regOperands(Current)) {
+      const MachineInstr *Instruction = Operand.parent();
+      if (Instruction == nullptr)
+        continue;
+      const std::optional<RegisterBankID> OperandBank =
+          requiredBankForOperand(*Instruction, Operand, Target);
+      if (OperandBank) {
+        if (Required && *Required != *OperandBank) {
+          return makeUnexpected(ErrorCode::VerificationError,
+                                "RISC-V virtual register has conflicting register-bank uses");
+        }
+        Required = *OperandBank;
+      }
+      appendConnectedRegisters(*Instruction, Operand, Pending);
+    }
+  }
+  return Required;
+}
+
+[[nodiscard]] bool hasFloatingCopyUse(const MachineFunction &Function, Register Reg,
+                                      const TargetDescription &Target) {
+  for (const MachineOperand &Operand : Function.registerInfo().regOperands(Reg)) {
+    const MachineInstr *Instruction = Operand.parent();
+    if (!Operand.isUse() || Instruction == nullptr ||
+        Instruction->opcode() != opcode(GenericOpcode::COPY))
+      continue;
+    for (const MachineOperand &Other : Instruction->operands()) {
+      const Register *OtherReg = Other.getIf<Register>();
+      if (!Other.isDef() || OtherReg == nullptr || !OtherReg->isPhysical())
+        continue;
+      if (physicalRegisterBank(*OtherReg, Target) == fprBank())
+        return true;
+    }
+  }
+  return false;
+}
+
 [[nodiscard]] Expected<RegisterClassID> requiredRegisterClass(const MachineFunction &Function,
                                                               Register Reg,
                                                               const TargetDescription &Target) {
@@ -41,12 +221,18 @@ namespace {
   const VirtualRegisterInfo &RegisterInfo = Info->get();
   const LLT Type = RegisterInfo.Type.value_or(LLT{});
   RegisterClassID RequiredClass;
-  if (isScalarGPRType(Type)) {
+  if ((RegisterInfo.Bank && *RegisterInfo.Bank == fprBank()) || Type.isFloat()) {
+    if (Type.scalarBitWidth() != 32 && Type.scalarBitWidth() != 64) {
+      return makeUnexpected(ErrorCode::Unsupported,
+                            "RISC-V RV64GC has no floating register class for type '" +
+                                Type.toString() + "'");
+    }
+    RequiredClass =
+        Type.scalarBitWidth() == 32
+            ? RegisterClassID::fromValue(static_cast<std::uint32_t>(RegisterClass::FPR32))
+            : RegisterClassID::fromValue(static_cast<std::uint32_t>(RegisterClass::FPR64));
+  } else if (isScalarGPRType(Type)) {
     RequiredClass = RegisterClassID::fromValue(static_cast<std::uint32_t>(RegisterClass::GPR));
-  } else if (Type.isFloat() && Type.scalarBitWidth() == 32) {
-    RequiredClass = RegisterClassID::fromValue(static_cast<std::uint32_t>(RegisterClass::FPR32));
-  } else if (Type.isFloat() && Type.scalarBitWidth() == 64) {
-    RequiredClass = RegisterClassID::fromValue(static_cast<std::uint32_t>(RegisterClass::FPR64));
   } else {
     return makeUnexpected(ErrorCode::Unsupported,
                           "RISC-V RV64GC has no register class for type '" + Type.toString() + "'");
@@ -70,26 +256,76 @@ namespace {
 
 } // namespace
 
-Expected<RegisterBankID> RISCVRegisterBankInfo::selectBank(const MachineFunction &, Register,
-                                                           const VirtualRegisterInfo &Info,
-                                                           const MachineInstr *,
-                                                           const TargetDescription &) const {
+Expected<RegisterBankSelection>
+RISCVRegisterBankInfo::selectBank(const MachineFunction &Function, Register Reg,
+                                  const VirtualRegisterInfo &Info, const MachineInstr *,
+                                  const TargetDescription &Target) const {
   if (!Info.Type) {
     return makeUnexpected(ErrorCode::VerificationError,
                           "RISC-V register-bank selection requires a virtual-register type");
   }
-  if (isScalarGPRType(*Info.Type))
-    return RegisterBankID::fromValue(static_cast<std::uint32_t>(RegisterBank::GPR));
-  if (isRVFloatType(*Info.Type))
-    return RegisterBankID::fromValue(static_cast<std::uint32_t>(RegisterBank::FPR));
-  return makeUnexpected(ErrorCode::Unsupported,
-                        "RISC-V RV64GC has no register bank for virtual register type '" +
-                            Info.Type->toString() + "'");
+
+  auto Required = requiredBankForRegister(Function, Reg, Target);
+  if (!Required)
+    return std::unexpected(Required.error());
+  const bool IsConstrained = Required->has_value() || Info.Bank.has_value();
+  const RegisterBankID Selected =
+      Required->value_or(isRVFloatType(*Info.Type) ? fprBank() : gprBank());
+  if (Info.Bank && *Info.Bank != Selected) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V virtual register has a conflicting register bank");
+  }
+  return RegisterBankSelection{.Bank = Selected, .IsConstrained = IsConstrained};
+}
+
+Expected<RegisterBankInstructionMapping>
+RISCVRegisterBankInfo::getInstructionMapping(const MachineInstr &Instruction,
+                                             const MachineFunction &Function,
+                                             const TargetDescription &Target) const {
+  if (Instruction.opcode() != opcode(GenericOpcode::G_SELECT))
+    return RegisterBankInstructionMapping{};
+  if (Instruction.operandCount() != 4) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V G_SELECT register-bank mapping requires four operands");
+  }
+
+  RegisterBankInstructionMapping Mapping(Instruction.operandCount());
+  Mapping[1] = gprBank();
+  unsigned FloatingConstraints = 0;
+
+  const Register *Result = Instruction.operandAt(0)->getIf<Register>();
+  if (Result == nullptr || !Result->isVirtual()) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V G_SELECT result must be a virtual register");
+  }
+  auto ResultRequirement = requiredBankForRegister(Function, *Result, Target);
+  if (!ResultRequirement)
+    return std::unexpected(ResultRequirement.error());
+  if (*ResultRequirement == fprBank() || hasFloatingCopyUse(Function, *Result, Target))
+    ++FloatingConstraints;
+
+  for (const std::size_t Index : {std::size_t{2}, std::size_t{3}}) {
+    const Register *Source = Instruction.operandAt(Index)->getIf<Register>();
+    if (Source == nullptr || !Source->isVirtual()) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V G_SELECT sources must be virtual registers");
+    }
+    auto SourceInfo = Function.registerInfo().virtualRegisterInfo(*Source);
+    if (!SourceInfo)
+      return std::unexpected(SourceInfo.error());
+    if (SourceInfo->get().Bank == fprBank())
+      ++FloatingConstraints;
+  }
+
+  const RegisterBankID ValueBank = FloatingConstraints >= 2 ? fprBank() : gprBank();
+  Mapping[0] = ValueBank;
+  Mapping[2] = ValueBank;
+  Mapping[3] = ValueBank;
+  return Mapping;
 }
 
-Expected<void> RISCVRegisterBankInfo::constrainInstruction(MachineInstr &Instruction,
-                                                           MachineFunction &Function,
-                                                           const TargetDescription &Target) const {
+Expected<void> RISCVRegisterBankInfo::constrainSelectedInstruction(
+    MachineInstr &Instruction, MachineFunction &Function, const TargetDescription &Target) const {
   std::vector<std::pair<Register, RegisterClassID>> PlannedClasses;
   for (const MachineOperand &Operand : Instruction.operands()) {
     const Register *Reg = Operand.getIf<Register>();
@@ -128,7 +364,7 @@ Expected<void> RISCVRegisterBankInfo::constrainRegister(MachineFunction &Functio
   if (!RequiredClass)
     return std::unexpected(RequiredClass.error());
   if (!Info->get().Bank) {
-    auto SetBank = Function.registerInfo().setRegisterBank(Reg, *Bank);
+    auto SetBank = Function.registerInfo().setRegisterBank(Reg, Bank->Bank);
     if (!SetBank)
       return std::unexpected(SetBank.error());
   }
diff --git a/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.h b/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.h
index 543bc4d..69eb539 100644
--- a/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.h
+++ b/backend/lib/Target/RISCV/GISel/RISCVRegisterBankInfo.h
@@ -9,13 +9,15 @@ namespace rcc::riscv {
 
 class RISCVRegisterBankInfo final : public RegisterBankInfo {
 public:
-  [[nodiscard]] Expected<RegisterBankID> selectBank(const MachineFunction &Function, Register Reg,
-                                                    const VirtualRegisterInfo &Info,
-                                                    const MachineInstr *Instruction,
-                                                    const TargetDescription &Target) const override;
-  [[nodiscard]] Expected<void> constrainInstruction(MachineInstr &Instruction,
-                                                    MachineFunction &Function,
-                                                    const TargetDescription &Target) const override;
+  [[nodiscard]] Expected<RegisterBankSelection>
+  selectBank(const MachineFunction &Function, Register Reg, const VirtualRegisterInfo &Info,
+             const MachineInstr *Instruction, const TargetDescription &Target) const override;
+  [[nodiscard]] Expected<RegisterBankInstructionMapping>
+  getInstructionMapping(const MachineInstr &Instruction, const MachineFunction &Function,
+                        const TargetDescription &Target) const override;
+  [[nodiscard]] Expected<void> constrainSelectedInstruction(MachineInstr &Instruction,
+                                                            MachineFunction &Function,
+                                                            const TargetDescription &Target) const;
   [[nodiscard]] Expected<void> constrainRegister(MachineFunction &Function, Register Reg,
                                                  const TargetDescription &Target) const;
   [[nodiscard]] Expected<void> validateInstruction(const MachineInstr &Instruction,