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 reduction(
self,
expr: expr,
neighborhood: expr,
op: str,
kwargs_keys: t.List[str],
kwargs_values: t.List[expr],
):
kwargs = dict(zip(kwargs_keys, kwargs_values))
assert len(kwargs) == len(kwargs_keys) == len(kwargs_values)
# TODO: what about these hard coded strings?
wrong_kwargs = kwargs.keys() - {"init", "weights"}
if 0 < len(wrong_kwargs):
raise DuskSyntaxError(
f"Unsupported kwargs '{wrong_kwargs}' in reduction!")
neighborhood = self.location_chain(neighborhood)
with self.ctx.location.reduction(neighborhood):
expr = self.expression(expr)
op_map = {"sum": "+", "mul": "*", "min": "min", "max": "max"}
if not op in op_map:
raise DuskSyntaxError(f"Invalid operator '{op}' for reduction!")
if "init" in kwargs:
init = self.expression(kwargs["init"])
else:
# TODO: "min" and "max" are still kinda stupid
# we should use something like this:
# https://en.cppreference.com/w/cpp/types/numeric_limits/max
# the current solution is:
# - appropriate for doubles
# - okish for floats (may trip floating point exceptions but correct outcome)
# - breaks for int (undefined behavior!)
init_map = {
"sum": "0",
"mul": "1",
"min": "1.79769313486231571e+308",
"max": "-1.79769313486231571e+308",
}
init = make_literal_access_expr(
init_map[op], sir.BuiltinType.TypeID.Value("Double"))
op = op_map[op]
weights = None
if "weights" in kwargs:
# TODO: check for `kwargs["weight"].ctx == Load`?
weights = [
self.expression(weight) for weight in kwargs["weights"].elts
]
return make_reduction_over_neighbor_expr(
op,
expr,
init,
neighborhood,
weights,
)
def reduction(
self,
expr: expr,
neighborhood: expr,
op: str,
kwargs_keys: t.List[str],
kwargs_values: t.List[expr],
):
kwargs = dict(zip(kwargs_keys, kwargs_values))
assert len(kwargs) == len(kwargs_keys) == len(kwargs_values)
# TODO: what about these hard coded strings?
wrong_kwargs = kwargs.keys() - {"init", "weights"}
if 0 < len(wrong_kwargs):
raise DuskSyntaxError(f"Unsupported kwargs '{wrong_kwargs}' in reduction!")
neighborhood = self.location_chain(neighborhood)
with self.ctx.location.reduction(neighborhood):
expr = self.expression(expr)
op_map = {"sum": "+", "mul": "*", "min": "min", "max": "max"}
if not op in op_map:
raise DuskSyntaxError(f"Invalid operator '{op}' for reduction!")
if "init" in kwargs:
init = self.expression(kwargs["init"])
else:
# TODO: "min" and "max" are kinda stupid
# we should use something like this:
# https://en.cppreference.com/w/cpp/types/numeric_limits/max
# but for double it should be 1.79769e+308
# FIXME: probably breaks for int
init_map = {
"sum": "0",
"mul": "1",
"min": "9" * 400,
"max": "-" + ("9" * 400),
}
init = make_literal_access_expr(
init_map[op], sir.BuiltinType.TypeID.Value("Double")
)
op = op_map[op]
weights = None
if "weights" in kwargs:
# TODO: check for `kwargs["weight"].ctx == Load`?
weights = [self.expression(weight) for weight in kwargs["weights"].elts]
return make_reduction_over_neighbor_expr(op, expr, init, neighborhood, weights,)
def make_unstructured_stencil_sir(name=None):
OUTPUT_NAME = name if name is not None else "unstructured_stencil"
OUTPUT_FILE = f"{OUTPUT_NAME}.cpp"
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"),
chain=[
SIR.LocationType.Value('Edge'),
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_utils.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))
],
)
],
)
return sir
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():
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():
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 sparse_temporary():
outputfile = "AlsoDemoteWeight"
interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End,
0, 0)
body_ast = serial_utils.make_ast([
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("test"),
),
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("outF"),
serial_utils.make_reduction_over_neighbor_expr(
"+",
serial_utils.make_unstructured_field_access_expr("inF"),
serial_utils.make_literal_access_expr("0.",
AST.BuiltinType.Double),
[
AST.LocationType.Value("Edge"),
AST.LocationType.Value("Cell")
],
weights=[
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF")
]), "="),
])
vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt(
body_ast, interval, AST.VerticalRegion.Forward)
sir = serial_utils.make_sir(
outputfile,
AST.GridType.Value("Unstructured"),
[
serial_utils.make_stencil(
"generated",
serial_utils.make_ast([vertical_region_stmt]),
[
serial_utils.make_field(
"inF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"outF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"test",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"tempF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
is_temporary=True),
],
),
],
)
sim = dawn4py.lower_and_optimize(sir, groups=[])
with open(outputfile, mode="w") as f:
f.write(MessageToJson(sim["generated"]))
os.rename(outputfile, "../input/" + outputfile + ".iir")
def 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)
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):
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)
