Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 5 commits into from
Nov 15, 2024
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0b43b52
Load backend dialects in 'IRSource' to make sure 'parse_mlir_module' …
anmyachev Nov 14, 2024
5f77793
use 'backend' instead of 'target'
anmyachev Nov 14, 2024
168d3e2
Merge branch 'main' into parse_mlir_module
anmyachev Nov 14, 2024
6dfc3fc
Merge branch 'main' into parse_mlir_module
anmyachev Nov 15, 2024
4dab704
return the comment
anmyachev Nov 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions python/test/unit/tools/test_irsource.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def test_mlir_attribute_parsing(tmp_path: pathlib.Path) -> None:
temp_file = tmp_path / "test_mlir_attribute_parsing0.ttgir"
temp_file.write_text(sample_ttgir)
context = ir.context()
src = IRSource(str(temp_file), context)
src = IRSource(str(temp_file), context, target)

# check name and type signature
# should match ty_to_cpp(...)
Expand Down Expand Up @@ -85,7 +85,7 @@ def test_mlir_attribute_parsing(tmp_path: pathlib.Path) -> None:
temp_file = tmp_path / "test_mlir_attribute_parsing1.ttgir"
temp_file.write_text(sample_ttgir_vector_add)
context = ir.context()
src = IRSource(str(temp_file), context)
src = IRSource(str(temp_file), context, target)

# now test compilation
triton.compile(str(temp_file), target=target)
14 changes: 8 additions & 6 deletions python/triton/compiler/compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,12 +91,17 @@ def parse_options(self):

class IRSource:

def __init__(self, path, context):
def __init__(self, path, context, target=None):
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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why the target must be a GPU target?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

backend = make_backend(target)
backend.load_dialects(context)

# We don't have a easy-to-use PTX parser that we can use, so keep that regex for now.
# TODO - replace with a proper parser
Expand Down Expand Up @@ -223,7 +228,7 @@ def compile(src, target=None, options=None):
if ir_source:
assert isinstance(src, str), "source must be either AST or a filepath"
context = ir.context()
src = IRSource(src, context)
src = IRSource(src, context, target)
peterbell10 marked this conversation as resolved.
Show resolved Hide resolved

extra_options = src.parse_options()
options = backend.parse_options(dict(options or dict(), **extra_options))
Expand Down Expand Up @@ -270,10 +275,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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please keep the comment, and just update it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

# just need to load the dialects for the backend.
backend.load_dialects(context)

codegen_fns = backend.get_codegen_implementation()
module_map = backend.get_module_map()
try:
Expand Down