Reinforcement learning for chip design

Daniel and Chris have a fascinating discussion with Anna Goldie and Azalia Mirhoseini from Google Brain about the use of reinforcement learning for chip floor planning - or placement - in which many new designs are generated, and then evaluated, to find an optimal component layout. Anna and Azalia also describe the use of graph convolutional neural networks in their approach. Sponsors: Linode https://linode.com/changelog – Our cloud of choice and the home of Changelog.com. Deploy a fast, efficient, native SSD cloud server for only $5/month. Get 4 months free using the code changelog2019 OR changelog2020. To learn more and get started head to linode.com/changelog https://linode.com/changelog. AI Classroom https://datadan.io – An immersive, 3 day virtual training in AI with Practical AI co-host Daniel Whitenack. Get 10% off using the code PRACTICALAI10. To learn more and purchase tickets go to datadan.io https://datadan.io/. Fastly https://www.fastly.com/?utm_source=changelog&utm_medium=podcast&utm_campaign=changelog-sponsorship – Our bandwidth partner. Fastly powers fast, secure, and scalable digital experiences. Move beyond your content delivery network to their powerful edge cloud platform. Learn more at fastly.com https://www.fastly.com/?utm_source=changelog&utm_medium=podcast&utm_campaign=changelog-sponsorship. Rollbar https://rollbar.com/changelog – We move fast and fix things because of Rollbar. Resolve errors in minutes. Deploy with confidence. Learn more at rollbar.com/changelog https://rollbar.com/changelog. Featuring: Anna Goldie – GitHub https://github.com/annagoldie, LinkedIn https://www.linkedin.com/in/adgoldie, X https://x.com/annadgoldieAzalia Mirhoseini – LinkedIn https://www.linkedin.com/in/azalia-mirhoseini-a5308522, X https://x.com/azaliamirhChris Benson – Website https://chrisbenson.com, GitHub https://github.com/chrisbenson, LinkedIn https://www.linkedin.com/in/chrisbenson, X https://x.com/chrisbensonDaniel Whitenack – Website https://www.datadan.io/, GitHub https://github.com/dwhitena, X https://x.com/dwhitenaShow Notes: Their research paper https://arxiv.org/ftp/arxiv/papers/2003/2003.08445.pdfGoogle Brain https://research.google/teams/brainGoogle is using AI to design chips that will accelerate AI | MIT Technology Review https://www.technologyreview.com/2020/03/27/950258/google-ai-chip-design-reinforcement-learningPractical AI episode #47: GANs, RL, and transfer learning oh my! https://changelog.com/practicalai/47Upcoming Events: Register for upcoming webinars here https://practicalai.fm/webinars!

Speaker names marked ? are inferred from speaking time and the show’s known hosts — they are not voice-verified.

• 0:00Speaker 00

  Being with your change log is provided by Fastly. Learn more at fastly.com. We move fast and fix things here at Change Law because of Rollbar. Check them out at rollbar.com. And we're hosted on Linode cloud servers. Head to linode.com slash change log.

• 0:17Speaker 00

  Do not underestimate the power of the independent open cloud for developers. Yes, I'm talking about Linode. Linode is our cloud of choice and it's the home of changelog.com. What we love most about Linode is their independence and their commitment to open cloud. Open cloud means being unencumbered by outside investment and maximizing value for the community, not shareholders. And that's exactly what Linode represents. No vendor lock-in, open at every layer. If you want to learn more, head to linode.com slash open. Again, linode.com slash open.

• 1:01Speaker 01

  Welcome to Practical AI, a weekly podcast that makes artificial intelligence practical, productive, and accessible to everyone. This is where conversations around AI, machine learning, and data science happen. Join the community and Slack with us around various topics of the show at changelog.com slash community and follow us on Twitter. We are at Practical AI FM. Okay, take it away, guys.

• 1:26Daniel Whitenack?

  Well, welcome to another episode of Practical AI. This is Daniel Whitenack. I'm a data scientist with S I L International. And I'm joined as always by my co-host, Chris, who is the Principal AI Strategist at Lockheed Martin. How are you doing, Chris?

• 1:43Chris Benson?

  I am doing very well. How's it going today, Daniel?

• 1:45Daniel Whitenack?

  It's going very well, staying very busy. As we were talking before the episode, it seems like after the crisis and after being everyone at home right now. So for future listeners, this is still during the COVID-19 crisis. It seems like I'm more busy, more busy work-wise now than even before. What about you?

