[Pallas:SC] Adds a test that verifies pl.kernel outputs are placed in the proper memory spaces.

PiperOrigin-RevId: 842659708

Author: brianwa84

tests/pallas/tpu_pallas_state_test.py

@@ -14,6 +14,7 @@
 
 import functools
 from absl.testing import absltest
+from absl.testing import parameterized
 import jax
 from jax._src import test_util as jtu
 from jax._src.state.primitives import pin, unpin
@@ -382,6 +383,28 @@ def body(x_ref, o1_ref, o2_ref, scratch_ref):
     np.testing.assert_array_equal(result1, x)
     np.testing.assert_array_equal(result2, x + 1)
 
+  @parameterized.named_parameters(
+      ("HBM", pltpu.HBM, 0),
+      ("VMEM", pltpu.VMEM, 1),
+      ("SMEM", pltpu.SMEM, 4),
+      ("SEMAPHORE", pltpu.SEMAPHORE, 2),
+  )
+  def test_kernel_with_output_memory_space(self, memory_space, color):
+    if not jtu.is_device_tpu_at_least(5):
+      self.skipTest("Only supported on TPU v5+")
+    mesh = pltpu.create_tensorcore_mesh("x", num_cores=1)
+    def body(x_ref, o_ref):
+      pltpu.sync_copy(x_ref, o_ref)
+    x = jnp.arange(8 * 128, dtype=jnp.int32).reshape((8, 128))
+    text = pl.kernel(
+        body, out_shape=memory_space(x.shape, x.dtype), mesh=mesh,
+    ).lower(x).as_text()
+    custom_call = [l for l in text.split("\n") if "@tpu_custom_call" in l]
+    self.assertLen(custom_call, 1)
+    custom_call = custom_call[0]
+    self.assertRegex(custom_call,
+                     r".*output_memory_colors\\22: \[" + str(color) + r"\].*")
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())