{"record":{"id":"05cf4f61d6e066a2","repo":"jax-ml/jax","slug":"subviewop-only-supports-a-single-tile-transform","errorCode":null,"errorMessage":"SubViewOp only supports a single tile transform.","messagePattern":"SubViewOp only supports a single tile transform\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2039,"sourceCode":"        )\n      new_sizes = tile_transform.transform_shape(list(op.static_sizes))\n      # TODO(bchetioui): support transposed offsets.\n      new_static_offsets, new_dynamic_offsets = _tile_transform_offsets(\n          tiling, list(op.static_offsets), list(op.offsets)\n      )\n\n      new_subview_op = memref.SubViewOp(\n          transform_type(ir.MemRefType(op.result.type), transforms),\n          unwrapped_source_ref,\n          new_dynamic_offsets,\n          sizes=[],\n          strides=[],\n          static_offsets=new_static_offsets,\n          static_sizes=new_sizes,\n          static_strides=[1] * len(in_transformed_ty.shape),\n      )\n    case _:\n      raise NotImplementedError(\n          \"SubViewOp only supports a single tile transform.\"\n      )\n\n  wrapped_ref = wrap_transformed_memref(\n      new_subview_op.result, op.result.type, out_transforms\n  )\n  return [wrapped_ref]\n\n\n# memref.cast shows up when we slice a ref with a dynamic index, that later gets\n# folded into a constant. At that time, the sliced ref type is simplified from\n# having a dynamic offset to a constant offset. However, downstream consumer ops\n# still expect the offset to be dynamic, forcing the insertion of a memref.cast\n# op to reintroduce the dynamic offset.\n@_register_lowering(memref.CastOp, support_warp_semantics=True)\ndef _memref_cast_op_lowering_rule(\n    ctx: LoweringContext, op: memref.CastOp\n) -> Sequence[ir.Value]:","sourceCodeStart":2021,"sourceCodeEnd":2057,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2021-L2057","documentation":"Mosaic GPU's lowering of memref.subview only handles a single tile transform on the source memref. When the input has zero or multiple transforms, the pattern match falls into the default case and lowering fails.","triggerScenarios":"Calling t.memref_subview (or building a SubViewOp) on a memref whose in_transforms contains more than one MemRefTransform (e.g. nested/stacked TileTransforms) during the Mosaic-to-MemRef dialect lowering pass.","commonSituations":"Composing multiple layout transforms (e.g. swizzle + tiling, or two tilings) on a shared-memory tensor and then slicing it in a Mosaic GPU kernel; typically after upgrading JAX where transform inference became stricter.","solutions":["Reduce the memref to a single tile transform before taking a subview (apply the other transform manually via explicit reshape/layout ops)","Check inference_utils.in_transforms(op) in a debug pass to see why multiple transforms were inferred","Avoid slicing transformed memrefs; instead slice first and then apply the tiling transform"],"exampleFix":"// before\nsliced = t.memref_subview(tiled_and_swizzled_ref, offsets, sizes)\n// after\nplain = t.memref_subview(plain_ref, offsets, sizes)\ntiled = t.tile(plain, tile_shape)  # single transform","handlingStrategy":"validation","validationCode":"in_ts = inference_utils.in_transforms(op)\nassert len(in_ts) <= 1, f'subview needs <=1 transform, got {len(in_ts)}'","typeGuard":"def has_single_tile_transform(op) -> bool:\n    ts = inference_utils.in_transforms(op)\n    return len(ts) == 1 and isinstance(ts[0], lc.TileTransform)","tryCatchPattern":null,"preventionTips":["Keep at most one tile transform per memref","Prefer slicing before applying transforms"],"tags":["jax","mosaic-gpu","memref","subview","tile-transform"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}