frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Simplify PrecompileContext to no longer be a CacheArtifactManager (#162886)

Browse Source 
Summary:
This diff does a big refactor of PrecompileContext to make it considerably simpler: instead of being a CacheArtifactManager and managing a bunch of bytes, it simply stores two things: dynamo cache entries and backend cache entries. When asked, it stitches them together into PrecompileCacheEntries, which are stored by DynamoCache.

This structure then allows us to register DynamoCache to the regular Megacache API, instead of having two separate APIs that are confusing. It also lets us remove the autotune cache integration, since MegaCache API will automatically store autotune cache entries.

The intent here is that users who want to use caching precompile will simply be able to use torch.compiler.save_cache_artifacts as before, just with `torch.dynamo.config.caching_precompile` set to True. They can also directly interact with PrecompileContext if they wish to specifically only load Precompile entries, using PrecompileContext.create_cache_entries().

Saving single entries and such with DynamoCache still works normally.

Test Plan:
All existing unit tests pass.

Rollback Plan:

Differential Revision: D82380307

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162886
Approved by: https://github.com/zhxchen17
This commit is contained in:
James Wu
2025-09-20 01:24:37 +00:00
committed by PyTorch MergeBot
parent 8225a26835
commit bfe9e60ffb
10 changed files with 146 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
torch/compiler/_cache.py
View File

@ -48,9 +48,6 @@ class CacheArtifact(ABC):
def populate_cache(self) -> None:
pass
def precompile_compatible(self) -> bool:
return False
@staticmethod
def type() -> str:
"""
@ -131,14 +128,6 @@ class CacheInfo:
def pgo_artifacts(self) -> list[str]: # type: ignore[empty-body]
...
@property
def precompile_aot_autograd_artifacts(self) -> list[str]: # type: ignore[empty-body]
...
@property
def precompile_dynamo_artifacts(self) -> list[str]: # type: ignore[empty-body]
...
def add(self, artifact: CacheArtifact) -> None:
self.artifacts[artifact.type()].append(artifact.key)