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 main(args: argparse.Namespace):
interval = sir_utils.make_interval(SIR.Interval.Start, SIR.Interval.End, 0,
0)
body_ast = sir_utils.make_ast([
sir_utils.make_loop_stmt([
sir_utils.make_assignment_stmt(
sir_utils.make_field_access_expr("out"),
sir_utils.make_field_access_expr("in"),
"=",
)
], [SIR.LocationType.Value("Cell"),
SIR.LocationType.Value("Edge")])
])
vertical_region_stmt = sir_utils.make_vertical_region_decl_stmt(
body_ast, interval, SIR.VerticalRegion.Forward)
sir = sir_utils.make_sir(
OUTPUT_FILE,
SIR.GridType.Value("Unstructured"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([vertical_region_stmt]),
[
sir_utils.make_field(
"out",
sir_utils.make_field_dimensions_unstructured([
SIR.LocationType.Value("Cell"),
SIR.LocationType.Value("Edge")
], 1),
),
sir_utils.make_field(
"in",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Edge")], 1),
),
],
)
],
)
# print the SIR
if args.verbose:
sir_utils.pprint(sir)
# compile
code = dawn4py.compile(sir, backend=dawn4py.CodeGenBackend.CXXNaiveIco)
# write to file
print(f"Writing generated code to '{OUTPUT_PATH}'")
with open(OUTPUT_PATH, "w") as f:
f.write(code)
def main(args: argparse.Namespace):
sir = sir_utils.make_sir(
OUTPUT_FILE,
SIR.GridType.Value("Unstructured"),
[
sir_utils.make_stencil(
OUTPUT_NAME,
sir_utils.make_ast([
create_vertical_region_stmt1(),
create_vertical_region_stmt2(),
create_vertical_region_stmt3(),
]),
[
sir_utils.make_field(
"a",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"b",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"c",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
sir_utils.make_field(
"d",
sir_utils.make_field_dimensions_unstructured(
[SIR.LocationType.Value("Cell")], 1),
),
],
)
],
)
# print the SIR
if args.verbose:
sir_utils.pprint(sir)
# compile
code = dawn4py.compile(sir, backend="c++-naive-ico")
# write to file
print(f"Writing generated code to '{OUTPUT_PATH}'")
with open(OUTPUT_PATH, "w") as f:
f.write(code)
def main(args: argparse.Namespace):
interval = 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", [0, 0, 0]),
sir_utils.make_field_access_expr("in", [1, 0, 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()),
],
)
],
)
# 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 test_sir_serialization(name):
sir = getattr(utils, f"make_{name}_sir")(name=name)
serialized_bytes = dawn4py.serialization.to_bytes(sir)
assert serialized_bytes is not None
sir_from_bytes = dawn4py.serialization.from_bytes(serialized_bytes, SIR.SIR)
assert sir == sir_from_bytes
serialized_json = dawn4py.serialization.to_json(sir)
assert serialized_json is not None
sir_from_json = dawn4py.serialization.from_json(serialized_json, SIR.SIR)
assert sir == sir_from_json
bytes_strio = io.StringIO()
sir_utils.pprint(sir_from_bytes, file=bytes_strio)
json_strio = io.StringIO()
sir_utils.pprint(sir_from_json, file=json_strio)
assert bytes_strio.getvalue() == json_strio.getvalue()
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):
sir = sir_utils.make_sir(
OUTPUT_FILE,
SIR.GridType.Value("Cartesian"),
[
sir_utils.make_stencil(
"global_indexing",
sir_utils.make_ast([
create_vertical_region_stmt(),
create_boundary_correction_region(
value="4",
i_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
),
create_boundary_correction_region(
value="8",
i_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
),
create_boundary_correction_region(
value="6",
j_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
),
create_boundary_correction_region(
value="2",
j_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
),
create_boundary_correction_region(
value="1",
j_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
i_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
),
create_boundary_correction_region(
value="3",
j_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
i_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
),
create_boundary_correction_region(
value="7",
j_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
i_interval=sir_utils.make_interval(
SIR.Interval.Start, SIR.Interval.Start, 0, 1),
),
create_boundary_correction_region(
value="5",
j_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
i_interval=sir_utils.make_interval(
SIR.Interval.End, SIR.Interval.End, -1, 0),
),
]),
[
sir_utils.make_field(
"in", sir_utils.make_field_dimensions_cartesian()),
sir_utils.make_field(
"out", sir_utils.make_field_dimensions_cartesian())
],
)
],
)
# print the SIR
if args.verbose:
sir_utils.pprint(sir)
# compile
code = dawn4py.compile(sir, backend="c++-naive")
# write to file
print(f"Writing generated code to '{OUTPUT_PATH}'")
with open(OUTPUT_PATH, "w") as f:
f.write(code)
def main(args: argparse.Namespace):
interval = serial_utils.make_interval(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)
