Refactor XLA

[mavenlin]

Description

Refactor the XLA code to simplify it and at the same time include GPU code.

Motivation and Context

XLA was only implemented for CPU, now GPU added.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Put an x in all the boxes that apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• New environment (non-breaking change which adds 3rd-party environment)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of example)

Implemented Tasks

• Common template for XLA custom calls
• Unify CPU/GPU XLA custom calls
• Fix bug in api test.

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code using make format (required)
• I have checked the code using make lint (required)
• I have ensured make bazel-test pass. (required)

[mavenlin]

This has been tested on my local GPU machine.
@Trinkle23897 could you install cuda library on the CI machine? We can test locally, but we need it to compile on CI.

cc @vwxyzjn, once this is in, we can jit entirely on GPU.

[vwxyzjn]

Awesome thank you. Any chance you could build a temporary wheel for me to give it a try?

[mavenlin]

Awesome thank you. Any chance you could build a temporary wheel for me to give it a try?

The machine I have access to is ubuntu with python3.8, would this work for you?

https://drive.google.com/file/d/1OWKgzE7MUz6vpCD9wC8cdCkF2TwGN1AB/view?usp=sharing