• 2:07Chris Benson?

  I think it's the same and we're just trying to not be paranoid. We're right in the middle of pollen season and everyone in my family suffers. So you know, you get sore throats and you're coughing and you're like, oh my God, that's not, you know, symptomatic of. Of COVID or something. So we're just trying to maintain our calm.

• 2:24Daniel Whitenack?

  Yep.

• 2:24Chris Benson?

  Get through this and all as well. Other than that, it is a lovely day here in Atlanta, Georgia.

• 2:30Daniel Whitenack?

  Yeah, it's a beautiful time to get outside as well. Of course, maintaining distance from others. But we've got something I think pretty interesting to chat about today. In the midst of all of the COVID-19 and coronavirus tweets and articles that I've been seeing and reading and all of that stuff, I was able to pick out this one story that seemed really interesting to me that was was not related to COVID-19. And that was this story about a team at Google who was using reinforcement learning to somehow design chips like hardware computing chips. And we were joined today by Anna Goldie and Azalea Mir.

• 3:15Daniel Whitenack?

  I'm sorry to mess up the name Mir Sany. You can correct me here in a second on that one. Sorry, Azalea. Welcome to the show.

• 3:24Speaker 04

  Thank you so much for having us.

• 3:25Speaker 05

  Yeah, thanks for having us.

• 3:27Daniel Whitenack?

  Yeah, excited to have you both and excited to chat about this amazing project. It was really interesting when I read it. But before we jump into the project itself, I would love to hear a little bit about both of your backgrounds and how you ended up doing what you're doing now. So maybe we could start with Azalea. Could you give us a little bit of information about your background?

• 3:49Speaker 05

  Yes, my PhD from Rice University in computer engineering. My work was or my thesis was focused on co-design of hardware software systems for machine learning applications. And then when I joined Google, I joined Google Brain to the residency program. I stayed at Google Brain as a resident for a year. And that was the time that I developed passion for the work at the intersection of ML for system. Like, how do we develop new machine learning algorithm and deep learning algorithms for system? And ever since then, I've been at Google Brain for like almost four years now and enjoy doing research and work on impactful projects.

• 4:37Daniel Whitenack?

  Awesome. And now you're a senior research scientist at Google Brain, is that correct?

• 4:41Speaker 05

  That's correct, yes.

• 4:43Daniel Whitenack?

  Awesome. Well, thank you for the background. Anna, do you want to give us a little bit of information about your background as well?

• 4:51Speaker 04

  Yeah, sure. So I studied computer science and linguistics, actually, at MIT. And I did my master's basically building a Mandarin speaking dialogue system. I joined Google about eight years ago, Google Research. And I've been working mostly before this on natural language processing applications. And about like two and a half, three years ago, I started working with Azalia. I actually saw some parallels to natural language processing in some of these systems problems. Because a lot of them can be formulated as sequential decision-making problems. And it's been just wonderful working with Azalia. And we have such an awesome team. Solving, basically trying to use machine learning to optimize and automate various problems in computer systems.

• 5:37Daniel Whitenack?

  So could you give us a little bit more of an idea about this team and how big it is? I know you both work closely together, but what's the sort of team like that you're working on?

• 5:47Speaker 04

  Actually, when we first started, it was basically just me and Azalia. And then we've been gradually growing. I think we have something like maybe eight, 10 people on the research side. And then we also partner with chip designers who are building the next generation of TPUs. And maybe there's about eight, 10 people on that side as well. So it's, you know, pretty substantial effort at this point.

• 6:08Chris Benson?

  So kind of wondering, you know, I come from a software-only background. And as we kind of dive in, you know, I'll probably be more comfortable on the reinforcement learning part. But I knew nothing about creating chips and stuff. And I guess if you could just kind of lay out the context of what that means. What does it mean? I've heard the phrase chip placement, and I've heard chip floor planning and stuff like that. Could you talk a little bit about the baseline, about what it is you're trying to do? And how that is, what the context of it from the hardware side, which I know nothing about?

• 6:39Speaker 04

  Uh, sure. Maybe I can take a stab at that. Um, so basically, this is just one of the stages of chip design. Um, there's already been a bunch, basically, computer architecture, stage and synthesis occurs first. But the problem that we were solving in our research was taking a graph of chip components, which is called a netlist. So it's basically a bunch of SRAMs, which are, or memory components, macros, and standard cells, which are logic gates, like NANDs and NORs. All of these are connected by wires, and so it's a graph. And we want to place that graph onto this two-dimensional grid, such that we minimize various costs, like latency of computation, power consumption, wire length area, while adhering to hard constraints on density and congestion.

