From ba02a5d2bc3e489e97813761d593bd7bb1057ad9 Mon Sep 17 00:00:00 2001 From: Collin Tod Date: Mon, 16 Dec 2024 15:33:12 -0600 Subject: [PATCH] Fix golden location bug (#1593) This change fixes a bug introduced by the previous attempt to use real file names into ops built via python bindings. The problem stemmed from using the `Location.file` constructor, instead of manually constructing a location string and using the `Location.name` constructor instead. This is a temporary fix, and in the future, pure locations shouldn't be used as op UUIDs as they can be inherited by their decomposed ops in some circumstances, thus making them non-unique. --- python/test_infra/test_utils.py | 3 +- python/test_infra/ttir_builder.py | 77 ++++++++++++++-------- runtime/tools/python/ttrt/common/golden.py | 12 +--- test/python/golden/test_ttir_ops.py | 2 +- 4 files changed, 53 insertions(+), 41 deletions(-) diff --git a/python/test_infra/test_utils.py b/python/test_infra/test_utils.py index 09e86db97..da1957b7f 100644 --- a/python/test_infra/test_utils.py +++ b/python/test_infra/test_utils.py @@ -4,9 +4,8 @@ import os import inspect -from typing import Callable, Dict, List, Optional +from typing import Callable, List, Optional -import torch from ttmlir.dialects import func from ttmlir.ir import * from ttmlir.passes import ( diff --git a/python/test_infra/ttir_builder.py b/python/test_infra/ttir_builder.py index 9c832d014..471c07ca7 100644 --- a/python/test_infra/ttir_builder.py +++ b/python/test_infra/ttir_builder.py @@ -3,12 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import inspect +import inspect from dataclasses import dataclass from typing import List, Optional, Union, Tuple, Callable, Dict from ttmlir.ir import * -from ttmlir.dialects import ttir, tt, func, tensor +from ttmlir.dialects import ttir, tt, tensor from ttmlir.passes import create_golden_tensor, DataType import torch @@ -17,7 +17,50 @@ Operand = Union[Value, OpView, Operation] # Convenience alias for shape -Shape = Union[List[int], Tuple[int]] +Shape = Union[List[int], Tuple[int, ...]] + + +def get_loc_of_extra_file_callee(id: int = 0) -> Location: + """When called, this function returns a `Location` referring to first + callee outside the file of the caller of this function. E.G., if a function + in `foo.py` called a function in `bar.py` that then called this function, + the location would be pointing to the call in `foo.py`. + + NOTE: this location is _NOT_ in the form of + {filename}:{line_number}:{col_number}, but instead in the form: + {filename}:{line_number}:id({id}), where id is supplied to this function as + a disambiguator for calls that happen on the same line + + Arguments + --------- + + id : int + An optional variable that defaults to 0 to be appended to the location, + disambiguating calls on the same line. + + Returns + ------- + + A `Location` referring to the first extra file callee of the caller of this function + + """ + + stack = inspect.stack() + + # find the innermost frame outside of this file + caller_filename = stack[1].filename + + while len(stack) > 0 and stack[0].filename == caller_filename: + stack = stack[1:] + + assert ( + len(stack) > 0 + ), "Top of callstack to builder funcs must be outside the caller's file" + + # FIXME: this should be a `Location.file`, but for some reason it causes + # strange decomposition inheritance behaviour that breaks using this as + # a key into the golden map + return Location.name(f"{stack[0].filename}:{str(stack[0].lineno)}:id({str(id)})") @dataclass(frozen=True) @@ -251,40 +294,20 @@ def eltwise_proxy( inputs: List[Operand], ) -> OpView: - # Snoop the location of the first caller outside of this file to - # annotate the MLIR with. NOTE that this location is _NOT_ row:col, but - # instead row:id, where id is a unique id given to all calls to builder - # funcs. See `get_next_global_id` for more details - stack = inspect.stack() - - # find the innermost frame outside of this file - cur_filename = stack[0].filename - - while len(stack) > 0 and stack[0].filename == cur_filename: - stack = stack[1:] - - assert ( - len(stack) > 0 - ), "Top of callstack to builder funcs must be outside this file" + id = self.get_next_global_id() + loc = get_loc_of_extra_file_callee(id=id) 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], - loc=Location.name(str(id)), - ) + op = op_ttir_function([self._get_type(output)], inputs, [output], 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 index 055d4c824..847942615 100644 --- a/runtime/tools/python/ttrt/common/golden.py +++ b/runtime/tools/python/ttrt/common/golden.py @@ -117,17 +117,7 @@ def golden_partial_function( print("-----------executing golden comparision-----------") try: - op_debug_str = ttrt.runtime.get_op_debug_str(op_context) - - # 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('"', "") + loc = ttrt.runtime.get_op_loc_info(op_context) print(f"found location={loc}") op_golden_tensor = binary.get_debug_info_golden(loc) diff --git a/test/python/golden/test_ttir_ops.py b/test/python/golden/test_ttir_ops.py index aa18e1036..e693196f5 100644 --- a/test/python/golden/test_ttir_ops.py +++ b/test/python/golden/test_ttir_ops.py @@ -2,7 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 -# RUN: %python %s +# RUN: SYSTEM_DESC_PATH=%system_desc_path% %python %s import inspect