Fix Constexpr method receivers in dynamic control-flow regions

[nodeeeeee]

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.