aot_precompile() — pytorch Function Reference

Architecture documentation for the aot_precompile() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration calls 1 called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ aot_precompile() — pytorch Function Reference

Dependency Diagram

graph TD
  5130216c_70e0_8051_b662_1aed31c12627["aot_precompile()"]
  04a3a4a6_8db3_854d_a893_02c9542bf9dd["speedup_experiment()"]
  04a3a4a6_8db3_854d_a893_02c9542bf9dd -->|calls| 5130216c_70e0_8051_b662_1aed31c12627
  1eff8423_1f23_d138_6815_07d8dc29a749["_normalize_bench_inputs()"]
  5130216c_70e0_8051_b662_1aed31c12627 -->|calls| 1eff8423_1f23_d138_6815_07d8dc29a749
  style 5130216c_70e0_8051_b662_1aed31c12627 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1517–1545

def aot_precompile(model, example_inputs):
    example_args, example_kwargs = _normalize_bench_inputs(example_inputs)

    with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
        save_path = f.name

    with fresh_cache(), torch._dynamo.config.patch("enable_aot_compile", True):
        compiled_fn = torch.compile(
            model,
            fullgraph=True,
            options={"guard_filter_fn": lambda guards: [False for _ in guards]},
        ).forward.aot_compile((example_args, example_kwargs))

        compiled_fn.save_compiled_function(save_path)

        torch._dynamo.reset()
        with open(save_path, "rb") as f:
            load_start_time = time.perf_counter()
            loaded_fn = torch.compiler.load_compiled_function(f)
            load_end_time = time.perf_counter()
            print(
                f"AOT Precompile loading time: {load_end_time - load_start_time} seconds"
            )

            def opt_aot_precompile(_, example_inputs, collect_outputs=False):
                example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
                return loaded_fn(model, *example_args, **example_kwargs)

            return opt_aot_precompile