Load backend dialects in IRSource to make sure parse_mlir_module works for third_party backends
#5146
Conversation
…works for third_party backends Signed-off-by: Anatoly Myachev <[email protected]>
python/triton/compiler/compiler.py
Outdated
| self.path = path | ||
| path = Path(path) | ||
| self.ext = path.suffix[1:] | ||
| self.src = path.read_text() | ||
| ir.load_dialects(context) | ||
| if target is None: | ||
| target = driver.active.get_current_target() | ||
| assert isinstance(target, GPUTarget), "target must be of GPUTarget type" |
There was a problem hiding this comment.
Why the target must be a GPU target?
There was a problem hiding this comment.
I took this check from compile function, just in case. However, looking at the code, it seems I can make target (or passing backend directly as suggested by @peterbell10 ) non-optional, then there will be no need to have assert here. Does it make sense?
Let me start with the fact that Nvidia and AMD layouts are used in the test file:
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 So at the moment everything works for Triton, but as soon as one want to use backend-specific layouts (namely those defined in |
Signed-off-by: Anatoly Myachev <[email protected]>
| @@ -270,10 +271,7 @@ def compile(src, target=None, options=None): | |||
| context = ir.context() | |||
| ir.load_dialects(context) | |||
| backend.load_dialects(context) | |||
| else: | |||
| # For IRSource, we have already grabbed the context + called ir.load_dialects | |||
There was a problem hiding this comment.
Please keep the comment, and just update it.
Signed-off-by: Anatoly Myachev <[email protected]>
|
Thanks everyone for review! |
The changes from #4924 do not take into account the situation when
ttgirlevel 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_modulefunction (viaparsefunction 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.