{"record":{"id":"11e54888ae717b04","repo":"jax-ml/jax","slug":"expected-smem-but-got-ref-ty-memory-space","errorCode":null,"errorMessage":"Expected SMEM but got: {ref_ty.memory_space}","messagePattern":"Expected SMEM but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":2406,"sourceCode":"    ref: ir.Value, dim: gpu.Dimension, idx: ir.Value, generic: bool = True\n):\n  i32 = ir.IntegerType.get_signless(32)\n  # We replace the offset in the ref type by 0, because memref_ptr always\n  # folds the offset into the pointer.\n  ref_ty = ir.MemRefType(ref.type)\n  strides, offset = ref_ty.get_strides_and_offset()\n  if offset != 0:\n    new_layout = ir.StridedLayoutAttr.get(0, strides)\n  else:\n    new_layout = ref_ty.layout\n  result_type = ir.MemRefType.get(\n      ref_ty.shape,\n      ref_ty.element_type,\n      new_layout,\n      None if generic else ir.IntegerAttr.get(i32, 7),\n  )\n  if not is_smem_ref(ref_ty):\n    raise ValueError(f\"Expected SMEM but got: {ref_ty.memory_space}\")\n  idxs: list[ir.Value] = [gpu.cluster_block_id(d) for d in gpu.Dimension]\n  idxs[dim] = idx\n  flat_block = arith.index_cast(i32, cluster_idx(dim_idx=idxs))\n  return ptr_as_memref(\n      get_cluster_ptr(memref_ptr(ref), flat_block, generic), result_type\n  )\n\n\ndef elements_to_bytes(offset: ir.Value, element_bitwidth: int) -> ir.Value:\n  \"\"\"Convert an element-based linear offset to a byte-based offset.\"\"\"\n  index_ty = offset.type\n\n  if element_bitwidth > 8:\n    return arith.muli(offset, c(element_bitwidth // 8, index_ty))\n  elif element_bitwidth < 8:\n    return arith.divsi(offset, c(8 // element_bitwidth, index_ty))\n  else:\n    return offset","sourceCodeStart":2388,"sourceCodeEnd":2424,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L2388-L2424","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Allocate the buffer in shared memory (smem) rather than global memory before remapping","Check ref.memory_space with utils.is_smem_ref before the call","If a generic pointer is intentional, pass generic=True so the 7 address-space attr is not forced, but verify semantics"],"exampleFix":"# before\nbuf = memref.alloc(...)\nptr = utils.get_cluster_ptr_ptr_as_memref(...)  # global-space ref\n# after\nbuf = smem_alloc(shape, dtype)  # #shared memory space\nresult = remap_cluster_ref(buf, ...)\nassert utils.is_smem_ref(result.type)","handlingStrategy":"validation","validationCode":"assert utils.is_smem_ref(ref), f'ref must be SMEM, got {ref.memory_space}'","typeGuard":"def is_smem_ref(v): return isinstance(v, ir.Value) and isinstance(v.type, ir.MemRefType) and v.type.memory_space is not None","tryCatchPattern":null,"preventionTips":["Allocate cluster buffers with the SMEM helper, not global alloc","Assert is_smem_ref before cluster ptr remapping"],"tags":["jax","mosaic-gpu","smem","memory-space","gpu-clusters"],"backgroundTag":"gpu-memory-space-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}