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_VarDecl(self, node: VarDecl) -> gtir.ScalarDecl:
# datatype conversion works via same ID
return gtir.ScalarDecl(
name=node.name,
dtype=common.DataType(int(node.data_type.value)),
loc=location_to_source_location(node.loc),
)
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_VarDecl(self, node: VarDecl):
# datatype conversion works via same ID
return gtir.ScalarDecl(name=node.name, dtype=common.DataType(int(node.data_type.value)))
def visit_Cast(self, node: Cast) -> gtir.Cast:
return gtir.Cast(dtype=common.DataType(node.data_type.value), expr=self.visit(node.expr))
def visit_ScalarLiteral(self, node: ScalarLiteral) -> gtir.Literal:
return gtir.Literal(value=str(node.value), dtype=common.DataType(node.data_type.value))
def visit_Cast(self, node: Cast) -> gtir.Cast:
return gtir.Cast(
dtype=common.DataType(node.data_type.value),
expr=self.visit(node.expr),
loc=location_to_source_location(node.loc),
)
