[hec-assembler]: Optimize linking at xinst/cinst/minst levels (#126)

* Linking up to SPAD boundary
* Adding unit tests
* Adding integration tests

Author: joserochh

assembler_tools/hec-assembler-tools/assembler/instructions/cinst/csyncm.py

@@ -17,10 +17,6 @@ class Instruction(CInstruction):
         https://github.com/IntelLabs/hec-assembler-tools/blob/master/docsrc/inst_spec/cinst/cinst_csyncm.md
     """
 
-    @classmethod
-    def get_throughput(cls) -> int:
-        return cls._OP_DEFAULT_THROUGHPUT
-
     @classmethod
     def _get_op_name_asm(cls) -> str:
         """

assembler_tools/hec-assembler-tools/assembler/instructions/instruction.py

@@ -117,6 +117,15 @@ class BaseInstruction(CycleTracker):
 
     # Class methods and properties
     # ----------------------------
+
+    @classmethod
+    def get_latency(cls) -> int:
+        return cls._OP_DEFAULT_LATENCY
+
+    @classmethod
+    def get_throughput(cls) -> int:
+        return cls._OP_DEFAULT_THROUGHPUT
+
     @classmethod
     def isa_spec_as_dict(cls) -> dict:
         """Returns attributes as dictionary."""

assembler_tools/hec-assembler-tools/he_link.py

@@ -50,7 +50,7 @@
 from linker.kern_trace.trace_info import TraceInfo
 from linker.linker_run_config import LinkerRunConfig
 from linker.loader import Loader
-from linker.steps import program_linker
+from linker.steps.program_linker import LinkedProgram
 from linker.steps.variable_discovery import check_unused_variables, scan_variables
 
 
@@ -100,6 +100,8 @@ def main(run_config: LinkerRunConfig, verbose_stream=None):
     print("", file=verbose_stream)
     print("Interpreting variable meta information...", file=verbose_stream)
 
+    p_linker = LinkedProgram(keep_spad_boundary=run_config.keep_spad_boundary, keep_hbm_boundary=run_config.keep_hbm_boundary)
+
     # Process kernel DInstructions when using trace file
     program_dinstrs = []
     if run_config.using_trace_file:
@@ -111,7 +113,7 @@ def main(run_config: LinkerRunConfig, verbose_stream=None):
         remap_vars(kernels_info, dinstrs_per_kernel, kernel_ops, verbose_stream)
 
         # Concatenate all mem info objects into one
-        program_dinstrs = program_linker.LinkedProgram.join_dinst_kernels(dinstrs_per_kernel)
+        program_dinstrs = p_linker.join_n_prune_dinst_kernels(dinstrs_per_kernel)
 
         # Write new program memory model to an output file
         if program_info.mem is None:
@@ -121,19 +123,22 @@ def main(run_config: LinkerRunConfig, verbose_stream=None):
     # Initialize memory model
     mem_model = initialize_memory_model(run_config, program_dinstrs, verbose_stream)
 
+    # Preload xinst and pre-process intermediate variables for future cinst pruning
+    p_linker.preload_kernels(kernels_info)
+
     # Discover variables
     print("  Finding all program variables...", file=verbose_stream)
     print("  Scanning", file=verbose_stream)
 
-    scan_variables(kernels_info=kernels_info, mem_model=mem_model, verbose_stream=verbose_stream)
+    scan_variables(p_linker, kernels_info, mem_model, verbose_stream)
 
     check_unused_variables(mem_model)
 
     print(f"    Variables found: {len(mem_model.variables)}", file=verbose_stream)
     print("Linking started", file=verbose_stream)
 
     # Link kernels and generate outputs
-    program_linker.LinkedProgram.link_kernels_to_files(kernels_info, program_info, mem_model, verbose_stream=verbose_stream)
+    p_linker.link_kernels_to_files(kernels_info, program_info, mem_model, verbose_stream=verbose_stream)
 
     # Flush cached kernels
     Loader.flush_cache()
@@ -253,7 +258,25 @@ def parse_args():
         "--no_hbm",
         dest="has_hbm",
         action="store_false",
-        help="If set, this flag tells he_prep there is no HBM in the target chip.",
+        help="If set, this flag tells he_link there is no HBM in the target chip.",
+    )
+    parser.add_argument(
+        "--keep_spad_boundary",
+        action="store_true",
+        help=(
+            "If used along with '--use_trace_file', this flag tells he_link to keep a data boundary "
+            "among kernels in the final cinst output. This can be useful for debugging purposes. "
+            "This flag is ignored if '--keep_hbm_boundary' is set."
+        ),
+    )
+    parser.add_argument(
+        "--keep_hbm_boundary",
+        action="store_true",
+        help=(
+            "If used along with '--use_trace_file', this flag tells he_link to keep a data boundary "
+            "among kernels in the final minst output. This can be useful for debugging purposes. "
+            "This flag will override the '--keep_spad_boundary' flag."
+        ),
     )
     parser.add_argument(
         "--suppress_comments",
@@ -282,6 +305,11 @@ def parse_args():
     if p_args.input_dir == "" and p_args.trace_file:
         p_args.input_dir = os.path.dirname(p_args.trace_file)
 
+    # If no trace file is used boundaries should be kept
+    p_args.keep_hbm_boundary = True if not p_args.using_trace_file else p_args.keep_hbm_boundary
+    # If HBM boundary is kept, SPAD boundary is kept as a consequence
+    p_args.keep_spad_boundary = True if p_args.keep_hbm_boundary else p_args.keep_spad_boundary
+
     # Enforce only if use_trace_file is not set
     if not p_args.using_trace_file:
         if p_args.input_mem_file == "":

assembler_tools/hec-assembler-tools/linker/__init__.py

@@ -214,8 +214,14 @@ def add_variable(self, var_name: str):
                 # with predefined HBM address
                 if var_name in self.__mem_info_fixed_addr_vars:
                     var_info.uses = float("inf")
+                else:
+                    var_info.uses = 0
                 self.hbm.force_allocate(var_info, self.__mem_info_vars[var_name].hbm_address)
+            else:
+                # Variables not explicitly marked in mem file are allocated on demand
+                var_info.hbm_address = -1
             self.variables[var_name] = var_info
+
         var_info.uses += 1
 
     def use_variable(self, var_name: str, kernel: int) -> int:
@@ -234,7 +240,6 @@ def use_variable(self, var_name: str, kernel: int) -> int:
 
         var_info.uses -= 1  # Mark the usage
         var_info.last_kernel_used = kernel
-
         if var_info.hbm_address < 0:
             # Find HBM address for variable
             self.hbm.allocate(var_info)

assembler_tools/hec-assembler-tools/linker/he_link_utils.py

@@ -111,7 +111,7 @@ def update_input_prefixes(kernel_ops, run_config):
 def remap_vars(kernels_info: list[KernelInfo], kernels_dinstrs, kernel_ops, verbose_stream):
     """
     @brief Process kernel DInstructions to remap variables based on kernel operations
-    and update KernelInfo with remap_dict.
+    and update KernelInfo with hbm_remap_dict.
 
     @param kernels_info List of input KernelInfo.
     @param kernels_dinstrs List of kernel DInstructions.
@@ -131,7 +131,7 @@ def remap_vars(kernels_info: list[KernelInfo], kernels_dinstrs, kernel_ops, verb
 
         # Remap dintrs' variables in kernel_dinstrs and return a mapping dict
         var_map = remap_dinstrs_vars(kernel_dinstrs, kernel_op)
-        kernel_info.remap_dict = var_map
+        kernel_info.hbm_remap_dict = var_map
 
 
 def initialize_memory_model(run_config, kernel_dinstrs=None, verbose_stream=None):

assembler_tools/hec-assembler-tools/linker/instructions/cinst/bload.py

@@ -50,15 +50,38 @@ def __init__(self, tokens: list, comment: str = ""):
         @throws ValueError If the number of tokens is invalid or the instruction name is incorrect.
         """
         super().__init__(tokens, comment=comment)
+        self._var_name = tokens[3]
+        # set spad_address to '0' if tokens[3] is a variable name
+        if not tokens[3].isdigit():
+            self.tokens[3] = "0"  # Should be set to SPAD address to write back.
 
     @property
-    def source(self) -> str:
+    def var_name(self) -> str:
         """
-        @brief Name of the source.
-        This is a Variable name when loaded. Should be set to HBM address to write back.
+        @brief Gets the name of the variable.
+
+        This is a Variable name when loaded with no_hbm.
+
+        @return The name of the variable.
+        """
+        return self._var_name
+
+    @var_name.setter
+    def var_name(self, value: str):
+        """
+        @brief Sets the name of the variable.
+
+        @param value The name of the variable to set.
+        """
+        self._var_name = value
+
+    @property
+    def spad_address(self) -> int:
+        """
+        @brief Source SPAD address.
         """
-        return self.tokens[3]
+        return int(self.tokens[3])
 
-    @source.setter
-    def source(self, value: str):
-        self.tokens[3] = value
+    @spad_address.setter
+    def spad_address(self, value: int):
+        self.tokens[3] = str(value)

assembler_tools/hec-assembler-tools/linker/instructions/cinst/bones.py

@@ -50,15 +50,35 @@ def __init__(self, tokens: list, comment: str = ""):
         @throws ValueError If the number of tokens is invalid or the instruction name is incorrect.
         """
         super().__init__(tokens, comment=comment)
+        self._var_name = tokens[2]
+        if not tokens[2].isdigit():
+            self.tokens[2] = "0"  # Should be set to SPAD address to write back.
 
     @property
-    def source(self) -> str:
+    def var_name(self) -> str:
         """
-        @brief Name of the source.
-        This is a Variable name when loaded. Should be set to HBM address to write back.
+        @brief Gets the name of the variable.
+
+        This is a Variable name when loaded and there is no HBM.
+        """
+        return self._var_name
+
+    @var_name.setter
+    def var_name(self, value: str):
+        """
+        @brief Sets the name of the variable.
+
+        @param value The name of the variable to set.
+        """
+        self._var_name = value
+
+    @property
+    def spad_address(self) -> int:
+        """
+        @brief Source SPAD address.
         """
-        return self.tokens[2]
+        return int(self.tokens[2])
 
-    @source.setter
-    def source(self, value: str):
-        self.tokens[2] = value
+    @spad_address.setter
+    def spad_address(self, value: int):
+        self.tokens[2] = str(value)

assembler_tools/hec-assembler-tools/linker/instructions/cinst/cinstruction.py

@@ -61,3 +61,23 @@ def to_line(self) -> str:
             str: The string representation of the instruction, excluding the first token.
         """
         return ", ".join(self.tokens[1:])
+
+    @property
+    def idx(self) -> int:
+        """
+        @brief Gets the instruction index.
+
+        This is the line number in the MInstruction file.
+
+        @return The instruction index.
+        """
+        return int(self.tokens[0])
+
+    @idx.setter
+    def idx(self, value: int):
+        """
+        @brief Sets the instruction index.
+
+        @param value The instruction index to set.
+        """
+        self.tokens[0] = str(value)

assembler_tools/hec-assembler-tools/linker/instructions/cinst/cload.py

@@ -50,15 +50,48 @@ def __init__(self, tokens: list, comment: str = ""):
         @throws ValueError If the number of tokens is invalid or the instruction name is incorrect.
         """
         super().__init__(tokens, comment=comment)
+        if not tokens[3].isdigit():
+            self._var_name = tokens[3]
+            self.tokens[3] = "-1"  # Should be set to SPAD address to write back.
+        else:
+            self._var_name = ""
 
     @property
-    def source(self) -> str:
+    def var_name(self) -> str:
         """
-        @brief Name of the source.
-        This is a Variable name when loaded. Should be set to HBM address to write back.
+        @brief Gets the name of the original source variable.
+
+        This is a Variable name when there is no HBM.
+        """
+        return self._var_name
+
+    @var_name.setter
+    def var_name(self, value: str):
+        self._var_name = value
+
+    @property
+    def spad_address(self) -> int:
+        """
+        @brief Source SPAD address.
+        """
+        return int(self.tokens[3])
+
+    @spad_address.setter
+    def spad_address(self, value: int):
+        self.tokens[3] = str(value)
+
+    @property
+    def register(self) -> str:
+        """
+        @brief Gets the name of the destination register.
         """
-        return self.tokens[3]
+        return self.tokens[2]
 
-    @source.setter
-    def source(self, value: str):
-        self.tokens[3] = value
+    @register.setter
+    def register(self, value: str):
+        """
+        @brief Sets the name of the destination register.
+
+        @param value The name of the register to set.
+        """
+        self.tokens[2] = value

assembler_tools/hec-assembler-tools/linker/instructions/cinst/cstore.py

@@ -50,22 +50,43 @@ def __init__(self, tokens: list, comment: str = ""):
         @throws ValueError If the number of tokens is invalid or the instruction name is incorrect.
         """
         super().__init__(tokens, comment=comment)
+        self._var_name = tokens[2]
+        if not tokens[2].isdigit():
+            self.tokens[2] = "0"  # Should be set to SPAD address to write back.
 
     @property
-    def dest(self) -> str:
+    def var_name(self) -> str:
         """
-        @brief Name of the destination.
-        This is a Variable name when loaded. Should be set to HBM address to write back.
+        @brief Gets the name of the variable.
 
-        @return The destination variable name or address.
+        @return The name of the variable.
         """
-        return self.tokens[2]
+        return self._var_name
 
-    @dest.setter
-    def dest(self, value: str):
+    @var_name.setter
+    def var_name(self, value: str):
+        """
+        @brief Sets the name of the variable.
+
+        @param value The name of the variable to set.
+        """
+        self._var_name = value
+
+    @property
+    def spad_address(self) -> int:
+        """
+        @brief SPAD address of destination.
+        Should be set to HBM address to write back.
+
+        @return The destination variable address.
+        """
+        return int(self.tokens[2])
+
+    @spad_address.setter
+    def spad_address(self, value: int):
         """
         @brief Sets the destination of the instruction.
 
         @param value The new destination value to set.
         """
-        self.tokens[2] = value
+        self.tokens[2] = str(value)