Skip to content

Commit

Permalink
feature[next]: Validate no dynamic offsets (#1504)
Browse files Browse the repository at this point in the history
Just a small PR that improves the error message of the temporary pass when there are dynamic offsets alla `shift(Ioff, deref(...))(...)`. This allows catching the error in Icon4Py and failing the respective codegen test gracefully.
  • Loading branch information
tehrengruber committed Mar 20, 2024
1 parent 0a9be51 commit be3728b
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions src/gt4py/next/iterator/transforms/global_tmps.py
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,15 @@ def convert_type(dtype):
)


def validate_no_dynamic_offsets(node: ir.Node):
"""Vaidate we have no dynamic offsets, e.g. `shift(Ioff, deref(...))(...)`"""
for call_node in node.walk_values().if_isinstance(ir.FunCall):
assert isinstance(call_node, ir.FunCall)
if call_node.fun == im.ref("shift"):
if any(not isinstance(arg, ir.OffsetLiteral) for arg in call_node.args):
raise NotImplementedError("Dynamic offsets not supported in temporary pass.")


# TODO(tehrengruber): Add support for dynamic shifts (e.g. the distance is a symbol). This can be
# tricky: For every lift statement that is dynamically shifted we can not compute bounds anymore
# and hence also not extract as a temporary.
Expand All @@ -661,6 +670,8 @@ def visit_FencilDefinition(
] = None,
symbolic_sizes: Optional[dict[str, str]],
) -> FencilWithTemporaries:
# Vaidate we have no dynamic offsets, e.g. `shift(Ioff, deref(...))(...)`
validate_no_dynamic_offsets(node)
# Split closures on lifted function calls and introduce temporaries
res = split_closures(
node, offset_provider=offset_provider, extraction_heuristics=extraction_heuristics
Expand Down

0 comments on commit be3728b

Please sign in to comment.