From 02c4aff4c399b06a31a452f02b612d5b8c89dddd Mon Sep 17 00:00:00 2001 From: Manfei <41607353+ManfeiBai@users.noreply.github.com> Date: Wed, 30 Aug 2023 12:33:28 -0700 Subject: [PATCH] [OpenXLA] Add openxla.md (#5507) * Create openxla.md --- docs/openxla.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 docs/openxla.md diff --git a/docs/openxla.md b/docs/openxla.md new file mode 100644 index 00000000000..5dd3c6a76ce --- /dev/null +++ b/docs/openxla.md @@ -0,0 +1,19 @@ +# OpenXLA + +As of June 28th, 2023, PyTorch/XLA now pulls XLA from OpenXLA. +OpenXLA is an [open source machine learning compiler XLA for GPUs, CPUs, and ML accelerators](https://github.com/openxla/xla). + +Previous to OpenXLA, PyTorch/XLA pulled XLA directly from [TensorFlow](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xla). With our [XLA to OpenXLA migration](https://github.com/pytorch/xla/pull/5202), PyTorch/XLA now pulls XLA from [OpenXLA](https://github.com/openxla/xla). + +# How to use OpenXLA + +For [PJRT runtime](https://github.com/pytorch/xla/blob/master/docs/pjrt.md) users, there is no change with this migration. For XRT runtime users, there is a separate [XRT branch of PyTorch/XLA](https://github.com/pytorch/xla/tree/xrt) since OpenXLA doesn't support XRT. + + +# Performance +Below is a performance visual comparison of throughput for ResNet50 pre and post the migration on different TPU hardwares. + +| | resnet50-pjrt-v2-8 | resnet50-pjrt-v4-8 | resnet50-pjrt-v4-32 | +| :------------ | :------------ | :------------ | :------------ | +| Pre Migration | 18.59 | 20.06 | 27.92 | +| Post Migration | 18.63 | 19.94 | 27.14 |