• 7:22Speaker 04

  So that, that's sort of the core problem that we're trying to solve.

• 7:25Daniel Whitenack?

  So do I have it right? Like, when you say that this is kind of graph-structured, you're meaning like, there's this component, like, something physical that has to go on the chip, and then there's this other components, and they need to be linked by an electrical connection, I guess?

• 7:43Speaker 04

  Exactly.

• 7:44Daniel Whitenack?

  Is that like a way of saying, so like, the graph is formed of these components and the electrical connections between them? Is that, is that kind of?

• 7:50Speaker 04

  Exactly, yeah. There's all these, like, sort of logic and memory components connected by wires, or like, yeah, electrical connections. And then we physically need to decide where to place them, so that's that we get better performance for that chip.

• 8:02Daniel Whitenack?

  Gotcha.

• 8:03Chris Benson?

  And can you talk a little bit about what that means? When you talk about, why does physical placement have an impact on performance? And what it is about that placement which affects performance?

• 8:14Speaker 04

  Um, so one way you could think about it is, the timing of computation, or the amount of time it takes to compute with this circuit, is affected by, like, the lengths of critical paths in this graph, this placed graph. So if the total wire length connecting these components is larger, then it's going to tend to be slower. It's going to consume more power, that sort of thing.

• 8:34Daniel Whitenack?

  Gotcha. And how big of a graph? Like, how many things are we needing to, like, place and, and optimize in, in general?

• 8:46Speaker 04

  So it's millions, like millions of standard cells. And then in a chip, there's typically many, many blocks. So hundreds of millions in total of components that you're placing.

• 8:57Daniel Whitenack?

  Gotcha. So Azalea, I'd love to get some context in terms of like, how has this kind of been this problem of figuring out the placement of all of these components of the graph? How has this been approached in the past? And what are the bottlenecks or problems in terms of creating a solution to this?

• 9:18Speaker 05

  Well, there are several approaches to this problem in the past. In fact, since 1960s, research in both academic community and industry started on doing the physical design or placement optimization. There are various approaches. For example, there are quantitative approaches. There are approaches based on greedy methods or simulated annealing or hill climbing approaches, genetic algorithms and such. I would say the way we came in and the way we basically deep learning and reinforcement learning is helping us taking a new step at this problem is that for the first time

• 10:01Speaker 05

  we can like learn the context of the problem and learn from experience. Meaning we think unlike all of the previous approaches, what we are doing is training agents that can accumulate experience and as they're optimizing more chips, they become better at placing new chips. This is an approach that's different from all the previous existing methods.

• 10:28Chris Benson?

  Gotcha. And for those who may not be very familiar with reinforcement learning as a technique, before we kind of dive into how you're using it in this, could you take a moment and kind of give listeners, either one of you, whoever would be, or both of you? Kind of what is reinforcement learning? And why is that in particular, a technique which lends itself? But even if it kind of starting with just the quick run through the fundamentals of what is it if you're not familiar with it?

• 10:55Speaker 04

  So basically, it's a way in normal like machine learning or supervised learning, you're trying to fit labels to input examples. In this case, you have this additional power, I guess. You can take actions in the world and then you receive feedback from your environment. And then you use that information to try to optimize the parameters of your own policy, which is generating these decisions to do better over time. So basically, it's composed of states, which is sort of the state of the world at a given moment in time. So for us, we're placing these chips, one of the nodes of this graph, one at a time onto the chip. So the state is kind of what is the placement so far? And then actions are decisions that you make at each point in time,

• 11:37Speaker 04

  which is for us, like where to place the next node. And then reward is the sort of final key component for reinforcement learning. It's the feedback that we get from our environment. So in our case, after we place all of the nodes, we have approximate signals on wireless congestion and now timing. And we use a weighted average of these to tell our policy how well it did. And so it can update itself and generate better placements over time.

• 12:08Daniel Whitenack?

  So I know a lot of people might have kind of heard about reinforcement learning, maybe with like agents that play Atari games or maybe more so in like robotics. In those types of scenarios, you have this agent, which you know, may be composed of one or more models. And it's trying to take actions like, you know, people like tend to maybe associate that with like taking actions in the video game or like moving the arm of your robot or something like that. In this case, like the quote unquote game you're playing is really the placement of these components. Right? So you're kind of, your agent is placing components and then getting feedback

• 12:52Daniel Whitenack?

  about how well it's placing those components. Is that a good way to put it?

• 12:57Speaker 04

  Yeah.

• 12:58Speaker 05

  Yeah, that's great. Yes, exactly.

