[BUG] Constexpr missing in dynamic control-flow regions

Which component has the problem?

CuTe DSL

Bug Report

Motivation

I encountered this while implementing FlashAttention. The implementation naturally groups the load and MMA paths into JIT methods on a compile-time configuration object, and selects between them using the runtime warp index. A reduced version looks like this:

class FlashAttention:
    @cute.jit
    def load(self, tensor, pipeline):
        ...

    @cute.jit
    def mma(self, tensor, pipeline):
        ...


@cute.jit
def fa4_device_body(
    fa: cutlass.Constexpr[FlashAttention],
    tensor: cute.Tensor,
    storage,
):
    warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())

    pipeline = make_pipeline(
        storage=storage,
        num_stages=fa.num_stages,
    )

    if warp_idx < 4:
        cute.arch.setmaxregister_decrease(fa.num_producer_regs)
        fa.load(tensor, pipeline)
    else:
        cute.arch.setmaxregister_increase(fa.num_mma_regs)
        fa.mma(tensor, pipeline)

Here, self in FlashAttention.load() and FlashAttention.mma() is a compile-time method receiver. The fa parameter in fa4_device_body() refers to the same kind of compile-time object, made explicit by its Constexpr annotation.

Writing the device body as an actual instance method makes this code compile, but only because the compiler has a special case that excludes a receiver whose name is literally self. Passing the same object as a named Constexpr[FlashAttention] parameter should have equivalent staging semantics, but currently does not.

Problem

CuTe DSL treats Constexpr[T] parameters as compile-time Python meta values. They do not have an MLIR representation and must not become arguments or results of runtime control-flow regions.

However, when a method is called through a Constexpr receiver inside a dynamic if, the receiver is incorrectly captured as a region argument:

class Receiver:
    @cute.jit
    def store(self, output):
        ...


@cute.kernel
def kernel(receiver: cutlass.Constexpr[Receiver], output: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()

    if tidx == Int32(0):
        receiver.store(output)

The region analyzer currently treats the base object of every method call as a mutable runtime value, except when its name is literally self.

Consequently, receiver is added to the values carried through the dynamic if. Lowering then attempts to flatten the plain Python Receiver instance into MLIR values and fails with:

error[TYPE_DYNAMIC_EXPR_UNSUPPORTED]:
A value carried through this `if` is a plain Python value (Meta value)
(a `Receiver`) that cannot be turned into a Runtime value (Staged value)

This means self.method() works because of the existing self special case, while an explicitly annotated receiver: Constexpr[Receiver] does not, despite having the same staging semantics.

Root cause

Constexpr is a property of the original function parameter binding. The AST preprocessor removes parameter annotations while transforming the function, so the later control-flow region analysis no longer knows which names refer to Constexpr parameters.

RegionAnalyzer.visit_Call() therefore sees only a method call on a Python object and adds its receiver to invoked_args.