Implement initial spilling algorithm for RISC-V

tests/filecheck/backend/riscv/convert_riscv_stack_to_riscv.mlir

@@ -0,0 +1,56 @@
+// RUN: xdsl-opt --split-input-file -p convert-riscv-stack-to-riscv %s | filecheck %s
+
+builtin.module {
+  riscv_func.func @main() {
+    %0 = "test.op"() : () -> !riscv.reg
+    %1 = riscv_stack.alloca : () -> !riscv_stack.ptr<i32>
+    riscv_stack.store %1, %0 : (!riscv_stack.ptr<i32>, !riscv.reg) -> ()
+    %2 = riscv_stack.load %1 : (!riscv_stack.ptr<i32>) -> !riscv.reg
+    %3 = riscv_stack.alloca : () -> !riscv_stack.ptr<i32>
+    riscv_stack.store %3, %2 : (!riscv_stack.ptr<i32>, !riscv.reg) -> ()
+    riscv_func.return
+  }
+}
+// CHECK:       builtin.module {
+// CHECK-NEXT:    riscv_func.func @main() {
+// CHECK-NEXT:      %0 = rv32.get_register : !riscv.reg<sp>
+// CHECK-NEXT:      %1 = riscv.addi %0, -8 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      %2 = "test.op"() : () -> !riscv.reg
+// CHECK-NEXT:      riscv.sw %0, %2, 0 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      %3 = riscv.lw %0, 0 : (!riscv.reg<sp>) -> !riscv.reg
+// CHECK-NEXT:      riscv.sw %0, %3, 4 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      %4 = riscv.addi %0, 8 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      riscv_func.return
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+
+// -----
+// Test floats
+builtin.module {
+  riscv_func.func @main() {
+    %0, %1 = "test.op"() : () -> (!riscv.reg, !riscv.freg)
+    %3 = riscv_stack.alloca : () -> !riscv_stack.ptr<f64>
+    %2 = riscv_stack.alloca : () -> !riscv_stack.ptr<i32>
+
+    riscv_stack.store %2, %0 : (!riscv_stack.ptr<i32>, !riscv.reg) -> ()
+    riscv_stack.store %3, %1 : (!riscv_stack.ptr<f64>, !riscv.freg) -> ()
+
+    %4 = riscv_stack.load %2 : (!riscv_stack.ptr<i32>) -> (!riscv.reg)
+    %5 = riscv_stack.load %3 : (!riscv_stack.ptr<f64>) -> (!riscv.freg)
+
+    riscv_func.return
+  }
+}
+// CHECK:       builtin.module {
+// CHECK-NEXT:    riscv_func.func @main() {
+// CHECK-NEXT:      %0 = rv32.get_register : !riscv.reg<sp>
+// CHECK-NEXT:      %1 = riscv.addi %0, -12 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      %2, %3 = "test.op"() : () -> (!riscv.reg, !riscv.freg)
+// CHECK-NEXT:      riscv.sw %0, %2, 8 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      riscv.fsd %0, %3, 0 : (!riscv.reg<sp>, !riscv.freg) -> ()
+// CHECK-NEXT:      %4 = riscv.lw %0, 8 : (!riscv.reg<sp>) -> !riscv.reg
+// CHECK-NEXT:      %5 = riscv.fld %0, 0 : (!riscv.reg<sp>) -> !riscv.freg
+// CHECK-NEXT:      %6 = riscv.addi %0, 12 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      riscv_func.return
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }

tests/filecheck/backend/riscv/prologue_epilogue_insertion.mlir