• 13:00Daniel Whitenack?

  Okay, great. So in doing this, I'm kind of curious, like, I don't know if anyone's tried to do this before. I assume maybe not in terms of reinforcement learning for this problem. How did you come to decide that reinforcement learning might be a good approach in this scenario versus maybe some other methodologies? How did you kind of come to the point where you say, oh, those things that people are doing in robotics or in these games or something else? How did you come to think that those methods, specifically reinforcement learning, might be suitable here?

• 13:38Speaker 05

  So before we started this project, we have been working on another project, which was doing device placement optimization with reinforcement learning. So that project had to do with taking a computational graph, such as a machine learning, like TensorFlow graph, and mapping it optimally to the hardware devices, such as GPUs or TPUs, such that the runtime or performance of the underlying ML algorithm becomes as fast as possible. So that problem was a combinatorial optimization problem and a very complex task. And started thinking about how ML and this context of our learning

• 14:19Speaker 05

  can help doing that optimization problem better than existing ones and be taught. And reinforcement learning is really like a natural thing to come in mind if we think about ML, because this task is not a supervised task. We don't have labels for it. We want to optimize this problem by doing several rounds of exploration, exploitations. So we did reinforcement learning for that, and we got a lot of interesting, very encouraging results on the device placement task. So then we came to a natural next step for us to try, okay, now, what if we tried the same kind of approaches?

• 14:59Speaker 05

  Or the chip placement problem, which is a much more complex problem than device placement. So that was the transition for us from devices to chips. But the interesting thing was that chip placement, when we came to it, we realized it's a way like orders of magnitude more complex problem than device placement. So it was very unclear to us in the beginning that we're gonna get gains with reinforcement learning for this problem that has been, there's so much research on it already. But after some trial and error and several rounds of improving our algorithms, it seems like it actually is helping a lot in this problem as well.

• 15:47Daniel Whitenack?

  What's up? This is Daniel Whitenack, one of your Practical AI co-hosts, and I hope you're enjoying this episode and staying healthy during these crazy times. I'm working on some pretty cool AI stuff here from my home office, but I've also found that I'm having to get a bit creative and be intentional when it comes to honing my AI skills and virtually connecting with the AI community. If you're in a similar situation or you've been inspired by the practical AI we talk about on this show, I want to invite you to a live online AI training event I'm hosting this May called AI Classroom. In AI Classroom, I'm going to teach you the practical skills I've learned over the years using the latest open source AI technology.

• 16:29Daniel Whitenack?

  You'll learn AI theory along with practical, hands-on implementations in both PyTorch and TensorFlow. And after the training, you'll be able to understand the latest AI models, implement your own models in code, train computer vision and NLP models, create model inference servers, and experiment with state-of-the-art methods like reinforcement learning. AI Classroom is taking place this May. It'll be taking place live and completely online in a high-quality virtual classroom, so no travel is required. There'll also be two cohorts with convenient time zones for Eastern and Western hemispheres. Don't miss out. Tickets and more information are available at datadan.io. That's datadan.io.

• 17:14Daniel Whitenack?

  And Practical AI listeners can use the code PRACTICALAI10 for 10% off. See you online in AI Classroom.

• 17:33Chris Benson?

  So I am curious, you mentioned a moment ago that there was, you know, like the data itself wasn't labeled, lack of labels and you know that it, that reinforcement learning was, seemed like a very good technique to lend. I am curious if you had not gone down this route, or maybe, you know, you know, not machine learning at all, what some of the other options, whether they be in the realm of machine learning or not, might have been, just to have a sense of what kind of the technique opportunity cost would have been. You know, how might others have done it had you not gone this, down this path?

• 18:05Speaker 04

  So we did experiment with some other techniques. They say evolutionary strategies, they tend to be less sample efficient. So it didn't really seem like too promising a path to go farther down. We also experiment with like using supervised learning as a way to basically ground our architecture search. The policy that, uh, architecture that we were able to achieve generalization with was tuned using a supervised learning objective. And then we used that as sort of the encoding stage of our full policy value net and achieve better generalization results.

• 18:42Daniel Whitenack?

  Yeah, I would love to, to follow up on a couple of those things. So like maybe digging into a couple of those pieces, just to break it down for listeners. So when you're talking about this encoding piece and the, the supervised stage that you, you did complete, does that have to do with getting the graph structured data into like, into another form, like a sort of embedding or a representation that you would use in other things? Could you kind of describe that a little bit more?

