Refactor compiler specializations to consider backend

[giuseros]

In this PR I am trying to refactor the specializations that we apply to the signature of a given function in Triton.

Basically, given a kernel there are some argument properties that can help compilation. E.g., divisibility by 16 and the fact that an integer is equal to 1.

In a previous PR: #4716, I needed other specializations to add buffer support in the AMD backend (and get back some performance when we were using unaligned masked loads).

So this is my attempt to redesign the specialization support to introduce per-backend specializations. The idea is that AttrsDescriptor is now the class that is taking care of doing the analysis of the parameters and adding the specialization. It also has a function table where more specializations can be added per-backend.

[giuseros]

Hi @ThomasRaoux , @antiagainst , this is my best attempt to redesign the specialization collection. Please, let me know what you think

[ThomasRaoux]

This doesn't seem to handle all the places where we deal with specialization like compute_spec_key (this one is a bit sensitive as it is used to compute the hash key and therefore is sensitive to execution time).
Also I was hoping for something where we call into the existing backend classes. I haven't looked in details at how easy it would be to do though.

[giuseros]

Hi @ThomasRaoux , I tried to apply your suggestions:

• I moved the class to store the properties of the different parameters into the common backend compiler class
• I added two methods to the BaseBackend class: get_attrs_descriptor and compute_spec_key. These will be used through out the different code generation phases (i.e., to generate the properties and create a key to retrieve those properties
• Now the properties are stored not as tuple anymore. The value of the property (16 or 1, in the current state of things) is stored in a different structure.

Please take another look and let me know what you think!

[antiagainst]

Thanks for this! A few comments inlined. Bigger questions are:

1. There are hardcoded logic to check current supported properties in runtime jit. We should think about the limitations on backend specific properties to be clear on what they can/cannot do and document. It's fine to only enable backend specific properties to just influence code generation alone IMO.
2. I'm wondering whether we can make the AttrsDescriptor better structured with __slots__.

[giuseros]

I had a go trying to reproduce NVIDIA failures without luck: it seems ok. Let's see if it was a glitch, otherwise I will work harder in reproducing the error

[antiagainst]

Hopefully the last batch of comments. :)

[giuseros]

Hi @antiagainst , I addressed your comments. PTAL. I still see and NVIDIA failure, but I want to get the code right before diving into that.

[antiagainst]

Nice, LGTM now! Thanks for bearing with my nitpicking. :) I'll leave the final approve to Thomas though.

[giuseros]

Thank you for your review @antiagainst !

[giuseros]

FYI, this is how I intend to use this refactor on the AMD backend

[ThomasRaoux]

I think that's fine, just added one comment. I do wonder about the runtime performance impact. I don't see any redflags but I'll check if this is something we have an easy way to benchmark

[giuseros]

Hi @ThomasRaoux , @antiagainst , any updates on this?

[ThomasRaoux]

LGTM