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 visit_BuiltinLiteral(self, node: gt_ir.BuiltinLiteral):
if node.value == gt_ir.Builtin.TRUE:
source = "true"
elif node.value == gt_ir.Builtin.FALSE:
source = "false"
else: # Builtin.NONE
source = "nullptr"
return sir_utils.make_literal_access_expr(source, type=SIR.BuiltinType.Boolean)
def function_call():
outputfile = "../input/test_set_stage_location_type_function_call.sir"
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0,
0)
fun_ast = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out"),
sir_utils.make_literal_access_expr(
value="2.0", type=sir_utils.BuiltinType.Float),
"=",
),
])
arg_field = sir_utils.make_field(
"out",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1))
fun = sir_utils.make_stencil_function(
name='f',
asts=[fun_ast],
intervals=[interval],
arguments=[sir_utils.make_stencil_function_arg(arg_field)])
body_ast = sir_utils.make_ast([
sir_utils.make_expr_stmt(expr=sir_utils.make_stencil_fun_call_expr(
callee="f",
arguments=[sir_utils.make_field_access_expr("out_cell")])),
])
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(
"out_cell",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
],
),
],
functions=[fun])
f = open(outputfile, "w")
f.write(MessageToJson(sir))
f.close()
def main(args: argparse.Namespace):
interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End,
0, 0)
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 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 visit_ScalarLiteral(self, node: gt_ir.ScalarLiteral, **kwargs):
assert node.data_type != gt_ir.DataType.INVALID
if node.data_type in (gt_ir.DataType.AUTO, gt_ir.DataType.DEFAULT):
sir_type = SIR.BuiltinType.type_id = SIR.BuiltinType.Auto
elif node.data_type in (
gt_ir.DataType.INT8,
gt_ir.DataType.INT16,
gt_ir.DataType.INT32,
gt_ir.DataType.INT64,
):
sir_type = SIR.BuiltinType.Integer
elif node.data_type in (gt_ir.DataType.FLOAT32, gt_ir.DataType.FLOAT64):
sir_type = SIR.BuiltinType.Float
else:
assert False, "Unrecognized data type"
return sir_utils.make_literal_access_expr(value=repr(node.value), type=sir_type)
def create_boundary_correction_region(
value="0", i_interval=None, j_interval=None
) -> SIR.VerticalRegionDeclStmt:
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0, 0)
boundary_body = sir_utils.make_ast(
[
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out", [0, 0, 0]),
sir_utils.make_literal_access_expr(value, SIR.BuiltinType.Float),
"=",
)
]
)
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
boundary_body, interval, SIR.VerticalRegion.Forward, IRange=i_interval, JRange=j_interval
)
return vertical_region_stmt
def create_vertical_region_stmt2():
""" create a vertical region statement for the stencil
"""
interval = sir_utils.make_interval(sir_utils.Interval.Start,
sir_utils.Interval.End, 1, 0)
body_ast = sir_utils.make_ast([
sir_utils.make_var_decl_stmt(
sir_utils.make_type(sir_utils.BuiltinType.Float), "m", 0, "=",
sir_utils.make_expr(
sir_utils.make_binary_operator(
sir_utils.make_literal_access_expr(
"1.0", sir_utils.BuiltinType.Float), "/",
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("b"),
"-",
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("a"),
"*",
sir_utils.make_unstructured_field_access_expr(
"c", sir_utils.make_unstructured_offset(False),
-1)))))),
sir_utils.make_assignment_stmt(
sir_utils.make_unstructured_field_access_expr("c"),
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("c"), "*",
sir_utils.make_var_access_expr("m")), "="),
sir_utils.make_assignment_stmt(
sir_utils.make_unstructured_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("d"), "-",
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("a"),
"*",
sir_utils.make_unstructured_field_access_expr(
"d", sir_utils.make_unstructured_offset(False),
-1))), "*", sir_utils.make_var_access_expr("m")),
"=")
])
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
body_ast, interval, SIR.VerticalRegion.Forward)
return vertical_region_stmt
def constant(self, value):
# TODO: properly distinguish between float and double
built_in_type_map = {bool: "Boolean", int: "Integer", float: "Double"}
if type(value) not in built_in_type_map.keys():
raise DuskSyntaxError(
f"Unsupported constant '{value}' of type '{type(value)}'!", value
)
_type = sir.BuiltinType.TypeID.Value(built_in_type_map[type(value)])
if isinstance(value, bool):
value = "true" if value else "false"
else:
# TODO: does `str` really work here? (what about NaNs, precision, 1e11 notation, etc)
value = str(value)
return make_literal_access_expr(value, _type,)
def main():
stencil_name = "ICON_laplacian_stencil"
gen_outputfile = f"{stencil_name}.cpp"
sir_outputfile = f"{stencil_name}.sir"
interval = sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.End, 0, 0)
body_ast = sir_utils.make_ast(
[
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("rot_vec"),
sir_utils.make_reduction_over_neighbor_expr(
op="+",
init=sir_utils.make_literal_access_expr(
"0.0", SIR.BuiltinType.Double),
rhs=sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("vec"),
"*",
sir_utils.make_field_access_expr("geofac_rot")),
chain=[SIR.LocationType.Value(
"Vertex"), SIR.LocationType.Value("Edge")]
),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("div_vec"),
sir_utils.make_reduction_over_neighbor_expr(
op="+",
init=sir_utils.make_literal_access_expr(
"0.0", SIR.BuiltinType.Double),
rhs=sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("vec"),
"*",
sir_utils.make_field_access_expr("geofac_div")),
chain=[SIR.LocationType.Value(
"Cell"), SIR.LocationType.Value("Edge")]
),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("nabla2t1_vec"),
sir_utils.make_reduction_over_neighbor_expr(
op="+",
init=sir_utils.make_literal_access_expr(
"0.0", SIR.BuiltinType.Double),
rhs=sir_utils.make_field_access_expr("rot_vec"),
chain=[SIR.LocationType.Value(
"Edge"), SIR.LocationType.Value("Vertex")],
weights=[sir_utils.make_literal_access_expr(
"-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr(
"1.0", SIR.BuiltinType.Double)]
),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("nabla2t1_vec"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr(
"tangent_orientation"),
"*",
sir_utils.make_field_access_expr("nabla2t1_vec")),
"/",
sir_utils.make_field_access_expr("primal_edge_length")),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("nabla2t2_vec"),
sir_utils.make_reduction_over_neighbor_expr(
op="+",
init=sir_utils.make_literal_access_expr(
"0.0", SIR.BuiltinType.Double),
rhs=sir_utils.make_field_access_expr("div_vec"),
chain=[SIR.LocationType.Value(
"Edge"), SIR.LocationType.Value("Cell")],
weights=[sir_utils.make_literal_access_expr(
"-1.0", SIR.BuiltinType.Double), sir_utils.make_literal_access_expr(
"1.0", SIR.BuiltinType.Double)]
),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("nabla2t2_vec"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("nabla2t2_vec"),
"/",
sir_utils.make_field_access_expr("dual_edge_length")),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("nabla2_vec"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("nabla2t2_vec"),
"-",
sir_utils.make_field_access_expr("nabla2t1_vec")),
"=",
),
]
)
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
body_ast, interval, SIR.VerticalRegion.Forward
)
sir = sir_utils.make_sir(
gen_outputfile,
SIR.GridType.Value("Unstructured"),
[
sir_utils.make_stencil(
stencil_name,
sir_utils.make_ast([vertical_region_stmt]),
[
sir_utils.make_field(
"vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"div_vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1
),
),
sir_utils.make_field(
"rot_vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Vertex")], 1
),
),
sir_utils.make_field(
"nabla2t1_vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"nabla2t2_vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"nabla2_vec",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"primal_edge_length",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"dual_edge_length",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"tangent_orientation",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"geofac_rot",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value(
"Vertex"), SIR.LocationType.Value("Edge")], 1
),
),
sir_utils.make_field(
"geofac_div",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value(
"Cell"), SIR.LocationType.Value("Edge")], 1
),
),
],
),
],
)
# write SIR to file (for debugging purposes)
f = open(sir_outputfile, "w")
f.write(MessageToJson(sir))
f.close()
# compile
code = dawn4py.compile(sir, backend="c++-naive-ico")
# write to file
print(f"Writing generated code to '{gen_outputfile}'")
with open(gen_outputfile, "w") as f:
f.write(code)
def make_hori_diff_stencil_sir(name=None):
OUTPUT_NAME = name if name is not None else "hori_diff_stencil"
OUTPUT_FILE = f"{OUTPUT_NAME}.cpp"
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0,
0)
# create the stencil body AST
body_ast = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("lap"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_literal_access_expr("-4.0",
SIR.BuiltinType.Float),
"*",
sir_utils.make_field_access_expr("in"),
),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("coeff"),
"*",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("in", [1, 0, 0]),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("in", [-1, 0, 0]),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr(
"in", [0, 1, 0]),
"+",
sir_utils.make_field_access_expr(
"in", [0, -1, 0]),
),
),
),
),
),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_literal_access_expr("-4.0",
SIR.BuiltinType.Float),
"*",
sir_utils.make_field_access_expr("lap"),
),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("coeff"),
"*",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("lap", [1, 0, 0]),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr(
"lap", [-1, 0, 0]),
"+",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr(
"lap", [0, 1, 0]),
"+",
sir_utils.make_field_access_expr(
"lap", [0, -1, 0]),
),
),
),
),
),
"=",
),
])
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
body_ast, interval, SIR.VerticalRegion.Forward)
sir = sir_utils.make_sir(
OUTPUT_FILE,
sir_utils.GridType.Value("Cartesian"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([vertical_region_stmt]),
[
sir_utils.make_field(
"in", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"out", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"coeff", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"lap",
sir_utils.make_field_dimensions_cartesian(),
is_temporary=True),
],
)
],
)
return sir
def make_tridiagonal_solve_stencil_sir(name=None):
OUTPUT_NAME = name if name is not None else "tridiagonal_solve_stencil"
OUTPUT_FILE = f"{OUTPUT_NAME}.cpp"
# ---- First vertical region statement ----
interval_1 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
0, 0)
body_ast_1 = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("c"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"),
"/",
sir_utils.make_field_access_expr("b"),
),
"=",
)
])
vertical_region_stmt_1 = sir_utils.make_vertical_region_decl_stmt(
body_ast_1, interval_1, SIR.VerticalRegion.Forward)
# ---- Second vertical region statement ----
interval_2 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
1, 0)
body_ast_2 = sir_utils.make_ast([
sir_utils.make_var_decl_stmt(
sir_utils.make_type(SIR.BuiltinType.Integer),
"m",
0,
"=",
sir_utils.make_expr(
sir_utils.make_binary_operator(
sir_utils.make_literal_access_expr("1.0",
SIR.BuiltinType.Float),
"/",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("b"),
"-",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("a"),
"*",
sir_utils.make_field_access_expr("c", [0, 0, -1]),
),
),
)),
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("c"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"), "*",
sir_utils.make_var_access_expr("m")),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("d"),
"-",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("a"),
"*",
sir_utils.make_field_access_expr("d", [0, 0, -1]),
),
),
"*",
sir_utils.make_var_access_expr("m"),
),
"=",
),
])
vertical_region_stmt_2 = sir_utils.make_vertical_region_decl_stmt(
body_ast_2, interval_2, SIR.VerticalRegion.Forward)
# ---- Third vertical region statement ----
interval_3 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
0, -1)
body_ast_3 = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"),
"*",
sir_utils.make_field_access_expr("d", [0, 0, 1]),
),
"-=",
)
])
vertical_region_stmt_3 = sir_utils.make_vertical_region_decl_stmt(
body_ast_3, interval_3, SIR.VerticalRegion.Backward)
sir = sir_utils.make_sir(
OUTPUT_FILE,
sir_utils.GridType.Value("Cartesian"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([
vertical_region_stmt_1, vertical_region_stmt_2,
vertical_region_stmt_3
]),
[
sir_utils.make_field(
"a", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"b", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"c", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"d", sir_utils.make_field_dimensions_cartesian()),
],
)
],
)
return sir
def main(args: argparse.Namespace):
interval = 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(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)
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 = "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 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(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):
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):
# ---- 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_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 sparse_temporary():
outputfile = "AlsoDemoteWeight"
interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End,
0, 0)
body_ast = serial_utils.make_ast([
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("test"),
),
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("outF"),
serial_utils.make_reduction_over_neighbor_expr(
"+",
serial_utils.make_unstructured_field_access_expr("inF"),
serial_utils.make_literal_access_expr("0.",
AST.BuiltinType.Double),
[
AST.LocationType.Value("Edge"),
AST.LocationType.Value("Cell")
],
weights=[
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF"),
serial_utils.make_unstructured_field_access_expr("tempF")
]), "="),
])
vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt(
body_ast, interval, AST.VerticalRegion.Forward)
sir = serial_utils.make_sir(
outputfile,
AST.GridType.Value("Unstructured"),
[
serial_utils.make_stencil(
"generated",
serial_utils.make_ast([vertical_region_stmt]),
[
serial_utils.make_field(
"inF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"outF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"test",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"tempF",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
is_temporary=True),
],
),
],
)
sim = dawn4py.lower_and_optimize(sir, groups=[])
with open(outputfile, mode="w") as f:
f.write(MessageToJson(sim["generated"]))
os.rename(outputfile, "../input/" + outputfile + ".iir")
def main():
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):
# ---- First vertical region statement ----
interval_1 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
0, 0)
body_ast_1 = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("c"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"),
"/",
sir_utils.make_field_access_expr("b"),
),
"=",
)
])
vertical_region_stmt_1 = sir_utils.make_vertical_region_decl_stmt(
body_ast_1, interval_1, SIR.VerticalRegion.Forward)
# ---- Second vertical region statement ----
interval_2 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
1, 0)
body_ast_2 = sir_utils.make_ast([
sir_utils.make_var_decl_stmt(
sir_utils.make_type(SIR.BuiltinType.Integer),
"m",
0,
"=",
sir_utils.make_expr(
sir_utils.make_binary_operator(
sir_utils.make_literal_access_expr("1.0",
SIR.BuiltinType.Float),
"/",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("b"),
"-",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("a"),
"*",
sir_utils.make_field_access_expr("c", [0, 0, -1]),
),
),
)),
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("c"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"), "*",
sir_utils.make_var_access_expr("m")),
"=",
),
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("d"),
"-",
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("a"),
"*",
sir_utils.make_field_access_expr("d", [0, 0, -1]),
),
),
"*",
sir_utils.make_var_access_expr("m"),
),
"=",
),
])
vertical_region_stmt_2 = sir_utils.make_vertical_region_decl_stmt(
body_ast_2, interval_2, SIR.VerticalRegion.Forward)
# ---- Third vertical region statement ----
interval_3 = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End,
0, -1)
body_ast_3 = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_field_access_expr("c"),
"*",
sir_utils.make_field_access_expr("d", [0, 0, 1]),
),
"-=",
)
])
vertical_region_stmt_3 = sir_utils.make_vertical_region_decl_stmt(
body_ast_3, interval_3, SIR.VerticalRegion.Backward)
sir = sir_utils.make_sir(
OUTPUT_FILE,
SIR.GridType.Value("Cartesian"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([
vertical_region_stmt_1, vertical_region_stmt_2,
vertical_region_stmt_3
]),
[
sir_utils.make_field(
"a", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"b", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"c", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"d", sir_utils.make_field_dimensions_cartesian()),
],
)
],
)
# print the SIR
if args.verbose:
sir_utils.pprint(sir)
# compile
code = dawn4py.compile(sir, backend="cuda")
# write to file
print(f"Writing generated code to '{OUTPUT_PATH}'")
with open(OUTPUT_PATH, "w") as f:
f.write(code)
def main(args: argparse.Namespace):
stencil_name = "ICON_laplacian_stencil"
gen_outputfile = f"{stencil_name}.cpp"
sir_outputfile = f"{stencil_name}.sir"
interval = serial_utils.make_interval(AST.Interval.Start, AST.Interval.End,
0, 0)
body_ast = serial_utils.make_ast([
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("rot_vec"),
serial_utils.make_reduction_over_neighbor_expr(
op="+",
init=serial_utils.make_literal_access_expr(
"0.0", AST.BuiltinType.Double),
rhs=serial_utils.make_binary_operator(
serial_utils.make_field_access_expr("vec", [True, 0]),
"*",
serial_utils.make_field_access_expr("geofac_rot"),
),
chain=[
AST.LocationType.Value("Vertex"),
AST.LocationType.Value("Edge")
],
),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("div_vec"),
serial_utils.make_reduction_over_neighbor_expr(
op="+",
init=serial_utils.make_literal_access_expr(
"0.0", AST.BuiltinType.Double),
rhs=serial_utils.make_binary_operator(
serial_utils.make_field_access_expr("vec", [True, 0]),
"*",
serial_utils.make_field_access_expr("geofac_div"),
),
chain=[
AST.LocationType.Value("Cell"),
AST.LocationType.Value("Edge")
],
),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("nabla2t1_vec"),
serial_utils.make_reduction_over_neighbor_expr(
op="+",
init=serial_utils.make_literal_access_expr(
"0.0", AST.BuiltinType.Double),
rhs=serial_utils.make_field_access_expr("rot_vec", [True, 0]),
chain=[
AST.LocationType.Value("Edge"),
AST.LocationType.Value("Vertex")
],
weights=[
serial_utils.make_literal_access_expr(
"-1.0", AST.BuiltinType.Double),
serial_utils.make_literal_access_expr(
"1.0", AST.BuiltinType.Double)
]),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("nabla2t1_vec"),
serial_utils.make_binary_operator(
serial_utils.make_binary_operator(
serial_utils.make_field_access_expr("tangent_orientation"),
"*",
serial_utils.make_field_access_expr("nabla2t1_vec"),
),
"/",
serial_utils.make_field_access_expr("primal_edge_length"),
),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("nabla2t2_vec"),
serial_utils.make_reduction_over_neighbor_expr(
op="+",
init=serial_utils.make_literal_access_expr(
"0.0", AST.BuiltinType.Double),
rhs=serial_utils.make_field_access_expr("div_vec", [True, 0]),
chain=[
AST.LocationType.Value("Edge"),
AST.LocationType.Value("Cell")
],
weights=[
serial_utils.make_literal_access_expr(
"-1.0", AST.BuiltinType.Double),
serial_utils.make_literal_access_expr(
"1.0", AST.BuiltinType.Double)
]),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("nabla2t2_vec"),
serial_utils.make_binary_operator(
serial_utils.make_field_access_expr("nabla2t2_vec"),
"/",
serial_utils.make_field_access_expr("dual_edge_length"),
),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_field_access_expr("nabla2_vec"),
serial_utils.make_binary_operator(
serial_utils.make_field_access_expr("nabla2t2_vec"),
"-",
serial_utils.make_field_access_expr("nabla2t1_vec"),
),
"=",
),
])
vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt(
body_ast, interval, AST.VerticalRegion.Forward)
sir = serial_utils.make_sir(
gen_outputfile,
AST.GridType.Value("Unstructured"),
[
serial_utils.make_stencil(
stencil_name,
serial_utils.make_ast([vertical_region_stmt]),
[
serial_utils.make_field(
"vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"div_vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Cell")], 1),
),
serial_utils.make_field(
"rot_vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Vertex")], 1),
),
serial_utils.make_field(
"nabla2t1_vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
is_temporary=True),
serial_utils.make_field(
"nabla2t2_vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
is_temporary=True),
serial_utils.make_field(
"nabla2_vec",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"primal_edge_length",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"dual_edge_length",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"tangent_orientation",
serial_utils.make_field_dimensions_unstructured(
[AST.LocationType.Value("Edge")], 1),
),
serial_utils.make_field(
"geofac_rot",
serial_utils.make_field_dimensions_unstructured([
AST.LocationType.Value("Vertex"),
AST.LocationType.Value("Edge")
], 1),
),
serial_utils.make_field(
"geofac_div",
serial_utils.make_field_dimensions_unstructured([
AST.LocationType.Value("Cell"),
AST.LocationType.Value("Edge")
], 1),
),
],
),
],
)
# print the SIR
if args.verbose:
print(MessageToJson(sir))
# compile
code = dawn4py.compile(sir,
groups=[],
backend=dawn4py.CodeGenBackend.CXXNaiveIco)
# write to file
print(f"Writing generated code to '{gen_outputfile}'")
with open(gen_outputfile, "w") as f:
f.write(code)