• 19:14Speaker 04

  Uh, yeah, sure. So I think basically for, in order to achieve generalization, it really, really is about the representation. Like as you said, what is the correct embedding for a given input graph? So basically we created this very large data set of different placements generated by different placement techniques, including reinforcement learning policies, but also like force directed methods, simulated annealing, greedy methods. And we used that to try different architectures on the task of predicting the approximate wire lanes and congestion for those placements.

• 19:51Speaker 04

  And the architectures that were better at this prediction task did a much better job of creating policies that were able to generalize across different chip netlists because they presumably had a better representation.

• 20:06Chris Benson?

  I am curious, you mentioned a little while ago that the thing that inspired you guys to kind of go down this particular path was device placement optimization. I would imagine, and correct me if I'm wrong, I would imagine that this is like a completely different scale in the sense of, you know, working in very, very small spaces. I would imagine compared to the original device placement optimization you were doing. If that's accurate, did the scale, you know, moving down to such small spaces make a difference? Or was it fundamentally the same? Did the approach hold up the same as you had experienced in the prior project? Azalea, do you have any thoughts on that one?

• 20:46Speaker 05

  Yes, so in both projects, we are still like doing reinforcement learning. So the meta approaches still remains the same. But like you said, the scale of these two problems are very different. For example, in device placement, we have like a dozen. Our action space is like tens of devices. But or less, or a few devices, few GPUs, CPUs. But here, our action space is the placement or cells of the canvas onto which we are placing the chip. And this canvas can have thousands or even more of locations. So our action space is orders of magnitude larger than the previous problem.

• 21:30Speaker 05

  At the same time, our input state, which is the graph that we are processing, a chip graph, like Anna mentioned, can have millions of nodes. Whereas a computational graph could have like tens of thousands or so. So here in this problem, we were dealing with a much more complex state and action space. And to enable RL agents that can optimize this problem, we had to do several Changes to the way that we present the environment to the agent. For example, we had to take, kind of take a hierarchical approach to the way we represent the input graph. So, for example, we grouped certain standard cells.

• 22:14Speaker 05

  We break down the complexity of the input state to a graph with like thousands of nodes that we were later on placing. And on the representation learning, we had to do a lot more work because in this problem, not only we were interested in placing one chip, we were also interested in creating agents that become better at placing unseen chips. Because that opens new opportunities for chip design optimization. If we can quickly, given a chip block, can upplace it, optimize it and see how well we are doing in it. So this generalization property that we wanted from this problem led us to really heavily focus on representation learning of the graph.

• 22:58Speaker 05

  And we created a lot of new like techniques for creating these generalized representations that we are hoping in future problems, better in other stacks of chip design or other kind of hard ML for combinator optimization that we are dealing with, can help us do better in those problems as well.

• 23:17Daniel Whitenack?

  So I'm, I'm really curious, I have a follow up from that. And as you were talking, I was thinking about how essentially you have all of these different possible arrangements of the graph onto the, the physical canvas, like you said. But also in this problem, as you're placing components, there is this like sequential nature to it. And maybe this is where, I think it was mentioned earlier that there were kind of even some parallels with natural language processing. And I was wondering how you deal with this situation where you're really not just taking like a one, at least in my understanding, you're not taking a one step approach of like,

• 23:58Daniel Whitenack?

  here's all my components and then here's my prediction for the placement of all of those components. You're kind of placing one component and then placing another and then placing another, kind of in a more iterative sort of way. So is that the, how do you deal with that sequential nature of this process? And does it involve kind of like sub graphs within the graph and then adding a component to that and kind of taking the last so many components and then trying to figure out how the next component comes in, kind of like placing characters or placing words when you're doing text generation? How does that sort of sequential thing come in?

• 24:39Speaker 04

  So the first architecture that we had that worked well, like we would actually pass images of the placement so far. And so the model was kind of like a human designer as they're maybe placing a graph. They could see, you know, what space is left on the canvas and such. And we had a basically an LSTM model for the policy head. Basically, that sort of stores information about the full like sequence of placement decisions that have been made up to that point. But in the end, actually, I think our current policy has a deconvolutional neural net that predicts a policy decision over this two dimensional grid.

• 25:17Chris Benson?

  And I'm kind of curious, and I'm also following up on the same thing, actually. And you may be starting to address that there, but I was kind of curious. You mentioned when you were talking about representation learning of the graph, that there were some new techniques that you got into and you made the comment. And then I was wondering, is there anything else that you kind of learned to apply to this? Or did you just cover it right there?

