[ez][dynamo] some code movement (#150753)

`optimize_assert` already does the lookup for `backend` and
`backend_ctx_ctor`. This simply moves the lookups within `optimize`
lower so we don't end up calling these functions twice unnecessarily
in the `optimize_assert` path.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150753
Approved by: https://github.com/anijain2305, https://github.com/jansel
This commit is contained in:
bobrenjc93 2025-04-06 23:52:59 -07:00 committed by PyTorch MergeBot
parent b0810168a3
commit 3f0931b1de

View File

@ -1018,11 +1018,6 @@ def _optimize(
):
return _NullDecorator()
backend = get_compiler_fn(backend)
# Find if backend has any extra context manager
backend_ctx_ctor = getattr(backend, "backend_ctx_ctor", null_context)
if nopython:
return optimize_assert(
backend,
@ -1030,6 +1025,12 @@ def _optimize(
hooks=hooks,
rebuild_ctx=rebuild_ctx,
)
backend = get_compiler_fn(backend)
# Find if backend has any extra context manager
backend_ctx_ctor = getattr(backend, "backend_ctx_ctor", null_context)
# The backend function is stashed in the callable returned by
# _optimize_catch_errors in the field _torchdynamo_orig_callable. This can
# be used by eval_frame.c to insert a guard on the backend.