Tensorflow plugin vs XLA backend

slai-nick · July 14, 2022, 4:21pm

Hi,

I have been working on pluggable device following this tutorial to integrate a new accelerator.
I just discovered the entry for developing a XLA backend in the Tensorflow documentation and I am now a bit confused as these seems to be two different ways to integrate a new device.

What are the different use cases?
What method is recommended/should be privileged?

slai-nick · July 21, 2022, 3:30pm

Can anyone explain to me the difference?

Bhack · July 22, 2022, 4:24pm

Follow the replies to my post:

github.com/tensorflow/community

RFC: TensorFlow on DirectML

tensorflow:master ← wchao1115:master

opened 06:24AM - 12 May 20 UTC

wchao1115

+116 -0

This RFC will be open for comment until Monday, May 25th, 2020. | Status … | Proposed | :-------------- |:---------------------------------------------------- | | **RFC #** | [243](https://github.com/tensorflow/community/pull/243) | | **Author(s)** | Chai Chaoweeraprasit (wchao@microsoft.com), Justin Stoecker (justoeck@microsoft.com), Adrian Tsai (adtsai@microsoft.com), Patrice Vignola (pavignol@microsoft.com) | | **Sponsor** | Penporn Koanantakool (penporn@google.com) | | **Updated** | 2020-06-08 | ## Objective Implement a new TensorFlow device type and a new set of kernels based on [DirectML](https://docs.microsoft.com/en-us/windows/win32/direct3d12/dml-intro), a hardware-accelerated machine learning library on the DirectX 12 Compute platform. This change broadens the reach of TensorFlow beyond its existing GPU footprint and enables high-performance training and inferencing on Windows devices with any DirectX12-capable GPU.

One important difference is that Tensorflow device plugin is realted to TF only and XLA/OpenXLA is multi frontend (Jax, TF, Pytorch, …):

github.com/tensorflow/community

RFC: Establish SIG OpenXLA

tensorflow:master ← tensorflow:theadactyl-patch-11

opened 08:35PM - 13 Jul 22 UTC

theadactyl

+82 -0

# RFC: Establish SIG OpenXLA | Status | Proposed | :-------------…- |:---------------------------------------------------- | | **RFC #** | [419](https://github.com/tensorflow/community/pull/419)| | **Author(s)** | Thea Lamkin (thealamkin@google.com), Mehdi Amini (aminim@google.com) | | **Sponsor** | Thea Lamkin (thealamkin@google.com) | | **Updated** | 2022-07-13 | We propose to create SIG OpenXLA to facilitate development of an open, state-of-art ML compiler, built collaboratively with ML hardware & framework developers, using the best of XLA & MLIR. ## Objective OpenXLA will be a community-driven and modular open source compiler. It will enable efficient lowering, optimization and deployment of ML models from most major frameworks to any hardware backend notably CPUs, GPUs, and ML ASICs. This work will be done collaboratively with major ML frameworks and hardware vendors. SIG OpenXLA will focus on creating the OpenXLA project, including the extraction of XLA from TensorFlow into a standalone project. SIG discussions will facilitate coordination around roadmap, design evolution, and new workflows to be created in OpenXLA. ## Goals * Accelerate industry collaboration around XLA and build a vibrant OSS community. * Share and receive feedback on the technical direction for OpenXLA and ensure it meets the needs of major users and contributors. * Set up a new XLA repository or organization with independent build/test, with infra to more easily accept PRs, and that is hardware and framework independent. * Ensure the extraction of XLA from TensorFlow is minimally disruptive to existing users and contributors. * Create a product identity with its own brand, website, docs, and communication channels. * Discuss establishment of governance outside TensorFlow.

penporn · July 28, 2022, 10:19pm

Thank you @Bhack for the quick answer!

What are the different use cases?

I can give more details:

Framework coverage: As @Bhack said, PluggableDevice is a TF-specific solution, while OpenXLA supports multiple frontends: TF, JAX, PyTorch, and possibly more.
Modularity:
- PluggableDevice is modular, meaning that you can work on your own repository, release your own plugin package to work with the standard TF package
  - For example, users call pip install my-tf-plug-in on top of normal pip install tensorflow.
  - See more details on advantages of modularity in the PluggableDevice blog post
- XLA currently doesn’t have C API for modular backends yet, so you will have to add code to XLA, and build and release your own “framework+XLA” binaries.
  - For example, users will have to call pip install tensorflow-with-my-xla-backend instead of the standard pip install tensorflow.
  - You might also need to ask PyPI to increase the standard package limit size for you since TensorFlow’s binary size is already over the standard limit by itself. (Your custom TF package would have the size of TF + your device backend.)
  - Device C API for XLA could be a discussion topic in SIG OpenXLA in the future.
Op fusion: XLA can support more aggressive op fusion compared to PluggableDevice.
- Vanilla TensorFlow performs limited op fusion (fused matmul and conv) in grappler remapper pass.
- PluggableDevice could support new fusion patterns through custom graph rewrite pass and custom op kernels, but it would be hard to match XLA’s subgraph clustering.
Op coverage:
- You may have to add more C API for PluggableDevice to support ops with DT_VARIANT data type, e.g., Variable, TensorList, AddNVariant & ZerosLikeVariant ops, etc.
- You don’t have this problem with XLA because it will be in the same build.
Development status:
- PluggableDevice is in semi-maintenance mode: We only add new features when necessary (e.g., to expand support for important ops).
- OpenXLA has been under active development. XLA will eventually be moved out of TensorFlow to its own repository, but we will try to minimize disruptions caused by the move.

What method is recommended/should be privileged?

It depends, based on the trade-offs listed above. If your goal is to support multiple frameworks in one go, I think XLA backend would be the best choice. (JAX doesn’t have custom kernels and only uses XLA.)

Bhack · July 28, 2022, 11:06pm

@penporn I suppose that a third theoretical alternative is to contribute another backend to the new TF runtime but I don’t think we are ready to accept contirbutions in this repository:

penporn · July 29, 2022, 10:48pm

@Bhack The TFRT backend support is still experimental/tentative. We recommend adding support for a new device using PluggableDevice or by building out an XLA backend. Both approaches will have better long-term support from the TF ecosystem. If you’re comfortable with compiler tooling, XLA has the advantage that it can also integrate with other frameworks such as JAX & PyTorch.

slai-nick · August 2, 2022, 10:44am

Thank you for the detailed answer.

Is it possible yet to get a rough time estimate for OpenXLA moving out of TensorFlow and for XLA having modular backends?

penporn · August 16, 2022, 12:49am

Is it possible yet to get a rough time estimate for OpenXLA moving out of TensorFlow and for XLA having modular backends?

Looping in @Mehdi_AMINI for OpenXLA questions. Should we move this discussion to the OpenXLA forum?

In the meanwhile, you might want to check out the slides and notes of the first SIG OpenXLA public meeting. Slide 10 has details on extracting XLA.

Mehdi_AMINI · August 16, 2022, 9:11am

XLA will be in its repo by the end of the year.
Modular backends in XLA isn’t designed yet, so I can’t provide an estimate, but contributions will be welcome on this!