Imarkov/conditional compilation ranges

[ProExpertProg]

Purpose

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[ProExpertProg]

A few initial thoughts. Could we also use a dataclass instead of a tuple for a compiled range? We can add utility methods (like is_single_size), names to the elements, and docs to make the code clearer.

Also, currently if there is more than one range above the cudagraphs capture size, I think we don't ever trigger it in the GPU model runner as the compilation only happens as the compiled model is invoked with the shape (with _dummy_run) - we should make sure to dummy run for each compile range.

I also think we should try to give hints to Inductor about the range, can be done as a follow-up.

[ProExpertProg]

Lost the compilation time update

Suggested change

	if compile_range is None:
	compilation_config.compilation_time += elapsed
	if compile_range is None:

[ProExpertProg]

This comment implies ranges are done inclusive-exclusive but in the code you use inclusive-inclusive. Can we standardize on inclusive-exclusive?

[ilmarkov]

We do inclusive-exclusive in the code in piecewise_backend.py.

self.is_in_range = (
            lambda x, range: range[0] <= x < range[1]
            if range[0] < range[1]
            else x == range[0]
        )

[ProExpertProg]

Add the current range to cache key and check the number of times the manager gets called (to make sure the bug you found doesn't manifest)