Use torch._dynamo.mark_static() API to allow tensor shape specialization outside of the kernel code

[yf225]

In general, baking in specific tensor dim sizes into the kernel is good for perf, but we can't do that for all scenarios because we need a catch-all dynamic kernel for rare sizes (otherwise we would have to autotune for those novel long-tail sizes which is generally prohibitive due to extra autotuning time).

Right now to bake in a specific shape, one must do hl.specialize(), which makes the kernel not suitable for the all-dynamic case (unless we copy-paste the kernel and remove specific hl.specialize() calls).

By reusing torch._dynamo.mark_static() API, user can achieve the above by having a Helion kernel impl without hl.specialize() calls, and then add specialization outside of the kernel code, making their codebase much cleaner.

Closes #1046.

cc. @Chillee

[jansel]

I'm a bit confused by the semantics of this. Does it mutate the behavior of the kernel? What if I call the kernel before calling this API? Can you provide some examples.

[yf225]

I'm a bit confused by the semantics of this. Does it mutate the behavior of the kernel? What if I call the kernel before calling this API? Can you provide some examples.

@jansel it doesn't mutate the behavior of the original kernel - kernel.specialize_args() will create a new kernel that only shares the config / settings / .fn / ._key_fn with the original kernel, and all other mutable states like _bound_kernels and _specialize_extra are separate

also updated the docs at https://github.com/pytorch/helion/pull/1210/files#diff-9fdadeb7f22ce14b3a9fa419dea27615a1e66a4faf75bc76942d3474eabc86c1R371-R413 and added test_specialize_args_does_not_mutate_original to add examples

[jansel]

Could we use the same mark dynamic API as torch.compile for this?

[jansel]

Should mention this in the docs