Skip to content

Commit

Permalink
[enums] monkey patch register_attribute_builder (nod-ai#670)
Browse files Browse the repository at this point in the history
Hack/patch to workaround overlapping `AttrBuilder` registrations. Needed
until iree-org/iree#19324 and
llvm/llvm-project#117918 get resolved.
  • Loading branch information
makslevental authored and IanNod committed Dec 17, 2024
1 parent d15f0c4 commit affaf53
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 0 deletions.
15 changes: 15 additions & 0 deletions tuner/tuner/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,18 @@
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from iree.compiler import ir # type: ignore


# Substitute `replace=True` so that colliding registration don't error.
# TODO(makslevental): remove after https://github.com/llvm/llvm-project/pull/117918 is resolved.
def register_attribute_builder(kind, replace=True):
def decorator_builder(func):
ir.AttrBuilder.insert(kind, func, replace=replace)
return func

return decorator_builder


ir.register_attribute_builder = register_attribute_builder
4 changes: 4 additions & 0 deletions tuner/tuner/libtuner_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -499,3 +499,7 @@ def test_validate_devices_with_invalid_device() -> None:
exit_program=True,
)
assert expected_call in mock_handle_error.call_args_list


def test_enum_collision():
from iree.compiler.dialects import linalg, vector, iree_gpu, iree_codegen, iree_input # type: ignore

0 comments on commit affaf53

Please sign in to comment.