def visit_BlockStmt(self, node: Node, **kwargs):
if self.isControlFlow:
for s in node.statements:
assert isinstance(s, sir.VerticalRegionDeclStmt)
return self.visit(s)
else:
horizontal_loops = []
declarations = []
for s in node.statements:
if isinstance(s, sir.VarDeclStmt):
# TODO this doesn't work: if we move the declaration out of the horizontal loop, we need to promote it to a field
[vardecl, initexpr] = self.visit(s)
declarations.append(vardecl)
transformed_stmt = naive.ExprStmt(
expr=naive.AssignmentExpr(
left=naive.FieldAccessExpr(
name=vardecl.name,
offset=(False, 0),
location_type=initexpr.location_type,
is_sparse=False,
),
right=initexpr,
))
else:
transformed_stmt = self.visit(s)
horizontal_loops.append(
naive.HorizontalLoop(ast=naive.BlockStmt(
statements=[transformed_stmt])))
return [declarations, horizontal_loops]
# CHECK: ReduceOverNeighbourExpr.sparse_access
# CHECK-NEXT: (m_sparse_dimension_idx = 0,
# CHECK-NEXT: reduceEdgeToVertex(mesh, iter, EXPR, [&](auto &lhs, auto const &redIdx) {
# CHECK-NEXT: lhs += field({{.*}}iter{{.*}} m_sparse_dimension_idx{{.*}},{{.*}}k);
# CHECK-NEXT: m_sparse_dimension_idx++;
# CHECK-NEXT: return lhs;
# CHECK-NEXT: }))
make_test(
naive.ReduceOverNeighbourExpr(
operation=common.BinaryOperator.ADD,
right=naive.FieldAccessExpr(name="field",
offset=[True, 0],
location_type=naive.LocationType.Edge,
is_sparse=True),
init=naive.Expr(location_type=naive.LocationType.Edge),
location_type=naive.LocationType.Node,
),
"sparse_access",
iter_var="iter",
)
# CHECK: ExprStmt
# CHECK-NEXT:
# CHECK-NEXT: EXPR = EXPR;
make_test(
naive.ExprStmt(expr=naive.AssignmentExpr(
left=naive.Expr(location_type=naive.LocationType.Edge),
right=naive.Expr(location_type=naive.LocationType.Edge),
)))
def visit_AssignmentExpr(self, node: Node, **kwargs):
assert node.op == "="
return naive.AssignmentExpr(left=self.visit(node.left),
right=self.visit(node.right))
location_type=naive.LocationType.Node)
in_acc = naive.FieldAccessExpr(name="in",
offset=[True, 0],
location_type=naive.LocationType.Edge)
red = naive.ReduceOverNeighbourExpr(
operation=common.BinaryOperator.ADD,
right=in_acc,
init=zero,
location_type=naive.LocationType.Node,
)
out_acc = naive.FieldAccessExpr(name="out",
offset=[False, 0],
location_type=naive.LocationType.Node)
assign = naive.ExprStmt(expr=naive.AssignmentExpr(left=out_acc, right=red))
hori = naive.HorizontalLoop(ast=naive.BlockStmt(statements=[assign]), )
vert = naive.ForK(horizontal_loops=[hori], loop_order=common.LoopOrder.FORWARD)
sten = naive.Stencil(name="reduce_to_node", k_loops=[vert])
comp = naive.Computation(params=[field_in, field_out], stencils=[sten])
# ---------------------------
generated_code = NaiveCodeGenerator.apply(comp)
print(generated_code)
# try compile the generated code
# p = Popen(
# [