【AI系统】核心计算之矩阵乘

简介: 本文探讨了AI模型中矩阵乘运算的优化实现及其在AI芯片设计中的重要性。文章首先介绍了卷积操作如何转化为矩阵乘,接着阐述了矩阵乘的分块(Tiling)技术以适应芯片内存限制,最后总结了几种常见的矩阵乘优化方法,包括循环优化、分块矩阵乘法、SIMD指令优化等,旨在提高计算效率和性能。

AI 模型中往往包含大量的矩阵乘运算,该算子的计算过程表现为较高的内存搬移和计算密度需求,所以矩阵乘的效率是 AI 芯片设计时性能评估的主要参考依据。本文我们一起来看一下矩阵乘运算在 AI 芯片的具体过程,了解它的执行性能是如何被优化实现的。

从卷积到矩阵乘

AI 模型中的卷积层的实现定义大家应该都已经比较熟悉了,卷积操作的过程大概可以描述为按照约定的窗口大小和步长,在 Feature Map 上进行不断地滑动取数,窗口内的 Feature Map 和卷积核进行逐元素相乘,再把相乘的结果累加求和得到输出 Feature Map 的每个元素结果。卷积到矩阵乘的的转换关系示意如下图。

05Matrix01.png

其中逐元素相乘,再累加的过程就是上文提到的一个计算单位:MACs,矩阵乘的 MACs 数对最终性能具有重要影响。通过将输入数据(Feature Map)和卷积核数据进行重排,卷积操作本质上可以等效理解为矩阵乘操作。

假设卷积的输入和输出的特征图维度用(IH, IW), (OH, OW)表示,卷积核窗口的数据维度用(KH, KW)表示,输入通道是 IC,输出通道是 OC,输入输出特征图和卷积核数据维度重排的转化对应关系如下公式,对输入数据的重排的过程称为 Im2Col,同理把转换后矩阵乘的数据排布方式再换回卷积输入的过程称为 Col2Im。

$$ \begin{align} &input:(IC, IH, IW)\rightarrow(OH*OW, KH*KW*IC)\\ &filter: (OC, KH, KW, IC)\rightarrow(OC, KH*KW*IC)\\ &output:(OC,OH, OW)\rightarrow(OC,OH*OW) \end{align}\\ $$

更具体的,假设卷积核的维度(2, 2),输入特征图维度(3, 3),输入和输出通道都是 1,对一个无 padding,stride=1 的卷积操作,输出特征图是(2, 2),所以输入卷积核转换为矩阵乘排布后的行数是 $2 * 2 = 4 \quad\quad\quad$,列数为 $2 * 2 * 1= 4 \quad\quad\quad\quad$。下图是对应的卷积到矩阵乘的转换示意,输入、输出特征图和卷积核都用不同的颜色表示,图中数字表示位置标记。

05Matrix02.png

比如输入特征图的排布转换过程:第 1 个输出对应输入特征图的窗口数据标记为 1, 2, 4, 5;第 2 个输出对应的输入特征图窗口数据标记为 2, 3, 5, 6;第 3 个输出对应的输入特征图窗口数据标记为 4, 5, 7, 8;第 4 个输出对应的输入特征图窗口数据标记为 5, 6, 8, 9。矩阵乘的维度对应关系如下。

$$ \begin{align} &input: (OH*OW, KH*KW*IC)\rightarrow (4,4)\\ &filter: (OC, KH*KW*IC)\rightarrow(1,4)\\ &output:(OC, OH*OW)\rightarrow(1,4) \end{align} $$

矩阵乘分块 Tilling

上面介绍了卷积到矩阵乘的转换过程,我们可以发现,转换后的矩阵乘的维度非常大,而芯片里的内存空间往往是有限的(成本高),表现为越靠近计算单元,带宽越快,内存越小。为了平衡计算和内存加载的时间,让算力利用率最大化,AI 芯片往往会进行由远到近,多级内存层级的设计方式,达到数据复用和空间换时间的效果。根据这样的设计,矩阵乘实际的数据加载和计算过程将进行分块 Tiling 处理。

假设用 CHW 表示上面转换公式中的 $KH * KW * IC \quad\quad\quad\quad\quad\quad$ 的值,M 表示 OC,N 表示 $OH * OW $,矩阵乘的输入特征图维度是 (CHW, N),矩阵乘的卷积核维度是(M, CHW),输出矩阵维度是(M, N),可以同时在 M,N,CHW 三个维度进行 Tiling,每次计算过程分别加载一小块的特征图和卷积核数据计算,比如在 M,N,CHW 三个维度各分了 2 小块,得到完成的输出特征图需要进行 8 次的数据加载和计算。下图中的 Step1, Step2 展示了两次数据加载可以完成一个输出 Tile 块的计算过程。

05Matrix03.png

矩阵乘的库

