Stream aware outputs #5684
Stream aware outputs #5684
Conversation
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 |
mzient
force-pushed
the
stream_aware_output
branch
from
acb54b0
to
672f6f9
Compare
|
CI MESSAGE: [19572115]: BUILD STARTED |
|
CI MESSAGE: [19572115]: BUILD FAILED |
|
CI MESSAGE: [19600864]: BUILD STARTED |
|
CI MESSAGE: [19600864]: BUILD FAILED |
| # 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?
mzient
force-pushed
the
stream_aware_output
branch
2 times, most recently
from
9b42d28
to
71854e7
Compare
* 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 <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
mzient
force-pushed
the
stream_aware_output
branch
from
71854e7
to
ee58ddf
Compare
Signed-off-by: Michał Zientkiewicz <mzient@gmail.com>
Signed-off-by: Michał Zientkiewicz <mzient@gmail.com>
|
CI MESSAGE: [19878075]: BUILD STARTED |
|
CI MESSAGE: [19878075]: BUILD FAILED |
Adjust dlpack tests. Signed-off-by: Michał Zientkiewicz <mzient@gmail.com>
|
CI MESSAGE: [19899374]: BUILD STARTED |
|
CI MESSAGE: [19899374]: BUILD PASSED |
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