Skip to content

Egaerly execute inplace ops if in eager mode#7322

Merged
JackCaoG merged 3 commits intomasterfrom
JackCaoG/in_place_update_egaer
Jun 21, 2024
Merged

Egaerly execute inplace ops if in eager mode#7322
JackCaoG merged 3 commits intomasterfrom
JackCaoG/in_place_update_egaer

Conversation

@JackCaoG
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG commented Jun 20, 2024

Even with functionization we still see some inplace ops, like optimization_barrier_, all_reduce_, currently in eager mode they won't be execute.

The eager mode today works by

xla/torch_xla/csrc/tensor.cpp

Lines 81 to 93 in 98dd99e

XLATensorPtr XLATensor::Create(
torch::lazy::Value ir_value, const torch::lazy::BackendDevice& device,
std::optional<at::ScalarType> logical_element_type) {
XLATensorPtr xtensor = c10::make_intrusive<XLATensor>(
XLATensor(std::move(ir_value), device, logical_element_type));
XLAGraphExecutor* graph_executor = XLAGraphExecutor::Get();
graph_executor->RegisterTensor(xtensor->data());
if (UseEagerDebugMode() || graph_executor->UseEagerMode()) {
std::vector<XLATensorPtr> xtensors({xtensor});
graph_executor->ApplyEagerSync(xtensors);
}
return xtensor;
}

when creating a new XLATensor with a IR, we will execute that IR. This doesn't handle the inplace update cases since there is no new XLATensor being created.

I also need to handle the two special cases in follow up pr

  1. random seed IR
  2. all_reduce_token IR

@JackCaoG JackCaoG added egaer usability Bugs/features related to improving the usability of PyTorch/XLA eager PyTorch/XLA eager-mode and removed egaer labels Jun 20, 2024
@JackCaoG JackCaoG marked this pull request as ready for review June 21, 2024 01:19
@JackCaoG JackCaoG merged commit 73923ec into master Jun 21, 2024
@JackCaoG JackCaoG mentioned this pull request Jul 2, 2024
Closed
JackCaoG added a commit that referenced this pull request Jul 10, 2024
@JackCaoG
Egaerly execute inplace ops if in eager mode (#7322)
2843d1a
JackCaoG added a commit that referenced this pull request Jul 10, 2024
@JackCaoG
Cherry-pick Egaerly execute inplace ops if in eager mode #7322 (#7666)
f84f7be
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@qihqi qihqi qihqi approved these changes

@wonjoo-wj wonjoo-wj wonjoo-wj approved these changes

@alanwaketan alanwaketan Awaiting requested review from alanwaketan

@will-cromar will-cromar Awaiting requested review from will-cromar

Assignees

No one assigned

Labels

eager PyTorch/XLA eager-mode usability Bugs/features related to improving the usability of PyTorch/XLA

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@JackCaoG @qihqi @wonjoo-wj