jax-ml/jax · error · NotImplementedError

The base ref for aliases must come from a slice_smem op.

Error message

The base ref for aliases must come from a slice_smem op.

What it means

In Warpgroup lowering semantics, an aliased SMEM ref's base must originate from a slice_smem op so the compiler can read its static base offset. If the ref is owned by another op (allocation, cast, layout op), the compiler cannot compute the alias offset and raises NotImplementedError; the source comment lists the memref-pointer plumbing needed to lift this.

Source

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

        assert isinstance(ref, ir.Value)  # make pyrefly happy
        input_ref_ty = ir.MemRefType(ref.type)
        if input_ref_ty.memory_space == mgpu_utils.smem():
          assert layout is None
          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(
              mlir_dtype
          )
          if ref_bits % 8:
            raise NotImplementedError("Only byte-aligned bitcasts are supported.")
          assert offset % gpu_core.SMEM_ALIGNMENT == 0

          if lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
            if not isinstance(ref.owner, mgpu.dialect.SliceSMEMOp):
              # This restriction can be lifted by:
              # - Using memref ops to get the pointer and offset of the base ref.
              # - Subtracting gpu_dialect.dynamic_shared_memory() from those to
              #   get the base offset relative to the beginning of SMEM.
              # - Implementing layout and lowering rules for all ops above.
              raise NotImplementedError(
                  "The base ref for aliases must come from a slice_smem op."
              )

            base_offset = ref.owner.offset.value
            total_offset = base_offset + offset

            ref_ty = ir.MemRefType.get(
                transformed_shape, mlir_dtype, memory_space=mgpu_utils.smem()
            )
            assert ref.owner.alias_id is not None
            alloc_id = ref.owner.alias_id.value
            # TODO(bchetioui): Use a scheme resilient to hash collisions.
            alias_id = hash((offset, alloc_id, alias_group_idx))
            # The composite key formed of `(offset, alloc_id, alias_group_idx)`
            # is a unique identifier across:
            #   - different RefUnions (different `alloc_id`, since two
            #     distinct RefUnions represent two SMEM allocations);
            #   - different ref_groups within a RefUnion (different

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure so the aliased base goes through a slice of the SMEM allocation (slice first, then alias)
2. Use Thread semantics for that part of the kernel or avoid aliasing the non-sliced ref
3. Upgrade JAX — support for additional base ops may be added via the documented memref-pointer approach
4. File an upstream issue with a minimal repro if your pattern is legitimate

Example fix

# before
base = smem_alloc(...)        # owner is alloc op
aliased = base.view(dtype)     # warpgroup: owner not SliceSMEMOp -> error
# after
base = smem_alloc(...)[0:n]   # owner is SliceSMEMOp
aliased = base.view(dtype)

# ensure the aliased base is produced by an SMEM slice
assert is_slice_smem_owner(ref), 'warpgroup alias base must come from slice_smem'

def is_slice_smem_owner(ref) -> bool:
    owner = getattr(ref, 'owner', None)
    return owner is not None and type(owner).__name__ == 'SliceSMEMOp'

try:
    aliased = base.view(dtype)
except NotImplementedError:
    base = base[0:n]  # route through a slice
    aliased = base.view(dtype)

When it happens

Trigger: Using aliased Refs under LoweringSemantics.Warpgroup where the base SMEM ref was produced by something other than mgpu.dialect.SliceSMEMOp — e.g. directly aliasing an smem alloc or a layout-cast result.

Common situations: Advanced kernels combining warpgroup MMA with dtype-view aliases; JAX version differences in how SMEM refs are materialized before aliasing.

Related errors

• All aliased Refs must have the same memory space (SMEM or TM
• Unsupported memory space.
• packed, collective and layout arguments are only supported f
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.

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