Adds a warning if a tensor is evaluated while compiling

In the case where a tensor is evaluated while compiling and then used in the
computation graph, we throw an error. However, there are cases where the result
of evaluation could make a round-trip through non-Tripy code, in which case we lose
visibility. An example of this, assuming `a` and `b` are tensors, is:
```py
b = b + int(a.shape[0])
```

Here, `a.shape[0]` will be evaluated due to the `int` conversion and then
used by the add operation with `b`. Because it was a Python integer in between,
Tripy has no way to track that it actually came from an evaluated tensor.

Hence, this change prints warnings in these ambiguous cases when a tensor
is evaluated while compiling. Here's an example of the warning messages:

```
[W] Tensor was evaluated while compiling which may cause unexpected behavior in the executable.
For example, this could cause values to be baked into the executable or dynamic shapes to become static.
If the result of the evaluation is not being used by other operations, you can safely ignore this warning.
[W] Note: Tensor was evaluated while compiling here:

--> /tripy/tests/backend/api/test_compile.py:174 in func()
      |
  174 |             print(a.shape)
      |             ^^^^^^^^^^^^^^

[2, 3]
[W] Note: Tensor was evaluated while compiling here:

--> /tripy/tests/backend/api/test_compile.py:176 in func()
      |
  176 |             c = a - int(a.shape[0])
      |                     ^^^^^^^^^^^^^^^

[W] Note: Tensor was evaluated while compiling here:

--> /tripy/tests/backend/api/test_compile.py:177 in func()
      |
  177 |             print(c)
      |             ^^^^^^^^
```

Author: pranavm-nvidia

tripy/tests/backend/api/test_compile.py

@@ -168,16 +168,21 @@ def func(a):
         with helper.raises(tp.TripyException, match="Cannot evaluate a tensor while compiling."):
             tp.compile(func, args=[tp.InputInfo((2, 3), dtype=tp.float32)])
 
-    def test_allow_eval_if_tensor_unused_in_compile(self):
+    def test_allow_eval_if_tensor_unused_in_compile(self, capsys):
         # If the tensor is not actually used in the computation graph then we don't care if it's eval'd.
         def func(a):
             print(a.shape)
 
-            c = a - 1
+            c = a - int(a.shape[0])
             print(c)
             return a
 
-        tp.compile(func, args=[tp.InputInfo((2, 3), dtype=tp.float32)])
+        tp.compile(func, args=[tp.InputInfo((2, 3), dtype=tp.int32)])
+        out, _ = capsys.readouterr()
+        print(f"\n{out}")
+
+        # Ensure that a warning is printed for each evaluation (2 prints + int).
+        assert out.count("Tensor was evaluated while compiling here:") == 3
 
     def test_allow_eval_for_non_input_to_compile(self):
         # We should allow non-inputs to be evaluated.

tripy/tests/common/test_exception.py

@@ -22,7 +22,7 @@
 from tests import helper
 
 import tripy as tp
-from tripy.common.exception import TripyException, _get_function_file_and_lines, _make_stack_info_message, raise_error
+from tripy.common.exception import TripyException, _get_function_file_and_lines, str_from_stack_info, raise_error
 from tripy.frontend.utils import convert_to_tensors
 from tripy.utils import StackInfo, get_stack_info
 from tripy.utils.stack_info import SourceInfo
@@ -112,7 +112,7 @@ def test_can_determine_column_range(self):
             ]
         )
 
