diff --git a/src/pystencils_autodiff/walberla.py b/src/pystencils_autodiff/walberla.py index 894feff0119b8e687cc9730da800c8e6c966951b..cf2839d35ea5b46fc01d45252bb2adf405f9f059 100644 --- a/src/pystencils_autodiff/walberla.py +++ b/src/pystencils_autodiff/walberla.py @@ -526,8 +526,7 @@ else { } {%- else %} timeloop.run(); -{%- endif %} - """) # noqa +{%- endif %}""") # noqa def __init__(self, block_forest, timeloop, with_gui=False, use_gui_default='false'): self.density_name = "DensityAdaptor" @@ -684,7 +683,7 @@ class InitBoundaryHandling(JinjaCppFile): {% endfor %} """) # noqa - def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions): + def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions, boundary_kernel: dict, field_allocations): self.fluid = FlagUidDefinition("fluid") ast_dict = {'fluid_uid_definition': self.fluid, 'geometry_initialization': BoundaryHandlingFromConfig(block_forest, @@ -694,20 +693,51 @@ class InitBoundaryHandling(JinjaCppFile): b, pdf_field_id, flag_field_id, - self.fluid.symbol) + self.fluid.symbol, + boundary_kernel[b.name], + field_allocations) for b in boundary_conditions] } super().__init__(ast_dict) headers = ['"cuda/FieldCopy.h"', '"geometry/InitBoundaryHandling.h"'] + @property + def undefined_symbols(self): + rtn = super().undefined_symbols + for b in self.ast_dict.generated_boundaries: + rtn = rtn | b.undefined_symbols + return rtn class GeneratedBoundaryInitialization(JinjaCppFile): - TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }} ); + TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }}{{ parameter_str }} ); {{ identifier }}.fillFromFlagField<FlagField_T>( {{ block_forest }}, {{ flag_field_id }}, FlagUID("{{ boundary_condition }}"), {{ fluid_uid }} ); """) # noqa - def __init__(self, block_forest, boundary_condition, pdf_field_id, flag_field_id, fluid_uid): + def __init__(self, + block_forest, + boundary_condition, + pdf_field_id, + flag_field_id, + fluid_uid, + kernel, + field_allocations): + def resolve_parameter(p): + if kernel.target == 'cpu': + dict = field_allocations._cpu_allocations + else: + dict = field_allocations._gpu_allocations + + return dict.get(p.symbol.name.replace('_data_', ''), p).symbol + parameters = kernel.get_parameters() + parameter_ids = [resolve_parameter(p) + for p in parameters + if (p.is_field_pointer or not p.is_field_parameter) and + p.symbol.name not in ('_data_indexVector', '_data_pdfs', 'indexVectorSize',)] + parameter_str = ', '.join(p.name for p in parameter_ids) + if parameter_str: + parameter_str = ', ' + parameter_str + self.fluid = FlagUidDefinition("fluid") ast_dict = {'block_forest': block_forest, 'boundary_condition': pascalcase(boundary_condition.name), @@ -715,6 +745,8 @@ class GeneratedBoundaryInitialization(JinjaCppFile): 'pdf_field_id': pdf_field_id, 'fluid_uid': fluid_uid, 'flag_field_id': flag_field_id, + 'parameter_ids': parameter_ids, + 'parameter_str': parameter_str } super().__init__(ast_dict) @@ -728,28 +760,41 @@ class GeneratedBoundaryInitialization(JinjaCppFile): # TypedSymbol(self.ast_dict.identifier, 'FlagUID'), return {TypedSymbol(self.ast_dict.identifier, f'lbm::{self.ast_dict.boundary_condition}')} + @property + def undefined_symbols(self): + return super().undefined_symbols | set(self.ast_dict.parameter_ids) + class SweepCreation(JinjaCppFile): TEMPLATE = jinja2.Template("""{{ sweep_class_name }}( {{ parameter_str }} )""") # noqa def __init__(self, sweep_class_name: str, field_allocation: AllocateAllFields, ast, parameters_to_ignore=None): + def resolve_parameter(p): + if ast.target == 'cpu': + dict = field_allocation._cpu_allocations + else: + dict = field_allocation._gpu_allocations + + return dict.get(p.symbol.name.replace('_data_', ''), p).symbol + parameters = ast.get_parameters() - parameter_ids = [field_allocation._cpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name - if ast.target == 'cpu' - else field_allocation._gpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name + parameter_ids = [resolve_parameter(p) for p in parameters - if p.is_field_pointer or not p.is_field_parameter - ] + if p.is_field_pointer or not p.is_field_parameter] ast_dict = {'sweep_class_name': sweep_class_name, 'parameter_ids': parameter_ids, - 'parameter_str': ', '.join(parameter_ids)} + 'parameter_str': ', '.join(p.name for p in parameter_ids)} super().__init__(ast_dict) @property def headers(self): return [f'"{self.ast_dict.sweep_class_name}.h"'] + @property + def undefined_symbols(self): + return set(self.ast_dict.parameter_ids) + class SweepOverAllBlocks(JinjaCppFile): # TEMPLATE = jinja2.Template("""std::for_each({{block_forest}}->begin(), {{block_forest}}->end(), {{functor}});""") # noqa diff --git a/src/pystencils_autodiff/wald_und_wiesen_simulation.py b/src/pystencils_autodiff/wald_und_wiesen_simulation.py index 0223c348b12c438fba44fd656c4b01899ac7ffab..662a6d0f0d04dcee9c46645f1913f67eb99bf38a 100644 --- a/src/pystencils_autodiff/wald_und_wiesen_simulation.py +++ b/src/pystencils_autodiff/wald_und_wiesen_simulation.py @@ -48,6 +48,7 @@ class WaldUndWiesenSimulation(): self._kernel_class_generator = WaldUndWiesenSimulation._get_sweep_class_name() self._with_gui = False self._with_gui_default = False + self._boundary_kernels = {} def _create_helper_files(self) -> Dict[str, str]: if self._lb_rule: @@ -56,8 +57,8 @@ class WaldUndWiesenSimulation(): refinement_scaling=self._refinement_scaling) if self._boundary_handling: for bc in self.boundary_conditions: - lbmpy_walberla.generate_boundary( - self._codegen_context, pascalcase(bc.name), bc, self._lb_rule.method) + self._boundary_kernels.update({bc.name: lbmpy_walberla.generate_boundary( + self._codegen_context, pascalcase(bc.name), bc, self._lb_rule.method)}) def _create_module(self): if self._lb_rule: @@ -87,8 +88,12 @@ class WaldUndWiesenSimulation(): ResolveUndefinedSymbols( Block([ field_allocations, - InitBoundaryHandling(self._block_forest.blocks, flag_field_id, - pdf_field_id, self.boundary_conditions) + InitBoundaryHandling(self._block_forest.blocks, + flag_field_id, + pdf_field_id, + self.boundary_conditions, + self._boundary_kernels, + self._field_allocations) if self._boundary_handling else EmptyLine(), LbCommunicationSetup(self._lb_model_name, pdf_field_id) diff --git a/tests/test_walberla.py b/tests/test_walberla.py index 162dea9bcf6437702d13f18aa206a20e70b061f5..a9954b779baf1969b22f5810b185de2e1232555a 100644 --- a/tests/test_walberla.py +++ b/tests/test_walberla.py @@ -78,7 +78,10 @@ def test_wald_wiesen_lbm(): with ManualCodeGenerationContext() as ctx: from test_graph_datahandling import ldc_setup opt_params = {'target': 'gpu'} - lbm_step = ldc_setup(domain_size=(30, 30), optimization=opt_params, fixed_loop_sizes=False) + import sympy as sp + lid_velocity = sp.symbols('lid_velocity') + lbm_step = ldc_setup(domain_size=(30, 30), optimization=opt_params, + fixed_loop_sizes=False, lid_velocity=lid_velocity) sim = WaldUndWiesenSimulation(lbm_step.data_handling, ctx, @@ -108,4 +111,3 @@ def test_resolve_parameters(): ]))) print(module) -