benchmark: adding ability to use the pooled allocator on benchmarks

src/pipedream/benchmark/common.py

@@ -1196,7 +1196,7 @@ class _Benchmark_Runner:
   ) -> ty.Tuple[ty.List['_ASM_Builder.ASM_Statement'], ty.List[ir.Instruction]]:
     """
       Generate inner benchmark loop.
-      Does *NOT* emit any code into :out:, code is returned an can later be spliced in
+      Does *NOT* emit any code into :out:, code is returned and can later be spliced in
     """
 
     if benchmark.kernel_iterations < 1:
@@ -1222,11 +1222,12 @@ class _Benchmark_Runner:
       insts = list(instructions) * benchmark.unroll_factor
 
       Allocator_Class = {
-        Benchmark_Kind.THROUGHPUT: pseudoalloc.MinimizeDepsPseudoAllocator,
-        Benchmark_Kind.LATENCY:    pseudoalloc.MaximizeDepsPseudoAllocator,
-      }[benchmark.kind]
-
-      alloc               = Allocator_Class(out.free_registers(), memory_reg=memory_reg)
+        (Benchmark_Kind.THROUGHPUT, True ): pseudoalloc.MinimizeDepsPseudoAllocator,
+        (Benchmark_Kind.THROUGHPUT, False): pseudoalloc.MinimizeDepsPooledPseudoAllocator,
+        (Benchmark_Kind.LATENCY,    True ): pseudoalloc.MaximizeDepsPseudoAllocator,
+      }[(benchmark.kind, benchmark.register_pools is None)]
+      alloc               = Allocator_Class(out.arch.make_register_pools(out.free_registers(), benchmark.register_pools),\
+                            memory_reg=memory_reg)
       kernel_instructions = alloc.allocate(insts, out.ir_builder)
 
       out.newline()
@@ -2142,6 +2143,9 @@ class _ASM_Builder:
       self._alloc.free(reg)
 
 
+
+
+
 def glob_instruction_tags(arch_or_instructions: ty.Union[ir.Architecture, ty.Iterable[ir.Instruction]],
                           tags: ty.Iterable[str]) -> ty.Iterable[ir.Instruction]:
   if isinstance(arch_or_instructions, ir.Architecture):

src/pipedream/benchmark/types.py

@@ -51,6 +51,7 @@ class Benchmark_Spec(yaml.YAML_Serializable):
   def __init__(self, name: str, kind: Benchmark_Kind,
                unroll_factor: int, kernel_iterations: int,
                arch: ir.Architecture, instructions: ty.List[ir.Instruction],
+               register_pools: ty.Optional[ty.Dict[ir.Instruction, int]],
                loop_overhead: Loop_Overhead):
     """
       ctor.
@@ -62,6 +63,7 @@ class Benchmark_Spec(yaml.YAML_Serializable):
       :param kernel               list of instructions inside the kernel (before unrolling)
       :param arch                 asm.Architecture for the CPU architecture this benchmark is for
       :param instructions         instructions in this benchmark kernel
+      :param register_pools       relative importance of the instruction in terms of amount of register reserved
       :param loop_overhead        number of instructions dynamically executed to run loop around
                                   benchmark kernel (does not include instructions of kernel itself)
     """
@@ -74,6 +76,7 @@ class Benchmark_Spec(yaml.YAML_Serializable):
     self.kernel_iterations = kernel_iterations
     self.arch              = arch
     self.instructions      = instructions
+    self.register_pools    = register_pools
     self.loop_overhead     = loop_overhead
 
   @classmethod
@@ -81,7 +84,8 @@ class Benchmark_Spec(yaml.YAML_Serializable):
                         kind: Benchmark_Kind,
                         arch: ir.Architecture,
                         instructions: ty.Sequence[ir.Instruction],
-                        name: str=None,
+                        register_pools: ty.Optional[ty.Dict[ir.Instruction, int]] = None,
+                        name: str = None,
                         num_dynamic_instructions: int,
                         unrolled_length: int = None,
                         unroll_factor: int = None):
@@ -136,6 +140,7 @@ class Benchmark_Spec(yaml.YAML_Serializable):
       kernel_iterations = kernel_iterations,
       arch              = arch,
       instructions      = instructions,
+      register_pools    = register_pools,
       loop_overhead     = arch.loop_overhead(kernel_iterations),
     )