Skip to content

Fix out-of-bounds access in nested select/if_then_else codegen#26

Merged
LeiWang1999 merged 1 commit intotilelang_mainfrom
fix/select-if-then-else-oob
Feb 3, 2026
Merged

Fix out-of-bounds access in nested select/if_then_else codegen#26
LeiWang1999 merged 1 commit intotilelang_mainfrom
fix/select-if-then-else-oob

Conversation

@LeiWang1999
Copy link
Member

Summary

  • Fix potential out-of-bounds memory access when if_then_else is nested inside select
  • Add select_condition_stack_ to track outer select conditions during code generation
  • Combine outer select condition with inner if_then_else condition to guard memory access

Problem

When generating code for select(cond1, if_then_else(cond2, A[index], default), fallback), the previous implementation would evaluate A[index] even when cond1 is false, potentially causing out-of-bounds access.

Before:

float condval;
if (cond2) {
    condval = A[index];  // Executed even when cond1=false!
} else {
    condval = default;
}
result = (cond1 ? condval : fallback);

After:

float condval;
if ((cond1) && (cond2)) {
    condval = A[index];  // Only executed when cond1=true
} else if (cond1) {
    condval = default;
}
result = (cond1 ? condval : fallback);

Test plan

  • Tested with select_oob.py example that triggers OOB when index < 0
  • Verified generated CUDA code correctly guards memory access

🤖 Generated with Claude Code

- Improved the handling of nested conditions in the if_then_else construct to prevent out-of-bounds access by combining outer select conditions.
- Added a stack to manage select conditions during code generation, ensuring proper evaluation order and safety.
- Updated comments for clarity and better understanding of the changes made.
@LeiWang1999 LeiWang1999 merged commit 391d3f7 into tilelang_main Feb 3, 2026
5 checks passed
LeiWang1999 added a commit to LeiWang1999/tilelang that referenced this pull request Feb 3, 2026
Update TVM to include fix for out-of-bounds memory access when
if_then_else is nested inside select during code generation.

See: tile-ai/tvm#26

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
LeiWang1999 added a commit to tile-ai/tilelang that referenced this pull request Feb 3, 2026
Update TVM submodule: fix select/if_then_else OOB access

Update TVM to include fix for out-of-bounds memory access when
if_then_else is nested inside select during code generation.

See: tile-ai/tvm#26

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant