jax-ml/jax · error · NotImplementedError

Expected TiledLayout, got {type(layout)}

Error message

Expected TiledLayout, got {type(layout)}

What it means

mgpu.async_store_smem requires the value's layout to be a TiledLayout; fragment-based layouts cannot be written to shared memory by this op.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:601

    raise ValueError(f"Unsupported memory space: {ref_type.memory_space}")

  if ctx.auto_barriers:
    utils.warpgroup_barrier()  # Make sure the writes have completed.

  return []


@_register_lowering(mgpu.AsyncStoreSmemOp)
def _async_store_smem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreSmemOp
) -> Sequence[ir.Value]:
  index = ir.IndexType.get()

  [to_store_layout] = inference_utils.in_layouts(op)
  value = _fragmented_array_from_ir(op.valueToStore, to_store_layout)
  layout = layouts_lib.from_layout_attr(to_store_layout)
  if not isinstance(layout, fa.TiledLayout):
    raise NotImplementedError(f"Expected TiledLayout, got {type(layout)}")

  ref = op.destination
  transforms_attr = inference_utils.in_transforms(op)[0]
  swizzle = swizzle_from_transforms_attr(transforms_attr)
  unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)
  tiling_transform, = memref_transforms_from_transforms_attr(transforms_attr)
  assert isinstance(tiling_transform, lc.TileTransform)

  dialect_barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)
  barrier_ref = dialect_barrier.barrier_ref

  cluster_dim = gpu.Dimension(op.cluster_dim.value)  # pyrefly: ignore[missing-attribute]
  cluster_idx = arith.index_cast(index, op.cluster_idx)
  cluster_barrier_ref = barrier_ref.remap_to_cluster(cluster_dim, cluster_idx)

  total_bits = math.prod(value.shape) * utils.bitwidth(value.mlir_dtype)
  if total_bits % (8 * utils.WARPGROUP_SIZE):
    raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the value to a TiledLayout (layout_cast / to_layout) before async_store_smem
2. Check isinstance(layouts_lib.from_layout_attr(layout), TiledLayout) before emitting the op

Example fix

// before
async_store_smem(value, smem_ref)
// after
tiled = layout_cast(value, tiled_layout)
async_store_smem(tiled, smem_ref)

layout = layouts_lib.from_layout_attr(in_layout)
assert isinstance(layout, fa.TiledLayout), 'async_store_smem requires TiledLayout'

def is_tiled(attr) -> bool:
    return isinstance(layouts_lib.from_layout_attr(attr), fa.TiledLayout)

When it happens

Trigger: Calling async_store_smem (or emitting mgpu.async_store_smem) where the in_layout attr decodes to a fragment layout like WGStridedFragLayout.

Common situations: Storing a matmul result (accumulator fragment layout) directly to SMEM without first casting to a tiled layout.

Related errors

• {op} has an unsupported layout: {out_layout_attr}
• Transfer of {total_bits} bits is not divisible by {8 * utils
• Output layout {out_layout} must match the accumulator layout
• Unsupported layout: {src.layout}
• Replicated dimensions are not supported

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