[BACKEND] Improve printf.

[jlebar]

[BACKEND] Improve printf.

Previously, we printed all of a GPU thread's values in a single printf()
call, and this, plus the user-specified prefix, was all we printed.

This caused a few problems.

• nvptx printf can only handle 32 arguments; if you pass more than
  that, it prints garbage. So if a thread had more than 32 values, you
  couldn't print them, issue issues with print function #2486.

• The order of the values within the Triton program (GPU thread block)
  is an implementation detail -- it depends on the layout the compiler
  assigns to a tensor. So this also prevented you from interpreting
  the printed output.

To address this, we now print the Triton pid and multi-dimensional
Tensor index for each value. And each value gets its own line to avoid
passing too many args to printf.

Example output:

```
pid (0, 1, 2) idx (36, 127) x: 42
```

If you want to observe all the values in a tensor in order, you can grep
and then sort the output.

We also make a UX enhancement to print: The printed label always ends
with ": "; you don't have to add it yourself.

Fixes #2486.

PR chain

1. 👉 [BACKEND] Improve printf. #2532 👈 YOU ARE HERE

[ThomasRaoux]

Caching will break if the new print function is not dominated by the value that has been cached. For instance it will always fail if we are printing from different functions. In general doing caching during rewrite patterns is not trivial, we could cache only the global value but this is probably a premature optimization.

[jlebar]

addStringToModule adds the string as a global, so shouldn't this work? https://github.com/openai/triton/blob/020f43d5a37c2f1cce5c70fcda17ec57b000755c/lib/Conversion/TritonGPUToLLVM/Utility.cpp#L354

(The reason I wanted to cache this is otherwise we get BLOCK_SIZE copies of the format-string. Or otherwise we have to rewrite it so we generate the format-string only once per block, but that's also bug-prone and harder to read.)

[ThomasRaoux]

but it creates operations that are within a block:

Value stringStart =
      rewriter.create<LLVM::GEPOp>(UnknownLoc::get(ctx), ptr_ty(i8_ty),
                                   globalPtr, SmallVector<Value>({zero, zero}));
  return stringStart;

if those are used by an instruction that isn't dominated by the cached ops it would break SSA. Can you change the code so that this gets called outside the loop?

[jlebar]

Indeed you're completely correct. Done, thanks.

[Jokeren]

Another comment regarding the output format. Is it possible to make the format consistent with the output of the assert which also annotates block ids and thread ids?

[jlebar]

Another comment regarding the output format. Is it possible to make the format consistent with the output of the assert which also annotates block ids and thread ids?

The assert message e.g.

/root/code/triton/python/test/unit/language/assert_helper.py:51: kernel_device_assert: block: [0,0,0], thread: [31,0,0] Assertion `x != 0` failed.

is actually coming from nvptx. The only part that we control is x != 0.

The block/thread-id here is CUDA, not Triton pid. AIUI Triton pid is only equivalent to the block id pre-Hopper; on Hopper, the Triton pid seems to be equal to the cluster ID.

Since we don't have control over the assertion message and since we can't get it to print out the actual Triton PID, I'm inclined not to try to match it. WDYT, @Jokeren ?

[Jokeren]

The block/thread-id here is CUDA, not Triton pid. AIUI Triton pid is only equivalent to the block id pre-Hopper; on Hopper, the Triton pid seems to be equal to the cluster ID.

OK, that sounds reasonable.

[ptillet]

thanks @jlebar ! It is a great improvement to our printer :)

[ThomasRaoux]

Looks great!