Stream aware outputs #5684
Open
Stream aware outputs #5684
+596
−103
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
* Add output order handling to exec2 * Add CUDA stream to Outputs and SharedOutputs in Python bindings for Pipeline. * Refactor stream pointer handling in Python Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michał Zientkiewicz <[email protected]>
Signed-off-by: Michał Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michał Zientkiewicz <[email protected]>
Signed-off-by: Michał Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michał Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
mzient
commented
Oct 21, 2024
dali/pipeline/data/dltensor_obj.h
Outdated
There was a problem hiding this comment.
TODO: Remove this class entirely.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
CI MESSAGE: [19571767]: BUILD STARTED |
Signed-off-by: Michal Zientkiewicz <[email protected]>
mzient
force-pushed
the
stream_aware_output
branch
from
acb54b0
to
672f6f9
Compare
Signed-off-by: Michal Zientkiewicz <[email protected]>
|
CI MESSAGE: [19572115]: BUILD STARTED |
|
CI MESSAGE: [19572115]: BUILD FAILED |
Signed-off-by: Michal Zientkiewicz <[email protected]>
|
CI MESSAGE: [19600864]: BUILD STARTED |
|
CI MESSAGE: [19600864]: BUILD FAILED |
szkarpinski
reviewed
Oct 23, 2024
Comment on lines
+32
to
+33
| # convert the tensors in the batch to DLPack | ||
| batch = [torch.from_dlpack(t) for t in out] |
There was a problem hiding this comment.
Nit: This comment is slightly confusing: "convert to DLPack", but the function is from_dlpack. Maybe rephrase it a bit to indicate that we're not really converting to DLPack, but DALI->DLPack->Torch (without a copy)
| for t in batch: | ||
| means[flat_idx] = torch.mean(t) | ||
| flat_idx += 1 | ||
| # those are meant to overwrite the results if synchronization fails |
There was a problem hiding this comment.
maybe check that we're actually sharing the memory:
batch_a = [torch.from_dlpack(t) for t in out]
batch_b = [torch.from_dlpack(t) for t in out]
# now change batch_b and make sure that batch_a is changed as well
| assert jax_array.device() == jax.devices()[0] | ||
| assert jax_array.device == jax.devices()[0] |
There was a problem hiding this comment.
Slipped in... That's a breaking change... in JAX 0.4.31 :\
| #include "dali/core/static_switch.h" | ||
|
|
||
| namespace dali { | ||
|
|
||
| class DLTensorGraveyard { |
There was a problem hiding this comment.
Maybe add some docs why we need this and how it works?
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.
Category:
New feature (non-breaking change which adds functionality)
Refactoring (Redesign of existing code that doesn't affect functionality)
Description:
This PR adds support for returning the pipeline outputs as DLPack without copying.
Additional information:
Affected modules and functionalities:
This PR adds the following features:
__dlpack__interface for TensorsIt removes the
_expose_dlpack_capsuleas it was incomplete.Key points relevant for the review:
Tests:
Checklist
Documentation
DALI team only
Requirements
REQ IDs: N/A
JIRA TASK: DALI-4075