jax-ml/jax · error · NotImplementedError

Unsupported TMEM ref

Error message

Unsupported TMEM ref {ref}.

What it means

For TMEM aliases, the lowering can only recover the base offset if the ref's owner chain is a SliceTmemOp (optionally behind a TmemLayoutCastOp). Any other op producing the TMEM ref (arbitrary layout casts, allocs, other dialect ops) makes offset computation impossible and NotImplementedError is raised.

Solutions

1. Ensure the aliased TMEM ref comes directly from a tmem slice op (slice before any other transform)
2. Apply the alias view earlier in the op chain, immediately after slicing
3. Simplify/reorder TMEM layout casts so a slice_tmem is the immediate (or one-hop) owner
4. Upgrade JAX / report upstream with a minimal kernel repro

Example fix

# before
r = tmem_alloc(...)            # owner not a slice op
v = r.view(dtype)
# after
r = tmem_alloc(...)
r = r[0:n]                     # slice_tmem owner
v = r.view(dtype)

owner = getattr(ref, 'owner', None)
ok = type(owner).__name__ == 'SliceTmemOp' or (
     type(owner).__name__ == 'TmemLayoutCastOp' and
     type(getattr(owner.operands[0], 'owner', None)).__name__ == 'SliceTmemOp')
assert ok, 'TMEM alias must come from slice_tmem (optionally behind tmem_layout_cast)'

def tmem_ref_aliasable(ref) -> bool:
    o = getattr(ref, 'owner', None)
    if type(o).__name__ == 'SliceTmemOp':
        return True
    if type(o).__name__ == 'TmemLayoutCastOp':
        return type(getattr(o.operands[0], 'owner', None)).__name__ == 'SliceTmemOp'
    return False

When it happens

Trigger: Aliasing a TMEM ref whose owner is not slice_tmem / tmem_layout_cast(slice_tmem) — e.g. aliasing a raw TMEM allocation or the result of another TMEM-transforming op.

Common situations: Blackwell tcgen05 kernels chaining multiple TMEM layout operations before an aliased view; evolving JAX versions adding new TMEM ops the alias path doesn't recognize.

Related errors

• TMEM aliasing only supported for Refs with the same first…
• a_scale must be a TMEM Ref
• Accumulator must be a TMEM Ref.
• Accumulator Ref must be collective if collective_axis is…
• All aliased Refs must have the same memory space (SMEM or…

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

Appendix: source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:1648

          else:
            ref_bytes = ref_bits // 8
            ref = mgpu.memref_slice(ref, slice(offset, offset + ref_bytes))
            ref = _handle_dtype_bitcast(
                ref,
                ir.MemRefType(ref.type).element_type,
                mlir_dtype,
            )
            ref = mgpu.memref_reshape(ref, transformed_shape)
        elif input_ref_ty.memory_space == mgpu_utils.tmem():

          if isinstance(ref.owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner
          elif isinstance(
              ref.owner, mgpu.dialect.TmemLayoutCastOp
          ) and isinstance(ref.owner.operands[0].owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner.operands[0].owner
          else:
            raise NotImplementedError(f"Unsupported TMEM ref {ref}.")

          base_offset = source_slice_op.offset.value
          assert isinstance(base_offset, int)  # make pyrefly happy
          total_offset = base_offset + offset
          ref_ty = ir.MemRefType.get(
              transformed_shape, mlir_dtype, memory_space=mgpu_utils.tmem()
          )
          alloc_id = source_slice_op.alias_id
          assert alloc_id is not None
          # TODO(bchetioui): Use a scheme resilient to hash collisions.
          alias_id = hash((offset, alloc_id.value, alias_group_idx))
          slice_op = mgpu.dialect.SliceTmemOp(
              ref_ty, source_slice_op.source, total_offset
          )
          i64 = ir.IntegerType.get_signless(64)
          slice_op.attributes["alias_id"] = ir.IntegerAttr.get(i64, alias_id)
          ref = slice_op.result
          assert layout is not None

View on GitHub (pinned to 1e1c6a8fc0)