# Copyright 2025 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for distributed pallas GPU operations."""

from collections.abc import Callable, Sequence
import dataclasses
import functools
import math
import os
import re
import tempfile
import types
from typing import Any, ClassVar, TYPE_CHECKING
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import lax
from jax._src import dtypes
from jax._src import test_multiprocess as jt_multiprocess
from jax._src import test_util as jtu
from jax._src.config import config
from jax._src.lib import cuda_versions
from jax.experimental import multihost_utils
from jax.experimental import pallas as pl
from jax.experimental.mosaic.gpu import launch_context as mgpu_launch_context
from jax.experimental.pallas import mosaic_gpu as _plgpu
from jax.experimental.pallas.ops.gpu.all_gather_mgpu import all_gather
from jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu import reduce_scatter
import jax.numpy as jnp
import numpy as np


P = jax.sharding.PartitionSpec
partial = functools.partial


# We don't want the user to call plgpu.kernel directly, so we monkey patch it
# in `plgpu`.
def do_not_call_me_directly(*args, **kwargs):
  raise RuntimeError("Use self.kernel instead of plgpu.kernel.")

if TYPE_CHECKING:
  plgpu = _plgpu
else:
  # Clone the module locally because the functions are called from other
  # other test files in OSS and the tests are not isolated.
  plgpu = types.ModuleType("_plgpu_local")
  plgpu.__dict__.update(_plgpu.__dict__)

  _kernel = _plgpu.kernel
  del _plgpu

  plgpu.kernel = do_not_call_me_directly


def is_multiprocess():
  return "MULTIPROCESS_TEST" in os.environ


def get_reduction_impl(reduction):
  match reduction:
    case "add":
      return jnp.add
    case "min":
      return jnp.minimum
    case "max":
      return jnp.maximum
    case "and":
      return jnp.bitwise_and
    case "or":
      return jnp.bitwise_or
    case "xor":
      return jnp.bitwise_xor
    case _:
      raise ValueError(reduction)


_TestCaseBase = (jt_multiprocess.MultiProcessTest
                 if is_multiprocess()
                 else parameterized.TestCase)


class MonkeyPatchTest(jtu.JaxTestCase):

  def test_calling_kernel_directly_raises(self):
    with self.assertRaises(RuntimeError):
      plgpu.kernel()


class PallasTestMetaclass(type(_TestCaseBase)):

  def __new__(mcs, *args, lowering_semantics=plgpu.LoweringSemantics.Lane):
    cls = super().__new__(mcs, *args)
    cls.LOWERING_SEMANTICS = lowering_semantics
    return cls


class TestCase(_TestCaseBase, metaclass=PallasTestMetaclass):
  LOWERING_SEMANTICS: ClassVar[plgpu.LoweringSemantics]
  # We track whether we called monkey patched APIs for every test.
  # In tearDown, we verify that the test actually called the monkey-patched APIs.
  # This is not a perfect way of ensuring we do not call plgpu.kernel and
  # plgpu.pallas_call, but is a helpful safeguard.
  monkey_patched_api_was_used: bool

  def setUp(self):
    if jtu.is_device_rocm():
      self.skipTest("Mosaic GPU is not supported on ROCm.")

    if (not jtu.is_device_cuda() or
        not jtu.is_cuda_compute_capability_at_least("9.0")):
      self.skipTest("Only works on GPU with capability >= sm90")

    self.monkey_patched_api_was_used = False
    super().setUp()

  def tearDown(self):
    self.assertTrue(self.monkey_patched_api_was_used)

  def kernel(self, *args, **kwargs):
    compiler_params = dataclasses.replace(
        kwargs.pop("compiler_params", plgpu.CompilerParams()),
        lowering_semantics=self.LOWERING_SEMANTICS,
    )
    result = _kernel(*args, compiler_params=compiler_params, **kwargs)
    self.monkey_patched_api_was_used = True
    return result

  def skipTest(self, msg):
    # Setting `monkey_patched_api_was_used` to true for skipped tests to prevent
    # the assertion failure on teardown.
    self.monkey_patched_api_was_used = True
    super().skipTest(msg)

  def default_transforms(
      self, *, swizzle: int = 128, dtype: jnp.dtype
  ) -> Sequence[plgpu.Transform]:
    swizzle_elems = 8 * swizzle // dtypes.itemsize_bits(dtype)
    return (
        plgpu.TilingTransform((8, swizzle_elems)),
        plgpu.SwizzleTransform(swizzle),
    )

  def assert_arrays_equal_per_shard(
      self,
      y: jax.Array,
      expected_fn: Callable[[int], Any],
      **kwargs,
  ):
    if hasattr(y.sharding, "mesh"):
      devices = list(y.sharding.mesh.devices.flat)
    elif hasattr(y.sharding, "_internal_device_list"):
      devices = list(y.sharding._internal_device_list)
    else:
      devices = list(y.sharding.device_set)
    self.assertNotEmpty(y.addressable_shards)
    if jax.process_count() == 1:
      self.assertLen(y.addressable_shards, len(devices))
    for shard in y.addressable_shards:
      dev_idx = devices.index(shard.device)
      np.testing.assert_array_equal(shard.data, expected_fn(dev_idx), **kwargs)


