[alwaysprince05]

Follow-up to #1563 (see #1563 (comment)).

With continues_on between two stream schedulers, the work before the transition borrowed the destination's stream, so both sides ran on the destination:

schedule(ctx_0.get_scheduler())
  | then(a)                              // ran on ctx_1's stream
  | continues_on(ctx_1.get_scheduler())
  | then(b);                             // ctx_1's stream

Now, if the source and destination schedulers compare unequal (another context, or the same context with another priority), the work before the transition gets its own stream from the source scheduler. Equal schedulers still share one stream.

No extra synchronization is added. The completion already reaches continues_on through schedule_from's task queue, which only runs after the source stream got there.

Tests (in test/nvexec/continues_on.cpp):

• a single-device version of @maartenarnst's test: two contexts on one GPU must record different streams before and after the transition
• low to high priority on one context switches streams, and each side runs at its own priority
• the step after the transition sees memory written by a slow step before it
• continues_on within one context keeps the same stream

Tested on a T4 (clang 20, CUDA 12.9): the switch and priority tests fail on main and pass with this change, all [continues_on] tests pass, the ordering test passed 20/20 runs, and the full test.nvexec passes 3/3 runs (124 test cases). compute-sanitizer memcheck is the same as on main: no leaks, and the only reported errors are the two intentional ones in "continues_on destroys host-constructed storage after a CUDA error".

I only have a single GPU, so I haven't tested contexts on two different devices.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maartenarnst]

These are a few bigger picture questions/thoughts:

• What should determine the "switching"? Is it the streams being different between two segments? Or the schedulers, or the contexts, or the devices? I think that in nvexec, a context can create schedulers with different priorities.
• How should the current polling mechanism coexist/be replaced with the event mechanism you propose in this pr?
• How should we think about making a fix in continues_on versus making a fix in schedule_from?

To move forward, it may make sense to split the pr into two separate pieces. A first pr would bring the fix to bring the switching behavior. And then a second pr would address the event mechanism.

[alwaysprince05]

Thanks for looking at this.

On what should decide the switch: I think the scheduler. stream_scheduler::operator== already compares the hub and the priority, so if the two schedulers compare unequal, the work before continues_on should stay on the source scheduler's stream. You're right about priorities, the first version only compared the contexts, so ctx.get_scheduler(low) to ctx.get_scheduler(high) ran everything on the high priority stream. The device alone is not enough since your case is two contexts on one GPU, and the streams are only known after connect, so they can't decide it.

On polling vs events: I looked at this again and the event is not needed for this fix. continues_on(sndr, sched) wraps the child in schedule_from, and in nvexec that goes through the task queue and the poller. The task is pushed from the source stream, so by the time the poller hands the values to continues_on the work on the source stream is already done. So in the current code the event was only on top of the host round trip. Replacing that round trip with an event when both sides are stream schedulers could be a separate change, and schedule_from to a non-stream receiver would keep polling.

On continues_on vs schedule_from: only continues_on knows both schedulers, so I kept the check there. schedule_from stays the generic "leave the stream" step. I know #2277 also touches continues_on, so I'm fine adjusting this to it.

So I followed your suggestion and reduced this PR to the switching only:

• compare schedulers instead of stream pools
• removed the event
• added a test for the priority case and one where the step after the transition reads memory written by a slow step before it

Ran it on a T4 again (clang 20, CUDA 12.9): main with the new tests: the switch test and the priority test fail. With the change: all 17 [continues_on] tests pass, the ordering test passed 20 out of 20 runs, and the full test.nvexec passes 3/3 runs (124 test cases). compute-sanitizer memcheck shows 0 errors except the two intentional ones in "continues_on destroys host-constructed storage after a CUDA error", same as on main.

If the event part is wanted I can open it as a separate PR after this one.

[ericniebler]

/ok to test 3e0d3b2

[ericniebler]

thanks!