{"record":{"id":"2337b73f3703d80f","repo":"jax-ml/jax","slug":"pointer-address-space-ptr-ty-address-space-does","errorCode":null,"errorMessage":"Pointer address space {ptr_ty.address_space} does not match memref memory space {memref_ty.memory_space}.","messagePattern":"Pointer address space (.+?) does not match memref memory space (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":94,"sourceCode":"def gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:\n  match address_space:\n    case gpu.AddressSpace.Global:\n      return 1\n    case gpu.AddressSpace.Workgroup:\n      return 3\n    case _:\n      raise NotImplementedError(f\"address_space not supported: {address_space}\")\n\n\nWORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(\n    gpu.AddressSpace.Workgroup\n)\n\n\ndef ptr_as_memref(ptr, memref_ty: ir.MemRefType):\n  ptr_ty = llvm.PointerType(ptr.type)\n  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):\n    raise ValueError(\n        f\"Pointer address space {ptr_ty.address_space} does not match \"\n        f\"memref memory space {memref_ty.memory_space}.\"\n    )\n\n  strides, offset = memref_ty.get_strides_and_offset()\n  if offset != 0:\n    raise ValueError(\"Non-zero offset is not supported for ptr_as_memref\")\n  i64 = ir.IntegerType.get_signless(64)\n  rank = len(memref_ty.shape)\n  desc_ty_fields = [ptr_ty, ptr_ty, i64]\n  if rank > 0:\n    desc_ty_fields += [llvm.ArrayType.get(i64, rank)] * 2\n  desc_ty = llvm.StructType.get_literal(desc_ty_fields)\n  desc = llvm.UndefOp(desc_ty).result\n  desc = llvm.InsertValueOp(desc, ptr, [0]).result  # Allocation\n  desc = llvm.InsertValueOp(desc, ptr, [1]).result  # Aligned Base\n  desc = llvm.InsertValueOp(\n      desc, llvm.ConstantOp(i64, ir.IntegerAttr.get(i64, 0)).result, [2]","sourceCodeStart":76,"sourceCodeEnd":112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L76-L112","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make the memref type's memory_space match the pointer's address space (e.g. use a workgroup memref for a workgroup pointer)","Re-derive the pointer in the correct address space (addrspacecast or re-allocate) before the cast","When doing dtype bitcasts, keep the address space unchanged and pass a memref_ty with the same space"],"exampleFix":"# before\nptr = llvm_ops.mlir_addresscast(wg_ptr, llvm.PointerType.get(f32, 0))\nmr = utils.ptr_as_memref(ptr, wg_memref_ty)  # space 3\n# after\nptr_ty = llvm.PointerType.get(f32, address_space=3)\nptr = llvm_ops.mlir_addresscast(wg_ptr, ptr_ty)\nmr = utils.ptr_as_memref(ptr, wg_memref_ty)","handlingStrategy":"validation","validationCode":"ptr_space = llvm.PointerType(ptr.type).address_space\nexpected = utils.get_memref_llvm_address_space(memref_ty) or 0\nassert ptr_space == expected, f'pointer space {ptr_space} != memref space {expected}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep pointer address space and memref memory_space in one config tuple so they can't drift","After any bitcast of a pointer, verify its address space before ptr_as_memref"],"tags":["gpu","mosaic","address-space","pointer","memref","llvm"],"backgroundTag":"address-space-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}