use 'parse_expr' in 'process_dimension_pragmas'

Author: MichaelSt98

loki/frontend/fparser.py

@@ -1864,7 +1864,7 @@ def visit_Subroutine_Subprogram(self, o, **kwargs):
 
         # Update array shapes with Loki dimension pragmas
         with pragmas_attached(routine, ir.VariableDeclaration):
-            routine.spec = process_dimension_pragmas(routine.spec)
+            routine.spec = process_dimension_pragmas(routine.spec, scope=routine)
 
         if isinstance(o, Fortran2003.Subroutine_Body):
             # Return the subroutine object along with any clutter before it for interface declarations

loki/frontend/ofp.py

@@ -1230,7 +1230,7 @@ def visit_subroutine(self, o, **kwargs):
 
         # Update array shapes with Loki dimension pragmas
         with pragmas_attached(routine, ir.VariableDeclaration):
-            routine.spec = process_dimension_pragmas(routine.spec)
+            routine.spec = process_dimension_pragmas(routine.spec, scope=routine)
 
         return routine
 

loki/frontend/omni.py

@@ -433,7 +433,7 @@ def visit_FfunctionDefinition(self, o, **kwargs):
 
         # Update array shapes with Loki dimension pragmas
         with pragmas_attached(routine, ir.VariableDeclaration):
-            routine.spec = process_dimension_pragmas(routine.spec)
+            routine.spec = process_dimension_pragmas(routine.spec, scope=routine)
 
         return routine
 

loki/frontend/tests/test_frontends.py

@@ -1970,3 +1970,64 @@ def test_source_sanitize_fp_module(preprocess):
 
 
 # TODO: Add tests for source sanitizer with other frontends
+
+
+@pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'OMNI does not like Loki pragmas, yet!')]))
+def test_frontend_routine_variables_dimension_pragmas(frontend):
+    """
+    Test that `!$loki dimension` pragmas can be used to verride the
+    conceptual `.shape` of local and argument variables.
+    """
+    fcode = """
+subroutine routine_variables_dimensions(x, y, v0, v1, v2, v3, v4)
+  integer, parameter :: jprb = selected_real_kind(13,300)
+  integer, intent(in) :: x, y
+
+  !$loki dimension(10)
+  real(kind=jprb), intent(inout) :: v0(:)
+  !$loki dimension(x)
+  real(kind=jprb), intent(inout) :: v1(:)
+  !$loki dimension(x,y,:)
+  real(kind=jprb), dimension(:,:,:), intent(inout) :: v2, v3
+  !$loki dimension(x,y)
+  real(kind=jprb), pointer, intent(inout) :: v4(:,:)
+  !$loki dimension(x+y,2*x)
+  real(kind=jprb), allocatable :: v5(:,:)
+  !$loki dimension(x/2, x**2, (x+y)/x)
+  real(kind=jprb), dimension(:, :, :), pointer :: v6
+
+end subroutine routine_variables_dimensions
+"""
+    def to_str(expr):
+        return str(expr).lower().replace(' ', '')
+
+    routine = Subroutine.from_source(fcode, frontend=frontend)
+    assert routine.variable_map['v0'].shape[0] == 10
+    assert isinstance(routine.variable_map['v0'].shape[0], sym.IntLiteral)
+    assert isinstance(routine.variable_map['v1'].shape[0], sym.Scalar)
+    assert routine.variable_map['v2'].shape[0] == 'x'
+    assert routine.variable_map['v2'].shape[1] == 'y'
+    assert routine.variable_map['v2'].shape[2] == ':'
+    assert isinstance(routine.variable_map['v2'].shape[0], sym.Scalar)
+    assert isinstance(routine.variable_map['v2'].shape[1], sym.Scalar)
+    assert isinstance(routine.variable_map['v2'].shape[2], sym.RangeIndex)
+    assert routine.variable_map['v3'].shape[0] == 'x'
+    assert routine.variable_map['v3'].shape[1] == 'y'
+    assert routine.variable_map['v3'].shape[2] == ':'
+    assert isinstance(routine.variable_map['v3'].shape[0], sym.Scalar)
+    assert isinstance(routine.variable_map['v3'].shape[1], sym.Scalar)
+    assert isinstance(routine.variable_map['v3'].shape[2], sym.RangeIndex)
+    assert routine.variable_map['v4'].shape[0] == 'x'
+    assert routine.variable_map['v4'].shape[1] == 'y'
+    assert isinstance(routine.variable_map['v4'].shape[0], sym.Scalar)
+    assert isinstance(routine.variable_map['v4'].shape[1], sym.Scalar)
+    assert to_str(routine.variable_map['v5'].shape[0]) == 'x+y'
+    assert to_str(routine.variable_map['v5'].shape[1]) == '2*x'
+    assert isinstance(routine.variable_map['v5'].shape[0], sym.Sum)
+    assert isinstance(routine.variable_map['v5'].shape[1], sym.Product)
+    assert to_str(routine.variable_map['v6'].shape[0]) == 'x/2'
+    assert to_str(routine.variable_map['v6'].shape[1]) == 'x**2'
+    assert to_str(routine.variable_map['v6'].shape[2]) == '(x+y)/x'
+    assert isinstance(routine.variable_map['v6'].shape[0], sym.Quotient)
+    assert isinstance(routine.variable_map['v6'].shape[1], sym.Power)
+    assert isinstance(routine.variable_map['v6'].shape[2], sym.Quotient)

