PyTorch Mega-Cache 实现深度解读

class CacheArtifactManager:
    """统一管理所有缓存产物类型"""
    artifact_types = {
        "inductor_artifacts": InductorCacheArtifact,
        "autotune_artifacts": AutotuneCacheArtifact,
        "aot_autograd_artifacts": AOTAutogradCacheArtifact,
        "pgo_artifacts": PGOArtifact
    }

def save_cache_artifacts() -> Optional[tuple[bytes, "CacheInfo"]]:
    """收集并序列化所有缓存产物"""
    artifacts = CacheArtifactManager.serialize()
    return pickle.dumps(artifacts), cache_info

def load_cache_artifacts(data: bytes) -> bool:
    """反序列化并热加载缓存"""
    artifacts = pickle.loads(data)
    return CacheArtifactManager.deserialize(artifacts)

class AOTAutogradCacheDetails:
    """捕获 AOTAutograd 缓存的所有相关信息"""
    def __init__(self, gm, example_inputs, config, fx_config):
        self.gm = gm                    # FX 图模块
        self.example_inputs = example_inputs  # 示例输入
        self.config = config            # AOT 配置
        self.fx_config = fx_config      # FX 编译配置

def check_autotune_cache(configs, filename, inductor_meta):
    """检查自动调优缓存命中情况"""
    if should_use_cache():
        autotune_cache = AutotuneCache.create(inductor_meta, filename, configs_hash)
        best_config = autotune_cache.read_best(inductor_meta, configs)
        if best_config:
            return [best_config], "hit"  # 缓存命中
    return configs, "miss"  # 缓存未命中

def autotune_to_one_config(self, *args, **kwargs):
    """执行实际的自动调优"""
    # 基准测试所有配置
    timings = self.benchmark_all_configs(*args, **kwargs)
    
    # 选择最优配置
    best_launcher = min(timings, key=timings.get)
    
    # 保存到缓存
    if self.save_cache_hook:
        self.save_cache_hook(best_launcher.config, autotune_time)
    
    return best_launcher

class FxGraphHashDetails:
    """FX 图缓存键生成器"""
    def __init__(self, gm, example_inputs, fx_config):
        self.gm = gm
        self.example_inputs = example_inputs
        self.fx_config = fx_config
        self.torch_version = torch.__version__
        self.system_info = get_system_info()

import torch

# 启用编译缓存
torch.compiler.config.enable_caching = True

# 编译函数
@torch.compile
def my_function(x, y):
    return torch.matmul(x, y) + torch.sum(x)

# 首次编译（慢）
result1 = my_function(x, y)

# 后续编译（快）
result2 = my_function(x, y)  # 使用缓存

# 保存缓存产物
cache_data, cache_info = torch.compiler.save_cache_artifacts()

# 加载缓存产物
success = torch.compiler.load_cache_artifacts(cache_data)

# 清除所有缓存
torch.compiler.clear_cache()

# 清除特定类型缓存
torch._inductor.FxGraphCache.clear()
torch._functorch.AOTAutogradCache.clear()

# 配置远程缓存
torch.compiler.config.remote_cache_url = "http://cache-server:8080"
torch.compiler.config.enable_remote_cache = True

torch.compiler.config.compress_cache = True
torch.compiler.config.compression_threshold = 1024 * 1024  # 1MB

# 获取缓存统计信息
stats = torch.compiler.get_cache_stats()
print(f"缓存命中率: {stats.hit_rate:.2%}")
print(f"缓存大小: {stats.total_size / 1024 / 1024:.1f} MB")

# 开发环境启用详细缓存日志
torch.compiler.config.cache_log_level = "DEBUG"
torch.compiler.config.enable_cache_metrics = True

# 生产环境配置
torch.compiler.config.enable_caching = True
torch.compiler.config.cache_dir = "/opt/pytorch_cache"
torch.compiler.config.max_cache_size = 10 * 1024 * 1024 * 1024  # 10GB

# 禁用特定缓存进行调试
torch.compiler.config.enable_fx_graph_cache = False
torch.compiler.config.enable_autotune_cache = False

# 启用缓存调试模式
torch.compiler.config.cache_debug_mode = True

# 导出缓存分析报告
torch.compiler.export_cache_report("cache_analysis.json")