jax-ml/jax · error · ValueError

tmem_addr_ref must be in shared memory, got: {addr_ref_ty}

Error message

tmem_addr_ref must be in shared memory, got: {addr_ref_ty}

What it means

TMEM allocation addresses are communicated through a single-element i32 memref that must live in shared memory (address space 3) because the tcgen05.alloc result is written by the hardware/CTA into smem. from_alloc checks utils.is_smem_ref and rejects global or generic address-space memrefs.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1223

  def __post_init__(self):
    self.layout.check_type(self.shape, utils.bitwidth(self.dtype))

  @classmethod
  def from_alloc(
      cls,
      tmem_addr_ref: ir.Value,
      shape: tuple[int, int],
      dtype,
      collective: bool | None = None,
      layout: TMEMLayout | None = None,
  ) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(tmem_addr_ref.type, ir.MemRefType):
      raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)
    if not utils.is_smem_ref(addr_ref_ty):
      raise ValueError(f"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}")
    if addr_ref_ty.element_type != i32:
      raise ValueError(f"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}")
    if math.prod(addr_ref_ty.shape) != 1:
      raise ValueError(f"tmem_addr_ref must contain a single element, got: {addr_ref_ty}")
    i0 = arith.ConstantOp.create_index(0)
    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)
    if shape[0] < 32:
      raise ValueError(f"TMEM refs must have at least 32 rows, got: {shape[0]}")
    if layout is None:
      if collective is None:
        raise ValueError(
            "collective argument must be provided when TMEM layout is inferred"
        )
      layout = _infer_tmem_layout(shape, collective, packing=1)
    # TODO: Do we have to do this??
    # warp_idx = utils.warp_idx(sync=False)
    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))
    return cls(tmem_addr, shape, dtype, layout)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate tmem_addr_ref in shared memory (use the Mosaic smem allocation utilities)
2. Verify the memref type string contains the workgroup/shared address space attribute
3. Reproduce the pattern from jax mosaic gpu tests/examples for tmem.alloc

Example fix

# before
addr_ref = memref.alloca(..., memref_type_in_global_space)
# after
addr_ref = smem_alloca(i32, [])  # shared-memory 0-d memref, then pass to from_alloc

from jax._src.experimental.mosaic.gpu import utils
assert utils.is_smem_ref(ir.MemRefType(tmem_addr_ref.type)), 'need smem address space'

def addr_ref_ok(ref) -> bool:
    return (isinstance(ref.type, ir.MemRefType)
            and utils.is_smem_ref(ir.MemRefType(ref.type)))

When it happens

Trigger: Passing a global-memory memref or a memref without the shared-memory address space attribute as tmem_addr_ref to from_alloc.

Common situations: Creating the address memref manually with default (global) address space instead of via smem helpers; refactoring that stripped the #gpu.address_space<workgroup> attribute.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Logical shape {self.logical_shape} cannot be bigger than con
• Attempting to deallocate allocation with key `{key}` that is
• Attempting to get contents of allocation with key `{key}` th
• Attempting to store into allocation with key `{key}` that is
• Attempting to swap into allocation with `key` {key} that is

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/db53512c3814a48d. Report an issue: GitHub.