Using Generative AI Models in Circuit Design

Generative models have been making big waves in the past few years, from intelligent text-generating large language models (LLMs) to creative image and video-generation models. At NVIDIA, we are exploring using generative AI models to speed up the circuit design process and deliver better designs to meet the ever-increasing demands for computational power.

Circuit design is a challenging optimization problem. Designers often need to balance several conflicting objectives, such as power and area, and satisfy constraints, such as hitting a specific timing. The design space is usually combinatorial, making it difficult to find optimal designs. Previous research into prefix circuit design used hand-crafted heuristics and reinforcement learning to explore the vast design space. For more details, see Towards Optimal Performance-Area Trade-Off in Adders by Synthesis of Parallel Prefix Structures and Cross-Layer Optimization for High Speed Adders: A Pareto Driven Machine Learning Approach.

While these methods help to overcome the vastness of the search space, they are associated with high computational costs to train and often have poor generalizability.

Our paper CircuitVAE: Efficient and Scalable Latent Circuit Optimization, recently published at the Design Automation Conference, provides a glimpse into the potential of generative models in circuit designs. We demonstrate that variational autoencoders (VAEs), a class of generative models, can produce better prefix adder designs at a fraction of the computational cost required by previous approaches.

The complexity of circuit design

The design space of prefix adders is vast: there are around $2^{n^2}$ possible designs for adding two n-bit numbers, so it is truly a needle-in-the-haystack challenge. A circuit consists of logic gates, such as AND, OR, and XOR. There are many ways of connecting them to implement the same functionality, though different connections yield changes in the observed power, performance, and area of the circuit.

In our paper, we focus on optimizing prefix adders, a class of circuits prevalent in modern GPUs. We represent a prefix adder with a tree, as shown in Figure 1. We minimize two metrics, area and delay, which we combine using a weighted sum into a single objective.

Three examples of prefix trees with their associated costs. — *Figure 1. Example 32-bit prefix trees and their costs*

What are variational autoencoders?

VAEs are generative models that estimate some data distribution. We can sample from the estimated distribution after training a VAE model. VAEs are versatile in modeling data of different modalities, from images to graphs. A VAE model consists of an encoder and a decoder.

A variational autoencoder model consists of an encoder, a latent space, and a decoder in a sequence. — *Figure 2. A variational autoencoder model. Credit: Wikipedia*

In the case of image generation, an encoder maps an input image to a distribution of vectors called a latent space. A decoder converts the vector of an encoded image back to an image. The VAE is trained by minimizing the reconstruction loss between inputs and outputs, along with a regularization loss on the latent space. VAEs are generative models because they can generate new outputs by sampling vectors from the latent space and decoding them with the learned decoder.

CircuitVAE: VAE for circuit design

CircuitVAE is a search algorithm that embeds computation graphs in a continuous space and optimizes a learned surrogate of physical simulation by gradient descent. It learns to embed circuits into a continuous latent space and predict quality metrics, such as area and delay, from latent representations. The cost predictor is fully differentiable when it is instantiated with a neural network. Thus, it’s possible to apply gradient descent in the latent space to optimize circuit metrics, circumventing the challenge of searching in a combinatorial design space.

A flowchart of the CircuitVAE algorithm. It consists of training a VAE with a cost predictor and optimizing the cost predictor to find new designs. — *Figure 3. The CircuitVAE algorithm flowchart*

CircuitVAE training

The CircuitVAE training loss has two parts:

The standard VAE reconstruction and regularization losses.
The mean squared error between the true and the predicted area and delay produced by the cost predictor model using the encoded circuit latent vectors.

While fitting the cost predictor, the latent space is organized according to costs, which is amenable to gradient-based optimization. A set of adders is generated through the genetic algorithm to bootstrap the training. One could also use a random sample of adders to start.

Gradient-based optimization

After training a CircuitVAE model, it’s used to find prefix tree structures that minimize costs. First choose a latent vector using cost-weighted sampling, a technique that ensures starting from a good design. This vector is then modified with gradient descent by minimizing the cost estimated by the cost predictor model. The final vector is decoded into a prefix tree and synthesized to obtain its actual cost.

Iteration

The full CircuitVAE algorithm interleaves the training and optimization stage. After each round of model training, more data is collected with gradient-based optimization and physical synthesis. Model fitting resumes with a growing dataset of circuits and associated metrics, resulting in a virtuous cycle where the cost predictor model increases in accuracy, leading to a more targeted optimization.

Results

We tested our approach on the design of circuits that have 32 inputs and 64 inputs (the width of the prefix tree circuit, corresponding to 32 bits and 64 bits). To supply our physical synthesis with components needed for simulation, we used an open-source cell library called Nangate45.

Figure 4 shows the cost progression while each method evaluates more designs through physical simulations. CircuitVAE consistently achieves the lowest costs compared to baseline methods. Both RL and GA optimize in the discrete domain and are slow to explore, while CircuitVAE is 2-3x faster, thanks to gradient-based optimization in the latent space.

Six plots with results on designing 32-bit and 64-bit prefix adders with three delay weights. — *Figure 4. Experiment results on designing 32-bit and 64-bit prefix adders*

We evaluated CircuitVAE in a real-world prefix adder task with a proprietary cell library with input-output timings captured from a complete datapath. Figure 5 shows that CircuitVAE generates designs that form a better Pareto frontier of area and delay than a commercial tool.

A Pareto plot of area and delay comparing CircuitVAE with a commercial tool and classical designs. — *Figure 5. CircuitVAE Pareto dominates a commercial tool in designing a real-world datapath design*

Conclusion

CircuitVAE showcases the power of generative models in circuit design tasks. Operating in a latent space, rather than the combinatorially large discrete space of circuit designs, reaps the benefits of continuous optimization in the form of reduced computational costs. We believe this transformation from discrete to continuous holds promise in other areas of hardware design, such as place-and-route. We anticipate that generative models will play an increasingly central role in hardware design.

For more information about CircuitVAE, see CircuitVAE: Efficient and Scalable Latent Circuit Optimization.