class PallasCallRemoteDMATest(TestCase):
  def setUp(self):
    super().setUp()
    if jax.device_count() < 2:
      self.skipTest("Needs at least two devices")

  def test_remote_dma_basic(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      y_ref[...] = x_ref[...]
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
          ],
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    y = jax.jit(
        jax.shard_map(
            body,
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),
            check_vma=False,
        )
    )(x)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[8:] if dev_idx == 0 else x[:8]
    )

  @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
  def test_remote_store_vectorized(self):
    if jax.process_index() > 2:
      self.monkey_patched_api_was_used = True
      return  # Only 2 processes needed.

    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = 1 - lax.axis_index("x")
      other_dev = {"x": other_dev_id}
      pl.semaphore_signal(ready_sem, device_id=other_dev)
      pl.semaphore_wait(ready_sem)
      for i in range(2):
        for j in range(2):
          idx = (
              pl.ds(lax.rem(other_dev_id + i, jnp.int32(2)) * 16, 16),
              pl.ds(j * 64, 64),
          )
          plgpu.remote_ref(y_ref, other_dev)[idx] = x_ref[idx]
      pl.semaphore_signal(recv_sem, device_id=other_dev)
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(2 * 32 * 128, dtype=jnp.bfloat16).reshape((2 * 32, 128))
    body = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((32, 128), jnp.bfloat16),
        scratch_types=[
            plgpu.SemaphoreType.REGULAR,
            plgpu.SemaphoreType.REGULAR,
        ],
    )
    mesh = jax.sharding.Mesh(jax.devices()[:2], ["x"])
    with jtu.set_env(MOSAIC_GPU_DUMP_PTX="1"), jtu.capture_stdout() as ptx:
      y = jax.block_until_ready(
          jax.jit(
              jax.shard_map(
                  body,
                  mesh=mesh,
                  in_specs=P("x"),
                  out_specs=P("x"),
                  check_vma=False,
              )
          )(x)
      )
    self.assertIn("st.global.v2.b32", ptx())
    self.assertNotIn("st.b16", ptx())
    self.assertNotIn("st.global.b16", ptx())

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[32:] if dev_idx == 0 else x[:32]
    )

  def test_skip_device_sync(self):
    if jax.process_index() > 2:
      self.monkey_patched_api_was_used = True
      return  # Only 2 processes needed.

    # Kernel with cross-device barrier which makes sure that the other device
    # has completed the previous kernel.
    def barrier_kernel(x_ref, y_ref):
      other_dev_id = {"x": 1 - lax.axis_index("x")}
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]

    # Kernel with no cross-device barrier.
    def skip_barrier_kernel(x_ref, y_ref):
      other_dev_id = {"x": 1 - lax.axis_index("x")}
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...] * 2

    x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))

    # The first kernel is executed with a device barrier, the second without.
    def body(x):
      y = self.kernel(
          barrier_kernel,
          out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
      )(x)
      y_ref = jax.new_ref(y)
      self.kernel(
          skip_barrier_kernel,
          compiler_params=plgpu.CompilerParams(skip_device_barrier=True),
      )(x, y_ref)
      return jax.freeze(y_ref)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    sharded_fn = jax.jit(
        jax.shard_map(
            body,
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),
            check_vma=False,
        )
    )

    hlo = sharded_fn.lower(x).compile().as_text()
    self.assertIn("skip_device_barrier = true", hlo)

  def test_remote_dma_tma_load(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref, ready_sem, recv_sem, scratch_ref, barrier):
      other_dev_id = {"x": 1 - lax.axis_index("x")}
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)

      neighbor_ptr = plgpu.remote_ref(x_ref, other_dev_id)
      plgpu.copy_gmem_to_smem(neighbor_ptr, scratch_ref, barrier)
      plgpu.barrier_wait(barrier)
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)

      plgpu.copy_smem_to_gmem(scratch_ref, y_ref)
      plgpu.wait_smem_to_gmem(0)
      # Make sure the other device's copy completed before exiting the kernel.
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(2 * 8 * 128, dtype=jnp.float32).reshape((2 * 8, 128))
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
              plgpu.SMEM((8, 128), jnp.float32),
              plgpu.Barrier(),
          ],
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    y = jax.jit(
        jax.shard_map(
            body,
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),
            check_vma=False,
        )
    )(x)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[8:] if dev_idx == 0 else x[:8]
    )

  def test_remote_tma_dynamic_peer_id(self):
    if jax.process_index() >= 2:
      self.monkey_patched_api_was_used = True
      return  # Only 2 processes needed.
    # Block `i` of device 0 writes `iota + i` to device `i`, so a wrong entry
    # in the dynamic descriptor table shows up as a wrong shard.
    def kernel(y_ref, scratch_ref, sem):
      dev_id = lax.axis_index("x")
      peer_id = lax.axis_index("peer_id")
      @pl.when(dev_id == 0)
      def _store():
        output = plgpu.layout_cast(
            lax.broadcasted_iota(jnp.int32, (128, 128), 1), plgpu.Layout.WGMMA
        )
        scratch_ref[...] = output + peer_id
        plgpu.copy_smem_to_gmem(
            scratch_ref, plgpu.remote_ref(y_ref, {"x": peer_id})
        )
        plgpu.wait_smem_to_gmem(0)
      pl.semaphore_signal(sem, device_id=1 - dev_id)
      pl.semaphore_wait(sem)
    transforms = self.default_transforms(dtype=jnp.int32)
    def body():
      return self.kernel(
          kernel,
          grid=(2,),
          grid_names=("peer_id",),
          out_type=jax.ShapeDtypeStruct((128, 128), jnp.int32),
          scratch_types=[
              plgpu.SMEM((128, 128), jnp.int32, transforms=transforms),
              plgpu.SemaphoreType.REGULAR,
          ],
      )()
    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=(), out_specs=P("x"), check_vma=False
        )
    )()
    expected = lax.broadcasted_iota(jnp.int32, (128, 128), 1)
    self.assert_arrays_equal_per_shard(y, lambda dev_idx: expected + dev_idx)

  @jtu.thread_unsafe_test()  # Monkey-patches LaunchContext._alloc_scratch.
  def test_recomputable_peer_id_does_not_allocate_dynamic_table(self):
    mesh = jax.sharding.Mesh(jax.devices()[:2], ["x"])
    def scratch_sizes(get_peer_id, **kernel_kwargs):
      """Lowers a remote TMA store and returns its GMEM scratch sizes."""
      def kernel(y_ref, scratch_ref):
        remote_y = plgpu.remote_ref(y_ref, {"x": get_peer_id()})
        plgpu.copy_smem_to_gmem(scratch_ref, remote_y)
        plgpu.wait_smem_to_gmem(0)
      transforms = self.default_transforms(dtype=jnp.int32)
      kernel_call = self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((128, 128), jnp.int32),
          scratch_types=[
              plgpu.SMEM((128, 128), jnp.int32, transforms=transforms)
          ],
          **kernel_kwargs,
      )
      sharded_call = jax.shard_map(
          kernel_call, mesh=mesh, in_specs=(), out_specs=P("x"), check_vma=False
      )
      sizes = []
      alloc_scratch = mgpu_launch_context.LaunchContext._alloc_scratch

      def recording_alloc_scratch(launch_ctx, size, *args, **kwargs):
        sizes.append(size)
        return alloc_scratch(launch_ctx, size, *args, **kwargs)

      with mock.patch.object(
          mgpu_launch_context.LaunchContext,
          "_alloc_scratch",
          recording_alloc_scratch,
      ):
        jax.jit(sharded_call).lower()
      return sizes
    desc_bytes = mgpu_launch_context.TMA_DESCRIPTOR_BYTES
    # `1 - device_id` can be recomputed on the host: a single descriptor.
    self.assertEqual(
        scratch_sizes(lambda: 1 - lax.axis_index("x")), [desc_bytes]
    )
    # A grid-derived peer id can't: a table with one descriptor per peer.
    self.assertEqual(
        scratch_sizes(
            lambda: lax.axis_index("peer_id"), grid=(2,), grid_names=("peer_id",)
        ),
        [desc_bytes * mesh.size],
    )

  def test_remote_dma_dynamic_other_device_id(self):
    # Regression test for device_collective_metadata being DCE'd under
    # WG lowering.
    #
    # We use a dynamic device_id (loaded from a memref)
    # which forces the lowering to use `device_collective_metadata` to resolve
    # the remote address.
    #
    # The choice of `copy_smem_to_gmem` instead of simple store is deliberate as
    # in the former case we do not handle remote ref GMEM transforms in Pallas
    # lowering. As a consequence we do not call `launch_ctx.to_remote` leaving
    # device collective metadata unused and prone to DCE.
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    @self.kernel(
      out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
      scratch_types=[
          plgpu.SMEM((8, 128), jnp.float32),
      ],
    )
    def kernel(x_ref, other_dev_id_ref, y_ref, x_smem):
      x_smem[...] = x_ref[...]
      other_dev_id = {"x": other_dev_id_ref[0]}
      remote_y_ref = plgpu.remote_ref(y_ref, other_dev_id)
      plgpu.copy_smem_to_gmem(x_smem, remote_y_ref)
      plgpu.wait_smem_to_gmem(0)

    x = jnp.arange(2 * 8 * 128, dtype=jnp.float32).reshape((2 * 8, 128))
    other_dev_ids = jnp.array([1, 0], dtype=jnp.int32)
    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    y = jax.jit(
        jax.shard_map(
            kernel,
            mesh=mesh,
            in_specs=(P("x"), P("x")),
            out_specs=P("x"),
            check_vma=False,
        )
    )(x, other_dev_ids)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[8:] if dev_idx == 0 else x[:8]
    )

  # Test verifies an execution of HLO with several slightly different mosaic
  # custom calls. The difference is needed to validate correct initialization
  # of the collective metadata before each kernel execution.
  def test_remote_dma_several_mosaic_ops(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(x_ref, y_ref, done_sem):
      other_dev_id = {"x": 1 - lax.axis_index("x")}
      pl.semaphore_signal(done_sem, device_id=other_dev_id)
      pl.semaphore_wait(done_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(done_sem, device_id=other_dev_id)
      pl.semaphore_wait(done_sem)

    def different_kernel(x_ref, y_ref, wait_sem, ready_sem):
      other_dev_id = {"x": 1 - lax.axis_index("x")}
      pl.semaphore_signal(wait_sem, device_id=other_dev_id)
      pl.semaphore_wait(wait_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)

    def body(x):
      result = x
      for _ in range(25):
        result = self.kernel(
            kernel,
            out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
            scratch_types=[
                plgpu.SemaphoreType.REGULAR,
            ],
        )(result)

        result = self.kernel(
            different_kernel,
            out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
            scratch_types=[
                plgpu.SemaphoreType.REGULAR,
                plgpu.SemaphoreType.REGULAR,
            ],
        )(result)

      return result

    x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))
    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    y = jax.jit(
        jax.shard_map(
            body,
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),
            check_vma=False,
        )
    )(x)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[:8] if dev_idx == 0 else x[8:]
    )

  def test_remote_dma_inline_mgpu(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)

      @plgpu.inline_mgpu(arg_types=(plgpu.RefType(),))
      def write_inline(ctx, ref):
        del ctx
        pass

      write_inline(neighbor_ptr)

    x = jnp.arange(128, dtype=jnp.float32)
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P('x'), out_specs=P('x'), check_vma=False,
        )
    )(x)
    y.block_until_ready()

  def test_remote_dma_atomic_store(self):
    if jax.process_index() > 2:
      self.skipTest("Needs at least two devices")
    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      y_ref[...] = jnp.zeros_like(y_ref)
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)

      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      plgpu.atomic_add(neighbor_ptr, x_ref[...])

      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(256, dtype=jnp.float32)
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((128,), jnp.float32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
          ],
          compiler_params=plgpu.CompilerParams(),
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P('x'), out_specs=P('x'), check_vma=False,
        )
    )(x)
    y.block_until_ready()

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[128:] if dev_idx == 0 else x[:128]
    )

  def test_remote_dma_with_retries(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      y_ref[...] = x_ref[...]
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
          ],
          compiler_params=plgpu.CompilerParams(),
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    jit_body = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P('x'), out_specs=P('x'),
            check_vma=False,
        )
    )
    y = x
    for _ in range(51):
      y = jit_body(y)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[8:] if dev_idx == 0 else x[:8]
    )

  def test_remote_dma_with_profiler(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      y_ref[...] = x_ref[...]
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    # Ignore the warning about the profile already existing since in a
    # single-process mode both device results will try to write to the same
    # profile file.
    with jtu.ignore_warning(category=UserWarning,
                            message=".*profile already exists.*"):
      with tempfile.TemporaryDirectory() as tmpdir:
        x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))
        def body(x):
          return self.kernel(
              kernel,
              out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
              scratch_types=[
                  plgpu.SemaphoreType.REGULAR,
                  plgpu.SemaphoreType.REGULAR,
              ],
              compiler_params=plgpu.CompilerParams(
                  profile_space=2, profile_dir=tmpdir
              ),
          )(x)

        devices = jax.devices()[:2]
        mesh = jax.sharding.Mesh(devices, ['x'])
        y = jax.jit(
            jax.shard_map(
                body, mesh=mesh, in_specs=P('x'), out_specs=P('x'), check_vma=False,
            )
        )(x)

      self.assert_arrays_equal_per_shard(
          y, lambda dev_idx: x[8:] if dev_idx == 0 else x[:8]
      )

  def test_remote_dma_in_loop(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      device_id = lax.axis_index('x')
      other_dev_id = {'x': 1 - device_id}
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      def body(i, _):
        y_ref.at[0, i].set(x_ref.at[0, i].get())
        pl.semaphore_signal(ready_sem, device_id=other_dev_id)
        pl.semaphore_wait(ready_sem)
        neighbor_ptr.at[0, i].set(x_ref.at[0, i].get())
        pl.semaphore_signal(recv_sem, device_id=other_dev_id)
        pl.semaphore_wait(recv_sem)

      lax.fori_loop(0, 128, body, init_val=None, unroll=False)

    x = jnp.arange(2 * 128.0, dtype=jnp.float32).reshape((2, 128))
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((1, 128), jnp.float32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
          ],
          compiler_params=plgpu.CompilerParams(),
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P('x'), out_specs=P('x'), check_vma=False,
        )
    )(x)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[1:] if dev_idx == 0 else x[:1]
    )

  def test_remote_dma_dynamic_index(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    def kernel(x_ref, y_ref, ready_sem, recv_sem):
      other_dev_id = {'x': jnp.sum(x_ref[...] == 1, dtype=jnp.int32)}
      y_ref[...] = x_ref[...]
      pl.semaphore_signal(ready_sem, device_id=other_dev_id)
      pl.semaphore_wait(ready_sem)
      neighbor_ptr = plgpu.remote_ref(y_ref, other_dev_id)
      neighbor_ptr[...] = x_ref[...]
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    x = jnp.zeros((2, 128), dtype=jnp.int32)
    x = x.at[0, 0].set(1)
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((1, 128), jnp.int32),
          scratch_types=[
              plgpu.SemaphoreType.REGULAR,
              plgpu.SemaphoreType.REGULAR,
          ],
          compiler_params=plgpu.CompilerParams(),
      )(x)

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P('x'), out_specs=P('x'), check_vma=False,
        )
    )(x)

    self.assert_arrays_equal_per_shard(
        y, lambda dev_idx: x[1:] if dev_idx == 0 else x[:1]
    )

  @parameterized.parameters(('x',), ('y',))
  def test_remote_dma_2d_mesh(self, axis):
    if (jax.process_count() < 4 and
        (jax.process_count() != 1 or jax.local_device_count() < 4)):
      self.skipTest('Test requires at least 4 devices either accessible by a'
                    'single process or 4 processes with 1 device each.')
    if jax.process_index() > 4:
      return  # Only 4 processes needed.
    def kernel(x_ref, y_ref, recv_sem):
      other_dev_id = {axis: 1 - lax.axis_index(axis)}
      other_y_ref = plgpu.remote_ref(y_ref, other_dev_id)
      other_y_ref[...] = x_ref[...]
      pl.semaphore_signal(recv_sem, device_id=other_dev_id)
      pl.semaphore_wait(recv_sem)

    x = jnp.arange(2 * 8 * 128.0, dtype=jnp.float32).reshape((2 * 8, 128))
    def body(x):
      return self.kernel(
          kernel,
          out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
          scratch_types=[plgpu.SemaphoreType.REGULAR],
          compiler_params=plgpu.CompilerParams(),
      )(x)

    devices = jax.devices()[:4]
    mesh = jax.sharding.Mesh(np.asarray(devices).reshape(2, 2), ['x', 'y'])
    y = jax.jit(
        jax.shard_map(
            body, mesh=mesh, in_specs=P(axis), out_specs=P(axis), check_vma=False,
        )
    )(x)

    axis_idx = ["x", "y"].index(axis)
    self.assert_arrays_equal_per_shard(
        y,
        lambda dev_idx: (
            x[8:] if np.unravel_index(dev_idx, (2, 2))[axis_idx] == 0 else x[:8]
        ),
    )

  def test_wait_twice(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 2, device_id=other_dev_id)
      pl.semaphore_wait(sem)
      pl.semaphore_wait(sem)
      y_ref[...] = jnp.ones_like(y_ref)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        compiler_params=plgpu.CompilerParams(),
    )

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P(None), check_vma=False,
        )
    )()
    np.testing.assert_allclose(y, jnp.ones_like(y))

  def test_wait_nodec(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 2, device_id=other_dev_id)
      pl.semaphore_wait(sem, decrement=False)
      pl.semaphore_wait(sem, 2, decrement=False)
      pl.semaphore_wait(sem, 2)
      y_ref[...] = jnp.ones_like(y_ref)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        compiler_params=plgpu.CompilerParams(),
    )

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P(None), check_vma=False,
        )
    )()
    np.testing.assert_allclose(y, jnp.ones_like(y))

  def test_signal_parallel(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, sem, sem2):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      plgpu.semaphore_signal_parallel(
          plgpu.SemaphoreSignal(sem, device_id=other_dev_id),
          plgpu.SemaphoreSignal(sem2, device_id=other_dev_id),
      )
      pl.semaphore_wait(sem)
      pl.semaphore_wait(sem2)
      y_ref[...] = jnp.ones_like(y_ref)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        scratch_types=[plgpu.SemaphoreType.REGULAR] * 2,
        compiler_params=plgpu.CompilerParams(),
    )

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P(None), check_vma=False,
        )
    )()
    np.testing.assert_allclose(y, jnp.ones_like(y))

  def test_semaphore_signal_collective_axes(self):
    if os.environ.get("NCCL_NVLS_ENABLE") == "0":
      self.skipTest(
          "Skipping multimem test when NCCL_NVLS_ENABLE='0' (b/541298091)."
      )
    if jax.process_index() > 2:
      self.monkey_patched_api_was_used = True
      return  # Only 2 processes needed.

    def kernel(y_ref, sem_out_ref, sem):
      plgpu.semaphore_signal_multicast(sem.at[0], collective_axes='x')
      # Wait for the multicast signal (each device gets signaled by all devices)
      pl.semaphore_wait(sem.at[0], 2)  # Wait for signals from both devices
      y_ref[...] = jnp.ones_like(y_ref)
      sem_out_ref[0] = pl.semaphore_read(sem.at[0])
      sem_out_ref[1] = pl.semaphore_read(sem.at[1])

    kernel_call = self.kernel(
        kernel,
        out_type=(
            jax.ShapeDtypeStruct((8, 128), jnp.float32),
            jax.ShapeDtypeStruct((2,), jnp.int32),
        ),
        scratch_types=[plgpu.SemaphoreType.REGULAR((2,))],
        compiler_params=plgpu.CompilerParams(),
    )

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    y, sem_out = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P(None), check_vma=False,
        )
    )()
    np.testing.assert_allclose(y, jnp.ones_like(y))

    try:
      np.testing.assert_allclose(sem_out, jnp.zeros_like(sem_out))
    except Exception:
      # On some CUDA versions there is a compiler bug where the predicate
      # on the multimem reduction is not respected.
      if cuda_versions.cuda_runtime_get_version() not in [12080, 12090, 13000]:
        raise

  def test_semaphore_signal_multicast_collective_axes_warp_level(self):
    if os.environ.get("NCCL_NVLS_ENABLE") == "0":
      self.skipTest(
          "Skipping multimem test when NCCL_NVLS_ENABLE='0' (b/541298091)."
      )
    if jax.process_index() > 2:
      self.monkey_patched_api_was_used = True
      return  # Only 2 processes needed.

    def kernel(y_ref, sem_out_ref, sem):
      dev_idx = jax.lax.axis_index("x")

      @plgpu.warp_map
      def _per_warp(warp_id):

        @pl.when(warp_id == dev_idx)
        def _():
          # One warp on each device signals. In total, both semaphores get signaled once.
          plgpu.semaphore_signal_multicast(sem.at[warp_id], collective_axes="x")

          y_ref[0, 0] = 1.0

      pl.semaphore_wait(sem.at[0], 1)
      pl.semaphore_wait(sem.at[1], 1)

      sem_out_ref[0] = pl.semaphore_read(sem.at[0])
      sem_out_ref[1] = pl.semaphore_read(sem.at[1])

    kernel_call = self.kernel(
        kernel,
        out_type=(
            jax.ShapeDtypeStruct((8, 128), jnp.float32),
            jax.ShapeDtypeStruct((2,), jnp.int32),
        ),
        scratch_types=[plgpu.SemaphoreType.REGULAR((2,))],
        compiler_params=plgpu.CompilerParams(),
    )

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    _, sem_out = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=(P(None), P(None)), check_vma=False,
        )
    )()

    try:
      np.testing.assert_allclose(sem_out, jnp.zeros_like(sem_out))
    except Exception:
      # On some CUDA versions there is a compiler bug where the predicate
      # on the multimem reduction is not respected.
      if cuda_versions.cuda_runtime_get_version() not in [12080, 12090, 13000]:
        raise

  def test_permuted_mesh(self):
    def kernel(y_ref, sem):
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 1, device_id=other_dev_id)
      pl.semaphore_wait(sem)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        compiler_params=plgpu.CompilerParams(),
    )
    mesh = jax.sharding.Mesh(jax.devices()[::-1], ['x'])  # Reverse the devices.
    f = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P(None), check_vma=False,
        )
    )
    msg = (
        'Mosaic GPU only supports meshes with device ordering that follows'
        ' row-major device ids.'
    )
    with self.assertRaisesRegex(NotImplementedError, msg):
      f()

  @parameterized.parameters(False, True)
  def test_copy_tma(self, use_dict):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, smem_ref, sem):
      dev_id = lax.axis_index("y")
      other_dev_id = 1 - dev_id
      if use_dict:
        ids = lambda x, y: dict(x=x, y=y)
      else:
        ids = lambda x, y: (x, y)

      # Device ID must be an int32.
      zero = jnp.int32(0)

      @pl.when(dev_id == zero)
      def _store():
        output = plgpu.layout_cast(lax.broadcasted_iota(jnp.int32, (128, 128), 1), plgpu.Layout.WGMMA)
        smem_ref[...] = output
        plgpu.copy_smem_to_gmem(smem_ref, plgpu.remote_ref(y_ref, ids(zero, dev_id)))
        plgpu.copy_smem_to_gmem(smem_ref, plgpu.remote_ref(y_ref, ids(zero, other_dev_id)))
        plgpu.wait_smem_to_gmem(0)
      pl.semaphore_signal(sem, 1, device_id=ids(zero, other_dev_id))
      pl.semaphore_wait(sem)

    transforms = self.default_transforms(dtype=jnp.int32)
    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.int32),
        scratch_types=[
            plgpu.SMEM((128, 128), jnp.int32, transforms=transforms),
            plgpu.SemaphoreType.REGULAR,
        ],
        compiler_params=plgpu.CompilerParams(),
    )
    mesh = jtu.create_mesh((1, 2), ("x", "y"))
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P("y"), check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = lax.broadcasted_iota(jnp.int32, (128, 128), 1)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0))

  def test_contiguous_copy_tma(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    shape = (1024,)
    tile = 512

    def kernel(y_ref, smem_ref, sem):
      dev_id = lax.axis_index("y")
      other_dev_id = 1 - dev_id

      # Device ID must be an int32.
      zero = jnp.int32(0)

      @pl.when(dev_id == zero)
      def _store():
        # Unrolling to test that offsets are applied correctly.
        for i in range(2):
          base_data = plgpu.layout_cast(
              jnp.arange(tile, dtype=jnp.int32),
              plgpu.Layout.WG_STRIDED((tile,), vec_size=1),
          )
          smem_ref[...] = base_data + (i * tile)
          plgpu.commit_smem()

          y_slice = pl.ds(i * tile, tile)
          plgpu.copy_smem_to_gmem(
              smem_ref,
              plgpu.remote_ref(y_ref, {"x": zero, "y": dev_id}).at[y_slice]
          )
          plgpu.copy_smem_to_gmem(
              smem_ref,
              plgpu.remote_ref(y_ref, {"x": zero, "y": other_dev_id}).at[y_slice]
          )
          plgpu.wait_smem_to_gmem(0)

      pl.semaphore_signal(sem, 1, device_id={"x": zero, "y": other_dev_id})
      pl.semaphore_wait(sem)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, jnp.int32),
        scratch_types=[
            plgpu.SMEM((tile,), jnp.int32),
            plgpu.SemaphoreType.REGULAR,
        ],
    )
    mesh = jtu.create_mesh((1, 2), ("x", "y"))
    y = jax.jit(
        jax.shard_map(
            kernel_call,
            mesh=mesh,
            in_specs=(),
            out_specs=P("y"),
            check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = jnp.arange(math.prod(shape)).reshape(shape)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0))


