# 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.

from __future__ import annotations

import abc
import collections
from collections.abc import Callable
from collections.abc import Sequence
import contextlib
import dataclasses
import gc
import logging
import threading
import traceback
from typing import Any, Literal

from jax._src.pallas.mosaic.interpret import vector_clock
import jax._src.pallas.mosaic.interpret.params as params
import jax._src.pallas.mosaic.interpret.utils as interpret_utils
import numpy as np

logger = logging.getLogger(__name__)


class Semaphore:

  def __init__(
      self,
      shared_memory: SharedMemory,
      semaphore_id: int,
      enable_logging: bool = False,
  ):
    self.shared_memory = shared_memory
    self.id: int = semaphore_id
    self.enable_logging: bool = enable_logging

    # TODO(jburnim): Use one Condition variable per device.  (Which will be
    # easier to do when we're using single integer device IDs.)
    self.cv = threading.Condition()

    self.count_by_core = np.zeros(self.shared_memory.num_cores, dtype=np.int32)

    # (global_core_id) -> tasks that will signal the semaphore on the core with
    #   the given ID and that should therefore be considered for execution when
    #   the semaphore is waiting (to be signalled).
    self.tasks: list[list[SemaphoreTask]] = [
        [] for _ in range(self.shared_memory.num_cores)
    ]

    if self.shared_memory.detect_races:
      # We associate a vector clock with each count in self.counts.  Whenever
      # self.count_by_core[i] is signaled, self.clocks[i] is updated with the
      # vector clock of the signaling core.  Whenever core i successfully waits
      # on self.count_by_core[i], the vector clock of core i is updated with
      # self.clocks[i].
      #
      # TODO(jburnim): Model happens-before more precisely for the case where
      # semaphores are over-signaled.
      self.clocks: list[SharedMemory.VectorClock | None] = [
          None
      ] * self.shared_memory.num_cores

  @property
  def num_cores(self) -> int:
    return self.shared_memory.num_cores

  @property
  def detect_races(self) -> bool:
    return self.shared_memory.detect_races

  @property
  def dma_execution_mode(self) -> str:
    return self.shared_memory.dma_execution_mode

  def _log(self, message: str):
    """Logs a message. To be called while holding the lock on `self.cv`."""
    # Log every line separately to make sure `absl.logging` adds the correct
    # prefix (i.e. I*** <time> ... <source.py>:<line_number>) to each line in
    # `message`. This should not lead to mangled output within the logging for
    # `self` since the lock on `self.cv` is expected to be held whenever this
    # method is called. However, nothing keeps logged output from being
    # interleaved with logging from other semaphores or from the global
    # `SharedMemory` object.
    for msg in message.split("\n"):
      logger.info(msg)

  def get_global_core_id(self, device_id: int, local_core_id: int) -> int:
    return self.shared_memory.get_global_core_id(device_id, local_core_id)

  def enqueue_task(self, task: SemaphoreTask, global_core_id: int):
    with self.cv:
      self.tasks[global_core_id].append(task)
      self.cv.notify_all()

  def signal(
      self,
      inc,
      global_core_id,
      clock,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ):
    """Signal the semaphore on `(device_id, core_id)` by `inc`.

    Args:
      inc: A positive integer.  The amount by which to increment the semaphore
        on the target device.
      global_core_id: The ID of the target core.
      clock: The vector clock of the signaling device at the time of the signal.
      logging_info: Information about the source of the signal.
    """
    global_core_id = int(global_core_id)
    with self.cv:
      self.count_by_core[global_core_id] += inc

      if self.enable_logging and logging_info is not None:
        self._log(
            logging_info.format(
                f"semaphore={self.id}, inc={inc},"
                f" count_on_core={int(self.count_by_core[global_core_id])}.\n"
                f" count_by_core={self.count_by_core.tolist()}.",
                line_prefix="`signal`",
            )
        )

      if self.detect_races:
        if (global_clock := self.clocks[global_core_id]) is None:
          self.clocks[global_core_id] = clock.copy()
        else:
          global_clock.update(clock)
      self.cv.notify_all()

  def read(self, global_core_id):
    with self.cv:
      return self.count_by_core[global_core_id]

  def wait(
      self,
      value,
      global_core_id,
      *,
      has_tasks=False,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ):
    global_core_id = int(global_core_id)
    thread = self.shared_memory.global_core_id_to_thread(global_core_id)

    # TODO(nrink): Update the comment below to generalize from DMAs and DMA
    # semaphores. We now have the concept of 'tasks' that can signal a
    # semaphore. At the moment, DMAs are the only tasks that occur; and what is
    # allowed to be a task may still change (because it should probably be more
    # restricted than allowing tasks to be arbitrary callables, as is currently
    # done).
    #
    # For DMA semaphores (when shared_memory.dma_execution_mode=='on_wait'),
    # while our count is not large enough we will select and partially execute
    # pending DMAs until our count is large enough.
    #
    # This approach will tend to run DMAs as late as possible, as well as
    # out-of-order.  This approach also lets us avoid the complexity of spinning
    # up separate threads to handle executing DMAs.
    while True:
      self.shared_memory.check_failed()

      clock = None
      done = False
      task = None
      with self.cv:
        # TODO(jburnim):
        #  - If the count is larger than value, raise an error?
        #  - If the count is equal to value, but there DMAs waiting to signal
        #    us, raise an error?
        if self.count_by_core[global_core_id] >= value:
          self.count_by_core[global_core_id] -= value
          done = True
          if self.enable_logging and logging_info is not None:
            self._log(
              logging_info.format(
                  f"semaphore={self.id} finished waiting for {value=}, "
                  f"count_on_core={int(self.count_by_core[global_core_id])}.\n"
                  f"count_by_core={self.count_by_core.tolist()}.",
                  line_prefix="`wait`",
              )
          )
          if self.detect_races:
            clock = self.clocks[global_core_id]
            assert clock is not None
            clock = clock.copy()
        elif len(self.tasks[global_core_id]) > 0:
          task = self.tasks[global_core_id].pop()
        else:
          # We need to set a timeout here so that we periodically check for
          # failures set in `self.shared_memory`.
          self.cv.wait(timeout=0.1)

      if clock is not None:
        with self.shared_memory.lock:
          self.shared_memory.clocks[thread].update(clock)

      if done:
        return

      if task is not None:
        task()


