{"record":{"id":"eca23782e9e727b3","repo":"jax-ml/jax","slug":"unsupported-memory-space-when-lowering-memref-cast","errorCode":null,"errorMessage":"Unsupported memory space when lowering memref.cast: {memory_space}","messagePattern":"Unsupported memory space when lowering memref\\.cast: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2109,"sourceCode":"      )\n    result = memref.cast(\n        transform_type(ir.MemRefType(op.result.type), out_transforms),\n        unwrap_transformed_memref(op.source, in_transforms),\n    )\n    return [wrap_transformed_memref(result, op.result.type, out_transforms)]\n\n  if memory_space == utils.tmem():\n    [in_tmem_layout] = inference_utils.in_tmem_layouts(op)\n    [out_tmem_layout] = inference_utils.out_tmem_layouts(op)\n    if in_tmem_layout != out_tmem_layout:\n      raise NotImplementedError(\n          \"memref.cast tmem layouts must be identical for both input and\"\n          f\" output but got {in_tmem_layout=} and {out_tmem_layout=}\"\n      )\n    return [_tmem_ref_to_ir(_tmem_ref_from_ir(op.source, in_tmem_layout),\n                            op.result.type)]\n\n  raise NotImplementedError(\n      f\"Unsupported memory space when lowering memref.cast: {memory_space}\"\n  )\n\n\ndef _permutation_to_affine_map_attr(\n    permutation: Sequence[int],\n) -> ir.AffineMapAttr:\n  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))\n\n\n@_register_lowering(memref.TransposeOp, support_warp_semantics=True)\ndef _memref_transpose_op_lowering_rule(\n    ctx: LoweringContext, op: memref.TransposeOp\n) -> Sequence[ir.Value]:\n  del ctx\n\n  in_transforms_attr = inference_utils.in_transforms(op)[0]\n  unwrapped_in_ref = unwrap_transformed_memref(op.in_, in_transforms_attr)","sourceCodeStart":2091,"sourceCodeEnd":2127,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2091-L2127","documentation":"memref.cast lowering in Mosaic only handles shared memory (smem) and tensor memory (tmem). A cast whose result memory space is anything else hits the terminal NotImplementedError.","triggerScenarios":"Constructing memref.cast where ir.MemRefType(op.result.type).memory_space is neither utils.smem() nor utils.tmem() (e.g. global/-generic memory space attribute).","commonSituations":"Hand-building Mosaic IR with custom memory space integer attributes, or a version change where a new memory space enum value isn't yet handled by the lowering.","solutions":["Restrict casts to smem or tmem results; for global memory use plain loads/stores or load_tensor/store_tensor ops","If you need a global cast, lower it before Mosaic (standard MLIR memref.cast is a no-op) and skip the Mosaic rule","Check the memory_space attribute value on the result type and fix the type construction"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def cast_memory_space_supported(ty) -> bool:\n    ms = ir.MemRefType(ty).memory_space\n    return ms == utils.smem() or ms == utils.tmem()","tryCatchPattern":null,"preventionTips":["Restrict casts to smem/tmem","Use loads/stores for global memory"],"tags":["jax","mosaic-gpu","memref","cast","memory-space"],"backgroundTag":"unsupported-memory-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}