loki/ir/pragma_utils.py

@@ -14,9 +14,8 @@
 from loki.ir.find import FindNodes
 from loki.ir.transformer import Transformer
 from loki.ir.visitor import Visitor
-
+# from loki.expression.parser import parse_expr
 from loki.tools.util import as_tuple, replace_windowed
-from loki.types import BasicType
 from loki.logging import debug, warning
 
 
@@ -51,7 +50,7 @@ def is_loki_pragma(pragma, starts_with=None):
     return True
 
 
-_get_pragma_parameters_re = re.compile(r'(?P<command>[\w-]+)\s*(?:\((?P<arg>.+?)\))?')
+_get_pragma_parameters_re = re.compile(r'(?P<command>[\w-]+)\s*(?:\((?P<arg>.*)\))?')
 
 def get_pragma_parameters(pragma, starts_with=None, only_loki_pragmas=True):
     """
@@ -94,7 +93,7 @@ def get_pragma_parameters(pragma, starts_with=None, only_loki_pragmas=True):
     return parameters
 
 
-def process_dimension_pragmas(ir):
+def process_dimension_pragmas(ir, scope=None):
     """
     Process any ``!$loki dimension`` pragmas to override deferred dimensions
 
@@ -106,21 +105,18 @@ def process_dimension_pragmas(ir):
     ir : :any:`Node`
         Root node of the (section of the) internal representation to process
     """
-    from loki.expression.symbols import Literal, Variable  # pylint: disable=import-outside-toplevel
+    from loki.expression.parser import parse_expr  # pylint: disable=import-outside-toplevel
 
     for decl in FindNodes(VariableDeclaration).visit(ir):
         if is_loki_pragma(decl.pragma, starts_with='dimension'):
             for v in decl.symbols:
                 # Found dimension override for variable
                 dims = get_pragma_parameters(decl.pragma)['dimension']
                 dims = [d.strip() for d in dims.split(',')]
-                shape = []
-                for d in dims:
-                    if d.isnumeric():
-                        shape += [Literal(value=int(d), type=BasicType.INTEGER)]
-                    else:
-                        shape += [Variable(name=d, scope=v.scope)]
-                v.scope.symbol_attrs[v.name] = v.type.clone(shape=as_tuple(shape))
+                # parse each dimension
+                shape = tuple(parse_expr(d, scope=scope) for d in dims)
+                # update symbol table
+                v.scope.symbol_attrs[v.name] = v.type.clone(shape=shape)
     return ir