# A `SemaphoreTask` is called when a semaphore is waiting to be signalled on a
# specific core. A `SemaphoreTask` will typically capture the `Semaphore` object
# that is waiting, so that when the task is called, it can signal the semaphore
# (by calling `Semaphore.signal` from within the task). When a `SemaphoreTask`
# object is called, it can be assumed that the call stack of the task will
# *not* hold the lock on the shared memory in the captured `Semaphore` object.
# This allows the task to use methods from `SharedMemory` to access and modify
# the global shared memory object.
SemaphoreTask = Callable[[], None]


@dataclasses.dataclass(init=False)
class Allocation:
  ...


class Buffer(Allocation):

  def __init__(
      self,
      content: np.ndarray,
      ref_count: int = 1,
      logical_shape: tuple[int, ...] | None = None,
  ):
    super().__init__()

    self._content = content
    self._ref_count = ref_count

    if logical_shape is None:
      self._logical_shape = tuple(self._content.shape)
    else:
      self._logical_shape = logical_shape

    for dim, ldim in zip(self.content.shape, self.logical_shape, strict=True):
      if ldim > dim:
        raise ValueError(
            f"Logical shape {self.logical_shape} cannot be bigger than content"
            f" shape {self.content.shape}."
        )

  @property
  def content(self) -> np.ndarray:
    return self._content

  @property
  def ref_count(self) -> int:
    return self._ref_count

  @property
  def logical_shape(self) -> tuple[int, ...]:
    return self._logical_shape

  def decrease_ref_count(self):
    # We should never decrease the `ref_count` to below zero.
    assert self.ref_count > 0
    self._ref_count -= 1

  def has_zero_ref_count(self) -> bool:
    return self.ref_count == 0

  @property
  def size(self) -> int:
    return self.content.itemsize * self.content.size

  @property
  def shape(self) -> tuple[int, ...]:
    return self.content.shape

  @property
  def dtype(self) -> np.dtype:
    return self.content.dtype

  def _normalize_range(
      self, rnge: tuple[slice | int, ...]
  ) -> tuple[slice | int, ...]:
    """Normalizes `rnge` by adding slices to match the `self.logical_shape`."""
    assert len(self.logical_shape) == len(self.shape)
    return rnge + tuple(slice(0, s) for s in self.logical_shape[len(rnge):])

  def __getitem__(self, rnge: tuple[slice | int, ...]) -> np.ndarray:
    """Returns the portion of `self.content` specified by `rnge`.

    Args:
      rnge: The range to read.

    Raises:
      IndexError: If `rnge` is entirely out of bounds for the allocated array in
        the `Buffer`, i.e. `rnge` is out of bounds for `self.shape`.

    Returns:
      The portion of `self.content` specified by `rnge`.
    """
    rnge = self._normalize_range(rnge)
    rnge_or_none = interpret_utils.clip_range_to_shape(rnge, self.shape)
    if rnge_or_none is None:
      # Raise if reading entirely outside of the allocated shape. We leave it to
      # the client to handle the case where out-of-bounds reads are allowed (and
      # should return uninitialized values).
      raise IndexError(
          f"Range {rnge} is entirely out of bounds for shape {self.shape}."
      )

    rnge = rnge_or_none
    return self.content[rnge]

  def _set_within_logical_shape(
      self, rnge: tuple[slice | int, ...], value: np.ndarray
  ):
    """Updates `self.content` with `value` for the portion of `rnge` within `self.logical_shape`."""
    rnge_or_none = interpret_utils.clip_range_to_shape(rnge, self.logical_shape)
    if rnge_or_none is None:
      return

    rnge = rnge_or_none
    shape_to_write = self.content[rnge].shape
    self.content[rnge] = value[tuple(slice(0, s) for s in shape_to_write)]

  def __setitem__(self, rnge: tuple[slice | int, ...], value: np.ndarray):
    """Updates `self.content` with `value`, if `rnge` is fully within `self.shape`.

    Args:
      rnge: The range to write.
      value: The value to write.

    Raises:
      IndexError: If any part of `rnge` is out of bounds for the allocated array
        in the `Buffer`, i.e. if any part of `rnge` is out of bounds for
        `self.shape`.
    """
    rnge = self._normalize_range(rnge)
    if interpret_utils.is_range_out_of_bounds_for_shape(rnge, self.shape):
      raise IndexError(
          f"Range {rnge} is (at least partially) out of bounds for"
          f" allocation shape {self.shape}."
      )

    self._set_within_logical_shape(rnge, value)

  def set_in_bounds_portion(
      self, rnge: tuple[slice | int, ...], value: np.ndarray
  ):
    """Updates `self.content` with `value` for the portion of `rnge` within `self.logical_shape`."""
    rnge = self._normalize_range(rnge)
    self._set_within_logical_shape(rnge, value)


