diff --git a/dace/frontend/fortran/ast_transforms.py b/dace/frontend/fortran/ast_transforms.py index 57508d6d90..9c81a43c61 100644 --- a/dace/frontend/fortran/ast_transforms.py +++ b/dace/frontend/fortran/ast_transforms.py @@ -467,6 +467,10 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No variable = self.scope_vars.get_var(child.parent, var_name) offset = variable.offsets[idx] + # it can be a symbol - Name_Node - or a value + if not isinstance(offset, ast_internal_classes.Name_Node): + offset = ast_internal_classes.Int_Literal_Node(value=str(offset)) + newbody.append( ast_internal_classes.BinOp_Node( op="=", @@ -474,7 +478,7 @@ def visit_Execution_Part_Node(self, node: ast_internal_classes.Execution_Part_No rval=ast_internal_classes.BinOp_Node( op="-", lval=i, - rval=ast_internal_classes.Int_Literal_Node(value=str(offset)), + rval=offset, line_number=child.line_number), line_number=child.line_number)) else: @@ -752,7 +756,11 @@ def par_Decl_Range_Finder(node: ast_internal_classes.Array_Subscript_Node, lower_boundary = None if offsets[idx] != 1: - lower_boundary = ast_internal_classes.Int_Literal_Node(value=str(offsets[idx])) + # support symbols and integer literals + if isinstance(offsets[idx], ast_internal_classes.Name_Node): + lower_boundary = offsets[idx] + else: + lower_boundary = ast_internal_classes.Int_Literal_Node(value=str(offsets[idx])) else: lower_boundary = ast_internal_classes.Int_Literal_Node(value="1") @@ -765,10 +773,17 @@ def par_Decl_Range_Finder(node: ast_internal_classes.Array_Subscript_Node, But since the generated loop has `<=` condition, we need to subtract 1. """ if offsets[idx] != 1: + + # support symbols and integer literals + if isinstance(offsets[idx], ast_internal_classes.Name_Node): + offset = offsets[idx] + else: + offset = ast_internal_classes.Int_Literal_Node(value=str(offsets[idx])) + upper_boundary = ast_internal_classes.BinOp_Node( lval=upper_boundary, op="+", - rval=ast_internal_classes.Int_Literal_Node(value=str(offsets[idx])) + rval=offset ) upper_boundary = ast_internal_classes.BinOp_Node( lval=upper_boundary, diff --git a/tests/fortran/offset_normalizer_test.py b/tests/fortran/offset_normalizer_test.py index b4138c1cac..9e942a97a6 100644 --- a/tests/fortran/offset_normalizer_test.py +++ b/tests/fortran/offset_normalizer_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import ast_internal_classes, ast_transforms, fortran_parser def test_fortran_frontend_offset_normalizer_1d(): """ @@ -48,6 +48,60 @@ def test_fortran_frontend_offset_normalizer_1d(): for i in range(0,5): assert a[i] == (50+i)* 2 +def test_fortran_frontend_offset_normalizer_1d_symbol(): + """ + Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct. + """ + test_string = """ + PROGRAM index_offset_test + implicit none + integer :: arrsize + integer :: arrsize2 + double precision :: d(arrsize:arrsize2) + CALL index_test_function(d, arrsize, arrsize2) + end + + SUBROUTINE index_test_function(d, arrsize, arrsize2) + integer :: arrsize + integer :: arrsize2 + double precision :: d(arrsize:arrsize2) + + do i=arrsize,arrsize2 + d(i) = i * 2.0 + end do + + END SUBROUTINE index_test_function + """ + + # Test to verify that offset is normalized correctly + ast, own_ast = fortran_parser.create_ast_from_string(test_string, "index_offset_test", True, True) + + for subroutine in ast.subroutine_definitions: + + loop = subroutine.execution_part.execution[1] + idx_assignment = loop.body.execution[1] + assert isinstance(idx_assignment.rval.rval, ast_internal_classes.Name_Node) + assert idx_assignment.rval.rval.name == "arrsize" + + # Now test to verify it executes correctly + + sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True) + sdfg.simplify(verbose=True) + sdfg.compile() + + from dace.symbolic import evaluate + arrsize=50 + arrsize2=54 + assert len(sdfg.data('d').shape) == 1 + assert evaluate(sdfg.data('d').shape[0], {'arrsize': arrsize, 'arrsize2': arrsize2}) == 5 + + arrsize=50 + arrsize2=54 + a = np.full([arrsize2-arrsize+1], 42, order="F", dtype=np.float64) + sdfg(d=a, arrsize=arrsize, arrsize2=arrsize2) + for i in range(0, arrsize2 - arrsize + 1): + assert a[i] == (50+i)* 2 + def test_fortran_frontend_offset_normalizer_2d(): """ Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct. @@ -103,6 +157,78 @@ def test_fortran_frontend_offset_normalizer_2d(): for j in range(0,3): assert a[i, j] == (50+i) * 2 + 3 * (7 + j) +def test_fortran_frontend_offset_normalizer_2d_symbol(): + """ + Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct. + """ + test_string = """ + PROGRAM index_offset_test + implicit none + integer :: arrsize + integer :: arrsize2 + integer :: arrsize3 + integer :: arrsize4 + double precision, dimension(arrsize:arrsize2,arrsize3:arrsize4) :: d + CALL index_test_function(d, arrsize, arrsize2, arrsize3, arrsize4) + end + + SUBROUTINE index_test_function(d, arrsize, arrsize2, arrsize3, arrsize4) + integer :: arrsize + integer :: arrsize2 + integer :: arrsize3 + integer :: arrsize4 + double precision, dimension(arrsize:arrsize2,arrsize3:arrsize4) :: d + + do i=arrsize, arrsize2 + do j=arrsize3, arrsize4 + d(i, j) = i * 2.0 + 3 * j + end do + end do + + END SUBROUTINE index_test_function + """ + + # Test to verify that offset is normalized correctly + ast, own_ast = fortran_parser.create_ast_from_string(test_string, "index_offset_test", True, True) + + for subroutine in ast.subroutine_definitions: + + loop = subroutine.execution_part.execution[1] + nested_loop = loop.body.execution[1] + + idx = nested_loop.body.execution[1] + assert idx.lval.name == 'tmp_index_0' + assert isinstance(idx.rval.rval, ast_internal_classes.Name_Node) + assert idx.rval.rval.name == "arrsize" + + idx2 = nested_loop.body.execution[3] + assert idx2.lval.name == 'tmp_index_1' + assert isinstance(idx2.rval.rval, ast_internal_classes.Name_Node) + assert idx2.rval.rval.name == "arrsize3" + + # Now test to verify it executes correctly + + sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True) + sdfg.simplify(verbose=True) + sdfg.compile() + + from dace.symbolic import evaluate + values = { + 'arrsize': 50, + 'arrsize2': 54, + 'arrsize3': 7, + 'arrsize4': 9 + } + assert len(sdfg.data('d').shape) == 2 + assert evaluate(sdfg.data('d').shape[0], values) == 5 + assert evaluate(sdfg.data('d').shape[1], values) == 3 + + a = np.full([5,3], 42, order="F", dtype=np.float64) + sdfg(d=a, **values) + for i in range(0,5): + for j in range(0,3): + assert a[i, j] == (50+i) * 2 + 3 * (7 + j) + def test_fortran_frontend_offset_normalizer_2d_arr2loop(): """ Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct. @@ -157,8 +283,81 @@ def test_fortran_frontend_offset_normalizer_2d_arr2loop(): for j in range(0,3): assert a[i, j] == (50 + i) * 2 +def test_fortran_frontend_offset_normalizer_2d_arr2loop_symbol(): + """ + Tests that the Fortran frontend can parse array accesses and that the accessed indices are correct. + """ + test_string = """ + PROGRAM index_offset_test + implicit none + integer :: arrsize + integer :: arrsize2 + integer :: arrsize3 + integer :: arrsize4 + double precision, dimension(arrsize:arrsize2,arrsize3:arrsize4) :: d + CALL index_test_function(d, arrsize, arrsize2, arrsize3, arrsize4) + end + + SUBROUTINE index_test_function(d, arrsize, arrsize2, arrsize3, arrsize4) + integer :: arrsize + integer :: arrsize2 + integer :: arrsize3 + integer :: arrsize4 + double precision, dimension(arrsize:arrsize2,arrsize3:arrsize4) :: d + + do i=arrsize,arrsize2 + d(i, :) = i * 2.0 + end do + + END SUBROUTINE index_test_function + """ + + # Test to verify that offset is normalized correctly + ast, own_ast = fortran_parser.create_ast_from_string(test_string, "index_offset_test", True, True) + + for subroutine in ast.subroutine_definitions: + + loop = subroutine.execution_part.execution[1] + nested_loop = loop.body.execution[1] + + idx = nested_loop.body.execution[1] + assert idx.lval.name == 'tmp_index_0' + assert isinstance(idx.rval.rval, ast_internal_classes.Name_Node) + assert idx.rval.rval.name == "arrsize" + + idx2 = nested_loop.body.execution[3] + assert idx2.lval.name == 'tmp_index_1' + assert isinstance(idx2.rval.rval, ast_internal_classes.Name_Node) + assert idx2.rval.rval.name == "arrsize3" + + # Now test to verify it executes correctly with no normalization + + sdfg = fortran_parser.create_sdfg_from_string(test_string, "index_offset_test", True) + sdfg.simplify(verbose=True) + sdfg.compile() + + from dace.symbolic import evaluate + values = { + 'arrsize': 50, + 'arrsize2': 54, + 'arrsize3': 7, + 'arrsize4': 9 + } + assert len(sdfg.data('d').shape) == 2 + assert evaluate(sdfg.data('d').shape[0], values) == 5 + assert evaluate(sdfg.data('d').shape[1], values) == 3 + + a = np.full([5,3], 42, order="F", dtype=np.float64) + sdfg(d=a, **values) + for i in range(0,5): + for j in range(0,3): + assert a[i, j] == (50 + i) * 2 + if __name__ == "__main__": test_fortran_frontend_offset_normalizer_1d() test_fortran_frontend_offset_normalizer_2d() test_fortran_frontend_offset_normalizer_2d_arr2loop() + test_fortran_frontend_offset_normalizer_1d_symbol() + test_fortran_frontend_offset_normalizer_2d_symbol() + test_fortran_frontend_offset_normalizer_2d_arr2loop_symbol()