Create executable buffer at the appropriate memory space on GPUs #14088
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jaro-sevcik
force-pushed
the
execute-output-buffers-with-memory-kind
branch
from
2d1bea6
to
fa7d607
Compare
|
cc @yashk2810 and @ezhulenev |
jyingl3
approved these changes
Jun 24, 2024
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 24, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646245598
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646245598
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646245598
copybara-service bot
pushed a commit
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR #14088 Copybara import of the project: -- fa7d607 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d607 PiperOrigin-RevId: 646386007
copybara-service bot
pushed a commit
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR #14088 Copybara import of the project: -- fa7d607 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d607 PiperOrigin-RevId: 646386007
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646245598
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646386007
copybara-service bot
pushed a commit
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR #14088 Copybara import of the project: -- fa7d607 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d607 PiperOrigin-RevId: 646386007
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646386007
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 645788858
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
Replaced by %s/in \[\(\d\+\), \(\d\+\)\]/\='in ['.(submatch(1)).', '.(submatch(2)+1).')'/g FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 645128569
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
…n GPUs Imported from GitHub PR openxla/xla#14088 Copybara import of the project: -- fa7d6071f64c3c614ccf78e49c359078a0c49604 by Jaroslav Sevcik <[email protected]>: Create executable buffer at the appropriate memory space Merging this change closes #14088 PiperOrigin-RevId: 646398973
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
The rewrite is only correct if there's a single division in the sum. Reverts a0da032 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646384770
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jun 25, 2024
Currently, we only canonicalize the order of LHS and RHS. In some very rare cases, the structure of an expression is such that `canonicalize(simplify-mlir(expr))` is not idempotent, leading to infinite loops in apply_indexing canonicalization. FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14088 from jaro-sevcik:execute-output-buffers-with-memory-kind fa7d6071f64c3c614ccf78e49c359078a0c49604 PiperOrigin-RevId: 646401010
yliu120
reviewed
Jun 27, 2024
| if (shape.has_layout() && | ||
| shape.layout().memory_space() == Layout::kHostMemorySpace) { | ||
| absl::StatusOr<PjRtMemorySpace*> memory_space_or = | ||
| device->memory_space_by_kind(PinnedHostMemorySpace::kKind); |
There was a problem hiding this comment.
Please use memory_space_by_kind_id whenever possible to accelerate lookup.
yliu120
reviewed
Jun 27, 2024
| shape.layout().memory_space() == Layout::kHostMemorySpace) { | ||
| absl::StatusOr<PjRtMemorySpace*> memory_space_or = | ||
| device->memory_space_by_kind(PinnedHostMemorySpace::kKind); | ||
| if (memory_space_or.ok()) { |
There was a problem hiding this comment.
This code looks a bit weird.
I feel like if the lookup failed, you will silently use the default_memory_space without any error or warning.
yliu120
reviewed
Jun 27, 2024
| @@ -2240,9 +2241,20 @@ std::unique_ptr<PjRtBuffer> OutputBufferHelper( | |||
| std::shared_ptr<TrackedDeviceBuffer> out_buffer = | |||
| TrackedDeviceBuffer::FromScopedShapedBuffer(result_buffer, | |||
| {definition_event}); | |||
| Shape shape = result_buffer->on_device_shape(); | |||
yliu120
reviewed
Jun 27, 2024
| @@ -2240,9 +2241,20 @@ std::unique_ptr<PjRtBuffer> OutputBufferHelper( | |||
| std::shared_ptr<TrackedDeviceBuffer> out_buffer = | |||
| TrackedDeviceBuffer::FromScopedShapedBuffer(result_buffer, | |||
| {definition_event}); | |||
| Shape shape = result_buffer->on_device_shape(); | |||
| PjRtMemorySpace* memory_space = | |||
| device->default_memory_space().value_or(nullptr); | |||
There was a problem hiding this comment.
So junwhan@ made it always returning a value. No need to .value_or(nullptr)?
There was a problem hiding this comment.
As far as I can see, the interpreter device can still return error, see
xla/xla/pjrt/interpreter_device.h
Line 26 in 60c9ecf
xla/xla/pjrt/pjrt_stream_executor_client.cc
Line 1364 in 60c9ecf
No?
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jul 4, 2024
…ut buffers Imported from GitHub PR openxla/xla#14328 This PR addresses @yliu120 's post-landing comments from openxla/xla#14088: 1. Error out on unknown memory_space in the output's layout (this requires a bit of plumbing to propagate the error). 2. Use `memory_space_by_kind_id` in place of `memory_space_by_kind` (this requires some ugly cast, not sure if it's worth it - performance should not matter much here). 3. Avoid shape copy, pass by reference. Copybara import of the project: -- b64a7f97e3cc080f3bba43ceac0f534467f1f2a6 by Jaroslav Sevcik <[email protected]>: Address reviewer comments Merging this change closes #14328 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14328 from jaro-sevcik:clean-up-after-pr-14088 b64a7f97e3cc080f3bba43ceac0f534467f1f2a6 PiperOrigin-RevId: 649324323
copybara-service bot
pushed a commit
that referenced
this pull request
Jul 4, 2024
…ut buffers Imported from GitHub PR #14328 This PR addresses @yliu120 's post-landing comments from #14088: 1. Error out on unknown memory_space in the output's layout (this requires a bit of plumbing to propagate the error). 2. Use `memory_space_by_kind_id` in place of `memory_space_by_kind` (this requires some ugly cast, not sure if it's worth it - performance should not matter much here). 3. Avoid shape copy, pass by reference. Copybara import of the project: -- b64a7f9 by Jaroslav Sevcik <[email protected]>: Address reviewer comments Merging this change closes #14328 COPYBARA_INTEGRATE_REVIEW=#14328 from jaro-sevcik:clean-up-after-pr-14088 b64a7f9 PiperOrigin-RevId: 649341102
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jul 4, 2024
…ut buffers Imported from GitHub PR openxla/xla#14328 This PR addresses @yliu120 's post-landing comments from openxla/xla#14088: 1. Error out on unknown memory_space in the output's layout (this requires a bit of plumbing to propagate the error). 2. Use `memory_space_by_kind_id` in place of `memory_space_by_kind` (this requires some ugly cast, not sure if it's worth it - performance should not matter much here). 3. Avoid shape copy, pass by reference. Copybara import of the project: -- b64a7f97e3cc080f3bba43ceac0f534467f1f2a6 by Jaroslav Sevcik <[email protected]>: Address reviewer comments Merging this change closes #14328 PiperOrigin-RevId: 649341102
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.