Support Fortran sum for arrays without explicit dimension access decl…

…aration
spcl · Oct 12, 2023 · 6544673 · 6544673
1 parent 81786a8
commit 6544673
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/dace/frontend/fortran/ast_transforms.py b/dace/frontend/fortran/ast_transforms.py
@@ -941,7 +941,12 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No
                             if isinstance(arg, ast_internal_classes.Name_Node):
                                 array_node = ast_internal_classes.Array_Subscript_Node(parent=arg.parent)
                                 array_node.name = arg
-                                array_node.indices = [ast_internal_classes.ParDecl_Node(type='ALL')]
+
+                                # If we access SUM(arr) where arr has many dimensions,
+                                # We need to create a ParDecl_Node for each dimension
+                                dims = len(self.scope_vars.get_var(node.parent, arg.name).sizes)
+                                array_node.indices = [ast_internal_classes.ParDecl_Node(type='ALL')] * dims
+
                                 rvals.append(array_node)
 
                             # supports syntax SUM(arr(:))

diff --git a/tests/fortran/sum_to_loop_offset.py b/tests/fortran/sum_to_loop_offset.py
@@ -98,7 +98,7 @@ def test_fortran_frontend_arr2loop_2d():
                     double precision, dimension(5,3) :: d
                     double precision, dimension(4) :: res
 
-                    !res(1) = SUM(d)
+                    res(1) = SUM(d)
                     res(2) = SUM(d(:,:))
                     res(3) = SUM(d(2:4, 2))
                     res(4) = SUM(d(2:4, 2:3))
@@ -121,10 +121,12 @@ def test_fortran_frontend_arr2loop_2d():
             cnt += 1
     res = np.full([4], 42, order="F", dtype=np.float64)
     sdfg(d=d, res=res)
+    assert res[0] == 105
     assert res[1] == 105
     assert res[2] == 21
     assert res[3] == 45
 
+
 if __name__ == "__main__":
 
     test_fortran_frontend_sum2loop_1d_without_offset()