矩阵乘作为 AI 模型中的重要性能算子,CPU 和 GPU 的平台上都有专门对其进行优化实现的库函数。比如 CPU 的 OpenBLAS, Intel MKL 等,GPU 的 cuBLAS, cuDNN 等。实现的方法主要有 Loop 循环优化 (Loop Tiling)和多级缓存 (Memory Hierarchy)。

其两者的实现逻辑大概分为如下 2 步,关于 Kernel 实现优化的技术细节在[推理引擎]篇章进一步展开。

  1. Lib 感知相乘矩阵的 Shape
  2. 选择最优的 Kernel 实现来执行

下图展示了对矩阵乘进行 Loop 循环优化和多级缓存结合的实现流程。

05Matrix04.png

左边是共 6 级 Loop 循环展开的伪代码,右边是 Loop 对应多级存储的数据 Tilling 和搬移过程,假设矩阵乘 A,B,C 对应维度是(m, k, n)。

  • Loop5, Loop4, Loop3 对应把矩阵在 n, k, m 维度进行 Tilling 的切分,Tilling 后维度大小分别是 nc, kc, mc。
  • Loop2, Loop1 分别将 Tilling 后的 nc, mc 维度再一次 Tilling,Tilling 后维度大小分别是 nr, mr。
  • Loop0 对 kc 维度进行展开,实现累加求和的过程,得到(mr, nr)大小输出矩阵的部分和。

图中不同的颜色框指代了在不同存储层级上的数据计算,不同颜色块表示该块数据的存储位置。结合不同存储层级的内存空间和数据搬移带宽大小,将不同大小的 A,B 矩阵的 Tilling 块放在不同的存储层级上,可以平衡 AI 芯片执行矩阵乘任务时的时间和空间开销,提升整体算力利用率。比如,对(mr, nr)的计算过程,通过将 B 矩阵的(kc,nr)加载 1 次到 L1 cache 中,每次从 L2 cache 加载 A 矩阵的(mr, kc)大小到计算模块,进行计算,假设 mc 切分了 3 个 mr,则 B 矩阵的(kc, nr)就在 L1 中被重复利用了 3 次。这种用空间换时间或者用时间换空间的策略是进行算子性能优化的主要方向。

矩阵乘的优化

矩阵乘作为计算机科学领域的一个重要基础操作,有许多优化算法可以提高其效率。下面我们对常见的矩阵乘法优化算法做一个整体的归类总结。

  1. 基本的循环优化:通过调整循环顺序、内存布局等手段,减少缓存未命中(cache miss)和数据依赖,提高缓存利用率,从而加速矩阵乘法运算。

  2. 分块矩阵乘法(Blocked Matrix Multiplication):将大矩阵划分成小块,通过对小块矩阵进行乘法运算,降低了算法的时间复杂度,并能够更好地利用缓存。

  3. SIMD 指令优化:利用单指令多数据(SIMD)指令集,如 SSE(Streaming SIMD Extensions)和 AVX(Advanced Vector Extensions),实现并行计算,同时处理多个数据,提高计算效率。

  4. SIMT 多线程并行化:利用多线程技术,将矩阵乘法任务分配给多个线程并行执行,充分利用多核处理器的计算能力。

  5. 算法改进:如 Fast Fourier Transform 算法,Strassen 算法、Coppersmith-Winograd 算法等,通过矩阵分解和重新组合,降低了算法的时间复杂度,提高了计算效率。

这些优化算法通常根据硬件平台、数据规模和计算需求选择不同的策略,以提高矩阵乘法运算的效率。在具体的 AI 芯片或其它专用芯片里面,对矩阵乘的优化实现主要就是减少指令开销,可以表现为两个方面:

  1. 让每个指令执行更多的 MACs 计算。比如 CPU 上的 SIMD/Vector 指令,GPU 上的 SIMT/Tensor 指令,NPU 上 SIMD/Tensor,Vector 指令的设计。

  2. 在不增加内存带宽的前提下,单时钟周期内执行更多的 MACs。比如英伟达的 Tensor Core 中支持低比特计算的设计,对每个 cycle 执行 512bit 数据的带宽前提下,可以执行 64 个 8bit 的 MACs,大于执行 16 个 32bit 的 MACs。

如果您想了解更多AI知识,与AI专业人士交流,请立即访问昇腾社区官方网站https://www.hiascend.com/ 或者深入研读《AI系统:原理与架构》一书,这里汇聚了海量的AI学习资源和实践课程,为您的AI技术成长提供强劲动力。不仅如此,您还有机会投身于全国昇腾AI创新大赛和昇腾AI开发者创享日等盛事,发现AI世界的无限奥秘~

