[BUG] visit_IfExp would execute both if and else statements

[henrylhtsang]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
If you run

ans = [cute.printf("if statement\n"), 1][1] if row_idx < self.seqlen_q else [cute.printf("else statement\n"), 1][1]

then you would see both executed

if statement

else statement

From @v0i0, this is due to https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py#L1637, which transforms the if else into select so it computes both sides.

This can be a bit unexpected, since this differs from the usual behavior:

if A:
  XXX
else:
  YYY

Steps/Code to reproduce bug
Include this code.

ans = [cute.printf("if statement\n"), 1][1] if row_idx < self.seqlen_q else [cute.printf("else statement\n"), 1][1]

Expected behavior
Same behavior as normal if else blocks.

Environment details (please complete the following information):
Linux, python 3.12, GB200, nvidia-cutlass-dsl 4.3.3

Additional context
NA

[anakinxc]

Yes, as ifExp with dynamic predicate is a select after rewrite, this is kind of unavoidable.

We will see what we can do

[ngimel]

That means that your rewrite is wrong, language semantics dictates that only one branch is executed.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[brandon-yujie-sun]

@henrylhtsang @ngimel this issue has been fixed in 4.4.0.dev1. Please let us know if you see any other issues.