gm is only used for @masnesral's specific freezing use case. In that case, we need the post compile gm to do things. But we don't have that for AOTAutogradCache, so AOTAutogradCache doesn't support the freezing workflow. It always passes in None here since the gm should never be needed for post compile.

(The only "gm" it has access to at this step is the post dynamo gm. So anything needed to load from FXGraphCache needs to be obtainable from that)

Hmm, are we talking about the same code? I see here in post_compile:

        if cudagraphs:
            # It's possible that cudagraphs is enabled, but was disabled
            # during a previous compilation we're loading from the cache.
            # If so, we need to disable it on this new process too.
            if self.disabled_cudagraphs_reason:
                if "cuda" in self.device_types:
                    log_cudagraph_skip_and_bump_counter(
                        f"skipping cudagraphs due to {self.disabled_cudagraphs_reason}"
                    )
                else:
                    counters["inductor"]["cudagraph_skips"] += 1
                BoxedBool.disable(cudagraphs)
            else:
                cudagraph_post_compile(
                    example_inputs,
                    self,
                    cudagraphs,
                    gm,
                )

So gm is needed to do some CUDA graphs stuff. I don't see how AOTAutograd cache proves that cudagraphs could never be enabled here. I can try to write a test case that tickles it... later, maybe.

Ah yeah, it's confusing because cudagraphs_post_compile only uses gm if freezing was enabled on the cache write, and CompiledFxGraph doesn't have the constants on hand already (which it always will in AOTAutogradCache scenarios)

IMO we should change this so the gm is not passed to cudagraphs_post_compile: only the constants themselves, and instead have whatever caller to cudagraphs_post_compile handle obtaining the constants from the gm.

The gm field is really confusing to work with because it is essentially untyped: it can represent both a dynamo gm and a post aotdispatch gm.

@masnesral It's really confusing what's happening here so I might refactor it some more: we shouldn't call CompoledFxGraph.get_constants() with an optional gm and change its behavior depending on some unspoken invariants. We should only call this helper function if we know freezing is on in the cache entry and we're not coming from AOTautogradcache.

The way it's written now there's no way to tell that this gm is always None from the exact code path where it's precisely correct, and a correct value otherwise

I see, this refactor is nice because OutputGraph is at a layer higher than a cache, so it can do the post compilation unconditionally. Without OutputCode, this refactor would be hard to write because there would be two different ways you could obtain all the needed inputs: through compile or through cache hit.

That said, maybe I'm missing context from the last stack but I don't remember how the deserialize step converts the stuff saved into the cache to an outputcode. I assume we still need something like that, to massage the cached result into the object we return from inductor.

Actually, you could have always done this refactor, I just forgot to update this call site when I initially introduced post_compile. It's not doing much.

That said, maybe I'm missing context from the last stack but I don't remember how the deserialize step converts the stuff saved into the cache to an outputcode. I assume we still need something like that, to massage the cached result into the object we return from inductor.

I'm not exactly sure what your question here is. But CompiledFxGraph is directly unpickled from the cache:

        def iterate_over_candidates() -> Generator[CompiledFxGraph, None, None]:
            if local:
                subdir = FxGraphCache._get_tmp_dir_for_key(key)
                if os.path.exists(subdir):
                    for path in sorted(os.listdir(subdir)):
                        try:
                            with open(os.path.join(subdir, path), "rb") as f:
                                yield pickle.load(f)
                        except Exception:
                            log.warning(
                                "fx graph cache unable to load compiled graph",
                                exc_info=True,
                            )

And, contrary to expectation, the current_callable is directly bong'ed in by codecache.py

        try:
            with dynamo_timed(
                "PyCodeCache.load_by_key_path", log_pt2_compile_event=True
            ):
                graph.current_callable = PyCodeCache.load_by_key_path(
                    graph.cache_key,
                    artifact_path,
                    graph.cache_linemap,
                    graph.get_constants(gm),
                ).call

This part I do plan to refactor as well but I haven't gotten to it yet.