def loop_stmt(self, neighborhood, body: t.List):
neighborhood = self.location_chain(neighborhood)
with self.ctx.location.loop_stmt(neighborhood):
body = self.statements(body)
return make_loop_stmt(body, neighborhood)
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 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():
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)
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)