{"record":{"id":"f92d0a7b5aa32221","repo":"jax-ml/jax","slug":"unsupported-register-bitwidth-reg-bitwidth","errorCode":null,"errorMessage":"Unsupported register bitwidth: {reg_bitwidth}","messagePattern":"Unsupported register bitwidth: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3850,"sourceCode":"      reg_ty = ir.VectorType(reg.type)\n      element_bitwidth = utils.bitwidth(reg_ty.element_type)\n      if (\n          isinstance(reg_ty.element_type, ir.FloatType)\n          and element_bitwidth <= 8\n      ):\n        narrow_int = ir.IntegerType.get_signless(element_bitwidth)\n        reg = vector.bitcast(ir.VectorType.get(reg_ty.shape, narrow_int), reg)\n      reg_bitwidth = utils.bitwidth(reg_ty)\n      if reg_bitwidth == 32:\n        ptx_constraint = \"r\"\n        ptx_type = \"b32\"\n        reg = utils.bitcast(reg, i32)\n      elif reg_bitwidth == 64:\n        ptx_constraint = \"l\"\n        ptx_type = \"b64\"\n        reg = utils.bitcast(reg, i64)\n      else:\n        raise NotImplementedError(f\"Unsupported register bitwidth: {reg_bitwidth}\")\n      llvm.inline_asm(\n          ir.Type.parse(\"!llvm.void\"),\n          [cluster_ptr, reg, cluster_barrier_ptr],\n          f\"st.async.cluster.shared::cluster.mbarrier::complete_tx::bytes.{ptx_type} [$0], $1, [$2];\",\n          f\"l,{ptx_constraint},l\",\n          has_side_effects=True,\n      )\n\n  def _store_register_atomic(\n      self,\n      base_ptr: ir.Value,\n      vreg: ir.Value,\n      atomic: Literal[\"add\", \"max\", \"min\", \"and\", \"or\", \"xor\"],\n      is_smem: bool,\n      multimem: bool = False,\n      cluster_barrier_ptr: ir.Value | None = None,\n  ):\n    i32 = ir.IntegerType.get_signless(32)","sourceCodeStart":3832,"sourceCodeEnd":3868,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3832-L3868","documentation":"To emit the inline st.async PTX instruction, store_tiled_async must bitcast each register to a 32- or 64-bit PTX register type (b32 with 'r' constraint or b64 with 'l'); any other register bitwidth (e.g. 8- or 128-bit registers) has no corresponding PTX encoding so it raises NotImplementedError.","triggerScenarios":"Storing values whose fragment registers are not 32 or 64 bits wide — e.g. i8/f16 vector fragments that were not vectorized to 32-bit granularity, or unusual wide types.","commonSituations":"Async-storing sub-32-bit element types without ensuring vec_len pairs them into 32-bit registers; layouts with odd vec_size producing non-standard register widths.","solutions":["Ensure sub-32-bit elements are packed to at least 32-bit registers (vec_len a multiple of 2 for 16-bit, 4 for 8-bit) via the layout's vec_size","Convert the array to a layout whose registers are 32/64-bit wide before store_tiled_async","Use the non-async store path for exotic widths"],"exampleFix":"// before\nfa8.store_tiled_async(ref)  # i8 regs < 32 bits\n// after\nfa8 = fa8.to_layout(TiledLayout(vec_size=4))\nfa8.store_tiled_async(ref)","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import utils\nreg_bitwidth = utils.bitwidth(fa.mlir_dtype) * fa.layout.vec_size  # approximate reg width\nif reg_bitwidth not in (32, 64):\n    fa = fa.to_layout(TiledLayout(vec_size=max(1, 32 // utils.bitwidth(fa.mlir_dtype))))","typeGuard":null,"tryCatchPattern":"try:\n    fa.store_tiled_async(ref, ...)\nexcept NotImplementedError as e:\n    if 'register bitwidth' in str(e).lower():\n        fa.store_tiled(ref)  # sync fallback\n    else:\n        raise","preventionTips":["Choose vec_size so elements pack into 32/64-bit registers","Prefer 32-bit-friendly element types for async stores"],"tags":["jax","mosaic","gpu","ptx","register-width","not-implemented"],"backgroundTag":"unsupported-datatype-width","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}