Repository: pytorch
Issue: AOTAutograd has high fixed overheads
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

A few months ago [I looked at the fixed overheads of torch.compile](https://dev-discuss.pytorch.org/t/micro-optimizations-for-the-most-micro-of-benchmarks/1836) using [this microbenchmark](https://github.com/pytorch/pytorch/blob/6d9588a12b5834f29bd8970936749a9725e7f609/benchmarks/dynamo/microbenchmarks/overheads.py).

I wanted to give a quick update to that, in that since I realized I didn’t test things with `requires_grad=True` which triggers a bunch of extra wrapper code in AOTAutograd.

Here are, updated numbers from the last post with latest main.  Sadly, we have regressed 3.7us since last time.
- eager    4.6us
- compiled 14.0us


As an aside, if I enable `torch.inference_mode()` things get a bit better:
- eager    4.1us
- compiled 11.1us

Now on the other end, if I add `requires_grad=True` to [this line](https://github.com/pytorch/pytorch/blob/6d9588a12b5834f29bd8970936749a9725e7f609/benchmarks/dynamo/microbenchmarks/overheads.py#L14) of the microbenchmark, I get:
- eager    5.6us
- compiled 39.0us

Almost 2.8x worse than compile without the need to track gradients.  Note I am not calling backward, just triggering extra wrapper code to deal with backward if it does get called.  

So where is it coming from?  Here is the top of a cProfile (run 1000 times):
```
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1000    0.010    0.000    0.026    0.000 jit_compile_runtime_wrappers.py:505(forward)
2000/1000    0.007    0.000    0.076    0.000 eval_frame.py:427(_fn)
    20000    0.007    0.000    0.007    0.000 function.py:346(__getattribute__)
     1000    0.005    0.000    0.033    0.000 {built-in method apply}
2000/1000    0.005    0.000    0.043    0.000 utils.py:105(call_func_at_runtime_with_args)
     1000    0.005    0.000    0.005    0.000 {built-in method torch.tensor}
     1000    0.005    0.000    0.055    0.000 runtime_wrappers.py:77(runtime_wrapper)
     1000    0.002    0.000    0.039    0.000 function.py:582(apply)
     2000    0.002    0.000    0.002    0.000 eval_frame.py:140(change)
     1000    0.001    0.000    0.004    0.000 <string>:2(guard)
     1000    0.001    0.000    0.003    0.000 clph2nks5ejst7jdpwldntojs4uhnqr57fddd5w6jpqaqposu7td.py:46(call)
```

It looks like the worst offender is coming form [jit_compile_runtime_wrappers.py](https://github.com/pytorch/pytorch/blob/6d9588a12b5834f29bd8970936749a9725e7f609/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L506).  As a quick experiment I wrote a patch that just skips that entire function:
```diff
diff --git a/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py b/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py
index ce087d44e25..e4328df32a9 100644
--- a/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py
+++ b/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py
@@ -504,6 +504,8 @@ def aot_dispatch_autograd(
 
         @staticmethod
         def forward(ctx, *deduped_flat_tensor_args):
+            return CompiledFunction.compiled_fw([*deduped_flat_tensor_args])[:1]
+
             args = deduped_flat_tensor_args
             if backward_state_indices:
                 bw_state = args[backward_state_indices[0]]
```

This improves things to:
- compiled 26.3us

Obviously we can’t land that, but it does show the maximum benefit that could be had by optimizing that function.  More than half of the overhead added by AOTAutograd is coming from that one function.  Looking over that function’s code, it seems like there are a lot of checks for various corner cases that will do nothing on most invocations.  In addition, nearly all of those corner case checks should be statically known at compile time.  One way to speed that up would be to codegen a custom wrapper function that only includes the corner cases that are actually relevant to the function being called.  Though some more basic trimming and code optimization would also likely help.

I don't plan to work on this, so this issue is up for grabs if someone is interested in some micro-optimizations.

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang