Simplify function implementation returned by thunder.jit for easier instrumentation of different stages

[IvanYashchuk]

This PR changes the thunder.jit.fn_ implementation to use just a few lines of code making it more readable and amenable to per-function instrumentation like cProfile, pyinstrument, or NVTX decorators.

I removed a few unused timer measurements or if there was just a stop with no start. Other than that all CompileStats time measurements are preserved in this PR.

There was a todo to rename "last_executed" to "last_computation", this change is included here.

Here is what the function looks like now:

def fn_(*args, **kwargs) -> Any:
    cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
    check_storage_aliases(cache_entry, inps)
    result = cache_entry.computation_fn(*inps)
    result = maybe_connect_to_autograd(cache_entry, result)
    result = maybe_call_epilogue(cache_entry, result, pro_to_epi)
    return result

[IvanYashchuk]

@tfogal, could you please review the changes? Would this PR work well with the planned nvtx marker decorators from #1268?

[IvanYashchuk]

CI failures that I need to take a look what's going on:

=========================== short test summary info ============================
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-llama1-like] - TypeError: tuple indices must be integers or slices, not tuple
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-codellama2-like] - TypeError: tuple indices must be integers or slices, not tuple
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-falcon-40b-like] - TypeError: tuple indices must be integers or slices, not tuple
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-long-context-like] - TypeError: tuple indices must be integers or slices, not tuple
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-llama2-like] - TypeError: tuple indices must be integers or slices, not tuple
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants_kvcache[cpu-falcon-7b-like] - TypeError: tuple indices must be integers or slices, not tuple

[tfogal]

Thank you!

[tfogal]

There was a todo to rename "last_executed" to "last_computation", this change is included here.

Is this one of the timers that mixology is using? If so, we may want to hold off on this part. @mpatel31415 ?

[mpatel31415]

There was a todo to rename "last_executed" to "last_computation", this change is included here.

Is this one of the timers that mixology is using? If so, we may want to hold off on this part. @mpatel31415 ?

No, we use iter_t0 = time.perf_counter() and t1 = time.perf_counter() from benchmark_litgpt.py script :)

[tfogal]

There was a todo to rename "last_executed" to "last_computation", this change is included here.

Is this one of the timers that mixology is using? If so, we may want to hold off on this part. @mpatel31415 ?

No, we use iter_t0 = time.perf_counter() and t1 = time.perf_counter() from benchmark_litgpt.py script :)

oh, great!

I thought you were using some timers, though---how are you measuring compilation time, if not using these CompileStats timers?

[mpatel31415]

There was a todo to rename "last_executed" to "last_computation", this change is included here.

Is this one of the timers that mixology is using? If so, we may want to hold off on this part. @mpatel31415 ?

No, we use iter_t0 = time.perf_counter() and t1 = time.perf_counter() from benchmark_litgpt.py script :)

oh, great!

I thought you were using some timers, though---how are you measuring compilation time, if not using these CompileStats timers?

We measure 2N iterations giving us iter_times = [t1, t2, .. t_2N] and we assume that compilation (and warmup) time is sum(iter_times[:N]) - sum(iter_times[N:]). In this way it's easy to measure it irrespective of the compilation method in the same way.

[tfogal]

I thought you were using some timers, though---how are you measuring compilation time [. . .]

We measure 2N iterations [. . .]

Oh! That's actually great news as I had (incorrectly) thought you were using the timers, and this makes it easier for us to change that interface.

Thanks!

[tfogal]

Looks like tests pass and this is good to go, AFAICT.

@t-vi merge?

[t-vi]

So when will we be able to drop the timers?

[t-vi]

Thank you @IvanYashchuk @mpatel31415 @tfogal
Seems good, I would want to drop all the timings gathered here very soon unless there is a user for them.
From what I understand, people interested in timings do profile things.

[tfogal]

So when will we be able to drop the timers?

I can start looking at this after #1268