Push the streams context onto the driver stack when calling cumemcpy

[Jacobfaib]

Description

The driver memcpy entry points validate their device pointer arguments against the current context, NOT against the context the supplied stream belongs to. In a multi-device program the current context is frequently some other device's context (e.g. nothing was pushed, or a previous operation left a different device current), in which case a perfectly valid device pointer is rejected with CUDA_ERROR_INVALID_VALUE. Pushing the stream's own context around the call guarantees the device pointer is validated against the context that actually owns it.

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Improved asynchronous memory copies so device pointers are validated against the CUDA context associated with the selected stream.
  • Prevented temporary context changes from affecting the caller’s current context or context stack.
  • Improved reliability for asynchronous transfers involving multiple GPUs or streams created on different devices.

• Tests

  • Added coverage for cross-device stream context handling and context restoration after asynchronous memory copies.

Walkthrough

Changes

The async memcpy wrapper now temporarily uses the context associated with its stream. A multi-device C2H test verifies the copy and preservation of the current context and driver context stack.

Stream Context Async Copy

Layer / File(s)	Summary
Stream context guard and memcpy integration libcudacxx/include/cuda/__driver/driver_api.h	Adds an RAII guard that pushes the stream context before cuMemcpyAsync and restores the previous context afterward.
Multi-device context preservation test libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu	Adds a conditional two-device test that performs asynchronous device-to-device copying and checks current-context and driver-stack preservation.

Suggested reviewers: pciolkosz

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu (1)

89-90: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: Declare current_device and stream_device as const; neither variable is reassigned. The repository guideline requires all unmodified variables to be declared const.

Source: Coding guidelines

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9a20e460-20e7-4b2f-83f0-acc1afb27e15

📥 Commits

Reviewing files that changed from the base of the PR and between e7c9b65 and 155a9ed.

📒 Files selected for processing (2)

• libcudacxx/include/cuda/__driver/driver_api.h
• libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu

[davebayer]

We already have __ensure_current_context, can't you use that instead?

[Jacobfaib]

We already have __ensure_current_context, can't you use that instead?

No, because it's a circular include. We would need to move __ensure_current_context inside driver_api.h

[github-actions]

😬 CI Workflow Results

🟥 Finished in 1h 49m: Pass: 99%/120 | Total: 2d 18h | Max: 1h 24m | Hits: 57%/684313

See results here.

[pciolkosz]

I think it's the first time I use request changes, but I disagree with this direction. I don't think this belongs in the driver layer. There might be users of this that already ensured the context is pushed onto the stack for other reasons and it also duplicates the infrastructure.
I think the caller of the __driver APIs should be expected to manage the current context instead.
Long term I would like to make the driver layer available for other projects to use and I think we should keep this layer simpler. I would instead implement an internal wrapper, something between __memcpyAsync and cuda::copy_bytes that does that

[Jacobfaib]

I think the caller of the __driver APIs should be expected to manage the current context instead.

Normally I agree, but this is clearly a bug in the driver. People will run into this issue, and now you will find everyone will litter the code with

auto _ = __ensure_current_context{stream};

__driver::__cuMemcpy(...);

which will get cargo-culted ad infinitum. This will then extend to other APIs because users will defensively assume that because cuMemcpy() seemingly doesn't infer the context that other APIs might not either.

Or, even worse, someone will not know this bug exists and forget to do this context management leading to cryptic errors down the line. Maybe end users report this, but probably not because they will figure this is their bug not ours.

I don't think that is a better alternative.

[pciolkosz]

I understand, but the flipside is sometimes the potential user of the driver layer might want to pass a context that is different from the stream's context to the copy routine, because the driver can use either the stream's or the current context to perform the copy as an optimization. They would need to now create their own driver API wrapper, because our wrapper now blocks the driver from using any context other than the stream's context to perform the copy.

I think the current memcpy situation is quite bad and I filed an issue to the driver to improve it, but I believe for now the __memcpyAsync function we have is the best we can do to navigate the traps in memcpy behavior.