Fix Constexpr method receivers in dynamic control-flow regions

[nodeeeeee]

Fixes #3632

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.

Fix

This change:

1. Inspects the original function signature before AST transformation.
2. Uses the existing is_arg_annotation_constexpr() helper to record Constexpr parameter names in the preprocessing session.
3. Excludes those names from method receivers carried through runtime control-flow regions.

Normal runtime receivers and decomposable DSL values continue to use the existing region argument handling.

Regression test

A regression test uses a deliberately non-MLIR-decomposable Python class as a Constexpr receiver and calls one of its JIT methods inside a dynamic if.

The test calls cute.compile(), so it verifies the complete lowering path that previously failed.

Validation

Both sides of the comparison used the same nvidia-cutlass-dsl 4.8.0.dev0 native runtime components and the same test.

Before this change

Base commit: 59e3a333 (v4.8.0dev-6)

FAILED test_constexpr_receiver_in_if.py
error[TYPE_DYNAMIC_EXPR_UNSUPPORTED]
1 failed

The failure occurs because _ConstexprReceiver is included in the values passed through the dynamic if.

After this change

Commit: 740f0a6c (v4.8.0dev-7)

1 passed

The new regression test together with the existing structured control-flow tests also passes:

7 passed, 6 warnings in 0.88s

The warnings are existing SmemAllocator deprecation warnings and are unrelated to this change.

[brandon-yujie-sun]

LGTM with a minor comment. cc @anakinxc

[brandon-yujie-sun]

check and confirm the self._session_data is valid before accessing it's attribute?

[nodeeeeee]

self._session_data is a property at ast_preprocessor.py:498, it has an assert that checks the validity of self._session_data. So I think we don't need to check again.