{"record":{"id":"906e13d6ca009617","repo":"jax-ml/jax","slug":"tmem-addr-ref-must-be-an-i32-memref-got-addr-re","errorCode":null,"errorMessage":"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}","messagePattern":"tmem_addr_ref must be an i32 memref, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1225,"sourceCode":"    self.layout.check_type(self.shape, utils.bitwidth(self.dtype))\n\n  @classmethod\n  def from_alloc(\n      cls,\n      tmem_addr_ref: ir.Value,\n      shape: tuple[int, int],\n      dtype,\n      collective: bool | None = None,\n      layout: TMEMLayout | None = None,\n  ) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    if not isinstance(tmem_addr_ref.type, ir.MemRefType):\n      raise ValueError(f\"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}\")\n    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)\n    if not utils.is_smem_ref(addr_ref_ty):\n      raise ValueError(f\"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}\")\n    if addr_ref_ty.element_type != i32:\n      raise ValueError(f\"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}\")\n    if math.prod(addr_ref_ty.shape) != 1:\n      raise ValueError(f\"tmem_addr_ref must contain a single element, got: {addr_ref_ty}\")\n    i0 = arith.ConstantOp.create_index(0)\n    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)\n    if shape[0] < 32:\n      raise ValueError(f\"TMEM refs must have at least 32 rows, got: {shape[0]}\")\n    if layout is None:\n      if collective is None:\n        raise ValueError(\n            \"collective argument must be provided when TMEM layout is inferred\"\n        )\n      layout = _infer_tmem_layout(shape, collective, packing=1)\n    # TODO: Do we have to do this??\n    # warp_idx = utils.warp_idx(sync=False)\n    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))\n    return cls(tmem_addr, shape, dtype, layout)\n\n  def slice(self, *idxs) -> TMEMRef:","sourceCodeStart":1207,"sourceCodeEnd":1243,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1207-L1243","documentation":"TMEMRef.from_alloc loads the TMEM base address via memref.load from the provided smem memref, which requires the element type to be exactly signless i32 — the tcgen05.alloc instruction returns a 32-bit address. Any other element type (i64, f32, i16, signed variants) is rejected.","triggerScenarios":"Passing an smem memref of i64/f32/i16 to from_alloc; using a signful integer type (e.g. si32) instead of signless i32.","commonSituations":"Building the address buffer with a helper defaulting to a different integer width; older snippets that used 64-bit address arithmetic.","solutions":["Make the address memref signless i32: ir.MemRefType.get([], i32) with smem address space","Cast: store an i32 truncation of your address value into the i32 smem memref","Use the library's tmem allocation helpers which produce the correct type"],"exampleFix":"# before\naddr_ref = smem_alloca(i64, [])\n# after\ni32 = ir.IntegerType.get_signless(32)\naddr_ref = smem_alloca(i32, [])","handlingStrategy":"type-guard","validationCode":"i32 = ir.IntegerType.get_signless(32)\nassert ir.MemRefType(tmem_addr_ref.type).element_type == i32","typeGuard":"def is_signless_i32_memref(ref) -> bool:\n    t = getattr(ref, 'type', None)\n    return isinstance(t, ir.MemRefType) and t.element_type == ir.IntegerType.get_signless(32)","tryCatchPattern":null,"preventionTips":["Use ir.IntegerType.get_signless(32); never assume default integer width","Copy the alloc boilerplate from jax's tcgen05 tests verbatim"],"tags":["jax","mosaic","tmem","dtype","memref"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}