[aarch64] handle stack frames that are larger than 4096 bytes

src/banjo/CMakeLists.txt

@@ -114,6 +114,7 @@ add_library(banjo STATIC
     target/aarch64/aarch64_instr_merge_pass.cpp
     target/aarch64/aarch64_reg_analyzer.cpp
     target/aarch64/aarch64_ssa_lowerer.cpp
+    target/aarch64/aarch64_stack_offset_fixup_pass.cpp
     target/aarch64/aarch64_target.cpp
     target/x86_64/ms_abi_calling_conv.cpp
     target/x86_64/sys_v_calling_conv.cpp

src/banjo/target/aarch64/aapcs_calling_conv.cpp

@@ -4,6 +4,7 @@
 #include "banjo/codegen/late_reg_alloc.hpp"
 #include "banjo/codegen/machine_pass_utils.hpp"
 #include "banjo/codegen/ssa_lowerer.hpp"
+#include "banjo/mcode/instruction.hpp"
 #include "banjo/mcode/register.hpp"
 #include "banjo/target/aarch64/aarch64_encoding_info.hpp"
 #include "banjo/target/aarch64/aarch64_opcode.hpp"
@@ -196,13 +197,10 @@ std::vector<mcode::Instruction> AAPCSCallingConv::get_prolog(mcode::Function *fu
          mcode::Operand::from_aarch64_addr(AArch64Address::new_base_offset_write(sp, -16))}
     ));
 
-    prolog.push_back(mcode::Instruction(
-        AArch64Opcode::SUB,
-        {mcode::Operand::from_register(sp, 8),
-         mcode::Operand::from_register(sp, 8),
-         mcode::Operand::from_int_immediate(size)},
-        mcode::Instruction::FLAG_ALLOCA
-    ));
+    modify_sp(AArch64Opcode::SUB, size, [&prolog](mcode::Instruction instr) {
+        instr.set_flag(mcode::Instruction::FLAG_ALLOCA);
+        prolog.push_back(std::move(instr));
+    });
 
     return prolog;
 
