If AI Writes All the Code, What Do the Programmers Do? | Probably Dance

Probably Dance

I can program and like games

If AI Writes All the Code, What Do the Programmers Do?<br>by Malte Skarupke

Eight months ago I was producing roughly 90% human written code and 10% AI written code. This has switched surprisingly rapidly and now my code is probably 90% AI. So what do I do all day?

I’ll go through a change that I made to a matrix-multiply kernel, for which I’ll sadly have to be a bit vague. Since GPUs are now giant matrix-multiply chips, where north of 96% of the flops are in the tensor cores (latest Nvidia GPUs have 2250 tflops in bfloat16 matmuls, compared to 75 tflops for everything that’s not a matmul), you’d think that they’d make it easy to use all those flops. But no, matrix multiply kernels are giant crazy beasts that are incredibly tricky to get right. Some quick googling finds this explanation on Nvidia hardware and this one on AMD hardware. Just open those two and scroll down both to get a feeling for how much work is involved. (and yes, these implement the simple O(n^3) matmul loop where you iterate along repeatedly and multiply and add all the numbers)

The particular kernel I was optimizing was using Cluster Launch Control (CLC), a complexity that’s covered in neither of the blog posts linked above, and it had some bad interactions with some particular inputs. The resulting change was 95% written by AI. So I’ll just go through what I did. AI asks are bolded:

0. Find and Understand the Problem

Before I started I had to find and understand the problem. Some coworkers had talked about the matmuls taking too long, and nsys showed that the tensor-cores were oddly idle while that matmul kernel was running, and then it took some more targeted benchmarking to confirm that slightly different inputs give big speedups, which convinced me that something silly must be going on. My coworkers actually had a really good theory, which brings me to the actual start of the work on this task:

1. Find the Code

The code was in a library that I was vaguely familiar with, but since matmul kernels are giant scary beasts, it was a bit hard to navigate. So before I even tried, I asked an AI to find the relevant code for me . I explain what the problem is and I explain our theory. Then I also start looking myself but the AI finds the relevant code first and even points me at exactly the right lines to look at. It also confirms that our theory for the problem sounds plausible.

2. Understand the Code

Next I decide to understand the code myself, so I intentionally don’t ask the AI anything more. As I look around I begin to think that our theory actually isn’t right. I mean it was partially right, but the code very much wasn’t doing the slow thing that we thought it was doing. It was doing something slower: returning out of CLC mode and giving control back to the hardware scheduler.

This is a little silly, so after I am convinced of it, I ask the AI to confirm , just to get a second opinion. It agrees with me (but I take that with a grain of salt, because it also agreed with the initial theory).

3. Monkey-patch the Code

Since neural network training happens in Python, you fix things by monkey-patching. At least initially this is the fastest way to iterate on this code. But monkey-patching has a bit of boiler-plate that’s easy to get wrong, so I ask the AI to set it up for me .

Then I try to fix the bad behavior, but my fix results in a deadlock. I look at the surrounding code again, but realize I don’t understand CLC enough (this is my first interaction with it, and I jumped straight into a complicated kernel where it’s part of other pipelining), so I could either spend the time to understand the surrounding code better, or I could just ask the AI to have a look.

4. Fix the Code

I ask the AI how it would fix the issue . It immediately tells me why my fix didn’t work: I was violating some invariant in the code where the pipelining logic was also used to reason about what state the CLC is in, and my change required keeping the CLC state separately. The AI helpfully tells me that there is one unused slot of shared memory that is reserved for unknown reasons and is entirely unused by the kernel, so it suggests just using that memory to pull out the state that now has to be tracked separately.

Understanding this would have taken me hours, maybe even a full day, so I tell the AI to just go ahead and implement the fix that it has in mind. The fix immediately works and makes the code much faster for the problematic inputs.

5. Reviewing the Code

At this point I spend some time to understand the code and the change better. It looks reasonable to me, but I also ask a second AI to review the work of the first AI. The second AI finds a problem: The matmul kernel was clobbering over some other state that I hadn’t fully reasoned through, and this was not a problem when the thread-block was shutting down and giving control...