Updates MLIR to include layer metadata (#669)

Updates our Trace->MLIR translation to include metadata for each layer
that includes stack information and details about the Trace operation.

Author: pranavm-nvidia

tripy/nvtripy/common/exception.py

@@ -15,8 +15,6 @@
 # limitations under the License.
 #
 
-import inspect
-from dataclasses import dataclass
 from textwrap import indent
 from typing import Any, List, Optional, Tuple
 
@@ -84,12 +82,14 @@ def apply_color(inp, color):
     return frame_info
 
 
-def str_from_stack_info(stack_info: "utils.stack_info.StackInfo", enable_color: bool = True) -> Optional[str]:
+def str_from_stack_info(
+    stack_info: "utils.stack_info.StackInfo", enable_color: bool = True, fetch_source_code: bool = True
+) -> Optional[str]:
     from nvtripy.frontend.module import module
 
     def should_exclude(source_info):
         return (
-            source_info.code is None
+            (fetch_source_code and source_info.code is None)
             or source_info.module in utils.stack_info.get_module_names_to_exclude_from_stack_info()
             # Exclude module.__call__ since it just invokes forward and clutters the stack trace
             or (source_info.module == module.__name__ and source_info.function == "__call__")
@@ -98,7 +98,8 @@ def should_exclude(source_info):
     frame_strs = []
     num_frames_printed = 0
 
-    stack_info.fetch_source_code()
+    if fetch_source_code:
+        stack_info.fetch_source_code()
     for index, source_info in enumerate(stack_info):
         if should_exclude(source_info):
             continue

tripy/nvtripy/trace/trace.py

@@ -28,7 +28,7 @@
     map_error_to_user_code_and_raise,
     redirect_stderr,
 )
-from nvtripy.common.exception import raise_error
+from nvtripy.common.exception import raise_error, str_from_stack_info
 from nvtripy.logging import logger
 from nvtripy.trace.tensor import TraceTensor
 from nvtripy.trace.utils import topological_sort
@@ -178,6 +178,25 @@ def num_known_dims(ranked_tensor_type):
                                 if num_known_dims(output_type) >= num_known_dims(mlir_output_op.type):
                                     mlir_output_op.set_type(output_type)
 
+                                mlir_output_op.owner.attributes["metadata"] = ir.StringAttr.get(
+                                    f"<TraceOp: {op}, Stack Info: "
+                                    + ", ".join(
+                                        (
+                                            str_from_stack_info(
+                                                # include_code_index points to the first "useful" frame, i.e. usually the API
+                                                # that the user calls, or a few frames below that:
+                                                out.stack_info[out.stack_info.include_code_index or 0 :],
+                                                enable_color=False,
+                                                fetch_source_code=False,
+                                            )
+                                            if out.stack_info
+                                            else ""
+                                        ).replace("\n", " ")
+                                        for out in op.outputs
+                                    )
+                                    + " >, "
+                                )
+
                             mlir_ops.update(zip([out.name for out in op.outputs], mlir_output_ops))
 
                         func_dialect.ReturnOp([mlir_ops[o.name] for o in self.outputs])