Metal backend: Add MPSGraph caching

[manuelcandales]

This pull request introduces a new caching infrastructure for compiled MPSGraph objects in the Metal backend, significantly improving performance for repeated matrix multiplication and convolution operations by reusing previously compiled graphs and their associated tensors. The changes also include cache statistics tracking and refactoring of the relevant code paths to leverage the cache. Note that caching for the attention operation is not yet implemented.

MPSGraph Caching Infrastructure

• Added a new global cache (graph_cache) for compiled MPSGraph objects and their input/output tensors, keyed by operation type and parameters, enabling reuse and reducing graph compilation overhead.
• Introduced GraphCacheKey and CachedGraph structures, along with a custom hash function, to uniquely identify and store cached graphs for matrix multiplication and convolution operations.

Matrix Multiplication and Convolution Refactoring

• Refactored the matrix multiplication (aoti_torch_mps_mm_out) and convolution (aoti_torch_mps_convolution) functions to check the cache for an existing compiled graph before creating a new one, and to store newly compiled graphs in the cache for future reuse. [1] [2]
• Removed manual release of MPSGraph objects after execution, as cached graphs are now retained for reuse, preventing unnecessary recompilation and memory leaks. [1] [2]

Cache Statistics and Monitoring

• Added CacheStats structure to track cache hits and misses.

[manuelcandales]

Stack from ghstack (oldest at bottom):

• -> Metal backend: Add MPSGraph caching #15346

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/15346

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 116 Pending

As of commit f66ac9c with merge base 526eb18 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.