Lightning Talk: Torch-Spyre: Compiling To a Multi-core Dataflow Accelerator... D. Grove & O. Tardieu

This talk covers the insights gathered by IBM's team regarding their implementation of PyTorch Inductor as part of the software stack for the IBM zEnterprise accelerator. The speaker discusses Torch-M Spire, an open-source PyTorch device geared towards enhancing the developer experience and enabling AI computations on IBM's hardware. Key features of the zEnterprise accelerator are highlighted, including its 32 compute cores optimized for AI processing and the use of specialized compute units. The session delves into the significance of memory layout and tensor tiling in enhancing performance and efficiency during operations such as matrix multiplication, and how these concepts are integrated into their compiler framework to optimize PyTorch for their hardware. Finally, the speaker addresses ongoing contributions to the PyTorch ecosystem and plans to maintain compatibility for future enhancements.

I agree, so thank you for joining us here this afternoon. I'm Dave Grove from IBM and I'm joined by my colleague Olivia Choudhury and we're excited to share with you some of the insights that our team has gathered gathered over the last year while adopting PyTorch Inductor as a key part of our software stack for IBM zEnterprise accelerator. Now our Inductor work is part of the a

larger project we call Torch-M Spire. So Torch-M Spire is an open source out-of-tree PyTorch device for IBM Uh it's built using the architecture recommended by Torch-M Spire and Open Neural Network Um and our goal with this is to really enable the entire PyTorch ecosystem and developer experience for the IBM zEnterprise accelerator and to bring that capability to our customers. And we're doing this by leveraging important open

source projects like PyTorch and LLVM. We're also trying to contribute back to those projects to give them make them more amenable to out-of-tree accelerators like ours and hopefully other people can do that. Um it's a large team effort. Uh Torch-M Spire is on GitHub. We have 50 contributors already and Torch-M Spire is a piece of the overall software stack that's being developed by IBM for by IBM

for this accelerator. Um so what is the IBM zEnterprise accelerator in a nutshell? So it's a system on a chip. Um it has 32 compute cores. Each of those compute cores has a private minus scratchpad memory. It's really optimized for AI models and the kind of computations you do in AI. So that the the cores have different precision floating points, there's specialized compute units and there's non-chip

interconnect that links them all together. To drill into a tiny bit more detail on this, uh we'll first look at the side inside the compute units. Um each one of those compute cores has a number of floating point units. Um the heart of it comes from the heart of it is a deep by eight systolic array which is where the major multiply happens. Um and every compute

unit each one of the small boxes in there is 16-bit SIMD engine for its different levels of precision floating point. A second key piece of the chip to get performance is making [clears throat] good use of the on-chip interconnect. In particular, the interconnect and the memory subsystem are optimized for multicast loads and for doing bulk loads contiguous data from the the off-chip So, let's see what that

means for matrix multiply. Um so, the compute width of the units for the unit matrix multiply we want to use the systolic array. Um the compute unit width here is 128 bytes, so there's 16 bytes in the for each in the units in a row, so that gives us 128 bytes. And so, to utilize the systolic array, for example for FP16 values we're really doing a matrix

multiply on 64 by 64 tiles. All right, so let's see what happens. So, if you have the full of first two tiles out of our matrix multiply, um there'll be 64 rows in each of these, each having 64 floating point values. And then as the as the we proceed to go through the matrix multiply, those rows will flow through the systolic array in the north-south and east-west

direction, and the various multiply accumulates will happen to actually build up the matrix multiply for us for this particular piece of the All right, so the second thing is we really want to be able to exploit the memory system subsystem. So, if we look at a small bit set of tensors here, [clears throat] we have our 64 by 64 tiles. Um the default uh you know, the

tensors have to be linearized in memory somehow. By default, PyTorch will do a row-major layout. So, it'll take our the first first row of our tile, first row of our tensor which has three tiles in it, the red, the orange, and blue, lay them out contiguously in memory. Um and they'll look something like that which is on the right. Um unfortunately now it means we need we

want to load the red tile, we need 64 separate loads. We're not able to use our bulk memory loading to do this. Um so, what we want to do instead is do something we call tile tensors on the device. So, we're going to go through the >> [clears throat] >> uh All right, let me go back for a second. There we go. All right, we go through

the personal layout of the red chips, the the rows in the red tile, uh down one after the other, and then the green tile, and so on. And by tiling the tensors in this way, interspersing the the rows and columns, breaking them up in the way that PyTorch does, we now have our nice all of our tiles are laid out exactly the way we need the tile

to be memory. For load and store it's all stride. It's all stride and compute efficient. Now I'm going to turn it over to Olivia to explain how this is actually implemented and what what it does. Okay. Thank you, Tate. So so Tate has explained why this is really important that we have tensors. So let's uh look at how we use these in our PyTorch uh PyTorch device.

So if I copy over let's say 3D tensor to the host and I load it from the device, I instead I want to copy back. So I took the 200 elements in the device as a dimension. I divided this into a multiple of 64 because it can pass our single width. And then I divided this into 4 * 64 and then I finally I essentially bring the

