— 机器之心编译 —
Composable Kernel（CK）库旨在提供一套在AMD GPU 上算子融合的后端方案，该研究希望未来能够移植到 AMD 的所有 GPU 上，并且最终也可以被移植到 AMD CPU 上，该项目已开源。
此外，CK库与Meta AITemplate的深度合作大幅提升了AI模型在AMD GPU的端到端性能。
AITemplate 是一个能把 AI 模型转换成高性能 C++ GPU 模板代码的 Python 框架。该框架在设计上专注于性能和简化系统。AITemplate 系统一共分为两层：前端部分进行图优化，后端部分针对目标 GPU 生成 C++ 模板代码。AITemplate 不依赖任何额外的库或 Runtime，如 cuBLAS、cudnn、rocBLAS、MIOpen、TensorRT、MIGraphX 等。
图1 BERT类模型AMD MI250性能 vs. FasterTransformer A100性能比较
图优化在降低 AI 模型的训练和推理使用的时间和资源方面起着重要作用。图优化的一个重要功能是模型中将可以融合的算子进行融合，通过降低内存占用和减少数据在低速内存中的搬运来提高计算效率。然而，实现一套能够提供各种算子融合的后端方案难度很大，导致在实际硬件上 AI 模型能够使用的算子融合非常有限。
Composable Kernel (CK)库旨在提供一套在 AMD GPU 上的算子融合的后端方案。CK 使用通用编程语言 HIP C++，完全开源。其设计理念包括：
- 高性能& 高生产力：CK 的核心是一组精心设计，高度优化，可复用的基础模块。CK 库内所有的算子都是通过组合这些基础模块实现的。复用这些基础模块大大缩短开发后端算法的周期，同时还能保证高性能。
- 精通当前的 AI 问题，快速适应未来的 AI 问题：CK 旨在提供一套完整的 AI 算子后端方案，这让复杂的算子融合成为可能，因为这样让整个后端都可以用 CK 实现，而不需依赖外部算子库。CK 的可复用基础模块足以实现常见 AI 模型（机器视觉，自然语言处理，等等）所需的各种算子及其融合。当新出现的 AI 模型需要新的算子时，CK 也将会提供所需的基础模块（复用已有模块需要对已有模块的API很熟悉，对新算子也需要人力跟进）。
- AI 系统专家的简单但强大的工具：CK 所有的算子都是用 HIP C++ 模版实现的。AI 系统专家可以通过实例化模版来定制这些算子的属性，比如数据类型，元操作类型，张量存储格式，等等。这通常只需要几行代码（听起来HIP C++的扩展和表达能力很强）。
- 友好的 HIP C++ 界面：HPC 算法开发者一直在推动着 AI 计算加速的前沿。CK 的一个重要设计理念就是要让 HPC 算法开发者更容易对 AI 加速作出贡献。因此CK 所有核心模块都是用 HIP C++ 实现，而不是 Intermediate Representation （IR）。HPC 算法开发者直接以他们熟悉的编写 C++ 代码的形式编写算法，而无需像基于 IR 的算子库那样，以通过编写针对某种特定算法的 Compiler Pass 来实现。这样做可以大大提高算法的迭代速度。
- 可移植性：今天使用 CK 作为后端的图优化将能够移植到未来 AMD 的所有的 GPU上，并且最终也可以被移植到AMD CPU上【2】。
核心概念
CK 引入了两个概念：Tensor Coordinate Transformation 和 Tile Tensor Operator ，二者共同组成了 CK 的可复用的基础模块。
1.开创性的引入“张量坐标变换”（Tensor Coordinate Transformation）降低 AI 算子的编写复杂度。该研究开创性地定义了一组可复用的 Tensor Coordinate Transformation 基础模块，并且用它们把复杂的 AI 算子（比如卷积，group normalization reduction，Depth2Space，等等）以数学严谨的方式重新表达成了最基础的 AI 算子（GEMM，2D reduction，tensor transfer，等等）。
图1 使用 CK 的 Tensor Coordinate Transformation 基础模块将 convolution 算子表达成 GEMM 算子
这项技术可以让为基础 AI 算子编写的算法直接被用到所有与之对应的复杂的 AI 算子上，而无需重写算法。
2.基于 Tile 的编程范式：开发算子融合的后端算法可以被看成先将每一个融合前的算子（独立算子）拆解成许多 “小块” 的数据操作，然后再把这些“小块” 操作组合成融合的算子。每一个这样的 “小块” 操作都对应一个原始的独立算子，但是被操作的数据只是原始张量的一部分（tile），因此这样的 “小块” 操作被称为 Tile Tensor Operator。
图2 CK 的组成（下：可复用的基础模块；上：独立算子与融合算子）
CK 库包含一组针对 Tile Tensor Operator 的高度优化的实现，CK 里所有的 AI 独立算子和融合算子都是用它们实现的。目前，这些Tile Tensor Operators 包括 Tile GEMM，Tile Reduction 和 Tile Tensor Transfer。每一个 Tile Tensor Operator 都有针对 GPU thread block，warp 和 thread 的实现。
图3 Composable Kernel Overview
代码结构
CK 库结构分为四层，从下到上分别是：Templated Tile Operator，Templated Kernel and Invoker，Instantiated Kernel and Invoker和Client API【3】。每一层对应不同的开发者：
- AI 系统专家：“我需要一个后端方案提供高性能的独立和融合算子让我可以直接使用”。这个例子【4】里用的Client API 和 Instantiated Kernel and Invoker提供了预先实例化并编译好的对象，以满足这类开发者的需求。
- AI 系统专家：“我为一个开源的 AI 框架做最先进的图优化工作。我需要一个能够为图优化所需的所有融合算子提供高性能 kernel 的后端方案。同时我也需要定制这些 kernel，所以像 “要么接受，要么弃用” 的黑盒解决方案不能满足我的需求”。Templated Kernel and Invoker 层能满足这类开发者。比如这个例子【5】中开发者可以自己使用 Templated Kernel and Invoker 层实例化出所需的 FP16 的 GEMM + Add + Add + FastGeLU 的 kernel。
- HPC 算法专家：“我的团队为公司内部不断迭代的 AI 模型开发高性能后端算法。我们团队中有 HPC 算法专家，但我们仍然希望可以通过复用和改进硬件供应商提供的高度优化的源代码来提高我们的生产力，并且让我们的代码可以被移植到未来的硬件构架上。我们希望可以不用通过与硬件提供商分享我们的代码来做到这点”。Templated Tile Operator层可以帮助到这一类开发者。比如这个代码【6】中开发者使用 Templated Tile Operator 来实现 GEMM 的优化管线。
图4 CK 库四层结构
基于 AITemplate + CK 的端到端模型推理
Meta 的 AITemplate 【7】（AIT）是一个统一 AMD 和 Nvidia GPU 的 AI 推理系统。AITemplate 使用 CK 作为其 AMD GPU 上的后端，它使用的是 CK 的 Templated Kernel and Invoker 层。
AITemplate + CK 在 AMD Instinct™ MI250 上取得了多个重要 AI 模型最先进的推理性能。CK 里大多数先进的融合算子的定义，都是在 AITemplate 团队的远见下推动的。许多融合算子的算法也是由 CK 和 AITemplate 团队共同设计。
本文比较了几个端到端模型在 AMD Instinct MI250 和同级别产品【8】的性能表现。本文中所有 AMD Instinct MI250 的 AI 模型的性能数据都是用 AITemplate【9】 + CK【10】取得的。
实验
ResNet-50
下图显示了 AMD Instinct MI250 上的 AIT + CK 与 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12 【11】（TRT）的性能比较。
结果显示AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 1.08 倍的加速。
BERT
一个基于 CK 实现的 Batched GEMM + Softmax + GEMM 融合算子模版，可以完全消除掉中间结果在 GPU 计算单元（Compute Unit）与 HBM 之间的搬运。通过使用这个融合算子模版，attention layer 许多原本是带宽瓶颈（bandwidth bound）的问题变成了计算瓶颈（compute bound）的问题，这样可以更好发挥 GPU 的计算能力。这个 CK 的实现深受 FlashAttention 【12】的启发，并比原始的 FlashAttention 的实现减少了更多的数据搬运。
下图显示了 AMD Instinct MI250 上的 AIT + CK 与 A100-PCIe-40GB 和 A100-DGX-80GB 上的 FasterTransformer v5.1.1 bug fix 【13】（FT）的 Bert Base 模型（uncased)的性能比较。当 Sequence 是 4096 时，FT 在 A100-PCIe-40GB 和 A100-DGX-80GB 上会在 Batch 32 时 GPU 内存溢出。因此，在 Sequence 是 4096 时，本文只显示 Batch 16 的结果。
结果显示AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 FT 3.28 倍，以及相比于 A100-DGX-80GB 上的 FT 2.91 倍的加速。
Vision Transformer （VIT）
下图显示了 AMD Instinct MI250 上的 AIT + CK 与 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12（TRT）的 Vision Transformer Base （224x224 图片）的性能比较。
结果显示AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 1.8 倍，以及相比于 A100-DGX-80GB 上的 TRT 1.4 倍的加速。
Stable Diffusion
端到端的 Stable Diffusion
下表显示 AIT + CK 在 AMD Instinct MI250 上 Stable Diffusion 端到端（Batch 1，2，4， 6）的性能数据。当 Batch 是 1 时，在 MI250 上只有一个 GCD 被使用，而在 Batch 2，4，6 时，两个 GCD 都被使用了。
Stable Diffusion 中的 UNet
不过本文还没有关于使用 TensorRT 运行 Stable Diffusion 端到端模型的公开的信息。但这篇文章“Make stable diffusion 25% faster using TensorRT” 【14】说明了怎么使用 TensorRT 加速 Stable Diffusion 中的 UNet 模型。UNet 是 Stable Diffusion 中最重要最花时间的部分，因此UNet 的性能大致反应了 Stable Diffusion 的性能。
下图显示了 AMD Instinct MI250 上的 AIT + CK 与 A100-PCIe-40GB 和 A100-DGX-80GB 上的 TensorRT v8.5.0.12（TRT）的 UNet 的性能比较。
结果显示AMD Instinct MI250 上的 AIT + CK 取得了相比于 A100-PCIe-40GB 上的 TRT 2.45 倍，以及相比于 A100-DGX-80GB 上的 TRT 2.03 倍的加速。
算子开发伪代码
图5 Transform a forward convolution operator into a GEMM operator
原文链接
作者：AMD
文章来源：NeuralTalk
推荐阅读
更多嵌入式AI干货请关注 嵌入式AI 专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
