rocm/composable_kernel

Composable Kernel Docker镜像文档

1. 镜像概述和主要用途

Composable Kernel（简称CK）是一款高性能计算内核开发库，旨在为开发者提供模块化、可组合的构建块，用于构建针对多硬件架构（如NVIDIA GPU、AMD GPU等）优化的计算内核。该Docker镜像基于CK库构建，提供预配置的、一致的开发与运行环境，旨在简化依赖管理、消除环境差异，支持开发者快速上手CK库的使用、内核开发、测试及部署。

2. 核心功能和特性

2.1 预配置开发环境

• 集成最新稳定版Composable Kernel库及其核心依赖（如CUDA Toolkit、ROCm SDK等，根据目标硬件架构适配）。
• 包含完整编译工具链：GCC/Clang编译器、CMake构建系统、Ninja等。

2.2 多架构支持

• 支持NVIDIA GPU环境（基于CUDA）和AMD GPU环境（基于ROCm），可通过镜像标签选择目标架构。
• 兼容主流GPU驱动版本，确保硬件加速功能正常启用。

2.3 工具集成

• 内置调试工具（如CUDA-GDB、ROCgdb）、性能分析工具（Nsight Systems、rocprof）及代码检查工具。
• 包含CK库示例代码和文档，便于快速学习和验证。

2.4 环境一致性

• 固定依赖版本，避免因系统环境差异导致的兼容性问题，确保开发、测试、部署流程的一致性。

3. 使用场景和适用范围

3.1 典型使用场景

• HPC应用内核开发：构建针对科学计算、数值模拟的高性能自定义内核。
• AI框架加速优化：为深度学习框架（如PyTorch、TensorFlow）开发定制化算子内核。
• 跨架构内核移植：在不同GPU架构间验证和移植优化后的计算内核。
• 教学与研究：作为高性能计算内核开发的教学环境，或用于硬件架构适配研究。

3.2 适用人群

• 高性能计算（HPC）开发者
• AI框架算子优化工程师
• 硬件架构适配研究人员
• 计算内核性能调优人员

4. 使用方法和配置说明

4.1 镜像拉取

Composable Kernel镜像托管于Docker Hub或GitHub Container Registry，可通过以下命令拉取：

bash
# 拉取最新版（默认架构，如CUDA）
docker pull composablekernel:latest

# 拉取指定版本（如v1.2.0）
docker pull composablekernel:v1.2.0

# 拉取特定架构版本（如ROCm支持）
docker pull composablekernel:rocm-latest

镜像标签说明：

• latest：最新稳定版，默认适配主流架构（如CUDA）；
• v<x.y.z>：指定CK版本（如v1.2.0）；
• cuda-<version>/rocm-<version>：指定底层计算栈版本（如cuda-12.1、rocm-5.6）。

4.2 容器运行

4.2.1 基础开发环境启动

bash
# 启动交互式终端，挂载当前目录至容器内/workspace（用于代码持久化）
docker run -it --rm \
  --name ck-dev-env \
  -v $(pwd):/workspace \  # 挂载本地代码目录
  composablekernel:latest \
  /bin/bash

4.2.2 GPU硬件访问（必要时）

若需使用GPU加速，需在运行时添加硬件访问参数：

• NVIDIA GPU（需安装nvidia-docker）：

  bash
  docker run -it --rm \
    --gpus all \  # 映射所有GPU设备
    -v $(pwd):/workspace \
    composablekernel:cuda-12.1 \
    /bin/bash
  

• AMD GPU（ROCm）：

  bash
  docker run -it --rm \
    --device /dev/kfd \  # 访问AMD GPU设备
    --device /dev/dri \
    -v $(pwd):/workspace \
    composablekernel:rocm-5.6 \
    /bin/bash
  

4.3 环境变量配置

容器支持通过环境变量自定义运行时行为，常用变量如下：

使用示例：指定Debug模式编译

bash
docker run -it --rm \
  -e CMAKE_BUILD_TYPE=Debug \
  composablekernel:latest \
  /bin/bash

4.4 典型操作示例

4.4.1 编译CK示例代码

容器内默认包含CK示例代码（路径：/opt/composablekernel/examples），可直接编译验证环境：

bash
# 进入容器后执行
cd /opt/composablekernel/examples/01_simple_gemm
mkdir build && cd build
cmake ..
make -j$(nproc)
# 运行编译后的示例（需GPU支持）
./simple_gemm

4.4.2 挂载本地代码开发

将本地CK项目挂载至容器，进行自定义内核开发：

bash
# 本地主机执行（挂载当前项目目录至容器/workspace）
docker run -it --rm \
  --gpus all \
  -v $(pwd):/workspace \
  composablekernel:latest \
  /bin/bash

# 容器内编译本地代码
cd /workspace
mkdir build && cd build
cmake .. -DCMAKE_PREFIX_PATH=/opt/composablekernel
make -j$(nproc)

5. 部署方案示例

5.1 基础开发环境部署（docker run）

bash
# 拉取并启动支持CUDA 12.1的CK开发环境，挂载代码目录并保留容器（--rm改为--name持久化）
docker run -it --name ck-dev-cuda12 \
  --gpus all \
  -v /path/to/local/code:/workspace \
  -v /path/to/datasets:/datasets \  # 挂载数据集
  composablekernel:cuda-12.1 \
  /bin/bash

5.2 持久化开发环境（docker-compose）

创建docker-compose.yml配置文件，定义持久化开发环境：

yaml
version: '3'
services:
  ck-dev:
    image: composablekernel:latest
    container_name: ck-dev-env
    runtime: nvidia  # NVIDIA GPU需添加（依赖nvidia-docker）
    environment:
      - CMAKE_BUILD_TYPE=Release
      - WORKSPACE=/workspace
    volumes:
      - ./local_code:/workspace  # 本地代码目录
      - ./datasets:/datasets    # 数据集目录
      - ./ck_cache:/root/.cache  # 缓存目录持久化
    tty: true  # 保持终端连接

启动服务：

bash
docker-compose up -d  # 后台启动
docker-compose exec ck-dev /bin/bash  # 进入容器

6. 注意事项

• 硬件兼容性：使用GPU时，需确保主机GPU驱动版本与镜像依赖的CUDA/ROCm版本匹配（如CUDA 12.1需驱动版本≥530.30.02）。
• 镜像标签选择：根据目标硬件架构（NVIDIA/AMD）和依赖版本选择对应标签，避免架构不匹配导致功能异常。
• 数据持久化：容器内数据默认非持久化，建议通过-v参数挂载本地目录保存代码、编译产物及数据。
• 性能调优：进行内核性能测试时，建议为容器分配足够CPU资源（如--cpus参数）和内存，避免资源限制影响测试结果。