jax-ml/jax · error · ValueError

Expected SMEM but got: {ref_ty.memory_space}

Error message

Expected SMEM but got: {ref_ty.memory_space}

What it means

Raised after a cluster-ptr memref cast (utils.py:2406) when the result type's memory space is not SMEM. The helper rewrites a ref so it points at another CTA's shared memory in the cluster; only SMEM refs can be remapped this way.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2406

    ref: ir.Value, dim: gpu.Dimension, idx: ir.Value, generic: bool = True
):
  i32 = ir.IntegerType.get_signless(32)
  # We replace the offset in the ref type by 0, because memref_ptr always
  # folds the offset into the pointer.
  ref_ty = ir.MemRefType(ref.type)
  strides, offset = ref_ty.get_strides_and_offset()
  if offset != 0:
    new_layout = ir.StridedLayoutAttr.get(0, strides)
  else:
    new_layout = ref_ty.layout
  result_type = ir.MemRefType.get(
      ref_ty.shape,
      ref_ty.element_type,
      new_layout,
      None if generic else ir.IntegerAttr.get(i32, 7),
  )
  if not is_smem_ref(ref_ty):
    raise ValueError(f"Expected SMEM but got: {ref_ty.memory_space}")
  idxs: list[ir.Value] = [gpu.cluster_block_id(d) for d in gpu.Dimension]
  idxs[dim] = idx
  flat_block = arith.index_cast(i32, cluster_idx(dim_idx=idxs))
  return ptr_as_memref(
      get_cluster_ptr(memref_ptr(ref), flat_block, generic), result_type
  )


def elements_to_bytes(offset: ir.Value, element_bitwidth: int) -> ir.Value:
  """Convert an element-based linear offset to a byte-based offset."""
  index_ty = offset.type

  if element_bitwidth > 8:
    return arith.muli(offset, c(element_bitwidth // 8, index_ty))
  elif element_bitwidth < 8:
    return arith.divsi(offset, c(8 // element_bitwidth, index_ty))
  else:
    return offset

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the buffer in shared memory (smem) rather than global memory before remapping
2. Check ref.memory_space with utils.is_smem_ref before the call
3. If a generic pointer is intentional, pass generic=True so the 7 address-space attr is not forced, but verify semantics

Example fix

# before
buf = memref.alloc(...)
ptr = utils.get_cluster_ptr_ptr_as_memref(...)  # global-space ref
# after
buf = smem_alloc(shape, dtype)  # #shared memory space
result = remap_cluster_ref(buf, ...)
assert utils.is_smem_ref(result.type)

assert utils.is_smem_ref(ref), f'ref must be SMEM, got {ref.memory_space}'

def is_smem_ref(v): return isinstance(v, ir.Value) and isinstance(v.type, ir.MemRefType) and v.type.memory_space is not None

When it happens

Trigger: Calling the get_cluster_ptr remapping path (typically via utils helpers used to address neighboring CTA SMEM) with a ref whose memory space is global or generic (non-SMEM).

Common situations: Passing a global-memory buffer where a cluster-shared buffer was expected; allocating with the generic address space instead of #shared memory space; TPOT-style collective code wired to the wrong buffer.

Related errors

• Only workgroup memory is supported but got {ref}.
• packed, collective and layout arguments are only supported f
• Only SMEM and TMEM refs are supported.
• All aliased Refs must have the same memory space (SMEM or TM
• Only byte-aligned bitcasts are supported.

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