def generate_implementation(self):
sources = gt_text.TextBlock(
indent_size=gt_backend.BaseModuleGenerator.TEMPLATE_INDENT_SIZE
)
args = []
for arg in self.implementation_ir.api_signature:
if arg.name not in self.implementation_ir.unreferenced:
args.append(arg.name)
if arg.name in self.implementation_ir.fields:
args.append("list(_origin_['{}'])".format(arg.name))
source = """
# Load or generate a GTComputation object for the current domain size
pyext_module.run_computation(list(_domain_), {run_args}, exec_info)
""".format(
run_args=", ".join(args)
)
if self.backend_name == "gtcuda":
source = (
source
+ """import cupy
cupy.cuda.Device(0).synchronize()
"""
)
sources.extend(source.splitlines())
return sources.text
def visit_While(self, node: gt_ir.While):
body_sources = gt_text.TextBlock()
body_sources.append("while {condition}:".format(condition=self.visit(node.condition)))
body_sources.indent()
for stmt in node.body.stmts:
body_sources.extend(self.visit(stmt))
return body_sources.text
def visit_While(self, node: gt_ir.While) -> gt_text.TextBlock:
body_sources = gt_text.TextBlock()
body_sources.append("while ({condition}) {{".format(condition=self.visit(node.condition)))
for stmt in node.body.stmts:
body_sources.extend(self.visit(stmt))
body_sources.append("}")
return body_sources
def generate_implementation(self) -> str:
block = gt_text.TextBlock(indent_size=self.TEMPLATE_INDENT_SIZE)
numpy_ir = NumpyIR.apply(self.builder.implementation_ir)
self.source_generator(numpy_ir, block)
if self.builder.options.backend_opts.get("ignore_np_errstate", True):
source = "with np.errstate(divide='ignore', over='ignore', under='ignore', invalid='ignore'):\n"
source += textwrap.indent(block.text, " " * self.TEMPLATE_INDENT_SIZE)
else:
source = block.text
return source
def visit_If(self, node: gt_ir.If) -> gt_text.TextBlock:
body_sources = gt_text.TextBlock()
body_sources.append("if ({condition}) {{".format(condition=self.visit(node.condition)))
for stmt in node.main_body.stmts:
body_sources.extend(self.visit(stmt))
if node.else_body:
body_sources.append("} else {")
for stmt in node.else_body.stmts:
body_sources.extend(self.visit(stmt))
body_sources.append("}")
return body_sources
def visit_If(self, node: gt_ir.If):
body_sources = gt_text.TextBlock()
body_sources.append("if {condition}:".format(condition=self.visit(node.condition)))
body_sources.indent()
for stmt in node.main_body.stmts:
body_sources.extend(self.visit(stmt))
body_sources.dedent()
if node.else_body:
body_sources.append("else:")
body_sources.indent()
for stmt in node.else_body.stmts:
body_sources.extend(self.visit(stmt))
body_sources.dedent()
return ["".join([str(item) for item in line]) for line in body_sources.lines]
def visit_ApplyBlock(
self, node: gt_ir.ApplyBlock
) -> Tuple[Tuple[Tuple[int, int], Tuple[int, int]], str]:
interval_definition = self.visit(node.interval)
body_sources = gt_text.TextBlock()
self.declared_symbols = set()
for name, var_decl in node.local_symbols.items():
assert isinstance(var_decl, gt_ir.VarDecl)
body_sources.append(self._make_cpp_variable(var_decl))
self.declared_symbols.add(name)
self.apply_block_symbols = {**self.stage_symbols, **node.local_symbols}
body_sources.extend(self.visit(node.body))
return interval_definition, body_sources.text
def visit_HorizontalIf(self, node: gt_ir.HorizontalIf) -> List[str]:
source = gt_text.TextBlock()
conditions = [
self._create_horizontal_conditional(axis_name, index, node.intervals[axis_name])
for index, axis_name in enumerate(
(axis.name for axis in gt_ir.Domain.LatLonGrid().parallel_axes)
)
]
source.append("if " + " and ".join([cond for cond in conditions if cond != ""]) + ":")
source.indent()
for stmt in node.body.stmts:
source.extend(self.visit(stmt))
return ["".join([str(item) for item in line]) for line in source.lines]
def generate_implementation(self) -> str:
sources = gt_text.TextBlock(
indent_size=gt_backend.BaseModuleGenerator.TEMPLATE_INDENT_SIZE)
args = []
empty_checks = []
none_checks = []
# Other backends allow None to be passed to unreferenced fields or scalars.
# Dawn backends do not check for null values passed to parameters, so the
# None values must be replaced with defaults for the unreferenced parameters,
# (0,0,0) for the origins, np.empty for fields, and zero for scalars.
field_info = self.args_data["field_info"]
unreferenced = self.args_data["unreferenced"]
for arg in self.builder.definition_ir.api_signature:
args.append(arg.name)
if arg.name in field_info:
if arg.name in unreferenced:
ndims = len(field_info[arg.name].boundary)
zeros = ", ".join(["0"] * ndims)
args.append("[{}]".format(zeros))
empty_checks.append(
f"{arg.name} = np.empty(({zeros})) if {arg.name} is None else {arg.name}"
)
else:
args.append("list(_origin_['{}'])".format(arg.name))
elif arg.name in self.builder.definition_ir.parameters and arg.default is None:
none_checks.append(
f"{arg.name} = 0 if {arg.name} is None else {arg.name}")
source = """
# Load or generate a Dawn Computation object for the current domain size
{empty_checks}
{none_checks}
pyext_module.run_computation(list(_domain_), {run_args}, exec_info)
""".format(
empty_checks="\n".join(empty_checks),
none_checks="\n".join(none_checks),
run_args=", ".join(args),
)
sources.extend(source.splitlines())
return sources.text
def visit_BlockStmt(self, node: gt_ir.BlockStmt) -> str:
body_sources = gt_text.TextBlock()
for stmt in node.stmts:
body_sources.extend(self.visit(stmt))
return body_sources.text
def generate_implementation(self):
sources = gt_text.TextBlock(indent_size=self.TEMPLATE_INDENT_SIZE)
self.source_generator(self.builder.implementation_ir, sources)
return sources.text