dimension to 5 elements in between. Right? So that's that's what we have to do. But of course we want to do this uh automatically and basically you know once we have a chance to the PyTorch code, right? So the first thing we have to do is to make sure we're going to keep the correspondence or the relationship between the host layout and the program it's written for

and the actual device layout that the program is going to be PyTorch. So the way we do that is the way uh PyTorch in general thinks about layouts, right? So uh PyTorch thinks about layouts in 3D structure. So for instance, you know that if we take [clears throat] the host coordinates and combine this with the host stride, we can get the host offsets of any dimension times

stride which is essentially displacement from the base address of the tensor in memory. So we can do the same with a device stride. But when we introduce here the third kind of stride which we call stride map which is a trick if you want, we say if I take the device coordinates and combine this with the stride map, I get actually the host stride. And that that

that thing that state that's that's stride map is what is going to keep the relationship between the host layout and the device layout and this >> [snorts] >> relationship between the two and that makes it possible for us to combine instructions. So, let's look maybe at one example of that. Here I just want to do something very simple. I want to add a tensor. We have our

tensors and we can point what adds. My second tensor is a 2D tensor and just to make this a little bit more interesting, my first tensor is actually the 3D tensor I have on the previous slide which I've using for the view which makes it look like the 3D tensor, right? Because I factor in the first dimension, the first dimension of the final tensor together which is

50. And my added 3D dimension. So, this is what it looks like on the host. But again, we said we're tiling tensors and that's why what we made on the device on the device, right? The first tensor we have explained by how we tile it and the second tensor was 2D tensor we just started so we probably like this. If we had the foresight, maybe we would

have made different choices in the way we used. But either we can have the foresight or we have good reasons to tile them differently and we have to do And what that is basically the trick because if you think about how PyTorch thinks or Inductor thinks about this operation on the host side, it says that this is what represents here. It says that this is just a

2D operation. This is essentially just the loop one of iteration by loop zero, we go from zero to 50. While from iteration by loop one, we go from zero to 200. And really the crux of this is that this this in these two loops I'm going to take one element of the first tensor and move another element of the second tensor to the sum and store it

somewhere in the output tensor. The key question which element do I get from the first tensor, which element do I get from the second, where do I write the result? And on the host side it's very simple. It's actually the same expression, the same index expression for everything. But now we've organized our tensors differently in memory. So, in on the device, it's not that simple. We need

actually to have different ways of addressing of you know, the access pattern for these tensors will be different for all two tensors. So, what is to go and take this nice expression here and convert it into something like a device index or coordinates on the device, but we have already explained in the previous slide that we use primat and that makes it possible to do that. So,

primat is a mechanism to map um device coordinates to host offset and then we use it in reverse to map host offset to device coordinates and so forth. In fact, more interesting than that is that we can actually do this symbolically. All right, so what I've done here is that I realized that this zero, which is maybe the coordinates on the second tensor of this of this

dimension, actually when I look at the processor, we need to use both axis of C and the fine dimension by using it from All right, so problem solved? Not entirely because our essentially code generator expects simple to be very very simple and this primat is more complex than what we can manage. So, what we really want to get to is some description of an operation that only

uses nice and simple expression like this. And the way we do this is to go out and find a 10 by 10 or 5 by 5 splits. That's something we can deal with. Uh but this 15 on the other hand we can't logically split into 5 and 10. And if we do that, now we have two 40 tensors, which are essentially the same tensor with dimension dimension

reversed, and we can have a nice and simple expression for this. So, I have 1 minute left. So, we talked a bit about tiling. So, one message here is that tiling is really interesting, it's really important. It's a bit of a headache, but it's fun to deal with, but that's not really the message for this talk here. But the message is over the course of the last

year, we've been rebuilding our compiler, our in-house compiler for uh for spire using inductor. And what that means is we now have the picture on the left. We have inductor based on compiler, and in that compiler we inject device specific knowledge, device specific optimizations and these things like this, which are the protocol for And uh what we want We can [clears throat] only say that we're very

uh happy with this friends because we think that it's actually working on a long run yes, but it's working we think that we're going to get all the way with this architecture getting all the performance and the support we need from the virtual and as we do that we we think we massively simplify compilers because rather than building building everything from scratch we build everything as extensions

of the compiler framework that is upstream and maintained by the community and it's strong and really working and updated for every version. That's what we think. We actually plan to merge using your tasks and maintain this compiler framework for you. Thank you. [clears throat] That's it for us for today unless you HAVE ANY QUESTIONS. SO I WAS LATE FOR YOUR TALK BUT DO YOU find that whatever

like you know back end registration is having in PyTorch case? All of I mean one of the things we're doing here is we're So we spend a lot of time discussing and we think that we think it works well there are few places where we need to hack we need to dig too much in PyTorch as in it doesn't have the configurability or doesn't expose the configurations

that we want to have. We also find some limitations when we want to say let's run let's say run compile for CPUs and a GPU or CPU and or device as part of the same run of PyTorch. That's also something that's not entirely supported as well as we should support it so we saw something that we're trying to contribute upstream to PyTorch. Yeah maybe this is easier

for us to maintain than any of those. Okay. Very good. >> Nothing else? Thank you very much.