class PallasCallMultimemTest(TestCase):
  def setUp(self):
    super().setUp()
    # TODO: b/541298091 - Skip when NCCL_NVLS_ENABLE='0' until b/541298091 is resolved.
    if os.environ.get("NCCL_NVLS_ENABLE") == "0":
      self.skipTest(
          "Skipping multimem test when NCCL_NVLS_ENABLE='0' (b/541298091)."
      )
    if jax.device_count() < 2:
      self.skipTest("Needs at least two devices")
    if any(
      not cuda_versions.cuda_supports_multicast(d.local_hardware_id)
      for d in jax.local_devices()
    ):
      self.skipTest("Not all local devices support multicast")

  def test_multimem_store_regs(self):
    # TODO(bchetioui): support for multimem store.
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, sem):
      @pl.when(lax.axis_index('x') == 0)
      def _store():
        output = plgpu.layout_cast(lax.broadcasted_iota(jnp.int32, (128, 128), 1), plgpu.Layout.WGMMA)
        plgpu.multimem_store(output, y_ref, 'x')
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 1, device_id=other_dev_id)
      pl.semaphore_wait(sem)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.int32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        compiler_params=plgpu.CompilerParams(),
    )
    mesh = jax.sharding.Mesh(jax.devices()[:2], ["x"])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P("x"), check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = lax.broadcasted_iota(jnp.int32, (128, 128), 1)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0))

  def test_multimem_store_scalar(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, sem):
      @pl.when(lax.axis_index('x') == 0)
      def _store():
        plgpu.multimem_store(jnp.int32(1), y_ref.at[0], 'x')
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 1, device_id=other_dev_id)
      pl.semaphore_wait(sem)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((1,), jnp.int32),
        scratch_types=[plgpu.SemaphoreType.REGULAR],
        compiler_params=plgpu.CompilerParams(),
    )
    mesh = jax.sharding.Mesh(jax.devices()[:2], ['x'])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P('x'), check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = jnp.array((1,), jnp.int32)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0), strict=True)

  def test_multimem_store_tma(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    def kernel(y_ref, smem_ref, sem):
      @pl.when(lax.axis_index('x') == 0)
      def _store():
        output = plgpu.layout_cast(lax.broadcasted_iota(jnp.int32, (128, 128), 1), plgpu.Layout.WGMMA)
        smem_ref[...] = output
        plgpu.copy_smem_to_gmem(smem_ref, plgpu.multicast_ref(y_ref, 'x'))
        plgpu.wait_smem_to_gmem(0)
      other_dev_id = {'x': 1 - lax.axis_index('x')}
      pl.semaphore_signal(sem, 1, device_id=other_dev_id)
      pl.semaphore_wait(sem)

    transforms = self.default_transforms(dtype=jnp.int32)
    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct((128, 128), jnp.int32),
        scratch_types=[
            plgpu.SMEM((128, 128), jnp.int32, transforms=transforms),
            plgpu.SemaphoreType.REGULAR,
        ],
        compiler_params=plgpu.CompilerParams(),
    )
    mesh = jax.sharding.Mesh(jax.devices()[:2], ["x"])
    y = jax.jit(
        jax.shard_map(
            kernel_call, mesh=mesh, in_specs=(), out_specs=P("x"), check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = lax.broadcasted_iota(jnp.int32, (128, 128), 1)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0))

  def test_multimem_store_contiguous_tma(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    shape = (1024,)
    tile = 512

    def kernel(y_ref, smem_ref, sem):
      @pl.when(lax.axis_index("x") == 0)
      def _store():
        for i in range(2):
          base_data = plgpu.layout_cast(
              jnp.arange(tile, dtype=jnp.int32),
              plgpu.Layout.WG_STRIDED((tile,), vec_size=1),
          )
          smem_ref[...] = base_data + (i * tile)
          plgpu.commit_smem()

          y_slice = pl.ds(i * tile, tile)
          plgpu.copy_smem_to_gmem(
              smem_ref, plgpu.multicast_ref(y_ref, "x").at[y_slice]
          )
          plgpu.wait_smem_to_gmem(0)

      other_dev_id = {"x": 1 - lax.axis_index("x")}
      pl.semaphore_signal(sem, 1, device_id=other_dev_id)
      pl.semaphore_wait(sem)

    kernel_call = self.kernel(
        kernel,
        out_type=jax.ShapeDtypeStruct(shape, jnp.int32),
        scratch_types=[
            plgpu.SMEM(
                (tile,),
                jnp.int32,
            ),
            plgpu.SemaphoreType.REGULAR,
        ],
    )
    mesh = jax.sharding.Mesh(jax.devices()[:2], ["x"])
    y = jax.jit(
        jax.shard_map(
            kernel_call,
            mesh=mesh,
            in_specs=(),
            out_specs=P("x"),
            check_vma=False,
        )
    )()
    y = multihost_utils.process_allgather(y, tiled=True)
    ref = jnp.arange(math.prod(shape)).reshape(shape)
    np.testing.assert_array_equal(y, np.concat([ref, ref], axis=0))

  @parameterized.parameters(
      (jnp.int32, 1, "add"),
      (jnp.int32, 1, "min"),
      (jnp.int32, 1, "max"),
      (jnp.int32, 1, "and"),
      (jnp.int32, 1, "or"),
      (jnp.int32, 1, "xor"),
      (jnp.int32, 2, "add"),
      (jnp.int32, 4, "min"),
      (jnp.int32, 2, "max"),
      (jnp.int32, 4, "and"),
      (jnp.int32, 2, "or"),
      (jnp.int32, 4, "xor"),
      (jnp.float32, 1, "add"),
      (jnp.float32, 2, "add", True),
      (jnp.float32, 4, "add"),
      (jnp.float16, 2, "add"),
      (jnp.float16, 2, "min"),
      (jnp.float16, 4, "max"),
      (jnp.float16, 8, "add", True),
      (jnp.bfloat16, 2, "max"),
      (jnp.bfloat16, 8, "add"),
      (jnp.float8_e5m2, 4, "add"),
      (jnp.float8_e5m2, 8, "min"),
      (jnp.float8_e5m2, 16, "max", True),
      (jnp.float8_e4m3fn, 4, "min", True),
      (jnp.float8_e4m3fn, 8, "max"),
      (jnp.float8_e4m3fn, 16, "add"),
  )
  def test_multimem_load_reduce(self, dtype, vector_length, reduction, tiled_layout=False):
    if dtype in (
        jnp.float8_e5m2,
        jnp.float8_e4m3fn,
    ) and not jtu.is_cuda_compute_capability_at_least("10.0"):
      self.skipTest("Only works on GPU with capability >= sm100")
    if jax.process_index() > 2:
      return  # Only 2 processes needed.
    devices = jax.devices()[:2]

    def kernel(x_ref, y_ref, sem_ref):
      if tiled_layout:
        layout = plgpu.Layout.TILED(
            plgpu.Tiling(
                (
                    (64, 2 * vector_length),
                    (16, 2 * vector_length),
                    (vector_length,),
                )
            ),
            warp_dims=(-5,),
            lane_dims=(-3, -2),
            vector_dim=-1,
        )
      else:
        layout = plgpu.Layout.WG_STRIDED((64, 32), vec_size=vector_length)
      plgpu.store(
          y_ref,
          plgpu.layout_cast(
              plgpu.multimem_load_reduce(
                  x_ref.at[16:-16], collective_axes="x", reduction_op=reduction,
              ),
              layout
          ),
          optimized=False,
      )
      my_device = lax.axis_index("x")
      other_device = {"x": 1 - my_device}
      pl.semaphore_signal(sem_ref, 1, device_id=other_device)
      pl.semaphore_wait(sem_ref)

    # The rounding we see in low precision types seems to be different from
    # what JAX/XLA use.
    match jnp.dtype(dtype).itemsize:
      case 4:
        bound = 800000
      case 2:
        bound = 128
      case 1:
        bound = 4
      case _:
        raise ValueError(f"Unsupported dtype: {dtype}")
    x_local = jax.random.randint(
        jax.random.key(1234), (128 + 64, 32), dtype=jnp.int32, minval=-bound, maxval=bound,
    ).astype(dtype)
    mesh = jax.sharding.Mesh(devices, ("x",))
    y_shape = jax.ShapeDtypeStruct((64, 32), dtype)
    y = jax.jit(
        jax.shard_map(
            self.kernel(
                kernel,
                out_type=y_shape,
                scratch_types=[plgpu.SemaphoreType.REGULAR],
                compiler_params=plgpu.CompilerParams(),
            ),
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),  # Not really, but lets us test.
            check_vma=False,
        )
    )(x_local)
    y = multihost_utils.process_allgather(y, tiled=True)
    np_reduction = get_reduction_impl(reduction)
    np.testing.assert_array_equal(
        y.astype(jnp.float32),
        np.tile(np_reduction(x_local[16:64+16], x_local[64+48:128+48]), (2, 1)),
    )

  def test_multimem_operand_is_placed_in_collective_memory(self):
    if jax.process_index() > 2:
      return  # Only 2 processes needed.

    # Temporary skip test until XLA switches to a different collective space
    # coloring.
    self.skipTest("Skip the test due to b/564807911.")
    devices = jax.devices()[:2]

    def kernel(x_ref, y_ref, sem_ref):
      y_ref[...] = plgpu.layout_cast(
          plgpu.multimem_load_reduce(
              x_ref.at[...], collective_axes="x", reduction_op="add"
          ),
          plgpu.Layout.WG_STRIDED((8, 128), vec_size=4),
      )
      other_device = {"x": 1 - lax.axis_index("x")}
      pl.semaphore_signal(sem_ref, 1, device_id=other_device)
      pl.semaphore_wait(sem_ref)

    mesh = jax.sharding.Mesh(devices, ("x",))
    x = jnp.zeros((2 * 8, 128), jnp.float32)
    f = jax.jit(
        jax.shard_map(
            self.kernel(
                kernel,
                out_type=jax.ShapeDtypeStruct((8, 128), jnp.float32),
                scratch_types=[plgpu.SemaphoreType.REGULAR],
                compiler_params=plgpu.CompilerParams(),
            ),
            mesh=mesh,
            in_specs=P("x"),
            out_specs=P("x"),
            check_vma=False,
        )
    )

    lowered = f.lower(x).as_text()

    # The multimem operand has to live in symmetric memory, so Mosaic asks XLA
    # to place it in the collective memory space.
    match = re.search(r'operands_memory_spaces = "\{([^}]*)\}"', lowered)
    self.assertIsNotNone(match, lowered)
    operand_memory_spaces = dict(
        tuple(int(i) for i in pair.split(":"))
        for pair in match.group(1).split(",")
        if pair
    )
    collective_memory_space = 1
    self.assertEqual(operand_memory_spaces.get(0), collective_memory_space)

    if jax.process_count() == 1:
      # Within a single process the multimem operand is the only buffer that
      # has to be symmetric. Across processes all of them do.
      self.assertEqual(operand_memory_spaces, {0: collective_memory_space})
      self.assertNotIn("results_memory_spaces", lowered)

    # And XLA has to honour the request by coloring the operand buffers S(1).
    compiled = f.lower(x).compile().as_text()
    custom_calls = [
        line
        for line in compiled.splitlines()
        if 'custom_call_target="mosaic_gpu_v2"' in line
    ]
    self.assertLen(custom_calls, 1, compiled)
    # Operands are printed by name only (and annotated with /*index=N*/ past
    # the first few), e.g. `%custom-call = f32[8,128]{1,0} custom-call(%copy)`.
    operands = re.search(r"custom-call\(([^)]*)\)", custom_calls[0])
    self.assertIsNotNone(operands, custom_calls[0])
    operand_names = [
        re.sub(r"/\*.*?\*/", "", operand).strip()
        for operand in operands.group(1).split(",")
    ]

    def memory_space(name: str) -> int:
      # Definitions carry the memory space in the layout of their shape, e.g.
      #   %copy = f32[16,128]{1,0:S(1)} copy(%param)
      definition = re.search(
          rf"^\s*{re.escape(name)} = (\S+)\s", compiled, re.MULTILINE
      )
      self.assertIsNotNone(
          definition, f"No definition of {name} in:\n{compiled}"
      )
      space = re.search(r"S\((\d+)\)", definition.group(1))
      return int(space.group(1)) if space else 0  # Default memory space.

    # Note that XLA is free to place more buffers in collective memory than we
    # asked for (it currently colors all buffers of a symmetric Mosaic kernel),
    # so we only check that the buffers we did ask for got colored.
    for index, requested_space in operand_memory_spaces.items():
      self.assertLess(index, len(operand_names), compiled)
      self.assertEqual(
          memory_space(operand_names[index]),
          requested_space,
          f"Operand {index} ({operand_names[index]}) is not in memory space"
          f" {requested_space}:\n{compiled}",
      )


