[STF] Add python bindings

ci/build_cuda_cccl_experimental_python_experimental.sh

@@ -0,0 +1,130 @@
+#!/bin/bash
+set -euo pipefail
+
+ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+usage="Usage: $0 -py-version <python_version> [additional options...]"
+
+source "$ci_dir/util/python/common_arg_parser.sh"
+parse_python_args "$@"
+
+# Check if py_version was provided (this script requires it)
+require_py_version "$usage" || exit 1
+
+echo "Docker socket: " $(ls /var/run/docker.sock)
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  # Prepare mount points etc for getting artifacts in/out of the container.
+  source "$ci_dir/util/artifacts/common.sh"
+  action_mounts=$(cat <<EOF
+    --mount type=bind,source=${ARTIFACT_ARCHIVES},target=${ARTIFACT_ARCHIVES} \
+    --mount type=bind,source=${ARTIFACT_UPLOAD_STAGE},target=${ARTIFACT_UPLOAD_STAGE}
+EOF
+)
+
+else
+  action_mounts=""
+fi
+
+readonly cuda12_version=12.9.1
+readonly cuda13_version=13.0.2
+readonly devcontainer_version=26.02
+readonly devcontainer_distro=rockylinux8
+
+if [[ "$(uname -m)" == "aarch64" ]]; then
+  readonly cuda12_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${py_version}-arm64
+  readonly cuda13_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${py_version}-arm64
+else
+  readonly cuda12_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${py_version}
+  readonly cuda13_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${py_version}
+fi
+
+mkdir -p wheelhouse_experimental
+
+for ctk in 12 13; do
+  image=$(eval echo \$cuda${ctk}_image)
+  echo "::group::⚒️ Building CUDA ${ctk} experimental wheel on ${image}"
+  (
+    set -x
+    docker pull $image
+    docker run --rm -i \
+        --workdir /workspace/python/cuda_cccl_experimental \
+        --mount type=bind,source=${HOST_WORKSPACE},target=/workspace/ \
+        ${action_mounts} \
+        --env py_version=${py_version} \
+        --env GITHUB_ACTIONS=${GITHUB_ACTIONS:-} \
+        --env GITHUB_RUN_ID=${GITHUB_RUN_ID:-} \
+        --env JOB_ID=${JOB_ID:-} \
+        $image \
+        /workspace/ci/build_cuda_cccl_experimental_wheel.sh
+    # Prevent GHA runners from exhausting available storage with leftover images:
+    if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+      docker rmi -f $image
+    fi
+  )
+  echo "::endgroup::"
+done
+
+echo "Merging CUDA experimental wheels..."
+
+# Needed for unpacking and repacking wheels.
+python -m pip install wheel
+
+# Find the built wheels
+cu12_wheel=$(find wheelhouse_experimental -name "*cu12*.whl" | head -1)
+cu13_wheel=$(find wheelhouse_experimental -name "*cu13*.whl" | head -1)
+
+if [[ -z "$cu12_wheel" ]]; then
+  echo "Error: CUDA 12 experimental wheel not found in wheelhouse_experimental/"
+  ls -la wheelhouse_experimental/
+  exit 1
+fi
+
+if [[ -z "$cu13_wheel" ]]; then
+  echo "Error: CUDA 13 experimental wheel not found in wheelhouse_experimental/"
+  ls -la wheelhouse_experimental/
+  exit 1
+fi
+
+echo "Found CUDA 12 wheel: $cu12_wheel"
+echo "Found CUDA 13 wheel: $cu13_wheel"
+
+# Merge the wheels
+python python/cuda_cccl_experimental/merge_cuda_wheels.py "$cu12_wheel" "$cu13_wheel" --output-dir wheelhouse_experimental_merged
+
+# Install auditwheel and repair the merged wheel
+python -m pip install patchelf auditwheel
+for wheel in wheelhouse_experimental_merged/cuda_cccl_experimental-*.whl; do
+    echo "Repairing merged wheel: $wheel"
+    python -m auditwheel repair \
+        --exclude 'libnvrtc.so.12' \
+        --exclude 'libnvrtc.so.13' \
+        --exclude 'libnvJitLink.so.12' \
+        --exclude 'libnvJitLink.so.13' \
+        --exclude 'libcuda.so.1' \
+        "$wheel" \
+        --wheel-dir wheelhouse_experimental_final
+done
+
+# Clean up intermediate files and move only the final merged wheel
+rm -rf wheelhouse_experimental/*
+mkdir -p wheelhouse_experimental
+
+if ls wheelhouse_experimental_final/cuda_cccl_experimental-*.whl 1> /dev/null 2>&1; then
+    mv wheelhouse_experimental_final/cuda_cccl_experimental-*.whl wheelhouse_experimental/
+    echo "Final merged experimental wheel moved to wheelhouse_experimental"
+else
+    echo "No final repaired wheel found, moving unrepaired merged wheel"
+    mv wheelhouse_experimental_merged/cuda_cccl_experimental-*.whl wheelhouse_experimental/
+fi
+
+# Clean up temporary directories
+rm -rf wheelhouse_experimental_merged wheelhouse_experimental_final
+
+echo "Final experimental wheels in wheelhouse_experimental:"
+ls -la wheelhouse_experimental/
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  wheel_artifact_name="$(ci/util/workflow/get_wheel_artifact_name.sh)_experimental"
+  ci/util/artifacts/upload.sh $wheel_artifact_name 'wheelhouse_experimental/.*'
+fi

ci/build_cuda_cccl_experimental_wheel.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+set -euo pipefail
+
+# Target script for `docker run` command in build_cuda_cccl_experimental_python_experimental.sh
+# The /workspace pathnames are hard-wired here.
+
+# Install GCC 13 toolset (needed for the build)
+/workspace/ci/util/retry.sh 5 30 dnf -y install gcc-toolset-13-gcc gcc-toolset-13-gcc-c++
+echo -e "#!/bin/bash\nsource /opt/rh/gcc-toolset-13/enable" >/etc/profile.d/enable_devtools.sh
+source /etc/profile.d/enable_devtools.sh
+
+# Check what's available
+which gcc
+gcc --version
+which nvcc
+nvcc --version
+
+# Set up Python environment
+source /workspace/ci/pyenv_helper.sh
+setup_python_env "${py_version}"
+which python
+python --version
+echo "Done setting up python env"
+
+# Figure out the version to use for the package, we need repo history
+if $(git rev-parse --is-shallow-repository); then
+  git fetch --unshallow
+fi
+export PACKAGE_VERSION_PREFIX="0.1."
+package_version=$(/workspace/ci/generate_version.sh)
+echo "Using package version ${package_version}"
+# Override the version used by setuptools_scm to the custom version
+export SETUPTOOLS_SCM_PRETEND_VERSION_FOR_CUDA_CCCL_EXPERIMENTAL="${package_version}"
+
+cd /workspace/python/cuda_cccl_experimental
+
+# Determine CUDA version from nvcc
+cuda_version=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+' | cut -d. -f1)
+echo "Detected CUDA version: ${cuda_version}"
+
+# Configure compilers:
+export CXX="$(which g++)"
+export CUDACXX="$(which nvcc)"
+export CUDAHOSTCXX="$(which g++)"
+
+# Build the wheel
+python -m pip wheel --no-deps --verbose --wheel-dir dist .
+
+# Rename wheel to include CUDA version suffix
+for wheel in dist/cuda_cccl_experimental-*.whl; do
+    if [[ -f "$wheel" ]]; then
+        base_name=$(basename "$wheel" .whl)
+        new_name="${base_name}.cu${cuda_version}.whl"
+        mv "$wheel" "dist/${new_name}"
+        echo "Renamed wheel to: ${new_name}"
+    fi
+done
+
+# Move wheel to output directory
+mkdir -p /workspace/wheelhouse_experimental
+mv dist/cuda_cccl_experimental-*.cu*.whl /workspace/wheelhouse_experimental/

ci/matrix.yaml

@@ -70,6 +70,8 @@ workflows:
     - {jobs: ['test'], project: 'python', ctk: ['12.X',        '13.X'], py_version: ['3.10'], gpu: 'l4', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X','13.0', '13.X'], py_version: ['3.13'], gpu: 'l4', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
+    # Python experimental (STF) -- pinned to gcc13, Linux only
+    - {jobs: ['test'], project: 'python_experimental', ctk: ['12.X', '13.X'], py_version: ['3.13'], gpu: 'l4', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
     - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'}
@@ -118,6 +120,7 @@ workflows:
     - {project: 'cccl_c_parallel', jobs: ['test'], ctk: '13.X', cxx: ['gcc13', 'msvc'], gpu: 'rtx2080', sm: 'gpu'}
     - {project: 'cccl_c_stf',      jobs: ['test'], ctk: '13.X', cxx: 'gcc13',           gpu: 'rtx2080', sm: 'gpu'}
     - {project: 'python', jobs: ['test'], ctk: '13.X', py_version: '3.13', gpu: 'l4', cxx: ['gcc13', 'msvc']}
+    - {project: 'python_experimental', jobs: ['test'], ctk: '13.X', py_version: '3.13', gpu: 'l4', cxx: 'gcc13'}
     # Packaging / install
     - {project: 'packaging', jobs: ['test'], ctk: '13.X', cxx: ['gcc', 'clang'], gpu: 'rtx2080', sm: 'gpu'}
     - {project: 'packaging', jobs: ['test'], args: '-min-cmake', gpu: 'rtx2080', sm: 'gpu'}
@@ -192,6 +195,9 @@ workflows:
     # Python -- pinned to gcc13 on Linux for consistency across CTK images
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13'], gpu: 'l4', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
+    # Python experimental (STF) -- pinned to gcc13, Linux only
+    - {jobs: ['test'], project: 'python_experimental', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test'], project: 'python_experimental', ctk: '13.X', py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
     - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'}
@@ -279,6 +285,9 @@ workflows:
     # Python -- pinned to gcc13 for consistency across CTK images
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13'], gpu: 'l4', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.13', gpu: 'h100', cxx: ['gcc13', 'msvc']}
+    # Python experimental (STF) -- pinned to gcc13, Linux only
+    - {jobs: ['test'], project: 'python_experimental', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test'], project: 'python_experimental', ctk: '13.X', py_version: '3.13', gpu: 'h100', cxx: 'gcc13'}
     # CCCL packaging:
     - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080', args: '-min-cmake'}
     - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'}
@@ -302,6 +311,7 @@ workflows:
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13'], gpu: 'l4', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.13', gpu: 'h100', cxx: ['gcc13', 'msvc']}
     - {jobs: ['test'], project: 'python', cpu: 'arm64', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13'], gpu: 'l4', cxx: 'gcc13'}
+    - {jobs: ['test'], project: 'python_experimental', ctk: ['12.X', '13.X'], py_version: ['3.10', '3.13'], gpu: 'l4', cxx: 'gcc13'}
 
 
   # This is just used to ensure that we generate devcontainers for all images we build.
@@ -325,6 +335,8 @@ workflows:
   exclude:
     # GPU runners are not available on Windows.
     - {jobs: ['test', 'test_gpu', 'test_nolid', 'test_lid0', 'test_lid1', 'test_lid2'], cxx: ['msvc2019', 'msvc14.39', 'msvc2022']}
+     # STF experimental Python bindings are not built for MSVC:
+    - {project: 'python_experimental', cxx: ['msvc2019', 'msvc14.39', 'msvc2022']}
     # cudax doesn't support C++17 on msvc:
     - {project: 'cudax', std: 17, cxx: ['msvc2019', 'msvc14.39', 'msvc2022']}
 
@@ -478,6 +490,9 @@ jobs:
   test_py_coop:     { name: "Test cuda.coop",  gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_coop'} }
   test_py_par:      { name: "Test cuda.compute",     gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_compute'} }
   test_py_examples: { name: "Test cuda.cccl.examples",     gpu: true,  needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_cccl_examples'} }
+  # Python experimental (cuda-cccl-experimental wheel):
+  build_py_experimental_wheel: { name: "Build cuda.cccl.experimental", gpu: false, invoke: { prefix: 'build_cuda_cccl_experimental'} }
+  test_py_stf:      { name: "Test cuda.stf",               gpu: true,  needs: 'build_py_experimental_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_stf'} }
 
   # Run jobs for 'target' project (ci/util/build_and_test_targets.sh):
   run_cpu: { gpu: false }
@@ -536,6 +551,11 @@ projects:
     job_map:
       build: ['build_py_wheel']
       test:  ['test_py_headers', 'test_py_coop', 'test_py_par', 'test_py_examples']
+  python_experimental:
+    name: "Python Experimental"
+    job_map:
+      build: ['build_py_experimental_wheel']
+      test:  ['test_py_stf']
   cccl_c_parallel:
     name: 'CCCL C Parallel'
     stds: [20]

ci/project_files_and_dependencies.yaml

@@ -128,8 +128,18 @@ projects:
     lite_dependencies: [cccl_c_parallel_public]
     full_dependencies: []
     include_regexes:
-    - "python/"
+    - "python/cuda_cccl/"
     - "pyproject.toml"
+    exclude_regexes:
+    - "python/cuda_cccl_experimental/"
+
+  python_experimental:
+    name: "Python Experimental"
+    matrix_project: "python_experimental"
+    lite_dependencies: [cccl_c_stf, python]
+    full_dependencies: []
+    include_regexes:
+    - "python/cuda_cccl_experimental/"
 
   packaging:
     name: "CCCL Packaging"

ci/test/inspect_changes/c2h_dependency.output

@@ -1,2 +1,2 @@
 FULL_BUILD=
-LITE_BUILD=libcudacxx cub cudax cccl_c_parallel cccl_c_stf packaging
+LITE_BUILD=libcudacxx cub cudax cccl_c_parallel cccl_c_stf python_experimental packaging

ci/test/inspect_changes/core_dirty.output

@@ -1,2 +1,2 @@
-FULL_BUILD=libcudacxx cub thrust cudax cccl_c_parallel cccl_c_stf python packaging stdpar nvbench_helper nvrtcc
+FULL_BUILD=libcudacxx cub thrust cudax cccl_c_parallel cccl_c_stf python python_experimental packaging stdpar nvbench_helper nvrtcc
 LITE_BUILD=

ci/test/inspect_changes/libcudacxx_both.output

@@ -1,2 +1,2 @@
 FULL_BUILD=libcudacxx
-LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_stf python packaging stdpar nvbench_helper
+LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_stf python python_experimental packaging stdpar nvbench_helper

ci/test/inspect_changes/libcudacxx_public_only.output

@@ -1,2 +1,2 @@
 FULL_BUILD=libcudacxx
-LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_stf python packaging stdpar nvbench_helper
+LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_stf python python_experimental packaging stdpar nvbench_helper

ci/test/inspect_changes/libcudacxx_thrust.output

@@ -1,2 +1,2 @@
 FULL_BUILD=libcudacxx thrust
-LITE_BUILD=cub cudax cccl_c_parallel cccl_c_stf python packaging stdpar nvbench_helper
+LITE_BUILD=cub cudax cccl_c_parallel cccl_c_stf python python_experimental packaging stdpar nvbench_helper

ci/test/inspect_changes/multiple_projects.output

@@ -1,2 +1,2 @@
 FULL_BUILD=python packaging
-LITE_BUILD=
+LITE_BUILD=python_experimental

ci/test_cuda_stf_python_experimental.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+set -euo pipefail
+
+ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "$ci_dir/pyenv_helper.sh"
+
+# Parse common arguments
+source "$ci_dir/util/python/common_arg_parser.sh"
+parse_python_args "$@"
+cuda_major_version=$(nvcc --version | grep release | awk '{print $6}' | tr -d ',' | cut -d '.' -f 1 | cut -d 'V' -f 2)
+
+# Setup Python environment
+setup_python_env "${py_version}"
+
+# Fetch or build the cuda_cccl wheel (base dependency):
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  wheel_artifact_name=$("$ci_dir/util/workflow/get_wheel_artifact_name.sh")
+  "$ci_dir/util/artifacts/download.sh" ${wheel_artifact_name} /home/coder/cccl/
+else
+  "$ci_dir/build_cuda_cccl_python.sh" -py-version "${py_version}"
+fi
+
+# Install cuda_cccl base wheel
+CUDA_CCCL_WHEEL_PATH="$(ls /home/coder/cccl/wheelhouse/cuda_cccl-*.whl)"
+python -m pip install "${CUDA_CCCL_WHEEL_PATH}[cu${cuda_major_version}]"
+
+# Fetch or build the experimental wheel:
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  experimental_artifact_name="${wheel_artifact_name}_experimental"
+  "$ci_dir/util/artifacts/download.sh" ${experimental_artifact_name} /home/coder/cccl/
+else
+  "$ci_dir/build_cuda_cccl_experimental_python_experimental.sh" -py-version "${py_version}"
+fi
+
+# Install cuda_cccl_experimental wheel
+EXPERIMENTAL_WHEEL_PATH="$(ls /home/coder/cccl/wheelhouse_experimental/cuda_cccl_experimental-*.whl)"
+python -m pip install "${EXPERIMENTAL_WHEEL_PATH}[test-cu${cuda_major_version}]"
+
+# Run tests for STF module
+cd "/home/coder/cccl/python/cuda_cccl_experimental/tests/"
+python -m pytest -n auto -v stf/

docs/python/index.rst

@@ -14,6 +14,9 @@ abstractions for CUDA Python developers.
 * :doc:`cuda.coop <coop>` — Cooperative block- and warp-level algorithms for
   writing highly efficient CUDA kernels with `Numba CUDA <https://nvidia.github.io/numba-cuda/>`_.
 
+* :doc:`cuda.stf <stf>` — Sequential Task Flow for CUDA: define logical data and
+  tasks with read/write annotations; STF orchestrates execution and data movement.
+
 These libraries expose the generic, highly-optimized algorithms from the
 `CCCL C++ libraries <https://nvidia.github.io/cccl/cpp.html>`_,
 which have been tuned to provide optimal performance across GPU architectures.
@@ -34,5 +37,6 @@ Who is this for?
    setup
    compute
    coop
+   stf
    resources
    api_reference