[Issue]: scf_if_dispatch doesn't support yielding return value from scf.if

[ruanjm]

Problem Description

Ref code: https://github.com/ROCm/aiter/blob/2fad199eebc07179de0b3e2d661122baf241400b/aiter/ops/flydsl/kernels/attn_reduce.py#L402-L432

Following code in main_loop() doesn't get compiled correctly. Current main_loop() always returns None which is not expected.

def __then_X():
    return False    # returns from __then_X, not main_loop
def __else_X():
    ...
    return True     # returns from __else_X, not main_loop
scf_if_dispatch(cond, __then_X, __else_X)

Operating System

all

CPU

all

GPU

all

ROCm Version

all

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response