From f80a6ac301622ac6165278f3cc9ed01193a87530 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 28 Jun 2025 09:09:54 -0700 Subject: [PATCH] Add stacktrace in metadata_props for rewriter created nodes Signed-off-by: Justin Chu --- onnxscript/ir/_tape.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/onnxscript/ir/_tape.py b/onnxscript/ir/_tape.py index 79312eaefa..0c0a259227 100644 --- a/onnxscript/ir/_tape.py +++ b/onnxscript/ir/_tape.py @@ -4,6 +4,7 @@ from __future__ import annotations +import traceback from typing import TYPE_CHECKING, Any, Optional, Sequence from onnx_ir import tape @@ -16,6 +17,12 @@ UsedOpsets = set[tuple[str, Optional[int]]] +def _get_stacktrace() -> str: + """Get the stack trace of the current execution context.""" + stack = traceback.extract_stack() + return "".join(traceback.format_list(stack[:-3])) + + class Builder(tape.Tape): """An extension of the tape that provides a more convenient API for constructing the IR.""" @@ -32,9 +39,18 @@ def _make_node(self, op_type: str, inputs: Sequence[ir.Value], kwargs: dict[str, assert isinstance(outputs, int) num_outputs = outputs + metadata_props = { + "pkg.onnxscript.rewriter.stacktrace": _get_stacktrace(), + } + if num_outputs == 1: value = super().op( - op_type, inputs=inputs, attributes=kwargs, domain=domain, version=version + op_type, + inputs=inputs, + attributes=kwargs, + domain=domain, + version=version, + metadata_props=metadata_props, ) if isinstance(outputs, Sequence): value.name = outputs[0] @@ -46,6 +62,7 @@ def _make_node(self, op_type: str, inputs: Sequence[ir.Value], kwargs: dict[str, domain=domain, version=version, num_outputs=num_outputs, + metadata_props=metadata_props, ) if isinstance(outputs, Sequence): for value, name in zip(values, outputs):