• 25:39Speaker 04

  So I think Azali was getting at this graph embeddings that were developed for this project. And I think at a high level, the insight there was that for most like sort of graph convolution neural net type applications, it's really about the features of the nodes themselves. And so you kind of represent nodes as some kind of average or other aggregation of their neighbors' features. But in our case, what really matters is the connections between these nodes, because it's about the past. And so our graph embeddings are much more focused on edge features.

• 26:12Daniel Whitenack?

  And kind of diving a little bit more into those embeddings. Again, I'm trying to make connections with maybe things that I've seen or heard about before. I know in like the NLP world with like, you know, these newer language models that are coming out and the word embeddings that they're generating. The thought is like, oh, we're going to train this model or learn this representation based on a one or more tasks, like, you know, replacing missing words or something like that. And then you learn this embedding and then kind of apply maybe some new layers onto the network to do a particular task, like question answering or whatever it is.

• 26:54Daniel Whitenack?

  Here, is it similar in that you were talking about how you use some supervised learning to train the train the embeddings, in my understanding. So you have these certain tasks that are supervised and you learn the graph embeddings and then you were able to apply those in a new scenario? Is that the strategy or do I have that wrong?

• 27:21Speaker 04

  Yeah, that's very much right.

• 27:22Speaker 05

  Yes, that's correct. So I think the way we can describe this is that we trained architectures to capture the representation or encode the embedding of the input by having a supervised model that with very easy to produce labels. We call them pseudo labels, right? Those labels were our proxy costs for the optimization that were very fast and not at all expensive to generate. So the motivation for us to train the architecture this way was if our agent, our policy is to generalize to unseen graphs, it should also have a

• 28:04Speaker 05

  Good understanding of predicting what the actual reward is for a given state. Like unless it can, so it's, but that's like a prerequisite for generalizing policies to answering graph is to have an idea how good a current state is. And that's what made us do the supervised approach first, where we predict these pseudo labels for a given graph. And once the architecture is tuned in a way that this prediction task is done at a high accuracy for the test set, then we take that and use that as the encoder part of the policy for further optimization of placement.

• 28:47Daniel Whitenack?

  Awesome. This is really interesting because we've brought up graph neural networks a couple of times on the show, but maybe not in the sort of applied way that we're talking about them here. I was wondering if you could just, before we get too much further, just mention like, what makes like a graph neural network, a graph neural network instead of just a like a normal neural network, I guess. And maybe like help clarify for people, because even in this episode in our conversation, we mentioned like computational graph, which people might, that might come to people's mind if they're thinking about TensorFlow, there's like this computational graph in the background.

• 29:28Daniel Whitenack?

  But here, like for a graph neural network, we're not talking necessarily about the computational graph. What makes the graph neural network a graph neural network? Is it just the input data and this sort of way of representing graph data?

• 29:44Speaker 05

  So what makes graph neural nets, graph neural nets is what the way they encode information. So they, in a typical graph neural net, we are learning representations of the nodes of a graph with respect to the properties of this node and the properties of its neighbor nodes and the neighbor of the neighbors and so on. So graph neural nets have this property that they can encode information about the one hop, two hop, like K hop adjacency information of a node. And you can also, on top of this adjacency information, like the connectivity graph,

• 30:24Speaker 05

  you can also add features per node and you can also, in our case, you can add features per. Edge of the graph. So basically, graph neurons are allowing us to capture all of this information about the graph structure of an input data and generate embeddings of the nodes and edges that kind of relate and can capture those graph structure and graph information.

• 30:53Chris Benson?

  So, you know, having gone through this, which is fascinating, it's entirely new to me, I'm curious what the results were like and as, you know, kind of where did you arrive? What surprised you along the way in the process? You know, what was not what you were expecting to see? And also, how did the larger organization at Google take the results? Is it something that is now becoming kind of standard at Google? Or was it just a test or an experiment? Or how did it affect the larger organization in terms of designing chips going forward?

• 31:25Speaker 05

  So we have definitely tested this method on chips that Google makes and have gotten superhuman results on a good portion of the, like, the complex chips that we tried placing them. But in terms of other questions you ask, I'm not sure if we can answer that at this point.

• 31:49Chris Benson?

  Okay. But nothing jumped out from a surprise standpoint, just kind of like, you got something? I was curious.

• 31:55Speaker 04

  Yeah, sure. I have something to offer on that. I don't know, surprising, maybe just exciting. In terms of those generalization results, we would say, take a policy and pre-train it on a larger number of chip netlists and then, you know, apply it to a new chip. So it's sort of surprised and excited us was that a pre-trained policy that was fine-tuned for, say, only 12 hours would outperform a policy that was trained from scratch on this netlist for 24 hours or more. So I think it was exciting to us that this new policy architecture generalized so well that it actually does better and it takes less time.

