[torch.compile] disable logging when cache is disabled (#12043)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao 2025-01-18 04:29:31 +08:00 committed by GitHub
parent b5b57e301e
commit 7b98a65ae6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -251,15 +251,27 @@ def wrap_inductor(graph: fx.GraphModule,
def _get_shape_env() -> AlwaysHitShapeEnv:
return AlwaysHitShapeEnv()
with patch(# for hijacking the hash of the compiled graph
"torch._inductor.codecache.compiled_fx_graph_hash",
hijack_compiled_fx_graph_hash), \
patch(# for providing a dummy shape environment
"torch._inductor.codecache.FxGraphCache._get_shape_env",
_get_shape_env), \
patch(# for forcing the graph to be cached
"torch._inductor.codecache.FxGraphCache._check_can_cache",
_check_can_cache):
with ExitStack() as stack:
if not cache_data.disabled:
# compilation cache is enabled, patch several functions
# for hijacking the hash of the compiled graph
stack.enter_context(
patch("torch._inductor.codecache.compiled_fx_graph_hash",
hijack_compiled_fx_graph_hash))
# for providing a dummy shape environment
stack.enter_context(
patch(
"torch._inductor.codecache.FxGraphCache._get_shape_env",
_get_shape_env))
# for forcing the graph to be cached
stack.enter_context(
patch(
"torch._inductor.codecache.FxGraphCache._check_can_cache",
_check_can_cache))
compiled_graph = compile_fx(graph,
example_inputs,
config_patches=current_config)