OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[Inductor] Fix an inductor_provenance bug (#166432)

Browse Source 
Summary: Fix an inductor_provenance related error seen when running TORCH_COMPILE_DEBUG generated fx_graph_runnable.py.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/166432
Approved by: https://github.com/mlazos
This commit is contained in:
Bin Bao 2025-10-29 08:15:00 -07:00 committed by PyTorch MergeBot
parent 3f1824742c
commit 08b0a8f11a
2 changed files with 25 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
test/dynamo/test_fx_graph_runnable.py
View File

@ -363,6 +363,25 @@ class FxGraphRunnableTest(TestCase):
self._exec_and_verify_payload() self._exec_and_verify_payload()
def test_metrics_context(self):
"""
When TORCH_COMPILE_DEBUG is set, provenance_tracking_level is set to 1, and
the generated fx_graph_runnable crashed with,
RuntimeError: Cannot add inductor_provenance outside of a MetricsContext
"""
import torch._inductor.config as inductor_config
def f(x):
return x * 2 + 1
# Enable provenance tracking to trigger the code path that adds metrics
with inductor_config.patch(
{"trace.enabled": True, "trace.provenance_tracking_level": 1}
):
x = torch.randn(4, 4)
torch.compile(f)(x)
self._exec_and_verify_payload()
if __name__ == "__main__": if __name__ == "__main__":
from torch._dynamo.test_case import run_tests from torch._dynamo.test_case import run_tests

9
torch/_inductor/compile_fx.py
View File

@ -1550,9 +1550,12 @@ class _InProcessFxCompile(FxCompile):
payload_fn=lambda: inductor_kernel_stack_trace_str, payload_fn=lambda: inductor_kernel_stack_trace_str,
) )
if inductor_kernel_stack_trace_str: if inductor_kernel_stack_trace_str:
get_metrics_context().add_to_set( metrics_context = get_metrics_context()
"inductor_provenance", inductor_kernel_stack_trace_str if metrics_context.in_progress():
) metrics_context.add_to_set(
"inductor_provenance",
inductor_kernel_stack_trace_str,
)
node_runtimes = None node_runtimes = None
if inductor_metrics_log.isEnabledFor(logging.INFO): if inductor_metrics_log.isEnabledFor(logging.INFO):