From d01390afb4217dad162fc2359b6dfe2f45e3246b Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Tue, 23 Jul 2024 04:57:32 +0000 Subject: [PATCH 1/5] IR: Add auto-casting of body tuples to Loop and Conditional --- loki/ir/nodes.py | 37 +++++++++++++++++++++++++--------- loki/ir/tests/test_ir_nodes.py | 31 ++++++++++++++++++++++++---- 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/loki/ir/nodes.py b/loki/ir/nodes.py index 50c732296..36401355b 100644 --- a/loki/ir/nodes.py +++ b/loki/ir/nodes.py @@ -20,6 +20,7 @@ from pymbolic.primitives import Expression from pydantic.dataclasses import dataclass as dataclass_validated +from pydantic import model_validator from loki.scope import Scope from loki.tools import flatten, as_tuple, is_iterable, truncate_string, CaseInsensitiveDict @@ -50,6 +51,14 @@ # Using this decorator, we can force strict validation dataclass_strict = partial(dataclass_validated, config=dataclass_validation_config) + +def _sanitize_tuple(t): + """ + Small helper method to ensure non-nested tuples without ``None``. + """ + return tuple(n for n in flatten(as_tuple(t)) if n is not None) + + # Abstract base classes @dataclass_strict(frozen=True) @@ -235,14 +244,17 @@ class InternalNode(Node): The nodes that make up the body. """ - # Certain Node types may contain Module / Subroutine objects - body: Tuple[Any, ...] = None + body: Tuple[Node, ...] = () _traversable = ['body'] - def __post_init__(self): - super().__post_init__() - assert self.body is None or isinstance(self.body, tuple) + @model_validator(mode='before') + @classmethod + def pre_init(cls, values): + """ Ensure non-nested tuples for body. """ + if 'body' in values.kwargs: + values.kwargs['body'] = _sanitize_tuple(values.kwargs['body']) + return values def __repr__(self): raise NotImplementedError @@ -560,7 +572,7 @@ class _ConditionalBase(): condition: Expression body: Tuple[Node, ...] - else_body: Optional[Tuple[Node, ...]] = None + else_body: Optional[Tuple[Node, ...]] = () inline: bool = False has_elseif: bool = False name: Optional[str] = None @@ -596,14 +608,19 @@ class Conditional(InternalNode, _ConditionalBase): _traversable = ['condition', 'body', 'else_body'] + @model_validator(mode='before') + @classmethod + def pre_init(cls, values): + values = super().pre_init(values) + # Ensure non-nested tuples for else_body + if 'else_body' in values.kwargs: + values.kwargs['else_body'] = _sanitize_tuple(values.kwargs['else_body']) + return values + def __post_init__(self): super().__post_init__() assert self.condition is not None - if self.body is not None: - assert isinstance(self.body, tuple) - assert all(isinstance(c, Node) for c in self.body) # pylint: disable=not-an-iterable - if self.has_elseif: assert len(self.else_body) == 1 assert isinstance(self.else_body[0], Conditional) # pylint: disable=unsubscriptable-object diff --git a/loki/ir/tests/test_ir_nodes.py b/loki/ir/tests/test_ir_nodes.py index 18ef52824..3cbef20fa 100644 --- a/loki/ir/tests/test_ir_nodes.py +++ b/loki/ir/tests/test_ir_nodes.py @@ -73,6 +73,7 @@ def test_loop(scope, one, i, n, a_i): assert isinstance(loop.bounds, Expression) assert isinstance(loop.body, tuple) assert all(isinstance(n, ir.Node) for n in loop.body) + assert loop.children == ( i, bounds, (assign,) ) # Ensure "frozen" status of node objects with pytest.raises(FrozenInstanceError) as error: @@ -82,9 +83,17 @@ def test_loop(scope, one, i, n, a_i): with pytest.raises(FrozenInstanceError) as error: loop.body = (assign, assign, assign) + # Test auto-casting of the body to tuple + loop = ir.Loop(variable=i, bounds=bounds, body=assign) + assert loop.body == (assign,) + loop = ir.Loop(variable=i, bounds=bounds, body=( (assign,), )) + assert loop.body == (assign,) + loop = ir.Loop(variable=i, bounds=bounds, body=( assign, (assign,), assign, None)) + assert loop.body == (assign, assign, assign) + # Test errors for wrong contructor usage with pytest.raises(ValidationError) as error: - ir.Loop(variable=i, bounds=bounds, body=assign) + ir.Loop(variable=i, bounds=bounds, body=n) with pytest.raises(ValidationError) as error: ir.Loop(variable=None, bounds=bounds, body=(assign,)) with pytest.raises(ValidationError) as error: @@ -108,6 +117,7 @@ def test_conditional(scope, one, i, n, a_i): assert all(isinstance(n, ir.Node) for n in cond.body) assert isinstance(cond.else_body, tuple) and len(cond.else_body) == 1 assert all(isinstance(n, ir.Node) for n in cond.else_body) + assert cond.children == ( condition, (assign, assign), (assign,) ) with pytest.raises(FrozenInstanceError) as error: cond.condition = parse_expr('k == 0', scope=scope) @@ -116,8 +126,21 @@ def test_conditional(scope, one, i, n, a_i): with pytest.raises(FrozenInstanceError) as error: cond.else_body = (assign, assign, assign) - # Test errors for wrong contructor usage - with pytest.raises(ValidationError) as error: - ir.Conditional(condition=condition, body=assign) + # Test auto-casting of the body / else_body to tuple + cond = ir.Conditional(condition=condition, body=assign) + assert cond.body == (assign,) and cond.else_body == () + cond = ir.Conditional(condition=condition, body=( (assign,), )) + assert cond.body == (assign,) and cond.else_body == () + cond = ir.Conditional(condition=condition, body=( assign, (assign,), assign, None)) + assert cond.body == (assign, assign, assign) and cond.else_body == () + + cond = ir.Conditional(condition=condition, body=(), else_body=assign) + assert cond.body == () and cond.else_body == (assign,) + cond = ir.Conditional(condition=condition, body=(), else_body=( (assign,), )) + assert cond.body == () and cond.else_body == (assign,) + cond = ir.Conditional( + condition=condition, body=(), else_body=( assign, (assign,), assign, None) + ) + assert cond.body == () and cond.else_body == (assign, assign, assign) # TODO: Test inline, name, has_elseif From 325f260c68a8dca8b9493a6de5fa0ad591cfd698 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 24 Jul 2024 09:44:22 +0000 Subject: [PATCH 2/5] IR: Allow `Scope` in `LeafNode` bodies and add test for `Section` --- loki/ir/nodes.py | 13 +--------- loki/ir/tests/test_ir_nodes.py | 47 ++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 12 deletions(-) diff --git a/loki/ir/nodes.py b/loki/ir/nodes.py index 36401355b..ad8518de3 100644 --- a/loki/ir/nodes.py +++ b/loki/ir/nodes.py @@ -244,7 +244,7 @@ class InternalNode(Node): The nodes that make up the body. """ - body: Tuple[Node, ...] = () + body: Tuple[Union[Node, Scope], ...] = () _traversable = ['body'] @@ -327,9 +327,6 @@ def __setstate__(self, s): class _SectionBase(): """ Type definitions for :any:`Section` node type. """ - # Sections may contain Module / Subroutine objects - body: Tuple[Any, ...] = () - @dataclass_strict(frozen=True) class Section(InternalNode, _SectionBase): @@ -337,14 +334,6 @@ class Section(InternalNode, _SectionBase): Internal representation of a single code region. """ - def __post_init__(self): - super().__post_init__() - assert self.body is None or isinstance(self.body, tuple) - - # Ensure we have no nested tuples in the body - if not all(not isinstance(n, tuple) for n in as_tuple(self.body)): - self._update(body=as_tuple(flatten(self.body))) - def append(self, node): """ Append the given node(s) to the section's body. diff --git a/loki/ir/tests/test_ir_nodes.py b/loki/ir/tests/test_ir_nodes.py index 3cbef20fa..dcac57105 100644 --- a/loki/ir/tests/test_ir_nodes.py +++ b/loki/ir/tests/test_ir_nodes.py @@ -14,6 +14,7 @@ from loki.expression import symbols as sym, parse_expr from loki.ir import nodes as ir from loki.scope import Scope +from loki.subroutine import Subroutine @pytest.fixture(name='scope') @@ -36,6 +37,10 @@ def fixture_n(scope): def fixture_a_i(scope, i): return sym.Array('a', dimensions=(i,), scope=scope) +@pytest.fixture(name='a_n') +def fixture_a_n(scope, n): + return sym.Array('a', dimensions=(n,), scope=scope) + def test_assignment(scope, a_i): """ @@ -144,3 +149,45 @@ def test_conditional(scope, one, i, n, a_i): assert cond.body == () and cond.else_body == (assign, assign, assign) # TODO: Test inline, name, has_elseif + + +def test_section(scope, one, i, n, a_n, a_i): + """ + Test constructors and behaviour of :any:`Section` nodes. + """ + assign = ir.Assignment(lhs=a_i, rhs=sym.Literal(42.0)) + decl = ir.VariableDeclaration(symbols=(a_n,)) + func = Subroutine( + name='F', is_function=True, spec=(decl,), body=(assign,) + ) + + # Test constructor for nodes and subroutine objects + sec = ir.Section(body=(assign, assign)) + assert isinstance(sec.body, tuple) and len(sec.body) == 2 + assert all(isinstance(n, ir.Node) for n in sec.body) + with pytest.raises(FrozenInstanceError) as error: + sec.body = (assign, assign) + + sec = ir.Section(body=(func, func)) + assert isinstance(sec.body, tuple) and len(sec.body) == 2 + assert all(isinstance(n, Scope) for n in sec.body) + with pytest.raises(FrozenInstanceError) as error: + sec.body = (func, func) + + # Test auto-casting of the body to tuple + sec = ir.Section(body=assign) + assert sec.body == (assign,) + sec = ir.Section(body=( (assign,), )) + assert sec.body == (assign,) + sec = ir.Section(body=( assign, (assign,), assign, None)) + assert sec.body == (assign, assign, assign) + + # Test prepend/insert/append additions + sec = ir.Section(body=func) + assert sec.body == (func,) + sec.prepend(assign) + assert sec.body == (assign, func) + sec.append((assign, assign)) + assert sec.body == (assign, func, assign, assign) + sec.insert(pos=3, node=func) + assert sec.body == (assign, func, assign, func, assign) From 99051d57e6ba8e5fd981ae694bb29443baf243d1 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 24 Jul 2024 13:32:28 +0000 Subject: [PATCH 3/5] IR: Fix tuple sanitising for non-keyword constructor use of `body` --- loki/ir/nodes.py | 17 +++++++++++++---- loki/ir/tests/test_ir_nodes.py | 5 +++++ 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/loki/ir/nodes.py b/loki/ir/nodes.py index ad8518de3..32de0b314 100644 --- a/loki/ir/nodes.py +++ b/loki/ir/nodes.py @@ -233,7 +233,14 @@ def uses_symbols(self): @dataclass_strict(frozen=True) -class InternalNode(Node): +class _InternalNode(): + """ Type definitions for :any:`InternalNode` node type. """ + + body: Tuple[Union[Node, Scope], ...] = () + + +@dataclass_strict(frozen=True) +class InternalNode(Node, _InternalNode): """ Internal representation of a control flow node that has a traversable `body` property. @@ -244,16 +251,18 @@ class InternalNode(Node): The nodes that make up the body. """ - body: Tuple[Union[Node, Scope], ...] = () - _traversable = ['body'] @model_validator(mode='before') @classmethod def pre_init(cls, values): """ Ensure non-nested tuples for body. """ - if 'body' in values.kwargs: + if values.kwargs and 'body' in values.kwargs: values.kwargs['body'] = _sanitize_tuple(values.kwargs['body']) + if values.args: + # ArgsKwargs are immutable, so we need to force it a little + new_args = (_sanitize_tuple(values.args[0]),) + values.args[1:] + values = type(values)(args=new_args, kwargs=values.kwargs) return values def __repr__(self): diff --git a/loki/ir/tests/test_ir_nodes.py b/loki/ir/tests/test_ir_nodes.py index dcac57105..cb17ea943 100644 --- a/loki/ir/tests/test_ir_nodes.py +++ b/loki/ir/tests/test_ir_nodes.py @@ -174,6 +174,9 @@ def test_section(scope, one, i, n, a_n, a_i): with pytest.raises(FrozenInstanceError) as error: sec.body = (func, func) + sec = ir.Section((assign, assign)) + assert sec.body == (assign, assign) + # Test auto-casting of the body to tuple sec = ir.Section(body=assign) assert sec.body == (assign,) @@ -181,6 +184,8 @@ def test_section(scope, one, i, n, a_n, a_i): assert sec.body == (assign,) sec = ir.Section(body=( assign, (assign,), assign, None)) assert sec.body == (assign, assign, assign) + sec = ir.Section((assign, (func,), assign, None)) + assert sec.body == (assign, func, assign) # Test prepend/insert/append additions sec = ir.Section(body=func) From a013e13e45c4e1252b67766df08b770ad19aa7d2 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 24 Jul 2024 15:23:21 +0000 Subject: [PATCH 4/5] IR: Appeasing the linting gods... --- loki/ir/nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/loki/ir/nodes.py b/loki/ir/nodes.py index 32de0b314..37c54194f 100644 --- a/loki/ir/nodes.py +++ b/loki/ir/nodes.py @@ -392,7 +392,7 @@ class _AssociateBase(): @dataclass_strict(frozen=True) -class Associate(ScopedNode, Section, _AssociateBase): +class Associate(ScopedNode, Section, _AssociateBase): # pylint: disable=too-many-ancestors """ Internal representation of a code region in which names are associated with expressions or variables. From 561d74898110bdac221915e7f3a71caa0fa9576e Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Fri, 26 Jul 2024 12:40:31 +0000 Subject: [PATCH 5/5] Loki: Add version dependency for pydantic>=2.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4fae38aad..61b4b4db8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ "coloredlogs", # optional for loki-build utility "junit_xml", # optional for JunitXML output in loki-lint "codetiming", # essential for scheduler and sourcefile timings - "pydantic", # type checking for IR nodes + "pydantic>=2.0", # type checking for IR nodes ] [project.optional-dependencies]