Updates exception throwing logic to correctly exclude decorators

Author: pranavm-nvidia

tripy/docs/pre0_user_guides/02-compiler.md

@@ -52,13 +52,13 @@ fast_geglu(inp).eval()
 In the example above, we assumed `inp` has a static shape of `(1, 2)`.
 Now, let's assume that the shape of `inp` can vary from `(1, 2)` to `(16, 2)`, with `(8, 2)`
 being the shape we'd like to optimize for. To express this constraint to the compiler,
-we can provide the range of shapes to `InputInfo` using `shape=((1, 8, 16), 2)`.
+we can provide the range of shapes to `InputInfo` using `shape=([1, 8, 16], 2)`.
 This indicates to the compiler that the first dimension can vary from 1 to 16,
 and it should optimize for a size of 8.
 
 ```py
 # doc: print-locals out out_change_shape
-inp_info = tp.InputInfo(shape=((1, 8, 16), 2), dtype=tp.float32)
+inp_info = tp.InputInfo(shape=([1, 8, 16], 2), dtype=tp.float32)
 fast_geglu = tp.compile(layer, args=[inp_info])
 out = fast_geglu(inp)
 

tripy/tests/README.md

@@ -10,13 +10,9 @@ The `tests/integration` directory captures the latter group of tests.
 
 You can run all tests locally in the development container by running:
 ```bash
-pytest tests/ -v -n 4 --dist worksteal --ignore tests/performance
-pytest tests/performance -v
+pytest tests/ -v
 ```
 
-Performance tests are run separately because they must run serially to ensure
-accurate measurements.
-
 You can also provide marker arguments to only run specific test cadences
 (see [the test cadence section](#test-cadence) below). For example, to run only
 L0 tests, use:
@@ -26,6 +22,9 @@ pytest tests/ -v -m "not l1 and not manual" -n 4 --dist worksteal --ignore tests
 pytest tests/performance -v -m "not l1 and not manual"
 ```
 
+Note that the L0/L1 tests can be parallelized, which is not necessarily
+true of `manual` tests. In that case, performance tests are run separately
+because they must run serially to ensure accurate measurements.
 
 ## Profiling
 

tripy/tests/common/test_exception.py

@@ -15,11 +15,15 @@
 # limitations under the License.
 #
 
+import re
 from dataclasses import dataclass
 from textwrap import dedent
 
 from tests import helper
-from tripy.common.exception import TripyException, _make_stack_info_message, raise_error
+
+import tripy as tp
+from tripy.common.exception import TripyException, _get_function_file_and_lines, _make_stack_info_message, raise_error
+from tripy.frontend.utils import convert_shape_inputs
 from tripy.utils import StackInfo, get_stack_info
 from tripy.utils.stack_info import SourceInfo
 
@@ -120,3 +124,37 @@ def test_can_determine_column_range(self):
             ).strip()
             in dedent(error_msg).strip()
         )
+
+    def test_convert_shape_inputs_is_excluded(self):
+        filename, start_line, end_line = _get_function_file_and_lines(convert_shape_inputs)
+        tensor = tp.ones((2, 3))
+
+        stack_info = tensor.stack_info
+
+        assert any(
+            frame.file == filename and frame.line >= start_line and frame.line <= end_line for frame in stack_info
+        )
+
+        # Make sure that no extraneous wrapper code is included
+        expected = dedent(
+            r"""
+                --> [a-z_/\.]+:[0-9]+ in full\(\)
+                    |
+                [0-9]+ |     return full_impl\(shape, value, dtype, output_rank\)
+                    |
+
+                --> [a-z_/\.]+:[0-9]+ in ones\(\)
+                    |
+                [0-9]+ |     return full\(shape, 1, dtype\)
+                    |            ^^^^^^^^^^^^^^^^^^^^^ --- required from here
+
+                --> [a-z_/\.]+:[0-9]+ in test_convert_shape_inputs_is_excluded\(\)
+                    |
+                [0-9]+ |         tensor = tp.ones\(\(2, 3\)\)
+                    |                  ^^^^^^^^^^^^^^^ --- required from here
+
+                """
+        ).strip()
+
+        actual = _make_stack_info_message(stack_info, enable_color=False)
+        assert re.search(expected, actual) is not None

tripy/tripy/common/exception.py

@@ -101,22 +101,30 @@ def apply_color(inp, color):
     return frame_info
 
 
+def _get_function_file_and_lines(func):
+    filename = inspect.getsourcefile(func)
+    lines, start_line = inspect.getsourcelines(func)
+    return filename, start_line, start_line + len(lines)
+
+
 def _make_stack_info_message(stack_info: "utils.StackInfo", enable_color: bool = True) -> Optional[str]:
+
     from tripy.frontend.utils import convert_inputs_to_tensors, convert_shape_inputs
 
     EXCLUDE_FUNCTIONS = [convert_inputs_to_tensors, convert_shape_inputs]
 
-    def should_exclude(frame):
-        for func in EXCLUDE_FUNCTIONS:
-            filename = inspect.getsourcefile(func)
-            lines, start_line = inspect.getsourcelines(func)
+    exclude_file_lines = {}  # Maps filenames to ranges of lines that should be ignored.
+    for func in EXCLUDE_FUNCTIONS:
+        filename, start_line, end_line = _get_function_file_and_lines(func)
+
+        exclude_file_lines[filename] = (start_line, end_line)
 
-            if frame.file != filename:
-                return False
+    def should_exclude(frame):
+        if frame.file not in exclude_file_lines:
+            return False
 
-            if frame.line < start_line or frame.line > (start_line + len(lines)):
-                return False
-            return True
+        start_line, end_line = exclude_file_lines[frame.file]
+        return frame.line >= start_line and frame.line <= end_line
 
     frame_strs = []
     num_frames_printed = 0

tripy/tripy/utils/stack_info.py

@@ -108,7 +108,7 @@ def get_stack_info(include_code_index: int = None) -> StackInfo:
             column_range=None,
         )
         if source_info.module == tripy.function_registry.__name__ and source_info.function == "wrapper":
-            source_info._dispatch_target = frame.f_locals.get("key", "")
+            source_info._dispatch_target = frame.f_locals["key"]
 
         try:
             # In Python 3.11, frames contain column offset information.