diff --git a/model/common/src/icon4py/model/common/dace_opt.py b/model/common/src/icon4py/model/common/dace_opt.py new file mode 100644 index 0000000000..df7eb01104 --- /dev/null +++ b/model/common/src/icon4py/model/common/dace_opt.py @@ -0,0 +1,15 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import dace +from dace.transformation import dataflow as dace_dataflow + + +def graupel_run_top_level_dataflow_post(sdfg: dace.SDFG) -> None: + sdfg.apply_transformations_repeated(dace_dataflow.TaskletFusion, validate=False) + sdfg.save("graupel_run_top_level_dataflow_post.sdfg") diff --git a/model/common/src/icon4py/model/common/model_options.py b/model/common/src/icon4py/model/common/model_options.py index 95aebf8d8c..70a025c1f5 100644 --- a/model/common/src/icon4py/model/common/model_options.py +++ b/model/common/src/icon4py/model/common/model_options.py @@ -15,7 +15,7 @@ from gt4py.next import backend as gtx_backend from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations -from icon4py.model.common import model_backends +from icon4py.model.common import dace_opt, model_backends log = logging.getLogger(__name__) @@ -52,6 +52,9 @@ def get_dace_options( backend_descriptor["use_zero_origin"] = True optimization_args["gpu_memory_pool"] = False optimization_args["make_persistent"] = True + optimization_hooks[gtx_transformations.GT4PyAutoOptHook.TopLevelDataFlowPost] = ( + dace_opt.graupel_run_top_level_dataflow_post + ) if optimization_hooks: optimization_args["optimization_hooks"] = optimization_hooks if optimization_args: