Load backend dialects in `IRSource` to make sure `parse_mlir_module` works for third_party backends

[anmyachev]

The changes from https://github.com/triton-lang/triton/pull/4924 do not take into account the situation when ttgir level contains dialects defined in third_party plugins (at least that's my understanding).

I'd also like to point out that the second use of parse_mlir_module function (via parse function call) happens after the dialects are loaded for the backend as well, which is why I thought my changes make sense.

I hope this implementation will suit Triton, or maybe one can suggest other options.

[anmyachev]

I would expect all the dialects to be loaded at this time. I don't understand why it is different for third_party dialects. We have some third_party dialect in tree and they work as far as I know

Let me start with the fact that Nvidia and AMD layouts are used in the test file:

• https://github.com/triton-lang/triton/blob/5ebd1e5a6630877e73cad547fd9495a5441c78be/python/test/unit/language/test_core.py#L138
• https://github.com/triton-lang/triton/blob/5ebd1e5a6630877e73cad547fd9495a5441c78be/python/test/unit/language/test_core.py#L128

For testing, we add our own layout (DPAS), which is defined as an intel-specific dialect and it is loaded during loading of backend dialects. However, those dialects are defined as the main gpu dialects and are loaded when ir.load_dialects(context) function is called. https://github.com/triton-lang/triton/blob/5ebd1e5a6630877e73cad547fd9495a5441c78be/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td#L1026

So at the moment everything works for Triton, but as soon as one want to use backend-specific layouts (namely those defined in third_party folder), be it NVIDIA or AMD, in tests, I believe the error will start to appear here too.

[anmyachev]

Thanks everyone for review!