search

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.07k stars 60 forks source link

Allowing static constraint in torch/__init__.py #613

Closed jjsjann123 closed 5 days ago

jjsjann123 commented 1 week ago

What does this PR do?

Fixes issues in comment https://github.com/Lightning-AI/lightning-thunder/issues/463#issuecomment-2161258626.

This PR includes:

  1. Adding prims.sink with a DONT_DCE tag to explicitly mark dependency of computation trace on inputs. This allows prologue trace to inject static constraint logic on NumberProxy inputs, which would otherwise not show up in computation trace.
  2. Adding constraint to device index in torch_device.
jjsjann123 commented 1 week ago

Note that the target branch here is static_constraint. So let's not merge this PR yet.

jjsjann123 commented 1 week ago

failure is not related to this PR. I'll just patch the test with an explicit dtype.