[ingress][torch] PyTorch MLIR backend #51
Conversation
Adds a custom torch.compile backend for JIT-compiling PyTorch models using MLIR. The backend accepts custom compilation function that provides explicit control over the MLIR IR lowering.
rolfmorel
left a comment
rolfmorel
left a comment
There was a problem hiding this comment.
LGTM!
Just a nit or two and a couple of questions.
|
|
||
| # Call jitted function directly with PyTorch tensors. | ||
| # | ||
| # The PyTorch model is specialized into an MLIR module with static shapes. |
There was a problem hiding this comment.
Just to check: this specialization, and hence compilation, can only happen upon the model(a, b) call, right? Before that the model doesn't know the static shapes. Maybe good to document somewhere (as it will be easy to think running inference is really slow, just because it is also compiling (the first time)).
There was a problem hiding this comment.
I'll expand on this. It's part of standard torch.compile flow.
Yes, the shapes are propagated at the execution call as example inputs are required.
They could be actually dynamic represented with symbolic torch tensors. By default, there are heuristics within torch.compile that make these decisions. However, torch-mlir didn't like these symbolic tensors out-of-the-box.
Support for symbolic/dynamic shapes is disabled for now.
There was a problem hiding this comment.
Expanded the example's description.
There was a problem hiding this comment.
Thanks for explaining!
Just a final check: if we/torch-mlir had support for the symbolic dimensions the compilation would happen before the first invocation on a concrete input? Or would it still be delayed to the first invocation?
Not that it matters much to us now. Just in case you happen to know.
2 participants
Adds a custom torch.compile backend for JIT-compiling PyTorch models using MLIR.
The backend accepts custom compilation function that provides explicit control over the MLIR IR lowering.