Skip to content

Add core infrastructure fixes for small-kernels pipeline#694

Open
MichaelSt98 wants to merge 1 commit into
mainfrom
nams-sk-infra
Open

Add core infrastructure fixes for small-kernels pipeline#694
MichaelSt98 wants to merge 1 commit into
mainfrom
nams-sk-infra

Conversation

@MichaelSt98
Copy link
Copy Markdown
Collaborator

  • Fix SubstituteExpressionsMapper to also register the underlying VariableSymbol when the expression map key is a MetaSymbol, so that derived-type parent substitution (e.g. geom -> geom_new propagating to geom%dim%nproma) works correctly.
  • Add a recurse_var_parent flag to LokiWalkMapper (default True, backward compatible) and a new FindUsedVariables expression finder that collects variables as they are used without decomposing derived-type parents into separate entries.
  • Add get_block_index to TemporariesPoolAllocatorTransformation so the pool allocator can resolve derived-type block-dimension indices (e.g. geom%blk_dim%ib) instead of only matching simple variable names. Add ignore_index_undefined flag to promote_variables so callers can keep an explicit index expression even when dataflow analysis cannot prove it is live at every use site.

Description

This pull request introduces several enhancements and fixes to variable handling in expression mapping, array promotion, and temporary pool allocator transformations. The main changes focus on providing more precise variable collection, supporting explicit index retention in array promotions, and improving the resolution of derived-type block indices. Comprehensive tests are added to ensure the correctness of these new behaviors.

Expression mapping and variable collection:

  • Added a new FindUsedVariables visitor in loki/ir/expr_visitors.py that collects only the used member expressions without recursing into their parent chains, improving the precision of variable usage analysis. ([loki/ir/expr_visitors.pyR187-R198](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-cd38c31e937a1478653d40a9f97adea800fadd922039b18b42b013e9b2b6c4b5R187-R198))
  • Updated LokiWalkMapper to allow controlling recursion into parent components via a new recurse_var_parent argument. ([loki/expression/mappers.pyR229-R236](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-86dcd26846fd8111831a312d147a885c713b6397e5d43bdf4180abebdba96b11R229-R236))
  • Added a test to verify that FindUsedVariables skips parent components as intended. ([loki/expression/tests/test_expression.pyR1381-R1402](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-edcf72b91a67e888e51fc7cabcd6cb87e6e4b561a8e2d60c38e2e3358d816025R1381-R1402))

Array promotion transformation:

  • Enhanced promote_variables to accept an ignore_index_undefined argument, allowing explicit retention of unresolved promotion indices when requested. ([[1]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-8afd2442c2d2f132b71a1eb4c2ba566fe746d19438462919cc0cd58b67cd2f3bL30-R30), [[2]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-8afd2442c2d2f132b71a1eb4c2ba566fe746d19438462919cc0cd58b67cd2f3bR59-R61), [[3]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-8afd2442c2d2f132b71a1eb4c2ba566fe746d19438462919cc0cd58b67cd2f3bR86-R90))
  • Added a test to ensure unresolved indices are kept verbatim when ignore_index_undefined=True. ([loki/transformations/array_indexing/tests/test_array_promote.pyR144-R170](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-25ca74992672348eefeb9c76bc8d9a890c46cdca3bf7b08cbbce4dd035e529d5R144-R170))

Temporary pool allocator transformation:

  • Added logic to resolve derived-type block indices via configured dimension aliases in get_block_index, improving support for more complex block index forms. ([[1]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-dee9adcde30167e25ddb24c4b6de20aae9bcba38017850a137261592e44be4bcR410-R423), [[2]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-dee9adcde30167e25ddb24c4b6de20aae9bcba38017850a137261592e44be4bcL858-R872))
  • Added a test to verify that derived-type block indices are resolved as intended. ([loki/transformations/temporaries/tests/test_pool_allocator.pyR1432-R1452](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-a999af48b93ac47f0c05921f255e19c790f6e7fd1a22a9927a6addf758d93400R1432-R1452))

Other improvements:

  • Improved SubstituteExpressionsMapper to handle MetaSymbol keys by also mapping their underlying symbols. ([loki/expression/mappers.pyR752-R757](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-86dcd26846fd8111831a312d147a885c713b6397e5d43bdf4180abebdba96b11R752-R757))
  • Minor import and test updates to support the new features. ([[1]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-edcf72b91a67e888e51fc7cabcd6cb87e6e4b561a8e2d60c38e2e3358d816025L27-R27), [[2]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-cd38c31e937a1478653d40a9f97adea800fadd922039b18b42b013e9b2b6c4b5L29-R29), [[3]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-25ca74992672348eefeb9c76bc8d9a890c46cdca3bf7b08cbbce4dd035e529d5R15), [[4]](https://github.com/ecmwf-ifs/loki/pull/694/files#diff-a999af48b93ac47f0c05921f255e19c790f6e7fd1a22a9927a6addf758d93400L11-R11))

Contributor Declaration

By opening this pull request, I affirm the following:

  • All authors agree to the Contributor License Agreement.
  • The code follows the project's coding standards.
  • I have performed self-review and added comments where needed.
  • I have added or updated tests to verify that my changes are effective and functional.
  • I have run all existing tests and confirmed they pass.

Fix SubstituteExpressionsMapper to also register the underlying
VariableSymbol when the expression map key is a MetaSymbol, so that
derived-type parent substitution (e.g. geom -> geom_new propagating to
geom%dim%nproma) works correctly.
Add a recurse_var_parent flag to LokiWalkMapper (default True, backward
compatible) and a new FindUsedVariables expression finder that collects
variables as they are used without decomposing derived-type parents into
separate entries.
Add get_block_index to TemporariesPoolAllocatorTransformation so the
pool allocator can resolve derived-type block-dimension indices (e.g.
geom%blk_dim%ib) instead of only matching simple variable names.
Add ignore_index_undefined flag to promote_variables so callers can keep
an explicit index expression even when dataflow analysis cannot prove it
is live at every use site.
@github-actions
Copy link
Copy Markdown

github-actions Bot commented Jun 1, 2026

Documentation for this branch can be viewed at https://sites.ecmwf.int/docs/loki/694/index.html

@codecov
Copy link
Copy Markdown

codecov Bot commented Jun 1, 2026

Codecov Report

❌ Patch coverage is 98.03922% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.19%. Comparing base (bfad794) to head (f72689b).
⚠️ Report is 13 commits behind head on main.

Files with missing lines Patch % Lines
loki/transformations/temporaries/pool_allocator.py 87.50% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #694      +/-   ##
==========================================
- Coverage   94.19%   94.19%   -0.01%     
==========================================
  Files         288      288              
  Lines       47496    47671     +175     
==========================================
+ Hits        44741    44904     +163     
- Misses       2755     2767      +12     
Flag Coverage Δ
lint_rules 96.40% <ø> (ø)
loki 94.16% <98.03%> (-0.01%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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.

1 participant