def create_vertical_region_stmt1(): """ create a vertical region statement for the stencil """ interval = serial_utils.make_interval(serial_utils.Interval.Start, serial_utils.Interval.Start, 0, 0) body_ast = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_unstructured_field_access_expr("c"), serial_utils.make_binary_operator( serial_utils.make_unstructured_field_access_expr("c"), "/", serial_utils.make_unstructured_field_access_expr("b"), ), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_unstructured_field_access_expr("d"), serial_utils.make_binary_operator( serial_utils.make_unstructured_field_access_expr("d"), "/", serial_utils.make_unstructured_field_access_expr("b"), ), "=", ), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) return vertical_region_stmt
def binop(self, left: expr, op: t.Any, right: expr): py_binops_to_sir_binops = { Add: "+", Sub: "-", Mult: "*", Div: "/", LShift: "<<", RShift: ">>", BitOr: "|", BitXor: "^", BitAnd: "&", } if type(op) in py_binops_to_sir_binops.keys(): op = py_binops_to_sir_binops[type(op)] return make_binary_operator(self.expression(left), op, self.expression(right)) elif isinstance(op, Pow): return make_fun_call_expr( "gridtools::dawn::math::pow", [self.expression(left), self.expression(right)], ) else: raise DuskSyntaxError(f"Unsupported binary operator '{op}'!", op)
def visit_BinOpExpr(self, node: gt_ir.BinOpExpr, **kwargs): left = self.visit(node.lhs) right = self.visit(node.rhs) if node.op.python_symbol == "**": sir = sir_utils.make_fun_call_expr("gridtools::dawn::math::pow", [left, right]) else: op = self._make_operator(node.op) sir = sir_utils.make_binary_operator(left, op, right) return sir
def boolop(self, op, values: t.List): py_boolops_to_sir_boolops = {And: "&&", Or: "||"} op = py_boolops_to_sir_boolops[type(op)] *remainder, last = values binop = self.expression(last) for value in reversed(remainder): binop = make_binary_operator(self.expression(value), op, binop) return binop
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_var_access_expr("dt", is_external=True), "*", serial_utils.make_field_access_expr("in")), "="), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, AST.VerticalRegion.Forward) globals = AST.GlobalVariableMap() globals.map["dt"].double_value = 0.5 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("Edge")], 1), ), serial_utils.make_field( "out", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Edge")], 1), ), ], ) ], global_variables=globals) # print the SIR if args.verbose: print(MessageToJson(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 compare(self, left: expr, op, right: expr): # FIXME: we should probably have a better answer when we need such mappings py_compare_to_sir_compare = { Eq: "==", NotEq: "!=", Lt: "<", LtE: "<=", Gt: ">", GtE: ">=", } if type(op) not in py_compare_to_sir_compare.keys(): raise DuskSyntaxError(f"Unsupported comparison operator '{op}'!", op) op = py_compare_to_sir_compare[type(op)] return make_binary_operator(self.expression(left), op, self.expression(right))
def create_vertical_region_stmt3(): """ create a vertical region statement for the stencil """ interval = sir_utils.make_interval(sir_utils.Interval.Start, sir_utils.Interval.End, 0, -1) body_ast = sir_utils.make_ast([ sir_utils.make_assignment_stmt( sir_utils.make_unstructured_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("c"), "*", sir_utils.make_unstructured_field_access_expr( "d", sir_utils.make_unstructured_offset(False), 1)), "-=") ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Backward) return vertical_region_stmt
def create_vertical_region_stmt2(): """ create a vertical region statement for the stencil """ interval = sir_utils.make_interval(sir_utils.Interval.Start, sir_utils.Interval.End, 1, 0) body_ast = sir_utils.make_ast([ sir_utils.make_var_decl_stmt( sir_utils.make_type(sir_utils.BuiltinType.Float), "m", 0, "=", sir_utils.make_expr( sir_utils.make_binary_operator( sir_utils.make_literal_access_expr( "1.0", sir_utils.BuiltinType.Float), "/", sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("b"), "-", sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("a"), "*", sir_utils.make_unstructured_field_access_expr( "c", sir_utils.make_unstructured_offset(False), -1)))))), sir_utils.make_assignment_stmt( sir_utils.make_unstructured_field_access_expr("c"), sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("c"), "*", sir_utils.make_var_access_expr("m")), "="), sir_utils.make_assignment_stmt( sir_utils.make_unstructured_field_access_expr("d"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("d"), "-", sir_utils.make_binary_operator( sir_utils.make_unstructured_field_access_expr("a"), "*", sir_utils.make_unstructured_field_access_expr( "d", sir_utils.make_unstructured_offset(False), -1))), "*", sir_utils.make_var_access_expr("m")), "=") ]) vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) return vertical_region_stmt
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 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 main(args: argparse.Namespace): interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) # out = in_1 on inner cells body_ast_1 = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_field_access_expr("in_1"), "=", ) ]) vertical_region_stmt_1 = serial_utils.make_vertical_region_decl_stmt( body_ast_1, interval, AST.VerticalRegion.Forward, serial_utils.make_magic_num_interval(0, 1, 0, 0)) # out = out + in_2 on inner cells # should be merge-able to last stage body_ast_2 = 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("out"), "+", serial_utils.make_field_access_expr("in_2"), ), "=") ]) vertical_region_stmt_2 = serial_utils.make_vertical_region_decl_stmt( body_ast_2, interval, AST.VerticalRegion.Forward, serial_utils.make_interval(2, 3, 0, 0)) # out = out + in_3 on lateral boundary cells # out = in_1 on inner cells body_ast_3 = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_field_access_expr("in_3"), "=", ) ]) vertical_region_stmt_3 = serial_utils.make_vertical_region_decl_stmt( body_ast_3, interval, AST.VerticalRegion.Forward, serial_utils.make_interval(3, 4, 0, 0)) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Unstructured"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([ vertical_region_stmt_1, vertical_region_stmt_2, vertical_region_stmt_3 ]), [ serial_utils.make_field( "in_1", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "in_2", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 1), ), serial_utils.make_field( "in_3", 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 if args.verbose: print(MessageToJson(sir)) # compile pass_groups = dawn4py.default_pass_groups() pass_groups.insert(1, dawn4py.PassGroup.MultiStageMerger) pass_groups.insert(1, dawn4py.PassGroup.StageMerger) # code = dawn4py.compile(sir, groups=pass_groups, # backend=dawn4py.CodeGenBackend.CXXNaiveIco, merge_stages=True, merge_do_methods=True) code = dawn4py.compile(sir, groups=pass_groups, 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(): stencil_name = "ICON_laplacian_diamond_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( [ # fill sparse dimension vn vert using the loop concept sir_utils.make_loop_stmt( [sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("vn_vert"), sir_utils.make_binary_operator( sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "u_vert", [True, 0]), "*", sir_utils.make_field_access_expr("primal_normal_x", [True, 0])), "+", sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "v_vert", [True, 0]), "*", sir_utils.make_field_access_expr("primal_normal_y", [True, 0])), ), "=")], [SIR.LocationType.Value( "Edge"), SIR.LocationType.Value("Cell"), SIR.LocationType.Value("Vertex")] ), # dvt_tang for smagorinsky sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("dvt_tang"), 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_binary_operator(sir_utils.make_field_access_expr( "u_vert", [True, 0]), "*", sir_utils.make_field_access_expr("dual_normal_x", [True, 0])), "+", sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "v_vert", [True, 0]), "*", sir_utils.make_field_access_expr("dual_normal_y", [True, 0])), ), chain=[SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), 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_literal_access_expr( "0.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double)] ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("dvt_tang"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("dvt_tang"), "*", sir_utils.make_field_access_expr("tangent_orientation")), "="), # dvt_norm for smagorinsky sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("dvt_norm"), 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_binary_operator(sir_utils.make_field_access_expr( "u_vert", [True, 0]), "*", sir_utils.make_field_access_expr("dual_normal_x", [True, 0])), "+", sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "v_vert", [True, 0]), "*", sir_utils.make_field_access_expr("dual_normal_y", [True, 0])), ), chain=[SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), SIR.LocationType.Value("Vertex")], weights=[sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "1.0", SIR.BuiltinType.Double)] ), "=", ), # compute smagorinsky sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("kh_smag_1"), 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_literal_access_expr( "-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double)] ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("kh_smag_1"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_field_access_expr("kh_smag_1"), "*", sir_utils.make_field_access_expr("tangent_orientation")), "*", sir_utils.make_field_access_expr("inv_primal_edge_length")), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("dvt_norm"), "*", sir_utils.make_field_access_expr("inv_vert_vert_length"))), "="), sir_utils.make_assignment_stmt(sir_utils.make_field_access_expr("kh_smag_1"), sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "kh_smag_1"), "*", sir_utils.make_field_access_expr("kh_smag_1"))), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("kh_smag_2"), 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_literal_access_expr( "0.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr( "0.0", SIR.BuiltinType.Double), 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("kh_smag_2"), sir_utils.make_binary_operator( sir_utils.make_binary_operator( sir_utils.make_field_access_expr("kh_smag_2"), "*", sir_utils.make_field_access_expr("inv_vert_vert_length")), "+", sir_utils.make_binary_operator( sir_utils.make_field_access_expr("dvt_tang"), "*", sir_utils.make_field_access_expr("inv_primal_edge_length"))), "="), sir_utils.make_assignment_stmt(sir_utils.make_field_access_expr("kh_smag_2"), sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "kh_smag_2"), "*", sir_utils.make_field_access_expr("kh_smag_2"))), # currently not able to forward a sqrt, so this is technically kh_smag**2 sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("kh_smag"), sir_utils.make_binary_operator(sir_utils.make_field_access_expr("diff_multfac_smag"), "*", sir_utils.make_fun_call_expr("math::sqrt", [sir_utils.make_binary_operator(sir_utils.make_field_access_expr( "kh_smag_1"), "+", sir_utils.make_field_access_expr("kh_smag_2"))])), "="), # 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_binary_operator(sir_utils.make_literal_access_expr( "4.0", SIR.BuiltinType.Double), "*", 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_binary_operator( sir_utils.make_field_access_expr( "inv_primal_edge_length"), '*', sir_utils.make_field_access_expr( "inv_primal_edge_length")), sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "inv_primal_edge_length"), '*', sir_utils.make_field_access_expr( "inv_primal_edge_length")), sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "inv_vert_vert_length"), '*', sir_utils.make_field_access_expr( "inv_vert_vert_length")), sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "inv_vert_vert_length"), '*', sir_utils.make_field_access_expr( "inv_vert_vert_length")), ] ), "=", ), sir_utils.make_assignment_stmt( sir_utils.make_field_access_expr("nabla2"), sir_utils.make_binary_operator( sir_utils.make_field_access_expr("nabla2"), "-", sir_utils.make_binary_operator( sir_utils.make_binary_operator(sir_utils.make_binary_operator(sir_utils.make_literal_access_expr( "8.0", SIR.BuiltinType.Double), "*", sir_utils.make_field_access_expr("vn")), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "inv_primal_edge_length"), "*", sir_utils.make_field_access_expr( "inv_primal_edge_length"))), "+", sir_utils.make_binary_operator(sir_utils.make_binary_operator(sir_utils.make_literal_access_expr( "8.0", SIR.BuiltinType.Double), "*", sir_utils.make_field_access_expr("vn")), "*", sir_utils.make_binary_operator( sir_utils.make_field_access_expr( "inv_vert_vert_length"), "*", sir_utils.make_field_access_expr( "inv_vert_vert_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( "diff_multfac_smag", 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( "inv_primal_edge_length", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "inv_vert_vert_length", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "u_vert", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Vertex")], 1 ), ), sir_utils.make_field( "v_vert", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Vertex")], 1 ), ), sir_utils.make_field( "primal_normal_x", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), SIR.LocationType.Value("Vertex")], 1 ), ), sir_utils.make_field( "primal_normal_y", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), SIR.LocationType.Value("Vertex")], 1 ), ), sir_utils.make_field( "dual_normal_x", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), SIR.LocationType.Value("Vertex")], 1 ), ), sir_utils.make_field( "dual_normal_y", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge"), SIR.LocationType.Value( "Cell"), SIR.LocationType.Value("Vertex")], 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( "vn", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "dvt_tang", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "dvt_norm", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "kh_smag_1", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "kh_smag_2", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 1 ), ), sir_utils.make_field( "kh_smag", sir_utils.make_field_dimensions_unstructured( [SIR.LocationType.Value("Edge")], 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 = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) # create the out = reduce(sparse_CE * in) statement body_ast = serial_utils.make_ast( [ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out"), serial_utils.make_reduction_over_neighbor_expr( "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("sparse_CE"), "*", serial_utils.make_field_access_expr("in"), ), serial_utils.make_literal_access_expr("1.0", AST.BuiltinType.Float), chain=[AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge")], ), "=", ) ] ) 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("Edge")], 1 ), ), serial_utils.make_field( "sparse_CE", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell"), AST.LocationType.Value("Edge")], 1 ), ), serial_utils.make_field( "out", serial_utils.make_field_dimensions_unstructured( [AST.LocationType.Value("Cell")], 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 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 = serial_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0) # create the laplace statement body_ast = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("out", [0, 0, 0]), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("in", [0, 0, 0]), "*", serial_utils.make_literal_access_expr( "-4.0", serial_utils.BuiltinType.Float), ), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("in", [1, 0, 0]), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "in", [-1, 0, 0]), "+", serial_utils.make_binary_operator( serial_utils.make_field_access_expr( "in", [0, 1, 0]), "+", serial_utils.make_field_access_expr( "in", [0, -1, 0]), ), ), ), ), "/", serial_utils.make_binary_operator( serial_utils.make_var_access_expr("dx", is_external=True), "*", serial_utils.make_var_access_expr("dx", is_external=True), ), ), "=", ), ]) vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt( body_ast, interval, SIR.VerticalRegion.Forward) stencils_globals = serial_utils.GlobalVariableMap() stencils_globals.map["dx"].double_value = 0.0 sir = serial_utils.make_sir( OUTPUT_FILE, SIR.GridType.Value("Cartesian"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([vertical_region_stmt]), [ serial_utils.make_field( "out", serial_utils.make_field_dimensions_cartesian()), serial_utils.make_field( "in", serial_utils.make_field_dimensions_cartesian()), ], ) ], global_variables=stencils_globals, ) # print the SIR if args.verbose: serial_utils.pprint(sir) # serialize the SIR to file sir_file = open("./laplacian_stencil_from_python.sir", "wb") sir_file.write(serial_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 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")), 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"), "*", 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"), 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"), 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(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)
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 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): 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(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)
def main(args: argparse.Namespace): # ---- First vertical region statement ---- interval_1 = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, 0) body_ast_1 = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("c"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("c"), "/", serial_utils.make_field_access_expr("b"), ), "=", ) ]) vertical_region_stmt_1 = serial_utils.make_vertical_region_decl_stmt( body_ast_1, interval_1, AST.VerticalRegion.Forward) # ---- Second vertical region statement ---- interval_2 = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 1, 0) body_ast_2 = serial_utils.make_ast([ serial_utils.make_var_decl_stmt( serial_utils.make_type(AST.BuiltinType.Integer), "m", 0, "=", serial_utils.make_expr( serial_utils.make_binary_operator( serial_utils.make_literal_access_expr( "1.0", AST.BuiltinType.Float), "/", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("b"), "-", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("a"), "*", serial_utils.make_field_access_expr( "c", [0, 0, -1]), ), ), )), ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("c"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("c"), "*", serial_utils.make_var_access_expr("m")), "=", ), serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("d"), serial_utils.make_binary_operator( serial_utils.make_binary_operator( serial_utils.make_field_access_expr("d"), "-", serial_utils.make_binary_operator( serial_utils.make_field_access_expr("a"), "*", serial_utils.make_field_access_expr("d", [0, 0, -1]), ), ), "*", serial_utils.make_var_access_expr("m"), ), "=", ), ]) vertical_region_stmt_2 = serial_utils.make_vertical_region_decl_stmt( body_ast_2, interval_2, AST.VerticalRegion.Forward) # ---- Third vertical region statement ---- interval_3 = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End, 0, -1) body_ast_3 = serial_utils.make_ast([ serial_utils.make_assignment_stmt( serial_utils.make_field_access_expr("d"), serial_utils.make_binary_operator( serial_utils.make_field_access_expr("c"), "*", serial_utils.make_field_access_expr("d", [0, 0, 1]), ), "-=", ) ]) vertical_region_stmt_3 = serial_utils.make_vertical_region_decl_stmt( body_ast_3, interval_3, AST.VerticalRegion.Backward) sir = serial_utils.make_sir( OUTPUT_FILE, AST.GridType.Value("Cartesian"), [ serial_utils.make_stencil( OUTPUT_NAME, serial_utils.make_ast([ vertical_region_stmt_1, vertical_region_stmt_2, vertical_region_stmt_3 ]), [ serial_utils.make_field( "a", serial_utils.make_field_dimensions_cartesian()), serial_utils.make_field( "b", serial_utils.make_field_dimensions_cartesian()), serial_utils.make_field( "c", serial_utils.make_field_dimensions_cartesian()), serial_utils.make_field( "d", serial_utils.make_field_dimensions_cartesian()), ], ) ], ) # print the SIR if args.verbose: print(MessageToJson(sir)) # compile pass_groups = dawn4py.default_pass_groups() pass_groups.insert(1, dawn4py.PassGroup.MultiStageMerger) code = dawn4py.compile(sir, groups=pass_groups, 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)