Preface
Human demand for computing power is endless. From the earliest mainframes manually wired to compute ballistic trajectories, to scientific programs encoded on punched paper tape, to the emergence of the first high-level language fortran, and to today’s diverse libraries, frameworks, and programming paradigms. Hardware performance has grown ever stronger, while software has gradually evolved more efficient and flexible techniques to fully leverage hardware capabilities.
In recent years, with the resurgence of deep learning and the massive datasets accumulated through the growth of the internet, traditional general-purpose processors can no longer meet the exponential increase in compute demand. Meanwhile, diverse model deployment scenarios in real-world production—such as combined constraints on latency, power consumption, and cost—have jointly posed new challenges for hardware and software technologies. Looking across mainstream deep learning frameworks on the market, one can see that they all contain substantial abstraction and code to simultaneously support CPU and GPU, and even more hardware architectures. “By examining the past, one can understand rise and fall.” Therefore, it is meaningful to review the design ideas and evolutionary history of these frameworks from the perspective of heterogeneous computing; in doing so, we can also roughly anticipate the future direction of foundational technologies related to deep learning.
First Generation
Before Caffe appeared, training relatively large convolutional neural networks required researchers to design heterogeneous computing programs for specific architectures themselves (for example, early convnet). This demanded programming skill and an understanding of the relevant architectures; in practice it was not only inefficient to develop, but also unfavorable for reuse. Subsequently, first-generation deep learning frameworks represented by Caffe used a general protobuf file as the model description format, and a structurally relatively separated, concise, high-quality C++/CUDA implementation, solving the problem of quickly building models and running them efficiently on heterogeneous hardware such as GPUs, thereby meeting the enormous compute demand in computer vision.
However, the wide variety of model topologies in deep learning cannot all be abstracted using Caffe’s simple “pipeline of layers” model. On the other hand, Caffe also cannot adequately meet the diverse and broad needs for operator customization: adding a new operator requires separately developing its Forward and Backward processes, and to accelerate it you still have to write CUDA kernels yourself, which continues to place demands on researchers’ heterogeneous programming skills. As a result, major companies and even individual algorithm teams would maintain their own forked versions of the caffe framework.
Second Generation
Taking MxNet as an example, second-generation deep learning frameworks improved the shortcomings of the previous generation in two ways. At the user API level, they primarily use static dataflow graphs—rather than the simple pipeline pattern—to build deep learning models. This abstraction allows users to construct models more flexibly without worrying about details such as how to compute gradients; from the perspective of internal framework implementation, it enables transformation and optimization of the entire dataflow graph at a lower and finer-grained level, such as memory allocation, data layout reordering, device planning, and so on.
From the perspective of how heterogeneous hardware is used, most frameworks of this generation adopted the Expression Template design pattern (for example, mshadow). The core idea is to decouple an operator’s numerical algorithm implementation from its scheduling and execution strategy, thereby using template language to encapsulate platform differences between heterogeneous hardware and CPU. This makes it possible, when developing new operators, to write heterogeneous code efficiently using existing template encapsulations, without implementing the same algorithm twice. However, the drawbacks of this technique are also quite obvious. Note the phrasing above: “write efficiently,” not “write efficient heterogeneous code.” This is because Expression Template is ultimately a C++ programming model and is statically determined at compile time—within the code context it can obtain neither runtime information nor model structure information—so operators implemented this way cannot be optimized at all. The root cause is that we are trying to implement a general numerical algorithm that supports arbitrary runtime parameters, but can only describe it using the extremely constrained C++ template language that is statically fixed at compile time; thus, beyond making it work, we cannot do more advanced things.
Overall, second-generation frameworks provided more flexible and user-friendly interfaces, added distributed support, and also improved performance and memory footprint. One could say they basically met the needs of typical model-training users.
Third Generation
It is commonly believed that Caffe and Tensorflow represent the first and second generations of deep learning frameworks, respectively, but there is some disagreement on how to define third-generation technology. From the user-interface perspective, many believe that support for the “imperative program paradigm,” the so-called “dynamic graph,” represented by Chainer and Pytorch, is the main feature of the next generation of frameworks. In fact, mature frameworks have also extended similar interfaces on top of prior-generation technology, such as Tensorflow Eager, MxNet Gluon, Paddle Fluid, and so on. However, these technologies essentially trade off performance to some extent in exchange for user flexibility, and cannot be considered true technological innovation.
From the heterogeneous computing perspective, Tensorflow XLA (and its derivatives JAX), TVM, and Tensor Comprehension (hereafter TC) and other deep learning compilers can be seen as representing third-generation deep learning framework technology. While retaining a flexible top-layer user interface, they attempt to thoroughly solve the development and optimization of heterogeneous algorithm programs through more advanced compilation techniques, fully extracting hardware compute performance.
Tensorflow XLA abstracts some commonly used operators in deep learning into HLO IR (High Level Optimizer), which is essentially a series of particularly fine-grained operators that can be composed into any practical operator. For example, the common batch_norm can be composed of a series of operations such as broadcast, reduce, and element_wise. In XLA, the idea of mapping each HLO operator to LLVM IR can be summarized as: writing C++ template code in a richer-information context and in a more abstract way. XLA’s main shortcoming also lies here: it essentially uses LLVM IR to describe the specific algorithm of each XLA OP. Although it can directly reuse a series of existing optimizations in LLVM, it limits to some extent its ability to extend to other hardware architectures. In fact, LLVM IR itself is designed to abstract different processor architectures, which means that from a low-level IR we have already completely lost the structural information of the algorithm itself.
TVM takes a more aggressive technical path. It not only includes a fine-grained computation graph representation similar to HLO, NNVM (later evolving into Relay IR), but also incorporates Halide IR as a lower-level IR than HLO, so that all algorithms can be described directly using a clearly defined domain-specific language (DSL, Domain Specific Language). The backend can create scheduling plans (Schedule) for existing algorithm descriptions and apply a series of scheduling primitives for optimization, thus achieving good performance. The optimization process can be manually written by developers who understand the hardware architecture, or it can use AutoTVM to abstract general search rules for the hardware architecture, enabling semi-automated code optimization. The biggest advantage of using an abstract DSL is that it is easy to extend to different types of hardware architectures.
TC is the most aggressive of these. It is also based on Halide IR, but during code generation it directly adopts polyhedral model optimization techniques from compiler theory: by affine transformations it maps nested loops to a high-dimensional vector space, and performs fully automated loop transformation and optimization based on hardware architectural characteristics. The entire process can require no human intervention.
The efforts above attempt to use different types of compiler technologies to address the optimization of heterogeneous programs, and in many scenarios can achieve better results than manual optimization. Whether in terms of technical depth or practical application scenarios, compared to user-interface-level improvements such as “dynamic graphs,” they better represent the technology of this generation of frameworks.
Looking Ahead
Today, with deep learning frameworks flourishing and contending (it always feels like the next line should be “lure the snake out of its hole”), and with the overall landscape trending toward stability and maturity, there are also views that this is merely major companies building their own wheels and pushing standards in the AI era in an attempt to seize an early advantage in the developer market. Moreover, most of them are nothing more than high-level wrappers around cuDNN and cuBlas; the differences lie only in whether the upper-layer API is easy and flexible to use, yet it is hard to conceal the overall severe homogenization and lack of new technologies and new ideas. But if we review the development history of deep learning frameworks above and think from the combined software-and-hardware perspective, we can clearly trace a throughline: how software frameworks and heterogeneous hardware have iterated and advanced together, and where the next innovation points may lie.
First-generation deep learning frameworks solved how to use heterogeneous hardware to meet the massive compute demand of the reborn deep learning field. In contrast, traditional scientific computing fields whose algorithmic characteristics are not particularly suitable for heterogeneous acceleration—for example, large-scale computational fluid dynamics numerical simulation in physical oceanography—still have to rely on traditional parallel computing techniques such as MPI clusters. The peak floating-point performance of an entire cluster may not even match that of a Kepler compute card. Second-generation frameworks adopted higher-level abstraction mechanisms and design patterns, greatly reducing the complexity of developing heterogeneous algorithms, but constrained by the characteristics of C++ itself, such implementations find it hard to reach optimal performance. Third-generation frameworks introduced automated code tuning and a broader range of heterogeneous hardware platforms, greatly reducing the costs of porting and deploying deep learning models.
From the perspective of user needs, the development history across these three generations of deep learning frameworks is the process of moving from meeting compute demand, to improving development efficiency and flexibility, to pursuing the best hardware utilization (across architectures). Looking back at the earlier process, it has essentially always been software developers trying every possible way to use(gui) heterogenous hardware well(tian), but existing hardware architectures may not necessarily meet needs well across all scenarios. Naturally, the next generation of deep learning frameworks should no longer be merely software frameworks, but a complete solution for maximizing heterogeneous compute across scenarios through hardware/software co-design, so as to fully unleash heterogeneous compute in diverse scenarios. In fact, during the last boom period of neural network development, academia had already attempted and explored in this direction. But at that time, general-purpose processors were advancing rapidly; if developers and users only needed to wait a year to get doubled model performance, they would naturally not choose specialized accelerators whose generality and usability were full of pitfalls. Today, with the breakdown of Moore’s Law, more and more scientific computing fields can no longer count on such free windfalls. Introducing domain-specific architecture (DSA, Domain Specific Architecture) is bound to be the future trend. Patterson’s recent reports and talks have repeatedly emphasized this, arguing that this trend marks the return of a golden age of computer architecture.
However, from a more realistic perspective, domain-specific architectures must achieve sufficiently low tape-out cost, sufficiently high hardware development efficiency, and a sufficiently usable software toolchain to be truly competitive in research and real applications. In many scenarios, computation is a means rather than an end; if the goal cannot be achieved at an acceptable development cost, users would rather switch to a general solution with lower performance but higher development efficiency. So far, in the scientific computing direction, only deep learning has produced a custom architecture competitive enough, such as TPU. This is because the enormous economic value in this field can sustain the massive investment required for chip and full software-stack R&D—indeed, the R&D cost of XLA may be even higher than the hardware.
DARPA’s Electronics Resurgence Initiative is precisely trying to address the first two issues above, aiming to drastically reduce tape-out costs over the next decade, improve the efficiency of hardware R&D processes, and establish an open-source IP ecosystem as prosperous as software. But even if we can reduce the cost of designing new hardware architectures, one remaining problem is that we cannot complete all algorithms before tape-out to ensure they can all run efficiently on that architecture, and we certainly cannot iterate a complete, easy-to-use software toolchain in the short term. Given such constraints, how can we ensure that domain-specific architectures are truly competitive? In this regard, Intel’s VLIW (Very Long Instruction Word) architecture is a classic failure example—hardware architects believed the burden of resolving data dependencies could be entirely handed over to the compiler, but in reality the problem is not that simple. For general-purpose computing scenarios, the compiler lacks the ability to span multiple contexts and perform large-scale instruction or even code-logic reordering to eliminate the randomness of off-chip memory accesses. Therefore, hardware architecture design cannot make unfounded simplifications; it must fully consider the upper bound of existing compiler capabilities.
From the software perspective, given a numerical algorithm, requiring it to run efficiently on different hardware architectures is essentially still a scheduling and planning problem in compilers. Different hardware architectures have different scheduling characteristics, and the design goal of some domain-specific languages that have emerged in recent years is precisely to abstract these characteristics, thereby decoupling algorithm implementation from hardware scheduling. Therefore, domain-specific architectures should be designed around the descriptive scope of domain-specific languages. In this way, existing DSL compilation toolchains can traverse all supported algorithm descriptions and introduce some machine learning and heuristic search algorithms to generate and optimize executable code for the target architecture, thereby quickly evaluating the architecture’s efficiency and performance. Ultimately, once a baseline hardware architecture is determined, a DSL-based design process will be able to form a fully automated closed-loop iteration like this:
- Tune hyperparameters in the hardware architecture via deep learning models
- Use high-level hardware description languages (such as Chisel) to rapidly generate hardware RTL and simulators
- The compiler, based on the algorithm’s DSL description, quickly generates all code implementations for the hardware architecture
- Evaluate algorithm performance and feed the corresponding results back to the model in step one, thereby starting the next iteration
In this way, we may be able to solve the problems of software usability and optimization efficiency, giving domain-specific architectures real vitality. In addition, one point worth noting is that DSL technologies are essentially a double-edged sword. They are not suitable in every scenario; for example, this blog post criticizes those scenarios where they do not apply. However, from the perspective of heterogeneous program development, if we can use a DSL to abstract most algorithms and significantly improve development and optimization efficiency, while also assisting the design of specialized hardware architectures, then there is reason to believe such a DSL is necessary and successful.
On a longer time horizon, for various scientific computing fields, it may be possible to abstract different DSLs for each domain in the future, covering the vast majority of algorithm descriptions, scheduling, and optimization needs within them. For example, the deep learning field, dominated by matrix multiply-accumulate operations, will certainly have a DSL representing its computational characteristics that is very different from that of the fluid dynamics field, which involves large amounts of iterative computation over partial differential equations. In the future, when the iteration process for specialized hardware design becomes more mature, tape-out costs drop significantly, and hardware is deeply integrated with software toolchains, we will truly usher in a new era in which heterogeneous computing replaces general-purpose computing.