jax-ml/jax · error · ValueError

Pointer address space {ptr_ty.address_space} does not match

Error message

Pointer address space {ptr_ty.address_space} does not match memref memory space {memref_ty.memory_space}.

What it means

Raised by ptr_as_memref, which reinterprets a raw LLVM pointer as a memref descriptor. It first verifies the pointer's LLVM address space equals the memref type's expected address space (Global->1, Workgroup->3, default->0). A mismatch means you're casting e.g. a global pointer to a workgroup memref type, which would silently produce invalid IR, so it's rejected.

Source

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

def gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:
  match address_space:
    case gpu.AddressSpace.Global:
      return 1
    case gpu.AddressSpace.Workgroup:
      return 3
    case _:
      raise NotImplementedError(f"address_space not supported: {address_space}")


WORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(
    gpu.AddressSpace.Workgroup
)


def ptr_as_memref(ptr, memref_ty: ir.MemRefType):
  ptr_ty = llvm.PointerType(ptr.type)
  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):
    raise ValueError(
        f"Pointer address space {ptr_ty.address_space} does not match "
        f"memref memory space {memref_ty.memory_space}."
    )

  strides, offset = memref_ty.get_strides_and_offset()
  if offset != 0:
    raise ValueError("Non-zero offset is not supported for ptr_as_memref")
  i64 = ir.IntegerType.get_signless(64)
  rank = len(memref_ty.shape)
  desc_ty_fields = [ptr_ty, ptr_ty, i64]
  if rank > 0:
    desc_ty_fields += [llvm.ArrayType.get(i64, rank)] * 2
  desc_ty = llvm.StructType.get_literal(desc_ty_fields)
  desc = llvm.UndefOp(desc_ty).result
  desc = llvm.InsertValueOp(desc, ptr, [0]).result  # Allocation
  desc = llvm.InsertValueOp(desc, ptr, [1]).result  # Aligned Base
  desc = llvm.InsertValueOp(
      desc, llvm.ConstantOp(i64, ir.IntegerAttr.get(i64, 0)).result, [2]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the memref type's memory_space match the pointer's address space (e.g. use a workgroup memref for a workgroup pointer)
2. Re-derive the pointer in the correct address space (addrspacecast or re-allocate) before the cast
3. When doing dtype bitcasts, keep the address space unchanged and pass a memref_ty with the same space

Example fix

# before
ptr = llvm_ops.mlir_addresscast(wg_ptr, llvm.PointerType.get(f32, 0))
mr = utils.ptr_as_memref(ptr, wg_memref_ty)  # space 3
# after
ptr_ty = llvm.PointerType.get(f32, address_space=3)
ptr = llvm_ops.mlir_addresscast(wg_ptr, ptr_ty)
mr = utils.ptr_as_memref(ptr, wg_memref_ty)

ptr_space = llvm.PointerType(ptr.type).address_space
expected = utils.get_memref_llvm_address_space(memref_ty) or 0
assert ptr_space == expected, f'pointer space {ptr_space} != memref space {expected}'

When it happens

Trigger: Calling utils.ptr_as_memref(ptr, memref_ty) where ptr lives in address space 0/1 but memref_ty is a workgroup memref (space 3), or vice versa — e.g. after _handle_dtype_bitcast or in distributed ops (to_remote, get_cluster_ref) that retype pointers without changing address space.

Common situations: Bitcasting a pointer's element type across memory spaces; writing custom DSMEM (distributed shared memory) code that mixes cluster pointers with local workgroup pointers; version changes in how Mosaic types workgroup memory.

Related errors

• Non-zero offset is not supported for ptr_as_memref
• Unsupported memory space: {orig_ref_ty.memory_space}
• address_space not supported: {address_space}
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}

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