Integration with TCP

[burmako]

TCP dialect is a super exciting initiative in MLIR upstream (previous Discourse thread, current Discourse thread, design document)! Let's contribute to this effort by providing converters to/from TCP, so that TCP work can potentially leverage them to connect with the existing MHLO and, in the future, StableHLO producers (which currently include JAX, PyTorch and TensorFlow).

[bhack]

I think the thread now is more at:
https://discourse.llvm.org/t/rfc-proposal-for-a-high-level-ml-dialect-in-mlir/64249

[burmako]

Thank you! Looks like I accidentally copy/pasted the wrong link from the design document :)

[burmako]

Current status: there is a proposal to add TCP to the LLVM Incubator (Discourse post, and there'll be an MLIR open design meeting to discuss the proposed spec next Thursday 9/1 (agenda). Nothing actionable for the StableHLO side just, but we're really looking forward to further progress!

[burmako]

Great progress on the TCP side which now has a dedicated category on Discourse and an initial draft of the spec which we discussed at an MLIR ODM last Thursday 9/1.

It looks like a code repository with the TCP dialect is going to be created very soon, and I'm looking forward to starting working on the conversion once it's there!

I've also looked into the logistics of integrating a potential GitHub-first dependency from StableHLO to TCP into the Google-first MLIR-HLO repository. On the technical side, the machinery is fairly straightforward. On the organizational side, this will depend on the eventual license that TCP adopts.

[burmako]

Exciting news! TCP is being bootstrapped in Torch-MLIR (llvm/torch-mlir#1366), and there's a pull request that introduces the needed boilerplate for the TCP dialect as well as a few ops - tanh, add and matmul (llvm/torch-mlir#1375).

As far as conversions between StableHLO and TCP go, the current plan is to wait until the above mentioned pull request lands and then start developing the conversions in Torch-MLIR. Torch-MLIR depends on MLIR-HLO, and StableHLO is shipped as part of MLIR-HLO, so we're already all set dependency-wise. I can't wait!

[bhack]

Do you mean that initially you will work on this conversion inside the torch-mlir incubator repo?

I suppose this cause I don't think we could create a dependency here between this project and torch-mlir repo.

[burmako]

"Do you mean that initially you will work on this conversion inside the torch-mlir incubator repo?". Yes, that's what I briefly discussed with Raghavan, Sean and Stella a few days ago. Planning to send a PR along those lines once TCP lands.

[burmako]

llvm/torch-mlir#1375 has just landed in the Torch-MLIR repository, creating a TCP dialect and setting up all the related boilerplate + a tcp.dummy op. I'm extremely excited for the future development that this PR unlocks! 👏

[burmako]

@navahgar What next steps for TCP do you have in mind? Anything we need to prepare to on the StableHLO side? In the the SIG MLIR Open Design Meeting document, I saw "Broadcasting op & broadcasting semantics in TCP (Raghavan Raman, Cruise)", and I'm looking forward to the ODM :)

[navahgar]

@burmako Our plan is to get to the elementwise ops in TCP first. But that requires broadcasting, which in turn touches several important design points. We are hoping to get through with a decision on broadcasting and elementwise ops soon, preferably in the ODM.

[burmako]

Here are the corresponding design documents: Broadcast and Elementwise. We've discussed the former at the ODM last Thursday 10/13 and came to an intermediate resolution which will be followed by prototyping. Elementwise proposal still needs to be discussed.

[bhack]

Is TCP still an active upstream project? As It seems to me not.

[sjarus]

Is TCP still an active upstream project? As It seems to me not.

It's very much in active development - we'd shifted focus to internally prototyping some of its important constructs and transformations for use in our deployment flows, and to enable us to provide feedback from that experience publicly. We'll be sharing some of this through additional design documents and code in the near future.

[bhack]

Ok thanks for the clarification as I saw no more public activities on Github.

[burmako]

llvm/torch-mlir#1932 introduces the boilerplate for the StableHLO => TCP pass as well as an implementation of converting stablehlo.tanh to tcp.tanh to make sure that all the boilerplate works. Once we're aligned on logistics, I'll implement the rest.

[bhack]

Is TCP actively maintained?

llvm/torch-mlir#1366

[navahgar]

Is TCP actively maintained?

Yes. This branch all the code: https://github.com/llvm/torch-mlir/tree/mlir-tcp

We have also been posting PRs to the discord channel : https://discord.com/channels/636084430946959380/1020532515720671332