Introduce tl.assume or use assert expression in non-debug builds to guide optimization?

[0x804d8000]

In certain cases compiler cannot infer some situation cannot happen, e.g.:

@triton.jit
def kernel(N: tl.constexpr, BLOCK_N: tl.constexpr):
  current_size = N - tl.program_id(0) * BLOCK_N
  if current_size >= 128:
    return impl_specialized_for_block_128(...)
  else:
    return impl_slower_but_safe_for_smaller_block(...)

kernel[(1024 // 128, )](1024, 128)

Note that both N and BLOCK_N is marked as constexpr so the compiler will be able to do constant propagation and dead code elimination as necessary.

Here the compiler won't be able to eliminate the if without external knowledge. This is because given a (wrongly) large grid, current_size could get negative, and thus the else branch would be executed.

I think we can either:

• Use assert expression as a hint in release build, and add a assert current_size > 0 before if. This way the compiler should be able to figure out else branch won't be needed. The caveat here is that if the assertion itself is wrong, previously such wrong assertion wouldn't hurt, while now it can lead to undefined behavior.
• Introduce tl.assume. Its usage would be largely the same as assert. However, this might result in people duplicating the same conditional expression for both assert and tl.assume all over the place, which would be a violation of DRY principle.

[Ghost-LZW]

tl.assume has been introduced. #4396

[0x804d8000]

Thanks, closing the issue.