如何在Kokkos环境中传递仿函数列表并并行化外层循环？

因为Kokkos的并行执行空间（如GPU）对可调用对象有严格要求：必须可设备复制、可内联，且无法直接使用主机函数指针或存储不同类型对象的容器。针对你的需求，这里提供两种可行方案：

方案1：使用std::variant统一仿函数类型（推荐，性能更优） #

利用C++17的std::variant对不同类型的仿函数做类型擦除，再通过std::visit执行对应的调用，这种方式能保留仿函数的内联特性，适合Kokkos的设备环境。

完整代码示例 #

#include <Kokkos_Core.hpp>
#include <variant>
#include <vector>

// 你的原仿函数定义
struct f1{
    KOKKOS_INLINE_FUNCTION
    void operator()() const{
        printf("Functor 1\n");
    }; 
};

struct f2{
    KOKKOS_INLINE_FUNCTION
    void operator()() const{
        printf("Functor 2\n");
    }; 
};

// 定义包含所有仿函数类型的variant
using FunctorVariant = std::variant<f1, f2>;

void parallel_function(int n, const std::vector<FunctorVariant>& host_functors) {
    // 将主机端的仿函数列表复制到设备可访问的Kokkos View中
    Kokkos::View<FunctorVariant*> device_functors("DeviceFunctors", host_functors.size());
    Kokkos::deep_copy(device_functors, host_functors);

    // 并行执行外层循环
    Kokkos::parallel_for("OuterLoop", Kokkos::RangePolicy<>(0, n), KOKKOS_LAMBDA(int i) {
        // 遍历设备端的仿函数列表并执行
        for (int j = 0; j < device_functors.size(); ++j) {
            std::visit([](const auto& func) {
                func();
            }, device_functors(j));
        }
    });
    Kokkos::fence(); // 等待所有并行任务完成
}

int main(int argc, char* argv[]) {
    Kokkos::initialize(argc, argv);
    {
        const int outer_loop_count = 10;
        std::vector<FunctorVariant> functor_list;
        functor_list.emplace_back(f1{});
        functor_list.emplace_back(f2{});
        
        parallel_function(outer_loop_count, functor_list);
    }
    Kokkos::finalize();
    return 0;
}

方案2：使用std::function（兼容性好，但性能略低） #

如果你的环境不支持C++17，可以用std::function做类型擦除，但要注意std::function依赖虚函数调用，在设备上可能带来额外开销，且部分老版本设备编译器对其支持有限。

代码示例片段 #

// 定义统一的函数签名
using FuncWrapper = std::function<void()>;

// 创建仿函数列表
std::vector<FuncWrapper> functor_list;
functor_list.emplace_back(f1{});
functor_list.emplace_back(f2{});

// 并行函数中的处理
void parallel_function(int n, const std::vector<FuncWrapper>& host_functors) {
    Kokkos::View<FuncWrapper*> device_functors("DeviceFunctors", host_functors.size());
    Kokkos::deep_copy(device_functors, host_functors);

    Kokkos::parallel_for("OuterLoop", Kokkos::RangePolicy<>(0, n), KOKKOS_LAMBDA(int i) {
        for (int j = 0; j < device_functors.size(); ++j) {
            device_functors(j)();
        }
    });
    Kokkos::fence();
}

关键注意事项 #

1. 仿函数必须满足设备可复制：你的f1、f2是空结构体，天然满足要求；如果仿函数有成员变量，需确保成员变量也能被设备复制。
2. 避免主机函数指针：设备执行空间无法直接调用主机函数指针，必须使用设备可见的仿函数或包装类型。
3. 内存复制：必须将主机端的仿函数列表通过Kokkos::deep_copy复制到设备可访问的Kokkos::View中，否则设备代码无法访问。

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