Learn Kokkos Module 4 - Hierarchical Parallelism📝 Kokkos 学习

📝 Kokkos 学习笔记：层级并行 (Module 4)

1. 核心架构：层级并行 (Hierarchical Parallelism) 🏗️

层级并行是为了解决“扁平并行”（如 RangePolicy）在处理小规模或嵌套任务时无法充分利用硬件资源的问题。

组织结构图

League (大组/联赛) ：所有团队的集合。
Team (团队) ：执行的基本单位，通常处理一个独立的工作块（如矩阵的一行）。
Thread (线程) ：团队内的成员，共同协作。
Vector (向量) ：最底层的并行，对应硬件的 SIMD 或 GPU 的 Warp 。

身份识别 API

级别	询问的问题	方法名
League	总共有多少个团队？	`teamMember.league_size()`
League	当前是第几个团队？	`teamMember.league_rank()`
Team	团队内有多少线程？	`teamMember.team_size()`
Team	我是团队第几个线程？	`teamMember.team_rank()`

2. 执行策略与嵌套并行 🔄

TeamPolicy

使用 TeamPolicy<ExecSpace>(league_size, team_size) 启动内核。

建议将 team_size 设置为 Kokkos::AUTO，由框架自动选择最优参数。

三级并行原语

Team 级：外层并行。
Thread 级：TeamThreadRange(teamMember, range) 分摊给团队线程。
Vector 级：ThreadVectorRange(teamMember, range) 分摊给向量通道。

3. 精确控制：Single 模式与同步 🛠️

由于在 GPU 等后端，代码会在多个向量通道中冗余执行，我们需要控制执行频率：

single(PerTeam(teamMember), ...) ：整个团队仅执行一次。
single(PerThread(teamMember), ...) ：每个线程仅执行一次（屏蔽向量通道冗余）。
teamMember.team_barrier() ：团队内所有线程必须在此集合，确保之前的任务（如数据搬运）全部完成。

4. 暂存内存 (Scratch Memory) 🧠

作为手动管理的缓存，用于减少全局内存访问压力。

关键特性

Level 0：极快、容量小（对应 GPU Shared Memory）。
Level 1：较大、速度适中（对应高带宽内存）。
使用流程：
1. 调用 View::shmem_size(dims) 计算字节。
2. policy.set_scratch_size(level, PerTeam(bytes)) 申请空间。
3. 内核内通过 teamMember.team_scratch(level) 建立映射。

5. 资源管理：唯一标识符 (Unique Token) 🔑

提供一种线程安全、可移植的方式来获取当前执行单元的唯一 ID 。

应用场景：访问预先分配的私有资源（如随机数生成器池）。
操作方法：
- token.acquire()：领令牌。
- token.release(id)：还令牌。
优化：通常设置 20% 的冗余空间以提升在高并发下的搜索性能。

6.example

#include <Kokkos_Core.hpp>

using TeamPolicy = Kokkos::TeamPolicy<>;
using member_type = TeamPolicy::member_type;
using ScratchView = Kokkos::View<double*, Kokkos::DefaultExecutionSpace::scratch_memory_space>;

struct EnhancedModule4Functor {
    int _ncols;
    Kokkos::View<double**> _A;
    Kokkos::View<double*> _x;
    Kokkos::View<double*> _y;
    Kokkos::Experimental::UniqueToken<> _tokens;

    EnhancedModule4Functor(int ncols, Kokkos::View<double**> A, Kokkos::View<double*> x, Kokkos::View<double*> y)
        : _ncols(ncols), _A(A), _x(x), _y(y), _tokens(Kokkos::Experimental::UniqueToken<>()) {}

    KOKKOS_INLINE_FUNCTION
    void operator()(const member_type& teamMember) const {
        int row = teamMember.league_rank();

        // 1. 获取暂存内存 (Scratch Space)
        // 每个团队领用一个快速“笔筒”来存放该行计算所需的 x 向量切片
        ScratchView x_shared(teamMember.team_scratch(0), _ncols);

        // 2. 协作加载数据 (TeamVectorRange)
        // 团队所有线程和向量通道协作，将全局内存中的 x 搬到暂存内存
        Kokkos::parallel_for(Kokkos::TeamVectorRange(teamMember, _ncols), [&](const int j) {
            x_shared(j) = _x(j);
        });

        // 3. 团队屏障 (team_barrier)
        // 确保所有组员都放好了“画笔”，才能开始画画
        teamMember.team_barrier();

        // 4. 执行计算 (嵌套并行：Thread 级 + Vector 级)
        double row_sum = 0;
        Kokkos::parallel_reduce(Kokkos::TeamThreadRange(teamMember, _ncols), [&](const int j, double& lsum) {
            double val = 0;
            // 第三层：向量并行
            Kokkos::parallel_reduce(Kokkos::ThreadVectorRange(teamMember, 1), [&](const int, double& vsum) {
                vsum += _A(row, j) * x_shared(j);
            }, val);
            lsum += val;
        }, row_sum);

        // 5. Single 模式：每个团队仅由一人更新结果
        Kokkos::single(Kokkos::PerTeam(teamMember), [&]() {
            _y(row) = row_sum;
        });

        // 6. UniqueToken：模拟访问受限资源
        // 假设我们要记录哪个硬件线程处理了这一行
        Kokkos::single(Kokkos::PerTeam(teamMember), [&]() {
            int token_id = _tokens.acquire(); // 领令牌
            // 这里可以进行一些线程安全的操作，例如打印或写入日志
            _tokens.release(token_id); // 还令牌
        });
    }
};

int main() {
    Kokkos::initialize();
    {
        int nrows = 100, ncols = 64;
        Kokkos::View<double**> A("A", nrows, ncols);
        Kokkos::View<double*> x("x", ncols);
        Kokkos::View<double*> y("y", nrows);

        // 计算暂存内存大小
        size_t scratch_size = ScratchView::shmem_size(ncols);

        // 配置 TeamPolicy: (大组数, 团队人数, 向量长度)
        TeamPolicy policy(nrows, Kokkos::AUTO, 32);
        
        // 核心：设置暂存内存级别和大小
        policy.set_scratch_size(0, Kokkos::PerTeam(scratch_size));

        Kokkos::parallel_for("EnhancedModule4", policy, EnhancedModule4Functor(ncols, A, x, y));
    }
    Kokkos::finalize();
    return 0;
}