What's missing here is validation that guards were added after loading from the cache. Still investigating, but does anyone know if there's a straightforward way to access the guards in a structured way so I can do some validation here?

@ezyang do you have a recommendation? If I have a torch.compiled function, is there a straightforward way to see that our guards were properly added even in the case of a cache hit?

The prefix character seems to be the existing scheme here to differentiate different kinds of hashes for different types of cached objects. Probably I should use an enum. TODO.

yeah, when I wrote the original PR that was my thought too

Is there any reason not to use a longer, more descriptive prefix? It's not like we have a file path length limitation or something

I was sticking with the convention already started, but can definitely turn it into a a more descriptive string. There's also an existing precedent for adding an extension, e.g., "cg" which I guessed means "code graph"?

Since we need to access the example_args, it was convenient to make the gm and example_args arguments explicit rather than packaging them up in a list.

I kept this implementation under the FxGraphCache class, but all the methods are static. I prefer the namespacing of putting these methods in the class, but lemme know if a class of all static methods hurts your eyeballs.

Sometimes people prefer @classmethod and cls.method for this reason i think. up to you.

Please comment. I don't know if clearing the whole cache is problematic, e.g., from a performance perspective. I could, for example, introduce a new context manager analogous to suppress_guards() above that selectively uses a cached vs. non-cached version of shape_env.evaluate_expr, but I thought it was worth asking if that's overkill.

Seems off to me that we would clear cache. leaving this to @ezyang to comment

I agree it seems kinda weird that we would need to clear the cache. If a cache hit is preventing you from adding a guard, then that should mean that that guard was already evaluated, no?

See some of the comments above. The approach here is to:

1. Load a possible match from the cache.
2. Evaluate the guards, but in a mode that does not modify the guards in the current environment, because: in the case that there's a miss, we don't actually want the current env to change.
3. If there's a hit, only then re-evaluate to cause the guards to be added.

The problem here is that the caching at evaluate_expr() interferes with #3

@masnesral I think a better solution would be to customize the caching on evaluate_expr so that it respects whether or not guards were suppressed or not. WDYT?

@ezyang is that correct though? If the expression has been evaluated before (even guards suppressed), don't we still want to evaluate it again in order to get the guards added?

Hmm, well I guess it didn't matter when the hints were ints but it does matter if the hints are symints. I mean, I wouldn't be opposed to just turning off this caching entirely when you're playing funny tricks with symint hints.

I preferred a ctor here rather than extracting all these fields at the caller.

See comment above. I changed this to be explicit about the two args (gm, example_inputs)

The changes here are just about splitting apart the evaluate_guards_for_args function into smaller pieces that I can use for the cache impl. Namely, I need separate phases for creating the guards expression and evaluating it (in a different context)

Seems like a good comment to put into the code comments

The diff looks weird here. load_graph is replaced by lookup_graph above and the implementation here is for save_graph. So ignore the red lines I guess.

It might be useful in the future to canonicalize symints so that hash(tensor([s1, s1, s2])) == hash(tensor([s3, s3, s4])).. not needed now

It's not clear to me that this is sound, given https://github.com/pytorch/pytorch/pull/111421/files#diff-c9b517f8db609ffa866804dfa2689188a4fee20abacaa0b0dca91625c1b5cb8dR705

if you say [s0, s1] == [s4, s3], you have to make sure that you know how to flip the SymInt arguments when eval'ing the guards. That sounds difficult.

Sometimes people prefer @classmethod and cls.method for this reason i think. up to you.

Seems off to me that we would clear cache. leaving this to @ezyang to comment

do we need FileLock here ?

It is held from the caller

Is it possible to end up loading a more generic graph from cache than we need?

Ohhhh, yes. That's a nice catch. So really I need to consider all versions and pick the "best" option, don't I? Hmm, is there any obvious criteria I can use to determine whether one option is more appropriate than the other? cc @eellison

I think maybe as a follow up - not needed in initial pr imo. Generally the behavior for symbolic shapes is to reuse a compilation if it works.

Suppose you have a cache for (s0, 4) and another for (5, s1). There is no unique best choice cache entry to load.

Probably the only way to actually do this is some sort of autotuning based thing, where for any given concrete size you've benchmarked which one runs fastest, and you load that particular one.

I agree it seems kinda weird that we would need to clear the cache. If a cache hit is preventing you from adding a guard, then that should mean that that guard was already evaluated, no?

Given that dynamic=None has a different behavior compared to True and False, it might be worth adding here as an option too

So it turns out there's a shortcoming with that change. None seems to be equivalent to False in terms of the FX graph that gets generated (at least for these tests). So I'd need to rearchitect the tests slightly to make sure each test gets a clean tmp directory. I can do that, but I sorta liked the current behavior of leaving the tmp directory intact for the run of the full set of tests because we'd catch cases of a cache hit when we shouldn't (see my comment on setUpClass)

Same as previous comment

this is equivalent to dynamic=True on the compile function