Enable remat lowering and tests on GPU

Author: jaro-sevcik

jax/_src/dispatch.py

@@ -445,7 +445,7 @@ def device_put_transpose_rule(ct, _, device, src):
 ad.deflinear2(device_put_p, device_put_transpose_rule)
 batching.defvectorized(device_put_p)
 
-def _tpu_device_put_lowering(ctx, x, *, device, src):
+def _tpu_gpu_device_put_lowering(ctx, x, *, device, src):
   if (isinstance(device, (XLACompatibleSharding, TransferToMemoryKind)) and
       device.memory_kind is not None):
     aval, = ctx.avals_in
@@ -456,7 +456,10 @@ def _tpu_device_put_lowering(ctx, x, *, device, src):
           ctx, x, out_aval, device._to_xla_hlo_sharding(aval.ndim).to_proto())
     return [x]
   return [x]
-mlir.register_lowering(device_put_p, _tpu_device_put_lowering, platform='tpu')
+mlir.register_lowering(
+  device_put_p, _tpu_gpu_device_put_lowering, platform='tpu')
+mlir.register_lowering(
+  device_put_p, _tpu_gpu_device_put_lowering, platform='gpu')
 
 
 def _common_device_put_lowering(ctx, x, *, device, src):

tests/memories_test.py

@@ -1118,8 +1118,8 @@ def test_device_put_python_int(self):
 class ActivationOffloadingTest(jtu.JaxTestCase):
 
   def setUp(self):
-    if not jtu.test_device_matches(["tpu"]):
-      self.skipTest("Memories do not work on CPU and GPU backends yet.")
+    if not jtu.test_device_matches(["tpu", "gpu"]):
+      self.skipTest("Memories do not work on CPU backend.")
     super().setUp()
     self.orig_memories_flag = config.enable_memories.value
     jax.config.update('jax_enable_memories', True)
@@ -1167,11 +1167,13 @@ def f(x):
       self.assertRegex(compiled_text, r"copy-done.*S\(5\)")
 
     compiled_stats = compiled_f.memory_analysis()
-    if compiled_stats is not None:
+    if compiled_stats is not None and jtu.test_device_matches(["tpu"]):
       if xla_extension_version >= 240 and jtu.pjrt_c_api_version_at_least(0, 43):
         self.assertGreater(compiled_stats.host_temp_size_in_bytes, 0)
 
   def test_remat_scan_jaxpr_offloadable(self):
+    if not jtu.test_device_matches(["tpu"]):
+      self.skipTest("Remat scan does not work on GPU backend.")
     mesh = jtu.create_global_mesh((2,), ("x",))
     shape = (256, 128)
     np_inp = np.arange(math.prod(shape), dtype=np.float32).reshape(shape)
@@ -1229,6 +1231,8 @@ def g(ys, _):
         self.assertGreater(compiled_stats.host_temp_size_in_bytes, 0)
 
   def test_remat_scan_layout_change_offloadable(self):
+    if not jtu.test_device_matches(["tpu"]):
+      self.skipTest("Remat scan does not work on GPU backend.")
     mesh = jtu.create_global_mesh((2,), ("x",))
     shape = (256, 128)
     np_inp = np.arange(math.prod(shape), dtype=np.float32).reshape(shape)
@@ -1296,7 +1300,7 @@ def f(x):
       self.assertRegex(compiled_text, r"copy-done.*S\(5\)")
 
     compiled_stats = compiled_f.memory_analysis()
-    if compiled_stats is not None:
+    if compiled_stats is not None and jtu.test_device_matches(["tpu"]):
       if xla_extension_version >= 240 and jtu.pjrt_c_api_version_at_least(0, 43):
         self.assertGreater(compiled_stats.host_temp_size_in_bytes, 0)
 

tests/pjit_test.py

@@ -3748,8 +3748,8 @@ def f(inp):
         ' manager.*SingleDeviceSharding'):
       jax.jit(jax.vmap(f, spmd_axis_name='x'))(arr)
 
-  @jtu.skip_on_devices("tpu")
-  def test_device_put_memory_kind_not_tpu(self):
+  @jtu.skip_on_devices("tpu", "gpu")
+  def test_device_put_memory_kind_not_tpu_gpu(self):
     @jax.jit
     def f(x):
       y = x * 2