使用Thrust庫實現異步操作與回調函數

在Thrust庫中，你可以通過CUDA流(stream)來實現異步操作，并在適當的位置插入回調函數。以下是如何實現的詳細說明：

基本異步操作

Thrust本身并不直接暴露CUDA流接口，但你可以通過以下方式使用流：

#include <thrust/device_vector.h>
#include <thrust/transform.h>
#include <cuda_runtime.h>// 定義一個簡單的仿函數
struct saxpy_functor {float a;saxpy_functor(float _a) : a(_a) {}__host__ __device__float operator()(float x, float y) const {return a * x + y;}
};void async_thrust_operations() {// 創建CUDA流cudaStream_t stream;cudaStreamCreate(&stream);// 分配設備向量thrust::device_vector<float> X(10000, 1.0f);thrust::device_vector<float> Y(10000, 2.0f);thrust::device_vector<float> Z(10000);// 使用thrust::cuda::par.on(stream)指定執行流thrust::transform(thrust::cuda::par.on(stream),X.begin(), X.end(),Y.begin(), Z.begin(),saxpy_functor(2.0f));// 其他操作可以繼續在這里執行，因為上面的transform是異步的// 等待流完成cudaStreamSynchronize(stream);// 銷毀流cudaStreamDestroy(stream);
}

插入回調函數

要在CUDA流中插入回調函數，你可以使用cudaStreamAddCallback：

#include <iostream>// 回調函數
void CUDART_CB myCallback(cudaStream_t stream, cudaError_t status, void* userData) {std::cout << "CUDA callback executed!" << std::endl;// 可以在這里處理userData
}void async_with_callback() {cudaStream_t stream;cudaStreamCreate(&stream);thrust::device_vector<float> vec(1000);thrust::fill(thrust::cuda::par.on(stream), vec.begin(), vec.end(), 42.0f);// 插入回調函數cudaStreamAddCallback(stream, myCallback, nullptr, 0);// 其他操作...cudaStreamSynchronize(stream);cudaStreamDestroy(stream);
}

更復雜的回調示例

下面是一個更完整的示例，展示了如何傳遞數據給回調函數：

struct CallbackData {int value;float* d_ptr;
};void CUDART_CB complexCallback(cudaStream_t stream, cudaError_t status, void* userData) {CallbackData* data = static_cast<CallbackData*>(userData);std::cout << "Callback received value: " << data->value << std::endl;// 可以在這里處理設備指針data->d_ptr// 注意：回調函數在主機端執行，不能直接訪問設備內存delete data; // 清理分配的內存
}void advanced_async_example() {cudaStream_t stream;cudaStreamCreate(&stream);thrust::device_vector<float> vec(1000);// 準備回調數據CallbackData* cbData = new CallbackData{42, thrust::raw_pointer_cast(vec.data())};// 異步操作thrust::sequence(thrust::cuda::par.on(stream), vec.begin(), vec.end());// 添加回調cudaStreamAddCallback(stream, complexCallback, cbData, 0);// 其他操作可以繼續在這里執行cudaStreamSynchronize(stream);cudaStreamDestroy(stream);
}

注意事項

1. 回調函數在主機線程上執行，不是在GPU上執行
2. 回調函數中不能調用任何可能阻塞或等待CUDA完成的函數
3. 回調函數應該盡快完成，避免阻塞后續的操作
4. 傳遞給回調函數的數據需要手動管理生命周期
5. Thrust的并行算法默認使用默認流(stream 0)，要使用異步必須顯式指定流

通過這種方式，你可以在Thrust操作中實現異步執行并在適當的時候插入回調函數來處理完成事件。