jax-ml/jax · error · ValueError

Non-zero offset is not supported for ptr_as_memref

Error message

Non-zero offset is not supported for ptr_as_memref

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:101

      raise NotImplementedError(f"address_space not supported: {address_space}")


WORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(
    gpu.AddressSpace.Workgroup
)


def ptr_as_memref(ptr, memref_ty: ir.MemRefType):
  ptr_ty = llvm.PointerType(ptr.type)
  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):
    raise ValueError(
        f"Pointer address space {ptr_ty.address_space} does not match "
        f"memref memory space {memref_ty.memory_space}."
    )

  strides, offset = memref_ty.get_strides_and_offset()
  if offset != 0:
    raise ValueError("Non-zero offset is not supported for ptr_as_memref")
  i64 = ir.IntegerType.get_signless(64)
  rank = len(memref_ty.shape)
  desc_ty_fields = [ptr_ty, ptr_ty, i64]
  if rank > 0:
    desc_ty_fields += [llvm.ArrayType.get(i64, rank)] * 2
  desc_ty = llvm.StructType.get_literal(desc_ty_fields)
  desc = llvm.UndefOp(desc_ty).result
  desc = llvm.InsertValueOp(desc, ptr, [0]).result  # Allocation
  desc = llvm.InsertValueOp(desc, ptr, [1]).result  # Aligned Base
  desc = llvm.InsertValueOp(
      desc, llvm.ConstantOp(i64, ir.IntegerAttr.get(i64, 0)).result, [2]
  ).result
  if rank > 0:
    for i, s in enumerate(memref_ty.shape):
      desc = llvm.InsertValueOp(
          desc, llvm.ConstantOp(i64, ir.IntegerAttr.get(i64, s)).result, [3, i]
      ).result
    for i, s in enumerate(strides):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use an offset-0 memref type: apply the offset to the pointer (getelementptr) instead of encoding it in the type
2. Strip the layout map or rebuild the type with ir.MemRefType.get(shape, elem_ty) (implicit row-major, offset 0)
3. If an offset view is needed, create it via memref.subview on the resulting memref, not in the target type

Example fix

# before
ty = ir.MemRefType.get((4,4), f32, layout=offset_layout)  # offset 8
mr = utils.ptr_as_memref(ptr, ty)
# after
ptr = utils.getelementptr(ptr, [8], f32)
mr = utils.ptr_as_memref(ptr, ir.MemRefType.get((4,4), f32))

_, offset = memref_ty.get_strides_and_offset()
assert offset == 0, 'use an offset-0 memref type; fold the offset into the pointer instead'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Pointer address space {ptr_ty.address_space} does not match
• Unsupported memory space: {orig_ref_ty.memory_space}
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Unsupported memory space: {ref_type.memory_space}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/778e475296586f7c. Report an issue: GitHub.