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( # [