@@ -258,12 +256,7 @@ std::vector<mcode::Instruction> AAPCSCallingConv::get_epilog(mcode::Function *fu
 
     std::vector<mcode::Instruction> epilog;
 
-    epilog.push_back(mcode::Instruction(
-        AArch64Opcode::ADD,
-        {mcode::Operand::from_register(sp, 8),
-         mcode::Operand::from_register(sp, 8),
-         mcode::Operand::from_int_immediate(size)}
-    ));
+    modify_sp(AArch64Opcode::ADD, size, [&epilog](mcode::Instruction instr) { epilog.push_back(std::move(instr)); });
 
     epilog.push_back(mcode::Instruction(
         AArch64Opcode::LDP,
@@ -319,6 +312,29 @@ std::vector<mcode::Instruction> AAPCSCallingConv::get_epilog(mcode::Function *fu
     */
 }
 
+void AAPCSCallingConv::modify_sp(
+    mcode::Opcode opcode,
+    unsigned value,
+    const std::function<void(mcode::Instruction)> &emit
+) {
+    mcode::Register sp = mcode::Register::from_physical(AArch64Register::SP);
+    mcode::Operand m_sp = mcode::Operand::from_register(sp, 8);
+
+    if (value < 4096) {
+        mcode::Operand m_imm = mcode::Operand::from_int_immediate(value);
+        emit(mcode::Instruction(opcode, {m_sp, m_sp, m_imm}));
+    } else if (value < 4096 * 4096) {
+        mcode::Operand m_imm_shifted = mcode::Operand::from_int_immediate(value >> 12);
+        mcode::Operand m_shift = mcode::Operand::from_aarch64_left_shift(12);
+        emit(mcode::Instruction(opcode, {m_sp, m_sp, m_imm_shifted, m_shift}));
+
+        mcode::Operand m_imm_remainder = mcode::Operand::from_int_immediate(value & 0xFFF);
+        emit(mcode::Instruction(opcode, {m_sp, m_sp, m_imm_remainder}));
+    } else {
+        ASSERT_UNREACHABLE;
+    }
+}
+
 mcode::InstrIter AAPCSCallingConv::fix_up_instr(
     mcode::BasicBlock &basic_block,
     mcode::InstrIter iter,

src/banjo/target/aarch64/aapcs_calling_conv.hpp

@@ -4,6 +4,8 @@
 #include "banjo/mcode/calling_convention.hpp"
 #include "banjo/mcode/register.hpp"
 
+#include <functional>
+
 namespace banjo {
 
 namespace target {
@@ -34,6 +36,8 @@ class AAPCSCallingConv : public mcode::CallingConvention {
 
 private:
     mcode::Operand get_arg_dst(ssa::Operand &operand, int index, codegen::SSALowerer &lowerer);
+
+    void modify_sp(mcode::Opcode opcode, unsigned value, const std::function<void(mcode::Instruction)> &emit);
 };
 
 } // namespace target

src/banjo/target/aarch64/aarch64_stack_offset_fixup_pass.cpp

@@ -0,0 +1,56 @@
+#include "aarch64_stack_offset_fixup_pass.hpp"
+
+#include "banjo/mcode/instruction.hpp"
+#include "banjo/target/aarch64/aarch64_opcode.hpp"
+#include "banjo/utils/utils.hpp"
+
+namespace banjo {
+namespace target {
+
+void AArch64StackOffsetFixupPass::run(mcode::Module &module_) {
+    for (mcode::Function *func : module_.get_functions()) {
+        for (mcode::BasicBlock &block : *func) {
+            run(block);
+        }
+    }
+}
+
+void AArch64StackOffsetFixupPass::run(mcode::BasicBlock &block) {
+    mcode::StackFrame &frame = block.get_func()->get_stack_frame();
+
+    for (mcode::InstrIter iter = block.begin(); iter != block.end(); ++iter) {
+        mcode::Instruction &instr = *iter;
+        mcode::Opcode opcode = instr.get_opcode();
+
+        if (!Utils::is_one_of<mcode::Opcode>(opcode, {AArch64Opcode::ADD, AArch64Opcode::SUB})) {
+            continue;
+        }
+
+        mcode::Operand &operand = instr.get_operand(2);
+        if (!operand.is_stack_slot_offset() || instr.get_operands().get_size() == 4) {
+            continue;
+        }
+
+        mcode::Operand::StackSlotOffset offset = operand.get_stack_slot_offset();
+        mcode::StackSlot &stack_slot = frame.get_stack_slot(offset.slot_index);
+        unsigned total_offset = stack_slot.get_offset() + offset.additional_offset;
+
+        if (total_offset < 4096) {
+            continue;
+        }
+
+        ASSERT(total_offset < 4096 * 4096);
+
+        operand = mcode::Operand::from_int_immediate(total_offset >> 12, operand.get_size());
+        instr.get_operands().append(mcode::Operand::from_aarch64_left_shift(12));
+
+        mcode::Operand m_imm_remainder = mcode::Operand::from_int_immediate(total_offset & 0xFFF);
+        block.insert_after(
+            iter,
+            mcode::Instruction(opcode, {instr.get_operand(0), instr.get_operand(0), m_imm_remainder})
+        );
+    }
+}
+
+} // namespace target
+} // namespace banjo

src/banjo/target/aarch64/aarch64_stack_offset_fixup_pass.hpp

@@ -0,0 +1,21 @@
+#ifndef BANJO_TARGET_AARCH64_AARCH64_STACK_OFFSET_FIXUP_PASS_H
+#define BANJO_TARGET_AARCH64_AARCH64_STACK_OFFSET_FIXUP_PASS_H
+
+#include "banjo/codegen/machine_pass.hpp"
+
+namespace banjo {
+namespace target {
+
+class AArch64StackOffsetFixupPass final : public codegen::MachinePass {
+
+public:
+    void run(mcode::Module &module_);
+
+private:
+    void run(mcode::BasicBlock &block);
+};
+
+} // namespace target
+} // namespace banjo
+
+#endif

src/banjo/target/aarch64/aarch64_target.cpp

@@ -3,6 +3,7 @@
 #include "banjo/emit/aarch64_asm_emitter.hpp"
 #include "banjo/target/aarch64/aarch64_instr_merge_pass.hpp"
 #include "banjo/target/aarch64/aarch64_ssa_lowerer.hpp"
+#include "banjo/target/aarch64/aarch64_stack_offset_fixup_pass.hpp"
 
 namespace banjo {
 
@@ -19,7 +20,7 @@ std::vector<codegen::MachinePass *> AArch64Target::create_pre_passes() {
 }
 
 std::vector<codegen::MachinePass *> AArch64Target::create_post_passes() {
-    return {};
+    return {new AArch64StackOffsetFixupPass()};
 }
 
 std::string AArch64Target::get_output_file_ext() {