nvexec: `when_all` followed by `continues_on` does not compile

[romintomasetti]

It seems that the customization of the "combo" when_all/continues_on is broken for nvexec . But transfer_when_all is fine though.

I think that using when_all, followed by continues_on is legal. Please correct me if I am wrong.

I'm using clang-21 and 3363435. See reproducer below.

To compile the transfer_when_all example, use the following command:

docker buildx build --progress=plain --target=transfer-when-all --tag=test - < dockerfile && docker run --rm -it --privileged test

It compiles and runs just fine.

To get the compile error with continues_on, use the following command:

docker buildx build --progress=plain --target=when-all --tag=test - < dockerfile

It leads to:

In file included from /stdexec/include/nvexec/stream_context.cuh:28:
#12 3.935 /stdexec/include/nvexec/stream/schedule_from.cuh:184:21: error: static assertion failed
#12 3.935   184 |       static_assert(stream_completing_sender<Sender, Env>);

It feels like the schedule_from is expecting when_all to have a completion scheduler, which is wrong ? Not sure though.

FROM nvidia/cuda:12.8.1-devel-ubuntu24.04 AS base

RUN apt update
RUN apt --yes --no-install-recommends install git wget lsb-release gpg software-properties-common

ARG CLANG_VERSION=21

RUN <<EOF
    set -ex

    wget https://apt.llvm.org/llvm.sh
    chmod +x llvm.sh
    ./llvm.sh ${CLANG_VERSION}
    rm llvm.sh

    apt --yes --no-install-recommends install clang-21
EOF

RUN git clone --branch=main https://github.com/NVIDIA/stdexec
RUN cd stdexec && git checkout 3363435259b7ffae43d3f2e5f6b7a7b36d7cd7d3

FROM base AS when-all

COPY <<EOF test_when_all.cpp
#include "nvexec/stream_context.cuh"

int main()
{
    ::nvexec::stream_context stream_ctx{};

    ::stdexec::scheduler auto sch = stream_ctx.get_scheduler();

    auto fork = stdexec::schedule(sch) | stdexec::then([]{printf("Hi from 'then' 0.\n");}) | stdexec::split();

    auto snd = stdexec::when_all(
        fork | stdexec::bulk(::stdexec::par, 4, [](const auto item){printf("Hi from 'bulk' 1, item %d.\n", item);}),
        fork | stdexec::then(                   []{printf("Hi from 'then' 2.\n");}),
        fork | stdexec::bulk(::stdexec::par, 4, [](const auto item){printf("Hi from 'bulk' 3, item %d.\n", item);}))
             | stdexec::continues_on(sch)
             | stdexec::then(                   []{printf("Hi.\n");});
    stdexec::sync_wait(std::move(snd));
}
EOF

RUN clang++-${CLANG_VERSION} -x cuda --cuda-gpu-arch=sm_86 -std=c++20 -I/stdexec/include -I/usr/local/cuda/include/ test_when_all.cpp

FROM base AS transfer-when-all

COPY <<EOF test_transfer_when_all.cpp
#include "nvexec/stream_context.cuh"

int main()
{
    ::nvexec::stream_context stream_ctx{};

    ::stdexec::scheduler auto sch = stream_ctx.get_scheduler();

    auto fork = stdexec::schedule(sch) | stdexec::then([]{printf("Hi from 'then' 0.\\n");}) | stdexec::split();

    auto snd = stdexec::transfer_when_all(sch,
        fork | stdexec::bulk(::stdexec::par, 4, [](const auto item){printf("Hi from 'bulk' 1, item %d.\\n", item);}),
        fork | stdexec::then(                   []{printf("Hi from 'then' 2.\\n");}),
        fork | stdexec::bulk(::stdexec::par, 4, [](const auto item){printf("Hi from 'bulk' 3, item %d.\\n", item);}))
             | stdexec::then(                   []{printf("Hi from 'then' 4.\\n");});
    stdexec::sync_wait(std::move(snd));
}
EOF

RUN <<EOF
    set -ex

    clang++-${CLANG_VERSION} -x cuda --cuda-gpu-arch=sm_86 -std=c++20 \
        -I/stdexec/include \
        -I/usr/local/cuda/include/ \
        -L/usr/local/cuda/lib64 -lcudart \
        test_transfer_when_all.cpp
EOF

CMD ./a.out

[romintomasetti]

@ericniebler Any idea ?

[ericniebler]

this is definitely a bug. thanks for the report. i need some time to think of a fix.

[ericniebler]

auto fork = stdexec::schedule(sch) | ... | stdexec::split();

by the way, this isn't doing what you think. a split sender starts the work once and then caches the results. so only the first call to start on a split sender will see any side-effects, like scheduling onto a CUDA scheduler.

[romintomasetti]

Do I understand correctly that split caches the result values from the upstream sender, but it does not propagate the execution context (like the completion scheduler) to downstream senders.

So if I wanted the senders that I have written in when_all to actually execute on the GPU but express that they can run concurrently, I should be writing:

auto sndr = stdexec::when_all(
    fork | stdexec::continues_on(sch) | stdexec::bulk(...),
    fork | stdexec::continues_on(sch) | stdexec::then(...),
    fork | stdexec::continues_on(sch) | stdexec::bulk(...)
);

or use transfer_when_all?