Как правильно использовать OpenMP Target Enter/Exit/Update для неструктурированных асинхронных вычислений на стороне устройства?

Цель

Я хочу использовать цель OpenMP в C++ так же, как я сейчас использую CUDA:

1. Выделите любое количество массивов на стороне устройства (подойдет зеркало хост-массива), при необходимости инициализируйте.
2. Выполните произвольную серию операций на стороне устройства, используя эти массивы, без копирования на хост.
3. После завершения скопируйте нужный массив (например, выходные данные) на хост.

В конце концов, я хочу иметь возможность использовать этот код в качестве библиотеки, например. PyBind11, но сейчас это не так важно. Для этого я бы подумал:

Создайте класс, содержащий этот массив и все необходимые метаданные и служебные методы. Затем добавьте операторы #pragma omp target enter/exit/update к таким методам, как конструктор, деструктор, to_host и set_data, например.

Затем напишите функции, которые принимают эти объекты массива в качестве входных данных и выполняют, например, операции над ними.

Затем просто вызовите указанные функции по желанию, а после этого вызовите мой метод #pragma omp target parallel for и используйте мои данные на хосте.

Проблема

При этом я столкнулся с несколькими проблемами: во-первых, я не совсем уверен, как правильно использовать целевые директивы после to_host, но документация IBM предлагает просто использовать операторы enter как обычно, а документы сообщества предлагают использование map и enter в методах вполне допустимо.

Итак, в результате моей асинхронной операции с массивами я получаю, например, следующее:

void async_op(AsyncArray2D result, AsyncArray2D a, AsyncArray2D b) {
    #pragma omp target map(to: a.data_ptr[:a.size], b.data_ptr[:b.size], result.data_ptr[:result.size])
    #pragma omp teams distribute parallel for
    for (int i = 0; i < result.size; i++) {
        result.data_ptr[i] = a.data_ptr[i] * b.data_ptr[i];
        result.data_ptr[i] -= 1 / b.data_ptr[i];
        result.data_ptr[i] += b.data_ptr[i] / (a.data_ptr[i] + i);
    }
}

NB: я знаю, что для асинхронной работы мне придется использовать exit, но сейчас это не моя проблема.

Проблема в том, что я верну свой массив на хост только в том случае, если вставлю nowait в эту же функцию. Я не могу вызвать его с помощью метода и не могу поместить его, например. #pragma omp target update from(result.data_ptr[:result.size]). Если я это сделаю, копирование не будет работать правильно (на выходе я получаю только нули). Я, конечно, также могу использовать main() или неявное сопоставление, но, насколько я могу судить, это всего лишь означает, что копирование будет выполняться автоматически, а мой to/from все равно ничего не делает.

Таким образом, я даже не уверен, что мой оператор update что-то делает, потому что это почти неотличимо от отсутствия такого оператора вообще (за исключением, возможно, того, что передача памяти более эффективна).

Я предполагаю, что основная проблема заключается в том, что OpenMP каким-то образом ограничен данным контекстом, но мне нужно понять, какие у меня есть варианты и в чем именно заключается проблема, чтобы найти решение, которое будет работать в моем случае. Могу ли я использовать методы только одного класса для выполнения асинхронных операций или чего-то в этом роде?

Код

Вот не совсем минимальный пример, я мог бы скомпилировать его с помощью enter, но полагаю, что результаты будут различаться в зависимости от системы.

#include <cstring>
#include <iostream>
#include <omp.h>


class AsyncArray2D {
public:
    float* data_ptr;
    size_t size;
    bool is_on_device;
    int shape[2], strides[2];

    AsyncArray2D(size_t rows, size_t cols) {
        size = rows * cols;
        shape[0] = rows;
        shape[1] = cols;
        strides[0] = cols;
        strides[1] = 0;
        data_ptr = new float[size];
        #pragma omp target enter data map(alloc: data_ptr[:size])
    }

    void to_host() {
        // This method seems to accomplish nothing
        #pragma omp target update from(data_ptr[0:size])
    }

    void set_data(float * buf_ptr) {
        std::memcpy(data_ptr, buf_ptr, size * sizeof(float));
        // unclear if this accomplishes anything
        #pragma omp target update to(data_ptr[:size])
    }

