OpenMP设备offload归约如何复用现有设备内存避免数据传输

问题描述 #

需求为指定OpenMP device offload的reduction操作直接使用设备内存中已有的存储位置，避免主机与设备间的冗余数据传输，归约结果仅在设备端访问，无需回传主机。
初始实现代码如下：

void reduce(const double *mi, const double *xi, const double *yi,                                                                                
    double *mo, double *xo, double *yo, long n)
{
    #pragma omp target teams distribute parallel for reduction(+: mo[0],xo[0],yo[0]) is_device_ptr(mi,xi,yi,mo,xo,yo)
    for (long i = 0; i < n; ++i)
    {   
        mo[0] += mi[i];
        xo[0] += mi[i]*xi[i];
        yo[0] += mi[i]*yi[i];
    }   

    #pragma omp target is_device_ptr(mo,xo,yo)
    {   
    xo[0] /= mo[0];
    yo[0] /= mo[0];
    }   
}

使用clang++ 15面向NVIDIA PTX目标编译时，抛出如下错误：

test.cpp:6:109: error: reduction variable cannot be in a is_device_ptr clause in '#pragma omp target teams distribute parallel for' directive
    #pragma omp target teams distribute parallel for reduction(+: mo[0],xo[0],yo[0]) is_device_ptr(mi,xi,yi,mo,xo,yo)
                                                                                                            ^
test.cpp:6:67: note: defined as reduction
    #pragma omp target teams distribute parallel for reduction(+: mo[0],xo[0],yo[0]) is_device_ptr(mi,xi,yi,mo,xo,yo)
                                                                  ^

报错原因 #

OpenMP规范明确规定：同一target构造中，被reduction子句引用的变量对应的基指针，不能出现在is_device_ptr子句中。
is_device_ptr的语义是告知编译器：列出的指针存储的值是设备端地址，不需要在进入target区域时做主机-设备指针转译，这类指针无法被reduction默认的隐式私有副本、拷贝逻辑接管——默认reduction逻辑会假设归约变量存储在主机端，自动完成进区拷贝、归约、出区回传流程，和纯设备指针的语义直接冲突，因此Clang 15会抛出编译错误。

可行实现方案 #

以下两种方案均满足零主机-设备数据传输要求，归约全程在设备端完成，结果直接写入设备内存，仅在设备端可访问。

方案1：局部累加器实现（兼容性最佳） #

该方案完全规避reduction与is_device_ptr的语义冲突，将归约累加器声明为target区域内部的局部变量（驻留设备寄存器/共享内存），归约完成后直接写入设备内存，所有逻辑包裹在同一个target区域内，无任何跨设备数据传输，支持所有OpenMP 4.5+版本的offload编译器。

void reduce(const double *mi, const double *xi, const double *yi,                                                                                
    double *mo, double *xo, double *yo, long n)
{
    // 标记所有传入指针均为设备端地址，进入区域无隐式数据拷贝
    #pragma omp target is_device_ptr(mi, xi, yi, mo, xo, yo)
    {
        // 声明设备端局部累加器，初始值为0
        double mo_acc = 0.0, xo_acc = 0.0, yo_acc = 0.0;
        // 设备端分发并行任务，归约操作仅使用局部累加器
        #pragma omp teams distribute parallel for reduction(+:mo_acc, xo_acc, yo_acc)
        for (long i = 0; i < n; ++i)
        {
            mo_acc += mi[i];
            xo_acc += mi[i] * xi[i];
            yo_acc += mi[i] * yi[i];
        }
        // 归约完成后直接写入设备端目标内存，同步完成后续除法计算
        mo[0] = mo_acc;
        xo[0] = xo_acc / mo_acc;
        yo[0] = yo_acc / mo_acc;
    }
}

方案2：has_device_addr原生归约（OpenMP 5.0+支持） #

如果需要直接让OpenMP原生reduction逻辑操作已分配的设备内存地址，可使用has_device_addr子句标记归约目标的设备存储，不需要将对应基指针放入is_device_ptr列表，Clang 15已完整支持该特性。

void reduce(const double *mi, const double *xi, const double *yi,                                                                                
    double *mo, double *xo, double *yo, long n)
{
    // 设备端初始化归约目标存储为0，无主机传输
    #pragma omp target is_device_ptr(mo, xo, yo)
    {
        mo[0] = 0.0;
        xo[0] = 0.0;
        yo[0] = 0.0;
    }
    // 输入指针标记为设备地址，归约目标标记为已存在于设备内存
    #pragma omp target teams distribute parallel for \
        reduction(+:mo[0], xo[0], yo[0]) \
        is_device_ptr(mi, xi, yi) \
        has_device_addr(mo[0], xo[0], yo[0])
    for (long i = 0; i < n; ++i)
    {
        mo[0] += mi[i];
        xo[0] += mi[i] * xi[i];
        yo[0] += mi[i] * yi[i];
    }
    // 设备端完成后续归一化计算
    #pragma omp target is_device_ptr(mo,xo,yo)
    {
        xo[0] /= mo[0];
        yo[0] /= mo[0];
    }
}

编译参数参考 #

使用Clang 15面向NVIDIA PTX目标编译时，需添加如下offload参数：

clang++ -O3 -fopenmp -fopenmp-targets=nvptx64-nvidia-cuda test.cpp -o test

方案1兼容Clang 13及以上版本，方案2需要Clang 15及以上版本支持。

内容的提问来源于stack exchange，提问作者user2267882