From 909f4434b246ec1c78794b7dc1b7aa772dbe8239 Mon Sep 17 00:00:00 2001 From: christopherngutierrez Date: Mon, 22 Sep 2025 09:59:06 -0700 Subject: [PATCH] modified assembler to enable command line controlled coloring strategy --- .../assembler/stages/prep/preprocessor.py | 8 ++++++-- .../hec-assembler-tools/he_prep.py | 20 ++++++++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/assembler_tools/hec-assembler-tools/assembler/stages/prep/preprocessor.py b/assembler_tools/hec-assembler-tools/assembler/stages/prep/preprocessor.py index 2ebeb57..93abf46 100644 --- a/assembler_tools/hec-assembler-tools/assembler/stages/prep/preprocessor.py +++ b/assembler_tools/hec-assembler-tools/assembler/stages/prep/preprocessor.py @@ -150,7 +150,9 @@ def reduce_var_deps_by_var(mem_model: MemoryModel, insts_list: list, var_name: s last_pos += 1 -def assign_register_banks_to_vars(mem_model: MemoryModel, insts_list: list, use_bank0: bool) -> str: +def assign_register_banks_to_vars( + mem_model: MemoryModel, insts_list: list, use_bank0: bool, verbose=False, strategy="largest_first", interchange=False +) -> str: """ Assigns register banks to variables using vertex coloring graph algorithm. @@ -172,6 +174,8 @@ def assign_register_banks_to_vars(mem_model: MemoryModel, insts_list: list, use_ should add corresponding `move` instructions to move variables from bank 0 to correct bank. verbose (bool, optional): If True, prints verbose output. Defaults to False. + strategy (str, optional): Strategy for greedy coloring algorithm. Defaults to "largest_first". + interchange (bool, optional): Whether to use interchange in greedy coloring. Defaults to False. Raises: ValueError: Thrown for these cases: @@ -193,7 +197,7 @@ def assign_register_banks_to_vars(mem_model: MemoryModel, insts_list: list, use_ # Extract the dependency graph for variables dep_graph_vars, dest_names, source_names = dependency_graph_for_vars(insts_list) only_sources = source_names - dest_names # Find which variables are ever only used as sources - color_dict = nx.greedy_color(dep_graph_vars) # Do coloring + color_dict = nx.greedy_color(dep_graph_vars, strategy=strategy, interchange=interchange) # Do coloring needs_reduction = False for var_name, bank in color_dict.items(): diff --git a/assembler_tools/hec-assembler-tools/he_prep.py b/assembler_tools/hec-assembler-tools/he_prep.py index ef20c24..591ce32 100644 --- a/assembler_tools/hec-assembler-tools/he_prep.py +++ b/assembler_tools/hec-assembler-tools/he_prep.py @@ -69,6 +69,9 @@ def main(args): - output_file_name - mem_file - verbose + optional: + - strategy + - interchange """ GlobalConfig.debugVerbose = args.verbose @@ -93,7 +96,9 @@ def main(args): num_input_instr: int = len(insts_listing) # track number of instructions in input kernel if args.verbose > 0: print("Assigning register banks to variables...") - preprocessor.assign_register_banks_to_vars(hec_mem_model, insts_listing, use_bank0=False) + preprocessor.assign_register_banks_to_vars( + hec_mem_model, insts_listing, use_bank0=False, verbose=args.verbose, strategy=args.strategy, interchange=args.interchange + ) # Determine output file name if not args.output_file_name: @@ -195,6 +200,17 @@ def parse_args(): "Increase level of verbosity by specifying flag multiple times, e.g. -vv" ), ) + parser.add_argument( + "--strategy", + default="largest_first", + help="Strategy for greedy coloring algorithm. Defaults to 'largest_first'.", + ) + parser.add_argument( + "--interchange", + action="store_true", + default=False, + help="Whether to use interchange in greedy coloring. Defaults to False.", + ) p_args = parser.parse_args() p_args.split_on = bool(p_args.split_inst_limit != float("inf") or p_args.split_vars_limit != float("inf")) if p_args.split_on: @@ -223,6 +239,8 @@ def parse_args(): print(f"Split Inst Limit: {args.split_inst_limit}") print(f"Split Vars Limit: {args.split_vars_limit}") print(f"Split On: {args.split_on}") + print(f"Graph Coloring Strategy: {args.strategy}") + print(f"Graph Coloring Interchange: {args.interchange}") main(args)