def test_cyclic_assignment(self): stencil = make_stencil( fields=[], statements=[ make_var_decl(name="var"), make_var_decl(name="var2"), make_assign_to_local_var("var", make_var_acc("var2")), make_assign_to_local_var("var2", make_var_acc("var")), ], ) with pytest.raises(AnalysisException): InferLocalVariableLocationTypeTransformation.apply(stencil)
def test_incompatible_location(self): stencil = make_stencil( fields=[make_field("field_edge", sir.LocationType.Edge), make_field("field_cell")], statements=[ make_var_decl(name="var_edge"), make_var_decl(name="var_cell"), make_assign_to_local_var("var_edge", make_field_acc("field_edge")), make_assign_to_local_var("var_cell", make_field_acc("field_cell")), make_assign_to_local_var("var_cell", make_var_acc("var_edge")), ], ) with pytest.raises(AnalysisException): InferLocalVariableLocationTypeTransformation.apply(stencil)
def test_simple_assignment(self): stencil = make_stencil( fields=[make_field("field")], statements=[ make_var_decl(name="var"), make_assign_to_local_var("var", make_field_acc("field")), ], ) result = InferLocalVariableLocationTypeTransformation.apply(stencil) vardecl = FindNodes.by_type(sir.VarDeclStmt, result)[0] assert vardecl.location_type == sir.LocationType.Cell
def test_simple_assignment(self): stencil = make_stencil( fields=[make_field("field")], statements=[ make_var_decl(name="var"), make_assign_to_local_var("var", make_field_acc("field")), ], ) result = InferLocalVariableLocationTypeTransformation.apply(stencil) vardecl = result.iter_tree().if_isinstance( sir.VarDeclStmt).to_list()[0] assert vardecl.location_type == sir.LocationType.Cell
def test_chain_assignment(self): stencil = make_stencil( fields=[make_field("field")], statements=[ make_var_decl(name="var"), make_assign_to_local_var("var", make_field_acc("field")), make_var_decl(name="another_var", dtype=float_type, init=make_var_acc("var")), ], ) result = InferLocalVariableLocationTypeTransformation.apply(stencil) vardecls = FindNodes.by_type(sir.VarDeclStmt, result) assert len(vardecls) == 2 for vardecl in vardecls: assert vardecl.location_type == sir.LocationType.Cell
def test_reduction(self): stencil = make_stencil( fields=[], statements=[ make_var_decl(name="var"), make_assign_to_local_var( "var", sir.ReductionOverNeighborExpr( op="+", rhs=make_literal(), init=make_literal(), chain=[sir.LocationType.Edge, sir.LocationType.Cell], ), ), ], ) result = InferLocalVariableLocationTypeTransformation.apply(stencil) vardecl = FindNodes.by_type(sir.VarDeclStmt, result)[0] assert vardecl.location_type == sir.LocationType.Edge
def test_var_type_not_deducible(self): stencil = make_stencil(fields=[], statements=[make_var_decl(name="var")]) with pytest.raises(AnalysisException): InferLocalVariableLocationTypeTransformation.apply(stencil)
# ), # ] # ) assign_pnabla_MYY_vol = sir.ExprStmt(expr=sir.AssignmentExpr( left=sir.FieldAccessExpr(name="pnabla_MYY", vertical_offset=0, horizontal_offset=sir.ZeroOffset()), op="=", right=sir.BinaryOperator( left=sir.FieldAccessExpr(name="pnabla_MYY", vertical_offset=0, horizontal_offset=sir.ZeroOffset()), op="/", right=sir.FieldAccessExpr( name="vol", vertical_offset=0, horizontal_offset=sir.ZeroOffset()), ), )) statements.append(assign_pnabla_MYY_vol) block = sir.BlockStmt(statements=statements) ast = sir.AST(root=sir.BlockStmt(statements=statements)) vert_decl_stmt = sir.VerticalRegionDeclStmt(vertical_region=sir.VerticalRegion( ast=ast, interval=sir.Interval(), loop_order=common.LoopOrder.FORWARD)) ctrl_flow_ast = sir.AST(root=sir.BlockStmt(statements=[vert_decl_stmt])) stencil = sir.Stencil(name="nabla", ast=ctrl_flow_ast, params=fields) var_loc_type_inferred = InferLocalVariableLocationTypeTransformation.apply( stencil) naive_ir = sir_to_naive.SirToNaive().visit(var_loc_type_inferred) print(naive_codegen.NaiveCodeGenerator.apply(naive_ir))