-        error_msg = _make_stack_info_message(stack_info, enable_color=False)
+        error_msg = str_from_stack_info(stack_info, enable_color=False)
         assert (
             dedent(
                 """
@@ -156,5 +156,5 @@ def test_convert_to_tensors_is_excluded(self):
                 """
         ).strip()
 
-        actual = _make_stack_info_message(stack_info, enable_color=False)
+        actual = str_from_stack_info(stack_info, enable_color=False)
         assert re.search(expected, actual) is not None

tripy/tests/helper.py

@@ -34,7 +34,7 @@
 
 import tripy as tp
 from tripy import utils
-from tripy.common.exception import _make_stack_info_message
+from tripy.common.exception import str_from_stack_info
 from tripy.frontend import Tensor
 from tripy.frontend.trace import Trace
 
@@ -74,7 +74,7 @@ def raises(ExcType: type, match: Optional[str] = None, has_stack_info_for: Seque
     has_stack_info_for = has_stack_info_for or []
     for tensor in has_stack_info_for:
         # Stack info is indented since it's part of the `details` block in `raise_error`
-        expected_stack_info = indent(_make_stack_info_message(tensor.stack_info).strip(), " " * 4)
+        expected_stack_info = indent(str_from_stack_info(tensor.stack_info).strip(), " " * 4)
         assert expected_stack_info in error_msg, f"Missing stack information for tensor:\n{expected_stack_info}"
 
 

tripy/tripy/common/exception.py

@@ -107,7 +107,7 @@ def _get_function_file_and_lines(func):
     return filename, start_line, start_line + len(lines)
 
 
-def _make_stack_info_message(stack_info: "utils.StackInfo", enable_color: bool = True) -> Optional[str]:
+def str_from_stack_info(stack_info: "utils.StackInfo", enable_color: bool = True) -> Optional[str]:
     from tripy.frontend.utils import convert_to_tensors
 
     EXCLUDE_FUNCTIONS = [convert_to_tensors]
@@ -187,9 +187,9 @@ def raise_error(summary: str, details: List[Any] = []):
     for detail in details:
         stack_info_message = None
         if hasattr(detail, "stack_info"):
-            stack_info_message = _make_stack_info_message(detail.stack_info)
+            stack_info_message = str_from_stack_info(detail.stack_info)
         elif isinstance(detail, utils.StackInfo):
-            stack_info_message = _make_stack_info_message(detail)
+            stack_info_message = str_from_stack_info(detail)
 
         if stack_info_message is not None:
             detail_msg += stack_info_message

tripy/tripy/frontend/module/module.py

@@ -20,7 +20,7 @@
 from typing import Any, Dict, Iterator, List, Tuple, Union, Set, Sequence, TypeVar
 
 from tripy import export, utils
-from tripy.common.exception import raise_error, _make_stack_info_message
+from tripy.common.exception import raise_error, str_from_stack_info
 from tripy.frontend.module.parameter import Parameter
 from tripy.logging import logger
 
@@ -111,7 +111,7 @@ def __setattr__(self, name: str, value: Any) -> None:
             ):
                 stack_info = utils.get_stack_info()
                 stack_info.fetch_source_code()
-                stack_info_msg = _make_stack_info_message(stack_info)
+                stack_info_msg = str_from_stack_info(stack_info)
 
                 logger.warning(
                     "A container of mixed types will not be registered with this module's state_dict()."

tripy/tripy/frontend/tensor.py

@@ -26,10 +26,11 @@
 from tripy import export, utils
 from tripy.backend.mlir import memref
 from tripy.common import datatype
-from tripy.common.exception import raise_error
+from tripy.common.exception import raise_error, str_from_stack_info
 from tripy.frontend.ops.registry import TENSOR_METHOD_REGISTRY
 from tripy.frontend.trace.ops import Storage
 from tripy.frontend.trace.tensor import TraceTensor
+from tripy.logging.logger import logger
 from tripy.utils.stack_info import StackInfo
 
 
@@ -198,6 +199,18 @@ def eval(self) -> runtime.MemRefValue:
         self.trace_tensor.device = flat_ir.outputs[0].device
 
         self.trace_tensor.eval_stack_info = utils.get_stack_info()
+        if self.trace_tensor.is_compile_tracer:
+            logger.warning(
+                f"Tensor was evaluated while compiling which may cause unexpected behavior in the executable.\n"
+                f"For example, this could cause values to be baked into the executable or dynamic shapes to become static.\n"
+                f"If the result of the evaluation is not being used by other operations, you can safely ignore this warning.",
+                mode="once",
+            )
+            logger.warning(
+                f"Note: Tensor was evaluated while compiling here: {str_from_stack_info(self.trace_tensor.eval_stack_info)}",
+                mode="once",
+            )
+
         return data
 
     def tolist(self):

tripy/tripy/logging/logger.py

@@ -105,6 +105,7 @@ def __init__(self) -> None:
         self._indentation = 0
         self.verbosity: Union[str, Set[str], Dict[str, str], Dict[str, Set[str]]] = "info"
         self.enable_color = True
+        self._already_logged_hashes = set()
 
     @property
     def verbosity(self):
@@ -167,7 +168,9 @@ def indent(self, level: int = 4):
         finally:
             self._indentation = old_indentation
 
-    def log(self, message: Union[str, Callable[[], str]], verbosity: str, stack_depth: int = 2) -> None:
+    def log(
+        self, message: Union[str, Callable[[], str]], verbosity: str, mode: str = "each", stack_depth: int = 2
+    ) -> None:
         """
         Logs a message to standard output.
 
@@ -178,6 +181,9 @@ def log(self, message: Union[str, Callable[[], str]], verbosity: str, stack_dept
             message: The message to log. This can be provided as a callable in which case it will not
                 be called unless the message actually needs to be logged.
             verbosity: The verbosity at which to log this message.
+            mode: Indicates when or how to log the message. Available modes are:
+                - "each": Log the message each time.
+                - "once": Only log a message the first time it is seen.
             stack_depth: The stack depth to use when determining which file the message is being logged from.
         """
         assert (
@@ -202,6 +208,7 @@ def get_rel_file_path():
             return module_path(file_path)
 
         def should_log():
+
             path = None
             # Don't actually need to get the path if there are no non-default entries in the trie.
             if self.verbosity.has_non_default_entries:
@@ -211,6 +218,12 @@ def should_log():
         if not should_log():
             return
 
+        if mode == "once":
+            message_hash = hash(message)
+            if message_hash in self._already_logged_hashes:
+                return
+            self._already_logged_hashes.add(message_hash)
+
         if callable(message):
             message = message()
 

tripy/tripy/utils/ast.py

@@ -16,7 +16,7 @@
 #
 
 import ast
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Set
 
 from tripy.utils.result import Result
 from tripy.utils.stack_info import SourceInfo
@@ -40,14 +40,25 @@ def get_parsed_ast(code: str) -> Result[Tuple[str, int]]:
     return Result.ok((parsed_ast, indentation))
 
 
-def get_callee_func_name(callee: SourceInfo):
-    callee_name = callee.function
+def get_callee_func_name_candidates(callee: SourceInfo) -> Set[str]:
     # Some functions (e.g. tensor methods) are routed through a function registry.
     # We don't actually care about the dispatch function, so we look at the `key`
     # to determine which underlying method we're actually calling.
     if callee._dispatch_target:
-        callee_name = callee._dispatch_target
-    return callee_name
+        candidates = {callee._dispatch_target}
+    else:
+        candidates = {callee.function}
+
+    # Some methods are called by other builtins:
+    SPECIAL_METHODS = {
+        "__repr__": {"repr", "print"},
+        "__str__": {"str"},
+        "__int__": {"int"},
+        "__bool__": {"bool"},
+    }
+    candidates.update(SPECIAL_METHODS.get(callee.function, set()))
+
+    return candidates
 
 
 def get_ast_node_func_name(node) -> Optional[str]:
@@ -144,7 +155,7 @@ def index_into_expr(node: ast.expr, index: int) -> ast.expr:
 # Grab column offsets for a given frame based on information from its callee.
 # This method is not perfect and is not required for Python 3.11+, where frames include column offsets.
 def get_candidate_column_offsets(cur_frame: SourceInfo, callee: SourceInfo) -> List[Tuple[int, int]]:
-    callee_name = get_callee_func_name(callee)
+    candidate_callee_names = get_callee_func_name_candidates(callee)
 
     candidate_column_offsets = []
 
@@ -165,8 +176,8 @@ def get_candidate_column_offsets(cur_frame: SourceInfo, callee: SourceInfo) -> L
 
         def check_name_matches():
             # We need special checking for __init__ methods since the AST node will just be the class name, e.g. `Tensor`.
-            if callee_name != "__init__":
-                return ast_node_name == callee_name
+            if "__init__" not in candidate_callee_names:
+                return ast_node_name in candidate_callee_names
 
             # We hardcode names of some common classes here to avoid creating an import dependency:
             if ast_node_name in {"Tensor"}: