def test_compilation(unstructured_sir_with_reference_code):
sir, reference_code = unstructured_sir_with_reference_code
dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco)
dawn4py.codegen(
dawn4py.optimize(dawn4py.lower_and_optimize(sir, groups=[]),
groups=dawn4py.default_pass_groups()),
backend=dawn4py.CodeGenBackend.CXXNaiveIco,
)
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),
),
],
),
],
)
dawn4py.compile(sir,
backend=backend,
serialize_iir=True,
output_file=outputfile)
os.rename(outputfile + ".0.iir", "../input/" + outputfile + ".iir")
def transpile(in_path: str,
out_path: str,
backend: str = default_backend,
dump_sir: bool = False) -> None:
with open(in_path, "r") as in_file:
in_str = in_file.read()
in_ast = ast.parse(in_str, filename=in_path, type_comments=True)
grammar = Grammar()
# TODO: handle errors in different stencils separately
stencils = [grammar.stencil(node) for node in iter_stencils(in_ast)]
sir = make_sir(in_path, GridType.Value("Unstructured"), stencils)
if dump_sir:
out_name = os.path.splitext(out_path)[0]
with open(out_name + ".json", "w+") as f:
f.write(sir_to_json(sir))
out_code = compile(sir, backend=backend_map[backend])
with open(out_path, "w") as out_file:
out_file.write(out_code)
def main(args: argparse.Namespace):
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0)
# create the out = in[i+1] statement
body_ast = sir_utils.make_ast(
[
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out"),
sir_utils.make_reduction_over_neighbor_expr(
"+",
sir_utils.make_literal_access_expr("1.0", SIR.BuiltinType.Float),
sir_utils.make_field_access_expr("in"),
lhs_location=SIR.LocationType.Value("Edge"),
rhs_location=SIR.LocationType.Value("Cell"),
),
"=",
)
]
)
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("Cell")], 1
),
),
sir_utils.make_field(
"out",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 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 sir_to_cpp(sir: SIR,
verbose: bool = False,
groups: List = [],
backend=default_backend) -> str:
if verbose:
set_verbosity(LogLevel.All)
# TODO: default pass groups are bugged in Dawn, need to pass empty list of groups
return compile(sir, groups=groups, backend=backend_map[backend])
def test_compilation(unstructure_sir_with_reference_code):
sir, reference_code = unstructure_sir_with_reference_code
backend = "c++-naive-ico"
unit_info = {}
code = dawn4py.compile(sir, backend=backend, unit_info=unit_info)
# with open("new_code.hpp", "w") as f:
# f.write(code)
# assert code == reference_code
assert {"filename", "pp_defines", "stencils", "globals"} == set(unit_info.keys())
unit_info = {}
code = dawn4py.compile(sir, backend=backend, unit_info=unit_info)
# assert code == reference_code
assert {"filename", "pp_defines", "stencils", "globals"} == set(unit_info.keys())
unit_info = {}
code = dawn4py.compile(sir, backend=backend, unit_info=unit_info)
# assert code == reference_code
assert {"filename", "pp_defines", "stencils", "globals"} == set(unit_info.keys())
def main(args: argparse.Namespace):
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0,
0)
body_ast = sir_utils.make_ast([
sir_utils.make_loop_stmt([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out"),
sir_utils.make_field_access_expr("in"),
"=",
)
], [SIR.LocationType.Value("Cell"),
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(
"out",
sir_utils.make_field_dimensions_unstructured([
SIR.LocationType.Value("Cell"),
SIR.LocationType.Value("Edge")
], 1),
),
sir_utils.make_field(
"in",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1),
),
],
)
],
)
# print the SIR
if args.verbose:
sir_utils.pprint(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 = 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("x"),
serial_utils.make_literal_access_expr("1", AST.BuiltinType.Double),
"=",
)
])
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(
"out",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"x",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1), True),
],
),
],
)
# 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 test_compilation(grid_sir_with_reference_code):
sir, reference_code = grid_sir_with_reference_code
for backend in (
dawn4py.CodeGenBackend.CXXNaive,
dawn4py.CodeGenBackend.GridTools,
dawn4py.CodeGenBackend.CUDA,
):
dawn4py.compile(sir, backend=backend)
dawn4py.codegen(
dawn4py.optimize(
dawn4py.lower_and_optimize(sir, groups=[]),
groups=[
dawn4py.PassGroup.SetStageName,
dawn4py.PassGroup.StageReordering,
# dawn4py.PassGroup.StageMerger,
dawn4py.PassGroup.SetCaches,
dawn4py.PassGroup.SetBlockSize,
],
),
backend=backend,
)
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 main(args: argparse.Namespace):
sir = serial_utils.make_sir(
OUTPUT_FILE,
AST.GridType.Value("Unstructured"),
[
serial_utils.make_stencil(
OUTPUT_NAME,
serial_utils.make_ast([
create_vertical_region_stmt1(),
create_vertical_region_stmt2(),
create_vertical_region_stmt3(),
]),
[
serial_utils.make_field(
"a",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"b",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"c",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"d",
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):
sir = sir_utils.make_sir(
OUTPUT_FILE,
SIR.GridType.Value("Unstructured"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([
create_vertical_region_stmt1(),
create_vertical_region_stmt2(),
create_vertical_region_stmt3(),
]),
[
sir_utils.make_field(
"a",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"b",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"c",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"d",
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):
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_field_access_expr("out", [0, 0, 0]),
serial_utils.make_field_access_expr("in", [1, 0, 0]),
"=",
)
])
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("Cartesian"),
[
serial_utils.make_stencil(
OUTPUT_NAME,
serial_utils.make_ast([vertical_region_stmt]),
[
serial_utils.make_field(
"in", serial_utils.make_field_dimensions_cartesian()),
serial_utils.make_field(
"out", serial_utils.make_field_dimensions_cartesian()),
],
)
],
)
# print the SIR
if args.verbose:
print(MessageToJson(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():
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 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 make_extension_sources(self, gt_backend_t: str) -> Dict[str, Any]:
stencil_short_name = self.builder.stencil_id.qualified_name.split(
".")[-1]
backend_opts = dict(**self.builder.options.backend_opts)
dawn_namespace = self.DAWN_BACKEND_NS
dump_sir_opt = backend_opts.get("dump_sir", False)
if dump_sir_opt:
if isinstance(dump_sir_opt, str):
dump_sir_file = dump_sir_opt
else:
assert isinstance(dump_sir_opt, bool)
dump_sir_file = f"{stencil_short_name}_gt4py.sir"
with open(dump_sir_file, "w") as f:
f.write(sir_utils.to_json(self.sir))
# Get list of pass groups
if "no_opt" in backend_opts:
pass_groups = []
elif "opt_groups" in backend_opts:
pass_groups = [
DAWN_PASS_GROUPS[k] for k in backend_opts["opt_groups"]
]
if "default_opt" in backend_opts:
raise ValueError(
"Do not add 'default_opt' when opt 'opt_groups'. " +
"Instead, append dawn4py.default_pass_groups()")
else:
pass_groups = dawn4py.default_pass_groups()
# If present, parse backend string
dawn_backend = DAWN_CODEGEN_BACKENDS[self.DAWN_BACKEND_NAME
or "GridTools"]
dawn_opts = {
key: value
for key, value in backend_opts.items()
if key in _DAWN_TOOLCHAIN_OPTIONS.keys()
}
source = dawn4py.compile(self.sir,
groups=pass_groups,
backend=dawn_backend,
run_with_sync=False,
**dawn_opts)
stencil_unique_name = self.pyext_class_name
module_name = self.pyext_module_name
pyext_sources = {f"_dawn_{stencil_short_name}.hpp": source}
arg_fields = [{
"name": field.name,
"dtype": self._DATA_TYPE_TO_CPP[field.data_type],
"layout_id": i
} for i, field in enumerate(self.builder.definition_ir.api_fields)]
header_file = "computation.hpp"
parameters = []
for parameter in self.builder.definition_ir.parameters:
if parameter.data_type in [gt_ir.DataType.BOOL]:
dtype = "bool"
elif parameter.data_type in [
gt_ir.DataType.INT8,
gt_ir.DataType.INT16,
gt_ir.DataType.INT32,
gt_ir.DataType.INT64,
]:
dtype = "int"
elif parameter.data_type in [
gt_ir.DataType.FLOAT32, gt_ir.DataType.FLOAT64
]:
dtype = "double"
else:
assert False, "Wrong data_type for parameter"
parameters.append({"name": parameter.name, "dtype": dtype})
template_args = dict(
arg_fields=arg_fields,
dawn_namespace=dawn_namespace,
gt_backend=gt_backend_t,
header_file=header_file,
module_name=module_name,
parameters=parameters,
stencil_short_name=stencil_short_name,
stencil_unique_name=stencil_unique_name,
)
for key, file_name in self.TEMPLATE_FILES.items():
with open(os.path.join(self.TEMPLATE_DIR, file_name), "r") as f:
template = jinja2.Template(f.read())
pyext_sources[key] = template.render(**template_args)
return pyext_sources
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 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 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):
# ---- 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)
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():
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 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(
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 = 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(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)
