[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: def _get_shape_env() -> AlwaysHitShapeEnv:
return AlwaysHitShapeEnv() return AlwaysHitShapeEnv()
with patch(# for hijacking the hash of the compiled graph with ExitStack() as stack:
"torch._inductor.codecache.compiled_fx_graph_hash", if not cache_data.disabled:
hijack_compiled_fx_graph_hash), \ # compilation cache is enabled, patch several functions
patch(# for providing a dummy shape environment
"torch._inductor.codecache.FxGraphCache._get_shape_env", # for hijacking the hash of the compiled graph
_get_shape_env), \ stack.enter_context(
patch(# for forcing the graph to be cached patch("torch._inductor.codecache.compiled_fx_graph_hash",
"torch._inductor.codecache.FxGraphCache._check_can_cache", hijack_compiled_fx_graph_hash))
_check_can_cache):
# 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, compiled_graph = compile_fx(graph,
example_inputs, example_inputs,
config_patches=current_config) config_patches=current_config)