From d36b297048ac7a9e328587eeaf10b7dc0fb0a938 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Tue, 30 Jul 2024 09:49:23 +0200 Subject: [PATCH] backend: (riscv) add option to insert regalloc stats (#2958) Inserts a little json that we can parse in the experiments repo. --- .../register-allocation/preallocated.mlir | 9 ++-- .../riscv-backend-paper/bottom_up.mlir | 9 ++++ .../projects/riscv-backend-paper/nsnet.mlir | 2 + xdsl/backend/riscv/register_allocation.py | 51 +++++++++++++++++-- xdsl/transforms/riscv_register_allocation.py | 9 +++- .../transforms/test_lower_linalg_to_snitch.py | 2 +- 6 files changed, 73 insertions(+), 9 deletions(-) diff --git a/tests/filecheck/backend/riscv/register-allocation/preallocated.mlir b/tests/filecheck/backend/riscv/register-allocation/preallocated.mlir index 0679f60871..41e49a0b0d 100644 --- a/tests/filecheck/backend/riscv/register-allocation/preallocated.mlir +++ b/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 - %3 = riscv.fcvt.s.w %0 : (!riscv.reg) -> !riscv.freg + %3 = riscv.fcvt.s.w %0 : (!riscv.reg) -> !riscv.freg %4 = riscv.fcvt.s.w %1 : (!riscv.reg) -> !riscv.freg - %5 = riscv.fadd.s %3, %4 : (!riscv.freg, !riscv.freg) -> !riscv.freg + %5 = riscv.fadd.s %3, %4 : (!riscv.freg, !riscv.freg) -> !riscv.freg %2 = riscv.add %0, %1 : (!riscv.reg, !riscv.reg) -> !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 // LIVE-BNAIVE-NEXT: %1 = riscv.li 5 : !riscv.reg // LIVE-BNAIVE-NEXT: %2 = riscv.fcvt.s.w %0 : (!riscv.reg) -> !riscv.freg // LIVE-BNAIVE-NEXT: %3 = riscv.fcvt.s.w %1 : (!riscv.reg) -> !riscv.freg -// LIVE-BNAIVE-NEXT: %4 = riscv.fadd.s %2, %3 : (!riscv.freg, !riscv.freg) -> !riscv.freg +// LIVE-BNAIVE-NEXT: %4 = riscv.fadd.s %2, %3 : (!riscv.freg, !riscv.freg) -> !riscv.freg // LIVE-BNAIVE-NEXT: %5 = riscv.add %0, %1 : (!riscv.reg, !riscv.reg) -> !riscv.reg // LIVE-BNAIVE-NEXT: riscv_func.return // LIVE-BNAIVE-NEXT: } diff --git a/tests/filecheck/projects/riscv-backend-paper/bottom_up.mlir b/tests/filecheck/projects/riscv-backend-paper/bottom_up.mlir index 8a9a0055fa..0baf5883c3 100644 --- a/tests/filecheck/projects/riscv-backend-paper/bottom_up.mlir +++ b/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 diff --git a/tests/filecheck/projects/riscv-backend-paper/nsnet.mlir b/tests/filecheck/projects/riscv-backend-paper/nsnet.mlir index 8fd60ead18..60fffb9f6d 100644 --- a/tests/filecheck/projects/riscv-backend-paper/nsnet.mlir +++ b/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 diff --git a/xdsl/backend/riscv/register_allocation.py b/xdsl/backend/riscv/register_allocation.py index e835c9545a..89ab399c96 100644 --- a/xdsl/backend/riscv/register_allocation.py +++ b/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]] diff --git a/xdsl/transforms/riscv_register_allocation.py b/xdsl/transforms/riscv_register_allocation.py index 939f0f1c01..48de247f81 100644 --- a/xdsl/transforms/riscv_register_allocation.py +++ b/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 + ) diff --git a/xdsl/transforms/test_lower_linalg_to_snitch.py b/xdsl/transforms/test_lower_linalg_to_snitch.py index 17dfef0d36..3ac918c7ed 100644 --- a/xdsl/transforms/test_lower_linalg_to_snitch.py +++ b/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(),