Support for IBM Power CPU backend

[Ben Huntsman]

Hi guys-

   How much work would it be to add support for IBM Power, especially Power 10?  The Power 10 architecture has the Matrix Math Accelerator instructions which according to IBM should let it keep up with GPUs for certain workloads.  It would be interesting to see if JAX+XLA could leverage these and what the performance delta would be.

   The documentation makes it sound like it shouldn't be terribly difficult given that there is already an LLVM backend for Power 10.  I assume there is more to the story than this, however, as there would likely need to be some awareness of the MMA instructions and to be able to emit code that the backend can use to target it.

   Thanks in advance for any pointers.

[Penporn Koanantakool]

Hi Ben,

How much work would it be to add support for IBM Power, especially Power 10?

 

For sufficiently large tensor contractions (e.g., matmuls, convolutions, einsums), XLA:CPU calls Eigen contraction kernel instead of codegen. Eigen already supports Power architecture, e.g., recent MRs [1, 2] adding support for Power 10's MMA feature. This should cover a good amount of the performance bottlenecks. Please build XLA with `--define tensorflow_mkldnn_contraction_kernel=0` to avoid using oneDNN matmul as a basic building block for the Eigen contraction kernel.

I'm hoping codegen (op fusions, etc) would work in general if we use the right LLVM backend. (I haven't tried and don't have access to the hardware.) If anyone would like to help with ppc64le codegen optimizations, we are always open for contributions. :)

It would be interesting to see if JAX+XLA could leverage these and what the performance delta would be.

I'd love to know too! If you happen to run some benchmarks there, would you mind sharing the results?

Thank you!

Penporn

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/1603f28e-0de0-4a34-93ed-b6f600e11dfdn%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Ben Huntsman]

Hi Penporn!

   Thank you for the reply!  I do have access to the hardware but I am less familiar with the XLA codebase.  I will try compiling with the arguments you suggested.

   I'm curious though, what front-end is being used with XLA for those recent Power10 contributions?  I know JAX doesn't currently support the Power platform.  The main reason is that JAX is dependent upon BoringSSL, and the BoringSSL team stripped out all the Power code.  I contacted them and they were not interested in adding it back either.  This is kind of a conundrum.  I was able to backport the code but it's obviously unofficial.

   I will report back here after I have XLA compiled on Power 10.

Thank you!

[Kevin Gleason]

Hello!

Is the problem that you can't get JAX to run at all? Or that you need to use the distributed APIs which rely on BoringSSL and those aren't supported on power?

If it's the former you could maybe take an export / model-server workflow, export the StableHLO from framework (JAX or other) then do the XLA compilation on a machine with power 10 support.

Best,

Kevin

[Peter Hawkins]

One thing that seems plausible to me is that we could allow the user to provide an external or system-provided BoringSSL/OpenSSL in our Bazel build, rather than having us trying to self-build one. It's then up to you to provide such an SSL library, however you choose. Work along those lines was previously done in TensorFlow, if I recall correctly, so you may even be able to find inspiration from TF's Bazel logic.

Peter

To view this discussion visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/95b6f298-830d-4718-81f1-810fd5b9f75bn%40openxla.org.