main() — pytorch Function Reference

Source Code

benchmarks/dynamo/microbenchmarks/cache_hit_microbenchmarks.py lines 19–54

def main():
    torch._inductor.config.fx_graph_cache = True
    torch._inductor.config.fx_graph_remote_cache = False

    with fresh_cache():
        a = torch.randn(4).cuda()
        compiled_fn = torch.compile(huge_graph, backend="inductor")

        # write to cache
        compiled_fn(a)
        if counters["inductor"]["fxgraph_cache_miss"] != 1:
            raise AssertionError(
                f"expected fxgraph_cache_miss == 1, got {counters['inductor']['fxgraph_cache_miss']}"
            )

        def setup():
            torch._dynamo.reset()
            clear_caches()
            for m in torch._inductor.codecache.PyCodeCache.cache.values():
                os.remove(m.__file__)
            counters.clear()

        def fn():
            result = compiled_fn(a)
            if counters["inductor"]["fxgraph_cache_miss"] != 0:
                raise AssertionError(
                    f"expected fxgraph_cache_miss == 0, got {counters['inductor']['fxgraph_cache_miss']}"
                )
            if counters["inductor"]["fxgraph_cache_hit"] != 1:
                raise AssertionError(
                    f"expected fxgraph_cache_hit == 1, got {counters['inductor']['fxgraph_cache_hit']}"
                )
            return result

        t = min(timeit.repeat(fn, setup=setup, number=K, repeat=3))
        print(f"took {t:.1f}s")