Merge pull request #20342 from ROCm:rocm-export_test-add-rocm-platform

PiperOrigin-RevId: 618179680

Author: jax authors

tests/export_harnesses_multi_platform_test.py

@@ -136,9 +136,12 @@ def export_and_compare_to_native(
         if d.platform not in unimplemented_platforms
     ]
     logging.info("Using devices %s", [str(d) for d in devices])
-    # lowering_platforms uses "cuda" instead of "gpu"
+    # lowering_platforms uses "cuda" or "rocm" instead of "gpu"
+    gpu_platform = "cuda"
+    if jtu.is_device_rocm():
+        gpu_platform = "rocm"
     lowering_platforms: list[str] = [
-        p if p != "gpu" else "cuda"
+        p if p != "gpu" else gpu_platform
         for p in ("cpu", "gpu", "tpu")
         if p not in unimplemented_platforms
     ]

tests/export_test.py

@@ -1031,8 +1031,8 @@ def f_jax(x):  # x: f32[10,20] -> f32[20,10]
   def test_multi_platform(self):
     x = np.arange(8, dtype=np.float32)
     exp = get_exported(_testing_multi_platform_func,
-                        lowering_platforms=("tpu", "cpu", "cuda"))(x)
-    self.assertEqual(exp.lowering_platforms, ("tpu", "cpu", "cuda"))
+                        lowering_platforms=("tpu", "cpu", "cuda","rocm"))(x)
+    self.assertEqual(exp.lowering_platforms, ("tpu", "cpu", "cuda", "rocm"))
     module_str = str(exp.mlir_module())
     expected_main_re = (
       r"@main\("
@@ -1054,14 +1054,14 @@ def test_multi_platform(self):
   def test_multi_platform_nested(self):
     x = np.arange(5, dtype=np.float32)
     exp = get_exported(lambda x: _testing_multi_platform_func(jnp.sin(x)),
-                        lowering_platforms=("cpu", "tpu", "cuda"))(x)
-    self.assertEqual(exp.lowering_platforms, ("cpu", "tpu", "cuda"))
+                        lowering_platforms=("cpu", "tpu", "cuda","rocm"))(x)
+    self.assertEqual(exp.lowering_platforms, ("cpu", "tpu", "cuda","rocm"))
 
     # Now serialize the call to the exported using a different sequence of
     # lowering platforms, but included in the lowering platforms for the
     # nested exported.
     exp2 = get_exported(export.call_exported(exp),
-                         lowering_platforms=("cpu", "cuda"))(x)
+                         lowering_platforms=("cpu", "cuda","rocm"))(x)
 
     # Ensure that we do not have multiple lowerings of the exported function
     exp2_module_str = str(exp2.mlir_module())
@@ -1080,8 +1080,8 @@ def test_multi_platform_nested(self):
   def test_multi_platform_nested_inside_single_platform_export(self):
     x = np.arange(5, dtype=np.float32)
     exp = get_exported(_testing_multi_platform_func,
-                        lowering_platforms=("cpu", "tpu", "cuda"))(x)
-    self.assertEqual(exp.lowering_platforms, ("cpu", "tpu", "cuda"))
+                        lowering_platforms=("cpu", "tpu", "cuda","rocm"))(x)
+    self.assertEqual(exp.lowering_platforms, ("cpu", "tpu", "cuda", "rocm"))
 
     # Now serialize the call for the current platform.
     exp2 = get_exported(export.call_exported(exp))(x)
@@ -1120,7 +1120,7 @@ def f_jax(b):  # b: f32[16 // DEVICES, 4]
 
     res_native = f_jax(a)
     exp = get_exported(f_jax,
-                        lowering_platforms=("cpu", "tpu", "cuda"))(a)
+                        lowering_platforms=("cpu", "tpu", "cuda", "rocm"))(a)
 
     # Call with argument placed on different plaforms
     for platform in self.__class__.platforms: