From 2a61063dbeb50b1d5f37d8055c3afb3acb6acf46 Mon Sep 17 00:00:00 2001
From: Michael Kuron <mkuron@icp.uni-stuttgart.de>
Date: Tue, 21 Apr 2020 16:22:32 +0200
Subject: [PATCH] pystencils_walberla: staggered boundary generation
---
python/pystencils_walberla/__init__.py | 3 +-
python/pystencils_walberla/boundary.py | 108 ++++++++++++++++++++++++-
2 files changed, 108 insertions(+), 3 deletions(-)
diff --git a/python/pystencils_walberla/__init__.py b/python/pystencils_walberla/__init__.py
index 767581a75..88eb440a6 100644
--- a/python/pystencils_walberla/__init__.py
+++ b/python/pystencils_walberla/__init__.py
@@ -1,3 +1,4 @@
+from .boundary import generate_staggered_boundary, generate_staggered_flux_boundary
from .cmake_integration import CodeGeneration
from .codegen import (
generate_pack_info, generate_pack_info_for_field, generate_pack_info_from_kernel,
@@ -5,4 +6,4 @@ from .codegen import (
__all__ = ['CodeGeneration',
'generate_sweep', 'generate_pack_info_from_kernel', 'generate_pack_info_for_field', 'generate_pack_info',
- 'generate_mpidtype_info_from_kernel']
+ 'generate_mpidtype_info_from_kernel', 'generate_staggered_boundary', 'generate_staggered_flux_boundary']
diff --git a/python/pystencils_walberla/boundary.py b/python/pystencils_walberla/boundary.py
index 79b7ed28a..00346cfd6 100644
--- a/python/pystencils_walberla/boundary.py
+++ b/python/pystencils_walberla/boundary.py
@@ -1,16 +1,120 @@
import numpy as np
from jinja2 import Environment, PackageLoader, StrictUndefined
-from pystencils_walberla.codegen import KernelInfo
from pystencils import Field, FieldType
+from pystencils.boundaries.boundaryhandling import create_boundary_kernel
from pystencils.boundaries.createindexlist import (
boundary_index_array_coordinate_names, direction_member_name,
numpy_data_type_for_boundary_object)
from pystencils.data_types import TypedSymbol, create_type
-from pystencils_walberla.codegen import default_create_kernel_parameters
+from pystencils_walberla.codegen import KernelInfo
from pystencils_walberla.jinja_filters import add_pystencils_filters_to_jinja_env
+def generate_staggered_boundary(generation_context, class_name, boundary_object,
+ dim, neighbor_stencil, index_shape, target='cpu'):
+ struct_name = "IndexInfo"
+ boundary_object.name = class_name
+
+ index_struct_dtype = numpy_data_type_for_boundary_object(boundary_object, dim)
+
+ staggered_field = Field.create_generic('field', dim,
+ np.float64 if generation_context.double_accuracy else np.float32,
+ index_dimensions=len(index_shape), layout='c', index_shape=index_shape,
+ field_type=FieldType.STAGGERED)
+
+ index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
+ shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
+
+ kernel = create_boundary_kernel(staggered_field, index_field, neighbor_stencil, boundary_object, target=target,
+ openmp=generation_context.openmp)
+ kernel.function_name = "boundary_" + boundary_object.name
+
+ # waLBerla is a 3D framework. Therefore, a zero for the z index has to be added if we work in 2D
+ if dim == 2:
+ stencil = ()
+ for d in neighbor_stencil:
+ d = d + (0,)
+ stencil = stencil + (d,)
+ else:
+ stencil = neighbor_stencil
+
+ stencil_info = [(i, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
+
+ context = {
+ 'class_name': boundary_object.name,
+ 'StructName': struct_name,
+ 'StructDeclaration': struct_from_numpy_dtype(struct_name, index_struct_dtype),
+ 'kernel': KernelInfo(kernel),
+ 'stencil_info': stencil_info,
+ 'dim': dim,
+ 'target': target,
+ 'namespace': 'pystencils',
+ }
+
+ env = Environment(loader=PackageLoader('pystencils_walberla'), undefined=StrictUndefined)
+ add_pystencils_filters_to_jinja_env(env)
+
+ header = env.get_template('Boundary.tmpl.h').render(**context)
+ source = env.get_template('Boundary.tmpl.cpp').render(**context)
+
+ source_extension = "cpp" if target == "cpu" else "cu"
+ generation_context.write_file("{}.h".format(class_name), header)
+ generation_context.write_file("{}.{}".format(class_name, source_extension), source)
+
+
+def generate_staggered_flux_boundary(generation_context, class_name, boundary_object,
+ dim, neighbor_stencil, index_shape, target='cpu'):
+ struct_name = "IndexInfo"
+ boundary_object.name = class_name
+
+ index_struct_dtype = numpy_data_type_for_boundary_object(boundary_object, dim)
+
+ staggered_field = Field.create_generic('flux', dim,
+ np.float64 if generation_context.double_accuracy else np.float32,
+ index_dimensions=len(index_shape), layout='c', index_shape=index_shape,
+ field_type=FieldType.STAGGERED_FLUX)
+
+ index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
+ shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
+
+ kernel = create_boundary_kernel(staggered_field, index_field, neighbor_stencil, boundary_object, target=target,
+ openmp=generation_context.openmp)
+ kernel.function_name = "boundary_" + boundary_object.name
+
+ # waLBerla is a 3D framework. Therefore, a zero for the z index has to be added if we work in 2D
+ if dim == 2:
+ stencil = ()
+ for d in neighbor_stencil:
+ d = d + (0,)
+ stencil = stencil + (d,)
+ else:
+ stencil = neighbor_stencil
+
+ stencil_info = [(i, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
+
+ context = {
+ 'class_name': boundary_object.name,
+ 'StructName': struct_name,
+ 'StructDeclaration': struct_from_numpy_dtype(struct_name, index_struct_dtype),
+ 'kernel': KernelInfo(kernel),
+ 'stencil_info': stencil_info,
+ 'dim': dim,
+ 'target': target,
+ 'namespace': 'pystencils',
+ }
+
+ env = Environment(loader=PackageLoader('pystencils_walberla'), undefined=StrictUndefined)
+ add_pystencils_filters_to_jinja_env(env)
+
+ header = env.get_template('Boundary.tmpl.h').render(**context)
+ source = env.get_template('Boundary.tmpl.cpp').render(**context)
+
+ source_extension = "cpp" if target == "cpu" else "cu"
+ generation_context.write_file("{}.h".format(class_name), header)
+ generation_context.write_file("{}.{}".format(class_name, source_extension), source)
+
+
def struct_from_numpy_dtype(struct_name, numpy_dtype):
result = "struct %s { \n" % (struct_name,)
--
GitLab