{"record":{"id":"64ff9a072fe00108","repo":"jax-ml/jax","slug":"only-workgroup-memory-is-supported-but-got-ref","errorCode":null,"errorMessage":"Only workgroup memory is supported but got {ref}.","messagePattern":"Only workgroup memory is supported but got (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":1485,"sourceCode":"    ref: the reference to bitcast.\n    src_dtype: the source element type.\n    dst_dtype: the destination element type.\n\n  Returns:\n    A bitcasted version of `ref` with element type `dst_dtype`.\n\n  Raises:\n    ValueError: if the source ref is not in SMEM.\n  \"\"\"\n  if src_dtype == dst_dtype:\n    return ref\n  if src_dtype != ir.IntegerType.get_signless(8):\n    raise NotImplementedError(\n        \"Data type bitcast is only supported from i8 to other types.\"\n    )\n  ref_ty = ir.MemRefType(ref.type)\n  if not mgpu_utils.is_smem_ref(ref_ty):\n    raise ValueError(f\"Only workgroup memory is supported but got {ref}.\")\n  if len(ref_ty.shape) != 1:\n    raise NotImplementedError(\n        \"Data type bitcast is only supported for 1D arrays.\"\n    )\n  [stride], _ = ref_ty.get_strides_and_offset()\n  if stride != 1:\n    raise ValueError(\n        \"Data type bitcast is only supported for contiguous 1D arrays, but got \"\n        f\"stride={stride}.\"\n    )\n  [shape_bytes] = ref_ty.shape\n  shape_bitwidth = shape_bytes * 8\n  target_bitwidth = mgpu_utils.bitwidth(dst_dtype)\n\n  if shape_bitwidth % target_bitwidth:\n    raise ValueError(\n        f\"Can not bitcast memory region of size {shape_bitwidth} bits to dtype \"\n        f\"with {target_bitwidth} bits.\"","sourceCodeStart":1467,"sourceCodeEnd":1503,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L1467-L1503","documentation":"During Pallas GPU lowering, a dtype bitcast was requested on a ref that does not live in workgroup (shared) memory. _handle_dtype_bitcast checks mgpu_utils.is_smem_ref and only knows how to reinterpret SMEM memrefs; TMEM, WMEM, or global-memory refs cannot be bitcast here.","triggerScenarios":"Aliasing a Ref stored in tensor memory (tcgen05.TMEMRef) or non-SMEM memory with a different dtype, causing _extract_aliased_ref to route through the bitcast path on a non-SMEM ref.","commonSituations":"Mixing tmem (tcgen05) buffers with dtype reinterpretation on Hopper/Blackwell; kernels migrated from SMEM scratch to TMEM allocations while keeping view/alias dtype tricks.","solutions":["Move the aliased buffer to SMEM (use regular SMEM scratch instead of TMEM/WMEM)","Copy the data from the non-SMEM ref into an SMEM i8 buffer, then bitcast that buffer","Convert values explicitly with jax.lax.bitcast_convert after loading"],"exampleFix":"// before\nbuf = alloc_tmem(...)  # or non-SMEM ref\nview = buf.view(jnp.float32)\n// after\nbuf = alloc_smem(...)  # SMEM-backed i8 scratch\nview = buf.view(jnp.float32)","handlingStrategy":"validation","validationCode":"# only alias refs allocated in SMEM\nassert buf.memory_space == 'smem', 'bitcast alias requires an SMEM buffer'","typeGuard":"def is_smem_backed(buf) -> bool:\n    return getattr(buf, 'memory_space', 'smem') == 'smem'","tryCatchPattern":null,"preventionTips":["Keep reinterpret casts on SMEM scratch only","Document which kernel buffers must be SMEM","Avoid mixing TMEM accumulators with dtype views"],"tags":["jax","pallas","mosaic-gpu","smem","memory-space","bitcast"],"backgroundTag":"unsupported-memory-space-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}