[AMD] Enable masked load and pointer canonicalization pass

[giuseros]

This PR is doing two things:

• We are using the new llvm.masked{load/store} intrinsics. This means that the backend will take responsibility to lower the stores/loads.
• We are enabling the canonicalization pointer pass on the Triton IR. I extensively run testing and corrected a couple of minor issues still present in the implementation.

The reason why I am enabling both at the same time is because I saw a minor regression with llvm.masked{load,store} which seems to go away when using the pointer canonicalization. Also, this combination seems to reduce the numbers of vgprs used (at least for GEMM kernels).

[zhanglx13]

Also, this combination seems to reduce the numbers of vgprs used (at least for GEMM kernels).

Can you give a concrete example about this effect?

[giuseros]

Can you give a concrete example about this effect?

On my local machine (gfx11 card), if I run:

TRITON_HIP_USE_NEW_STREAM_PIPELINE=1 AMDGCN_ENABLE_DUMP=1 TRITON_ALWAYS_COMPILE=1 python3 python/tutorials/03-matrix-multiplication.py | grep vgpr_

I get:

So a 13% reduction on number of registers for the best config.

[giuseros]

I also want to underline that this PR is blocked by a recent PR #4369 . We will need to add volatile support to masked.load in LLVM before we can merge this.

[giuseros]

Upon further thought, I sank the emission of masked.load and masked.store in llLoad and llStore (respectively). So for now we will use masked intrinsics when we can use them

[antiagainst]

Cool thanks! Just two small nits.

[antiagainst]

Could you fix the test and update the pull request message too?

[giuseros]

Unfortunately the issue with the test revealed something more problematic of the pass. I might do a separate PR about this.

[giuseros]

Hi @antiagainst , this is the PR that fixes the issue: #4678

[giuseros]

Hi @antiagainst , I rebased, fixed the nits and also addressed the fact that we are now using llLoad also for LDS access. This mean that I had to add support for non-vector types when I emit the masked.load. I ran few tests and everything seems ok (also the FA test). Let's see how the CI does.

[giuseros]

Hi @antiagainst , could you start the tests again ?