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 = 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_literal_access_expr(
"1.0", AST.BuiltinType.Float),
serial_utils.make_unstructured_field_access_expr("in"),
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,
serial_utils.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('Edge')], 1))
],
)
],
)
return sir
def create_vertical_region_stmt3():
""" create a vertical region statement for the stencil
"""
interval = sir_utils.make_interval(sir_utils.Interval.Start,
sir_utils.Interval.End, 0, -1)
body_ast = sir_utils.make_ast([
sir_utils.make_assignment_stmt(
sir_utils.make_unstructured_field_access_expr("d"),
sir_utils.make_binary_operator(
sir_utils.make_unstructured_field_access_expr("c"), "*",
sir_utils.make_unstructured_field_access_expr(
"d", sir_utils.make_unstructured_offset(False), 1)), "-=")
])
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
body_ast, interval, SIR.VerticalRegion.Backward)
return vertical_region_stmt
def create_vertical_region_stmt1():
""" create a vertical region statement for the stencil
"""
interval = serial_utils.make_interval(serial_utils.Interval.Start,
serial_utils.Interval.Start, 0, 0)
body_ast = serial_utils.make_ast([
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("c"),
serial_utils.make_binary_operator(
serial_utils.make_unstructured_field_access_expr("c"),
"/",
serial_utils.make_unstructured_field_access_expr("b"),
),
"=",
),
serial_utils.make_assignment_stmt(
serial_utils.make_unstructured_field_access_expr("d"),
serial_utils.make_binary_operator(
serial_utils.make_unstructured_field_access_expr("d"),
"/",
serial_utils.make_unstructured_field_access_expr("b"),
),
"=",
),
])
vertical_region_stmt = serial_utils.make_vertical_region_decl_stmt(
body_ast, interval, AST.VerticalRegion.Forward)
return vertical_region_stmt
def 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 field_access_expr(self, field: DuskField, index: expr = None):
if not self.ctx.location.in_vertical_region:
raise DuskSyntaxError(
f"Invalid field access {name} outside of a vertical region!")
return make_unstructured_field_access_expr(
field.sir.name, *self.field_index(index, field=field))
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 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)