https://nvidia.github.io/cccl/cudax/stf.html

2025/05 CUDASTF属于[CCCL (CUDA Core Compute Libraries) > cudax (CUDA Experimental)]
所以它的 Path 是：https://github.com/NVIDIA/cccl/tree/main/cudax

因为需要用cuda::std，所以需要手动include libcudacxx（也包含在cccl中）

CUDASTF是Header-only的lib，源代码包含在cccl中，clone即可使用

# 编译
nvcc -std=c++17 --expt-relaxed-constexpr --extended-lambda -I/home/zh-ge/git/cccl/libcudacxx/include -I/home/zh-ge/git/cccl/cudax/include -lcuda main.cu

// Generated by AI
 
// 通用 task：可以执行任意 CUDA 代码，包括多个 kernel launch、host-side 逻辑等
// lambda 会立即运行，调用 kernel launch 并将它入队到 s 所指定的 CUDA stream 中
ctx.task(lX.read(), lY.rw())->*[&](cudaStream_t s, auto dX, auto dY) {
  axpy<<<16, 128, 0, s>>>(alpha, dX, dY);
};
 
// 专用链式 kernel launch：将多个 kernel launch 合并为一个 task，自动串行执行
ctx.cuda_kernel_chain(lX.read(), lY.rw())->*[&](auto dX, auto dY) {
  return std::vector<cuda_kernel_desc> {
    { axpy, 16, 128, 0, alpha, dX, dY },
    { axpy, 16, 128, 0, beta,  dX, dY },
    { axpy, 16, 128, 0, gamma, dX, dY }
  };
};
 
// 单 kernel 优化接口：针对执行一个 kernel 的场景，内部直接发起 kernel，无中间开销
ctx.cuda_kernel(lX.read(), lY.rw())->*[&](auto dX, auto dY) {
  return cuda_kernel_desc{axpy, 16, 128, 0, alpha, dX, dY};
};
 
// 自动生成 CUDA kernel：你在 lambda 中写类似逐元素操作，STF 会自动编译成 device kernel
ctx.launch(lX.read(), lY.rw())->*[=] _CCCL_DEVICE(auto t, auto dX, auto dY) {
  for (auto ind : t.apply_partition(shape(dX))) {
    dY(ind) += alpha * dX(ind);
  }
};
 
// 最简洁、类型安全的迭代接口：通过 index 和 capture 参数写 device logic，STF 根据 shape 自动调度执行
ctx.parallel_for(lY.shape(), lX.read(), lY.rw())->*[alpha] __device__(size_t i, auto dX, auto dY) {
  dY(i) += alpha * dX(i);
};

// Host端的callback
// 这个lambda不会立即执行
ctx.host_launch(lX.read(), lY.read())->*[](auto sX, auto sY) {
    // 此处 garantueed 数据已经从 GPU 拷回到主内存
    callback_process(sX, sY);
};

TODO

TODO

除了CUDA原本的调试方法外，CUDASTF有一些其他的调试方式

# 执行./a.out 生成DAG图（dot文件）
# 直接在Terminal里面执行这句命令就行，虽然我也不太清楚这是什么语法，是设置环境变量吗？
CUDASTF_DOT_FILE=main.dot ./a.out

# 将dot格式绘制为png图像
dot -Tpng main.dot -o main.png

# 执行./a.out 生成DAG图（为不同device标注颜色、绘制Task以外的节点，比如memcpy）
CUDASTF_DOT_COLOR_BY_DEVICE=1 CUDASTF_DOT_IGNORE_PREREQS=0 CUDASTF_DOT_FILE=main2.dot ./a.out

# 生成结果中，task fence表示「同步屏障任务」，基本上可以理解为「同步节点」？

# 未测试
export CUDA_STF_DEBUG=1

1. 可以通过选择不同的 Context Backend 来决定使用 Streams 还是 CUDA Graphs 模式

#include <iostream>
#include <cstdint>
#include <vector>
 
#include <cuda/experimental/stf.cuh>
 
using namespace cuda::experimental::stf;
 
const int32_t N = 2e8;
 
__global__ void add(slice<const double> a1, slice<const double> a2, slice<double> res) {
    int32_t tid      = blockIdx.x * blockDim.x + threadIdx.x;
    int32_t nthreads = gridDim.x * blockDim.x;
    for(int32_t i = tid; i < res.size(); i += nthreads) {
        res(i) = a1(i) + a2(i);
    }
}
 
__global__ void sqrt(slice<double> a) {
    int32_t tid      = blockIdx.x * blockDim.x + threadIdx.x;
    int32_t nthreads = gridDim.x * blockDim.x;
    for(int32_t i = tid; i < a.size(); i += nthreads) {
        a(i) = sqrt(a(i));
    }
}
 
inline void test() {
    context ctx;
    std::vector<double> X(N, 4.0), Y(N, 4.0), Z(N, 4.0), S(N, 0.0);
 
    auto lX = ctx.logical_data(make_slice(X.data(), N));
    auto lY = ctx.logical_data(make_slice(Y.data(), N));
    auto lZ = ctx.logical_data(make_slice(Z.data(), N));
    auto lS = ctx.logical_data(make_slice(S.data(), N));
 
    ctx.task(exec_place::device(0), lX.rw()).set_symbol("sqrt(X)")->*[&](cudaStream_t s, auto dX) {
        sqrt<<<16, 128, 0, s>>>(dX);
    };
    ctx.task(exec_place::device(1), lY.rw()).set_symbol("sqrt(Y)")->*[&](cudaStream_t s, auto dY) {
        sqrt<<<16, 128, 0, s>>>(dY);
    };
    ctx.task(exec_place::device(0), lZ.rw()).set_symbol("sqrt(Z)")->*[&](cudaStream_t s, auto dZ) {
        sqrt<<<16, 128, 0, s>>>(dZ);
    };
    ctx.task(exec_place::device(0), lX.read(), lY.read(), lS.rw()).set_symbol("S = X + Y")->*[&](cudaStream_t s, auto dX, auto dY, auto dS) {
        add<<<16, 128, 0, s>>>(dX, dY, dS);
    };
    ctx.task(exec_place::device(1), lZ.read(), lS.rw()).set_symbol("S = S + Z")->*[&](cudaStream_t s, auto dZ, auto dS) {
        add<<<16, 128, 0, s>>>(dZ, dS, dS);
    };
 
    ctx.finalize();
}
 
int main() {
    test();
    return 0;
}