Arranging ‘billions’ of components on a tiny surface area of a computer chip is a complicated process. It calls for precise decision-making at every step of the way and requires a designer with years of experience in laying out circuits that squeeze power efficiency from nanoscopic devices.
Designers now tap the latest AI advancements to learn the processes involved in chip designing to help draw up more powerful blueprints in less time. It allows engineers to co-design an AI software to find the optimal configuration with different designer perspectives.
Nvidia
Principal research scientist, Haoxing Ren, spoke to Wired about testing reinforcement learning on AI chips to arrange components and wiring. The process includes exploring chip designs in simulation, training a neural network to recognise the efficient decisions to produce a high performing chip etc. According to Ren, this approach can create an excellent chip while cutting the engineer’s work in half.
At Google, designers are using ML to create chips at speed far superior to humans. Algorithms are finishing months worth of work in just six hours. Google has been successful in applying AI commercially in chipmaking. The study has been published in Nature. The paper talks about the upcoming version of Google’s TPU chips optimised for AI computation.
In the paper, the researchers explained the task of ‘floorplanning’ that usually requires a human designer but is now tackled by Google’s algorithms. The planning involves CPUs, GPUs and memory cores connected with tens of kilometres of writing- with a lot of importance on the placement on the chip. It is essential to decide and place each component on a die to maintain the eventual efficacy of the chip; even a nanometer of change can have huge effects.
Google’s engineers called this floor plan designing a ‘board game’ for machines. The ‘game’ here is a silicon die instead of a board game, and CPUs & GPUs instead of game pieces- with AI tasked with finding the board’s ‘win conditions’ aka computational efficiency.
Google engineers have trained algorithms with a dataset of 10,000 chip floor plans- each tagged with a reward function depending on its efficacy across different metrics. Then, through reinforcement learning, the algorithm used this data to differentiate between good and bad floor plans.
Reinforcement learning
According to Song Han, an assistant professor of electrical engineering at MIT, reinforcement learning has significant potential for improving the design of chips since it is difficult to predict ‘good’ design without experience. The input to our model is the chip netlist (node types and graph adjacency information), the ID of the current node to be placed, and some netlist metadata, such as the total number of wires, macros, and standard cell clusters, according to a Google AI blog. The netlist graph and the current node are passed through an edge-based graph neural network to encode the input state. This generates embeddings of the partially placed graph and the candidate node’.
Credits: Google AI Blog
The embeddings are concatenated to form a single state embedding and passed to a feedforward neural network whose output is a learned representation of the valuable features. These are the input for the policy network probability distribution for node placement. ‘RL training is guided by a fast-but-approximate reward signal calculated for each of the agent’s chip placements using the weighted average of approximate wire length and approximate congestion’. The study observed that pre-training improved the sample efficiency and placement quality.