@jtu.thread_unsafe_test_class()
class PallasCallMultimemThreadUnsafeTest(TestCase):
  """
  This class is thread-unsafe because the tests monkey patch loaded modules.
  """

  def setUp(self):
    super().setUp()
    # TODO: b/541298091 - Skip when NCCL_NVLS_ENABLE='0' until b/541298091 is resolved.
    if os.environ.get("NCCL_NVLS_ENABLE") == "0":
      self.skipTest(
          "Skipping multimem test when NCCL_NVLS_ENABLE='0' (b/541298091)."
      )
    if jax.device_count() < 2:
      self.skipTest("Needs at least two devices")
    if any(
      not cuda_versions.cuda_supports_multicast(d.local_hardware_id)
      for d in jax.local_devices()
    ):
      self.skipTest("Not all local devices support multicast")

  def _test_reduce_scatter(
      self,
      shape,
      dtype,
      reduction,
      scatter_dimension=0,
      tile_size=None,
      vec_size=None,
      num_blocks=None,
  ):
    if jax.process_index() > 2:
      return

    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    if jnp.issubdtype(dtype, jnp.floating):
      x = jax.random.uniform(jax.random.key(42), shape, dtype=dtype, minval=-1.0, maxval=1.0)
    else:
      x = jax.random.randint(jax.random.key(42), shape, dtype=dtype, minval=-1000, maxval=1000)

    def body(x):
      return reduce_scatter(
          x,
          axis_name="x",
          scatter_dimension=scatter_dimension,
          reduction=reduction,
          vec_size=vec_size,
          tile_size=tile_size,
          num_blocks=num_blocks,
      )

    spec = P(*([None] * scatter_dimension), "x")
    with mock.patch.object(
        jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu.plgpu,
        "kernel",
        self.kernel,
    ):
      y = jax.jit(
          jax.shard_map(
              body, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False
          )
      )(x)

    y = multihost_utils.process_allgather(y, tiled=True)
    np_reduction = get_reduction_impl(reduction)

    split_idx = x.shape[scatter_dimension] // 2
    slices_first = [slice(None)] * len(shape)
    slices_first[scatter_dimension] = slice(None, split_idx)
    slices_second = [slice(None)] * len(shape)
    slices_second[scatter_dimension] = slice(split_idx, None)
    expected = np_reduction(x[tuple(slices_first)], x[tuple(slices_second)])
    tol = 1e-5 if reduction == "add" else 0
    np.testing.assert_allclose(y, expected, rtol=tol, atol=tol)

  @parameterized.parameters(
      (jnp.float32, "add", 1),
      (jnp.float16, "add", 2),
      (jnp.bfloat16, "add", 2),
      (jnp.float16, "min", 4),
      (jnp.float16, "max", 8),
      (jnp.int32, "add", 1),
  )
  def test_reduce_scatter(self, dtype, reduction, vec_size):
    # 16 rows * 64 cols = 1024 elements = 8 elements per thread
    self._test_reduce_scatter(
        (1024, 64), dtype, reduction, tile_size=1024, vec_size=vec_size, num_blocks=4
    )

  def test_reduce_scatter_large_minor_dims(self):
    self._test_reduce_scatter(
        (512, 32768), jnp.float16, "add", tile_size=8192, vec_size=4, num_blocks=4
    )

  @parameterized.parameters(2048, 256, None)
  def test_reduce_scatter_auto_vec_size(self, tile_size):
    self._test_reduce_scatter(
        (1024, 64), jnp.float16, "add", tile_size=tile_size, vec_size=None, num_blocks=4
    )

  @parameterized.parameters(2048, 256, None)
  def test_reduce_scatter_auto_vec_size_int(self, tile_size):
    self._test_reduce_scatter(
        (1024, 64), jnp.int32, "add", tile_size=tile_size, vec_size=None, num_blocks=4
    )

  @parameterized.parameters(1, 2)
  def test_reduce_scatter_different_axes(self, axis):
    if axis == 1:
      shape = (64, 1024, 32)
      tile_size = 2048
    else:  # axis == 2
      shape = (32, 64, 1024)
      tile_size = 2048
    self._test_reduce_scatter(
        shape, jnp.float16, "add", scatter_dimension=axis, tile_size=tile_size, vec_size=None, num_blocks=4
    )

  @parameterized.parameters(
      (jnp.float16, "add"),
      (jnp.float32, "add"),
      (jnp.bfloat16, "max"),
  )
  def test_all_reduce(self, dtype, reduction):
    """Test all-reduce functionality when scatter_dimension=None."""
    self._test_all_reduce(
        (1024, 1024), dtype, reduction, tile_size=512, vec_size=None, num_blocks=4
    )

  def _test_all_reduce(
      self,
      shape,
      dtype,
      reduction,
      tile_size=None,
      vec_size=None,
      num_blocks=None,
  ):
    """Helper function to test all-reduce functionality."""
    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ['x'])
    x = jax.random.normal(jax.random.key(42), (2, *shape), dtype)

    def body(x):
      return reduce_scatter(
          x,
          axis_name="x",
          scatter_dimension=None,  # All-reduce mode
          reduction=reduction,
          vec_size=vec_size,
          tile_size=tile_size,
          num_blocks=num_blocks,
      )

    spec = P("x")
    with mock.patch.object(
        jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu.plgpu,
        "kernel",
        new=self.kernel,
    ):
      y = jax.jit(
          jax.shard_map(
              body, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False
          )
      )(x)
    y = multihost_utils.process_allgather(y, tiled=True)
    np_reduction = get_reduction_impl(reduction)
    expected = np_reduction(x[0], x[1])
    tol = 1e-5 if reduction == "add" else 0
    for ys in y:
      # It seems that the rounding used by the switch is different from what
      # XLA uses.
      y_rounded = np.nextafter(ys, expected)
      np.testing.assert_allclose(y_rounded, expected, rtol=tol, atol=tol)

  def _test_all_gather(
      self,
      shape,
      dtype,
      gather_dimension=0,
      tile_size=None,
      vec_size=None,
      num_blocks=None,
  ):
    if jax.process_index() > 2:
      return

    if jnp.issubdtype(dtype, jnp.floating):
      x = jax.random.uniform(jax.random.key(42), shape, dtype=dtype, minval=-1.0, maxval=1.0)
    else:
      x = jax.random.randint(jax.random.key(42), shape, dtype=dtype, minval=-1000, maxval=1000)

    def body(x):
      return all_gather(
          x,
          axis_name="x",
          gather_dimension=gather_dimension,
          vec_size=vec_size,
          tile_size=tile_size,
          num_blocks=num_blocks,
      )

    spec = P(*([None] * gather_dimension), "x")
    devices = jax.devices()[:2]
    mesh = jax.sharding.Mesh(devices, ["x"])
    with mock.patch.object(
        jax.experimental.pallas.ops.gpu.all_gather_mgpu.plgpu,
        "kernel",
        self.kernel,
    ):
      y = jax.jit(
          jax.shard_map(
              body, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False
          )
      )(x)
    y = multihost_utils.process_allgather(y, tiled=True)
    repeats = [1] * len(x.shape)
    repeats[gather_dimension] = 2

    try:
      np.testing.assert_array_equal(y, np.tile(x, repeats))
    except Exception:
      # On some CUDA versions there is a compiler bug where the predicate
      # on the multimem reduction is not respected.
      if cuda_versions.cuda_runtime_get_version() not in [12080, 12090, 13000]:
        raise

  @parameterized.parameters(
      (jnp.float32, 1),
      (jnp.float16, 2),
      (jnp.bfloat16, 2),
      (jnp.float16, 4),
      (jnp.float16, 8),
      (jnp.int32, 1),
  )
  def test_all_gather(self, dtype, vec_size):
    # 16 rows * 64 cols = 1024 elements = 8 elements per thread
    self._test_all_gather(
        (1024, 64), dtype, tile_size=1024, vec_size=vec_size, num_blocks=4
    )

  def test_all_gather_large_minor_dims(self):
    self._test_all_gather(
        (512, 32768), jnp.float16, tile_size=8192, vec_size=4, num_blocks=4
    )

  @parameterized.parameters(2048, 256, None)
  def test_all_gather_auto_vec_size(self, tile_size):
    self._test_all_gather(
        (1024, 64), jnp.float16, tile_size=tile_size, vec_size=None, num_blocks=4
    )

  @parameterized.parameters(2048, 256, None)
  def test_all_gather_auto_vec_size_int(self, tile_size):
    self._test_all_gather(
        (1024, 64), jnp.int32, tile_size=tile_size, vec_size=None, num_blocks=4
    )

  @parameterized.parameters(1, 2)
  def test_all_gather_different_axes(self, axis):
    if axis == 1:
      shape = (64, 1024, 32)
      tile_size = 2048
    else:  # axis == 2
      shape = (32, 64, 1024)
      tile_size = 2048
    self._test_all_gather(
        shape, jnp.float16, gather_dimension=axis, tile_size=tile_size, vec_size=None, num_blocks=4
    )


class PallasCallRemoteDMAWGTest(
    PallasCallRemoteDMATest,
    lowering_semantics=plgpu.LoweringSemantics.Warpgroup,
):
  ...


class PallasCallMultimemThreadUnsafeWGTest(
    PallasCallMultimemThreadUnsafeTest,
    lowering_semantics=plgpu.LoweringSemantics.Warpgroup,
):
  ...


class PallasCallMultimemWGTest(
    PallasCallMultimemTest,
    lowering_semantics=plgpu.LoweringSemantics.Warpgroup,
):
  ...

if __name__ == '__main__':
  # This test doesn't work with the platform allocator, so we override it
  # if it's ran alone. If it's part of a larger test suite and the platform
  # allocator is used, setUp will skip the test.
  os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.01'
  os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'default'

  if is_multiprocess():
    jt_multiprocess.main()
  else:
    config.config_with_absl()
    absltest.main(testLoader=jtu.JaxTestLoader())