{"record":{"id":"778e475296586f7c","repo":"jax-ml/jax","slug":"non-zero-offset-is-not-supported-for-ptr-as-memref","errorCode":null,"errorMessage":"Non-zero offset is not supported for ptr_as_memref","messagePattern":"Non-zero offset is not supported for ptr_as_memref","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":101,"sourceCode":"      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]\n  ).result\n  if rank > 0:\n    for i, s in enumerate(memref_ty.shape):\n      desc = llvm.InsertValueOp(\n          desc, llvm.ConstantOp(i64, ir.IntegerAttr.get(i64, s)).result, [3, i]\n      ).result\n    for i, s in enumerate(strides):","sourceCodeStart":83,"sourceCodeEnd":119,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L83-L119","documentation":"Raised by ptr_as_memref when the target memref type has a non-zero symbolic/constant offset. The helper builds a memref descriptor whose offset field is implicitly 0 (it only fills the pointer), so a memref type like memref<4x4xf32, offset:8, ...> cannot be faithfully represented and is rejected.","triggerScenarios":"Calling utils.ptr_as_memref(ptr, memref_ty) where memref_ty was created with an explicit offset (ir.MemRefType.get(..., memory_space=...) on a layout map with offset, or a layout from get_strides_and_offset returning offset != 0), e.g. a subview type or affine layout with non-zero origin.","commonSituations":"Reusing memref types derived from subviews (which carry offsets) as the target of a pointer cast; constructing layouts with shifted affine maps for tiled buffers.","solutions":["Use an offset-0 memref type: apply the offset to the pointer (getelementptr) instead of encoding it in the type","Strip the layout map or rebuild the type with ir.MemRefType.get(shape, elem_ty) (implicit row-major, offset 0)","If an offset view is needed, create it via memref.subview on the resulting memref, not in the target type"],"exampleFix":"# before\nty = ir.MemRefType.get((4,4), f32, layout=offset_layout)  # offset 8\nmr = utils.ptr_as_memref(ptr, ty)\n# after\nptr = utils.getelementptr(ptr, [8], f32)\nmr = utils.ptr_as_memref(ptr, ir.MemRefType.get((4,4), f32))","handlingStrategy":"validation","validationCode":"_, offset = memref_ty.get_strides_and_offset()\nassert offset == 0, 'use an offset-0 memref type; fold the offset into the pointer instead'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never reuse subview-derived memref types as ptr_as_memref targets","Apply offsets via getelementptr on the pointer, not via the type's layout map"],"tags":["gpu","mosaic","memref","offset","llvm"],"backgroundTag":"nonzero-memref-offset","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}