{"record":{"id":"2769d450adfdf025","repo":"jax-ml/jax","slug":"only-2d-gathers-scatters-for-async-load-store-are","errorCode":null,"errorMessage":"Only 2D gathers/scatters for async load/store are supported.","messagePattern":"Only 2D gathers/scatters for async load/store are supported\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2228,"sourceCode":") -> ConstraintSystemDerivationRuleResult:\n  # We only support 2D gathers/scatters along the leading dimension. Tiling\n  # either keeps the gather/scatter dimension leading or allows\n  # collapsing leading dimensions to maintain contiguity without\n  # transforming global memory.\n  tiling_multiple = []\n  for i, (size, index) in enumerate(zip(op.slice_lengths, op.indices, strict=True)):\n    if size == -1:\n      # This dimension does not appear in the final smem memref shape.\n      continue\n    if isinstance(index.type, ir.VectorType):\n      if i != 0:\n        raise NotImplementedError(\"Only leading gather dimensions allowed.\")\n      if isinstance(op, mgpu.AsyncStoreOp):\n        gmem_shape = ir.MemRefType(op.destination.type).shape\n      else:\n        gmem_shape = ir.MemRefType(op.source.type).shape\n      if len(gmem_shape) != 2:\n        raise NotImplementedError(\"Only 2D gathers/scatters for async load/store are supported.\")\n      tiling_multiple.append(size)\n      continue\n    tiling_multiple.append(dynamic_gcd(size, index))\n\n  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0\n  operand = ValueSite(op, VariableType.OPERAND, operand_index)\n  var = ctx.producer_ref(operand)\n  constraints: list[cs.Constraint] = [\n      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))\n  ]\n  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):\n    element_bitwidth = utils.bitwidth(op.source.type.element_type)\n    # This constraint enforces sufficient SMEM-alignment.\n    # The transfer chunk needs to be 1024 bit-aligned. For each write in the\n    # lowering we transfer 4 rows, so each row must be 256 bit-aligned.\n    divisor = (1024 // 4) // element_bitwidth\n    slice_lengths = [s for s in op.slice_lengths if s != -1]\n    if slice_lengths and (slice_lengths[-1] % divisor):","sourceCodeStart":2210,"sourceCodeEnd":2246,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2210-L2246","documentation":"Vector-indexed (gather/scatter) async loads and stores in Mosaic are only implemented for 2D global-memory tensors: if the source (async_load) or destination (async_store) memref does not have exactly 2 dimensions, a NotImplementedError is raised before constraints are built.","triggerScenarios":"Calling mgpu.async_load/async_store with a VectorType index where the gmem source/destination memref has rank != 2, e.g. a 1D tensor gather or a 3D batched gather.","commonSituations":"Gathering rows from a flattened 1D array, gathering from 3D+ activations in attention-style kernels, or reusing gather code written for 2D against higher-rank tensors after a refactor.","solutions":["Reshape the global tensor to 2D before the gather (e.g. flatten leading dims into rows: [B, N, D] -> [B*N, D]) and adjust offsets accordingly","For 1D gathers, add a trivial second dimension of size 1 (memref.expand_shape / reshape to [N, 1] or [1, N]) so the rank is 2","For genuinely higher-rank access patterns, loop over the extra dims with 2D gathers per slice"],"exampleFix":"# before: src is memref<10000xf32> (rank 1) -> NotImplementedError\nmgpu.async_load(src, smem, indices=vec_rows, slice_lengths=(1,))\n\n# after: reshape to 2D first\nsrc2 = reshape_to_2d(src)  # memref<10000x1xf32>\nmgpu.async_load(src2, smem, indices=vec_rows, slice_lengths=(1, 1))","handlingStrategy":"validation","validationCode":"gmem = op.destination if isinstance(op, mgpu.AsyncStoreOp) else op.source\nif any(isinstance(i.type, ir.VectorType) for i in op.indices) and len(ir.MemRefType(gmem.type).shape) != 2:\n    gmem = reshape_to_2d(gmem)","typeGuard":"def is_2d_gather_compatible(op) -> bool:\n    gmem = op.destination if isinstance(op, mgpu.AsyncStoreOp) else op.source\n    has_vec = any(isinstance(i.type, ir.VectorType) for i in op.indices)\n    return (not has_vec) or len(ir.MemRefType(gmem.type).shape) == 2","tryCatchPattern":null,"preventionTips":["Flatten gmem tensors to rank 2 before vector-indexed async copies","Batch higher-rank gathers as a loop of per-slice 2D gathers"],"tags":["jax","mosaic-gpu","tma","gather","async-load","rank","layout-inference"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}