CUDA:如何使用推力进行矩阵乘法？

#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/transform.h>
#include <thrust/functional.h>

using namespace thrust;

const int N = 1024;  // 矩阵维度
const int SIZE = N * N;  // 矩阵大小

struct matrix_multiply : public thrust::binary_function<float, float, float>
{
    const int N;  // 矩阵维度

    matrix_multiply(int _N) : N(_N) {}

    __host__ __device__
    float operator()(const float& a, const float& b) const
    {
        int row = threadIdx.y + blockIdx.y * blockDim.y;
        int col = threadIdx.x + blockIdx.x * blockDim.x;

        float sum = 0.0f;
        for (int i = 0; i < N; ++i)
        {
            float element_a = a[row * N + i];
            float element_b = b[i * N + col];
            sum += element_a * element_b;
        }

        return sum;
    }
};

device_vector<float> d_input_a(SIZE);
device_vector<float> d_input_b(SIZE);
device_vector<float> d_output(SIZE);

host_vector<float> h_input_a(SIZE);
host_vector<float> h_input_b(SIZE);

// 填充输入矩阵数据
// ...

// 将输入数据从主机内存复制到设备内存
d_input_a = h_input_a;
d_input_b = h_input_b;

dim3 block_size(16, 16);
dim3 grid_size(N / block_size.x, N / block_size.y);

transform(
    make_zip_iterator(make_tuple(d_input_a.begin(), d_input_b.begin())),
    make_zip_iterator(make_tuple(d_input_a.end(), d_input_b.end())),
    d_output.begin(),
    matrix_multiply(N)
);

host_vector<float> h_output = d_output;