def make_horizontal_loop_with_init(field: Str): write_access = nir.FieldAccess( name=field, primary=no_extent, location_type=default_location, ) return ( nir.HorizontalLoop( stmt=nir.BlockStmt( declarations=[], statements=[ nir.AssignStmt( left=write_access, right=nir.Literal( value=common.BuiltInLiteral.ONE, vtype=default_vtype, location_type=default_location, ), ) ], ), location_type=default_location, ), write_access, )
def visit_NeighborReduce(self, node: gtir.NeighborReduce, *, last_block, **kwargs): loc_comprehension = copy.deepcopy(kwargs["location_comprehensions"]) assert node.neighbors.name not in loc_comprehension loc_comprehension[node.neighbors.name] = node.neighbors kwargs["location_comprehensions"] = loc_comprehension body_location = node.neighbors.chain.elements[-1] reduce_var_name = "local" + str(node.id_) last_block.declarations.append( nir.LocalVar( name=reduce_var_name, vtype=common.DataType.FLOAT64, # TODO location_type=node.location_type, )) last_block.statements.append( nir.AssignStmt( left=nir.VarAccess(name=reduce_var_name, location_type=node.location_type), right=nir.Literal( value=self.REDUCE_OP_INIT_VAL[node.op], location_type=node.location_type, vtype=common.DataType.FLOAT64, # TODO ), location_type=node.location_type, ), ) body = nir.BlockStmt( declarations=[], statements=[ nir.AssignStmt( left=nir.VarAccess(name=reduce_var_name, location_type=body_location), right=nir.BinaryOp( left=nir.VarAccess(name=reduce_var_name, location_type=body_location), op=self.REDUCE_OP_TO_BINOP[node.op], right=self.visit(node.operand, in_neighbor_loop=True, **kwargs), location_type=body_location, ), location_type=body_location, ) ], location_type=body_location, ) last_block.statements.append( nir.NeighborLoop( neighbors=self.visit(node.neighbors.chain), body=body, location_type=node.location_type, )) return nir.VarAccess(name=reduce_var_name, location_type=node.location_type) # TODO
def make_init(field: Str): write_access = nir.FieldAccess(name=field, primary=no_extent, location_type=default_location) return ( nir.AssignStmt( left=write_access, right=nir.Literal( value=common.BuiltInLiteral.ONE, vtype=default_vtype, location_type=default_location, ), ), write_access, )
def visit_Literal(self, node: gtir.Literal, **kwargs): return nir.Literal(value=node.value, vtype=node.vtype, location_type=node.location_type)