Accelerating Complex-Valued Tensors With Torch.compile - Hameer Abbasi, OpenTeams Inc.

This talk addresses the issue of compiling complex tensors in PyTorch, which currently cannot be compiled due to limitations within the Inductor and Dynamo components. The speaker explains the performance advantages that JAX and TensorFlow have in this area and discusses their approach to solving this problem through the creation of a subclass specifically for complex numbers. They detail the mechanics of storing the real and imaginary parts of tensors and the intricacies of ensuring performance while maintaining compatibility with view semantics. The session also touches on the challenges of supporting standard operations like FFT and the considerations around binary size when adding support for various complex number formats. Overall, the speaker emphasizes ongoing efforts to enhance PyTorch's capabilities in handling complex tensor operations.

Hi, good morning everyone. Nice to see you all. Um, right. So, I have some good news and some bad news. A, um, you will notice that I don't match my profile picture because of the head. positively hairless. Enjoy it while you can. I'm all for a good joke. Yeah, good news is it's not cancer. So, uh, bleep cancer anyway, but It'll grow back. Okay, so, um, my

talk is about about me. I I love my boo. You know all don't know who that is, my partner. Um, I love my two cats that I have at home. I love cooking. I love hiking and like many of us here wasting a lot of my free time coding. uh, I'd like for this presentation to be interactive. So, if you have any question, anything to say, just

stick your hand up. I'll point it to you and you can you can say what you want to. Okay? Um, right. Oh, wait. Sorry about that. What the heck did it go? It's not on this. Ah, there. Right. um, the problem uh, statement that we're or I'm trying to solve with uh, with PyTorch. Uh, complex tensors cannot be compiled. I mean like a tensor with a complex

data type. Right? Uh, those cannot be compiled at the moment. If you throw those into a torch.compile block, it will graph break at the moment uh, uh, because while Dynamo knows how to trace them, Inductor doesn't know how compile them, treat deal with them. And this is one of the places where JAX and TensorFlow are better in performance simply because they do know how to compile uh,

tensors with a complex data type. We're trying to solve that and some application examples are, you know, rotary embeddings that you can find in many common LLMs or or, you know, world models or anything physics-y that you might want to do with your PyTorch installation, right? So, you obviously don't want to have to either, A, have your code be slow or really, really write your code in

a way that that you have to, you know, really use real tensors all the way. That That would be a nightmare, I think, in terms of code quality. So, um, the solution that the first step of the solution that we have is a subclass for If you've had some experience with torch.compile, you know that you can write subclasses. Alban Desmaison from the meta team has a subclass

zoo, uh, a repository that details examples of how to write different subclasses. I think there's even a complex subclass in there. But that one is for So, good news, there's one in-tree now inside PyTorch and you can use it actually today. Um, assuming you're on PyTorch nightly to compile your complex code. Um, please try it out. If you run into any like limitations or anything, uh, any

rough edges, feel free to reach out. Um, my GitHub handle's at the end of talk. You'll have it soon. Right. Um, the the way we do this is we store the real and the imaginary part of the tensor and we just use decomposition to produce another subclass tensor. Or, if it's a real tensor, then we produce like a regular PyTorch tensor and not a subclass. Right. That's

part A. Then, the second part is, uh, there is this open PR to the PyTorch uh, where I try to >> [clears throat] >> work with Inductor. Sorry, my throat is a bit bad because I have a cough and a flu. Uh, where I try to, uh, teach Inductor to kind of like recognize where a complex data type is getting passed in and try to, you know,

thing, uh, to the subclass. And the same on the way out, right? Um, that's pretty simple to do. It's a pretty small PR. Right. Then, there's a number of, I don't know, caveats or decision points that are still open. Uh, the the main one is that well, you want this thing to be performant, right? I mean, torch.compile is all about performance. >> if you want to do

matmuls to interleave format, the default way that complex numbers are stored, you cannot do that in a performant on GPUs because the matmul units on a GPU do not need know how to deal with complex numbers or complex matrix multiplications. uh, while we do deal with it in the subclass in a way that that you actually, you know, can uh, real matmuls, the bad thing about that

