[FRONTEND][BACKEND] Implement tl.device_assert and rename tl.printf to tl.device_print

[Jokeren]

Note that tl.device_print and print accepts different arguments than the normal print. The first argument must be a string, following by variables.

Device side:

• tl.device_print
• tl.device_assert
• print
• assert

Compilation time:

• tl.static_assert
• tl.static_print

Usage example:

tl.device_assert(x == 0, "x != 0")

Output:

...
python/test/unit/language/assert_helper.py:18: kernel: block: [0,0,0], thread: [33,0,0] Assertion `x != 0` failed.
...

tl.device_print("hello ", x)

Output:

...
hello 1
...

[ptillet]

Do you think we should do things the other way around to be more consistent with Python:

• assert is evaluated at runtime; tl.static_assert is evaluated at compile-time.
• print is evaluated at runtime; tl.static_print is evaluated at compile-time.

Also, maybe we could enable debug mode by default when there kernel contains a print?

[Jokeren]

assert is evaluated at runtime; tl.static_assert is evaluated at compile-time.
print is evaluated at runtime; tl.static_print is evaluated at compile-time.

For "assert", I think the answer is that we could do it this way.

Unfortunately we don't fully support the print semantic yet. Right now tl.device_print only accepts a str prefix with all args appended to the end. If you want to print something like print(f"msg {k}: {v}"), $k and $v will be evaluated at compile time I think.

@christopherhesse What's your suggestion? Long story short, option (1) device_assert evaluates dynamic values, assert evaluates static (i.e., compile-time) values; option (2) assert evaluates dynamic values, static_assert evaluates static values.

Also, maybe we could enable debug mode by default when there kernel contains a print?

Yeah, we could. I haven't sorted out what features should be disabled in the debug mode though. Could we do it in another PR?

[ptillet]

Unfortunately we don't fully support the print semantic yet.

Hmmm. This may be indeed be a little premature. Right now tl.printf is not very useful and would confuse users anyway as it prints once per CUDA thread instead of once per program.

Yeah, we could. I haven't sorted out what features should be disabled in the debug mode though. Could we do it in another PR?

Yes, we can do it in another PR. We just need to create a github issue to make sure we don't lose track of that :)

[ptillet]

@Jokeren I think we should go with assert being evaluated at runtime and static_assert being evaluated statically. This also matches the current semantic of static_range vs range.

[christopherhesse]

I'll have to defer to @ptillet here. I do think if you allow the use of the built-in assert and print it could be weird if they behave differently than the normal versions of those. Here's an example from jax where they have to explain how assert only runs at trace time: jax-ml/jax#2273 (comment) (though at least that is consistent with print)

[Jokeren]

@ptillet @christopherhesse ready for review

[yiakwy-xpu-ml-framework-team]

I

I'll have to defer to @ptillet here. I do think if you allow the use of the built-in assert and print it could be weird if they behave differently than the normal versions of those. Here's an example from jax where they have to explain how assert only runs at trace time: google/jax#2273 (comment) (though at least that is consistent with print)

This print works as on device print op which will be executed by a warps of 32 threads. This is not necessary.

Could we instead develop a LLVM printop for TritonIR: just print during tiles mapping stage and evaluate the result with CPU for verification purpose ?