Sanitise: Resolve free range indices when resolving associates

[mlange05]

The current Associate resolution does not account for array index ranges if they are specific in the associate clauses (index remapping). This PR adds a partial resolution step to ResolveAssociates, which resolves free index ranges (:) in association clause against explicit indices from the original expression, should should both exist. The respective changes in the expression mapper require some care with respect to early termination and recursion into .dimension and .type properties of the given array expression.

Note: It does not resolve index shifts, where offsets are introduced by using bounded ranges (eg. 3:8). Instead it warns about such (ab)use.

[github-actions]

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

[codecov]

Codecov Report

Attention: Patch coverage is 98.00000% with 1 line in your changes missing coverage. Please review.

Project coverage is 93.29%. Comparing base (e9760e8) to head (e350f50).
Report is 19 commits behind head on main.

Files with missing lines	Patch %	Lines
loki/transformations/sanitise/associates.py	95.23%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #455   +/-   ##
=======================================
  Coverage   93.28%   93.29%           
=======================================
  Files         220      220           
  Lines       41224    41269   +45     
=======================================
+ Hits        38457    38502   +45     
  Misses       2767     2767           

Flag	Coverage Δ
lint_rules	96.39% <ø> (ø)
loki	93.25% <98.00%> (+<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.

[reuterbal]

Thanks for this! I have some questions about the recursion and potentially a missing test but not sure on either, so marking this as a comment only for now

[reuterbal]

I'm not sure I'm following the logic here. The first recursion to expr_dims seems to only be used if the array's scope isn't an associate - could we not move it into the conditional just before the return statement?

Or alternatively, in which situation do we expect the return value of the recursion on new.dimensions to differ from the first recursion? Could we maybe use the previous recursion value here?

I suspect the second wouldn't work but the first should allow us to save some recursion maybe?

[mlange05]

Ah, yes, this was quite fiddly. The initial recursion is required for resolving associations in the dimensions of arrays that are not actually scope to the Associate; that's also why it has to happen before the short-circuiting. For example, consider

associate(a => b)
  some%obj(i, :, a) = ...
end associate

The second recursion is then done on the potentially replaced symbol (new), which now might have entirely different .dimensions. For example consider

associate(a => b(:, : i))
  a = 42.0
  ! or even
  a(j, k) = 66.6
end associate

In the first line of this example expr == 'a' and expr.dimensions == None and new.dimensions == (:, :, i), the latter of which we have not recursed into yet. In the second example line new.dimensions then also needs to be matched against (j, k) to yield b(j, k, i). Does that make sense?

[reuterbal]

EDIT: Scratch this, I'm seeing my error 🤦

---

Thanks, I suspected something along these lines.

My question was mostly on the grounds of avoiding redundant recursion, purely judged from the dependencies within this method. Something like this:

        # Recurse over the type's shape
        _type = expr.type
        if _type.shape:
            new_shape = self.rec(_type.shape, *args, **kwargs)
            _type = expr.type.clone(shape=new_shape)

        # Stop if scope is not an associate, but still recurse on the dimensions
        # in case an associated symbol is in there
        if not isinstance(expr.scope, ir.Associate):
            expr_dims = self.rec(expr.dimensions, *args, **kwargs)
            return expr.clone(dimensions=expr_dims, type=_type)

        # For arrays that represent an associate symbol, re-use the
        # scalar implementation to resolve the symbol
        new = self.map_scalar(expr, *args, **kwargs)

        # Recurse over array dimensions
        if isinstance(new, sym.Array) and new.dimensions:
            # Resolve unbound range symbols form existing indices
            new_dims = self.rec(new.dimensions, *args, **kwargs)
            new_dims = self._match_range_indices(new_dims, expr_dims)
        else:
            new_dims = expr_dims
        return new.clone(dimensions=new_dims, type=_type)

This should save the need to do the recursion on the dimensions for symbols from the associate since you're recursing on the new_dims anyway

[reuterbal]

Do we also have a test that covers the reverse situation, where a dimension expression contains an associated symbol but the array is defined outside?

Something like:

real, intent(inout) :: arr2d(:,:)

associate(var_idx=>some%type%idx)
  do i=1,5
    arr2d(i, var_idx) = 1.0
  end do
end associate

[mlange05]

Yes, agreed, good point! I've extended the array slicing test to cover this example.

[mlange05]

Hi @reuterbal, thanks for the review and apologies for the delay. In addition to the things you flagged, I've independently found another bug in this, which I've now fixed and extended the slice array test for. This pertains to the matching of ':' when matching the free range symbols in the utility method of the mapper class. Please have another look. 🙏

[reuterbal]

Many thanks for the additional testing and explanations. I was a bit slow but finally got the double recursion... (scratch my comment)