[nvidia] Support passing TMA descriptors by-value

[embg]

Motivation

Currently, Triton passes TMA descriptors by-ref through global memory. This has a number of problems:

• Significant launch overhead (5-10us) for the host-to-device memcpy
• Users must insert fences for TMA descriptor cache flush (see Short-term solution for TMA descriptor cache management #4342). When users don't insert these fences correctly, they run into very strange bugs: run into dead loop when tuning the tma persistent kernel #4332
• The memcpy makes it nearly impossible to use cudagraphs

There are two possible solutions:

• Pass the descriptor by-value
• Create the descriptor on-device

Because of the tricky memory model for TMA descriptors on H100, creating a descriptor on-device requires moving data back and forth from L2 cache. This is relatively expensive (100s of cycles at least) and requires the user or compiler to correctly insert release/acquire fences.

In some cases, there is no way to avoid creating the descriptor on-device. But for many use-cases, it's perfectly fine to set up the descriptor on the host and pass by-value, avoiding both performance and correctness issues. This PR implements the by-value functionality.

User-level API

Whenever the user provides a kernel param which implements the method tma_desc_cpu_ptr(), Triton will lower that argument to a __grid_constant__ by-value param. The existing helper methods create_[1d/2d]_tma_descriptor were modified to return such a type, so existing code does not need any changes to take advantage of the new feature.

Implementation details

When a kernel param with tma_desc_cpu_ptr() is detected, we attach an attribute to that param at the TTIR level. The attribute is passed through to TTGIR. When lowering TTGIR to LLIR, we use code ported from Mosaic (jax-ml/jax#22175) to set up the correct LLVM attributes. The runtime is also modified to pass by-value TMA descriptors properly.

Limitations

This feature is currently broken when compiling an IRSource directly (which is useful for editing IR and re-compiling). That would require updating some regexes which infer the function signature from the IR. IRSource compilation still works fine for kernels which do not use the new feature.

Once the approach I'm taking here is reviewed, I plan to fix that limitation, either in this PR or in a follow-up PR.

[ThomasRaoux]

what is causing the extra warning?

[embg]

GCC doesn't like the CUtensorMap struct. This is called out in the CUDA C++ Programming Guide as a false warning:

When passing the tensor map as a parameter, some versions of the GCC C++ compiler issue the warning “the ABI for passing parameters with 64-byte alignment has changed in GCC 4.6”. This warning can be ignored.

I don't think it can be suppressed inline via pragma, it has to be suppressed on the command line: https://godbolt.org/z/f5n5crhjG

[htyu]

Thanks for the good work. LGTM in general. Left a couple minor feedbacks.

[htyu]

Is 64 a required alignment value?

[embg]

Yes. Here is the definition of CUtensorMap in <cuda.h>:

typedef struct CUtensorMap_st {
    alignas(64)
    unsigned long long opaque[16];
} CUtensorMap;

[ThomasRaoux]

LGTM once the other comments are addressed