def main(args: argparse.Namespace): interval = serial_utils.make_interval( AST.Interval.Start, AST.Interval.End, 0, 0) # create the out = in[i+1] statement body_ast = serial_utils.make_ast( [ serial_utils.make_assignment_stmt( serial_utils.make_unstructured_field_access_expr("out"), serial_utils.make_reduction_over_neighbor_expr( "+", serial_utils.make_unstructured_field_access_expr( "in", horizontal_offset=serial_utils.make_unstructured_offset(False)), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Float), chain=[AST.LocationType.Value( "Cell"), AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell")], ), "=", ) ] ) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward ) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "in", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1 ), ), serial_utils.make_field( "out", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1 ), ), ], ), ], ) # print the SIR f = open("unstructured_stencil.sir", "w") f.write(MessageToJson(sir)) f.close() # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def main(args: argparse.Namespace): interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) # create the stencil body AST body_ast = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("lap"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("-4.0", SIR.BuiltinType.Float), "*", sir_utils.make_field_access_expr("in"), ), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("coeff"), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [-1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "in", [0, 1, 0]), "+", sir_utils.make_field_access_expr( "in", [0, -1, 0]), ), ), ), ), ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("out"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("-4.0", SIR.BuiltinType.Float), "*", sir_utils.make_field_access_expr("lap"), ), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("coeff"), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("lap", [1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "lap", [-1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "lap", [0, 1, 0]), "+", sir_utils.make_field_access_expr( "lap", [0, -1, 0]), ), ), ), ), ), "=", ), ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) sir = sir_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Cartesian"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "in", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "out", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "coeff", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "lap", sir_utils.make_field_dimensions_cartesian(), is_temporary=True), ], ) ], ) # print the SIR if args.verbose: sir_utils.pprint(sir) # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CUDA) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def main(args: argparse.Namespace): interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) # create the out = reduce(sparse_CE * in) statement body_ast = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("out"), sir_utils.make_reduction_over_neighbor_expr( "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("sparse_CE"), "*", sir_utils.make_field_access_expr("in")), sir_utils.make_literal_access_expr("1.0", SIR.BuiltinType.Float), lhs_location=SIR.LocationType.Value('Cell'), rhs_location=SIR.LocationType.Value('Edge')), "=", ) ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) sir = sir_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Unstructured"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "in", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value('Edge')], 1)), sir_utils.make_field( "sparse_CE", sir_utils.make_field_dimensions_unstructured([ SIR.LocationType.Value('Cell'), SIR.LocationType.Value('Edge') ], 1)), sir_utils.make_field( "out", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value('Cell')], 1)) ], ), ], ) # print the SIR if args.verbose: sir_utils.pprint(sir) # compile code = dawn4py.compile(sir, backend="c++-naive-ico") # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
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", [True, 0]), "*", sir_utils.make_field_access_expr("geofac_rot"), ), chain=[SIR.LocationType.Value("Vertex"), 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", [True, 0]), "*", sir_utils.make_field_access_expr("geofac_div"), ), chain=[SIR.LocationType.Value("Cell"), 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", [True, 0]), chain=[SIR.LocationType.Value( "Edge"), SIR.LocationType.Value("Vertex")], weights=[sir_utils.make_literal_access_expr( "-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "1.0", SIR.BuiltinType.Double)] ), "=", ), 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", [True, 0]), chain=[SIR.LocationType.Value( "Edge"), SIR.LocationType.Value("Cell")], weights=[sir_utils.make_literal_access_expr( "-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "1.0", SIR.BuiltinType.Double)] ), "=", ), 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(sir_utils.to_json(sir)) f.close() # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{gen_outputfile}'") with open(gen_outputfile, "w") as f: f.write(code)
def main(): stencil_name = "general_weights" 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([ # compute nabla2 using the diamond reduction sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("nabla2"), 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("vn_vert"), chain=[ SIR.LocationType.Value("Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Vertex") ], weights=[ sir_utils.make_field_access_expr("inv_primal_edge_length"), sir_utils.make_field_access_expr("inv_primal_edge_length"), sir_utils.make_field_access_expr("inv_primal_edge_length"), sir_utils.make_field_access_expr("inv_primal_edge_length") ]), "=", ), ]) 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( "inv_primal_edge_length", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1), ), sir_utils.make_field( "vn_vert", sir_utils.make_field_dimensions_unstructured([ SIR.LocationType.Value("Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Vertex") ], 1), ), sir_utils.make_field( "nabla2", sir_utils.make_field_dimensions_unstructured( [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=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{gen_outputfile}'") with open(gen_outputfile, "w") as f: f.write(code)
def main(args: argparse.Namespace): interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) # create the laplace statement body_ast = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("out", [0, 0, 0]), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [0, 0, 0]), "*", sir_utils.make_literal_access_expr( "-4.0", sir_utils.BuiltinType.Float), ), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [-1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "in", [0, 1, 0]), "+", sir_utils.make_field_access_expr( "in", [0, -1, 0]), ), ), ), ), "/", sir_utils.make_binary_operator( sir_utils.make_var_access_expr("dx", is_external=True), "*", sir_utils.make_var_access_expr("dx", is_external=True), ), ), "=", ), ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) stencils_globals = sir_utils.GlobalVariableMap() stencils_globals.map["dx"].double_value = 0.0 sir = sir_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Cartesian"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "out", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "in", sir_utils.make_field_dimensions_cartesian()), ], ) ], global_variables=stencils_globals, ) # print the SIR if args.verbose: sir_utils.pprint(sir) # serialize the SIR to file sir_file = open("./laplacian_stencil_from_python.sir", "wb") sir_file.write(sir_utils.to_json(sir)) sir_file.close() # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaive) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def function_call(): outputfile = "../input/test_set_stage_location_type_function_call.sir" interval = serial_utils.make_interval( SIR.Interval.Start, SIR.Interval.End, 0, 0) fun_ast = serial_utils.make_ast( [ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_literal_access_expr( value="2.0", type=serial_utils.BuiltinType.Float), "=", ), ] ) arg_field = serial_utils.make_field( "out", serial_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Cell")], 1 ) ) fun = serial_utils.make_stencil_function( name='f', asts=[fun_ast], intervals=[interval], arguments=[serial_utils.make_stencil_function_arg(arg_field)]) body_ast = serial_utils.make_ast( [ serial_utils.make_expr_stmt(expr=serial_utils.make_stencil_fun_call_expr( callee="f", arguments=[serial_utils.make_field_access_expr("out_cell")])), ] ) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward ) sir = serial_utils.make_sir( outputfile, SIR.GridType.Value("Unstructured"), [ serial_utils.make_stencil( "generated", serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "out_cell", serial_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Cell")], 1 ), ), ], ), ], functions=[fun] ) f = open(outputfile, "w") f.write(MessageToJson(sir)) f.close()
def main(args: argparse.Namespace): interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) line_1 = serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("a"), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("b"), "/", serial_utils.make_field_access_expr("c"), ), "+", serial_utils.make_literal_access_expr("5", AST.BuiltinType.Float)), "=") line_2 = serial_utils.make_block_stmt( serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("a"), serial_utils.make_field_access_expr("b"), "=")) line_3 = serial_utils.make_block_stmt( serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("c"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("a"), "+", serial_utils.make_literal_access_expr("1", AST.BuiltinType.Float)), "=")) body_ast = serial_utils.make_ast([ line_1, serial_utils.make_if_stmt( serial_utils.make_expr_stmt( serial_utils.make_field_access_expr("d")), line_2, serial_utils.make_block_stmt( serial_utils.make_if_stmt( serial_utils.make_expr_stmt( serial_utils.make_field_access_expr("e")), line_3))) ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "a", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "b", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "c", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "d", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "e", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), ], ), ], ) # print the SIR if args.verbose: print(MessageToJson(sir)) # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def sparse_temporary(): outputfile = "DontDemoteSparse" interval = serial_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_loop_stmt( serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("sparseF"), serial_utils.make_literal_access_expr("1.", SIR.BuiltinType.Double), "="), [ SIR.LocationType.Value("Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Edge") ]), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("outF"), serial_utils.make_reduction_over_neighbor_expr( "+", serial_utils.make_binary_operator( serial_utils.make_unstructured_field_access_expr( "inF", horizontal_offset=serial_utils. make_unstructured_offset(False)), "*", serial_utils.make_unstructured_field_access_expr( "sparseF", horizontal_offset=serial_utils. make_unstructured_offset(True))), serial_utils.make_literal_access_expr("0.", SIR.BuiltinType.Double), [ SIR.LocationType.Value("Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Edge") ]), "="), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) sir = serial_utils.make_sir( outputfile, SIR.GridType.Value("Unstructured"), [ serial_utils.make_stencil( "generated", serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "inF", serial_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "outF", serial_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "sparseF", serial_utils.make_field_dimensions_unstructured([ SIR.LocationType.Value("Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Edge") ], 1), is_temporary=True), ], ), ], ) sim = dawn4py.lower_and_optimize(sir, groups=[]) with open(outputfile, mode="w") as f: f.write(MessageToJson(sim["generated"])) os.rename(outputfile, "../input/" + outputfile + ".iir")
def copy_vars(): outputfile = "../input/test_set_stage_location_type_copy_vars.sir" interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) body_ast = sir_utils.make_ast([ sir_utils.make_var_decl_stmt( sir_utils.make_type(sir_utils.BuiltinType.Float), "out_var_cell"), sir_utils.make_var_decl_stmt( sir_utils.make_type(sir_utils.BuiltinType.Float), "out_var_edge"), sir_utils.make_var_decl_stmt( sir_utils.make_type(sir_utils.BuiltinType.Float), "out_var_vertex"), sir_utils.make_assignment_stmt( sir_utils.make_var_access_expr("out_var_cell"), sir_utils.make_field_access_expr("in_cell"), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_var_access_expr("out_var_edge"), sir_utils.make_field_access_expr("in_edge"), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_var_access_expr("out_var_vertex"), sir_utils.make_field_access_expr("in_vertex"), "=", ) ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) sir = sir_utils.make_sir( outputfile, SIR.GridType.Value("Unstructured"), [ sir_utils.make_stencil( "generated", sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "in_cell", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Cell")], 1), ), sir_utils.make_field( "in_edge", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1), ), sir_utils.make_field( "in_vertex", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Vertex")], 1), ), ], ), ], ) f = open(outputfile, "w") f.write(MessageToJson(sir)) f.close()
def main(args: argparse.Namespace): interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_loop_stmt(serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("geofac_grg"), serial_utils.make_literal_access_expr("2.", AST.BuiltinType.Double)), [ AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], include_center=True), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("p_grad"), serial_utils.make_reduction_over_neighbor_expr( "+", serial_utils.make_binary_operator( serial_utils.make_unstructured_field_access_expr( "geofac_grg"), "*", serial_utils.make_unstructured_field_access_expr( "p_ccpr", horizontal_offset=serial_utils. make_unstructured_offset(True))), init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), chain=[ AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], include_center=True, ), "=", ) ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "p_grad", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "p_ccpr", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "geofac_grg", serial_utils.make_field_dimensions_unstructured( [ AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], 1, include_center=True), ), ], ), ], ) # print the SIR # if args.verbose: f = open(SIR_OUTPUT_FILE, "w") f.write(MessageToJson(sir)) f.close() # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CUDAIco) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def two_copies_mixed(): outputfile = "StageMergerTestTwoCopiesMixed" 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("out_cell"), sir_utils.make_field_access_expr("in_cell"), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("out_edge"), sir_utils.make_field_access_expr("in_edge"), "=", ), ] ) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward ) sir = sir_utils.make_sir( outputfile, SIR.GridType.Value("Unstructured"), [ sir_utils.make_stencil( "generated", sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "in_cell", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Cell")], 1 ), ), sir_utils.make_field( "out_cell", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Cell")], 1 ), ), sir_utils.make_field( "in_edge", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "out_edge", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), ], ), ], ) sim = dawn4py.lower_and_optimize(sir) with open(outputfile, mode="w") as f: f.write(sim["generated"]) os.rename(outputfile + ".0.iir", "../input/" + outputfile + ".iir")
def main(args: argparse.Namespace): interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("full"), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("horizontal"), "+", serial_utils.make_field_access_expr("vertical"))), "="), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Float), rhs=serial_utils.make_field_access_expr( "horizontal_sparse", [True, 0]), chain=[ AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], ), "=", ) ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "out", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "full", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "horizontal", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 0), ), serial_utils.make_field( "horizontal_sparse", serial_utils.make_field_dimensions_unstructured([ AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], 0), ), serial_utils.make_vertical_field("vertical"), ], ), ], ) # print the SIR if args.verbose: print(MessageToJson(sir)) with open("out.json", "w+") as f: f.write(serial_utils.to_json(sir)) # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CUDAIco) # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def main(args: argparse.Namespace): stencil_name = "ICON_laplacian_diamond_stencil" gen_outputfile = f"{stencil_name}.cpp" sir_outputfile = f"{stencil_name}.sir" interval = serial_utils.make_interval( AST.Interval.Start, AST.Interval.End, 0, 0) body_ast = serial_utils.make_ast( [ # fill sparse dimension vn vert using the loop concept serial_utils.make_loop_stmt( [serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("vn_vert"), serial_utils.make_binary_operator( serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "u_vert", [True, 0]), "*", serial_utils.make_field_access_expr("primal_normal_x", [True, 0])), "+", serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "v_vert", [True, 0]), "*", serial_utils.make_field_access_expr("primal_normal_y", [True, 0])), ), "=")], [AST.LocationType.Value( "Edge"), AST.LocationType.Value("Cell"), AST.LocationType.Value("Vertex")] ), # dvt_tang for smagorinsky serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("dvt_tang"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_binary_operator( serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "u_vert", [True, 0]), "*", serial_utils.make_field_access_expr("dual_normal_x", [True, 0])), "+", serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "v_vert", [True, 0]), "*", serial_utils.make_field_access_expr("dual_normal_y", [True, 0])), ), chain=[AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], weights=[serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double)] ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("dvt_tang"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("dvt_tang"), "*", serial_utils.make_field_access_expr("tangent_orientation")), "="), # dvt_norm for smagorinsky serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("dvt_norm"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_binary_operator( serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "u_vert", [True, 0]), "*", serial_utils.make_field_access_expr("dual_normal_x", [True, 0])), "+", serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "v_vert", [True, 0]), "*", serial_utils.make_field_access_expr("dual_normal_y", [True, 0])), ), chain=[AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], weights=[serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Double)] ), "=", ), # compute smagorinsky serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("kh_smag_1"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_field_access_expr("vn_vert"), chain=[AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], weights=[serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double)] ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("kh_smag_1"), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("kh_smag_1"), "*", serial_utils.make_field_access_expr("tangent_orientation")), "*", serial_utils.make_field_access_expr("inv_primal_edge_length")), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("dvt_norm"), "*", serial_utils.make_field_access_expr("inv_vert_vert_length"))), "="), serial_utils.make_assignment_stmt(serial_utils.make_field_access_expr("kh_smag_1"), serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "kh_smag_1"), "*", serial_utils.make_field_access_expr("kh_smag_1"))), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("kh_smag_2"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_field_access_expr("vn_vert"), chain=[AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], weights=[serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( " 1.0", AST.BuiltinType.Double)] ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("kh_smag_2"), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("kh_smag_2"), "*", serial_utils.make_field_access_expr("inv_vert_vert_length")), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("dvt_tang"), "*", serial_utils.make_field_access_expr("inv_primal_edge_length"))), "="), serial_utils.make_assignment_stmt(serial_utils.make_field_access_expr("kh_smag_2"), serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "kh_smag_2"), "*", serial_utils.make_field_access_expr("kh_smag_2"))), # currently not able to forward a sqrt, so this is technically kh_smag**2 serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("kh_smag"), serial_utils.make_binary_operator(serial_utils.make_field_access_expr("diff_multfac_smag"), "*", serial_utils.make_fun_call_expr("math::sqrt", [serial_utils.make_binary_operator(serial_utils.make_field_access_expr( "kh_smag_1"), "+", serial_utils.make_field_access_expr("kh_smag_2"))])), "="), # compute nabla2 using the diamond reduction serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_binary_operator(serial_utils.make_literal_access_expr( "4.0", AST.BuiltinType.Double), "*", serial_utils.make_field_access_expr("vn_vert")), chain=[AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], weights=[ serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_primal_edge_length"), '*', serial_utils.make_field_access_expr( "inv_primal_edge_length")), serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_primal_edge_length"), '*', serial_utils.make_field_access_expr( "inv_primal_edge_length")), serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_vert_vert_length"), '*', serial_utils.make_field_access_expr( "inv_vert_vert_length")), serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_vert_vert_length"), '*', serial_utils.make_field_access_expr( "inv_vert_vert_length")), ] ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("nabla2"), "-", serial_utils.make_binary_operator( serial_utils.make_binary_operator(serial_utils.make_binary_operator(serial_utils.make_literal_access_expr( "8.0", AST.BuiltinType.Double), "*", serial_utils.make_field_access_expr("vn")), "*", serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_primal_edge_length"), "*", serial_utils.make_field_access_expr( "inv_primal_edge_length"))), "+", serial_utils.make_binary_operator(serial_utils.make_binary_operator(serial_utils.make_literal_access_expr( "8.0", AST.BuiltinType.Double), "*", serial_utils.make_field_access_expr("vn")), "*", serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "inv_vert_vert_length"), "*", serial_utils.make_field_access_expr( "inv_vert_vert_length"))))), "=") ] ) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward ) sir = serial_utils.make_sir( gen_outputfile, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( stencil_name, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "diff_multfac_smag", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value( "Edge")], 1 ), ), serial_utils.make_field( "tangent_orientation", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "inv_primal_edge_length", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "inv_vert_vert_length", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "u_vert", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "v_vert", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "primal_normal_x", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "primal_normal_y", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "dual_normal_x", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "dual_normal_y", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "vn_vert", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge"), AST.LocationType.Value( "Cell"), AST.LocationType.Value("Vertex")], 1 ), ), serial_utils.make_field( "vn", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "dvt_tang", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "dvt_norm", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "kh_smag_1", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "kh_smag_2", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "kh_smag", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "nabla2", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1 ), ), ], ), ], ) # print the SIR if args.verbose: print(MessageToJson(sir)) # compile code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{gen_outputfile}'") with open(gen_outputfile, "w") as f: f.write(code)
def make_tridiagonal_solve_stencil_sir(name=None): OUTPUT_NAME = name if name is not None else "tridiagonal_solve_stencil" OUTPUT_FILE = f"{OUTPUT_NAME}.cpp" # ---- First vertical region statement ---- interval_1 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) body_ast_1 = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("c"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "/", sir_utils.make_field_access_expr("b"), ), "=", ) ]) vertical_region_stmt_1 = sir_utils.make_vertical_region_decl_stmt( body_ast_1, interval_1, SIR.VerticalRegion.Forward) # ---- Second vertical region statement ---- interval_2 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 1, 0) body_ast_2 = sir_utils.make_ast([ sir_utils.make_var_decl_stmt( sir_utils.make_type(SIR.BuiltinType.Integer), "m", 0, "=", sir_utils.make_expr( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("1.0", SIR.BuiltinType.Float), "/", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("b"), "-", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("a"), "*", sir_utils.make_field_access_expr("c", [0, 0, -1]), ), ), )), ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("c"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "*", sir_utils.make_var_access_expr("m")), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_field_access_expr("d"), "-", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("a"), "*", sir_utils.make_field_access_expr("d", [0, 0, -1]), ), ), "*", sir_utils.make_var_access_expr("m"), ), "=", ), ]) vertical_region_stmt_2 = sir_utils.make_vertical_region_decl_stmt( body_ast_2, interval_2, SIR.VerticalRegion.Forward) # ---- Third vertical region statement ---- interval_3 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, -1) body_ast_3 = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "*", sir_utils.make_field_access_expr("d", [0, 0, 1]), ), "-=", ) ]) vertical_region_stmt_3 = sir_utils.make_vertical_region_decl_stmt( body_ast_3, interval_3, SIR.VerticalRegion.Backward) sir = sir_utils.make_sir( OUTPUT_FILE, sir_utils.GridType.Value("Cartesian"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([ vertical_region_stmt_1, vertical_region_stmt_2, vertical_region_stmt_3 ]), [ sir_utils.make_field( "a", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "b", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "c", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "d", sir_utils.make_field_dimensions_cartesian()), ], ) ], ) return sir
def main(args: argparse.Namespace): sir = sir_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Cartesian"), [ sir_utils.make_stencil( "global_indexing", sir_utils.make_ast([ create_vertical_region_stmt(), create_boundary_correction_region( value="4", i_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), ), create_boundary_correction_region( value="8", i_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), ), create_boundary_correction_region( value="6", j_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), ), create_boundary_correction_region( value="2", j_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), ), create_boundary_correction_region( value="1", j_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), i_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), ), create_boundary_correction_region( value="3", j_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), i_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), ), create_boundary_correction_region( value="7", j_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), i_interval=sir_utils.make_interval( SIR.Interval.Start, SIR.Interval.Start, 0, 1), ), create_boundary_correction_region( value="5", j_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), i_interval=sir_utils.make_interval( SIR.Interval.End, SIR.Interval.End, -1, 0), ), ]), [ sir_utils.make_field( "in", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "out", sir_utils.make_field_dimensions_cartesian()) ], ) ], ) # print the SIR if args.verbose: sir_utils.pprint(sir) # compile code = dawn4py.compile(sir, backend="c++-naive") # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def make_hori_diff_stencil_sir(name=None): OUTPUT_NAME = name if name is not None else "hori_diff_stencil" OUTPUT_FILE = f"{OUTPUT_NAME}.cpp" interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) # create the stencil body AST body_ast = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("lap"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("-4.0", SIR.BuiltinType.Float), "*", sir_utils.make_field_access_expr("in"), ), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("coeff"), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("in", [-1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "in", [0, 1, 0]), "+", sir_utils.make_field_access_expr( "in", [0, -1, 0]), ), ), ), ), ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("out"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("-4.0", SIR.BuiltinType.Float), "*", sir_utils.make_field_access_expr("lap"), ), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("coeff"), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("lap", [1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "lap", [-1, 0, 0]), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "lap", [0, 1, 0]), "+", sir_utils.make_field_access_expr( "lap", [0, -1, 0]), ), ), ), ), ), "=", ), ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) sir = sir_utils.make_sir( OUTPUT_FILE, sir_utils.GridType.Value("Cartesian"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([vertical_region_stmt]), [ sir_utils.make_field( "in", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "out", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "coeff", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field( "lap", sir_utils.make_field_dimensions_cartesian(), is_temporary=True), ], ) ], ) return sir
def sparse_temporary(): outputfile = "AlsoDemoteWeight" interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_unstructured_field_access_expr("tempF"), serial_utils.make_unstructured_field_access_expr("test"), ), serial_utils.make_assignment_stmt( serial_utils.make_unstructured_field_access_expr("outF"), serial_utils.make_reduction_over_neighbor_expr( "+", serial_utils.make_unstructured_field_access_expr("inF"), serial_utils.make_literal_access_expr("0.", AST.BuiltinType.Double), [ AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], weights=[ serial_utils.make_unstructured_field_access_expr("tempF"), serial_utils.make_unstructured_field_access_expr("tempF"), serial_utils.make_unstructured_field_access_expr("tempF"), serial_utils.make_unstructured_field_access_expr("tempF") ]), "="), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) sir = serial_utils.make_sir( outputfile, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( "generated", serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "inF", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "outF", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "test", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "tempF", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), is_temporary=True), ], ), ], ) sim = dawn4py.lower_and_optimize(sir, groups=[]) with open(outputfile, mode="w") as f: f.write(MessageToJson(sim["generated"])) os.rename(outputfile, "../input/" + outputfile + ".iir")
def main(args: argparse.Namespace): # ---- First vertical region statement ---- interval_1 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) body_ast_1 = sir_utils.make_ast( [ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("c"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "/", sir_utils.make_field_access_expr("b"), ), "=", ) ] ) vertical_region_stmt_1 = sir_utils.make_vertical_region_decl_stmt( body_ast_1, interval_1, SIR.VerticalRegion.Forward ) # ---- Second vertical region statement ---- interval_2 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 1, 0) body_ast_2 = sir_utils.make_ast( [ sir_utils.make_var_decl_stmt( sir_utils.make_type(SIR.BuiltinType.Integer), "m", 0, "=", sir_utils.make_expr( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr("1.0", SIR.BuiltinType.Float), "/", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("b"), "-", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("a"), "*", sir_utils.make_field_access_expr("c", [0, 0, -1]), ), ), ) ), ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("c"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "*", sir_utils.make_var_access_expr("m") ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_field_access_expr("d"), "-", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("a"), "*", sir_utils.make_field_access_expr("d", [0, 0, -1]), ), ), "*", sir_utils.make_var_access_expr("m"), ), "=", ), ] ) vertical_region_stmt_2 = sir_utils.make_vertical_region_decl_stmt( body_ast_2, interval_2, SIR.VerticalRegion.Forward ) # ---- Third vertical region statement ---- interval_3 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, -1) body_ast_3 = sir_utils.make_ast( [ sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("c"), "*", sir_utils.make_field_access_expr("d", [0, 0, 1]), ), "-=", ) ] ) vertical_region_stmt_3 = sir_utils.make_vertical_region_decl_stmt( body_ast_3, interval_3, SIR.VerticalRegion.Backward ) sir = sir_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Cartesian"), [ sir_utils.make_stencil( OUTPUT_NAME, sir_utils.make_ast([vertical_region_stmt_1, vertical_region_stmt_2, vertical_region_stmt_3]), [ sir_utils.make_field("a", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field("b", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field("c", sir_utils.make_field_dimensions_cartesian()), sir_utils.make_field("d", sir_utils.make_field_dimensions_cartesian()), ], ) ], ) # print the SIR if args.verbose: sir_utils.pprint(sir) # compile code = dawn4py.compile(sir, backend="cuda") # write to file print(f"Writing generated code to '{OUTPUT_PATH}'") with open(OUTPUT_PATH, "w") as f: f.write(code)
def main(args: argparse.Namespace): stencil_name = "ICON_laplacian_stencil" gen_outputfile = f"{stencil_name}.cpp" sir_outputfile = f"{stencil_name}.sir" interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("rot_vec"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_binary_operator( serial_utils.make_field_access_expr("vec", [True, 0]), "*", serial_utils.make_field_access_expr("geofac_rot"), ), chain=[ AST.LocationType.Value("Vertex"), AST.LocationType.Value("Edge") ], ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("div_vec"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_binary_operator( serial_utils.make_field_access_expr("vec", [True, 0]), "*", serial_utils.make_field_access_expr("geofac_div"), ), chain=[ AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge") ], ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2t1_vec"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_field_access_expr("rot_vec", [True, 0]), chain=[ AST.LocationType.Value("Edge"), AST.LocationType.Value("Vertex") ], weights=[ serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Double) ]), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2t1_vec"), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("tangent_orientation"), "*", serial_utils.make_field_access_expr("nabla2t1_vec"), ), "/", serial_utils.make_field_access_expr("primal_edge_length"), ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2t2_vec"), serial_utils.make_reduction_over_neighbor_expr( op="+", init=serial_utils.make_literal_access_expr( "0.0", AST.BuiltinType.Double), rhs=serial_utils.make_field_access_expr("div_vec", [True, 0]), chain=[ AST.LocationType.Value("Edge"), AST.LocationType.Value("Cell") ], weights=[ serial_utils.make_literal_access_expr( "-1.0", AST.BuiltinType.Double), serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Double) ]), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2t2_vec"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("nabla2t2_vec"), "/", serial_utils.make_field_access_expr("dual_edge_length"), ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("nabla2_vec"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("nabla2t2_vec"), "-", serial_utils.make_field_access_expr("nabla2t1_vec"), ), "=", ), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) sir = serial_utils.make_sir( gen_outputfile, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( stencil_name, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "div_vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "rot_vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Vertex")], 1), ), serial_utils.make_field( "nabla2t1_vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), is_temporary=True), serial_utils.make_field( "nabla2t2_vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), is_temporary=True), serial_utils.make_field( "nabla2_vec", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "primal_edge_length", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "dual_edge_length", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "tangent_orientation", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), serial_utils.make_field( "geofac_rot", serial_utils.make_field_dimensions_unstructured([ AST.LocationType.Value("Vertex"), AST.LocationType.Value("Edge") ], 1), ), serial_utils.make_field( "geofac_div", serial_utils.make_field_dimensions_unstructured([ AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge") ], 1), ), ], ), ], ) # print the SIR if args.verbose: print(MessageToJson(sir)) # compile code = dawn4py.compile(sir, groups=[], backend=dawn4py.CodeGenBackend.CXXNaiveIco) # write to file print(f"Writing generated code to '{gen_outputfile}'") with open(gen_outputfile, "w") as f: f.write(code)