Skip to content

Conversation

@alanwaketan
Copy link
Collaborator

Summary:
This patch inherits LazyGraphExecutor::DeviceContextArena and overrides the follow methods:

  1. GetLiveTensors to return XLATensorPtrs.
  2. GetRngSeed to use our own + and * for torch::lazy::Value.
  3. IrValueFromScalar to use TensorToXlaData.

In addition, it has an extra method: GetBaseSeedData that is used by dynamo.

This patch needs pytorch/pytorch#90531 to function.

Test Plan:
CI.

@alanwaketan alanwaketan force-pushed the alanwaketan/lazy_graph_3 branch from 30cb306 to a132b4c Compare December 10, 2022 20:43
@alanwaketan alanwaketan mentioned this pull request Dec 11, 2022
@alanwaketan alanwaketan force-pushed the alanwaketan/lazy_graph_3 branch from a132b4c to ad8696a Compare December 12, 2022 19:33
@alanwaketan
Copy link
Collaborator Author

Thanks Jack and Wonjoo.

pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Dec 13, 2022
Summary:
This patch makes DeviceContextArena protected such that XLAGraphExecutor can reuse it. In addition, it makes all methods that utilize DeviceContextArena virtual such that XLAGraphExecutor can override them to provide its own DeviceContextArena.

P.S. This patch depends on pytorch/xla#4307 too.

Test Plan:
CI.

Pull Request resolved: #90531
Approved by: https://github.com/antoniojkim, https://github.com/JackCaoG
@alanwaketan alanwaketan merged commit c570f7f into master Dec 13, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tracing Lazy Tensor tracing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants