[GlobalISel] Lower stack operations

5 files changed, 200 insertions, 8 deletions

diff --git a/backend/include/rcc/CodeGen/GlobalISel/LegalizerHelper.h b/backend/include/rcc/CodeGen/GlobalISel/LegalizerHelper.h
index 71ead8b..3e14311 100644
--- a/backend/include/rcc/CodeGen/GlobalISel/LegalizerHelper.h
+++ b/backend/include/rcc/CodeGen/GlobalISel/LegalizerHelper.h
@@ -39,6 +39,9 @@ public:
   [[nodiscard]] Expected<void> lowerBitCount(MachineInstr &Instruction);
   [[nodiscard]] Expected<void> lowerBswap(MachineInstr &Instruction);
   [[nodiscard]] Expected<void> lowerRotate(MachineInstr &Instruction);
+  [[nodiscard]] Expected<void> lowerDynStackAlloc(MachineInstr &Instruction);
+  [[nodiscard]] Expected<void> lowerStackSave(MachineInstr &Instruction);
+  [[nodiscard]] Expected<void> lowerStackRestore(MachineInstr &Instruction);
   [[nodiscard]] Expected<void> lowerVAArg(MachineInstr &Instruction);
   [[nodiscard]] Expected<void> libcall(MachineInstr &Instruction);
   [[nodiscard]] Expected<void> widenScalar(MachineInstr &Instruction, std::size_t TypeIndex,
diff --git a/backend/include/rcc/CodeGen/TargetDescription.h b/backend/include/rcc/CodeGen/TargetDescription.h
index 56d328e..6504ba8 100644
--- a/backend/include/rcc/CodeGen/TargetDescription.h
+++ b/backend/include/rcc/CodeGen/TargetDescription.h
@@ -36,6 +36,16 @@ struct TargetBooleanContentInfo final {
   BooleanContent Vector = BooleanContent::Undefined;
 };
 
+enum class StackGrowthDirection : std::uint8_t {
+  Down,
+  Up,
+};
+
+struct TargetStackInfo final {
+  PhysicalRegister StackPointerRegisterToSaveRestore;
+  StackGrowthDirection GrowthDirection = StackGrowthDirection::Down;
+};
+
 class TargetDescription final {
 public:
   TargetDescription(const TargetDescription &) = delete;
@@ -44,9 +54,11 @@ public:
   TargetDescription &operator=(TargetDescription &&) noexcept = default;
   ~TargetDescription() = default;
 
-  [[nodiscard]] static Expected<TargetDescription>
-  create(InstructionSet Instructions, TargetRegisterInfo Registers,
-         TargetMIRTextInfoDesc MIRText = {}, TargetBooleanContentInfo BooleanInfo = {});
+  [[nodiscard]] static Expected<TargetDescription> create(InstructionSet Instructions,
+                                                          TargetRegisterInfo Registers,
+                                                          TargetMIRTextInfoDesc MIRText = {},
+                                                          TargetBooleanContentInfo BooleanInfo = {},
+                                                          TargetStackInfo StackInfo = {});
 
   [[nodiscard]] const InstructionSet &instructions() const noexcept {
     return Instructions;
@@ -67,6 +79,14 @@ public:
     return IsFloatingPoint ? BooleanInfo.ScalarFloatingPoint : BooleanInfo.ScalarInteger;
   }
 
+  [[nodiscard]] PhysicalRegister stackPointerRegisterToSaveRestore() const noexcept {
+    return StackInfo.StackPointerRegisterToSaveRestore;
+  }
+
+  [[nodiscard]] StackGrowthDirection stackGrowthDirection() const noexcept {
+    return StackInfo.GrowthDirection;
+  }
+
   [[nodiscard]] const InstructionDescriptor *
   instructionDescriptor(Opcode InstructionOpcode) const noexcept {
     return Instructions.descriptor(InstructionOpcode);
@@ -107,9 +127,10 @@ public:
 
 private:
   TargetDescription(InstructionSet Instructions, TargetRegisterInfo Registers,
-                    TargetMIRTextInfo MIRText, TargetBooleanContentInfo BooleanInfo)
+                    TargetMIRTextInfo MIRText, TargetBooleanContentInfo BooleanInfo,
+                    TargetStackInfo StackInfo)
       : Instructions(std::move(Instructions)), Registers(std::move(Registers)),
-        MIRText(std::move(MIRText)), BooleanInfo(BooleanInfo) {}
+        MIRText(std::move(MIRText)), BooleanInfo(BooleanInfo), StackInfo(StackInfo) {}
 
   [[nodiscard]] Expected<void> validateTargetOperand(const OperandDescriptor *Descriptor,
                                                      const MachineOperand &Operand,
@@ -122,6 +143,7 @@ private:
   TargetRegisterInfo Registers;
   TargetMIRTextInfo MIRText;
   TargetBooleanContentInfo BooleanInfo;
+  TargetStackInfo StackInfo;
 };
 
 } // namespace rcc
diff --git a/backend/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/backend/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index 33ec1ce..b169c55 100644
--- a/backend/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/backend/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3345,6 +3345,12 @@ Expected<void> LegalizerHelper::lower(MachineInstr &Instruction,
   case GenericOpcode::G_ROTL:
   case GenericOpcode::G_ROTR:
     return lowerRotate(Instruction);
+  case GenericOpcode::G_DYN_STACKALLOC:
+    return lowerDynStackAlloc(Instruction);
+  case GenericOpcode::G_STACKSAVE:
+    return lowerStackSave(Instruction);
+  case GenericOpcode::G_STACKRESTORE:
+    return lowerStackRestore(Instruction);
   case GenericOpcode::G_ABDS:
   case GenericOpcode::G_ABDU: {
     const bool IsSigned = Instruction.opcode() == opcode(GenericOpcode::G_ABDS);
@@ -3649,6 +3655,153 @@ Expected<void> LegalizerHelper::lower(MachineInstr &Instruction,
   }
 }
 
+Expected<void> LegalizerHelper::lowerDynStackAlloc(MachineInstr &Instruction) {
+  if (Instruction.operandCount() != 3)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "G_DYN_STACKALLOC requires destination, size, and alignment operands");
+  if (MIRBuilder.target().stackGrowthDirection() != StackGrowthDirection::Down)
+    return makeUnexpected(ErrorCode::Unsupported,
+                          "dynamic stack allocation requires a downward-growing stack");
+
+  const PhysicalRegister StackPointer = MIRBuilder.target().stackPointerRegisterToSaveRestore();
+  if (!StackPointer.isValid())
+    return makeUnexpected(ErrorCode::Unsupported,
+                          "target does not define a stack pointer for dynamic allocation");
+
+  const MachineOperand *DestinationOperand = Instruction.operandAt(0);
+  const MachineOperand *SizeOperand = Instruction.operandAt(1);
+  const MachineOperand *AlignmentOperand = Instruction.operandAt(2);
+  const Register *Destination =
+      DestinationOperand == nullptr ? nullptr : DestinationOperand->getIf<Register>();
+  const Register *AllocationSize =
+      SizeOperand == nullptr ? nullptr : SizeOperand->getIf<Register>();
+  if (Destination == nullptr || AllocationSize == nullptr || AlignmentOperand == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError, "G_DYN_STACKALLOC has invalid operands");
+
+  auto PointerType = MRI.getType(*Destination);
+  auto AllocationSizeType = MRI.getType(*AllocationSize);
+  auto AlignmentValue = signedImmediate(*AlignmentOperand);
+  if (!PointerType)
+    return std::unexpected(PointerType.error());
+  if (!AllocationSizeType)
+    return std::unexpected(AllocationSizeType.error());
+  if (!AlignmentValue)
+    return std::unexpected(AlignmentValue.error());
+  if (!PointerType->isPointer() || !AllocationSizeType->isCategoryNeutralScalar() ||
+      PointerType->scalarBitWidth() != AllocationSizeType->scalarBitWidth()) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "G_DYN_STACKALLOC requires pointer-width integer allocation size");
+  }
+  if (*AlignmentValue <= 0 || !std::has_single_bit(static_cast<std::uint64_t>(*AlignmentValue))) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "G_DYN_STACKALLOC alignment must be a positive power of two");
+  }
+
+  auto StackPointerCopy = MRI.createVirtualRegister(*PointerType);
+  auto IntegerStackPointer = MRI.createVirtualRegister(*AllocationSizeType);
+  auto Allocation = MRI.createVirtualRegister(*AllocationSizeType);
+  if (!StackPointerCopy)
+    return std::unexpected(StackPointerCopy.error());
+  if (!IntegerStackPointer)
+    return std::unexpected(IntegerStackPointer.error());
+  if (!Allocation)
+    return std::unexpected(Allocation.error());
+
+  if (auto Point = MIRBuilder.setInsertPointBefore(Instruction); !Point)
+    return std::unexpected(Point.error());
+  if (auto Built = MIRBuilder.buildCopy(*StackPointerCopy, Register::fromPhysical(StackPointer));
+      !Built) {
+    return std::unexpected(Built.error());
+  }
+  if (auto Built = MIRBuilder.buildPtrToInt(*IntegerStackPointer, *StackPointerCopy); !Built)
+    return std::unexpected(Built.error());
+  if (auto Built = MIRBuilder.buildSub(*Allocation, *IntegerStackPointer, *AllocationSize); !Built)
+    return std::unexpected(Built.error());
+
+  Register AlignedAllocation = *Allocation;
+  if (*AlignmentValue > 1) {
+    auto LowBits = APInt::fromUnsigned(AllocationSizeType->scalarBitWidth(),
+                                       static_cast<std::uint64_t>(*AlignmentValue) - 1U);
+    if (!LowBits)
+      return std::unexpected(LowBits.error());
+    auto AlignmentMask = LowBits->bitwiseNot();
+    if (!AlignmentMask)
+      return std::unexpected(AlignmentMask.error());
+    auto Mask = buildConstantRegister(MIRBuilder, *AllocationSizeType, std::move(*AlignmentMask));
+    if (!Mask)
+      return std::unexpected(Mask.error());
+    auto MaskedAllocation = MRI.createVirtualRegister(*AllocationSizeType);
+    if (!MaskedAllocation)
+      return std::unexpected(MaskedAllocation.error());
+    if (auto Built = MIRBuilder.buildAnd(*MaskedAllocation, *Allocation, *Mask); !Built)
+      return std::unexpected(Built.error());
+    AlignedAllocation = *MaskedAllocation;
+  }
+
+  auto TargetPointer = MRI.createVirtualRegister(*PointerType);
+  if (!TargetPointer)
+    return std::unexpected(TargetPointer.error());
+  if (auto Built = MIRBuilder.buildIntToPtr(*TargetPointer, AlignedAllocation); !Built)
+    return std::unexpected(Built.error());
+  if (auto Built = MIRBuilder.buildCopy(Register::fromPhysical(StackPointer), *TargetPointer);
+      !Built) {
+    return std::unexpected(Built.error());
+  }
+  if (auto Built = MIRBuilder.buildCopy(*Destination, *TargetPointer); !Built)
+    return std::unexpected(Built.error());
+
+  MachineBasicBlock *Block = Instruction.parent();
+  if (Block == nullptr)
+    return makeUnexpected(ErrorCode::InternalError,
+                          "G_DYN_STACKALLOC became detached during lowering");
+  return Block->removeInstruction(Instruction);
+}
+
+Expected<void> LegalizerHelper::lowerStackSave(MachineInstr &Instruction) {
+  const PhysicalRegister StackPointer = MIRBuilder.target().stackPointerRegisterToSaveRestore();
+  if (!StackPointer.isValid())
+    return makeUnexpected(ErrorCode::Unsupported,
+                          "target does not define a stack pointer for stack save");
+  const MachineOperand *DestinationOperand = Instruction.operandAt(0);
+  const Register *Destination =
+      DestinationOperand == nullptr ? nullptr : DestinationOperand->getIf<Register>();
+  if (Instruction.operandCount() != 1 || Destination == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "G_STACKSAVE requires one register destination");
+  if (auto Point = MIRBuilder.setInsertPointBefore(Instruction); !Point)
+    return std::unexpected(Point.error());
+  if (auto Built = MIRBuilder.buildCopy(*Destination, Register::fromPhysical(StackPointer));
+      !Built) {
+    return std::unexpected(Built.error());
+  }
+  MachineBasicBlock *Block = Instruction.parent();
+  if (Block == nullptr)
+    return makeUnexpected(ErrorCode::InternalError, "G_STACKSAVE became detached during lowering");
+  return Block->removeInstruction(Instruction);
+}
+
+Expected<void> LegalizerHelper::lowerStackRestore(MachineInstr &Instruction) {
+  const PhysicalRegister StackPointer = MIRBuilder.target().stackPointerRegisterToSaveRestore();
+  if (!StackPointer.isValid())
+    return makeUnexpected(ErrorCode::Unsupported,
+                          "target does not define a stack pointer for stack restore");
+  const MachineOperand *SourceOperand = Instruction.operandAt(0);
+  const Register *Source = SourceOperand == nullptr ? nullptr : SourceOperand->getIf<Register>();
+  if (Instruction.operandCount() != 1 || Source == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "G_STACKRESTORE requires one register source");
+  if (auto Point = MIRBuilder.setInsertPointBefore(Instruction); !Point)
+    return std::unexpected(Point.error());
+  if (auto Built = MIRBuilder.buildCopy(Register::fromPhysical(StackPointer), *Source); !Built)
+    return std::unexpected(Built.error());
+  MachineBasicBlock *Block = Instruction.parent();
+  if (Block == nullptr) {
+    return makeUnexpected(ErrorCode::InternalError,
+                          "G_STACKRESTORE became detached during lowering");
+  }
+  return Block->removeInstruction(Instruction);
+}
+
 Expected<void> LegalizerHelper::lowerAbsToAddXor(MachineInstr &Instruction) {
   const Register Destination = *Instruction.operandAt(0)->getIf<Register>();
   const Register Source = *Instruction.operandAt(1)->getIf<Register>();
diff --git a/backend/lib/CodeGen/TargetDescription.cpp b/backend/lib/CodeGen/TargetDescription.cpp
index 2a64a02..a867a5b 100644
--- a/backend/lib/CodeGen/TargetDescription.cpp
+++ b/backend/lib/CodeGen/TargetDescription.cpp
@@ -52,11 +52,18 @@ namespace {
 Expected<TargetDescription> TargetDescription::create(InstructionSet Instructions,
                                                       TargetRegisterInfo Registers,
                                                       TargetMIRTextInfoDesc MIRTextDescription,
-                                                      TargetBooleanContentInfo BooleanInfo) {
+                                                      TargetBooleanContentInfo BooleanInfo,
+                                                      TargetStackInfo StackInfo) {
   auto MIRText = TargetMIRTextInfo::create(std::move(MIRTextDescription), Registers);
   if (!MIRText)
     return std::unexpected(MIRText.error());
 
+  if (StackInfo.StackPointerRegisterToSaveRestore.isValid() &&
+      Registers.physicalRegister(StackInfo.StackPointerRegisterToSaveRestore) == nullptr) {
+    return makeUnexpected(ErrorCode::InvalidArgument,
+                          "target stack pointer is not a known physical register");
+  }
+
   for (const InstructionDescriptor &Instruction : Instructions.descriptors()) {
     for (const OperandDescriptor &Operand : Instruction.Operands) {
       if (Operand.AllowAnyTargetFlags && Operand.AllowedTargetFlags != 0) {
@@ -124,7 +131,7 @@ Expected<TargetDescription> TargetDescription::create(InstructionSet Instruction
   }
 
   return TargetDescription(std::move(Instructions), std::move(Registers), std::move(*MIRText),
-                           BooleanInfo);
+                           BooleanInfo, StackInfo);
 }
 
 Expected<void> TargetDescription::validateVirtualRegister(Register Reg,
diff --git a/backend/lib/Target/RISCV/RISCVTarget.cpp b/backend/lib/Target/RISCV/RISCVTarget.cpp
index 1735386..0c202a2 100644
--- a/backend/lib/Target/RISCV/RISCVTarget.cpp
+++ b/backend/lib/Target/RISCV/RISCVTarget.cpp
@@ -245,8 +245,15 @@ Expected<TargetDescription> createTargetDescription() {
       .ScalarInteger = BooleanContent::ZeroOrOne,
       .ScalarFloatingPoint = BooleanContent::ZeroOrOne,
   };
+  auto StackPointer = physicalRegister(PhysicalRegisterID::X2);
+  if (!StackPointer)
+    return std::unexpected(StackPointer.error());
+  TargetStackInfo StackInfo{
+      .StackPointerRegisterToSaveRestore = *StackPointer,
+      .GrowthDirection = StackGrowthDirection::Down,
+  };
   return TargetDescription::create(std::move(*Instructions), std::move(*Registers),
-                                   std::move(MIRText), BooleanInfo);
+                                   std::move(MIRText), BooleanInfo, StackInfo);
 }
 
 std::shared_ptr<const RegisterMask> createCallPreservedMask(CallingConvention Convention) {