is is you still get a stride of two on on every single element because, you know, you're still using the interleave format. so, you know, you either have to make it contiguous and then make it interleave again or it's a whole mess. So, the way we're considering dealing with that right now is is to make it like non non-contiguous all the way through. So, we just have

two separate tensors for the real and the imaginary part. And then, when we're when it actually comes time to, you know, go into the compile block, come out of the compile block, then we reflect to the interleave format so that, you know, you don't have to deal with the But there's a small problem with that, right? If you want view semantics to work, uh, if there's any

view relationship between the incoming uh, just outgoing tensors that or any incoming and any outgoing torch tensors, then you so, for example, you're you imagine the following code, right? You pass in a simple complex torch tensor. You mutate and then you return it, right? This and an input have a view relationship. And if we're we're going an interleave a non-interleave format storing the uh, the real and

imaginary parts separately and then coming back to uh, an interleave format, you know that that So, this is a correctness issue. It needs to be fixed. Uh, there are a couple of options being discussed. The first one is just disallow views on input tensors, right? Or at least disallow mutations on input tensors. The second thing you can do is detect when there mutations that need to be

preserved and actually stick with the with an interleave format for those. Have a separate subclass that actually works with the interleave format in a way that, you know, you have this uh, extra two dimension that that can actually work at the same memory layout. and when that view semantics is not needed, then fall back to the separate real and imaginary parts so that you have the nice

matmul performance that we all know and love from our GPUs. Um, there are number of other caveats so that there are a number of ops which are like very hard to decompose for complex numbers. Examples are FFT. Well, you can decompose it in terms of the discrete discrete cosine transform, cosine transform, but even those, you know, they're they're present in in many like CUBLAS or CUFFT libraries

or something like that. So, so you have end up having to do the entire thing anyway. Um, for for preserving performance on, you know, operations like these, there is no way out. we're doing before with the subclass approach was that Inductor backends like like anything that Inductor emits that takes in an Inductor graph and uh, you know, emits code for it, that doesn't need to know anything

about complex numbers. It can just quote unquote work. that for this particular problem, it is uh, to actually have that work because, you know, discrete sine and cosine transforms, the vanilla forms anyway, are not available in many libraries. unfortunately for this one for those of you that Inductor backend code, um, I'm sorry. For this one op, we cannot make it so you don't have to modify your

backend to work with complex numbers. What that means is that you will need to support the FFT ops in your Inductor backends or those that write Inductor code anyway. We can actually get a number of small bonuses out of this. So, right now uh the the torch team is like balking on in introducing support for a lot of different complex number formats or a lot of different

d-types in general. they blow up the binary size, and that's a valid concern because uh you know, if you talk about torch CUDA, it's just I think about 2 gigs right now on on most CUDA platforms. And if you talk about ROCm, it's close to like 4 gigs. Um if you're compiling for every single uh GPU But, they don't want to blow up the binary size even

more for for good reason. So, So, you know, we have to get around that somehow. uh I do have this PR up that adds torch.complex32 support, which is like imaginary parts of complex 16s. The complex bfloat 16s. Pardon me. Now comes time for the bonus. The bonus we're thinking about of we're thinking about a scheme in which, you know, each time the first time you call just-in-time

compiled instead of, you know, being baked [snorts] into the binary. This makes it so that, you know, the binary size doesn't blow up when you download it, but you also get, you know, a nice uh you also get a not do not get a not implemented error when, you know, you try to, for example, Um this is uh so, currently, this is just talks. It's It's not

something that's being implemented uh inside this PR, but it is something that we're planning on doing. Uh I know this talk has been short. I expected the reviewer bandwidth. There has been limited progress. Uh but I hope to have some soon. Um thanks for your ear. If you have any any questions, you know, I'm happy to take them. Otherwise, I'll take your ear your with me like

Hannibal or something. Home. Okay? >> [applause] >> Right. Any questions? Otherwise, I'll just yield to the next presenter. Going once. Going twice. Thanks for your time, everybody.