From acea937f594dbf7f29a8524ad053549fa9d3412b Mon Sep 17 00:00:00 2001
From: Collin Tod <collintod@tenstorrent.com>
Date: Tue, 10 Dec 2024 10:51:34 -0600
Subject: [PATCH] Embed proper location information in GoldenTensors (#1488)

Closes #1300
---
 python/test_infra/ttir_builder.py             |  9 +++++----
 runtime/tools/python/ttrt/common/golden.py    | 11 +----------
 runtime/tools/python/ttrt/runtime/__init__.py |  1 +
 3 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/python/test_infra/ttir_builder.py b/python/test_infra/ttir_builder.py
index a6a302d4a6..0921db3fca 100644
--- a/python/test_infra/ttir_builder.py
+++ b/python/test_infra/ttir_builder.py
@@ -290,6 +290,9 @@ 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"
@@ -297,14 +300,12 @@ def eltwise_proxy(
         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 = []
@@ -312,7 +313,7 @@ def eltwise_proxy(
                 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
index 59adc97a19..7bab624f65 100644
--- 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
index 0376c07b58..5f1841ebeb 100644
--- 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(