• 32:34Chris Benson?

  That's pretty amazing.

• 32:35Daniel Whitenack?

  Yeah, was that having to do, I know when I was looking at the paper, you talked about like domain adaptation, which I, I remember we talked about with the OpenAI team. And also we've talked about in relation to like robotics and moving hands. So is that key to that sort of generalizability is adapting the domain or the, the environment during this training? If so, did you have to like, create a bunch of simulated data for various environment changes and that sort of thing? What was your approach there?

• 33:14Speaker 04

  So we actually just used a real chip netlists for all of the pre-training. But so we'd say trained on 20 real chip netlists and then we were able to achieve those results where we have much better and faster results, but we probably could do some kind of data augmentation. Where we could maybe turn those 20 into many more, or you know, source more netlists in some other way and you would do much better.

• 33:39Daniel Whitenack?

  And, and what is your feeling in terms of, you know, how specific this pre-trained policy is for the sorts of chips that are included in the set of chip nets, I think you called them, that, that you use during pre-training? How specific do you think the, the pre-trained version of the policy is for, for that kind of family of, of chips, or do you think it's generalizable beyond that?

• 34:06Speaker 04

  Um, I think it, it's definitely affected. The, the policy's performance on a new netlist is definitely affected by the types of netlists that it's trained on in the past. At the same time, it's a pretty general problem. So yeah, I think as long as you trained on a representative set of netlists, you could do well on a new one.

• 34:26Daniel Whitenack?

  Gotcha. And what are some of the challenges that maybe you faced during this project that you maybe didn't have time to address in the initial version of this project? What are some of the things that you want to explore more going forward?

• 34:40Speaker 04

  I mean, there's just so many other stages of this process. And kind of what's exciting about, you know, developing policies that can more quickly generate high-quality placements is that we can kind of explore feedback or interactions between, say, previous stages, like upstream choices, like this choice of SRAM. Basically, there's a certain amount of memory that needs to be in this chip, but the choice of how to slice it up into these macros is somewhat arbitrary. And if you can say, try one, like slice up the macros a particular way and then see what kind of placement, what level of quality you can get in terms of timing and other properties from that quickly, you could do all sorts of explorations upstream.

• 35:23Chris Benson?

  So I want to follow up on something that you were saying before. And just to make sure I understand, when you're looking at these different types of chips that you want to apply this to, and going back, I know we had someone talk about some chips from a previous company earlier, but, you know, they were talking about basically different types, you know, from GPU, TPU, FPGAs and such as that. Do those different architectures dramatically change the problem for you? I know that we were talking about the domain adaptation a moment ago, but, I mean, in a practical sense, do you have a substantially different RL approach every time you change

• 36:04Chris Benson?

  out the chips, given that like, I believe that like a GPU will have a whole bunch of things beyond what a TPU might have on it, you know, because it's being able to address problems, whereas a TPU is very specific to the matrix multiplication. How does that affect your approach on that? I personally wasn't clear enough on it because of trying to learn this as we go.

• 36:29Speaker 05

  Yes, so we have tried our method on a bunch of different types of chips that were available inside of Google and also chips that were available open source. And the way we did our RL operators didn't need to change going from one set of chips to the other, but definitely like Anna was mentioning, you, the larger, if you have a chip that is drastically different from anything you have seen before, then it could affect the performance of the agent. But at the same time, the input space of our problem is very abstract. We don't deal with the specifics of a chip. Rather, we are dealing with a generic, like netlist representation of the chip with these

• 37:14Speaker 05

  nodes that have certain connectivities and the node have different sizes and different shapes and we are placing them, optimizing for the cost that we have developed. So the problem, if we don't think about what chips they are, is very abstract in a sense that it can really can handle different sorts of input from different chips. And so far, we didn't have a chip that was drastically different from our training set that we had to change the RL algorithm for. There's always a modification of the algorithm for improving it overall, but like I said, the input state is pretty standard among different types of chips.

• 37:57Daniel Whitenack?

  So I'm, I'm kind of curious, the more I, I think about this problem, it seems kind of like we're using an AI method to help design a chip on which AI will hopefully operate or, or be trained or run inference on or, or, or that sort of thing. I'm kind of curious on a more general sense, how you see, you know, AI as, as we move into the future and AI development continues to accelerate. Are we going to kind of need these sorts of methods more and more because more specialized chips are going to be needed for these types of AI problems moving forward?

