diff --git a/src/gt4py/next/iterator/pretty_parser.py b/src/gt4py/next/iterator/pretty_parser.py index fc6802178e..b58f66436f 100644 --- a/src/gt4py/next/iterator/pretty_parser.py +++ b/src/gt4py/next/iterator/pretty_parser.py @@ -16,8 +16,7 @@ visitors as lark_visitors, ) -from gt4py.next.iterator import ir -from gt4py.next.iterator.ir_utils import ir_makers as im +from gt4py.next.iterator import ir, pretty_printer from gt4py.next.type_system import type_specifications as ts @@ -84,6 +83,7 @@ | "c⟨" ( prec0 "," )* prec0? "⟩" -> cartesian_domain ?prec9: _literal + | typed_literal | SYM_REF | named_range | cartesian_offset @@ -94,10 +94,15 @@ else_branch_seperator: "else" if_stmt: "if" "(" prec0 ")" "{" ( stmt )* "}" else_branch_seperator "{" ( stmt )* "}" + typed_literal: ( INT_LITERAL | FLOAT_LITERAL | SYM_REF ) ":" TYPE_LITERAL + ?type_expr: TYPE_LITERAL + | TYPE_LITERAL "[" INT_LITERAL ("," INT_LITERAL)* "]" -> shaped_scalar_type + | "{" type_expr ("," type_expr)* "}" -> tuple_type + named_range: AXIS_LITERAL ":" "[" prec0 "," prec0 "[" cartesian_offset: AXIS_LITERAL "→" AXIS_LITERAL function_definition: ID_NAME "=" "λ(" ( SYM "," )* SYM? ")" "→" prec0 ";" - declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" TYPE_LITERAL ")" ";" + declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" type_expr ")" ";" stencil_closure: prec0 "←" "(" prec0 ")" "(" ( SYM_REF ", " )* SYM_REF ")" "@" prec0 ";" fencil_definition: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( stencil_closure )+ "}" program: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( declaration )* ( stmt )+ "}" @@ -107,6 +112,11 @@ """ # noqa: RUF001 [ambiguous-unicode-character-string] +def _bare_literal(value: str) -> ir.Literal: + """A literal written without a type annotation.""" + return ir.Literal(value=value, type=pretty_printer.implied_literal_type(value)) + + @lark_visitors.v_args(inline=True) class ToIrTransformer(lark_visitors.Transformer): def SYM(self, value: lark_lexer.Token) -> ir.Sym: @@ -114,19 +124,35 @@ def SYM(self, value: lark_lexer.Token) -> ir.Sym: def SYM_REF(self, value: lark_lexer.Token) -> Union[ir.SymRef, ir.Literal]: if value.value in ("True", "False"): - return im.literal(value.value, "bool") + return _bare_literal(value.value) return ir.SymRef(id=value.value) def INT_LITERAL(self, value: lark_lexer.Token) -> ir.Literal: - return im.literal_from_value(int(value.value)) + return _bare_literal(str(int(value.value))) def FLOAT_LITERAL(self, value: lark_lexer.Token) -> ir.Literal: - return im.literal(value.value, "float64") - - def TYPE_LITERAL(self, value: lark_lexer.Token) -> ts.TypeSpec: - if hasattr(ts.ScalarKind, value.upper()): - return ts.ScalarType(kind=getattr(ts.ScalarKind, value.upper())) - raise NotImplementedError(f"Type {value} not supported.") + return _bare_literal(value.value) + + def TYPE_LITERAL(self, value: lark_lexer.Token) -> ts.ScalarType: + if (kind := pretty_printer.SCALAR_TYPE_KINDS.get(value.value)) is None: + raise ValueError( + f"Invalid type '{value}'; expected one of " + f"{', '.join(sorted(pretty_printer.SCALAR_TYPE_KINDS))}." + ) + return ts.ScalarType(kind=kind) + + def shaped_scalar_type(self, type_: ts.ScalarType, *shape: ir.Literal) -> ts.ScalarType: + return ts.ScalarType(kind=type_.kind, shape=[int(s.value) for s in shape]) + + def tuple_type(self, *types: ts.DataType) -> ts.TupleType: + return ts.TupleType(types=list(types)) + + def typed_literal( + self, value: Union[ir.Literal, ir.SymRef], type_: ts.ScalarType + ) -> ir.Literal: + if not isinstance(value, ir.Literal): + raise ValueError(f"Only a literal can carry a type annotation, got '{value.id}'.") + return ir.Literal(value=value.value, type=type_) def OFFSET_LITERAL(self, value: lark_lexer.Token) -> ir.OffsetLiteral: v: Union[int, str] = value.value[:-1] diff --git a/src/gt4py/next/iterator/pretty_printer.py b/src/gt4py/next/iterator/pretty_printer.py index 44d5a9546a..5fbba8920e 100644 --- a/src/gt4py/next/iterator/pretty_printer.py +++ b/src/gt4py/next/iterator/pretty_printer.py @@ -11,14 +11,16 @@ Inspired by P. Yelland, “A New Approach to Optimal Code Formatting”, 2015 """ -# TODO(tehrengruber): add support for printing the types of itir.Sym, itir.Literal nodes +# TODO(tehrengruber): add support for printing the types of itir.Sym nodes from __future__ import annotations -from collections.abc import Iterator, Sequence +import types as _types +from collections.abc import Iterator, Mapping, Sequence from typing import Final from gt4py.eve import NodeTranslator from gt4py.next.iterator import ir +from gt4py.next.type_system import type_specifications as ts, type_translation # replacements for builtin binary operations @@ -63,12 +65,81 @@ } +#: Surface spelling of the scalar types, LLVM/MLIR style. +SCALAR_TYPE_NAMES: Final[Mapping[ts.ScalarKind, str]] = _types.MappingProxyType( + { + ts.ScalarKind.BOOL: "i1", + ts.ScalarKind.INT8: "i8", + ts.ScalarKind.UINT8: "u8", + ts.ScalarKind.INT16: "i16", + ts.ScalarKind.UINT16: "u16", + ts.ScalarKind.INT32: "i32", + ts.ScalarKind.UINT32: "u32", + ts.ScalarKind.INT64: "i64", + ts.ScalarKind.UINT64: "u64", + ts.ScalarKind.FLOAT32: "f32", + ts.ScalarKind.FLOAT64: "f64", + } +) + +SCALAR_TYPE_KINDS: Final[Mapping[str, ts.ScalarKind]] = _types.MappingProxyType( + {name: kind for kind, name in SCALAR_TYPE_NAMES.items()} +) + + +def format_type(type_: ts.TypeSpec) -> str: + if isinstance(type_, ts.TupleType): + return f"{{{', '.join(format_type(t) for t in type_.types)}}}" + if isinstance(type_, ts.ScalarType) and type_.kind in SCALAR_TYPE_NAMES: + name = SCALAR_TYPE_NAMES[type_.kind] + if type_.shape is None: + return name + return f"{name}[{', '.join(str(s) for s in type_.shape)}]" + raise NotImplementedError(f"No pretty-printed form for type '{type_}'.") + + +def implied_literal_type(value: str) -> ts.ScalarType: + """Type the bare, unannotated lexeme `value` denotes. + + `pretty_parser` assigns exactly this to an unannotated lexeme, so the printer + leaves the annotation off when it would agree. The lexeme is read back to a + Python value and typed by `type_translation.from_value`, which is what the + frontend applies to a constant (`ffront/foast_passes/type_deduction.py`, + `TypeDeducer.visit_Constant`), so the two cannot disagree. + + Raises: + ValueError: If `value` is not a literal lexeme; neither a bare nor an + annotated spelling of it parses back. + """ + py_value: bool | int | float + if value in ("True", "False"): + py_value = value == "True" + else: + try: + py_value = int(value) + except ValueError: + try: + py_value = float(value) + except ValueError: + raise ValueError( + f"Invalid literal value '{value}'; expected a boolean, an integer or a" + " floating point lexeme." + ) from None + type_ = type_translation.from_value(py_value) + assert isinstance(type_, ts.ScalarType) + return type_ + + DEFAULT_INDENT: Final = 2 DEFAULT_WIDTH: Final = 100 class PrettyPrinter(NodeTranslator): - def __init__(self, indent: int = DEFAULT_INDENT, width: int = DEFAULT_WIDTH) -> None: + def __init__( + self, + indent: int = DEFAULT_INDENT, + width: int = DEFAULT_WIDTH, + ) -> None: super().__init__() self.indent: int = indent self.width: int = width @@ -131,7 +202,12 @@ def visit_Sym(self, node: ir.Sym, *, prec: int) -> list[str]: return [node.id] def visit_Literal(self, node: ir.Literal, *, prec: int) -> list[str]: - return [str(node.value)] + # Unconditional, as it also rejects a value that is no literal lexeme: printing such a + # value, annotated or not, yields text that does not parse back. + implied_type = implied_literal_type(node.value) + return [ + node.value if implied_type == node.type else f"{node.value}:{format_type(node.type)}" + ] def visit_InfinityLiteral(self, node: ir.InfinityLiteral, *, prec: int) -> list[str]: if node == ir.InfinityLiteral.POSITIVE: @@ -268,7 +344,7 @@ def visit_Temporary(self, node: ir.Temporary, *, prec: int) -> list[str]: if node.domain is not None: args.append(self._hmerge(["domain="], self.visit(node.domain, prec=0))) if node.dtype is not None: - args.append(self._hmerge(["dtype="], [str(node.dtype)])) + args.append(self._hmerge(["dtype="], [format_type(node.dtype)])) hargs = self._hmerge(*self._hinterleave(args, ", ")) vargs = self._vmerge(*self._hinterleave(args, ",")) oargs = self._optimum(hargs, vargs) diff --git a/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py b/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py index aac1096388..83a16a980c 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py @@ -63,6 +63,67 @@ def cmp(builtin): assert pparse("a >= b") == cmp("greater_equal") +def test_typed_literal(): + assert pparse("1.0:f32") == im.literal("1.0", "float32") + assert pparse("1:i64") == im.literal("1", "int64") + assert pparse("True:i1") == im.literal("True", "bool") + + +def test_typed_literal_binds_tighter_than_arithmetic(): + assert pparse("1.0:f32 + 2.0") == ir.FunCall( + fun=ir.SymRef(id="plus"), + args=[im.literal("1.0", "float32"), im.literal("2.0", "float64")], + ) + + +def test_typed_literal_scalar_kind_names_are_not_accepted(): + with pytest.raises(ValueError, match="float32"): + pparse("1.0:float32") + + +def test_only_a_literal_can_be_annotated(): + with pytest.raises(ValueError, match="'x'"): + pparse("x:i32") + + +def test_typed_literal_does_not_shadow_named_range(): + assert pparse("c⟨ IDimₕ: [0:i64, 4:i64[ ⟩") == ir.FunCall( + fun=ir.SymRef(id="cartesian_domain"), + args=[ + ir.FunCall( + fun=ir.SymRef(id="named_range"), + args=[ + ir.AxisLiteral(value="IDim"), + im.literal("0", "int64"), + im.literal("4", "int64"), + ], + ) + ], + ) + + +def test_type_name_lexing_prefers_the_longest_match(): + # `i1` is a proper prefix of `i16`; `TYPE_LITERAL` is a single greedy `CNAME`, + # so the shorter name must never win. + assert pparse("1:i1") == im.literal("1", "bool") + assert pparse("1:i16") == im.literal("1", "int16") + # `[2]` is a tuple index, not a shape + assert pparse("1:i16[2]") == ir.FunCall( + fun=ir.SymRef(id="tuple_get"), + args=[im.literal("2", "int32"), im.literal("1", "int16")], + ) + assert pparse("t = temporary(domain=domain, dtype={i1, i16});") == ir.Temporary( + id="t", + domain=ir.SymRef(id="domain"), + dtype=ts.TupleType( + types=[ + ts.ScalarType(kind=ts.ScalarKind.BOOL), + ts.ScalarType(kind=ts.ScalarKind.INT16), + ] + ), + ) + + def test_deref(): testee = "·x" expected = ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="x")]) @@ -256,13 +317,36 @@ def test_function_definition(): def test_temporary(): - testee = "t = temporary(domain=domain, dtype=float64);" + testee = "t = temporary(domain=domain, dtype=f64);" float64_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) expected = ir.Temporary(id="t", domain=ir.SymRef(id="domain"), dtype=float64_type) actual = pparse(testee) assert actual == expected +def test_temporary_compound_dtype(): + assert pparse("t = temporary(domain=domain, dtype={i1, i16});") == ir.Temporary( + id="t", + domain=ir.SymRef(id="domain"), + dtype=ts.TupleType( + types=[ + ts.ScalarType(kind=ts.ScalarKind.BOOL), + ts.ScalarType(kind=ts.ScalarKind.INT16), + ] + ), + ) + assert pparse("t = temporary(domain=domain, dtype=f64[3]);") == ir.Temporary( + id="t", + domain=ir.SymRef(id="domain"), + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=[3]), + ) + + +def test_scalar_kind_names_are_not_accepted(): + with pytest.raises(ValueError, match="float64"): + pparse("t = temporary(domain=domain, dtype=float64);") + + def test_set_at(): testee = "y @ cartesian_domain() ← x;" expected = ir.SetAt( @@ -306,7 +390,7 @@ def test_if_stmt(): def test_program(): - testee = "f(d, x, y) {\n g = λ(x) → x;\n tmp = temporary(domain=cartesian_domain(), dtype=float64);\n y @ cartesian_domain() ← x;\n}" + testee = "f(d, x, y) {\n g = λ(x) → x;\n tmp = temporary(domain=cartesian_domain(), dtype=f64);\n y @ cartesian_domain() ← x;\n}" expected = ir.Program( id="f", function_definitions=[ @@ -333,5 +417,5 @@ def test_program(): def test_transformer_error_is_not_wrapped(): - with pytest.raises(NotImplementedError, match="nonesuch"): + with pytest.raises(ValueError, match="nonesuch"): pparse("t = temporary(domain=cartesian_domain(), dtype=nonesuch);") diff --git a/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py b/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py index f9b3c7eb1f..1144b727b2 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py @@ -6,7 +6,9 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from gt4py.next.iterator import ir, builtins +import pytest + +from gt4py.next.iterator import builtins, ir, pretty_printer from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next.iterator.pretty_printer import PrettyPrinter, pformat from gt4py.next.type_system import type_specifications as ts @@ -97,6 +99,40 @@ def test_offset_literal(): assert actual == expected +def test_literal_type_annotation(): + assert pformat(im.literal("1.0", "float32")) == "1.0:f32" + assert pformat(im.literal("1", "int64")) == "1:i64" + assert pformat(im.literal("1", "int16")) == "1:i16" + assert pformat(im.literal("2147483648", "int32")) == "2147483648:i32" + + +def test_literal_type_annotation_elided_when_implied(): + assert pformat(im.literal("1", "int32")) == "1" + assert pformat(im.literal("1.0", "float64")) == "1.0" + assert pformat(im.literal("True", "bool")) == "True" + assert pformat(im.literal("2147483648", "int64")) == "2147483648" + + +def test_literal_value_that_is_no_lexeme_is_rejected(): + # A value with no literal spelling would be printed into text that does not parse back + # (`hello:i32`, `0x10:i32`, `:i32`), so it is rejected instead. + for value in ("hello", "0x10", ""): + with pytest.raises(ValueError, match="Invalid literal value"): + pformat(im.literal(value, "int32")) + with pytest.raises(ValueError, match="Invalid literal value"): + pretty_printer.implied_literal_type(value) + + +def test_implied_literal_type(): + def implied(value): + return pretty_printer.implied_literal_type(value) + + assert implied("True") == ts.ScalarType(kind=ts.ScalarKind.BOOL) + assert implied("1") == ts.ScalarType(kind=ts.ScalarKind.INT32) + assert implied("2147483648") == ts.ScalarType(kind=ts.ScalarKind.INT64) + assert implied("1.0") == ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + + def test_arithmetic(): testee = ir.FunCall( fun=ir.SymRef(id="divides"), @@ -117,7 +153,7 @@ def test_arithmetic(): im.literal("4", "int64"), ], ) - expected = "(1 + 2) × 3 / 4" + expected = "(1:i64 + 2:i64) × 3:i64 / 4:i64" actual = pformat(testee) assert actual == expected @@ -136,7 +172,7 @@ def test_associativity(): ), ], ) - expected = "1 + 2 + (3 + 4)" + expected = "1:i64 + 2:i64 + (3:i64 + 4:i64)" actual = pformat(testee) assert actual == expected @@ -317,11 +353,56 @@ def test_temporary(): testee = ir.Temporary( id="t", domain=ir.SymRef(id="domain"), dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) ) - expected = "t = temporary(domain=domain, dtype=float64);" + expected = "t = temporary(domain=domain, dtype=f64);" actual = pformat(testee) assert actual == expected +def test_scalar_type_names_cover_the_type_builtins(): + assert set(pretty_printer.SCALAR_TYPE_NAMES) == { + getattr(ts.ScalarKind, name.upper()) for name in builtins.TYPE_BUILTINS + } + assert ts.ScalarKind.STRING not in pretty_printer.SCALAR_TYPE_NAMES + + +def test_format_type(): + def scalar(kind, shape=None): + return pretty_printer.format_type(ts.ScalarType(kind=kind, shape=shape)) + + assert scalar(ts.ScalarKind.BOOL) == "i1" + assert scalar(ts.ScalarKind.INT16) == "i16" + assert scalar(ts.ScalarKind.UINT8) == "u8" + assert scalar(ts.ScalarKind.FLOAT32) == "f32" + assert scalar(ts.ScalarKind.FLOAT64, [3, 4]) == "f64[3, 4]" + assert ( + pretty_printer.format_type( + ts.TupleType( + types=[ + ts.ScalarType(kind=ts.ScalarKind.BOOL), + ts.ScalarType(kind=ts.ScalarKind.INT16), + ] + ) + ) + == "{i1, i16}" + ) + with pytest.raises(NotImplementedError): + scalar(ts.ScalarKind.STRING) + + +def test_temporary_compound_dtype(): + def temp(dtype): + return pformat(ir.Temporary(id="t", domain=ir.SymRef(id="domain"), dtype=dtype)) + + assert ( + temp(ts.TupleType(types=[ts.ScalarType(kind=ts.ScalarKind.FLOAT32)])) + == "t = temporary(domain=domain, dtype={f32});" + ) + assert ( + temp(ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=[3])) + == "t = temporary(domain=domain, dtype=f64[3]);" + ) + + def test_set_at(): testee = ir.SetAt( expr=ir.SymRef(id="x"), @@ -356,5 +437,5 @@ def test_program(): ], ) actual = pformat(testee) - expected = "f(d, x, y) {\n g = λ(x) → x;\n tmp = temporary(domain=cartesian_domain(), dtype=float64);\n y @ cartesian_domain() ← x;\n}" + expected = "f(d, x, y) {\n g = λ(x) → x;\n tmp = temporary(domain=cartesian_domain(), dtype=f64);\n y @ cartesian_domain() ← x;\n}" assert actual == expected diff --git a/tests/next_tests/unit_tests/iterator_tests/test_pretty_roundtrip.py b/tests/next_tests/unit_tests/iterator_tests/test_pretty_roundtrip.py index 27f4336d93..81396d8c59 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_pretty_roundtrip.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_pretty_roundtrip.py @@ -23,11 +23,6 @@ from gt4py.next.type_system import type_specifications as ts -_XFAIL_LITERAL_TYPE = pytest.mark.xfail( - reason="`Literal.type` is not printed; the parser re-types the lexeme", strict=True -) - - _SET_AT = ir.SetAt(expr=im.ref("x"), domain=im.call("cartesian_domain")(), target=im.ref("y")) @@ -50,7 +45,6 @@ im.literal("4", "int64"), ), id="arithmetic", - marks=_XFAIL_LITERAL_TYPE, ), pytest.param( im.plus( @@ -58,7 +52,6 @@ im.plus(im.literal("3", "int64"), im.literal("4", "int64")), ), id="associativity", - marks=_XFAIL_LITERAL_TYPE, ), pytest.param(im.deref("x"), id="deref"), pytest.param(im.call("lift")("x"), id="lift"), @@ -138,6 +131,32 @@ ), id="program", ), + pytest.param(im.literal("1.0", "float32"), id="literal_float32"), + pytest.param(im.literal("1", "int64"), id="literal_int64"), + pytest.param(im.literal("1", "int8"), id="literal_int8"), + pytest.param(im.literal("1", "int16"), id="literal_int16"), + pytest.param(im.literal("True", "bool"), id="literal_bool"), + pytest.param( + ir.Temporary( + id="t", + domain=im.ref("domain"), + dtype=ts.TupleType( + types=[ + ts.ScalarType(kind=ts.ScalarKind.BOOL), + ts.ScalarType(kind=ts.ScalarKind.INT16), + ] + ), + ), + id="temporary_tuple_dtype", + ), + pytest.param( + ir.Temporary( + id="t", + domain=im.ref("domain"), + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=[3]), + ), + id="temporary_shaped_dtype", + ), pytest.param(ir.InfinityLiteral.POSITIVE, id="infinity_positive"), pytest.param(ir.InfinityLiteral.NEGATIVE, id="infinity_negative"), pytest.param(