[user-cuda-streams] Add fork/join custom ops

[mlazos]

Creates the fork/join stream ops. These ops are passthrough ops which mutate all of their args (without actually performing any computation on them) so that during functionalization, implicit dependencies are added on all of their args. This allows us to prevent reordering during our pre/post grad graph passes.

Stack from ghstack (oldest at bottom):

• [dynamo] Remove retrieving objects by ID #162905
• [user-streams] Add basic stream tests #164523
• [user-streams] Handle returning the current stream with/without device index #165356
• [user-streams] Track symbolic current stream #165212
• [user-streams] Add current stream source #165211
• [user-streams] Fix stream graph output semantics #164819
• [User-streams] Make torch.Event weakref compatible #164522
• [user-streams] Make device-agnostic streams weakref compatible #164304
• [user-cuda-streams] Add cuda streams test suite #162901
• [user-streams] Support streams as contexts #164507
• [user-streams] Have StreamVariable inherit from StreamContextVariable #164344
• [user-streams] Move StreamContextVariable into streams module #164343
• [user-streams] update stream context to use fork/join #162903
• -> [user-cuda-streams] Add fork/join custom ops #162900
• [user-streams] Handle aliasing properly #163028
• [user-cuda-streams] Pass streams/events to the graph via lookup table #162899
• [user-streams] Move stream code to streams module #163027

Make custom ops inplace

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/162900

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm][CI] Machines under the label linux.rocm.gpu.2 are undergoing maintenance.

✅ You can merge normally! (3 Unrelated Failures)

As of commit d2053df with merge base ecb5307 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-jammy-py3.13-clang12 / test (default, 5, 5, linux.4xlarge) (gh) (detected as infra flaky with no log or failing log classifier)
• pull / linux-jammy-py3.13-clang12 / test (einops, 1, 1, linux.2xlarge) (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• trunk / macos-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh) (similar failure)
  RuntimeError: doctests 1/1 failed!

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mlazos]

After discussing offline with Brian and Elias, I opted to have a simpler custom op here which just records the from/to streams and we will do more of the dependency tracking after joint graph tracing and functionalization. We no longer have "fake mutation" to get functionalization to stream dependencies properly.

[mlazos]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here