@dataclasses.dataclass(frozen=True)
class ShapeAndDtype:
  shape: Sequence[int]
  dtype: np.dtype

  def __iter__(self):
    return iter((self.shape, self.dtype))


@dataclasses.dataclass(kw_only=True)
class GenericSharedMemory[
    MemKey, ThreadKey, VectorClock: vector_clock.VectorClockProto
](abc.ABC):

  num_devices: int
  out_of_bounds_reads: Literal["raise", "uninitialized"]
  dma_execution_mode: str
  uninitialized_memory: Literal["nan", "zero"]
  detect_races: bool
  vector_clock_size: int

  clocks: dict[ThreadKey, VectorClock]
  barrier: threading.Barrier
  clean_up_barrier: threading.Barrier

  buffer_bounds: Literal["logical", "padded"] | None = None

  logging_mode: params.LoggingMode | None = None

  mem: dict[MemKey, Allocation] = dataclasses.field(
      default_factory=dict
  )

  lock: threading.Lock = dataclasses.field(default_factory=threading.Lock)

  next_buffer_id: dict[ThreadKey, int] = dataclasses.field(
      default_factory=lambda: collections.defaultdict(lambda: 100)
  )
  deallocated_bytes: int = 0

  # ThreadKey -> [(grid_index, [range])]
  output_ranges: dict[ThreadKey, list] = dataclasses.field(
      default_factory=lambda: collections.defaultdict(list)
  )

  _failure: Exception | None = None
  # The thread on which `_failure` occurred, if known. Only used in error
  # messages, so subclasses may record any object with an informative `repr`
  # (e.g. a device, if the thread is not known).
  _failed_thread: Any = None

  @property
  def enable_logging(self) -> bool:
    return (
        self.logging_mode is not None
        and params.LoggingMode.SHARED_MEMORY in self.logging_mode
    )

  def _log(self, message: str):
    """Logs a message. To be called while holding `self.lock`."""
    # Log every line separately to make sure `absl.logging` adds the correct
    # prefix (i.e. I*** <time> ... <source.py>:<line_number>) to each line in
    # `message`. This should not lead to mangled output as `self.lock` is
    # expected to be held whenever this method is called.
    for msg in message.split("\n"):
      logger.info(msg)

  @abc.abstractmethod
  def thread_to_vc_position(self, thread: ThreadKey) -> int:
    ...

  def incr_clock(self, thread: ThreadKey, take_lock: bool = True):
    """Increments a threads's own index within its clock by one."""
    with self.lock if take_lock else contextlib.nullcontext():
      pos = self.thread_to_vc_position(thread)
      self.clocks[thread].inc(pos)
      return self.clocks[thread].copy()

  def get_next_buffer_id(self, key: ThreadKey) -> int:
    """Returns the next buffer ID for the given device and thread."""
    with self.lock:
      buffer_id = self.next_buffer_id[key]
      self.next_buffer_id[key] = buffer_id + 1
      return buffer_id

  def set_failed(
      self,
      exception: Exception,
      failed_thread: Any = None,
      top_level: bool = True,
  ):
    with self.lock:
      if self._failure is None:
        self._failure = exception
        self._failed_thread = failed_thread
      # Abort while still holding lock, so no new waiters can appear between
      # recording the failure and aborting existing waiters.
      self._abort_waiters()

    self.barrier.abort()

    # If we are interpreting a kernel over N devices, we must wait N times
    # on the clean-up barrier.  As the computation on this device has failed
    # and will soon raise an exception, this device will not reach its final
    # call to `_clean_up_shared_memory` (which waits on the clean-up barrier).
    # So, we wait on the barrier here.
    #
    # Unless we are running inside a thread_map (i.e., top_level=False) -- in
    # which case we are one of the num_cores_per_device computations running for
    # one of the N devices.  In this case, do not wait on the barrier here.
    # The thread_map will detect that one or more of its threads failed and,
    # once all of its threads have completed (successfully or with an error),
    # it will call `set_failed` with `top_level=True`.
    if top_level:
      self.clean_up_barrier.wait()

  def _abort_waiters(self):
    """Wakes up all threads that are waiting on other threads.

    Lock should be held when called.
    """
    pass

  def check_failed(self):
    with self.lock:
      if self._failure is not None:
        # __cause__ information is lost when an exception from a callback goes
        # through XLA and back to Python, so we stuff the information into the
        # exception message.
        #
        # TODO(jburnim): The top-level exception (that the user sees) will not
        # always contain the original exception info/message.  It currently
        # depends on which raise (the raise here or the raise in the thread
        # that calls `set_failed`) happens first. (And maybe some other details
        # in XLA.)  We should figure out how to reliably propagate the
        # original exception.
        failure_str = "".join(traceback.format_exception(self._failure))
        failed_info = (
            f" on {self._failed_thread}"
            if self._failed_thread is not None
            else ""
        )
        raise RuntimeError(
            f"Computation failed{failed_info} with exception:\n\n{failure_str}"
        ) from None

  def allocate_buffer(
      self,
      key: MemKey,
      ref_count: int,
      value: np.ndarray,
      logical_shape: tuple[int, ...] | None = None,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ):
    """Allocates a memory buffer with the given key unless it already exists."""
    with self.lock:
      if key not in self.mem:
        buff = Buffer(
            value, ref_count=ref_count, logical_shape=logical_shape
        )
        self.mem[key] = buff

        if self.enable_logging and logging_info is not None:
          self._log(
              logging_info.format(
                  f"{key=}, {ref_count=}.\nvalue_shape={value.shape},"
                  f" logical_shape={buff.logical_shape},"
                  f" content_shape={buff.shape}",
                  line_prefix="`allocate_buffer`",
              )
          )

  def deallocate_buffer(
      self, key: MemKey, logging_info: interpret_utils.LoggingInfo | None = None
  ):
    """Decreases the ref count for the buffer with `key` and deallocates the buffer if the ref count is zero."""
    with self.lock:
      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to deallocate allocation with key `{key}` that is not"
            " a `Buffer`."
        )

      buff.decrease_ref_count()
      if buff.has_zero_ref_count():
        # TODO(paulbib): delete buffer from race detection state as well
        self.mem.pop(key)
        self.deallocated_bytes += buff.size
        del buff

        if self.enable_logging and logging_info is not None:
          self._log(
              logging_info.format(f"{key=}.", line_prefix="`deallocate_buffer`")
          )

      should_collect = self.deallocated_bytes > 100_000_000
      if should_collect:
        self.deallocated_bytes = 0

    if should_collect:
      # Periodic garbage collection here prevents OOMs -- although it's not clear
      # why arrays are not getting freed without this.
      gc.collect()

  def get_buffer_content(
      self,
      key: MemKey,
      rnge: tuple[slice | int, ...],
      thread: ThreadKey | None,
      increment_clock: bool = True,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ) -> tuple[np.ndarray | None, ShapeAndDtype, VectorClock | None]:
    """Reads contents of a memory buffer.

    Args:
      key: The key of the buffer to read.
      rnge: The range to read within the buffer.
      thread: The thread reading the buffer.
      increment_clock: Whether to increment the given thread's vector clock
      logging_info: Information about the source of the read.

    Returns:
      - The contents of the read range of the buffer, or None if reading
        entirely out of bounds.
      - The shape and dtype of the full content array of the buffer.
      - The incremented vector clock for the given thread.
        None if race detection is not enabled or if `increment_clock` is False.
    """
    clock = None
    with self.lock:
      if self.detect_races and increment_clock and thread is not None:
        clock = self.incr_clock(thread, take_lock=False)

      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to get contents of allocation with key `{key}` that is"
            " not a `Buffer`."
        )
      shape_and_dtype = ShapeAndDtype(buff.logical_shape, buff.dtype)

      try:
        result = buff[rnge].copy()
      except IndexError:
        # `buf` was accessed with `rnge` entirely out of bounds.
        result = None

      if self.enable_logging and logging_info is not None:
        self._log(
            logging_info.format(
                f"{key=}, {rnge=},"
                f" in_bounds={result is not None}.\n"
                f"logical_shape={buff.logical_shape},"
                f" content_shape={buff.shape},"
                f" {f'{result.shape=}' if result is not None else ''}.",
                line_prefix="`get_buffer_content`",
            )
        )

    return result, shape_and_dtype, clock

  def store_buffer_content(
      self,
      key: MemKey,
      rnge: tuple[slice | int, ...],
      value: np.ndarray,
      thread: ThreadKey,
      increment_clock: bool = True,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ) -> tuple[bool, ShapeAndDtype, VectorClock | None]:
    """Stores contents into a memory buffer.

    Args:
      key: The key of the buffer to store into.
      rnge: The range within the buffer contents that `value` is written to.
      value: The array to store into the buffer.
      thread: The thread writing into the buffer.
      increment_clock: Whether to increment the given thread's vector clock.
      logging_info: Information about the source of the store.

    Returns:
      - True if the store was entirely in bounds, False otherwise (i.e. if the
        store was at least partially out of bounds).
      - The shape and dtype of the full content array of the buffer.
      - The incremented vector clock for the given thread.
        None if race detection is not enabled or if `increment_clock` is False.
    """
    clock = None
    with self.lock:
      if self.detect_races and increment_clock:
        clock = self.incr_clock(thread, take_lock=False)

      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to store into allocation with key `{key}` that is not"
            " a `Buffer`."
        )
      shape_and_dtype = ShapeAndDtype(buff.logical_shape, buff.dtype)

      assert buff.dtype == value.dtype  # TODO(jburnim): Catch this statically.

      try:
        buff[rnge] = value
        is_in_bounds = True
      except IndexError:
        # `buf` was accessed with `rnge` at least partially out of bounds.
        is_in_bounds = False

      if self.enable_logging and logging_info is not None:
        self._log(
            logging_info.format(
                f"{key=}, {rnge=},"
                f" in_bounds={is_in_bounds}.\n"
                f"logical_shape={buff.logical_shape},"
                f" content_shape={buff.shape}, {value.shape=}.",
                line_prefix="`store_buffer_content`",
            )
        )

      return is_in_bounds, shape_and_dtype, clock

  def swap_buffer_content(
      self,
      key: MemKey,
      rnge: tuple[slice | int, ...],
      value: np.ndarray,
      mask: np.ndarray | None,
      thread: ThreadKey,
      increment_clock: bool = True,
      logging_info: interpret_utils.LoggingInfo | None = None,
  ) -> tuple[np.ndarray | None, ShapeAndDtype, VectorClock | None]:
    """Swaps contents of a memory buffer.

    Args:
      key: The key of the buffer to swap into.
      rnge: The range within the buffer contents that `value` is swapped into.
      value: The array to be written into the buffer.
      mask: The mask to apply to the swap operation.
      increment_clock: Whether to increment the given thread's vector clock.
      thread: The thread that's writing into the buffer.
      logging_info: Information about the source of the swap.

    Returns:
      - The contents of the range of the buffer (prior to the swap), or None if
        accessing buffer contents bounds.
      - The shape and dtype of the full content array of the buffer.
      - The incremented vector clock for the given thread.
        None if race detection is not enabled or if `increment_clock` is False.
    """
    clock = None
    with self.lock:
      if self.detect_races and increment_clock:
        clock = self.incr_clock(thread, take_lock=False)

      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to swap into allocation with `key` {key} that is not a"
            " `Buffer`."
        )

      shape_and_dtype = ShapeAndDtype(buff.logical_shape, buff.dtype)

      assert buff.dtype == value.dtype  # TODO(jburnim): Catch this statically.
      # TODO(jburnim): Better error message if this raises?

      try:
        result = buff[rnge].copy()
      except IndexError:
        # `buf` was accessed with `rnge` entirely out of bounds.
        result = None

      if result is not None:
        in_bounds_shape = result.shape

        if mask is None:
          if in_bounds_shape != value.shape:
            # Runs past the end of the buffer, and there is nowhere to put the
            # part that does not fit. Reported the same way a wholly
            # out-of-bounds swap is.
            result = None
          else:
            buff[rnge] = value
        else:
          in_bounds_mask = np.full(mask.shape, True)
          for i in range(len(in_bounds_shape)):
            in_bounds_mask[in_bounds_shape[i] :] = False
          if (~in_bounds_mask & mask).any():
            result = None
          else:
            in_bounds_idx = tuple(slice(i) for i in in_bounds_shape)
            raw_result = result
            result = value.copy()
            result[in_bounds_idx] = np.where(
                mask[in_bounds_idx], raw_result, value[in_bounds_idx]
            )
            buff.set_in_bounds_portion(
                rnge,
                np.where(mask[in_bounds_idx], value[in_bounds_idx], raw_result),
            )
            # Assert that `np.where` is not a view, and hence, in particular,
            # does not share underlying memory with `buff.content`.
            assert result.base is None

      if self.enable_logging and logging_info is not None:
        self._log(
            logging_info.format(
                f"{key=}, {rnge=},"
                f" in_bounds={result is not None}.\n"
                f"logical_shape={buff.logical_shape},"
                f" content_shape={buff.shape},"
                f" {f'{result.shape=}' if result is not None else ''},"
                f" {value.shape=}.",
                line_prefix="`swap_buffer_content`",
            )
        )

      return result, shape_and_dtype, clock

  def update_clocks(self, threads: Sequence[ThreadKey]):
    """Synchronizes the vector clocks for all of the given threads."""
    # Despite only updating the vector clocks for some threads, we still need to
    # hold the global lock to ensure that no other devices are concurrently
    # accessing the same vector clocks.
    assert len(threads) > 0
    init_thread = threads[0]
    with self.lock:
      for thread in threads[1:]:
        self.clocks[init_thread].update(self.clocks[thread])
      for thread in threads[1:]:
        self.clocks[thread].update(self.clocks[init_thread])


