Composable Kernel（CK）库旨在提供一套在 AMD GPU 上算子融合的后端方案，该研究希望未来能够移植到 AMD 的所有 GPU 上，并且最终也可以被移植到 AMD CPU 上，该项目已开源。与Meta AITemplate的深度合作大幅提升了AI模型在AMD GPU的端到端性能。

图优化在降低 AI 模型的训练和推理使用的时间和资源方面起着重要作用。图优化的一个重要功能是模型中将可以融合的算子进行融合，通过降低内存占用和减少数据在低速内存中的搬运来提高计算效率。然而，实现一套能够提供各种算子融合的后端方案难度很大，导致在实际硬件上 AI 模型能够使用的算子融合非常有限。

Composable Kernel (CK)库旨在提供一套在 AMD GPU 上的算子融合的后端方案。CK 使用通用编程语言 HIP C++，完全开源。其设计理念包括：

  • 高性能 & 高生产力：CK 的核心是一组精心设计，高度优化，可复用的基础模块。CK 库内所有的算子都是通过组合这些基础模块实现的。复用这些基础模块大大缩短开发后端算法的周期，同时还能保证高性能。

  • 精通当前的 AI 问题，快速适应未来的 AI 问题：CK 旨在提供一套完整的 AI 算子后端方案，这让复杂的算子融合成为可能，因为这样让整个后端都可以用 CK 实现，而不需依赖外部算子库。CK 的可复用基础模块足以实现常见 AI 模型（机器视觉，自然语言处理，等等）所需的各种算子及其融合。当新出现的 AI 模型需要新的算子时，CK 也将会提供所需的基础模块。

  • AI 系统专家的简单但强大的工具：CK 所有的算子都是用 HIP C++ 模版实现的。AI 系统专家可以通过实例化模版来定制这些算子的属性，比如数据类型，元操作类型，张量存储格式，等等。这通常只需要几行代码。

  • 友好的 HIP C++ 界面：HPC 算法开发者一直在推动着 AI 计算加速的前沿。CK 的一个重要设计理念就是要让 HPC 算法开发者更容易对 AI 加速作出贡献。因此 CK 所有核心模块都是用 HIP C++ 实现，而不是 Intermediate Representation （IR）。HPC 算法开发者直接以他们熟悉的编写 C++ 代码的形式编写算法，而无需像基于 IR 的算子库那样，以通过编写针对某种特定算法的 Compiler Pass 来实现。这样做可以大大提高算法的迭代速度。

  • 可移植性：今天使用 CK 作为后端的图优化将能够移植到未来 AMD 的所有的 GPU 上，并且最终也可以被移植到 AMD CPU 上。

CK 源代码：https://github.com/ROCmSoftwarePlatform/composable_kernel

内容中包含的图片若涉及版权问题，请及时与我们联系删除