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.

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

feat[next]: Infer as_fieldop type without domain #1853

Open
wants to merge 107 commits into
base: main
Choose a base branch
from

Conversation

SF-N
Copy link
Contributor

@SF-N SF-N commented Feb 7, 2025

This extends the GTIR type inference to infer the type of as_fieldop calls without a domain.

TODOs:

  • Fix the unnecessary promotion in the domain inference by using the now available type information.
  • as_fieldop with scan without domain, e.g., test_fieldop_from_scan:
    __field_operator_simple_scan_operator(out, __out_0_range) { out @ c⟨ KDimᵥ: [__out_0_range[0], __out_0_range[1][ ⟩ ← as_fieldop(scan(λ(carry) → carry + 1.0, True, 1.0), cartesian_domain())();
    has an empty domain. A special case is implemented for that, which ensures, that the vertical dimensions are kept.
  • test_cse.py:test_scalar_extraction_inside_as_fieldop results in an Assertion Error for assert isinstance(subexpr.type, ts.TypeSpec) because subexpr.type = None when running it or debugging without breakpoints. There is however no error when setting a breakpoint in type_synthesizer.py:applied_as_fieldop.
  • test_program.py::test_index_builtin[None] - TypeError: broadcast() missing 1 required positional argument: 'dimensions'
  • test_icon_like_scan.py::test_solve_nonhydro_stencil_52_like_z_q_tup[gtfn.run_gtfn] - NotImplementedError: Tuple expressions with different domains is not supported yet.
  • test_type_inference.py::test_function_definition - AssertionError: Node bar already has a type Field[[IDim], float64] which differs from (IteratorType(position_dims=[Dimension(value='IDim', kind=<DimensionKind.HORIZONTAL: 'horizontal'>)], defined_dims=[Dimension(value='IDim', kind=<DimensionKind.HORIZONTAL: 'horizontal'>)], element_type=ScalarType(kind=<ScalarKind.FLOAT64: 11>, shape=None))) -> float64.
  • test_domain_inference.py::test_nested_let_arg_shadowed2 - TypeError: Error during coercion of given value 'in_field1' for field 'Sym.id'
  • test_domain_inference.py everything with tuple_get
  • test_execution.py::test_zero_dim_tuple_arg[simple_mesh-dace.run_dace_cpu] - TypeError: Expected an array of index expressions: got empty array or None
  • test_execution.py::test_implicit_broadcast_mixed_dim[dace.run_dace_cpu] - ValueError: not enough values to unpack (expected 3, got 0)
  • test_gt4py_builtins.py::test_reduction_expression_in_call[simple_mesh-dace.run_dace_cpu_noopt] - ValueError: not enough values to unpack (expected 3, got 0)

@@ -492,10 +492,7 @@ def visit_SetAt(self, node: itir.SetAt, *, ctx) -> None:
# probably just change the behaviour of the lowering. Until then we do this more
# complicated comparison.
if isinstance(target_type, ts.FieldType) and isinstance(expr_type, ts.FieldType):
assert (
set(expr_type.dims).issubset(set(target_type.dims))
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 think the subset requirement is not necessary anymore, e.g. in the case

itir.SetAt(
        domain=unstructured_domain,
        expr=im.as_fieldop(
            im.lambda_("it")(im.reduce("plus", 0.0)(im.deref("it"))),
            unstructured_domain,
        )(im.ref("inp")),
        target=im.ref("out"),
)

with
im.sym("inp", float_vertex_v2e_k_field) and im.sym("out", float_vertex_k_field)

the expr_type is

[Dimension(value='Vertex', kind=<DimensionKind.HORIZONTAL: 'horizontal'>), Dimension(value='KDim', kind=<DimensionKind.VERTICAL: 'vertical'>), Dimension(value='V2E', kind=<DimensionKind.LOCAL: 'local'>)]

and the target_type is

[Dimension(value='Vertex', kind=<DimensionKind.HORIZONTAL: 'horizontal'>), Dimension(value='KDim', kind=<DimensionKind.VERTICAL: 'vertical'>)]

which is correct in my opinion.

cf. test_fencil_with_nb_field_input

)(im.ref("inp1", float_i_field), im.ref("inp2", float_i_field)),
ts.TupleType(types=[float_i_field, float_i_field]),
)(im.ref("inp1", float_i_field), im.ref("inp2", float_j_field)),
ts.TupleType(types=[float_ij_field, float_ij_field]),
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Is ts.TupleType(types=[float_ij_field, float_ij_field] correct here?

@SF-N SF-N requested a review from tehrengruber February 13, 2025 11:54
return (
sorted(unique_dims, key=lambda dim: (kind_order[dim.kind], dim.value))
if unique_dims
else []
Copy link
Contributor Author

Choose a reason for hiding this comment

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

check this

SF-N added a commit that referenced this pull request Mar 19, 2025
Previously, a broadcast generated an `as_fieldop`. In
[PR#1853](#1853), the domain
restriction returns a smaller domain for the `as_fieldop`.
In order to avoid this, we introduce the broadcast builtin here and
transform it into an `im.as_fieldop("deref", ...)` call after the domain
inference.

---------

Co-authored-by: Till Ehrengruber <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants