We’ve been having parts of a running design discussion on the #mlir-npcomp discord channel, which is a bad habit that I’m breaking here. Including some excerpts for context and to continue.
Folks - just catching up on the discussions here and trying to understand the current and future design so heads up on some noob questions. We currently have a path down from TF2XLA -> MLIR -> Our Custom Dialects -> “Our Distributed / fine grained parallel” runtime -> Hardware. We currently have a functional Torchscript–>Hacks–>runtime that works well. To get rid of the hacks we been playing with npcomp–>MLIR (Aten) with pytorch 1.3 frontend -->…–> Our Custom Dialects --> Runtime -> Hardware. However we don’t particularly want to expose a pseudo-device since our runtime tries to abstract the algorithm from the concept of a device since that is a late binding runtime construct. Reading the comments earlier in this channel there was some interest in Torchscript IR <> Aten MLIR Dialect --> . Is that still an option ?
Also playing with Torchscript IR we have hit a few limitations like https://github.com/facebookresearch/dlrm/issues/69 because the supported types are a subset https://pytorch.org/docs/master/jit_language_reference.html#supported-type .
@powderluv yep, that is “4b” that I mentioned in https://discord.com/channels/636084430946959380/742573221882364009/753347232610123847 (AST Extraction is kind of a synonym for TorchScript in this context). My/Stella’s team are very interested in that direction, since it is important for mobile deployment. I really like how you described the device concept as “late binding runtime construct” – I’m going to steal that from now on
Yes, going directly from TorchScript IR to MLIR entirely at the C++ level is an option (and a good one in my opinion). The TorchScript IR is actually very, very similar to MLIR structurally – a straightforward traveral / conversion is very possible.
about the conversion just FYI here is a python version of torchscript to RelayIR https://github.com/apache/incubator-tvm/blob/master/python/tvm/relay/frontend/pytorch.py . Maybe a good example to use to generate MLIR using the python binding of MLIR for a proof of concept