• 38:40Daniel Whitenack?

  How do you see kind of AI influencing the hardware that AI runs on, I guess, is my question?

• 38:49Speaker 05

  So a chip design is a really complicated task. And making customized chips is definitely also very complicated. We are witnessing the more that we are going to need more and more of these customized chips because of various computational demands of especially AI algorithms. And our vision is AI can help the design of these chips because of its ability to learn and improve over time and its ability to kind of optimize over a very large optimization space. For example, if you look into the chip design process,

• 39:31Speaker 05

  there are various stages of optimization from architecture design to logic design to verification and physical design and placement. Each of these stages are very complicated, are combinatorially hard. And so our goal or our vision is AI can help us finding globally optimized solutions across all of these stages. Then we are going to have hopefully a lot more performance improvement over what we have right now where we optimize each stage separately and then just cascade them together. And the reason we think AI can help with it, we mentioned this a couple of times in this conversation was that AI can improve over time. And this property is something that's very different

• 40:16Speaker 05

  from what we have seen in any other existing methods. So the policy, the agents can become better, more experienced at doing newer tasks. So if we accumulate this experience over time, then we are having, we are dealing with these agents that become much better than any like single person or single algorithm that has ever optimized a chip.

• 40:43Chris Benson?

  Yeah, in that way, it almost seems like so many other areas that we're applying AI techniques to and that you take it to that superhuman level and just continue on. It makes me wonder, as you guys, and I don't know, this may be almost an organizational question to some degree, but I'm curious whether having pioneered this, you know, being able to apply reinforcement learning to this particular problem, is this something that the two of you are expecting to continue working on for some time? Or have you kind of done your experiment and, you know, you got your results and you're going to move on to other problems? If it's the latter, what might those other problems be? Or if it's staying on this, what are you looking to next? What's the next step, whether it be on this problem or doing something else for each of you?

• 41:28Chris Benson?

  Anna, you want to go first?

• 41:29Speaker 04

  I think that there are definitely other stages of the chip design process that, you know, have a lot of impact. I think getting to your last question a little bit in terms of like, how can this affect kind of AI for AI chips? The current chip designing process takes like nearly two years. And so there could be certain types of machine learning architectures that just aren't computationally feasible on today's hardware. But if we could more quickly design chips for them, they might become more viable approaches. But the problem is that, say, chip floor planning is just one of these stages. So if we wanted to say dramatically accelerate this process, we would have to tackle these other stages. Say like architectural exploration or design verification.

• 42:12Chris Benson?

  Awesome, and you've built one of the building blocks of that process, but you could be exploring some of those other building blocks as well, is that right?

• 42:21Speaker 04

  That's right.

• 42:23Chris Benson?

  Awesome. What about you, Azalea?

• 42:25Speaker 05

  Yeah, I think I'm in a similar boat. I think I'm very excited about the research on RL and ML for optimization tasks in general. And I think chip design is a very critical and important application of optimization. Something that's gonna enable, like if you have better chips, we're gonna have better next generation AI algorithms as well, because chips are key enablers of those algorithms. So I would say both research on RL for optimization and with applications in chip designs is something I'm very excited about and look forward to continue working on.

• 43:09Daniel Whitenack?

  Awesome. Well, thank you both for taking time to join us. This has been super fascinating, and it's been great to dive into some of these subjects like graph neural networks and chip design and these things that we haven't talked a lot about on the show. So I really appreciate both of you taking time and joining us for the conversation. It was great to talk.

• 43:30Speaker 05

  Thanks for having us. Thank you so much for having us.

• 43:36Speaker 01

  Thank you for listening to this episode of Practical AI. More like this at changelog.com/practicalai. There you'll find our latest as well as lists of our most popular episodes and the ones we recommend. If this show has helped you on your AI journey, please leave us a five-star review on Apple Podcasts, heart us on Spotify, star us on Overcast, and tell a friend what they're missing out on. Practical AI is hosted by Daniel Whiteneck and Chris Benson. It's produced by me, Jared Santo, and our music is brought to you by The Beat Freak, Breakmaster Cylinder. We have awesome sponsors. Please support them. They support us. Thanks again to Fastly, Linode, and Rollbar. If you and your organization could benefit from speaking directly to all the AI practitioners out there, you should sponsor the show. Podcast

• 44:20Speaker 01

  advertising is one of the most effective ways to spread your message in an authentic way. Plus, you get the added bonus of supporting something you love. That's all for now. We'll talk to you next time.