@@ -1,5 +1,5 @@
-// RUN: xdsl-opt --split-input-file -p "riscv-prologue-epilogue-insertion" %s | filecheck %s
-// RUN: xdsl-opt --split-input-file -p "riscv-prologue-epilogue-insertion{flen=4}" %s | filecheck %s --check-prefix=CHECK-SMALL-FLEN
+// RUN: xdsl-opt --split-input-file -p "riscv-prologue-epilogue-insertion,convert-riscv-stack-to-riscv" %s | filecheck %s
+// RUN: xdsl-opt --split-input-file -p "riscv-prologue-epilogue-insertion{flen=4},convert-riscv-stack-to-riscv" %s | filecheck %s --check-prefix=CHECK-SMALL-FLEN
 
 // CHECK: func @main
 riscv_func.func @main() {

tests/filecheck/dialects/riscv/riscv_stack.mlir

@@ -0,0 +1,12 @@
+// RUN: xdsl-opt --split-input-file --verify-diagnostics %s | filecheck %s
+
+%0 = "test.op"(): () -> !riscv.reg
+
+%1 = riscv_stack.alloca : () -> !riscv_stack.ptr<i32>
+// CHECK: %{{.*}} = riscv_stack.alloca : () -> !riscv_stack.ptr<i32>
+
+riscv_stack.store %1, %0: (!riscv_stack.ptr<i32>, !riscv.reg) -> ()
+// CHECK-NEXT: riscv_stack.store %{{.*}}, %{{.*}} : (!riscv_stack.ptr<i32>, !riscv.reg) -> ()
+
+riscv_stack.load %1: (!riscv_stack.ptr<i32>) -> !riscv.reg
+// CHECK-NEXT: %{{.*}} = riscv_stack.load %{{.*}} : (!riscv_stack.ptr<i32>) -> !riscv.reg

tests/filecheck/transforms/test-riscv-spilling.mlir

@@ -0,0 +1,27 @@
+// RUN: xdsl-opt --split-input-file -p test-riscv-spilling,convert-riscv-stack-to-riscv %s | filecheck %s
+builtin.module {
+  riscv_func.func @main() {
+    %0, %1, %2 = "test.op"() : () -> (!riscv.reg, !riscv.reg, !riscv.reg)
+    %3 = riscv.add %0, %1 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+    %4 = riscv.add %2, %3 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+    %5 = riscv.add %0, %1 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+    riscv_func.return
+  }
+}
+// CHECK:       builtin.module {
+// CHECK-NEXT:    riscv_func.func @main() {
+// CHECK-NEXT:      %0 = rv32.get_register : !riscv.reg<sp>
+// CHECK-NEXT:      %1 = riscv.addi %0, -12 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      %2, %3, %4 = "test.op"() : () -> (!riscv.reg, !riscv.reg, !riscv.reg)
+// CHECK-NEXT:      riscv.sw %0, %4, 0 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      %5 = riscv.add %2, %3 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+// CHECK-NEXT:      riscv.sw %0, %2, 4 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      %6 = riscv.lw %0, 0 : (!riscv.reg<sp>) -> !riscv.reg
+// CHECK-NEXT:      %7 = riscv.add %6, %5 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+// CHECK-NEXT:      riscv.sw %0, %6, 8 : (!riscv.reg<sp>, !riscv.reg) -> ()
+// CHECK-NEXT:      %8 = riscv.lw %0, 4 : (!riscv.reg<sp>) -> !riscv.reg
+// CHECK-NEXT:      %9 = riscv.add %8, %3 : (!riscv.reg, !riscv.reg) -> !riscv.reg
+// CHECK-NEXT:      %10 = riscv.addi %0, 12 : (!riscv.reg<sp>) -> !riscv.reg<sp>
+// CHECK-NEXT:      riscv_func.return
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }

