Fix non-deterministic CUDA custom-op test by sequencing H2D copy on the compute stream

[Copilot]

Description

MyCustomKernelSecondInputOnCpu::Compute copied its CPU input to the GPU with a synchronous cudaMemcpy (default null stream) but launched cuda_add on the session's compute stream. When that stream is non-blocking, there is no ordering between the copy and the add, so cuda_add could read Y_cuda before the copy finished — causing CApiTest.custom_op_set_input_memory_type to fail intermittently.

• Replaced cudaMemcpy with cudaMemcpyAsync issued on the same compute stream used by cuda_add, so the copy is ordered before the add.
• Both the copy and the add now share one stream variable; the trailing cudaFree(Y_cuda) still provides implicit device sync, so the buffer is not freed early.

cudaStream_t compute_stream = compute_stream_ == nullptr ? 0 : reinterpret_cast<cudaStream_t>(compute_stream_);
cudaMemcpyAsync(Y_cuda, Y, y_size * sizeof(float), cudaMemcpyHostToDevice, compute_stream);
...
cuda_add(size, out, X, Y_cuda, compute_stream);

Motivation and Context

onnxruntime_shared_lib_test failed ~40% of the time on CUDA builds (--gtest_filter=CApiTest.custom_op_set_input_memory_type --gtest_repeat=20). The flakiness stemmed from the missing stream ordering in the test custom op rather than a library defect; sequencing the copy on the compute stream removes the race.

[tianleiwu]

Verdict: Approve — clean, correct, minimal fix.

Correctness

The root-cause analysis is accurate. A host→device cudaMemcpy from pageable memory returns to the host once the source is staged, but the actual DMA is enqueued on the null stream. Because cuda_add is launched on compute_stream_ (which may be created with cudaStreamNonBlocking), the null-stream DMA and the add kernel were unordered — exactly the race that made CApiTest.custom_op_set_input_memory_type flaky. Switching to cudaMemcpyAsync on the same compute_stream makes the copy and the add stream-ordered, which is the right fix.

Lifetime / teardown

Hoisting the stream into a single compute_stream variable and reusing it for both the copy and the add is correct. With the copy now asynchronous, both the device buffer Y_cuda and the host source Y must stay valid until the stream work completes — the trailing cudaFree(Y_cuda) performs an implicit full-device synchronization, so neither is released early. Y is owned by ORT for the duration of Compute, so that holds too.

Minor (optional, non-blocking, pre-existing)

None of the CUDA calls in this function (cudaMalloc, cudaMemcpyAsync, cudaFree) check their return status. This predates the PR and this is test/sample custom-op code, so it's out of scope — just noting it for completeness.

No blocking concerns.

[Copilot]

Pull request overview

This PR fixes intermittent failures in the CUDA shared-lib custom-op test by ensuring the host-to-device copy of the CPU-resident second input is enqueued on the same (potentially non-blocking) compute stream as the subsequent cuda_add kernel, guaranteeing correct stream ordering.

Changes:

• Replace synchronous cudaMemcpy (default stream) with cudaMemcpyAsync on the session/custom-op compute stream.
• Reuse a single compute_stream variable for both the H2D copy and the CUDA kernel launch to avoid mismatched stream selection.