mlir-npcomp intersects with torch-xla

[byronyi]

Background

ATen device capture is an attempt to produce MLIR from tracing PyTorch program running on CPU. It largely modeled after torch/xla in its very first version contributed by @stephenneuendorffer from Xilinx.

Old API (pseudo-device, similar to torch/xla):

import npcomp.frontends.pytorch as torch_mlir

dev = torch_mlir.mlir_device()
t0 = torch.randn((4,4), device=dev)
t1 = torch.randn((4,4)).to(dev)
t2 = t0 + t1
t2_mlir = torch_mlir.get_mlir( t2 )
t2_cpu = t2.to('cpu')

The device API is still recommended in the official tutorial for adding backend support.

@stellaraccident refactored (and greatly reduce the LOC from ~10k to <1k) the code to utilize the c10 dispatcher. Stella also changed the user-facing API from pseudo device to a Python context associated with local dispatch key.

New API:

import torch
import torch_mlir

lhs = torch.rand(2, 3)
rhs = torch.rand(3, 4)

mb = torch_mlir.ModuleBuilder()
with mb.capture_function("mm", [lhs, rhs]) as f:
  result = torch.mm(lhs, rhs)
  f.returns([result])

mb.module.operation.print()

Output:

module  {
  func @mm(%arg0: !numpy.ndarray<[2,3]:f32>, %arg1: !numpy.ndarray<[3,4]:f32>) -> !numpy.ndarray<[2,4]:f32> {
    %0 = torch.kernel_call "aten::mm" %arg0, %arg1 : (!numpy.ndarray<[2,3]:f32>, !numpy.ndarray<[3,4]:f32>) -> !numpy.ndarray<[2,4]:f32> {sigArgTypes = ["Tensor", "Tensor"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
    return %0 : !numpy.ndarray<[2,4]:f32>
  }
}

More examples could be found here: https://github.com/llvm/mlir-npcomp/tree/main/frontends/pytorch/test/acap_export

Under the hood of the new API, c10 dispatcher picks up a local dispatch key associated with mb.capture_function context manager, and a backend fallback function picks up the boxed kernel call on torch::jit::stack and produces generic MLIR (1:1 mapping to boxed kernel call). It then re-dispatch to CPU backend to get the shape and dtype of ATen tensors.

Despite that, the c10 dispatcher has some caveats and we suspect that the it is not really how the mode-based backend fallback mechanism is supposed to be used, notably on convolution, copy_ and factory functions like arange. For us it also record the shape and dtype too early (more on that later to support dynamic shape) but this could be altered.

Design goal

The reason we look into MLIR is two-fold. XLA HLO IR has some nice properties, but it does not support dynamic shape (except for padding) and MHLO will probably fix that in the future. In the meanwhile we would like to plug in backend specific intrinsics for some of the custom ops (notably torchvision ops, or ctc loss in seq2seq models) when migrating existing PyTorch users, and we are looking into better alternatives to XLA custom call.

Xilinx and some of the custom training ASIC vendors (AFAIK) are also moving towards MLIR for easier interoperability between frameworks and their software/hardware stack. They are satisfied to export MLIR some where in the front end stack.

Proposed changes

Option 1

Export HLO graph from torch/xla and translate into MHLO. Not really what we look into for adding dynamic shape and custom ops support as it does not enhance the expressiveness of current torch/xla frontend.

Option 2

Adding back the pseudo-device API in mlir-npcomp/acap to probably workaround the caveats. It would be similar to the wrapper-based backend fallback in contrast to mode-based backend fallback. It retains the CLOC advantage while switching to (IMHO) canonical extension points for backends. We also would like to re-visit the shape and dtype static inference vs runtime tracing to support dynamic shape.

Option 3

Migrate torch/xla frontend to MLIR (some overlap with option 2), e.g. adding an XLA dialect and directly go to MHLO (probably mixed with custom ops untouched). Basically a re-write of current XLA type dispatch using backend fallback, and XLA IR into MLIR dialect (probably using chlo which provides similar functionalities to XlaBuilder). Different fallback strategy could be employed: either using "copy tensors back to CPU and call CPU kernels" similar to what AtenXlaTypeDefault does today, or plugin backend specific intrinsics (require support from runtime side).

Discussion

Option 1 could be a quick shot and doable in 2-4 weeks. I am leaning into option 2 in the mid-term (1-2 month?) and gradually move to option 3 in the long term (6-12 months) and it does enhance the UX of torch/xla. I guess the XLA/TPU stack is also moving to MLIR so we are looking for some early feedbacks here.

PS: great thanks to @silvasean for initiating and coordinating the discussion. Overall roadmap of mlir-npcomp could be found here: https://github.com/llvm/mlir-npcomp/blob/main/docs/roadmap.md

[stellaraccident]

+1 on option 2, which I don't think will be that hard (earlier on, I did make an attempt at it and so had issues, but I think I may have been holding it wrong).

[silvasean]

+1 on option 2 as well.

[ailzhang]

Thanks for providing detailed context in the issue! I haven't gone through all details here yet, will take another look on Monday!
cc: @ezyang @wconstab @asuhan @bdhirsh for more visibility.

[byronyi]

I just found the symbolic shape DSL and structured kernel RFCs as I was thinking that dynamic shape support would need something like meta tensors where you can run the operator statically with type and shape but no real data or computation.

Not sure how long it would take to migrate most existing kernels into structured ones (or even in symbolic shape DSL), but it looks very promising, especially the part that "allowing build of TS compiler independently without aten library".

[ezyang]

So if I understand correctly, the menu of choices from Option 1 - Option 3 represent most conservative to most radical changes to the torch_xla project. Although no one seems to be advocating for it in the short-mid term, Option 3 in particular seems like a really radical change for the torch_xla project, which has from the start been HLO first. Because torch_xla is a real project that people are using to do actual training, we have to keep the plane flying as we change out the engine.

I guess the XLA/TPU stack is also moving to MLIR so we are looking for some early feedbacks here.

This seems to imply that there is organizational buy in from Google? But I'd feel better hearing from, e.g., @JackCaoG, that this is indeed globally Google's plan on record.

Something else that is a little awkward is, independently of this thrust, there's been a recent effort headed by @asuhan and @wconstab to figure out how to factor torch_xla's codebase into an independent, functional lazy tracing mechanism that can be used in other contexts. We'd only been talking about internally recently, but I just put up our working design pytorch/rfcs#18 that discusses some of the considerations. The short answer is that @asuhan has a split out version of XLA that desugars to a simple ATen dialect IR, which then can be further lowered (currently, it only lowers to HLO). This rewrite, however, was done expressly with the intent of otherwise being as low risk as much as possible (so for example, it doesn't do any of the refactors that you've done in your fork, including removal of the ATen operators classes to reduce LOC, or switch to using backend fallback as proposed in Option 2).

Speaking as a core PyTorch developer, here are my thoughts:

• I think it is totally reasonable to want to target MLIR/MHLO as a distinct target from HLO; although PyTorch's infrastructure isn't there yet, we generally want to make it possible for accelerators to target whatever lower level language makes sense for them. But I have less familiarity with the question of whether or not this should be the sole target.
• You wrote about being interested in avoiding baking in dtype/size information. This is a fundamental technical limitation of how XLA integrates, due to the fact that all of our implementation code is in C++ in monomorphic types (e.g., if we ask for a dim, in C++, its an int, not some symbolically traceable object). Torch FX is also interested in this problem, and I have a very vague master evil plan (involving a Python transpiler) but we're a pretty long way from it, it will require major infrastructure investments from PyTorch core. The only workaround we have right now is to use __torch_function__ to capture things like sizes before they get zapped into C++ form so you can retain symbolicness, but we don't really have any existing examples of combining both low level dispatching and __torch_function__ (there's a cluster of issues about this, including Make it possible to skip only one hop of __torch_function__ override pytorch#55093 )

Also, cc @smessmer @ngimel ; these two folks are trying to resolve a major bug in backend fallbacks afflicting master (perhaps that's what stopped @stellaraccident from successfully making this work last time).

[ezyang]

@byronyi

I just found the symbolic shape DSL and structured kernel RFCs as I was thinking that dynamic shape support would need something like meta tensors where you can run the operator statically with type and shape but no real data or computation.

Unfortunately I'm not sure how much this actually helps you. In particular meta tensors still have exact shape/dtype information everywhere (I did want to support symbolic, but it was too hard to figure out how to do that at the same time as other the design constraints), so you won't be able to get dynamic shapes without rerunning the network, at which point XLA's current strategy works perfectly fine.

[eellison]

I think it's possible that JIT symbolic shape analysis could help you here, but there's a lot left to be designed...

[shauheen]

Thanks @ezyang, I will need some more time to read through the RFC and the other proposal but just to comment on this:

I guess the XLA/TPU stack is also moving to MLIR so we are looking for some early feedbacks here.

This seems to imply that there is organizational buy in from Google? But I'd feel better hearing from, e.g., @JackCaoG, that this is indeed globally Google's plan on record.

Want to clarify that MLIR will be used for optimizing many things including XLA however there are no current plans to replace XLA.

[silvasean]

If the goal is to avoid excessive recompilation for dynamic shapes, it probably makes more sense to capture multiple traces and opportunistically erasing shapes to a common supertype if the traces are structurally the same -- should converge to a most de-refined supertype for all ops in the steady state. This could be done "under the hood".

[byronyi]

Thanks @ezyang, I will need some more time to read through the RFC and the other proposal but just to comment on this:

I guess the XLA/TPU stack is also moving to MLIR so we are looking for some early feedbacks here.

This seems to imply that there is organizational buy in from Google? But I'd feel better hearing from, e.g., @JackCaoG, that this is indeed globally Google's plan on record.

Want to clarify that MLIR will be used for optimizing many things including XLA however there are no current plans to replace XLA.

I was just skimming through TPU headers and found XRT variants of TpuCompile: tensorflow/tensorflow@c31e582

It also accepts MLIR module directly: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/tpu/tpu_ops_c_api.h#L111-L120

So TPU runtime still accepts XLA/HLO, and whether frontends other than TF migrates to MLIR should largely depend on if we would like to share some of the optimization passes between these lowering paths (which is exactly what MLIR is primarily designed for IMHO).

[ailzhang]

@byronyi @stellaraccident @silvasean
For option 2, I wonder what do you expect to gain compared to your current mode-base approach? I took a quick look at your mode-based approach (which toggles the included keyset of TLS), although it's slightly different from how xla carries keys on the tensor, it's still a way to tweak our dispatcher and they're not fundamentally different. I feel issues you're facing about conv etc might not go away with option 2, so I'm curious what's the main motivation for it.
(Following our device extension guide could be one reason ;) but we are also interested in learning about issues that are blocking you so that we can see if there's anything we can help :D)

[stellaraccident]

For option 2, I wonder what do you expect to gain compared to your current mode-base approach?

This is my question as well. I've swapped a lot of the context out and would have a hard time articulating it, but ~6 months ago, I did read/trace through all of the PyTorch code in this area, thought I understood it, and did try a version that was more true to the device extension guide. As you note, this doesn't change a lot of the mechanics, and indeed didn't magically fix anything. But I also discount this experience because a) it isn't memorialized anywhere and I may have been holding it wrong even still, and b) there were a lot of other issues with the dispatcher at that time that it appears have been somewhat fixed in the intervening time.

I still suspect that there are "just bugs" here, either on the PyTorch or npcomp side. The areas that were particularly fiddly were the fallback handler (which was, surprisingly, a somewhat leaky abstraction in practice for doing what I actually wanted -- a way to snoop on the dispatch stream with generality) and the bits on the npcomp that are trying to track tensor equality to stitch the graph back together.

If I were putting more cycles into this, I would strip out all of the special cases from the npcomp tracer, see how things work now without them and then walk it back forward, fixing issues upstream vs patching around them. I'd also like someone who knows better to look at the way that the fallback handlers are being used and at least give an ack that it should work. The other examples of using the dispatcher that I could find all use explicit kernel registration, which at least seems to be the more trodden path.

[ailzhang]

@stellaraccident Overall we're interested in unblocking you if there's any dispatcher related bugs/issues. :D So I'm personally interested in what's your current blockers in mode-based approach before we choose among option1/2/3 :D.
@smessmer recently fixed some issues for backend fallback (details in https://dev-discuss.pytorch.org/t/backend-fallbacks/195 ) and it could enable xla to AtenXlaTypeDefault codegen I think. I plan to take a look there in 1-2 weeks to see if we can take advantage of it.
I think in your approach which uses backend fallback it might have hit some unsupported cases in backend fallback in the past, maybe worth a try after this fix?

[sanjoy]

We (TensorFlow/XLA GPU team) are incrementally refactoring XLA GPU to use MHLO and LMHLO (latter being bufferized version of MHLO), with the intent of supporting unbounded dynamic shapes, better custom calls (both of which are important for TensorFlow as well) amongst other things.

OOC, can you share why you need more flexibility in XLA custom calls?

The current "cut point" for this transition is here where we translate XLA HLO graph into LMHLO and go from there.

This is very much a work in progress, but in the spirit of "keep[ing] the plane flying as we change out the engine", we're trying to land all of the work in the default load bearing execution codepath for XLA GPU.