    ~AsyncArray2D() {
        #pragma omp target exit data map(delete: data_ptr[:size])            
    }
};

void async_op(AsyncArray2D result, AsyncArray2D a, AsyncArray2D b) {
    #pragma omp target map(to: a.data_ptr[:a.size], b.data_ptr[:b.size], result.data_ptr[:result.size])
    #pragma omp teams distribute parallel for
    for (int i = 0; i < result.size; i++) {
        result.data_ptr[i] = a.data_ptr[i] * b.data_ptr[i];
        result.data_ptr[i] -= 1 / b.data_ptr[i];
        result.data_ptr[i] += b.data_ptr[i] / (a.data_ptr[i] + i);
    }
   // uncomment the below line and we get the correct result.
   // #pragma omp target update from(result.data_ptr[:result.size])
}

int main() {
    size_t dim_size = 8;
    
    float * buffer_a = new float[dim_size * dim_size];
    float * buffer_b = new float[dim_size * dim_size];
    float * buffer_result = new float[dim_size * dim_size];
    std::fill_n(buffer_a, dim_size * dim_size, 1.2);
    std::fill_n(buffer_b, dim_size * dim_size, 2.7);
    std::fill_n(buffer_result, dim_size * dim_size, 0);
    std::cout << "Created buffer" << std::endl;
    AsyncArray2D a(dim_size, dim_size);
    AsyncArray2D b(dim_size, dim_size);
    AsyncArray2D result(dim_size, dim_size);
    std::cout << "Created arrays" << std::endl;
    a.set_data(buffer_a);
    b.set_data(buffer_b);
    result.set_data(buffer_result);
    std::cout << "Set data" << std::endl;
    // run for a while so that we can see if the GPU is doing anything...
    for (int i = 0; i < 100; i++) {
        async_op(result, a, b);
    }
    std::cout << "Added results" << std::endl;
    result.to_host();
    std::cout << "Moved to host" << std::endl;
    for (int j = 0; j < dim_size; j++) {
        for (int i = 0; i < dim_size; i++) {
            std::cout << result.data_ptr[j * dim_size + i] << ' ' ;
        }
        std::cout << std::endl;
    }
    return 0;
}

Итак, с помощью этого кода я получаю результат

0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0 
0 0 0 0 0 0 0 0

но если я раскомментирую строку g++-12 -O2 -fopenmp -fno-stack-protector -fcf-protection=none -foffload=nvptx-none example.cpp -o example или изменю отображение update на result, я получу ожидаемый результат:

5.11963 4.0969 3.71338 3.51249 3.38886 3.30511 3.24463 3.1989 
3.16311 3.13434 3.1107 3.09094 3.07418 3.05977 3.04726 3.0363 
3.02661 3.01798 3.01025 3.00329 2.99699 2.99125 2.98601 2.9812 
2.97677 2.97268 2.96889 2.96537 2.9621 2.95903 2.95617 2.95348 
2.95096 2.94858 2.94633 2.94422 2.94221 2.94031 2.93851 2.93679 
2.93516 2.93361 2.93213 2.93072 2.92936 2.92807 2.92683 2.92565 
2.92451 2.92341 2.92236 2.92135 2.92038 2.91945 2.91854 2.91767 
2.91683 2.91602 2.91524 2.91448 2.91375 2.91304 2.91235 2.91169

Насколько я могу судить, метод from никогда ничего не делает, прагма просто ничего не выполняет.

Есть ли у кого-нибудь указания, как мне поступить с этим? Я бы предпочел не переключаться на другую структуру, поскольку мне нужно что-то как можно более общее, и я не хочу застревать в чем-то вроде экосистемы Intel или использовать что-то еще, что слишком сильно усложнило бы распространение кода.

Обновлено: Просто добавлю: я только что протестировал создание to_host метода AsyncArray2D, вызываемого массивом async_op, и это работает так, как ожидалось (result делает то, что должно). Однако необходимость превращать все операции в методы немного ограничивает меня (и приведет к очень уродливому кодированию с огромными списками методов...), так что это не мой идеальный выбор...

Спасибо!