@dataclasses.dataclass
class SharedMemory(
    GenericSharedMemory[
        tuple[str, int, int, int], tuple[int, int], vector_clock.NpVectorClock
    ]
):
  """The shared state of memory for TPU interpret mode."""

  # (memory_space, buffer_id, device_id, local_core_id)
  MemKey = tuple[str, int, int, int]
  # (device_id, local_core_id)
  ThreadKey = tuple[int, int]
  VectorClock = vector_clock.NpVectorClock

  num_cores_per_device: int

  # semaphore_id -> Semaphore
  sem: dict[int, Semaphore]

  # global_core_id -> next semaphore ID
  next_semaphore_id: dict[int, int]

  # semaphore_id -> Semaphore, where the semaphore_id is user-specified.
  fixed_id_sem: dict[int, Semaphore]

  def __init__(
      self,
      *,
      num_devices: int,
      out_of_bounds_reads: Literal["raise", "uninitialized"],
      dma_execution_mode: str,
      uninitialized_memory: Literal["nan", "zero"],
      detect_races: bool,
      vector_clock_size: int,
      barrier: threading.Barrier,
      clean_up_barrier: threading.Barrier,
      buffer_bounds: Literal["logical", "padded"] | None = None,
      logging_mode: params.LoggingMode | None = None,
      num_cores_per_device: int,
  ):
    clocks = {
        (device_id, core_id): self.VectorClock(vector_clock_size)
        for device_id in range(num_devices)
        for core_id in range(num_cores_per_device)
    }
    super().__init__(
        num_devices=num_devices,
        out_of_bounds_reads=out_of_bounds_reads,
        dma_execution_mode=dma_execution_mode,
        uninitialized_memory=uninitialized_memory,
        detect_races=detect_races,
        vector_clock_size=vector_clock_size,
        clocks=clocks,
        barrier=barrier,
        clean_up_barrier=clean_up_barrier,
        buffer_bounds=buffer_bounds,
        logging_mode=logging_mode,
    )
    self.num_cores_per_device = num_cores_per_device
    self.sem = {}
    self.next_semaphore_id = collections.defaultdict(lambda: 2000)
    self.fixed_id_sem = {}

  @property
  def num_cores(self) -> int:
    return self.num_devices * self.num_cores_per_device

  def get_global_core_id(self, device_id: int, local_core_id: int) -> int:
    """Computes the global core ID from the given device and local core ID."""
    return device_id * self.num_cores_per_device + local_core_id

  def get_global_core_ids(self, device_id: int) -> Sequence[int]:
    """Computes the global core IDs for all cores in the given device."""
    return tuple(
        self.get_global_core_id(device_id, core_id)
        for core_id in range(self.num_cores_per_device)
    )

  def thread_to_vc_position(self, thread: ThreadKey) -> int:
    return self.get_global_core_id(*thread)

  def global_core_id_to_thread(self, global_core_id: int) -> ThreadKey:
    device_id = global_core_id // self.num_cores_per_device
    local_core_id = global_core_id % self.num_cores_per_device
    return (device_id, local_core_id)

  def print(self, device_id: int):
    device_id = int(device_id)
    if device_id == 0:
      with self.lock:
        print(self.mem)

  def _unsafe_get_semaphore(self, sem_id: int) -> Semaphore:
    """Returns the semaphore with the given ID. `self.lock` must be held."""

    if sem_id in self.fixed_id_sem:
      if sem_id in self.sem:
        # TODO(nrink): For now we make it the responsibility of the client to
        # ensure that fixed-ID semaphores do not collide with internal
        # semaphore IDs.
        raise ValueError(
            f'Semaphore {sem_id} occurs as both fixed-id and internal.'
        )
      return self.fixed_id_sem[sem_id]
    else:
      return self.sem[sem_id]

  def append_semaphore_task(
      self,
      semaphore_id: int,
      global_core_id: int,
      task: SemaphoreTask,
  ):
    """Appends a task to be executed if the semaphore with the given sempahore ID is waiting to be signalled on the core with the given global core ID."""
    with self.lock:
      sem = self._unsafe_get_semaphore(semaphore_id)
    sem.enqueue_task(task, global_core_id)

  def get_random_virtual_device_id(self) -> int:
    # Virtual device IDs are needed for DMAs. Conceptually, each DMA runs on its
    # own, independent device. Representing this precisely would require vector
    # clocks to have sizes linear in the number of DMAs.
    #
    # Instead, we use approximate vector clocks of fixed size. We assign each
    # DMA a virtual core ID in the range
    #
    #   [num_cores, self.vector_clock_size - 1],
    #
    # and each operation of a DMA increments the corresponding coordinate in its
    # vector clock. (So the "virtual" part of a vector clock is effectively
    # counting, for each virtual core, the number of DMAs that happened-before
    # the vector clock and were assigned to that virtual core.)
    #
    # If two approximate clocks are unordered, then their corresponding events
    # are not ordered by the happens-before relation. So this approximation will
    # not introduce any false positives in detecting data races. But we may fail
    # to detect some true data races because there can be cases where two
    # approximate clocks are ordered, and we will treat the corresponding events
    # as ordered by the happens-before relation, but the corresponding events
    # are not actually ordered.
    return np.random.randint(self.num_cores, self.vector_clock_size)

  def get_semaphores_and_increment_clock(
      self, sem_ids: Sequence[int | None], global_core_id: int
  ) -> tuple[list[Semaphore | None], VectorClock | None]:
    """Returns the semaphores with the given `sem_ids` and increments the vector clock for the core with `global_core_id`.

    If race detection is enabled, this method increments the vector clock for
    the core with the given `global_core_id` (while holding the lock on `self`).
    We do this so that we can associate a (vector clock) time with the shared
    memory operation of looking up the semaphores, which in turn can be used as
    a proxy for the time when the returned semaphores are used by the client of
    the `SharedMemory` class without acquiring the lock on `self`. (For the
    purpose of encapsulation, we prefer to think of `self.lock` as a private
    attribute of the `SharedMemory` class; hence clients of the class should not
    attempt to acquire this lock explicitly.)

    Args:
      sem_ids: The IDs of the semaphores to return or None.
      global_core_id: The ID of the core whose vector clock should be
        incremented (if race detection is enabled).

    Returns:
      - The semaphores with the given `sem_ids` or None if the corresponding
        entry in `sem_ids` is None.
      - The incremented vector clock for the core with the given
        `global_core_id`, or None if race detection is not enabled.
    """
    thread = self.global_core_id_to_thread(global_core_id)
    clock = None
    with self.lock:
      if self.detect_races:
        clock = self.incr_clock(thread, take_lock=False)

      sems = []
      for sem_id in sem_ids:
        if sem_id is None:
          sem = None
        else:
          sem = self._unsafe_get_semaphore(sem_id)
        sems.append(sem)

    return sems, clock

  def get_sempahores_with_nonzero_count(
      self, device_id: int
  ) -> list[tuple[Semaphore, int]]:
    """Returns tuples (semaphore, global_core_id) for all semaphores with a nonzero count for the core with `global_core_id`."""
    result = []
    with self.lock:
      sems = self.sem.items() | self.fixed_id_sem.items()
    for _, sem in sems:
      with sem.cv:
        for gci in self.get_global_core_ids(device_id):
          if sem.count_by_core[gci] != 0:
            result.append((sem, gci))
    return result

  def allocate_semaphores(self, key: Any, num_semaphores: int) -> int:
    """Returns the next semaphore ID and ensures that the next `num_semaphores` are allocated."""
    with self.lock:
      semaphore_id = self.next_semaphore_id[key]
      self.next_semaphore_id[key] = semaphore_id + num_semaphores

      for i in range(semaphore_id, semaphore_id + num_semaphores):
        if i not in self.sem:
          self.sem[i] = Semaphore(
              shared_memory=self,
              semaphore_id=i,
              enable_logging=(
                  self.logging_mode is not None
                  and params.LoggingMode.SEMAPHORE in self.logging_mode
              ),
          )

    return semaphore_id

  def guarantee_semaphore_with_fixed_id(self, semaphore_id: int):
    """Ensures that a semaphore with the given `semaphore_id` exists.

    If the semaphore with the given ID does not exist, it is allocated. Note
    that semaphores that are allocated with this method live in their own
    address space (internally, they are mapped in a separate dictionary) from
    the sempahores allocated with the `allocate_sempahores` method above.

    This methods is intended to be used for barrier semaphores, where the
    _collective_ semaphore ID is specified by the interpreter (i.e. by the
    client of the `SharedMemory` class). This simulates sempahores that exist
    prior to any Pallas kernels being run.

    Args:
      semaphore_id: The ID of the semaphore to ensure exists, i.e. is allocated.
    """
    with self.lock:
      if semaphore_id not in self.fixed_id_sem:
        self.fixed_id_sem[semaphore_id] = Semaphore(
            semaphore_id=semaphore_id,
            shared_memory=self,
            enable_logging=(
                self.logging_mode is not None
                and params.LoggingMode.SEMAPHORE in self.logging_mode
            ),
        )

  def update_clocks_for_device_barrier(self, device_id: int):
    """Synchronizes the vector clocks for the cores on the given device."""
    threads = [(device_id, core_id) for core_id in range(self.num_cores_per_device)]
    self.update_clocks(threads)