Embed proper location information in GoldenTensors (#1488)

Closes #1300
tenstorrent · Dec 16, 2024 · acea937 · acea937
1 parent 7fae99b
commit acea937
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 14 deletions.
diff --git a/python/test_infra/ttir_builder.py b/python/test_infra/ttir_builder.py
@@ -290,29 +290,30 @@ def eltwise_proxy(
         while len(stack) > 0 and stack[0].filename == cur_filename:
             stack = stack[1:]
 
+        id = self.get_next_global_id()
+        loc = Location.file(stack[0].filename, stack[0].lineno, id)
+
         assert (
             len(stack) > 0
         ), "Top of callstack to builder funcs must be outside this file"
 
         with self._ctx, self._loc:
             output = self.empty(self.get_shape(inputs[0]))
 
-            id = self.get_next_global_id()
-
             op = op_ttir_function(
                 [self._get_type(output)],
                 inputs,
                 [output],
                 self._get_operand_constraint_attr(3),
-                loc=Location.name(str(id)),
+                loc=loc,
             )
 
             goldens = []
             for input in inputs:
                 goldens.append(self._get_golden_tensor(input))
 
             golden = Golden(op_golden_function(*goldens))
-            self.id_golden_map[str(id)] = golden
+            self.id_golden_map[str(loc)] = golden
             self._store_golden(op, golden)
             self._override_golden(output, golden)
 

diff --git a/runtime/tools/python/ttrt/common/golden.py b/runtime/tools/python/ttrt/common/golden.py
@@ -111,17 +111,8 @@ def golden(binary, programContext, opContext):
     print("-----------executing golden comparision-----------")
 
     try:
-        op_debug_str = ttrt.runtime.get_op_debug_str(opContext)
+        loc = ttrt.runtime.get_op_loc_info(opContext)
 
-        # find matching golden tensor based on loc in op debug string
-        match = re.search(r"loc\(([^)]+)\)", op_debug_str)
-
-        if not match:
-            print(f"debug_str={op_debug_str}")
-            print("No location found in debug string - skipping golden comparison")
-            return
-
-        loc = match.group(1).replace('"', "")
         print(f"found location={loc}")
 
         op_golden_tensor = binary.get_debug_info_golden(loc)

diff --git a/runtime/tools/python/ttrt/runtime/__init__.py b/runtime/tools/python/ttrt/runtime/__init__.py
@@ -30,6 +30,7 @@
         memcpy,
         deallocate_tensor,
         WorkaroundEnv,
+        get_op_loc_info,
     )
 except ModuleNotFoundError:
     raise ImportError(