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]
def visit_FieldAccessExpr(self, node: Node, **kwargs):
horizontal_offset = False # TODO
return naive.FieldAccessExpr(
name=node.name,
offset=(horizontal_offset, node.vertical_offset),
location_type=self._get_field_location_type(self.sir_stencil_params[node.name]),
is_sparse=self._is_sparse_field(self.sir_stencil_params[node.name]),
)
def visit_VarAccessExpr(self, node: Node, **kwargs):
loctype = ""
if node.location_type:
loctype = self.SIR_TO_NAIVE_LOCATION_TYPE[node.location_type]
elif self.current_loc_type_stack:
loctype = self.current_loc_type_stack[-1]
else:
raise ValueError("no location type")
return naive.FieldAccessExpr(
name=node.name, offset=(False, 0), location_type=loctype, is_sparse=False,
)
naive.UnstructuredField(name="in",
location_type=naive.LocationType.Edge,
data_type=common.DataType.FLOAT64), )
# CHECK: LiteralExpr
# CHECK-NEXT: (double)0.0
make_test(
naive.LiteralExpr(value="0.0",
data_type=common.DataType.FLOAT64,
location_type=naive.LocationType.Node))
# CHECK: FieldAccessExpr
# CHECK-NEXT: field(deref(LibTag{}, iter), k)
make_test(
naive.FieldAccessExpr(name="field",
offset=[True, 0],
location_type=naive.LocationType.Edge),
iter_var="iter",
)
# CHECK: ReduceOverNeighbourExpr.dense_access
# CHECK-NEXT: (m_sparse_dimension_idx = 0,
# CHECK-NEXT: reduceEdgeToVertex(mesh, iter, EXPR, [&](auto &lhs, auto const &redIdx) {
# CHECK-NEXT: lhs += field({{.*}}redIdx{{.*}} 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",
# Op::plus, b.at(in_f, HOffsetType::withOffset, 0),
# b.lit(0.), LocType::Cells, LocType::Edges))))))));
# ```
field_in = naive.UnstructuredField(name="in",
location_type=naive.LocationType.Edge,
data_type=common.DataType.FLOAT64)
field_out = naive.UnstructuredField(name="out",
location_type=naive.LocationType.Node,
data_type=common.DataType.FLOAT64)
zero = naive.LiteralExpr(value="0.0",
data_type=common.DataType.FLOAT64,
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]), )