def expression(self, expr: expr):
return make_expr(
dispatch(
{
Constant: self.constant,
Name: self.var,
Subscript: self.subscript,
UnaryOp: self.unop,
BinOp: self.binop,
BoolOp: self.boolop,
Compare: self.compare,
IfExp: self.ifexp,
Call: self.funcall,
},
expr,
))
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 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):
# ---- 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 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