xdsl/backend/riscv/lowering/convert_riscv_stack_to_riscv.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass, field
+
+from xdsl.builder import Builder
+from xdsl.context import Context
+from xdsl.dialects import riscv_func, rv32
+from xdsl.dialects.builtin import FixedBitwidthType, IntegerType, ModuleOp
+from xdsl.dialects.riscv.ops import AddiOp, FLdOp, FSdOp, LwOp, SwOp
+from xdsl.dialects.riscv.registers import FloatRegisterType, IntRegisterType, Registers
+from xdsl.dialects.riscv.stack import AllocaOp, LoadOp, StoreOp
+from xdsl.ir import Use
+from xdsl.passes import ModulePass
+from xdsl.rewriter import InsertPoint, Rewriter
+from xdsl.utils.exceptions import PassFailedException
+
+
+def get_type_size(value_type: FixedBitwidthType):
+    return (value_type.bitwidth + 7) // 8
+
+
+@dataclass(frozen=True)
+class ConvertRiscvStackToRiscvPass(ModulePass):
+    name = "convert-riscv-stack-to-riscv"
+
+    align_sp: bool = field(default=False)
+
+    def apply(self, ctx: Context, op: ModuleOp) -> None:
+        for func_op in op.walk():
+            if not isinstance(func_op, riscv_func.FuncOp):
+                continue
+
+            cur_offset = 0
+            # Need to instantiate generator so erasing op doesn't stop iteration
+            alloca_ops = tuple(op for op in func_op.walk() if isinstance(op, AllocaOp))
+
+            if not alloca_ops:
+                return
+
+            stack_ptr = rv32.GetRegisterOp(
+                Registers.SP
+            )  # inserted in insert_prologue_epilogue
+            # For each function, give each alloca its stack offset
+            for alloca_op in alloca_ops:
+                stack_slot = alloca_op.ref.type
+                value_size = get_type_size(stack_slot.value_type)
+
+                for use in alloca_op.ref.uses:
+                    self.rewrite_stack_slot_use(
+                        use, cur_offset, stack_slot.value_type, stack_ptr
+                    )
+
+                Rewriter.erase_op(alloca_op)
+                cur_offset += value_size
+
+            # Adjust stack pointer
+            self.insert_prologue_epilogue(cur_offset, func_op, stack_ptr)
+
+    def rewrite_stack_slot_use(
+        self,
+        use: Use,
+        stack_offset: int,
+        val_type: FixedBitwidthType,
+        stack_ptr: rv32.GetRegisterOp,
+    ):
+        if isinstance(use.operation, StoreOp):
+            if isinstance(val_type, IntegerType):
+                Rewriter.replace_op(
+                    use.operation,
+                    (SwOp(stack_ptr, use.operation.rs, stack_offset),),
+                )
+            else:
+                Rewriter.replace_op(
+                    use.operation,
+                    (FSdOp(stack_ptr, use.operation.rs, stack_offset),),
+                )
+
+        elif isinstance(use.operation, LoadOp):
+            if isinstance(val_type, IntegerType):
+                assert isinstance(use.operation.rd.type, IntRegisterType)
+                Rewriter.replace_op(
+                    use.operation,
+                    (LwOp(stack_ptr, stack_offset, rd=use.operation.rd.type),),
+                )
+            else:
+                assert isinstance(use.operation.rd.type, FloatRegisterType)
+                Rewriter.replace_op(
+                    use.operation,
+                    (FLdOp(stack_ptr, stack_offset, rd=use.operation.rd.type),),
+                )
+
+        else:
+            raise PassFailedException("Invalid use of StackSlotType: ", use.operation)
+
+    def insert_prologue_epilogue(
+        self,
+        total_offset: int,
+        func_op: riscv_func.FuncOp,
+        stack_ptr: rv32.GetRegisterOp,
+    ):
+        # Align SP to 16 bytes (from RISC-V calling convention)
+        if self.align_sp:
+            total_offset = (total_offset + 15) & ~15
+
+        if total_offset > 0:
+            # prologue
+            builder = Builder(InsertPoint.at_start(func_op.body.blocks[0]))
+            builder.insert(stack_ptr)
+            builder.insert(AddiOp(stack_ptr, -total_offset, rd=Registers.SP))
+            # epilogue
+            # using logic from prologue_epilogue_insertion.py
+            for block in func_op.body.blocks:
+                ret_op = block.last_op
+                if not isinstance(ret_op, riscv_func.ReturnOp):
+                    continue
+
+                builder = Builder(InsertPoint.before(ret_op))
+                builder.insert(AddiOp(stack_ptr, total_offset, rd=Registers.SP))

xdsl/backend/riscv/prologue_epilogue_insertion.py

