diff --git a/dawn/examples/python/unstructured_stencil.py b/dawn/examples/python/unstructured_stencil.py index 64b8c24d7..efdd1e2c2 100644 --- a/dawn/examples/python/unstructured_stencil.py +++ b/dawn/examples/python/unstructured_stencil.py @@ -36,7 +36,6 @@ def main(args: argparse.Namespace): interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) - # create the out = in[i+1] statement body_ast = sir_utils.make_ast( [ sir_utils.make_assignment_stmt( diff --git a/dawn/test/integration-test/dawn4py-tests/CMakeLists.txt b/dawn/test/integration-test/dawn4py-tests/CMakeLists.txt index 446cf2a6f..701988951 100644 --- a/dawn/test/integration-test/dawn4py-tests/CMakeLists.txt +++ b/dawn/test/integration-test/dawn4py-tests/CMakeLists.txt @@ -40,6 +40,7 @@ endfunction() add_python_example(NAME copy_stencil VERIFY) add_python_example(NAME hori_diff_stencil VERIFY) add_python_example(NAME tridiagonal_solve_stencil VERIFY) +add_python_example(NAME ICON_laplacian_stencil VERIFY) add_python_example(NAME tridiagonal_solve_unstructured) add_python_example(NAME global_index_stencil) add_python_example(NAME unstructured_stencil) diff --git a/dawn/test/integration-test/dawn4py-tests/ICON_laplacian_stencil.py b/dawn/test/integration-test/dawn4py-tests/ICON_laplacian_stencil.py new file mode 100644 index 000000000..b7b77f1c8 --- /dev/null +++ b/dawn/test/integration-test/dawn4py-tests/ICON_laplacian_stencil.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python + +##===-----------------------------------------------------------------------------*- Python -*-===## +## _ +## | | +## __| | __ ___ ___ ___ +## / _` |/ _` \ \ /\ / / '_ | +## | (_| | (_| |\ V V /| | | | +## \__,_|\__,_| \_/\_/ |_| |_| - Compiler Toolchain +## +## +## This file is distributed under the MIT License (MIT). +## See LICENSE.txt for details. +## +##===------------------------------------------------------------------------------------------===## + +"""Generate input for the ICON Laplacian stencil test""" + +import os + +import dawn4py +from dawn4py.serialization import SIR +from dawn4py.serialization import utils as sir_utils +from google.protobuf.json_format import MessageToJson, Parse + + + +def main(): + stencil_name = "ICON_laplacian_stencil" + gen_outputfile = f"{stencil_name}.cpp" + sir_outputfile = f"{stencil_name}.sir" + + interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) + + body_ast = sir_utils.make_ast( + [ + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("rot_vec"), + sir_utils.make_reduction_over_neighbor_expr( + op="+", + init=sir_utils.make_literal_access_expr("0.0", SIR.BuiltinType.Double), + rhs=sir_utils.make_binary_operator( + sir_utils.make_field_access_expr("vec"), + "*", + sir_utils.make_field_access_expr("geofac_rot")), + lhs_location=SIR.LocationType.Value("Vertex"), + rhs_location=SIR.LocationType.Value("Edge"), + ), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("div_vec"), + sir_utils.make_reduction_over_neighbor_expr( + op="+", + init=sir_utils.make_literal_access_expr("0.0", SIR.BuiltinType.Double), + rhs=sir_utils.make_binary_operator( + sir_utils.make_field_access_expr("vec"), + "*", + sir_utils.make_field_access_expr("geofac_div")), + lhs_location=SIR.LocationType.Value("Cell"), + rhs_location=SIR.LocationType.Value("Edge"), + ), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("nabla2t1_vec"), + sir_utils.make_reduction_over_neighbor_expr( + op="+", + init=sir_utils.make_literal_access_expr("0.0", SIR.BuiltinType.Double), + rhs=sir_utils.make_field_access_expr("rot_vec"), + lhs_location=SIR.LocationType.Value("Edge"), + rhs_location=SIR.LocationType.Value("Vertex"), + weights=sir_utils.make_weights([-1.0, 1.0]) + ), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("nabla2t1_vec"), + sir_utils.make_binary_operator( + sir_utils.make_binary_operator( + sir_utils.make_field_access_expr("tangent_orientation"), + "*", + sir_utils.make_field_access_expr("nabla2t1_vec")), + "/", + sir_utils.make_field_access_expr("primal_edge_length")), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("nabla2t2_vec"), + sir_utils.make_reduction_over_neighbor_expr( + op="+", + init=sir_utils.make_literal_access_expr("0.0", SIR.BuiltinType.Double), + rhs=sir_utils.make_field_access_expr("div_vec"), + lhs_location=SIR.LocationType.Value("Edge"), + rhs_location=SIR.LocationType.Value("Cell"), + weights=sir_utils.make_weights([-1.0, 1.0]) + ), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("nabla2t2_vec"), + sir_utils.make_binary_operator( + sir_utils.make_field_access_expr("nabla2t2_vec"), + "/", + sir_utils.make_field_access_expr("dual_edge_length")), + "=", + ), + sir_utils.make_assignment_stmt( + sir_utils.make_field_access_expr("nabla2_vec"), + sir_utils.make_binary_operator( + sir_utils.make_field_access_expr("nabla2t2_vec"), + "-", + sir_utils.make_field_access_expr("nabla2t1_vec")), + "=", + ), + ] + ) + + vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( + body_ast, interval, SIR.VerticalRegion.Forward + ) + + sir = sir_utils.make_sir( + gen_outputfile, + SIR.GridType.Value("Unstructured"), + [ + sir_utils.make_stencil( + stencil_name, + sir_utils.make_ast([vertical_region_stmt]), + [ + sir_utils.make_field( + "vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "div_vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Cell")], 1 + ), + ), + sir_utils.make_field( + "rot_vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Vertex")], 1 + ), + ), + sir_utils.make_field( + "nabla2t1_vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "nabla2t2_vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "nabla2_vec", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "primal_edge_length", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "dual_edge_length", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "tangent_orientation", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "geofac_rot", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Vertex"), SIR.LocationType.Value("Edge")], 1 + ), + ), + sir_utils.make_field( + "geofac_div", + sir_utils.make_field_dimensions_unstructured( + [SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Edge")], 1 + ), + ), + ], + ), + ], + ) + + # write SIR to file (for debugging purposes) + f = open(sir_outputfile, "w") + f.write(MessageToJson(sir)) + f.close() + + # compile + code = dawn4py.compile(sir, backend="c++-naive-ico") + + # write to file + print(f"Writing generated code to '{gen_outputfile}'") + with open(gen_outputfile, "w") as f: + f.write(code) + + +if __name__ == "__main__": + main() diff --git a/dawn/test/integration-test/dawn4py-tests/data/ICON_laplacian_stencil_ref.cpp b/dawn/test/integration-test/dawn4py-tests/data/ICON_laplacian_stencil_ref.cpp new file mode 100644 index 000000000..b0cb1e198 --- /dev/null +++ b/dawn/test/integration-test/dawn4py-tests/data/ICON_laplacian_stencil_ref.cpp @@ -0,0 +1,106 @@ +//---- Preprocessor defines ---- +#define DAWN_GENERATED 1 +#undef DAWN_BACKEND_T +#define DAWN_BACKEND_T CXXNAIVEICO +#include + +//---- Globals ---- + +//---- Stencils ---- +namespace dawn_generated{ +namespace cxxnaiveico{ +template +class ICON_laplacian_stencil { +private: + + struct stencil_64 { + dawn::mesh_t const& m_mesh; + int m_k_size; + dawn::edge_field_t& m_vec; + dawn::cell_field_t& m_div_vec; + dawn::vertex_field_t& m_rot_vec; + dawn::edge_field_t& m_nabla2t1_vec; + dawn::edge_field_t& m_nabla2t2_vec; + dawn::edge_field_t& m_nabla2_vec; + dawn::edge_field_t& m_primal_edge_length; + dawn::edge_field_t& m_dual_edge_length; + dawn::edge_field_t& m_tangent_orientation; + dawn::sparse_vertex_field_t& m_geofac_rot; + dawn::sparse_cell_field_t& m_geofac_div; + public: + + stencil_64(dawn::mesh_t const &mesh, int k_size, dawn::edge_field_t&vec, dawn::cell_field_t&div_vec, dawn::vertex_field_t&rot_vec, dawn::edge_field_t&nabla2t1_vec, dawn::edge_field_t&nabla2t2_vec, dawn::edge_field_t&nabla2_vec, dawn::edge_field_t&primal_edge_length, dawn::edge_field_t&dual_edge_length, dawn::edge_field_t&tangent_orientation, dawn::sparse_vertex_field_t&geofac_rot, dawn::sparse_cell_field_t&geofac_div) : m_mesh(mesh), m_k_size(k_size), m_vec(vec), m_div_vec(div_vec), m_rot_vec(rot_vec), m_nabla2t1_vec(nabla2t1_vec), m_nabla2t2_vec(nabla2t2_vec), m_nabla2_vec(nabla2_vec), m_primal_edge_length(primal_edge_length), m_dual_edge_length(dual_edge_length), m_tangent_orientation(tangent_orientation), m_geofac_rot(geofac_rot), m_geofac_div(geofac_div){} + + ~stencil_64() { + } + + void sync_storages() { + } + static constexpr dawn::driver::unstructured_extent vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent div_vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent rot_vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent nabla2t1_vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent nabla2t2_vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent nabla2_vec_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent primal_edge_length_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent dual_edge_length_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent tangent_orientation_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent geofac_rot_extent = {false, 0,0}; + static constexpr dawn::driver::unstructured_extent geofac_div_extent = {false, 0,0}; + + void run() { + using dawn::deref; +{ + for(int k = 0+0; k <= ( m_k_size == 0 ? 0 : (m_k_size - 1)) + 0+0; ++k) { + for(auto const& loc : getCells(LibTag{}, m_mesh)) { +int m_sparse_dimension_idx = 0; +m_div_vec(deref(LibTag{}, loc),k+0) = reduceEdgeToCell(LibTag{}, m_mesh,loc, (::dawn::float_type) 0.0, [&](auto& lhs, auto const& red_loc) { lhs += (m_vec(deref(LibTag{}, red_loc),k+0) * m_geofac_div(deref(LibTag{}, loc),m_sparse_dimension_idx, k+0)); +m_sparse_dimension_idx++; +return lhs; +}); + } for(auto const& loc : getEdges(LibTag{}, m_mesh)) { +int m_sparse_dimension_idx = 0; +m_nabla2t2_vec(deref(LibTag{}, loc),k+0) = reduceCellToEdge(LibTag{}, m_mesh,loc, (::dawn::float_type) 0.0, [&](auto& lhs, auto const& red_loc, auto const& weight) { +lhs += weight * m_div_vec(deref(LibTag{}, red_loc),k+0); +m_sparse_dimension_idx++; +return lhs; +}, std::vector({-1.000000, 1.000000})); + } for(auto const& loc : getEdges(LibTag{}, m_mesh)) { +m_nabla2t2_vec(deref(LibTag{}, loc),k+0) = (m_nabla2t2_vec(deref(LibTag{}, loc),k+0) / m_dual_edge_length(deref(LibTag{}, loc),k+0)); + } for(auto const& loc : getVertices(LibTag{}, m_mesh)) { +int m_sparse_dimension_idx = 0; +m_rot_vec(deref(LibTag{}, loc),k+0) = reduceEdgeToVertex(LibTag{}, m_mesh,loc, (::dawn::float_type) 0.0, [&](auto& lhs, auto const& red_loc) { lhs += (m_vec(deref(LibTag{}, red_loc),k+0) * m_geofac_rot(deref(LibTag{}, loc),m_sparse_dimension_idx, k+0)); +m_sparse_dimension_idx++; +return lhs; +}); + } for(auto const& loc : getEdges(LibTag{}, m_mesh)) { +int m_sparse_dimension_idx = 0; +m_nabla2t1_vec(deref(LibTag{}, loc),k+0) = reduceVertexToEdge(LibTag{}, m_mesh,loc, (::dawn::float_type) 0.0, [&](auto& lhs, auto const& red_loc, auto const& weight) { +lhs += weight * m_rot_vec(deref(LibTag{}, red_loc),k+0); +m_sparse_dimension_idx++; +return lhs; +}, std::vector({-1.000000, 1.000000})); + } for(auto const& loc : getEdges(LibTag{}, m_mesh)) { +m_nabla2t1_vec(deref(LibTag{}, loc),k+0) = ((m_tangent_orientation(deref(LibTag{}, loc),k+0) * m_nabla2t1_vec(deref(LibTag{}, loc),k+0)) / m_primal_edge_length(deref(LibTag{}, loc),k+0)); + } for(auto const& loc : getEdges(LibTag{}, m_mesh)) { +m_nabla2_vec(deref(LibTag{}, loc),k+0) = (m_nabla2t2_vec(deref(LibTag{}, loc),k+0) - m_nabla2t1_vec(deref(LibTag{}, loc),k+0)); + } }} sync_storages(); + } + }; + static constexpr const char* s_name = "ICON_laplacian_stencil"; + stencil_64 m_stencil_64; +public: + + ICON_laplacian_stencil(const ICON_laplacian_stencil&) = delete; + + // Members + + ICON_laplacian_stencil(const dawn::mesh_t &mesh, int k_size, dawn::edge_field_t& vec, dawn::cell_field_t& div_vec, dawn::vertex_field_t& rot_vec, dawn::edge_field_t& nabla2t1_vec, dawn::edge_field_t& nabla2t2_vec, dawn::edge_field_t& nabla2_vec, dawn::edge_field_t& primal_edge_length, dawn::edge_field_t& dual_edge_length, dawn::edge_field_t& tangent_orientation, dawn::sparse_vertex_field_t& geofac_rot, dawn::sparse_cell_field_t& geofac_div) : m_stencil_64(mesh, k_size,vec,div_vec,rot_vec,nabla2t1_vec,nabla2t2_vec,nabla2_vec,primal_edge_length,dual_edge_length,tangent_orientation,geofac_rot,geofac_div){} + + void run() { + m_stencil_64.run(); +; + } +}; +} // namespace cxxnaiveico +} // namespace dawn_generated