This will, in essence, increase the average compile time logged to aot_dispatcher_function right? In the sense that, if there was significant compile time after this function returned, it would now be recorded under create_aot_dispatcher_function.

To be clear, I don't think there's actually any work being done after this function returns, but that's just in theory.

I believe that I end the ExitStack context manager at the same point in time it would have previously ended when we returned from the recursive calls. You should be able to verify this.

It almost feels like AOTState itself should either be a context manager that you can grab information from, or a functional state that you return from each stage like:

(aot_state, aot_graph_capture) = aot_stage1()

This in between where aot_state is mutable as a side effect of calling the function is... okay I guess lol

I don't think it's a huge antipattern to pass in AOTState as an argument to functions. Actually, I wanted to make it a proper object with methods on it for stages (so you access the state via self) but this would have required a lot of reindenting and code motion, so I didn't do it to make the PR easier to review.

A functional AOTState will be error prone, because there are a lot of places where we were previously relying on variable shadowing to ensure you get the "latest" version of any given variable. So I actually had a more functional version of this class that I replaced with the mutating one to stop people from accidentally using stale versions of variables.

I would expect fw_metadata to be in here, no? Isn't it generated during graph capture?

Weeeell it's already in AOTState, so I didn't put it in this struct lol. We could have this struct have a pointer to AOTState so you have everything available in the same place.

Actually I change my mind, fw_metadata gets changed by wrappers, so it being in AOTState makes more sense than putting it in AOTGraphCapture, because it's not a pure result of AOTGraphCapture

This comment is not completely true in that there are wrappers that happen after graph capture too, tehre's just some wrappers that happen even before graph capture. Hmm, maybe these wrappers in the new model should be in AOTState instead?

AOTDedupWrapper and AOTSyntheticBaseWrapper used to happen right before calling aot_dispatch_{base|autograd}, but also come with a post_compile after the inner compiler returns. So they're called as the first step in aot_dispatch_base. Not sure exactly where they belong in this model.

Other wrappers, like FunctionalizedRngWrapper and RuntimeWrapper itself, are created only for aot_stage2_compile. I don't think anything needs to track these, since they're all compile specific, so they can just be handled in aot_stage2_compile_{base|autograd}

While I'm happy to delete this comment for clarity, I think this comment is directionally correct (and it is certainly how @bdhirsh and I conceptualized AOTAutograd when we originally worked on v2 of the implementation). For example, you focus on "when" AOTDedupWrapper happens. But intuitively, what it actually does is it takes a graph where there some inputs that may occur multiple times, and gives us a new graph where the input only occurs once, so that it is safe to run autograd on it. There are things that have to happen before and after here: before we have to create the new function; after we have to apply a runtime wrapper that will eliminate the duplicate arguments.

Sorry I actually meant to make this comment on the wrapper itself and was reasoning about where these wrappers belong.

I think the reason I'm a bit hung up on the timing is precisely because wrappers execute code both before and after the graph capture, so it was a little confusing that it was being returned by AOTGraphCapture itself. But it makes sense why it's attached to the return value; you need to track the post compiles of each of the wrappers. I think the awkward part is, because they are wrappers, the actual precise ordering of wraps is like:

AOTDedupWrapper -> aot_dispatch wrappers (pre_compile) -> fw/bw_compiler output -> aot_dispatch wrappers (post_compile) -> AOTDedupWrapper (post compile)

So in essence, even though stage2 is taking in AOTGraphCapture and these wrappers are part of AOTGraphCapture, it's basically just passing them through without applying them until the very end.

But honestly that level of detail of understanding is probably not that important to fully encapsulate in the data model or this comment, I just needed to wrap my head around it. It's probably fine as is.