Uh oh!
There was an error while loading. Please reload this page.
Use caller CUDA stream for D2H and H2D copies (#20498) - #20498
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20498
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit d3f5462 with merge base b331ebd ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Contributor
@Conarnar has exported this pull request. If you are a Meta employee, you can view the originating Diff in D109590531. |
This PR needs a |
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 24, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync` and synchronize the stream before returning — preserving the blocking API contract while allowing work to be issued on the caller's stream. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Conarnarforce-pushed
the
export-D109590531
branch
2 times, most recently
from
June 25, 2026 17:10
3d8da75 to
07765c3CompareConarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 25, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync` and synchronize the stream before returning — preserving the blocking API contract while allowing work to be issued on the caller's stream. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Comment on lines
+161
to
+168
| cudaError_t err = cudaSuccess; | ||
| const auto caller_stream = executorch::extension::cuda::getCallerStream(); | ||
| if (caller_stream) { | ||
| err = cudaMemcpyAsync( | ||
| dst, src, nbytes, cudaMemcpyHostToDevice, *caller_stream); | ||
| } else { | ||
| err = cudaMemcpy(dst, src, nbytes, cudaMemcpyHostToDevice); | ||
| } |
Comment on lines
+144
to
+150
| // TODO: validate caller stream device matches index. | ||
| // For now assert single-GPU case. | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| index == -1 || index == 0, | ||
| InvalidArgument, | ||
| "CudaAllocator::copy_host_to_device only supports device 0, got %d", | ||
| static_cast<int>(index)); |
Comment on lines
+202
to
+208
| // TODO: validate caller stream device matches index. | ||
| // For now assert single-GPU case. | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| index == -1 || index == 0, | ||
| InvalidArgument, | ||
| "CudaAllocator::copy_device_to_host only supports device 0, got %d", | ||
| static_cast<int>(index)); |
Comment on lines
+78
to
+90
| cudaStream_t s; | ||
| ASSERT_EQ(cudaStreamCreate(&s), cudaSuccess); | ||
| executorch::extension::cuda::CallerStreamGuard g(s); | ||
| CudaAllocator& a = CudaAllocator::instance(); | ||
| auto res = a.allocate(256, 0); | ||
| ASSERT_TRUE(res.ok()); | ||
| void* d = res.get(); | ||
| std::vector<uint8_t> h(256, 7); | ||
| // should take async branch internally, still return Ok | ||
| EXPECT_EQ(a.copy_host_to_device(d, h.data(), 256, 0), Error::Ok); | ||
| a.deallocate(d, 0); | ||
| cudaStreamDestroy(s); |
Comment on lines
+103
to
+117
| cudaStream_t s; | ||
| ASSERT_EQ(cudaStreamCreate(&s), cudaSuccess); | ||
| executorch::extension::cuda::CallerStreamGuard g(s); | ||
| CudaAllocator& a = CudaAllocator::instance(); | ||
| auto res = a.allocate(256, 0); | ||
| ASSERT_TRUE(res.ok()); | ||
| void* d = res.get(); | ||
| std::vector<uint8_t> h_src(256, 5), h_dst(256, 0); | ||
| ASSERT_EQ(a.copy_host_to_device(d, h_src.data(), 256, 0), Error::Ok); | ||
| EXPECT_EQ(a.copy_device_to_host(h_dst.data(), d, 256, 0), Error::Ok); | ||
| EXPECT_EQ(h_src, h_dst); | ||
| a.deallocate(d, 0); | ||
| cudaStreamDestroy(s); |
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 25, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 25, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Comment on lines
+144
to
+150
| // TODO: validate caller stream device matches index. | ||
| // For now assert index is -1 or 0. | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| index == -1 || index == 0, | ||
| InvalidArgument, | ||
| "CudaAllocator::copy_host_to_device only supports device 0 or -1 (current), got %d", | ||
| static_cast<int>(index)); |
Comment on lines
+204
to
+210
| // TODO: validate caller stream device matches index. | ||
| // For now assert index is -1 or 0. | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| index == -1 || index == 0, | ||
| InvalidArgument, | ||
| "CudaAllocator::copy_device_to_host only supports device 0 or -1 (current), got %d", | ||
| static_cast<int>(index)); |
Comment on lines
+161
to
+166
| cudaError_t err = cudaSuccess; | ||
| const auto caller_stream = executorch::extension::cuda::getCallerStream(); | ||
| if (caller_stream) { | ||
| err = cudaMemcpyAsync( | ||
| dst, src, nbytes, cudaMemcpyHostToDevice, *caller_stream); | ||
| // We don't synchronize the stream here because the caller is expected to |
Comment on lines
+223
to
+228
| if (caller_stream) { | ||
| err = cudaMemcpyAsync( | ||
| dst, src, nbytes, cudaMemcpyDeviceToHost, *caller_stream); | ||
| if (err == cudaSuccess) { | ||
| err = cudaStreamSynchronize(*caller_stream); | ||
| } |
Comment on lines
+116
to
+118
| ASSERT_EQ(a.copy_host_to_device(d, h_src.data(), 256, 0), Error::Ok); | ||
| EXPECT_EQ(a.copy_device_to_host(h_dst.data(), d, 256, 0), Error::Ok); | ||
| EXPECT_EQ(h_src, h_dst); |
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 25, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 25, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 27, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 27, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnarforce-pushed
the
export-D109590531
branch
2 times, most recently
from
June 27, 2026 10:24
32968c0 to
056c25cCompareConarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 27, 2026
Summary: Pull Request resolved: pytorch#20498 CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 27, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 27, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 28, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 28, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 28, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. - Added CIs for unit tests. Reviewed By: Gasoonjia Differential Revision: D109590531
Conarnar added a commit
to Conarnar/executorch
that referenced
this pull request
Jun 28, 2026
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. - Added CIs for unit tests. Reviewed By: Gasoonjia Differential Revision: D109590531
Summary: CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via `getCallerStream()`), `copy_host_to_device` and `copy_device_to_host` use `cudaMemcpyAsync`. When no caller stream is set, the synchronous `cudaMemcpy` path is used as before. Additionally: - Added null pointer and zero-byte validation — null `dst`/`src` return `Error::InvalidArgument` instead of aborting in `cudaMemcpy`, and zero-byte copies return `Error::Ok` early. - Assert single-GPU case (index 0 or -1) until multi-GPU stream validation is added. - Wired `//executorch/extension/cuda:caller_stream` dependency in TARGETS. - Added `extension_cuda` dependencies to CMakeLists.txt. - Added `test_cuda_allocator` with coverage for sync/async paths and error handling. - Added CIs for unit tests. Reviewed By: Gasoonjia Differential Revision: D109590531
This file contains hidden or 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
Sign up for freeto 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.
Summary:
CudaAllocator memory copies now support async copy on a caller-provided CUDA stream. When a caller stream is available (via
getCallerStream()),copy_host_to_deviceandcopy_device_to_hostusecudaMemcpyAsync. When no caller stream is set, the synchronouscudaMemcpypath is used as before.Additionally:
dst/srcreturnError::InvalidArgumentinstead of aborting incudaMemcpy, and zero-byte copies returnError::Okearly.//executorch/extension/cuda:caller_streamdependency in TARGETS.extension_cudadependencies to CMakeLists.txt.test_cuda_allocatorwith coverage for sync/async paths and error handling.Reviewed By: Gasoonjia
Differential Revision: D109590531