ROCm/HIP中的协作组(Cooperative Groups)编程教程-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00517/article/details/148549905

ROCm/HIP中的协作组(Cooperative Groups)编程教程

HIP HIP: C++ Heterogeneous-Compute Interface for Portability 项目地址: https://gitcode.com/gh_mirrors/hi/HIP

前言

在GPU编程中，协作组(Cooperative Groups)是一种强大的线程同步和协作机制，它扩展了传统的线程块(thread block)概念，允许开发者定义更灵活的线程分组方式。本文将详细介绍如何在ROCm/HIP环境中使用协作组功能。

协作组基础概念

协作组是HIP编程模型中的一个重要特性，它提供了比传统线程块更细粒度的线程分组控制。主要优势包括：

可以在线程块内创建更小的线程组（称为tile）
支持跨线程块的协作
提供统一的同步机制
简化复杂并行算法的实现

开发环境准备

在使用协作组功能前，需要确保：

已安装最新版本的ROCm平台
确认GPU设备支持协作组功能
配置好HIP编译工具链

协作组核心组件

线程组类型

HIP协作组提供了几种主要的线程组类型：

thread_block：代表整个线程块
thread_block_tile：线程块内的子分区（tile）
grid_group：代表整个网格（跨线程块协作）

关键操作

tiled_partition()：创建线程块内的子分区
sync()：同步线程组内所有线程
thread_rank()：获取线程在组内的排名
size()：获取线程组的大小

实践示例：并行归约

下面我们通过一个并行归约的例子来演示协作组的实际应用。

主机端参考实现

std::vector<unsigned int> ref_reduced(const unsigned int partition_size,
                                    std::vector<unsigned int> input)
{
    const unsigned int input_size = input.size();
    const unsigned int result_size = input_size / partition_size;
    std::vector<unsigned int> result(result_size);

    for(unsigned int i = 0; i < result_size; i++)
    {
        unsigned int partition_result = 0;
        for(unsigned int j = 0; j < partition_size; j++)
        {
            partition_result += input[partition_size * i + j];
        }
        result[i] = partition_result;
    }
    return result;
}

设备端归约核心

__device__ unsigned int reduce_sum(thread_group g, unsigned int* x, unsigned int val)
{
    const unsigned int group_thread_id = g.thread_rank();
    
    for(unsigned int i = g.size() / 2; i > 0; i /= 2)
    {
        x[group_thread_id] = val;
        g.sync();
        
        if(group_thread_id < i)
        {
            val += x[group_thread_id + i];
        }
        
        g.sync();
    }

    return (g.thread_rank() == 0) ? val : 0;
}

内核实现分解

初始化阶段

thread_block thread_block_group = this_thread_block();
__shared__ unsigned int workspace[2048];
unsigned int output;
const unsigned int input = d_vector[thread_block_group.thread_rank()];

// 创建16线程的自定义分区
thread_block_tile<16> custom_partition = tiled_partition<16>(thread_block_group);

线程块级归约

output = reduce_sum(thread_block_group, workspace, input);
if(thread_block_group.thread_rank() == 0)
{
    d_block_reduced_vector[0] = output;
}

自定义分区级归约

output = reduce_sum(custom_partition, &workspace[group_offset], input);
if(custom_partition.thread_rank() == 0)
{
    const unsigned int partition_id = thread_block_group.thread_rank() / 16;
    d_partition_reduced_vector[partition_id] = output;
}

主机端关键步骤

1. 检查设备支持

int device = 0;
int supports_coop_launch = 0;
HIP_CHECK(hipGetDevice(&device));
HIP_CHECK(hipDeviceGetAttribute(&supports_coop_launch, 
                              hipDeviceAttributeCooperativeLaunch, 
                              device));
if(!supports_coop_launch)
{
    std::cout << "设备不支持协作组" << std::endl;
    return 0;
}

2. 配置协作组参数

constexpr unsigned int num_blocks = 1;
constexpr unsigned int threads_per_block = 64;
constexpr unsigned int partition_size = 16;

static_assert(threads_per_block % partition_size == 0,
             "线程数必须是分区大小的整数倍");

3. 启动协作内核

void* params[] = {&d_vector, &d_block_reduced, &d_partition_reduced};
HIP_CHECK(hipLaunchCooperativeKernel(vector_reduce_kernel<16>,
                                   dim3(num_blocks),
                                   dim3(threads_per_block),
                                   params,
                                   0,
                                   hipStreamDefault));