feat[next]: SDFGConvertible Program for dace_fieldview backend #1742
Conversation
src/gt4py/next/program_processors/runners/dace_fieldview/program.py
Outdated
Show resolved
Hide resolved
src/gt4py/next/program_processors/runners/dace_fieldview/program.py
Outdated
Show resolved
Hide resolved
| if conn_id not in self.connectivity_tables_data_descriptors: | ||
| conn = self.connectivities[name] | ||
| self.connectivity_tables_data_descriptors[conn_id] = dace.data.Array( | ||
| dtype=dace.int64 if conn.index_type == np.int64 else dace.int32, |
There was a problem hiding this comment.
I think that dace dtypes has a utility to parse numpy types:
dace.dtypes.typeclass(conn.index_type)
There was a problem hiding this comment.
Trying that lead to KeyError: dtype('int64') in typeclass().
There was a problem hiding this comment.
and trying dace.dtypes.typeclass(conn.index_type.type) leads to wrong stencil results. Edit: this might have instead been caused by auto_optimize=False.
There was a problem hiding this comment.
edit: after merging main, it passes. However according to type hints of NeighborTableOffsetProvider, your version should work and mine should fail. Yet, in the tests, the connectivities are set up in a way so that my version works and yours fails. The current version on the other hand is safe in both cases.
If I can get the test to produce the correct types so the official version works, I will change it. Otherwise I would prefer to leave this one until the neighbor table construction is anyway changed by what @havogt is working on.
There was a problem hiding this comment.
Could not get it to work, it seems like the type hints are wrong, but there is no point investing time in this when the API will change heavily anyway.
src/gt4py/next/program_processors/runners/dace_fieldview/program.py
Outdated
Show resolved
Hide resolved
| for in_field in closure.inputs | ||
| if str(in_field.id) in fields # type: ignore[union-attr] # ensured by assert | ||
| ] | ||
| sdfg.gt4py_program_input_fields = { |
There was a problem hiding this comment.
Could you please port this dynamic property to GTIR as well.
There was a problem hiding this comment.
Working on it, Program nodes will require a different approach.
| for arg in output.args: | ||
| if str(arg.id) in fields: # type: ignore[attr-defined] | ||
| output_fields.append(str(arg.id)) # type: ignore[attr-defined] | ||
| sdfg.gt4py_program_output_fields = { |
There was a problem hiding this comment.
Same for this property.
There was a problem hiding this comment.
*edit, moved this comment to the offset_provider_per_input_field property below.
| continue | ||
| if param.id not in sdfg.gt4py_program_input_fields: | ||
| continue | ||
| sdfg.offset_providers_per_input_field.setdefault(param.id, []).extend(list(shifts)) |
There was a problem hiding this comment.
Same here.
All these properties are needed for the automatic halo exchange placement.
There was a problem hiding this comment.
Results from discussion in a separate channel for "offset_providers_per_input_field":
- it is correct, as only used in
icon4pywhich uses "unstructured" fields. There is max one horizontal dimension per field. - trace_shifts should work™️ partially on GTIR (but not accross
as_fieldop). - currently not actively used in
icon4py, so defer implementation and put a TODO instead.
| return (args, []) | ||
|
|
||
|
|
||
| def _crosscheck_dace_parsing(dace_parsed_args: list[Any], gt4py_program_args: list[Any]) -> bool: |
There was a problem hiding this comment.
This function also needs to be ported to GTIR as well.
src/gt4py/next/program_processors/runners/dace_fieldview/program.py
Outdated
Show resolved
Hide resolved
| # Add them as dynamic properties to the SDFG | ||
| program = typing.cast( | ||
| itir.Program, gtir_stage.data | ||
| ) # we already checked that our backend uses GTIR |
There was a problem hiding this comment.
we already checked that our backend uses GTIR
where? otherwise we could make an assert isinstance().
| } | ||
|
|
||
| input_fields = (field_params[name] for name in InputNamesExtractor.only_fields(program)) | ||
| # TODO (ricoh): This will associate the last horizontal dimension of each field with it's name |
There was a problem hiding this comment.
My suspect is that this works under the assumption that all fields have a single horizontal dimension.
@kotsaloscv If this is the case, should we make a check and raise an exception if this assumption is not met?
Description
Add a
decrator.Programsubclass, which implementsSDFGConvertibletodace_fieldviewbackend, analogous to the one indace_iterator. Conditionally shadowdecorator.Programwith it and reactivate the orchestration tests by usingdace_fieldviewinstead ofdace_iterator.One caveat: The toolchain is not ready for pure
CompileTimeConnectivitiesin all cases yet, so thetest_sdfgConvertible_connectivitieshad to be adjusted for the moment.Requirements
If this PR contains code authored by new contributors please make sure:
AUTHORS.mdfile adding the names of all the new contributors.