def apply(cls, gtir, sdfg: dace.SDFG):
self = cls()
code_objects = sdfg.generate_code()
computations = code_objects[[co.title for co in code_objects
].index("Frame")].clean_code
lines = computations.split("\n")
computations = "\n".join(
lines[0:2] + lines[3:]) # remove import of not generated file
computations = codegen.format_source("cpp", computations, style="LLVM")
interface = cls.template.definition.render(
name=sdfg.name,
dace_args=self.generate_dace_args(gtir, sdfg),
functor_args=self.generate_functor_args(sdfg),
tmp_allocs=self.generate_tmp_allocs(sdfg),
)
generated_code = f"""#include <gridtools/sid/sid_shift_origin.hpp>
#include <gridtools/sid/allocator.hpp>
#include <gridtools/stencil/cartesian.hpp>
namespace gt = gridtools;
{computations}
{interface}
"""
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def apply(cls, root, **kwargs) -> str:
assert "stencil_code" in kwargs
generated_code = cls().visit(root, stencil_code=kwargs["stencil_code"])
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def _postprocess_dace_code(code_objects, is_gpu):
lines = code_objects[[co.title for co in code_objects
].index("Frame")].clean_code.split("\n")
if is_gpu:
regex = re.compile("struct [a-zA-Z_][a-zA-Z0-9_]*_t {")
for i, line in enumerate(lines):
if regex.match(line.strip()):
j = i + 1
while "};" not in lines[j].strip():
j += 1
lines = lines[0:i] + lines[j + 1:]
break
for i, line in enumerate(lines):
if "#include <dace/dace.h>" in line:
cuda_code = [
co.clean_code for co in code_objects
if co.title == "CUDA"
][0]
lines = lines[0:i] + cuda_code.split("\n") + lines[i + 1:]
break
def keep_line(line):
line = line.strip()
if line == '#include "../../include/hash.h"':
return False
if line.startswith("DACE_EXPORTED") and line.endswith(");"):
return False
if line == "#include <cuda_runtime.h>":
return False
return True
lines = filter(keep_line, lines)
return codegen.format_source("cpp", "\n".join(lines), style="LLVM")
def apply(cls, node: oir.HorizontalExecution, **kwargs: Any) -> str:
preprocessed_node = cls.RemoveCastInIndexVisitor().visit(node)
if not isinstance(node, oir.HorizontalExecution):
raise ValueError("apply() requires oir.HorizontalExecution node")
generated_code = super().apply(preprocessed_node)
formatted_code = codegen.format_source("python", generated_code)
return formatted_code
def apply(cls, root, **kwargs) -> str:
symbol_tbl_resolved = SymbolTblHelper().visit(root)
generated_code = super().apply(symbol_tbl_resolved, **kwargs)
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def apply(cls, root: LeafNode, **kwargs: Any) -> str:
if not isinstance(root, gtcpp.Program):
raise ValueError("apply() requires gtcpp.Progam root node")
if "gt_backend_t" not in kwargs:
raise TypeError("apply() missing 1 required keyword-only argument: 'gt_backend_t'")
generated_code = super().apply(root, offset_limit=_offset_limit(root), **kwargs)
formatted_code = codegen.format_source("cpp", generated_code, style="LLVM")
return formatted_code
def apply(cls, root, *, module_name="stencil", **kwargs) -> str:
generated_code = cls().visit(root, module_name=module_name, **kwargs)
if kwargs.get("format_source", True):
generated_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return generated_code
def apply(cls, root: LeafNode, **kwargs: Any) -> str:
if not isinstance(root, cuir.Program):
raise ValueError("apply() requires gtcpp.Progam root node")
generated_code = super().apply(root, **kwargs)
if kwargs.get("format_source", True):
generated_code = codegen.format_source("cpp", generated_code, style="LLVM")
return generated_code
def apply(cls, root, *, module_name="stencil", backend, **kwargs) -> str:
generated_code = cls(backend).visit(root,
module_name=module_name,
**kwargs)
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def apply(cls, gtir: gtir.Stencil, sdfg: dace.SDFG, module_name: str, *,
backend) -> str:
generated_code = cls(backend).generate_sdfg_bindings(
gtir, sdfg, module_name=module_name)
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def make_test(node, test_name="", **kwargs):
print("{test_name}\n{generated_code}".format(
test_name=".".join(item
for item in [node.__class__.__name__, test_name]
if item),
generated_code=codegen.format_source("cpp",
NaiveCodeGenerator().visit(
node, **kwargs),
style="LLVM"),
))
def generate_computation(self) -> Dict[str, Union[str, Dict]]:
computation_name = (self.builder.caching.module_prefix +
"computation" +
self.builder.caching.module_postfix + ".py")
source = NpirCodegen.apply(self.npir)
if self.builder.options.format_source:
source = format_source("python", source)
return {computation_name: source}
def generate_computation(self) -> Dict[str, Union[str, Dict]]:
computation_name = (self.builder.caching.module_prefix +
"computation" +
self.builder.caching.module_postfix + ".py")
ignore_np_errstate = self.builder.options.backend_opts.get(
"ignore_np_errstate", True)
source = NpirCodegen.apply(self.npir,
ignore_np_errstate=ignore_np_errstate)
if self.builder.options.format_source:
source = format_source("python", source)
return {computation_name: source}
def generate_computation(self) -> Dict[str, Union[str, Dict]]:
computation_name = (self.builder.caching.module_prefix +
"computation" +
self.builder.caching.module_postfix + ".py")
return {
computation_name:
format_source(
"python",
NpirGen.apply(self.npir,
field_extents=compute_legacy_extents(
self.builder.gtir)),
),
}
def apply(cls, stencil_ir: gtir.Stencil, sdfg: dace.SDFG):
self = cls()
with dace.config.temporary_config():
dace.config.Config.set("compiler",
"cuda",
"max_concurrent_streams",
value=-1)
dace.config.Config.set("compiler",
"cpu",
"openmp_sections",
value=False)
code_objects = sdfg.generate_code()
is_gpu = "CUDA" in {co.title for co in code_objects}
computations = cls._postprocess_dace_code(code_objects, is_gpu)
interface = cls.template.definition.render(
name=sdfg.name,
dace_args=self.generate_dace_args(stencil_ir, sdfg),
functor_args=self.generate_functor_args(sdfg),
tmp_allocs=self.generate_tmp_allocs(sdfg),
allocator="gt::cuda_util::cuda_malloc"
if is_gpu else "std::make_unique",
)
generated_code = f"""#include <gridtools/sid/sid_shift_origin.hpp>
#include <gridtools/sid/allocator.hpp>
#include <gridtools/stencil/cartesian.hpp>
{"#include <gridtools/common/cuda_util.hpp>" if is_gpu else ""}
namespace gt = gridtools;
{computations}
{interface}
"""
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def apply(cls, root, **kwargs) -> str:
generated_code = super().apply(root, **kwargs)
formatted_code = codegen.format_source("cpp",
generated_code,
style="LLVM")
return formatted_code
def apply(cls, node: oir.HorizontalExecution, **kwargs: Any) -> str:
if not isinstance(node, oir.HorizontalExecution):
raise ValueError("apply() requires oir.HorizontalExecution node")
generated_code = super().apply(node)
formatted_code = codegen.format_source("python", generated_code)
return formatted_code