@@ -9,8 +9,8 @@
     FloatRegisterType,
     IntRegisterType,
     Registers,
-    RISCVRegisterType,
 )
+from xdsl.dialects.riscv.stack import AllocaOp, LoadOp, StoreOp
 from xdsl.passes import ModulePass
 
 
@@ -47,27 +47,26 @@ def _process_function(self, func: riscv_func.FuncOp) -> None:
         if not used_callee_preserved_registers:
             return
 
-        def get_register_size(r: RISCVRegisterType):
-            if isinstance(r, IntRegisterType):
-                return self.xlen
-            return self.flen
-
         # Build the prologue at the beginning of the function.
         builder = Builder(InsertPoint.at_start(func.body.blocks[0]))
-        sp_register = builder.insert(rv32.GetRegisterOp(Registers.SP))
-        stack_size = sum(get_register_size(r) for r in used_callee_preserved_registers)
-        builder.insert(riscv.AddiOp(sp_register, -stack_size, rd=Registers.SP))
-        offset = 0
+        memrefs: list[AllocaOp] = []
         for reg in used_callee_preserved_registers:
             if isinstance(reg, IntRegisterType):
+                alloca_op = builder.insert(AllocaOp(builtin.i32))
                 reg_op = builder.insert(rv32.GetRegisterOp(reg))
-                op = riscv.SwOp(rs1=sp_register, rs2=reg_op, immediate=offset)
             else:
+                match self.flen:
+                    case 8:
+                        float_size = builtin.f64
+                    case 4:
+                        float_size = builtin.f32
+                    case _:
+                        raise NotImplementedError("flen must be 4 or 8")
+
+                alloca_op = builder.insert(AllocaOp(float_size))
                 reg_op = builder.insert(riscv.GetFloatRegisterOp(reg))
-                op = riscv.FSdOp(rs1=sp_register, rs2=reg_op, immediate=offset)
-
-            builder.insert(op)
-            offset += get_register_size(reg)
+            builder.insert(StoreOp(alloca_op, reg_op))
+            memrefs.append(alloca_op)
 
         # Now build the epilogue right before every return operation.
         for block in func.body.blocks:
@@ -76,16 +75,8 @@ def get_register_size(r: RISCVRegisterType):
                 continue
 
             builder = Builder(InsertPoint.before(ret_op))
-            offset = 0
-            for reg in used_callee_preserved_registers:
-                if isinstance(reg, IntRegisterType):
-                    op = riscv.LwOp(rs1=sp_register, rd=reg, immediate=offset)
-                else:
-                    op = riscv.FLdOp(rs1=sp_register, rd=reg, immediate=offset)
-                builder.insert(op)
-                offset += get_register_size(reg)
-
-            builder.insert(riscv.AddiOp(sp_register, stack_size, rd=Registers.SP))
+            for memref, reg in zip(memrefs, used_callee_preserved_registers):
+                builder.insert(LoadOp(memref, rd=reg))
 
     def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
         for func in op.walk():

xdsl/dialects/__init__.py

@@ -288,6 +288,11 @@ def get_riscv_snitch():
 
         return RISCV_Snitch
 
+    def get_riscv_stack():
+        from xdsl.dialects.riscv.stack import RISCVStack
+
+        return RISCVStack
+
     def get_scf():
         from xdsl.dialects.scf import Scf
 
@@ -445,6 +450,7 @@ def get_x86_scf():
         "riscv_scf": get_riscv_scf,
         "riscv_cf": get_riscv_cf,
         "riscv_snitch": get_riscv_snitch,
+        "riscv_stack": get_riscv_stack,
         "scf": get_scf,
         "seq": get_seq,
         "shard": get_shard,

xdsl/dialects/riscv/__init__.py

@@ -9,6 +9,7 @@
 
 # The imports with re-export are to preserve the API of the dialect, as it was when it
 # was a single file. They will be deprecated and removed in the future.
+from . import stack as stack
 from .abstract_ops import (
     AssemblyInstructionArg as AssemblyInstructionArg,
 )