backend: (riscv) add option to insert regalloc stats (#2958)

Inserts a little json that we can parse in the experiments repo.

tests/filecheck/backend/riscv/register-allocation/preallocated.mlir

@@ -1,22 +1,23 @@
-// RUN: xdsl-opt -p "riscv-allocate-registers{allocation_strategy=LivenessBlockNaive}" %s | filecheck %s --check-prefix=LIVE-BNAIVE
+// RUN: xdsl-opt -p "riscv-allocate-registers{allocation_strategy=LivenessBlockNaive add-regalloc-stats=true}" %s | filecheck %s --check-prefix=LIVE-BNAIVE
 
 riscv_func.func @main() {
   %0 = riscv.li 6 : !riscv.reg
   %1 = riscv.li 5 : !riscv.reg<t0>
-  %3 = riscv.fcvt.s.w %0 : (!riscv.reg) -> !riscv.freg
+  %3 = riscv.fcvt.s.w %0 : (!riscv.reg) -> !riscv.freg<ft0>
   %4 = riscv.fcvt.s.w %1 : (!riscv.reg<t0>) -> !riscv.freg
-  %5 = riscv.fadd.s %3, %4 : (!riscv.freg, !riscv.freg) -> !riscv.freg
+  %5 = riscv.fadd.s %3, %4 : (!riscv.freg<ft0>, !riscv.freg) -> !riscv.freg
   %2 = riscv.add %0, %1 : (!riscv.reg, !riscv.reg<t0>) -> !riscv.reg
   riscv_func.return
 }
 
 // LIVE-BNAIVE:       builtin.module {
+// LIVE-BNAIVE-NEXT:    riscv.comment {"comment" = "Regalloc stats: {\"preallocated_float\": 1, \"preallocated_int\": 1, \"allocated_float\": 5, \"allocated_int\": 7}"} : () -> ()
 // LIVE-BNAIVE-NEXT:    riscv_func.func @main() {
 // LIVE-BNAIVE-NEXT:      %0 = riscv.li 6 : !riscv.reg<t1>
 // LIVE-BNAIVE-NEXT:      %1 = riscv.li 5 : !riscv.reg<t0>
 // LIVE-BNAIVE-NEXT:      %2 = riscv.fcvt.s.w %0 : (!riscv.reg<t1>) -> !riscv.freg<ft0>
 // LIVE-BNAIVE-NEXT:      %3 = riscv.fcvt.s.w %1 : (!riscv.reg<t0>) -> !riscv.freg<ft1>
-// LIVE-BNAIVE-NEXT:      %4 = riscv.fadd.s %2, %3 : (!riscv.freg<ft0>, !riscv.freg<ft1>) -> !riscv.freg<ft0>
+// LIVE-BNAIVE-NEXT:      %4 = riscv.fadd.s %2, %3 : (!riscv.freg<ft0>, !riscv.freg<ft1>) -> !riscv.freg<ft1>
 // LIVE-BNAIVE-NEXT:      %5 = riscv.add %0, %1 : (!riscv.reg<t1>, !riscv.reg<t0>) -> !riscv.reg<t1>
 // LIVE-BNAIVE-NEXT:      riscv_func.return
 // LIVE-BNAIVE-NEXT:    }

tests/filecheck/projects/riscv-backend-paper/bottom_up.mlir

@@ -24,6 +24,7 @@ func.func public @ssum(
 // CHECK:       .text
 // CHECK-NEXT:  .globl ssum
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 6, "allocated_int": 21}
 // CHECK-NEXT:  ssum:
 // CHECK-NEXT:      mv t2, a0
 // CHECK-NEXT:      mv t1, a1
@@ -80,6 +81,7 @@ func.func public @conv_2d_nchw_fchw_d1_s1_3x3(
 // CHECK:       .text
 // CHECK-NEXT:  .globl conv_2d_nchw_fchw_d1_s1_3x3
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 57, "allocated_int": 109}
 // CHECK-NEXT:  conv_2d_nchw_fchw_d1_s1_3x3:
 // CHECK-NEXT:      mv t2, a0
 // CHECK-NEXT:      mv t1, a1
@@ -203,6 +205,7 @@ func.func public @conv_2d_nchw_fchw_d1_s1_3x3(
 // CHECK:       .text
 // CHECK-NEXT:  .globl ddot
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 15, "allocated_int": 25}
 // CHECK-NEXT:  ddot:
 // CHECK-NEXT:      mv t2, a0
 // CHECK-NEXT:      mv t1, a1
@@ -249,6 +252,7 @@ func.func public @conv_2d_nchw_fchw_d1_s1_3x3(
 // CHECK:       .text
 // CHECK-NEXT:  .globl dsum
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 6, "allocated_int": 21}
 // CHECK-NEXT:  dsum:
 // CHECK-NEXT:      mv t2, a0
 // CHECK-NEXT:      mv t1, a1
@@ -292,6 +296,7 @@ func.func public @conv_2d_nchw_fchw_d1_s1_3x3(
 // CHECK:       .text
 // CHECK-NEXT:  .globl fill
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 4, "preallocated_int": 2, "allocated_float": 5, "allocated_int": 15}
 // CHECK-NEXT:  fill:
 // CHECK-NEXT:      fmv.d ft3, fa0
 // CHECK-NEXT:      mv t0, a0
@@ -348,6 +353,7 @@ func.func public @conv_2d_nchw_fchw_d1_s1_3x3(
 // CHECK-NEXT:  .text
 // CHECK-NEXT:  .globl matmul
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 57, "allocated_int": 67}
 // CHECK-NEXT:  matmul:
 // CHECK-NEXT:      mv t0, a0
 // CHECK-NEXT:      mv t1, a1
@@ -457,6 +463,7 @@ func.func public @pooling_nchw_max_d1_s2_3x3(
 // CHECK-NEXT:  .text
 // CHECK-NEXT:  .globl pooling_nchw_max_d1_s2_3x3
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 3, "allocated_float": 49, "allocated_int": 92}
 // CHECK-NEXT:  pooling_nchw_max_d1_s2_3x3:
 // CHECK-NEXT:      mv t1, a0
 // CHECK-NEXT:      mv t0, a1
@@ -552,6 +559,7 @@ func.func public @pooling_nchw_max_d1_s2_3x3(
 // CHECK:       .text
 // CHECK-NEXT:  .globl relu
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 3, "allocated_float": 6, "allocated_int": 22}
 // CHECK-NEXT:  relu:
 // CHECK-NEXT:      mv t1, a0
 // CHECK-NEXT:      mv t0, a1
@@ -610,6 +618,7 @@ func.func public @pooling_nchw_sum_d1_s2_3x3(
 // CHECK-NEXT:  .text
 // CHECK-NEXT:  .globl pooling_nchw_sum_d1_s2_3x3
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 3, "allocated_float": 49, "allocated_int": 94}
 // CHECK-NEXT:  pooling_nchw_sum_d1_s2_3x3:
 // CHECK-NEXT:      mv t1, a0
 // CHECK-NEXT:      mv t0, a1

tests/filecheck/projects/riscv-backend-paper/nsnet.mlir

@@ -15,6 +15,7 @@ func.func @main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1(%
 // CHECK:       .text
 // CHECK-NEXT:  .globl main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1
 // CHECK-NEXT:  .p2align 2
+// CHECK-NEXT:  # Regalloc stats: {"preallocated_float": 3, "preallocated_int": 4, "allocated_float": 75, "allocated_int": 44}
 // CHECK-NEXT:  main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1:
 // CHECK-NEXT:      mv t2, a0
 // CHECK-NEXT:      mv t1, a1
@@ -62,6 +63,7 @@ func.func @main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1(%
 // CHECK-OPT:       .text
 // CHECK-OPT-NEXT:  .globl main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1
 // CHECK-OPT-NEXT:  .p2align 2
+// CHECK-OPT-NEXT:  # Regalloc stats: {"preallocated_float": 0, "preallocated_int": 4, "allocated_float": 8, "allocated_int": 63}
 // CHECK-OPT-NEXT:  main$async_dispatch_0_matmul_transpose_b_1x400x161_f64$xdsl_kernel1:
 // CHECK-OPT-NEXT:      mv t4, a0
 // CHECK-OPT-NEXT:      mv t3, a1

xdsl/backend/riscv/register_allocation.py

@@ -1,5 +1,6 @@
 import abc
-from collections.abc import Sequence
+import json
+from collections.abc import Iterable, Sequence
 from itertools import chain
 from typing import cast
 
@@ -14,7 +15,8 @@
     RISCVAsmOperation,
     RISCVRegisterType,
 )
-from xdsl.ir import Block, Operation, SSAValue
+from xdsl.ir import Attribute, Block, Operation, SSAValue
+from xdsl.rewriter import InsertPoint, Rewriter
 from xdsl.transforms.canonicalization_patterns.riscv import get_constant_value
 from xdsl.transforms.snitch_register_allocation import get_snitch_reserved
 from xdsl.utils.exceptions import DiagnosticException
@@ -64,6 +66,18 @@ def allocate_func(self, func: riscv_func.FuncOp) -> None:
         raise NotImplementedError()
 
 
+def count_reg_types(regs: Iterable[Attribute]) -> tuple[int, int]:
+    """
+    Returns a tuple containing the count of IntRegister and FloatRegister in the iterable.
+    """
+    num_complex = sum(
+        isinstance(reg, IntRegisterType) + 1j * isinstance(reg, FloatRegisterType)
+        for reg in regs
+    )
+
+    return int(num_complex.real), int(num_complex.imag)
+
+
 class RegisterAllocatorLivenessBlockNaive(RegisterAllocator):
     """
     It traverses the use-def SSA chain backwards (i.e., from uses to defs) and:
@@ -299,7 +313,16 @@ def allocate_frep_loop(self, loop: riscv_snitch.FRepOperation) -> None:
             for op in reversed(loop.body.block.ops):
                 self.process_operation(op)
 
-    def allocate_func(self, func: riscv_func.FuncOp) -> None:
+    def allocate_func(
+        self, func: riscv_func.FuncOp, *, add_regalloc_stats: bool = False
+    ) -> None:
+        """
+        Allocates values in function passed in to registers.
+        The whole function must have been lowered to the relevant riscv dialects
+        and it must contain no unrealized casts.
+        If `add_regalloc_stats` is set to `True`, then a comment op will be inserted
+        before the function op passed in with a json containing the relevant data.
+        """
         if not func.body.blocks:
             # External function declaration
             return
@@ -333,6 +356,28 @@ def allocate_func(self, func: riscv_func.FuncOp) -> None:
         for op in reversed(block.ops):
             self.process_operation(op)
 
+        if add_regalloc_stats:
+            num_preallocated_int, num_preallocated_float = count_reg_types(preallocated)
+            num_allocated_int, num_allocated_float = count_reg_types(
+                val.type
+                for op in block.walk()
+                for vals in (op.results, op.operands)
+                for val in vals
+            )
+            stats = {
+                "preallocated_float": num_preallocated_float,
+                "preallocated_int": num_preallocated_int,
+                "allocated_float": num_allocated_float,
+                "allocated_int": num_allocated_int,
+            }
+
+            stats_str = json.dumps(stats)
+
+            Rewriter.insert_op(
+                riscv.CommentOp(f"Regalloc stats: {stats_str}"),
+                InsertPoint.before(func),
+            )
+
 
 def _live_ins_per_block(
     block: Block, acc: dict[Block, OrderedSet[SSAValue]]

xdsl/transforms/riscv_register_allocation.py

@@ -30,6 +30,11 @@ class RISCVRegisterAllocation(ModulePass):
     exclude_snitch_reserved: bool = True
     """Excludes floating-point registers that are used by the Snitch ISA extensions."""
 
+    add_regalloc_stats: bool = False
+    """
+    Inserts a comment with register allocation info in the IR.
+    """
+
     def apply(self, ctx: MLContext, op: ModuleOp) -> None:
         allocator_strategies = {
             "LivenessBlockNaive": RegisterAllocatorLivenessBlockNaive,
@@ -54,4 +59,6 @@ def apply(self, ctx: MLContext, op: ModuleOp) -> None:
                     allocator.available_registers.limit_registers(self.limit_registers)
                 allocator.exclude_preallocated = self.exclude_preallocated
                 allocator.exclude_snitch_reserved = self.exclude_snitch_reserved
-                allocator.allocate_func(inner_op)
+                allocator.allocate_func(
+                    inner_op, add_regalloc_stats=self.add_regalloc_stats
+                )

xdsl/transforms/test_lower_linalg_to_snitch.py

@@ -70,7 +70,7 @@
     canonicalize.CanonicalizePass(),
     riscv_scf_loop_range_folding.RiscvScfLoopRangeFoldingPass(),
     canonicalize.CanonicalizePass(),
-    riscv_register_allocation.RISCVRegisterAllocation(),
+    riscv_register_allocation.RISCVRegisterAllocation(add_regalloc_stats=True),
     canonicalize.CanonicalizePass(),
     convert_riscv_scf_to_riscv_cf.ConvertRiscvScfToRiscvCfPass(),
     canonicalize.CanonicalizePass(),