目录
相关文章
|
23天前
|
人工智能
AI对话网站一键生成系统源码
可以添加进自己的工具箱,也可以嵌入自己博客的页面中,引流效果杠杠的,新拟态设计风格,有能力的大佬可以进行二开,仅提供学习,用户可输入网站名称、AI默认的开场白、AI头像昵称、AI网站中引流的你的网站等等内容,所有生成的网页全部保存到你的服务器上
58 27
AI对话网站一键生成系统源码
|
8天前
|
机器学习/深度学习 人工智能 搜索推荐
BioEmu:微软黑科技炸场!生成式AI重构蛋白质模拟:千倍效率碾压传统计算,新药研发周期砍半
BioEmu 是微软推出的生成式深度学习系统,可在单个 GPU 上每小时生成数千种蛋白质结构样本,支持模拟动态变化、预测热力学性质,并显著降低计算成本。
21 2
BioEmu:微软黑科技炸场!生成式AI重构蛋白质模拟:千倍效率碾压传统计算,新药研发周期砍半
|
2天前
|
数据采集 人工智能 编解码
算法系统协同优化,vivo与港中文推出BlueLM-V-3B,手机秒变多模态AI专家
BlueLM-V-3B是由vivo与香港中文大学共同研发的多模态大型语言模型,专为移动设备优化。它通过算法和系统协同优化,实现了高效部署和快速生成速度(24.4 token/s),并在OpenCompass基准测试中取得优异成绩(66.1分)。模型小巧,语言部分含27亿参数,视觉编码器含4000万参数,适合移动设备使用。尽管如此,低端设备可能仍面临资源压力,实际应用效果需进一步验证。论文链接:https://arxiv.org/abs/2411.10640。
18 9
|
1天前
|
人工智能 自然语言处理 机器人
对话阿里云CIO蒋林泉:AI时代,企业如何做好智能化系统建设?
对话阿里云CIO蒋林泉:AI时代,企业如何做好智能化系统建设?
|
4天前
|
存储 人工智能 搜索推荐
Memobase:开源AI长期记忆系统,让AI真正记住每个用户的秘密武器
Memobase 是一个开源的长期记忆系统,专为生成式 AI 应用设计,通过用户画像和时间感知记忆功能,帮助 AI 记住、理解并适应用户需求。
39 0
|
2月前
|
人工智能 自然语言处理 数据可视化
校企合作|TsingtaoAI携手潍坊学院,共建AI驱动的党建信息化系统
TsingtaoAI与潍坊学院近日达成合作,正式签署《人工智能党建信息化系统开发》技术开发合同,计划在未来两年内联合开发一套集党员教育、党务管理、党建活动智能化以及数据可视化于一体的智能党建系统。本次合作将充分结合TsingtaoAI在AI大模型领域的技术优势和潍坊学院的学术资源,为推动党建工作的数字化、智能化和高效化注入新的动力。
53 10
|
2月前
|
机器学习/深度学习 存储 人工智能
AI实践:智能工单系统的技术逻辑与应用
智能工单系统是企业服务管理的核心工具,通过多渠道接入、自然语言处理等技术,实现工单自动生成、分类和分配。它优化了客户服务流程,提高了效率与透明度,减少了运营成本,提升了客户满意度。系统还依托知识库和机器学习,持续改进处理策略,助力企业在竞争中脱颖而出。
135 5
|
2月前
|
人工智能 运维 Serverless
云端问道8期方案教学-基于Serverless计算快速构建AI应用开发
本文介绍了基于Serverless计算快速构建AI应用开发的技术和实践。内容涵盖四个方面:1) Serverless技术价值,包括其发展趋势和优势;2) Serverless函数计算与AI的结合,探讨AIGC应用场景及企业面临的挑战;3) Serverless函数计算AIGC应用方案,提供一键部署、模型托管等功能;4) 业务初期如何低门槛使用,介绍新用户免费额度和优惠活动。通过这些内容,帮助企业和开发者更高效地利用Serverless架构进行AI应用开发。
|
5天前
|
人工智能 Java API
Java也能快速搭建AI应用?一文带你玩转Spring AI可落地性
Java语言凭借其成熟的生态与解决方案,特别是通过 Spring AI 框架,正迅速成为 AI 应用开发的新选择。本文将探讨如何利用 Spring AI Alibaba 构建在线聊天 AI 应用,并实现对其性能的全面可观测性。
|
3天前
|
人工智能 边缘计算 运维
容器化浪潮下的AI赋能:智能化运维与创新应用
近年来,容器技术以其轻量、高效、可移植的特性成为云原生时代的基石,推动应用开发和部署方式革新。随着容器化应用规模扩大,传统运维手段逐渐力不从心。AI技术的引入为容器化生态带来新活力,实现智能监控、自动化故障诊断与修复及智能资源调度,提升运维效率和可靠性。同时,AI驱动容器化创新应用,如模型训练、边缘计算和Serverless AI服务,带来更多可能性。未来,AI与容器技术的融合将更加紧密,推动更智能、高效的运维平台和丰富的创新应用场景,助力数字化转型。

热门文章

最新文章