New Feature - Pass coloring strategy via command line for he_prep

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():

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)