Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid unnecessary copy in TensorSource #8849

Merged
merged 9 commits into from
Mar 26, 2025
Merged

Avoid unnecessary copy in TensorSource #8849

Changes from 3 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
495cb2f
avoid unnecessary copy in tensorsource
lsy323 Mar 18, 2025
e483f51
remove print
lsy323 Mar 19, 2025
407fb1c
avoid .to if not needed
lsy323 Mar 19, 2025
eeb715e
add comment, tests
lsy323 Mar 20, 2025
7b0689a
add comment
lsy323 Mar 20, 2025
888b0a4
fix cuda test
lsy323 Mar 20, 2025
a11744a
set copy as false
lsy323 Mar 25, 2025
2909950
remove testg
lsy323 Mar 25, 2025
170d8b0
revert cuda test
lsy323 Mar 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions torch_xla/csrc/runtime/tensor_source.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,23 @@ class AtenSource : public TensorSource {
if (target_torch_type != tensor.type().scalarType()) {
TORCH_LAZY_COUNTER("AtenSourceDowncasts", 1);
}
// TODO(ysiraichi): check, first, if tensor lives in a device that the
// current PjRt client has access. If so, we don't need to go through the
// CPU.
tensor_ = std::move(
tensor.to(at::TensorOptions().device(at::kCPU).dtype(target_torch_type),
/*non_blocking=*/false,
/*copy=*/true, at::MemoryFormat::Contiguous));
// The purposes of copy are:
// 1. Ensure the memory is contiguous, which is expected by PJRT.
// 2. Move CUDA tensor to CPU since we cannot pass CUDA memory to PJRT now.
// 3. Cast data type.
// We can avoid if copy is not needed.
if (tensor.device() == at::kCPU && tensor.is_contiguous() &&
tensor.dtype() == target_torch_type) {
tensor_ = std::move(tensor);
} else {
// TODO(ysiraichi): check, first, if tensor lives in a device that the
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NIT: I personally prefer to have TODOs linked to issues, and then having those assigned to people. That way, things can be more easily followed-up if a contributor is no longer active

// current PjRt client has access. If so, we don't need to go through the
// CPU.
tensor_ = std::move(tensor.to(
at::TensorOptions().device(at::kCPU).dtype(target_torch_type),
/*non_blocking=*/false,
/*copy=*/true, at::MemoryFormat::Contiguous));
}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As far as I understand it, tensor.to(...) (without the copy argument) already checks whether it should actually copy or not. So, what do you think of reverting to the old tensor.to(...) usage, but removing the copy argument, instead?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @ysiraichi, I didn't find a tensor.to(...) without the copy arg in C++, is it only in python?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. But, I think we can just /* copy= */false.

}

const void* data() const override { return tensor_.const_data_ptr(); }
Expand Down
Loading