resolve vector notation: refactored and extended#653
Draft
MichaelSt98 wants to merge 11 commits intomainfrom
Draft
resolve vector notation: refactored and extended#653MichaelSt98 wants to merge 11 commits intomainfrom
MichaelSt98 wants to merge 11 commits intomainfrom
Conversation
|
Documentation for this branch can be viewed at https://sites.ecmwf.int/docs/loki/653/index.html |
There was a problem hiding this comment.
Pull request overview
Refactors and extends the array-indexing vector-notation resolution utilities to support more cases (including offset handling) and to avoid transforming assignments that involve literal array constructors, with small backend formatting and test updates.
Changes:
- Extend
ResolveVectorNotationTransformer.visit_Assignmentwith new RHS/LHS range alignment logic (including shifted index mapping) and additional early-exit conditions. - Add
FindLiteralListsexpression visitor and use it to skip transforming assignments containing literal lists. - Adjust Fortran backend inline-IF formatting heuristic and add/extend vector-notation tests.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 9 comments.
| File | Description |
|---|---|
loki/transformations/array_indexing/vector_notation.py |
Major refactor of vector-notation resolution logic; adds new helper and literal-list skipping. |
loki/transformations/array_indexing/tests/test_vector_notation.py |
Adds new test scenarios (currently includes debug output and invalid Fortran fixtures). |
loki/ir/expr_visitors.py |
Introduces FindLiteralLists for locating LiteralList nodes in expression trees. |
loki/backend/fgen.py |
Tweaks when conditionals are formatted as inline single-line IF statements. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
Comment on lines
+387
to
+392
| # TODO: make this an option | ||
| if True: | ||
| rel_dim_indices = lhs_dims_ri | ||
| else: | ||
| rel_dim_indices = [i for i, j in enumerate(lhs_dims_ri) if all(array_dims[array_dims_gri[i]] != sym.RangeIndex((None, None)) | ||
| for array_dims, array_dims_gri in zip(arrays_dims, arrays_dims_gri))] |
Comment on lines
+429
to
+430
| # map this to the rhs arrays | ||
| rel_arrays_rhs_dims = [[array_dims[array_dims_gri[i]] for i in rel_dim_indices] for array_dims, array_dims_gri in zip(arrays_dims, arrays_dims_gri)] |
| for array, rel_rhs_dims in zip(arrays, rel_arrays_rhs_dims): | ||
| new_rel_rhs_dims = [] | ||
| for i, lhs_dim, new_lhs_dim, rhs_dim in zip(count(), rel_lhs_dims, new_lhs_dims, rel_rhs_dims): | ||
| if lhs_dim == rhs_dim or rhs_dim == sym.RangeIndex((None, None)) or isinstance(rhs_dim, sym.RangeIndex) and rhs_dim.lower == 1: |
Comment on lines
+449
to
+459
| horizontal = Dimension(name='horizontal', index='jl', lower=['start', 'alt_start'], upper=['end', 'alt_end'], size=['nlon']) | ||
| resolve_vector_dimension(routine, dimension=horizontal, derive_qualified_ranges=True) | ||
| # resolve_vector_dimension(routine, dimension=horizontal) | ||
| resolve_vector_notation(routine) | ||
| loops = FindNodes(ir.Loop).visit(routine.body) | ||
| print(f"{routine.to_fortran()}") | ||
| print(f"loops: {loops}") | ||
| print(f"{len(loops)}") | ||
| print(f"{loop.index: loop.bounds for loop in loops}") | ||
|
|
||
|
|
Comment on lines
+516
to
+520
| if (TRUE) local_var1(:, 1, ibl) = ptr_var1(5:klon+4, ibl) | ||
|
|
||
| local_var1(1:2, 1, ibl) = (/ 2.739,4.043 /) | ||
| local_var1(:, 1, ibl) = ptr_var1(5:klon+4, ibl) | ||
| local_var1(1, :, ibl) = ptr_var1(5:klon+4, ibl) |
Comment on lines
+538
to
+542
| print(f"{routine.to_fortran()}") | ||
| loops = FindNodes(ir.Loop).visit(routine.body) | ||
| print(f"{routine.to_fortran()}") | ||
| print(f"loops: {loops}") | ||
| print(f"{len(loops)}") |
Comment on lines
+135
to
+144
| Check loop bounds for a particular :any:`Dimension` in a | ||
| :any:`Subroutine`. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| routine : :any:`Subroutine` | ||
| Subroutine to perform checks on. | ||
| dimension : :any:`Dimension` | ||
| :any:`Dimension` object describing the variable conventions | ||
| used to define the data dimension and iteration space. |
Comment on lines
+491
to
+503
| integer :: jl, jk, ibl, jkglo, icend | ||
| integer :: ngptot | ||
|
|
||
| ngptot = ygeometry%ydim%ngptot | ||
|
|
||
| do jkglo=1,ngptot,nproma | ||
| ibl = (jgklo-1) / nproma+1 | ||
| icend = min(nproma, ngptot - jkglo + 1) | ||
| dims%klev = ygeometry%ydimv%len | ||
| dims%kidia = 1 | ||
| dims%kfdia = icend | ||
| dims%klon = ygeometry%ydim%nproma | ||
|
|
| """ | ||
| if o.inline: | ||
| # TODO: hack ..., this is a necessary but not a sufficient criteria | ||
| if o.inline and len(o.body) <= 1: |
…ved-type bound substitution - Add per-routine config flag 'resolve_vector_notation' to SCCBaseTransformation, following the same item.config pattern as demote_locals/preserve_arrays - Call resolve_vector_dimension + resolve_vector_notation in process_driver - Substitute derived-type members in shape-derived loop bounds with plain scalars (reuse existing scalar extractions or create new ones), so that generated loops are device-safe - Update test_scc_multiple_acc_pragmas pragma count (6->7) for the new loop generated from local_dims%wrk = 0 - Add tests: config-based disabling, derived-type bound substitution (existing scalar, new scalar, explicit range not substituted)
…ernel Re-add get_loop_bounds call in SCCBaseTransformation.process_kernel, wrapped in try/except RuntimeError -> TransformationError, so that routines lacking the expected horizontal loop bounds fail with a clear error rather than silently producing incorrect output. Fixes test_scc_base_horizontal_bounds_checks.
- Disable resolve_vector_notation in test_scc_hoist_multiple_kernels_loops config (full-array init 'tmp = 0.0' in driver was resolved into 3 extra loops, breaking the driver loop count assertion) - Fix TransformationError call in base.py: supply all 3 required arguments (message, transformation, source) to avoid pylint E1120 - Remove unused get_loop_bounds import from vector_notation.py - Remove trailing whitespace on blank lines (pylint C0303) - Replace dict comprehension with dict(zip(...)) (pylint R1721) - Remove unused 'assigns' variables in test_vector_notation.py - Remove spurious f-prefix from non-interpolated assertion strings
There was a problem hiding this comment.
Pull request overview
This PR refactors and extends Loki’s vector-notation resolution utilities, and wires them into the single-column (SCC) transformation flow with an opt-out per routine.
Changes:
- Add per-routine SCC config (
resolve_vector_notation) to enable/disable vector-notation resolution. - Major refactor/extension of
resolve_vector_notation/ResolveVectorNotationTransformerwith additional behaviors (eg. shifted ranges, derived-type bound handling) plus substantial new tests. - Minor backend formatting tweak for inline conditionals; adjust SCC tests for new pragma/loop outcomes.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
loki/transformations/single_column/vertical.py |
Adds runtime output when vertical dimension is set (currently via print). |
loki/transformations/single_column/base.py |
Adds per-routine config to skip vector-notation resolution and integrates resolve_vector_notation. |
loki/transformations/single_column/scc.py |
Adds SCCDemoteVerticalTransformation to SCC pipeline (but module appears missing). |
loki/transformations/single_column/__init__.py |
Re-exports demote_vertical (but module appears missing). |
loki/transformations/single_column/tests/test_scc.py |
Adds config-based test for disabling vector-notation resolution; updates pragma expectations. |
loki/transformations/single_column/tests/test_scc_hoist.py |
Updates scheduler config to disable vector-notation resolution for this test setup. |
loki/transformations/array_indexing/vector_notation.py |
Core refactor: new transformer logic, new options, derived-type bound substitutions, literal-list early exit. |
loki/transformations/array_indexing/tests/test_vector_notation.py |
Adds extensive new coverage for IFS-like patterns and edge cases. |
loki/ir/expr_visitors.py |
Introduces FindLiteralLists visitor used by the resolver. |
loki/backend/fgen.py |
Limits inline-conditional formatting to single-statement bodies. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if self.vertical is None: | ||
| info('[SCCFuseVerticalLoops] is not applied as the vertical dimension is not defined!') | ||
| else: | ||
| print(f"[SCCFuseVerticalLoops] vertical: {self.vertical}") |
There was a problem hiding this comment.
Avoid using print in transformations; it will spam stdout during pipelines and tests. Please use the existing Loki logger (eg. info/debug) or remove this debug output entirely.
Suggested change
| print(f"[SCCFuseVerticalLoops] vertical: {self.vertical}") | |
| info(f'[SCCFuseVerticalLoops] vertical: {self.vertical}') |
Comment on lines
19
to
23
| from loki.transformations.single_column.base import SCCBaseTransformation | ||
| from loki.transformations.single_column.annotate import SCCAnnotateTransformation | ||
| from loki.transformations.single_column.demote import SCCDemoteTransformation | ||
| from loki.transformations.single_column.demote_vertical import SCCDemoteVerticalTransformation | ||
| from loki.transformations.single_column.hoist import SCCHoistTemporaryArraysTransformation |
There was a problem hiding this comment.
This import will fail at runtime: loki/transformations/single_column/demote_vertical.py (and SCCDemoteVerticalTransformation) does not exist in this PR/repo tree. Either add the missing module/class or remove the import and pipeline entry.
| from loki.transformations.single_column.annotate import * # noqa | ||
| from loki.transformations.single_column.base import * # noqa | ||
| from loki.transformations.single_column.demote import * # noqa | ||
| from loki.transformations.single_column.demote_vertical import * # noqa |
There was a problem hiding this comment.
from loki.transformations.single_column.demote_vertical import * will raise ModuleNotFoundError because demote_vertical.py is not present in this PR/repo tree. Please add the missing module or remove this export.
Suggested change
| from loki.transformations.single_column.demote_vertical import * # noqa |
| ): | ||
| if (lhs_range == rhs_range | ||
| or rhs_range == sym.RangeIndex((None, None)) | ||
| or isinstance(rhs_range, sym.RangeIndex) and rhs_range.lower == 1): |
There was a problem hiding this comment.
rhs_range.lower == 1 is not a safe criterion for deciding that LHS/RHS ranges are unshifted. If LHS is e.g. 2:n and RHS is 1:n-1, this branch would incorrectly use the plain loop index (no offset) and generate wrong indexing. Consider comparing lhs_range.lower vs rhs_range.lower (and/or full range equality) instead.
Suggested change
| or isinstance(rhs_range, sym.RangeIndex) and rhs_range.lower == 1): | |
| or isinstance(lhs_range, sym.RangeIndex) | |
| and isinstance(rhs_range, sym.RangeIndex) | |
| and lhs_range.lower == rhs_range.lower): |
Comment on lines
+209
to
215
| _lower = as_tuple(dimension.lower) + ('1',) | ||
| _upper = as_tuple(dimension.upper) + as_tuple(dimension.sizes) | ||
| bounds = _get_all_valid_loop_bounds(routine, lower=_lower, upper=_upper) | ||
|
|
||
| # Map any range indices to the given loop index variable | ||
| loop_map = {sym.RangeIndex(bounds): index} | ||
| loop_map = {sym.RangeIndex(_bounds): index for _bounds in bounds} | ||
|
|
There was a problem hiding this comment.
resolve_vector_dimension no longer errors when no valid (lower, upper) bound pair is found (empty bounds => empty loop_map), so callers may silently get no transformation. This used to raise via get_loop_bounds; consider raising a RuntimeError/TransformationError when bounds is empty to avoid hard-to-debug no-ops.
Remove changes that were accidentally included from the unrelated demote_vertical.py development work: - Remove 'from demote_vertical import *' from single_column __init__.py - Remove SCCDemoteVerticalTransformation import from scc.py - Remove SCCDemoteVerticalTransformation from SCCSStackPipeline in scc.py - Remove debug print() from SCCFuseVerticalLoops.__init__ in vertical.py demote_vertical.py is not staged, committed, or planned for upstreaming.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #653 +/- ##
==========================================
+ Coverage 96.39% 96.41% +0.02%
==========================================
Files 266 266
Lines 46417 46930 +513
==========================================
+ Hits 44744 45249 +505
- Misses 1673 1681 +8
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Extend the pool of loop bound candidates to mirror resolve_vector_dimension:
- lower candidates: dimension.lower + '1' (numeric literal fallback)
- upper candidates: dimension.upper + dimension.sizes
This means get_loop_bounds succeeds in exactly the same cases where
resolve_vector_dimension can produce loops, e.g. for routines that
only carry a size variable (n) instead of an explicit upper loop bound.
The new extended_candidates parameter defaults to True; pass False to
restore the original strict behaviour (only dimension.lower/upper).
The bounds check in SCCBaseTransformation.process_kernel uses
extended_candidates=False since it validates that the actual iteration
bounds (start/end) are present, not merely that some bound can be found.
Updated test_transform_utilites_get_loop_bounds:
- The existing RuntimeError assertion now passes extended_candidates=False
(dimension 'y' with bounds ('a','b') not in routine now succeeds with
extended candidates, since '1' and 'n' are found as fallbacks)
- New cases: verify size-based upper bound fallback and '1' lower bound
fallback both with extended_candidates=True (default) and False
… improve coverage
- Fix unsafe rhs_range.lower == 1 condition in visit_Assignment Step 6:
replace with lhs_range.lower == rhs_range.lower comparison, which correctly
determines whether an offset is needed (if both ranges start at the same
value, the plain loop index can be reused without any offset computation).
The old condition incorrectly skipped offset computation when rhs.lower == 1
regardless of what the lhs lower bound was (e.g. LHS=2:n, RHS=1:n-1).
- Add warning() in resolve_vector_dimension when _get_all_valid_loop_bounds
returns an empty tuple (none of the dimension's lower/upper candidates exist
in the routine's scope). Previously, callers would silently get no
transformation; now a Loki warning message is emitted.
- Add three new tests to improve line coverage:
- test_resolve_vector_notation_derived_type_name_collision: exercises the
name-collision skip in _substitute_derived_type_bounds (line 543) where
a derived-type member basename collides with a different existing variable.
- test_resolve_vector_notation_no_implicit_rhs: exercises the
resolve_implicit_rhs_ranges=False branch (lines 622-628) via
resolve_vector_dimension, verifying bare ':' RHS assignments are skipped.
- test_resolve_vector_dimension_empty_bounds_warning: verifies that calling
resolve_vector_dimension with unknown bounds emits the new warning and
leaves the routine unchanged.
… loop_map When resolve_vector_dimension (map_unknown_ranges=False) encountered an LHS range not present in loop_map, _map_ranges_to_indices left new_lhs_dim as the RangeIndex itself. Step 6 then passed that RangeIndex into _compute_shifted_index, producing corrupt expressions like '-1 + (1:klevsn)' instead of '-1 + i_var'. Fix: after Step 5, filter out any (resolved_dim_index, new_lhs_dim) pair where new_lhs_dim is still a RangeIndex. If nothing survives the filter the assignment is returned unchanged, exactly as if the range had never been considered. This fixes both the MAX(zdsnnewg(jl,jk-1), zdsng(jl,0:klevsn-1)) pattern and the beta2alpha(overlap_param(jcol,:), frac(jcol,1:nlev-1), frac(jcol,2:nlev)) pattern. Two regression tests are added covering both cases with the two-step pipeline and resolve_vector_notation alone.
…types + optional flag) Bug 2a: _substitute_derived_type_bounds incorrectly included intermediate struct nodes (e.g. yrdimv from ydg%yrdimv%nflevg) in the derived_members filter because FindVariables returns the full parent chain. Fix: require v.type.dtype == BasicType.INTEGER so only leaf integer scalars are extracted. Bug 2b: _substitute_derived_type_bounds was always active, but should only run in driver contexts. Add substitute_derived_type_bounds=False parameter to resolve_vector_notation, resolve_vector_dimension, and ResolveVectorNotationTransformer. process_driver in base.py passes True. Tests: four existing derived-type tests updated to pass substitute_derived_type_bounds=True; new regression test test_resolve_vector_notation_derived_type_nested covers both kernel path (no scalar extraction) and driver path (only leaf integers extracted).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refactored and extended resolve vector notation utility.
Makes a bunch of assumptions which probably should be at least be optional to allow for more conservative transformations. Further, this does not work strided array dimensions.
Obviously some more rigorous testing required too.