Switch JAX experimental on IPU to new PjRt client. (#21)

The new IPU PjRt client supports:
* Multi IPUs, with collectives (supported by GCL);
* Asynchronous dispatch;
* Slicing/low flops ops fallback on host/cpu backend;

Author: balancap

.github/workflows/jax-ci-ipu-internal.yaml

@@ -20,7 +20,7 @@ jobs:
     timeout-minutes: 10
     steps:
       - name: Cancel previous
-        uses: styfle/cancel-workflow-action@0.10.0
+        uses: styfle/cancel-workflow-action@0.11.0
         with:
           access_token: ${{ github.token }}
         if: ${{github.ref != 'refs/head/jax-v0.3.16-ipu'}}
@@ -67,9 +67,9 @@ jobs:
           python3 setup.py bdist_wheel --universal
           pip3 install dist/*.whl
       # Run IPU specific unit tests
-      - name: Run JAX IPU unit tests
+      - name: Run JAX IPU model unit tests
         run: |
-          XLA_IPU_PLATFORM_DEVICE_COUNT=2 JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=8 pytest --tb=short -vv --log-cli-level=INFO ./tests/ipu/
+          JAX_IPU_DEVICE_COUNT=2 JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=8 pytest --tb=short -vv --log-cli-level=INFO ./tests/ipu/
       # Dockerized workflow known to create issues with self-hosted servers.
       # Solution is to fully cleanup the workspace for the next action.
       # See: https://stackoverflow.com/questions/70483902/how-to-actually-clean-up-the-repository-on-self-hosted-runner-after-github-actio

.github/workflows/jax-ci-ipu-public.yaml

@@ -20,7 +20,7 @@ jobs:
     timeout-minutes: 10
     steps:
       - name: Cancel previous
-        uses: styfle/cancel-workflow-action@0.10.0
+        uses: styfle/cancel-workflow-action@0.11.0
         with:
           access_token: ${{ github.token }}
         if: ${{github.ref != 'refs/head/jax-v0.3.16-ipu'}}
@@ -68,7 +68,7 @@ jobs:
       # Run IPU specific unit tests
       - name: Run JAX IPU unit tests
         run: |
-          XLA_IPU_PLATFORM_DEVICE_COUNT=2 JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=8 pytest --tb=short -vv --log-cli-level=INFO ./tests/ipu/
+          JAX_IPU_DEVICE_COUNT=2 JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=8 pytest --tb=short -vv --log-cli-level=INFO ./tests/ipu/
       # Dockerized workflow known to create issues with self-hosted servers.
       # Solution is to fully cleanup the workspace for the next action.
       # See: https://stackoverflow.com/questions/70483902/how-to-actually-clean-up-the-repository-on-self-hosted-runner-after-github-actio

ipu/docs/build.md

@@ -13,7 +13,7 @@ Note: the build process will work with more recent versions of NumPy, but that w
 Building `jaxlib` currently requires the branch `jax-v0.3.15-ipu`. Once the branch checked out, the build process is similar to other backends:
 ```bash
 export TF_POPLAR_BASE=#...poplar install directory
-python build/build.py --bazel_options=--override_repository=org_tensorflow=PATH/tensorflow-jax-experimental
+python build/build.py --enable_ipu --bazel_options=--override_repository=org_tensorflow=PATH/tensorflow-jax-experimental
 ```
 The `override_repository` config is optional. By default, the build process will pull the experimental IPU TensorFlow XLA code from the repository https://github.com/graphcore-research/tensorflow-jax-experimental.
 

jax/_src/lib/xla_bridge.py

@@ -129,9 +129,19 @@ def ipu_flag_update_hook(name: str, value: Any):
     update_hook=ipu_model_tiles_update_hook)
 flags.DEFINE_bool(
     'jax_ipu_use_legacy_client',
-    bool_env('JAX_IPU_LEGACY_CLIENT', True),
+    bool_env('JAX_IPU_USE_LEGACY_CLIENT', False),
     'Use legacy IPU PjRt client, not supporting multiple IPUs.',
     update_hook=partial(ipu_flag_update_hook, 'jax_ipu_use_legacy_client'))
+flags.DEFINE_integer(
+    'jax_ipu_device_count',
+    int_env('JAX_IPU_DEVICE_COUNT', -1),
+    'Number of IPUs attached to JAX.',
+    update_hook=partial(ipu_flag_update_hook, 'jax_ipu_device_count'))
+flags.DEFINE_string(
+    'jax_ipu_visible_devices',
+    os.getenv('JAX_IPU_VISIBLE_DEVICES', '').lower(),
+    'Specific IPUs visible and attached to JAX.',
+    update_hook=partial(ipu_flag_update_hook, 'jax_ipu_visible_devices'))
 
 def get_compile_options(
     num_replicas: int,

tests/ipu/basics_test.py

@@ -16,6 +16,7 @@
 from jax._src import test_util as jtu
 from unittest import SkipTest
 from functools import partial
+import time
 
 import os
 import numpy as np
@@ -24,16 +25,24 @@
 
 from jax import lax
 from jax.config import config
+from jaxlib.ipu_xla_client import IpuPjRtDevice
+
+# Skipping tests on legacy IPU backend.
+is_ipu_legacy_backend = not isinstance(jax.devices("ipu")[0], IpuPjRtDevice)
 
 
 class IpuBasicsTest(jtu.JaxTestCase):
+  def setUp(self):
+    super().setUp()
+    self.is_ipu_model = config.FLAGS.jax_ipu_use_model
 
   def test_device_count(self):
-    expected = os.getenv('XLA_IPU_PLATFORM_DEVICE_COUNT')
+    expected = os.getenv('JAX_IPU_DEVICE_COUNT')
     expected = int(expected) if expected else 1
-
-    assert jax.device_count(backend='ipu') == expected
-    assert len(jax.devices("ipu")) == expected
+    # Only testing on IPU model for device count.
+    if self.is_ipu_model:
+      assert jax.device_count(backend='ipu') == expected
+      assert len(jax.devices("ipu")) == expected
 
   def test_default_backend(self):
     config.FLAGS.jax_platform_name = 'ipu'
@@ -67,6 +76,31 @@ def add(a,b):
     c = jit_add(a, b)
     self.assertAllClose(c, a + b)
 
+  def test_asynchronous_backend(self):
+    @partial(jax.jit, backend="ipu")
+    def fn(x):
+        y = x * x
+        return y * (y - x)
+
+    N = 1000000
+    x = np.arange(N).astype(np.float32)
+    # First dummy call to compile.
+    fn(x)
+
+    num_iters = 10
+    start = time.perf_counter()
+    for _ in range(num_iters):
+        x = fn(x)
+    async_timing = time.perf_counter() - start
+    x.block_until_ready()
+
+    start = time.perf_counter()
+    for _ in range(num_iters):
+        x = fn(x).block_until_ready()
+    block_timing = time.perf_counter() - start
+
+    # At least 10x faster without blocking.
+    self.assertLessEqual(async_timing * 10, block_timing)
 
   def test_lax_argmin_argmax(self):
     @partial(jax.jit, backend="ipu")

tests/ipu/donate_argnums_test.py

@@ -21,7 +21,7 @@
 import jax
 import numpy as np
 
-class DonateArgnumsTest(jtu.JaxTestCase):
+class IpuDonateArgnumsTest(jtu.JaxTestCase):
 
   def testSingleDonateBufferFirstArgument(self):
 

tests/ipu/infeed_outfeed_test.py

@@ -14,6 +14,7 @@
 
 
 import threading
+import unittest
 
 from absl.testing import absltest
 from jax._src import test_util as jtu
@@ -25,11 +26,19 @@
 from jax.lib import xla_client
 import numpy as np
 
+from jaxlib.ipu_xla_client import IpuPjRtDevice
+
 config.parse_flags_with_absl()
 FLAGS = config.FLAGS
 
-class InfeedTest(jtu.JaxTestCase):
 
+# Skipping tests on new IPU backend.
+is_ipu_legacy_backend = not isinstance(jax.devices("ipu")[0], IpuPjRtDevice)
+
+
+class IpuInfeedTest(jtu.JaxTestCase):
+
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testInfeed(self):
 
     @jax.jit
@@ -45,6 +54,7 @@ def f(x):
     device.transfer_to_infeed((y,))
     self.assertAllClose(f(x), x + y)
 
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testInfeedPytree(self):
 
     x = np.float32(1.5)
@@ -64,6 +74,7 @@ def f(x):
     device.transfer_to_infeed(tuple(flat_to_infeed))
     self.assertAllClose(f(x), to_infeed)
 
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testOutfeed(self):
     hcb.stop_outfeed_receiver()
 
@@ -86,6 +97,7 @@ def f(x):
     device.transfer_from_outfeed(
       xla_client.shape_from_pyval((x,)).with_major_to_minor_layout_if_absent())
 
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testOutfeedPytree(self):
     hcb.stop_outfeed_receiver()
 
@@ -109,6 +121,7 @@ def f(x, y):
     device.transfer_from_outfeed(
       xla_client.shape_from_pyval((x, y)).with_major_to_minor_layout_if_absent())
 
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testInfeedThenOutfeed(self):
     hcb.stop_outfeed_receiver()
 
@@ -131,6 +144,7 @@ def f(x):
     execution.join()
     self.assertAllClose(out, y + np.float32(1))
 
+  @unittest.skipUnless(is_ipu_legacy_backend, "Infeed/outfeed not yet supported on IPU.")
   def testInfeedThenOutfeedInALoop(self):
     hcb.stop_outfeed_receiver()
 

tests/ipu/multi_device_test.py

@@ -12,27 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
 from absl.testing import absltest
 from jax._src import test_util as jtu
-from unittest import SkipTest
+from functools import partial
 
 import jax
 from jax import jit
 import jax.numpy as jnp
-from jax.config import config
+from jaxlib.ipu_xla_client import IpuTargetType
 
 import numpy as np
 
+ipu_num_devices = len(jax.devices("ipu"))
+is_ipu_model = len(jax.devices("ipu")) > 0 and jax.devices("ipu")[0].target_type == IpuTargetType.IPU_MODEL
 
-# Set XLA_IPU_PLATFORM_DEVICE_COUNT=N os env to attach multiple IPUs
-class MultiDeviceTest(jtu.JaxTestCase):
-  def test_jit_with_multi_devices(self):
-    config.FLAGS.jax_platform_name = 'ipu'
-    self.assertEqual(jax.default_backend(), 'ipu')
 
-    devices = jax.devices()
-    if len(jax.devices()) < 2:
-      raise SkipTest("IPU test requires multiple devices")
+# Set JAX_IPU_DEVICE_COUNT=N os env to attach multiple IPUs
+class IpuMultiDeviceTest(jtu.JaxTestCase):
+  @unittest.skipIf(ipu_num_devices < 2, "Requires multiple IPU devices")
+  def test_jit_on_multi_devices(self):
+    self.assertEqual(jax.default_backend(), 'ipu')
+    ipu_devices = jax.devices()
 
     def func(x, w, b):
       return jnp.matmul(w, x) + b
@@ -41,11 +42,26 @@ def func(x, w, b):
     w = np.random.normal(size=[3, 2])
     b = np.random.normal(size=[3, 3])
 
-    for i in range(len(devices)):
-      jit_func = jit(func, device=devices[i])
+    # Just testing on 2 IPU devices, to reduce compilation time.
+    # TODO: support portable executable?
+    for i in range(2):
+      jit_func = jit(func, device=ipu_devices[i])
       r = jit_func(x, w, b)
       self.assertAllClose(r, w @ x + b)
 
+  @unittest.skipIf(ipu_num_devices < 2 or is_ipu_model, "Requires multiple IPU hardware devices")
+  def test_pmap_simple_reduce(self):
+    N = 3
+    data = np.arange(2 * N, dtype=np.float32).reshape((-1, N))
+
+    @partial(jax.pmap, axis_name='i', donate_argnums=(1,), backend="ipu")
+    def parallel_fn(x, y):
+        z = x + jax.lax.psum(y, 'i')
+        return z
+
+    output = parallel_fn(data**2, data)
+    self.assertAllClose(output, data**2 + np.sum(data, axis=0))
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())