{"record":{"id":"6232bd81dc115db0","repo":"jax-ml/jax","slug":"the-base-ref-for-aliases-must-come-from-a-slice-sm","errorCode":null,"errorMessage":"The base ref for aliases must come from a slice_smem op.","messagePattern":"The base ref for aliases must come from a slice_smem op\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":1603,"sourceCode":"        assert isinstance(ref, ir.Value)  # make pyrefly happy\n        input_ref_ty = ir.MemRefType(ref.type)\n        if input_ref_ty.memory_space == mgpu_utils.smem():\n          assert layout is None\n          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(\n              mlir_dtype\n          )\n          if ref_bits % 8:\n            raise NotImplementedError(\"Only byte-aligned bitcasts are supported.\")\n          assert offset % gpu_core.SMEM_ALIGNMENT == 0\n\n          if lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n            if not isinstance(ref.owner, mgpu.dialect.SliceSMEMOp):\n              # This restriction can be lifted by:\n              # - Using memref ops to get the pointer and offset of the base ref.\n              # - Subtracting gpu_dialect.dynamic_shared_memory() from those to\n              #   get the base offset relative to the beginning of SMEM.\n              # - Implementing layout and lowering rules for all ops above.\n              raise NotImplementedError(\n                  \"The base ref for aliases must come from a slice_smem op.\"\n              )\n\n            base_offset = ref.owner.offset.value\n            total_offset = base_offset + offset\n\n            ref_ty = ir.MemRefType.get(\n                transformed_shape, mlir_dtype, memory_space=mgpu_utils.smem()\n            )\n            assert ref.owner.alias_id is not None\n            alloc_id = ref.owner.alias_id.value\n            # TODO(bchetioui): Use a scheme resilient to hash collisions.\n            alias_id = hash((offset, alloc_id, alias_group_idx))\n            # The composite key formed of `(offset, alloc_id, alias_group_idx)`\n            # is a unique identifier across:\n            #   - different RefUnions (different `alloc_id`, since two\n            #     distinct RefUnions represent two SMEM allocations);\n            #   - different ref_groups within a RefUnion (different","sourceCodeStart":1585,"sourceCodeEnd":1621,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L1585-L1621","documentation":"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.","triggerScenarios":"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.","commonSituations":"Advanced kernels combining warpgroup MMA with dtype-view aliases; JAX version differences in how SMEM refs are materialized before aliasing.","solutions":["Restructure so the aliased base goes through a slice of the SMEM allocation (slice first, then alias)","Use Thread semantics for that part of the kernel or avoid aliasing the non-sliced ref","Upgrade JAX — support for additional base ops may be added via the documented memref-pointer approach","File an upstream issue with a minimal repro if your pattern is legitimate"],"exampleFix":"# before\nbase = smem_alloc(...)        # owner is alloc op\naliased = base.view(dtype)     # warpgroup: owner not SliceSMEMOp -> error\n# after\nbase = smem_alloc(...)[0:n]   # owner is SliceSMEMOp\naliased = base.view(dtype)","handlingStrategy":"fallback","validationCode":"# ensure the aliased base is produced by an SMEM slice\nassert is_slice_smem_owner(ref), 'warpgroup alias base must come from slice_smem'","typeGuard":"def is_slice_smem_owner(ref) -> bool:\n    owner = getattr(ref, 'owner', None)\n    return owner is not None and type(owner).__name__ == 'SliceSMEMOp'","tryCatchPattern":"try:\n    aliased = base.view(dtype)\nexcept NotImplementedError:\n    base = base[0:n]  # route through a slice\n    aliased = base.view(dtype)\n","preventionTips":["Slice SMEM allocations before aliasing under warpgroup semantics","Keep alias chains simple: alloc -> slice -> view","Pin the JAX version you validated warpgroup kernels against"],"tags":["jax","pallas","mosaic-gpu","warpgroup","smem","aliasing","not-implemented"],"backgroundTag":"unsupported-lowering-path","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}