[hec-assembler]: Optimize linking at cinst/minst levels

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

@@ -21,6 +21,10 @@ class Instruction(CInstruction):
         spad_src (int): SPAD address of the metadata word to load.
     """
 
+    @classmethod
+    def get_latency(cls) -> int:
+        return cls._OP_DEFAULT_LATENCY
+
     @classmethod
     def _get_op_name_asm(cls) -> str:
         """

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

@@ -19,6 +19,10 @@ class Instruction(CInstruction):
         spad_src (int): SPAD address of the metadata variable to load.
     """
 
+    @classmethod
+    def get_latency(cls) -> int:
+        return cls._OP_DEFAULT_LATENCY
+
     @classmethod
     def _get_op_name_asm(cls) -> str:
         """

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

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

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

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

assembler_tools/hec-assembler-tools/he_link.py

@@ -44,7 +44,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
 
 
@@ -94,6 +94,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=True, keep_hbm_boundary=run_config.keep_hbm_boundary)
+
     # Process kernel DInstructions when using trace file
     program_dinstrs = []
     if run_config.using_trace_file:
@@ -105,7 +107,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:
@@ -119,15 +121,15 @@ def main(run_config: LinkerRunConfig, verbose_stream=None):
     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()
@@ -247,7 +249,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",
@@ -276,6 +296,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/cload.py

@@ -50,15 +50,46 @@ 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]
+        if not tokens[3].isdigit():
+            self.tokens[3] = "-1"  # 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 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)