Bugfixes and extended testing for Fortran SUM (#1390)

* Fix incorrect generation of sum to loop code for Fortran frontend

* Support passing array with no bounds in Fortran sum()

* Add test case for Foftran sum

* Fix bug in offset normalization and support Fortran SUM for arrays with offsets

* Expand tests for array2loop in Fortran

* Add more tests covering 2D sum in Fortran

* Support Fortran sum for arrays without explicit dimension access declaration

* Add more tests for Fortran sum over 2D arrays

---------

Co-authored-by: acalotoiu <[email protected]>

dace/frontend/fortran/ast_transforms.py

@@ -268,7 +268,7 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No
                                 ast_internal_classes.Var_Decl_Node(
                                     name="tmp_call_" + str(temp),
                                     type=res[i].type,
-                                    sizes=None,
+                                    sizes=None
                                 )
                             ]))
                         newbody.append(
@@ -284,7 +284,7 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No
                                 ast_internal_classes.Var_Decl_Node(
                                     name="tmp_call_" + str(temp),
                                     type=res[i].type,
-                                    sizes=None,
+                                    sizes=None
                                 )
                             ]))
                     newbody.append(
@@ -458,7 +458,11 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No
                             if self.normalize_offsets:
 
                                 # Find the offset of a variable to which we are assigning
-                                var_name = child.lval.name.name
+                                var_name = ""
+                                if isinstance(j, ast_internal_classes.Name_Node):
+                                    var_name = j.name
+                                else:
+                                    var_name = j.name.name
                                 variable = self.scope_vars.get_var(child.parent, var_name)
                                 offset = variable.offsets[idx]
 
@@ -737,8 +741,7 @@ def par_Decl_Range_Finder(node: ast_internal_classes.Array_Subscript_Node,
                           count: int,
                           newbody: list,
                           scope_vars: ScopeVarsDeclarations,
-                          declaration=True,
-                          is_sum_to_loop=False):
+                          declaration=True):
     """
     Helper function for the transformation of array operations and sums to loops
     :param node: The AST to be transformed
@@ -753,6 +756,7 @@ def par_Decl_Range_Finder(node: ast_internal_classes.Array_Subscript_Node,
 
     currentindex = 0
     indices = []
+
     offsets = scope_vars.get_var(node.parent, node.name.name).offsets
 
     for idx, i in enumerate(node.indices):
@@ -926,14 +930,36 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No
 
                 current = child.lval
                 val = child.rval
-                rvals = [i for i in mywalk(val) if isinstance(i, ast_internal_classes.Array_Subscript_Node)]
+
+                rvals = []
+                for i in mywalk(val):
+                    if isinstance(i, ast_internal_classes.Call_Expr_Node) and i.name.name == '__dace_sum':
+
+                        for arg in i.args:
+
+                            # supports syntax SUM(arr)
+                            if isinstance(arg, ast_internal_classes.Name_Node):
+                                array_node = ast_internal_classes.Array_Subscript_Node(parent=arg.parent)
+                                array_node.name = arg
+
+                                # 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(:))
+                            if isinstance(arg, ast_internal_classes.Array_Subscript_Node):
+                                rvals.append(arg)
+
                 if len(rvals) != 1:
                     raise NotImplementedError("Only one array can be summed")
                 val = rvals[0]
                 rangeposrval = []
                 rangesrval = []
 
-                par_Decl_Range_Finder(val, rangesrval, rangeposrval, self.count, newbody, self.scope_vars, False, True)
+                par_Decl_Range_Finder(val, rangesrval, rangeposrval, self.count, newbody, self.scope_vars, True)
 
                 range_index = 0
                 body = ast_internal_classes.BinOp_Node(lval=current,

tests/fortran/array_to_loop_offset.py

@@ -112,8 +112,112 @@ def test_fortran_frontend_arr2loop_2d_offset():
         for j in range(7,10):
             assert a[i-1, j-1] == i * 2
 
+def test_fortran_frontend_arr2loop_2d_offset2():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(5,7:9) :: d
+                    CALL index_test_function(d)
+                    end
+
+                    SUBROUTINE index_test_function(d)
+                    double precision, dimension(5,7:9) :: d
+
+                    d(:,:) = 43
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", False)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    assert len(sdfg.data('d').shape) == 2
+    assert sdfg.data('d').shape[0] == 5
+    assert sdfg.data('d').shape[1] == 3
+
+    a = np.full([5,9], 42, order="F", dtype=np.float64)
+    sdfg(d=a)
+    for i in range(1,6):
+        for j in range(7,10):
+            assert a[i-1, j-1] == 43
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    a = np.full([5,3], 42, order="F", dtype=np.float64)
+    sdfg(d=a)
+    for i in range(0,5):
+        for j in range(0,3):
+            assert a[i, j] == 43
+
+def test_fortran_frontend_arr2loop_2d_offset3():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(5,7:9) :: d
+                    CALL index_test_function(d)
+                    end
+
+                    SUBROUTINE index_test_function(d)
+                    double precision, dimension(5,7:9) :: d
+
+                    d(2:4, 7:8) = 43
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", False)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    assert len(sdfg.data('d').shape) == 2
+    assert sdfg.data('d').shape[0] == 5
+    assert sdfg.data('d').shape[1] == 3
+
+    a = np.full([5,9], 42, order="F", dtype=np.float64)
+    sdfg(d=a)
+    for i in range(2,4):
+        for j in range(7,9):
+            assert a[i-1, j-1] == 43
+        for j in range(9,10):
+            assert a[i-1, j-1] == 42
+
+    for i in [1, 5]:
+        for j in range(7,10):
+            assert a[i-1, j-1] == 42
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    a = np.full([5,3], 42, order="F", dtype=np.float64)
+    sdfg(d=a)
+    for i in range(1,4):
+        for j in range(0,2):
+            assert a[i, j] == 43
+        for j in range(2,3):
+            assert a[i, j] == 42
+
+    for i in [0, 4]:
+        for j in range(0,3):
+            assert a[i, j] == 42
+
 if __name__ == "__main__":
 
     test_fortran_frontend_arr2loop_1d_offset()
     test_fortran_frontend_arr2loop_2d_offset()
+    test_fortran_frontend_arr2loop_2d_offset2()
+    test_fortran_frontend_arr2loop_2d_offset3()
     test_fortran_frontend_arr2loop_without_offset()

tests/fortran/sum_to_loop_offset.py

@@ -0,0 +1,176 @@
+# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved.
+
+import numpy as np
+
+from dace.frontend.fortran import ast_transforms, fortran_parser
+
+def test_fortran_frontend_sum2loop_1d_without_offset():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(7) :: d
+                    double precision, dimension(3) :: res
+                    CALL index_test_function(d, res)
+                    end
+
+                    SUBROUTINE index_test_function(d, res)
+                    double precision, dimension(7) :: d
+                    double precision, dimension(3) :: res
+
+                    res(1) = SUM(d(:))
+                    res(2) = SUM(d)
+                    res(3) = SUM(d(2:6))
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", False)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    size = 7
+    d = np.full([size], 0, order="F", dtype=np.float64)
+    for i in range(size):
+        d[i] = i + 1
+    res = np.full([3], 42, order="F", dtype=np.float64)
+    sdfg(d=d, res=res)
+    assert res[0] == (1 + size) * size / 2
+    assert res[1] == (1 + size) * size / 2
+    assert res[2] == (2 + size - 1) * (size - 2)/ 2
+
+def test_fortran_frontend_sum2loop_1d_offset():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(2:6) :: d
+                    double precision, dimension(3) :: res
+                    CALL index_test_function(d,res)
+                    end
+
+                    SUBROUTINE index_test_function(d, res)
+                    double precision, dimension(2:6) :: d
+                    double precision, dimension(3) :: res
+
+                    res(1) = SUM(d)
+                    res(2) = SUM(d(:))
+                    res(3) = SUM(d(3:5))
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    size = 5
+    d = np.full([size], 0, order="F", dtype=np.float64)
+    for i in range(size):
+        d[i] = i + 1
+    res = np.full([3], 42, order="F", dtype=np.float64)
+    sdfg(d=d, res=res)
+    assert res[0] == (1 + size) * size / 2
+    assert res[1] == (1 + size) * size / 2
+    assert res[2] == (2 + size - 1) * (size - 2) / 2
+
+def test_fortran_frontend_arr2loop_2d():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(5,3) :: d
+                    double precision, dimension(4) :: res
+                    CALL index_test_function(d,res)
+                    end
+
+                    SUBROUTINE index_test_function(d, res)
+                    double precision, dimension(5,3) :: d
+                    double precision, dimension(4) :: res
+
+                    res(1) = SUM(d)
+                    res(2) = SUM(d(:,:))
+                    res(3) = SUM(d(2:4, 2))
+                    res(4) = SUM(d(2:4, 2:3))
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    sizes = [5, 3]
+    d = np.full(sizes, 42, order="F", dtype=np.float64)
+    cnt = 0
+    for i in range(sizes[0]):
+        for j in range(sizes[1]):
+            d[i, j] = cnt
+            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
+
+def test_fortran_frontend_arr2loop_2d_offset():
+    """
+    Tests that the generated array map correctly handles offsets.
+    """
+    test_string = """
+                    PROGRAM index_offset_test
+                    implicit none
+                    double precision, dimension(2:6,7:10) :: d
+                    double precision, dimension(3) :: res
+                    CALL index_test_function(d,res)
+                    end
+
+                    SUBROUTINE index_test_function(d, res)
+                    double precision, dimension(2:6,7:10) :: d
+                    double precision, dimension(3) :: res
+
+                    res(1) = SUM(d)
+                    res(2) = SUM(d(:,:))
+                    res(3) = SUM(d(3:5, 8:9))
+
+                    END SUBROUTINE index_test_function
+                    """
+
+    # Now test to verify it executes correctly with no offset normalization
+
+    sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True)
+    sdfg.simplify(verbose=True)
+    sdfg.compile()
+
+    sizes = [5, 4]
+    d = np.full(sizes, 42, order="F", dtype=np.float64)
+    cnt = 0
+    for i in range(sizes[0]):
+        for j in range(sizes[1]):
+            d[i, j] = cnt
+            cnt += 1
+    res = np.full([3], 42, order="F", dtype=np.float64)
+    sdfg(d=d, res=res)
+    assert res[0] == 190
+    assert res[1] == 190
+    assert res[2] == 57
+
+if __name__ == "__main__":
+
+    test_fortran_frontend_sum2loop_1d_without_offset()
+    test_fortran_frontend_sum2loop_1d_offset()
+    test_fortran_frontend_arr2loop_2d()
+    test_fortran_frontend_arr2loop_2d_offset()