AI芯片“软硬件协同设计”的理想与实践
最近看到一篇文章“Hardware-Software Co-Design Reappears”[1]，很好的总结了软硬件协同设计“失败”的历史，也提出一些新的想法，这个话题在AI芯片的背景下还是非常值得探讨的。
首先，文章对于软硬件协同设计做了如下描述：“The initial idea behind co-design was that a single language could be used to describe hardware and software. With a single description, it would be possible to optimize the implementation, partitioning off pieces of functionality that would go into accelerators, pieces that would be implemented in custom hardware and pieces that would run as software on the processor—all at the touch of a button.”
我觉得这是对于软硬件协同设计的理想非常精准的描述。其中几个核心概念包括：1. 一种能够描述软件和硬件的单一语言；2. 优化实现；3.功能划分，包括不同功能在不同类型硬件上的实现；3. 所有功能“一键”完成（自动化）。
遗憾的是，到目前为止，上述理想并没有实现。[1]文章对失败的原因做了一些分析。首先，在问题提出的时代，软硬件系统相对比较简单。做功能划分的时候，为了抵消通信开销，一般粒度都比较粗。考虑这些环境的因素和限制，软硬件划分和优化往往都是非常直观的（比较显而易见的），实现自动化的软硬件协同设计的价值不大。另一个方面的原因主要在于工作文化。软硬件协同设计希望能够打破软硬件团队间的界限。但实际上，要求软件和硬件团队采用协同设计方法绝对不是件容易的事情。虽然目前软硬件系统的复杂度越来越高，但似乎还没有到必须颠覆传统方法的程度。对于这些问题，我在之前的文章“在体系结构黄金期，ESL设计方法学能否“焕发青春”?”也有过类似的讨论。在一个组织中实施新的设计方法的还会有很多其它的阻力，这里就不再赘述。
回到现在，在AI芯片领域，在一些公司宣传自己的技术优势或者大家讨论主要挑战的时候，我们经常可以听到软硬件协同设计的说法。应该说，对于它在AI芯片领域的重要意义，大家是有共识的。而在实际项目中，不同团队也正在或多或少的践行类似思想。
即使到今天，要实现理想的软硬件协同设计还是非常困难的。首先，”一种语言“就几乎是不可能的。当年，SystemC的提出是希望能够架起软硬件的桥梁，通过一种语言来描述软件和硬件，支持从算法到硬件的自动设计和优化过程。但它并没有实现设计的初衷。软件和硬件在思维方式上的差异就决定”一种语言“的尝试很难成功。而现在流行的各种DSL（Domain-specific Language），似乎是反其道而行之，一个任务就来一种更高效的语言。所以我们思考的是，做软硬件协同设计，是不是真的需要统一到一种语言？。第二，优化实现，软硬件划分和功能映射，是架构探索的问题。不管对人工还是自动工具而言，都存在一个探索粒度的平衡点。对人工方法来说，探索的粒度不可能太细。使用自动化的方法，这个平衡点其实也不是一概而论的，不同目标，不同约束，就会发生变化。我多次说过，AI芯片是个Domain-specific问题，相对而言架构探索的空间缩小了许多。在一定粒度（比如Data Transaction层面）和架构假设下，实现自动探索是有可能的。但即便如此，要实现软硬件联动的优化，即算法实现（比如Graph层面的优化），硬件架构和映射方法（比如Operator层面的优化）都允许发生变化，还是相当困难的。第三，理想的软硬件协同设计是一个自动化的过程，这也是我们一直追求的目标。但真正的no-human-in-the-loop是非常困难的，在很多时候也不是刚需。我们可以先在一些小的loop里实现自动化，而在关键点还是要靠人把握（也可以避免大家面临失业问题
总得来说，经典的（或者说狭义的）软硬件协同设计是一个非常理想的愿景，但实现起来困难重重。不过我们回到问题本身，软硬件协同设计的本质是希望把软件和硬件的设计优化过程统一起来，而不是孤立对待，以在整个系统上获得更多的收益。在这个意义上，即使我们不能实现完全自动的协同设计，我们还是有很多事情可以先做起来的。
文章[1]也提到，随着软硬件协同设计概念的提出，一些相关技术不仅生存了下来，并且得到了很好的发展。其中一个最好的例子就是虚拟原型技术（Virtual Platform）。我们做AI芯片，可以在没有物理硬件的情况下快速搭建一个虚拟的平台，让软件工程师能够更早的在”目标硬件“（模型）上调试自己的程序。这种平台一般包括CPU（比如ARM）的编程模型，程序员可以像使用真正的开发板一样在这个平台上进行调试。现在的虚拟原型工具已经非常好用。我们自己搭建一个ARM+NVDLA（SystemC model）的简单SoC虚拟原型，运行起来操作系统和NVDLA的简单inference实例只用不到两周的时间。由于虚拟原型主要面向软件开发，你的AI加速部分只需要建一个编程接口的模型就可以，这个建模的代价也不算大。除了能在硬件开发的同时就可以给软件工程师一个实际的软件开发调试环境，虚拟原型在很大程度上也是软硬件工程师交流的接口。软件工程师拿到的是一个可运行程序的模型，而不是简单的spec文档。硬件工程师可以得到软件运行的反馈，比如编程模型是不是合理。更进一步，虽然虚拟原型的主要目标是功能性模拟，但也可以增加一些的粗粒度的性能分析，并把这些信息反馈给硬件设计者，或者架构设计工具，作为硬件优化的参考。而对于硬件实现来说，虚拟原型也是整个验证过程的开始，可以给芯片验证打下很好的基础。总得来说，对于一个希望建立软硬件协同设计机制的团队来说，虚拟原型方法可能是最合适的第一步实践。
软硬件协同设计的一个主要障碍是软硬件工程师的思维方式和使用的语言的差异。这也是为什么最初大家希望能用一个统一的语言来描述问题。虽然这个尝试并不成功，但我们是不是可以找到更好的描述问题的方法呢？也许我们可以用一个多层次的语言来渐进式细化问题。这里我们可以参考一下Google前段时间推出的MLIR（Multi-Level Intermediate Representation）。MLIR的基本目标是提供一个统一的IR和编译器框架（主要针对机器学习应用），包括了设计编译器需要的工具。对于它的详细讨论，大家可以参考知乎上的问题”如何看待Google的MLIR项目？“，其中杨军的回答已经相当全面。这里我想讨论的重点是Multi-Level的概念。我们先看一下MLIR试图解决的问题：
The ability to represent all TensorFlow graphs, including dynamic shapes, the user-extensible op ecosystem, TensorFlow variables, etc.
Optimizations and transformations typically done on a TensorFlow graph, e.g. in Grappler.
Quantization and other graph transformations done on a TensorFlow graph or the TF Lite representation.
Representation of kernels for ML operations in a form suitable for optimization.
Ability to host high-performance-computing-style loop optimizations across kernels (fusion, loop interchange, tiling, etc) and to transform memory layouts of data.
Code generation "lowering" transformations such as DMA insertion, explicit cache management, memory tiling, and vectorization for 1D and 2D register architectures.
Ability to represent target-specific operations, e.g. the MXU on TPUs.
虽然这些问题主要是从软件视角出发，但很多已经显然踏入了硬件领域。换句话说，MLIR是要在软件优化中引入硬件信息的。而对于硬件设计者来说，也可以通过适当层次的IR更好的理解软件（或者编译器）的需求。相对而言，MLIR可能实现一种软硬件设计者都能看懂（或者看懂关键部分）的语言。在另一篇介绍MLIR的文章中[2]，我们也可以看到它的目标人群同样覆盖了软硬件设计者：
Compiler researchers and implementers looking to optimize performance and memory consumption of machine learning models
Hardware makers looking for a way to connect their hardware to TensorFlow, such as TPUs, portable neural hardware in phones, and other custom ASICs
People writing language bindings that want to take advantage of optimizing compilers and hardware acceleration.
MLIR实现多种层次的表达主要是通过”方言“（Dialect）的形式，不同的方言共享基本的格式和基本操作，增加用于描述目标软件和硬件的内容。其实从某种程度来说，SystemC也可以看作是C++的一种方言。但相对而言，MLIR提供了在多个层次上描述问题的能力和在不同层级间转换的便利。
目前MLIR还在快速更新当中，主要在做的事情是统一各种神经网络framework的前端和后端间的接口。但我感觉它的一个潜在价值在于对软硬件协同设计的支持，也欢迎感兴趣的朋友和我交流。
前面分析的基本上是比较传统的软硬件协同设计方法思考的问题。如果更极端一些，未来的软硬件架构和关系也可能会发生新的变化，那么软硬件协同设计就可能需要新的思考。首先，”重软件，轻硬件“的趋势下，未来硬件可能更加趋向于简单和同构。TVM/VTA就是一种新的软硬件协同设计的方法（“全栈”开源的VTA会给AI芯片产业带来什么？），先做软件栈，并由软件栈的经验导出硬件设计。我们也听到其它一些团队在倡导这种方法。在这种模式下，所谓的协同更多的偏向软件这边，软件占据绝对的主导权，这其实也是一种”简单粗暴“的软硬件协同设计思路（当然很多硬件工程师会反对）。另一种模式是软件定义硬件（Software Defined Hardware），相对没有这么极端。硬件要考虑的主要是提供一个效率和性能的平衡点，而面向更细的应用需求的优化则靠软件完成。另外，文章[1]中提出的各种新机会，比如开源硬件，嵌入FPGA等等也给我们一些新的选择。这些不同的方向和思路也都是值得关注的。
Reference:
题图来自网络，版权归原作者所有
本文为个人兴趣之作，仅代表本人观点，与就职单位无关