def visit_ComputationBlock(self, node: ComputationBlock) -> gtir.VerticalLoop:
stmts = []
temporaries = []
for s in node.body.stmts:
# FieldDecl or VarDecls in the body are temporaries
if isinstance(s, FieldDecl) or isinstance(s, VarDecl):
dtype = common.DataType(int(s.data_type.value))
if dtype == common.DataType.DEFAULT:
# TODO this will be a frontend choice later
# in non-GTC parts, this is set in the backend
dtype = cast(
common.DataType, common.DataType.FLOAT64
) # see https://github.com/GridTools/gtc/issues/100
temporaries.append(
gtir.FieldDecl(
name=s.name,
dtype=dtype,
dimensions=(True, True, True),
)
)
else:
stmts.append(self.visit(s))
start, end = self.visit(node.interval)
interval = gtir.Interval(start=start, end=end)
return gtir.VerticalLoop(
interval=interval,
loop_order=self.GT4PY_ITERATIONORDER_TO_GTIR_LOOPORDER[node.iteration_order],
body=stmts,
temporaries=temporaries,
)
def visit_FieldDecl(self, node: FieldDecl):
dimension_names = ["I", "J", "K"]
dimensions = [dim in node.axes for dim in dimension_names]
# datatype conversion works via same ID
return gtir.FieldDecl(name=node.name,
dtype=common.DataType(int(node.data_type.value)),
dimensions=dimensions)
def visit_FieldDecl(
self, node: gtir.FieldDecl, new_symbols: Dict[str, Any], **kwargs: Any
) -> gtir.FieldDecl:
if node.dtype == DataType.AUTO:
dtype = new_symbols[node.name].dtype
return gtir.FieldDecl(name=node.name, dtype=dtype, dimensions=node.dimensions)
else:
return node
def visit_FieldDecl(self, node: FieldDecl) -> gtir.FieldDecl:
dimension_names = ["I", "J", "K"]
dimensions = [dim in node.axes for dim in dimension_names]
# datatype conversion works via same ID
return gtir.FieldDecl(
name=node.name,
dtype=common.DataType(int(node.data_type.value)),
dimensions=dimensions,
data_dims=node.data_dims,
loc=location_to_source_location(node.loc),
)
def visit_FieldDecl(self, node: FieldDecl):
# datatype conversion works via same ID
return gtir.FieldDecl(name=node.name, dtype=common.DataType(int(node.data_type.value)))