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)