jax-ml/jax · error · ValueError

gpu.launch op not found.

Error message

gpu.launch op not found.

What it means

The lowering pass searches the module body for the single gpu.launch op that wraps the kernel; if none exists it cannot proceed. This is an internal structure assumption of the Mosaic GPU pipeline.

Source

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

  return (
      # pyrefly: ignore[missing-attribute]
      op.OPERATION_NAME.startswith("mosaic_gpu.")
      or inference_utils.should_have_layout(op)
      or inference_utils.should_have_transforms(op)
      or inference_utils.should_have_tmem_layout(op)
      # Does it have subblocks?
      or any(bool(b) for r in op.regions for b in r)
  )


def _gpu_launch_op(module: ir.Module) -> gpu.LaunchOp:
  for op in module.body.operations:
    for region in op.operation.regions:
      for block in region.blocks:
        for sub_op in block.operations:
          if isinstance(sub_op, gpu.LaunchOp):
            return sub_op
  raise ValueError("gpu.launch op not found.")


def _lowering_context(
    module: ir.Module,
    launch_context: lc.LaunchContext | None,
    auto_barriers: bool,
) -> LoweringContext:
  """Returns a `LoweringContext` for the given `LaunchContext`."""
  # TODO(bchetioui): fix tests to not have a test-only path polluting the API.
  if launch_context is None:  # this case is used in some tests
    return LoweringContext(None, None, None, None, None, auto_barriers, 10**9)

  gpu_launch_op = _gpu_launch_op(module)
  with ir.InsertionPoint.at_block_begin(gpu_launch_op.regions[0].blocks[0]):
    eq = arith.CmpIPredicate.eq
    i32 = ir.IntegerType.get_signless(32)
    single_warp_per_block_predicate = arith.cmpi(
        eq, utils.warp_idx(sync=False), utils.c(0, i32)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the standard entry point (e.g. mosaic gpu kernel launch utilities) that creates the gpu.launch wrapper before lowering
2. If building IR manually, wrap your func in a gpu.launch op before invoking the lowering
3. Check pass ordering — ensure launch materialization runs before the lowering pass

assert any(isinstance(o, gpu.LaunchOp) for f in module.body.operations for r in f.operation.regions for b in r.blocks for o in b.operations), 'missing gpu.launch'

When it happens

Trigger: Calling the lower_mosaic_module pass on a module whose top-level function contains no gpu.launch op — e.g. hand-built IR, partially constructed modules, or a pass ordering where the launch wrapper hasn't been created yet.

Common situations: Running lower_mosaic_module directly on user IR without going through the normal kernel-build path; incorrect pass pipeline ordering after refactor/version changes.

Related errors

• {ref} is not a memref.
• {ref} has a memory space {mem_ref_ty.memory_space} that is n
• A custom return op must terminate the block.
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}

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